Simple Tcp Server

Use a simple TCP server to highlight the extra code that can be added to make a process better.

%%{init: {"flowchart": {"htmlLabels": false}} }%%
flowchart LR
    main("MAIN")

    main --> echoListener("echo server\n:2007")
    echoListener --> echoHandler1("echo client")
    echoListener --> echoHandler2("echo client")

    main --> daytimeListener("daytime server\n:2013")
    daytimeListener --> daytimeHandler1("daytime client")
    daytimeListener --> daytimeHandler2("daytime client")

This server is a TCP server that implements two common TCP services. The echo service (historically on port 7) and the daytime service (historically on port 13). The services will run concurrently. There is a minimal implementation of the server in the Minimal Server page.

  • Refactor into separate packages to provide some isolation.
  • Refactor the listener to make it more general. (DRY, function values)
  • Control-C handling to start cleanup. (signal, channels and context)
  • Add wait groups to allow the server to wait for all services to complete. (sync)
  • Add metrics. (expvar)
  • Add a web server to expose the metrics.
  • Add a logging module to roll logs (lumberjack and vendoring)
  • Signal to roll the logs.
  • Tests
  • Add information during build (Makefile, git, go link commands)
  • Move code to our own module (ex. testing code ioBuffer)

Subsections of Simple Tcp Server

Minimal Server

The purpose of this code is to highlight what is needed to create a TCP server in Go.

I have implemented a simple tcp server without any error checking, logging, cleanup, or anything extra. There are two TCP servers in a single program. An echo server listening on port 2007 run in the background and a daytime server listening on port 2013 run in the foreground.

package main

import (
	"io"
	"net"
	"time"
)

func main() {
	// the echo server is put in the background
	go func() {
		echoListener, _ := net.Listen("tcp", ":2007")
		for {
			connection, _ := echoListener.Accept()
			go func(c net.Conn) {
				io.Copy(c, c)
				c.Close()
			}(connection)
		}
	}()
	// the daytime server is in the foreground so we don't exit early.
	daytimeListener, _ := net.Listen("tcp", ":2013")
	for {
		connection, _ := daytimeListener.Accept()
		go func(c net.Conn) {
			c.Write([]byte(time.Now().Format("Monday, January 2, 2006 15:04:05-MST\r\n")))
			c.Close()
		}(connection)
	}
}

Testing

package echo

import (
	"bytes"
	"io"
	"testing"
)

type ioBuffer struct {
	// I need a separate place to read from and write to
	in       bytes.Buffer
	out      bytes.Buffer
	isClosed bool
}

func (b *ioBuffer) Read(p []byte) (n int, err error) {
	if b.isClosed {
		return 0, io.ErrClosedPipe
	}
	return b.in.Read(p)
}

func (b *ioBuffer) Write(p []byte) (n int, err error) {
	if b.isClosed {
		return 0, io.ErrClosedPipe
	}
	return b.out.Write(p)
}

func (b *ioBuffer) Close() error {
	b.isClosed = true
	return nil
}

func TestHandler(t *testing.T) {
	connection := &ioBuffer{in: *bytes.NewBufferString("testing")}
	want := "testing"

	err := Handler(connection)
	if err != nil {
		t.Fatalf("Handler failed: %s", err.Error())
	}
	if connection.out.String() != want {
		t.Errorf(`got %q, want %q`, connection.out.String(), want)
	}
	if !connection.isClosed {
		t.Error("Did not close the connection")
	}
}

func TestHandlerErr(t *testing.T) {
	connection := &ioBuffer{in: *bytes.NewBufferString("testing")}
	connection.Close()
	err := Handler(connection)
	if err != io.ErrClosedPipe {
		t.Fatalf("Handler should fail with closed pipe. got: %v\n", err)
	}
}
package daytime

import (
	"bytes"
	"io"
	"testing"
	"time"
)

type ioBuffer struct {
	// I need a separate place to read from and write to
	in       bytes.Buffer
	out      bytes.Buffer
	isClosed bool
}

func (b *ioBuffer) Read(p []byte) (n int, err error) {
	if b.isClosed {
		return 0, io.ErrClosedPipe
	}
	return b.in.Read(p)
}

func (b *ioBuffer) Write(p []byte) (n int, err error) {
	if b.isClosed {
		return 0, io.ErrClosedPipe
	}
	return b.out.Write(p)
}

