Originally published on Code Beneath.
When you shut down a network service, you cannot simply close the listening socket and exit. Clients might have in-flight requests, half-written buffers, or active reads waiting for your response. A golang TCP socket graceful shutdown example requires coordinating context cancellation, draining pending connections, and enforcing timeouts so no data disappears. This is not theoretical: lose proper shutdown and you will drop customer transactions, corrupt state, and make debugging hell.
The Problem: Why Abrupt Shutdown Fails
A naive shutdown looks like this: close the listener, exit the main function, the OS kills all goroutines, sockets close hard. Any client mid-read gets an EOF. Any server trying to write gets a broken pipe error. Uncommitted state, incomplete transactions, nothing graceful about it.
The correct pattern is: stop accepting new connections, give existing connections time to finish, close those connections cleanly, only then exit. TCP and UDP sockets have different shutdown semantics, and TCP requires a coordinated close where both sides send FIN. Go's context package and sync.WaitGroup are the tools to choreograph this.
Architecture: Context, WaitGroup, and Timeouts
The pattern uses three layers. First, a context tied to the service lifecycle. Second, a WaitGroup that tracks active connection handlers. Third, a shutdown timeout that forces closure if handlers don't finish fast enough. When shutdown is triggered, you close the listener (which unblocks Accept calls), drain remaining handlers by waiting on the WaitGroup, then close sockets with a deadline.
Here is a complete server and client demonstrating proper golang TCP socket graceful shutdown:
package main
import (
"bufio"
"context"
"fmt"
"io"
"log"
"net"
"os"
"os/signal"
"sync"
"syscall"
"time"
)
// Server holds the listener and connection tracking.
type Server struct {
listener net.Listener
wg sync.WaitGroup
ctx context.Context
cancel context.CancelFunc
}
// NewServer creates a server bound to addr.
func NewServer(addr string) (*Server, error) {
listener, err := net.Listen("tcp", addr)
if err != nil {
return nil, err
}
ctx, cancel := context.WithCancel(context.Background())
return &Server{
listener: listener,
ctx: ctx,
cancel: cancel,
}, nil
}
// Start accepts connections in a goroutine.
func (s *Server) Start() {
s.wg.Add(1)
go func() {
defer s.wg.Done()
for {
conn, err := s.listener.Accept()
if err != nil {
// Listener closed or accept failed.
select {
case <-s.ctx.Done():
// Shutdown in progress, exit cleanly.
return
default:
log.Printf("accept error: %v\n", err)
return
}
}
// Handle each connection in its own goroutine.
s.wg.Add(1)
go s.handleConnection(conn)
}
}()
}
// handleConnection reads a message, processes it, and sends a response.
func (s *Server) handleConnection(conn net.Conn) {
defer s.wg.Done()
defer conn.Close()
// Wrap the connection with a reader and writer.
reader := bufio.NewReader(conn)
writer := bufio.NewWriter(conn)
for {
// Set read deadline using context or a fixed timeout.
// This prevents reads from blocking forever during shutdown.
deadline, ok := s.ctx.Deadline()
if !ok {
// No deadline on context, use a reasonable read timeout.
deadline = time.Now().Add(10 * time.Second)
}
conn.SetReadDeadline(deadline)
// Read a line from the client.
line, err := reader.ReadString('\n')
if err != nil {
if err == io.EOF {
// Client closed cleanly.
return
}
// Check if this is a deadline exceeded error during shutdown.
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
select {
case <-s.ctx.Done():
// Context canceled, shutdown in progress.
return
default:
log.Printf("read timeout: %v\n", err)
return
}
}
log.Printf("read error: %v\n", err)
return
}
// Check if shutdown was initiated.
select {
case <-s.ctx.Done():
// Write a goodbye message and close.
fmt.Fprintf(writer, "Server shutting down\n")
writer.Flush()
return
default:
}
// Echo the message back with a timestamp.
response := fmt.Sprintf("Echo: %s received at %s\n", line[:len(line)-1], time.Now().Format(time.RFC3339))
if _, err := fmt.Fprint(writer, response); err != nil {
log.Printf("write error: %v\n", err)
return
}
if err := writer.Flush(); err != nil {
log.Printf("flush error: %v\n", err)
return
}
}
}
// Shutdown gracefully closes the server.
// It stops accepting connections, waits for existing handlers to finish,
// and enforces a timeout to prevent indefinite hangs.
func (s *Server) Shutdown(timeout time.Duration) error {
log.Println("initiating graceful shutdown...")
// Close the listener. This unblocks Accept() calls.
if err := s.listener.Close(); err != nil {
log.Printf("listener close error: %v\n", err)
}
// Signal all handlers that shutdown is in progress.
s.cancel()
// Wait for all handlers to finish, with a timeout.
done := make(chan struct{})
go func() {
s.wg.Wait()
close(done)
}()
select {
case <-done:
log.Println("all connections closed")
return nil
case <-time.After(timeout):
log.Printf("shutdown timeout after %v\n", timeout)
return fmt.Errorf("shutdown timeout")
}
}
// Client connects to the server and sends messages.
func runClient(addr string, messages []string) error {
conn, err := net.Dial("tcp", addr)
if err != nil {
return err
}
defer conn.Close()
writer := bufio.NewWriter(conn)
reader := bufio.NewReader(conn)
for _, msg := range messages {
// Send message.
if _, err := fmt.Fprintf(writer, "%s\n", msg); err != nil {
return err
}
if err := writer.Flush(); err != nil {
return err
}
// Read response with a timeout.
conn.SetReadDeadline(time.Now().Add(5 * time.Second))
response, err := reader.ReadString('\n')
if err != nil {
return err
}
fmt.Print(response)
time.Sleep(500 * time.Millisecond)
}
return nil
}
func main() {
// Start server.
server, err := NewServer("127.0.0.1:9999")
if err != nil {
log.Fatalf("failed to create server: %v\n", err)
}
server.Start()
log.Println("server listening on 127.0.0.1:9999")
// If an argument is passed, run as client. Otherwise run as server.
if len(os.Args) > 1 && os.Args[1] == "client" {
time.Sleep(1 * time.Second) // Give server time to start.
err := runClient("127.0.0.1:9999", []string{"hello", "world", "test"})
if err != nil {
log.Fatalf("client error: %v\n", err)
}
return
}
// Server: wait for interrupt signal.
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
<-sigChan
// Shutdown with a 10-second timeout.
if err := server.Shutdown(10 * time.Second); err != nil {
log.Printf("shutdown error: %v\n", err)
}
}
To run this code:
go run main.go &
sleep 0.5
go run main.go client
kill %1
In another terminal, press Ctrl+C to trigger the graceful shutdown. The server closes the listener, prevents new accepts, waits for the current connection handler to finish reading and writing, then exits cleanly.
Why This Works: The Three Phases
Phase 1: Stop accepting. Close the listener. The Accept loop receives an error and exits if the context is already canceled, or it checks the context on the next iteration.
Phase 2: Signal handlers. Call cancel() on the context. Every goroutine listening on the context exits cleanly. Handlers that are mid-read will eventually see the context done, then they break the read loop and defer close the connection.
Phase 3: Wait with timeout. Use a WaitGroup to block until all handlers are done. If they are not done after the timeout, the shutdown returns an error and the process exits anyway (because main exits), but at least we gave handlers a chance.
Key Details and Edge Cases
Use SetReadDeadline and SetWriteDeadline on the connection. If you leave reads and writes undeadlined, a handler can hang forever waiting for a client, and your shutdown will block forever waiting for that handler. The deadline forces the read to unblock, you check the context, and you exit if shutdown is in progress.
Check the context immediately after returning from a blocking I/O call. A read might timeout naturally (not due to shutdown). You need to distinguish: did the timeout happen because shutdown was signaled, or did the client just go silent? If the context is done, exit. If not, it might be a transient issue; log and return.
Use sync.WaitGroup carefully. Add to the group before starting a goroutine, not after. Use defer Done() so the count decreases even if a goroutine panics. Forget to Add or Done and you will have goroutine leaks or a panic in Wait.
Buffered I/O (bufio.Writer) requires explicit Flush before the connection closes. A writer might hold data in a 4096-byte buffer. If you close the connection without flushing, the last chunk of the response is lost. Always defer Flush or call it after the final Write in critical sections.
Go's goroutine model makes this pattern clean because goroutines are cheap and the select statement is built for multiplexing context. In languages without goroutines (Java threads are heavier, C requires manual event loops), this requires more infrastructure.
Common Mistakes
Mistake 1: Closing the listener but not canceling the context. The Accept loop exits because Accept fails, but handlers still run. You wait on a WaitGroup that never decrements because no handler checks for shutdown. Result: deadlock on Shutdown.
Mistake 2: Not setting read/write deadlines. A handler blocked on a read never sees the context. The WaitGroup.Wait() call hangs forever. You hit the shutdown timeout and kill the process, dropping the in-flight data from that handler.
Mistake 3: Forgetting to flush buffered writers before closing. The response is never sent. The client sees EOF mid-response and logs an error. The data is lost even though the handler thought it succeeded.
Mistake 4: Too short a shutdown timeout. If you set timeout to 1 millisecond, a handler doing any non-trivial work (database query, file I/O) cannot finish. You force-close and lose data. Start with 10-30 seconds and tune based on your handler latency.
Comparing Approaches
| Approach | Pros | Cons |
|---|---|---|
| Abrupt close (no shutdown) | Simple, no coordination needed | Data loss, client errors, no cleanup |
| Close listener, immediate exit | Faster than waiting, some handlers finish naturally | Pending handlers still run but may be interrupted |
| Context cancel with WaitGroup and timeout (recommended) | Clean signal, handlers have time to finish, timeout prevents hangs | More code, requires discipline on deadlines |
| HTTP server with Server.Shutdown() | Built-in, handles connection draining, request context automatically canceled | Limited to HTTP, abstracts the mechanics |
For raw TCP services, you implement the pattern yourself. For HTTP, use the standard library's Server.Shutdown() method, which does all this under the hood.
Testing Graceful Shutdown
Write a test that connects to the server, starts a slow operation (sleep for a second), then sends a shutdown signal. Verify the slow operation completes (no data loss). Then test with a very short timeout: verify the shutdown function returns a timeout error but the process exits.
func TestGracefulShutdown(t *testing.T) {
server, err := NewServer("127.0.0.1:0") // Port 0 = OS chooses
if err != nil {
t.Fatalf("failed to create server: %v", err)
}
server.Start()
addr := server.listener.Addr().String()
// Connect and send a message.
go func() {
conn, _ := net.Dial("tcp", addr)
defer conn.Close()
fmt.Fprintf(conn, "test\n")
io.ReadAll(conn) // Read until EOF
}()
time.Sleep(100 * time.Millisecond)
// Shutdown should complete without timeout.
err = server.Shutdown(5 * time.Second)
if err != nil {
t.Errorf("shutdown failed: %v", err)
}
}
FAQ
What if a handler is stuck in an infinite loop and never checks the context?
The shutdown timeout expires and you force-exit. This is a bug in your handler code, not the shutdown mechanism. Every handler must have checkpoints where it checks the context or has a deadline. Handlers doing long-running compute (no I/O) should explicitly check the context periodically using select <-ctx.Done().
Can I use context.WithTimeout instead of context.WithCancel?
Yes, if you know exactly how long the service should run. But WithCancel is cleaner for shutdown because the timeout is not about the service lifetime, it is about the shutdown grace period. Use WithCancel for the service, then pass a separate timeout to the Shutdown function.
Do I need a WaitGroup if I am not spawning many goroutines?
No, but do not skip it. Even if you think your service only handles a few connections, adding WaitGroup costs nothing and prevents bugs when load increases. It is defensive programming.
What happens if a client reconnects after I close the listener?
If the listener is closed, new connection attempts will fail with "connection refused". Existing connections are not affected until you close them. This is correct behavior: you stop accepting new work but let current work finish.
Should I set deadlines on the server side or let the client timeout?
Set them on both. Server-side deadlines prevent handlers from hanging indefinitely, especially during shutdown. Client-side timeouts are for client resilience. If either one fires, the connection is closed and resources are freed.
Does this pattern work for UDP?
UDP is connectionless, so there is no connection to gracefully close. You close the socket and any pending receives return an error immediately. The pattern is simpler: just cancel the context and exit the receive loop.
How do I know if my shutdown was successful?
Check the error return from Shutdown. If it is nil, all handlers finished before the timeout. If it is a timeout error, handlers were force-closed. Log this and monitor it in production. A persistent timeout error might indicate handlers are not checking the context properly.
Top comments (0)