Build a concurrent TCP server in Go with graceful shutdown include unit tests

Concurrent TCP server in Go with graceful shutdown and unit tests

TCP (Transmission Control Protocol) is a protocol that provides reliable, ordered, and error-checked delivery of data between applications running on hosts communicating over an IP network. A TCP server is a program that listens for incoming TCP connections and handles them by providing services or sending data to the client. In this article, we will discuss how to implement a TCP server in Golang that is concurrent and supports graceful shutdown. We will also provide an example of a unit test for this implementation.

If you’re new to socket programming in Go, you may want to check out my previous article on the basics of socket programming in Go, which covers the fundamental concepts and includes some sample code.

TCP Server Implementation

The TCP server implementation consists of the following steps:

1. Create a new net.Listener object that listens for incoming connections on a specified address using the net.Listen function.
2. Start two goroutines to handle incoming connections concurrently: one to accept new connections and another to handle them.
3. In the acceptConnections goroutine, use a for loop and a select statement to listen for incoming connections on the listener and send them over a channel.
4. In the handleConnections goroutine, use a for loop and a select statement to receive connections from the channel and handle them in separate goroutines.
5. In the handleConnection function, handle the incoming connection by performing any necessary processing or sending data to the client.
6. Implement graceful shutdown by creating a shutdown channel that signals to the goroutines that they should stop processing connections. Close the listener and wait for the goroutines to finish using a sync.WaitGroup.

Here is the complete code for the TCP server implementation:

package main

import (
 "fmt"
 "net"
 "os"
 "os/signal"
 "sync"
 "syscall"
 "time"
)

type server struct {
 wg         sync.WaitGroup
 listener   net.Listener
 shutdown   chan struct{}
 connection chan net.Conn
}

func newServer(address string) (*server, error) {
 listener, err := net.Listen("tcp", address)
 if err != nil {
  return nil, fmt.Errorf("failed to listen on address %s: %w", address, err)
 }

 return &server{
  listener:   listener,
  shutdown:   make(chan struct{}),
  connection: make(chan net.Conn),
 }, nil
}

func (s *server) acceptConnections() {
 defer s.wg.Done()

 for {
  select {
  case <-s.shutdown:
   return
  default:
   conn, err := s.listener.Accept()
   if err != nil {
    continue
   }
   s.connection <- conn
  }
 }
}

func (s *server) handleConnections() {
 defer s.wg.Done()

 for {
  select {
  case <-s.shutdown:
   return
  case conn := <-s.connection:
   go s.handleConnection(conn)
  }
 }
}

func (s *server) handleConnection(conn net.Conn) {
 defer conn.Close()

 // Add your logic for handling incoming connections here
 fmt.Fprintf(conn, "Welcome to my TCP server!\n")
 time.Sleep(5 * time.Second)
 fmt.Fprintf(conn, "Goodbye!\n")
}

func (s *server) Start() {
 s.wg.Add(2)
 go s.acceptConnections()
 go s.handleConnections()
}

func (s *server) Stop() {
 close(s.shutdown)
 s.listener.Close()

 done := make(chan struct{})
 go func() {
  s.wg.Wait()
  close(done)
 }()

 select {
 case <-done:
  return
 case <-time.After(time.Second):
  fmt.Println("Timed out waiting for connections to finish.")
  return
 }
}

func main() {
 s, err := newServer(":8080")
 if err != nil {
  fmt.Println(err)
  os.Exit(1)
 }

 s.Start()

 // Wait for a SIGINT or SIGTERM signal to gracefully shut down the server
 sigChan := make(chan os.Signal, 1)
 signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 <-sigChan

 fmt.Println("Shutting down server...")
 s.Stop()
 fmt.Println("Server stopped.")
}

In this example, the server creates two goroutines to handle incoming connections and manage them concurrently. The acceptConnections method listens for new connections and sends them over a channel, while the handleConnections method receives these connections and handles them in separate goroutines.

The server also implements graceful shutdown by using a shutdown channel to signal to the goroutines that they should stop processing connections, and a WaitGroup to wait for them to finish before closing the listener and terminating the program.

To start the server, we create a new instance of the server struct and call its Start method. To stop the server, we wait for a SIGINT or SIGTERM signal, and then call the Stop method.

Note that this code is a simple example, and you may need to modify it to suit your specific use case.

Unit Test Implementation

To test the TCP server implementation, we will use the standard Golang testing framework. We will create a test case that starts the server, connects to it using a TCP client, sends a message, and verifies that the server received the message.

Here is the code for the unit test implementation:

package main

import (
 "net"
 "testing"
 "time"
)

func TestServer(t *testing.T) {
 // Start the server
 s, err := newServer(":8080")
 if err != nil {
  t.Fatal(err)
 }
 s.Start()

 // Connect to the server and send a message
 conn, err := net.Dial("tcp", "localhost:8080")
 if err != nil {
  t.Fatal(err)
 }
 defer conn.Close()

 expected := "Welcome to my TCP server!\n"
 actual := make([]byte, len(expected))
 if _, err := conn.Read(actual); err != nil {
  t.Fatal(err)
 }
 if string(actual) != expected {
  t.Errorf("expected %q, but got %q", expected, actual)
 }

 // Wait for the server to handle the connection
 time.Sleep(6 * time.Second)

 expected = "Goodbye!\n"
 actual = make([]byte, len(expected))
 if _, err := conn.Read(actual); err != nil {
  t.Fatal(err)
 }
 if string(actual) != expected {
  t.Errorf("expected %q, but got %q", expected, actual)
 }

 // Stop the server
 s.Stop()
}

In the TestTCP function, we start the TCP server using the newServer function and connect to it using a TCP client using the net.Dial function. We then send a message to the server using fmt.Fprintf and verify that the server received the message by reading from the connection and checking the response using an if statement.

To run the test, we can use the go test command in the terminal:

$ go test
PASS
ok   command-line-arguments 0.003s

Conclusion

In this article, we discussed how to implement a TCP server in Golang that is concurrent and supports graceful shutdown. We also provided an example of a unit test for this implementation. The implementation uses goroutines to handle incoming connections and gracefully shutdown, and the unit test verifies that the server can handle incoming connections and messages correctly. This implementation can be used as a starting point for building more complex TCP servers in Golang.