func (b *ioBuffer) Close() error {
	b.isClosed = true
	return nil
}

func TestHandler(t *testing.T) {
	const format = "Monday, January 2, 2006 15:04:05-MST\r\n"
	var connection ioBuffer
	err := Handler(&connection)
	if err != nil {
		t.Fatalf("Handler failed: %s", err.Error())
	}
	// Check that we write a correctly formatted date string to the outputt
	_, err = time.Parse(format, connection.out.String())
	if err != nil {
		t.Errorf("Error: %s\n", err.Error())
	}
	// Check that we closed the output.
	if !connection.isClosed {
		t.Error("Did not close the connection")
	}
}

func TestHandlerErr(t *testing.T) {
	var connection ioBuffer
	connection.Close()
	err := Handler(&connection)
	if err != io.ErrClosedPipe {
		t.Fatalf("Handler should fail with closed pipe. got: %v\n", err)
	}
}

func Test_formatResponse(t *testing.T) {
	const format = "Monday, January 2, 2006 15:04:05-MST\r\n"
	denver, err := time.LoadLocation("America/Denver")
	if err != nil {
		t.Fatalf("Can not find location info for America/Denver: %v", err)
	}
	type args struct {
		t time.Time
	}
	tests := []struct {
		name string
		args args
		want []byte
	}{
		{
			// result should match the format string.
			name: "format-pattern",
			args: args{
				t: time.Date(2006, 1, 2, 15, 4, 5, 0, denver),
			},
			want: []byte(format),
		},
		{
			// we don't need a lot of tests but a couple ensures that  we are
			// formatting as expected.
			name: "unix-epoch",
			args: args{
				t: time.UnixMilli(0).In(denver),
			},
			want: []byte("Wednesday, December 31, 1969 17:00:00-MST\r\n"),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := formatResponse(tt.args.t); !bytes.Equal(got, tt.want) {
				t.Errorf("formatResponse() = %v, want %v", got, tt.want)
			}
		})
	}
}

Daytime Handler

package daytime

import (
	"io"
	"log"
	"time"
)

func formatResponse(t time.Time) []byte {
	return []byte(t.Format("Monday, January 2, 2006 15:04:05-MST\r\n"))
}

func Handler(connection io.ReadWriteCloser) error {
	defer connection.Close()

	_, err := connection.Write(formatResponse(time.Now()))
	if err != nil {
		log.Print("Error writing date:", err)
	}
	return err
}

Echo Handler

package echo

import (
	"io"
	"log"
)

func Handler(connection io.ReadWriteCloser) error {
	defer connection.Close()

	_, err := io.Copy(connection, connection)
	if err != nil {
		log.Print("Error in copy: ", err)
	}
	return err
}

Web Package

This code is almost a direct copy of the example code at server.Shutdown. I used a context to signal the shutdown, and I added a timeout so the shutdown couldn’t hang forever.

While the global ListenAndServe() function is more convenient than creating a server object, if you are worried about shutting it down you will likely be worried about the TLS configuration as well.

A note about the implementation: I’ve see a few comments online that the idleConnsClosed channel is not needed because the Shutdown() call doesn’t return until everything is cleaned up. That is good in sequential code but this call to Shutdown is in a goroutine. The call to ListenAndServe will exit almost immediately after Shutdown is called, well before it is complete.

 1package web
 2
 3import (
 4	"context"
 5	"log"
 6	"net/http"
 7	"sync"
 8	"time"
 9)
10
11func Run(ctx context.Context, addr string, wg *sync.WaitGroup) {
12	defer wg.Done()
13
14	srv := &http.Server{Addr: addr}
15	idleConnsClosed := make(chan struct{})
16
17	// This goroutine will Wait for ctx to be triggered
18	go func() {
19		<-ctx.Done()
20		timeout, cancel := context.WithTimeout(context.Background(), 30*time.Second)
21		defer cancel()
22		if err := srv.Shutdown(timeout); err != nil {
23			log.Printf("HTTP server Shutdown: %v", err)
24		}
25		close(idleConnsClosed)
26	}()
27
28	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
29		log.Printf("HTTP server ListenAndServe: %v", err)
30	}
31	<-idleConnsClosed
32	log.Print("HTTP Server closed")
33}

Listener Package

To make this routine more general we pass in the handler for the client to run and the metrics object to be updated. So we can cleanly exit we also pass in the context to inform us when to quit, and a wait group so we can indicate that we are done.

The listener package gets touched by a lot of the addons that we are discussing.

  • Most of the metrics get set and updated in here.
  • The Ctrl-C gets used in here
  • A chain of waiting happens here. This will wait for all client handlers to exit, and then inform our parent that we are done.
Tip

Note that the handler is defined as handler func(io.ReadWriteCloser) error. I define the connection’s type as an io.ReadWriteCloser instead of net.Conn. The handlers don’t need the extra values available on the net.Conn. I remove (actually just hide) them so that I don’t have to consider them during testing or debugging.

Plain

Take all the addons away and we return to the minimal

func Run(address string, handler func(io.ReadWriteCloser) error) {
	listener, _ := net.Listen("tcp", address)
	defer listener.Close()

	for {
		connection, _ := listener.Accept()
		spawnClientHandler(connection, handler)
	}
}

Full

The full code for the package

 1package listener
 2
 3import (
 4	"context"
 5	"expvar"
 6	"io"
 7	"log"
 8	"net"
 9	"sync"
10)
11
12func Run(
13	ctx context.Context,
14	address string,
15	handler func(io.ReadWriteCloser) error,
16	wg *sync.WaitGroup,
17	metrics *expvar.Map) {
18
19	initializeMetrics(metrics, address)
20
21	// Create group to track our client handlers. When leaving, wait for the
22	// clients then tell our parent that we are done.
23	var clientsWg sync.WaitGroup
24	defer func() {
25		clientsWg.Wait()
26		wg.Done()
27	}()
28
29	// Open the listen port. I could get errors if, for example, the port is
30	// already in use
31	listener, err := net.Listen("tcp", address)
32	if err != nil {
33		log.Printf("Error opening listen socket %q: %s", address, err.Error())
34		return
35	}
36	defer listener.Close()
37
38	// The best way to break out of the Accept is to close the socket. If the
39	// context gets canceled, we close the socket so we can cleanup.
40	go func() {
41		<-ctx.Done()
42		listener.Close()
43	}()
44
45	
46	metrics.Add("listening", 1)
47	defer metrics.Add("listening", -1)
48
49	for {
50		connection, err := listener.Accept()
51		if err != nil {
52			log.Printf("Error in accept on %q: %s", address, err)
53			return
54		}
55
56		spawnClientHandler(connection, handler, metrics, &clientsWg)
57	}
58}
59
60// Start the client.  Here I can handle all of the basic metrics
61func spawnClientHandler(
62	connection net.Conn,
63	handler func(io.ReadWriteCloser) error,
64	metrics *expvar.Map,
65	clientsWg *sync.WaitGroup) {
66
67	clientsWg.Add(1)
68	go func() {
69		defer clientsWg.Done()
70
71		metrics.Add("busy", 1)
72		metrics.Add("count", 1)
73		defer metrics.Add("busy", -1)
74
75		err := handler(connection)
76		if err != nil {
77			metrics.Add("error", 1)
78		}
79	}()
80}
81
82// make the metric counters exist
83func initializeMetrics(metrics *expvar.Map, address string) {
84	metrics.Add("busy", 0)
85	metrics.Add("count", 0)
86	metrics.Add("error", 0)
87	metrics.Add("listening", 0)
88	addr := new(expvar.String)
89	addr.Set(address)
90	metrics.Set("address", addr)
91}

Control C Handling

There are a couple of options in the signal package for handling interrupts.

  • Have the signal written to a channel.
  • Have the signal cancel a context.

I’ve used the second method for handling control-c. The context returned is from context.WithCancel(..) which is a common (standard) way to notify child processing that it can stop. The context has a channel built in that is closed when the context is canceled. Anyone waiting on that channel will block until the channel is closed. The context can be read by multiple goroutines. You can also put the channel in a select{} clause for more flexibility.

    ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
    defer stop()

This context will then be used by the listener to close its connection

    go func() {
        <-ctx.Done()
        listener.Close()
    }()

And used by the web server to have it shutdown.

    go func() {
        <-ctx.Done()
        ...
        if err := srv.Shutdown(timeout); err != nil {...}
        ...
    }()