Why My TCP Server Wouldn't Shut Down ?
Building graceful shutdown for a TCP server in Go taught me why listeners and active connections have completely different lifetime.
The Setup
So the other day , i was building a small redis recreation in Go. After getting a basic server running, I wanted to support graceful shutdown when process recieved SIGINT (Ctrl+C) .
So i start like this :-
type Server struct {
lnAddr string
ln net.Listener
}
and the main function :-
func main(){
// listen for SIGINT (CTRL+C) and
// SIGTERM - process kill
ctx, stop := signal.NotifyContext(
context.Background(),
os.Interrupt,
syscall.SIGTERM,
)
defer stop()
// wait group to wait for all the goroutines
// to finish
var wg sync.WaitGroup
//initialize a new server
svr := server.NewServer(server.SERVER_DEFAULT_PORT)
// start the server
if err := svr.Start(ctx , &wg); err != nil {
log.Fatal(err)
}
// this here blocks until the context is cancelled
<- ctx.Done()
// wait for all the goroutine to finish
wg.Wait()
}
Initially , my intuation was to pass context.Context around in anything that "might hang forever" or "needs cancellation".
So my start function looked like this :-
func (s *Server) Start(ctx context.Context, wg *sync.WaitGroup) error {
// initialize a TCP listener
ln, err := net.Listen("tcp", s.lnAddr)
if err != nil {
return err
}
// add it to the server struct
s.ln = ln
// spawn a goroutine for context cancellation
// this will close the listener when context cancels
// in our case , it's Ctrl+c (SIGINT).
go func() {
<-ctx.Done()
ln.Close()
}()
// spawn will contineously listen for connections
wg.Go(func() {
s.connLoop(ctx, wg)
})
return nil
}
and the connection loop like this :-
func (s *Server) connLoop(ctx context.Context, wg *sync.WaitGroup) {
for {
conn, err := s.ln.Accept() // <- blocks until new connection arrives or listener is closed
if err != nil {
// this here checks for if the listener is closed
// if it is closed we stop listening for new connections
if errors.Is(err, net.ErrClosed) {
return
}
// if one connection is not accepted it's fine
// continue listening for newer connections
continue
}
// for each new connection,
// spawn a handler -> this will read from the connection
// and do what is supposed to be done
wg.Add(1)
go s.handleConn(ctx, conn, wg)
}
}
And the last piece we will need , the handleConn function which is something like this this :-
func (s *Server) handleConn(ctx context.Context, conn net.Conn, wg *sync.WaitGroup) {
defer wg.Done()
defer conn.Close()
buf := make([]byte, 4096)
for {
select {
case <-ctx.Done():
return
default:
n, err := conn.Read(buf)
// parse request
// execute command
// write response
}
}
}
The Problem
Everything look fine right ? Nope we have problems.
This is what out app flow looks like right now :
We stopped listening for new incoming connections, what about active connections ? They continue to serve.
But didn't we write in handleConn function ? It should return for each connection when we cancel the context right ?
NOPE.
The culprit is this line right here
n, err := conn.Read(buf) //<- blocks until some bytes arrive or the connection is closed
Since conn.Read() is blocking, execution never reaches the next iteration of the loop where ctx.Done() is checked.
case <-ctx.Done():
return
And since we have an incremented counter in the weight group in the connection Loop wg.Wait() in the main
function waits forever if idle connections are lingering around.
wg.Add(1)
go s.handleConn(ctx, conn, wg)
Fix Attempts (THE GOOD STUFF)
Okay, now that we know the problem is handleConn isn't actively listening for the context cancellation .
Let's see how we can solve this.
- Read deadlines:
The simplest fix is to introduce a read deadline. We can simply do something like :-
func (s *Server) handleConn(ctx context.Context, conn net.Conn, wg *sync.WaitGroup) {
defer wg.Done()
defer conn.Close()
buf := make([]byte, 4096)
for {
select {
case <-ctx.Done():
return
default:
// sets a read deadline for 1 sec
// when this expires we get a timeout error
conn.SetReadDeadline(time.Now().Add(time.Second))
n, err := conn.Read(buf)
// timeout error is handled here and continue
// just loops again to check the context cancellation
if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
continue
}
// parse request
// execute command
// write response
}
}
}
This was easy and our problem is basically solved.... Or is it ?
Let's find out, now our code is basically polling the context cancellation status every second.
And let's be honest context isn't going to cancel frequently it happens once in a while. So all these checks and
constant looping is just wasted work. Which doesn't look like much , but then we suddenly have million connections
and every connection handler is looping every second for nothing.
Well, THAT ESCALATED QUICKLY!!!
- Go routine listener:
The second most obvious solution that comes to the mind is to add something like this
func (s *Server) handleConn(ctx context.Context, conn net.Conn, wg *sync.WaitGroup) {
defer wg.Done()
// closes the connection
// when context cancels in a separate goroutine
go func(){
<- ctx.Done()
conn.Close()
}()
buf := make([]byte, 4096)
for {
n, err := conn.Read(buf)
if err != nil {
if errors.Is(err , net.ErrClosed) {
return
}
}
// parse request
// execute command
// write response
}
}
Now when the context listener goroutine cancels the conn.Close conn.Read errors with net.ErrClosed and exits the function
which calls wg.Done() for each connection and the server gets to exit.
Okay now it is solved right ? .... SIKEEEEEE! NOT YET
There is an issue with this one as well , suppose we spin up 1 gouroutine for each connection handler and I know this seems unreal
but HIPOTHETICALLY let's assume your shitty app somehow blows up and get's 10,00,000 active connections and each connection
fires one context cancellation listener ?
Do you see the problem ? The number of active grourountine per connection doubled.
Yeah yeah i know , BUT GOROUTINES ARE REALLY LIGHTWEIGHT YOU CAN SPAWN THOUSANDS OF THEM . They're cheap, not free and 20,00,000 goroutines
is SOMETHING .
More importantly, we already maintain a central Server type. Having a single shutdown coordinator is simpler than distributing one cancellation goroutine to every connection.
So how do we solve this ? Let's see in the next section.
The solution
So the actual actual fix goes like this :-
type Server struct {
lnAddr string
ln net.Listener
conns map[net.Conn]struct{} // <- connection set
mu sync.Mutex // <- a mutex cuz now we are modifying this
}
Here we are trying to manage connections and do SOMETHING WITH THEM. Keep reading for something.
Next we will add some helper functions :-
// helper function to register connection
// adds connection to the connections set
func (s *Server) RegisterConnection(conn net.Conn) {
s.mu.Lock()
defer s.mu.Unlock()
s.conns[conn] = struct{}{}
}
// i won't write this you know what this is ? Right ?
func (s *Server) UnregisterConnection(conn net.Conn) {
s.mu.Lock()
defer s.mu.Unlock()
delete(s.conns, conn)
}
Now in out connLoop we can do :-
func (s *Server) connLoop(ctx context.Context ,wg *sync.WaitGroup) {
for {
conn, err := s.ln.Accept() // <- blocks until new connection arrives or listener is closed
if err != nil {
// this here checks for if the listener is closed
// if it is closed we stop listening for new connections
if errors.Is(err, net.ErrClosed) {
return
}
// if one connection is not accepted it's fine
// continue listening for newer connections
continue
}
s.RegisterConnection(conn) // <- THIS RIGHT HERE , WE REGISTER THIS CONNECTION
// for each new connection,
// spawn a handler -> this will read from the connection
// and do what is supposed to be done
wg.Go(func() { // <- this function will do wg.Done() internally when the routine returns
// so we don't do it in function now
s.handleConn(ctx,conn)
})
}
}
And in our handleConn function we can simplify things :-
func (s *Server) handleConn(conn net.Conn) { // <- removed the context , not listening to it anymore
// defer wg.Done() <- removed this because wg.Go() does it for us this routine just has to return
defer conn.Close()
defer s.UnregisterConnection(conn) //<- when the function returns we unregister the connection.
buf := make([]byte, 4096)
for {
n, err := conn.Read(buf)
if err != nil {
return
}
// parse request
// execute command
// write response
}
}
}
Once the server became responsible for managing the lifetime of every active connection, handleConn no longer needed to know anything about the shutdown context. Its lifetime became naturally tied to the lifetime of the TCP connection itself. When the connection is closed, conn.Read() returns and the handler exits on its own.
Now out connection loop becomes :-
func (s *Server) Start(ctx context.Context, wg *sync.WaitGroup) error {
// ... exisitng code
//....
wg.Go(func() {
s.connLoop(wg) // <- removed context
})
return nil
func (s *Server) connLoop(wg *sync.WaitGroup) {
//...existing content
wg.Go(func() {
s.handleConn(ctx,conn) //<- no longer needs context
})
}
Okay but this is an infinite loop right ? . How do you exit when context cancels ? Where is that mechanism ?
Hold your horses, WE AIN'T DONE YET YOU OFFSPRING OF A MAMMAL!!!!!!!.
We will add another helper function Shutdown now.
It goes something like this....
func (s *Server) Shutdown() {
// go read about critical section handling in concurrent applications
// if you don't know this.
s.mu.Lock()
defer s.mu.Unlock()
for conn := range s.conns { //<- loop throughout connections and close them
conn.Close()
}
}
Now that we have this, let's wire things up together.
func main(){
// ..... existing logic
// ......
// this here blocks until the context is cancelled
<- ctx.Done()
svr.Shutdown() // <- THIS GOES TO THROUGH CONNECTIONS AND CLOSES THEM AND
// NEW CONNECTIONS ARE NOT BEING ACCEPTED
// wait for all the goroutine to finish
wg.Wait()
}
So when shutdown closes the connections in the registered connections.
In this implementation , when Shutdown() closes every active connection, the blocked conn.Read() returns with net.ErrClosed, allowing the connection handler to exit.
If you want you can do something this in handleConn :-
func (s *Server) handleConn(conn net.Conn) { // <- removed the context , not listening to it anymore
// defer wg.Done() <- removed this because wg.Go() does it for us this routine just has to return
defer conn.Close()
defer s.UnregisterConnection(conn) //<- when the function returns we unregister the connection.
buf := make([]byte, 4096)
for {
n, err := conn.Read(buf)
if err != nil {
if errors.Is(err, net.ErrClosed) {
log.Printf("connection closed during shutdown: %s", conn.RemoteAddr()) //<- ADDED LOG
return
}
return
}
// parse request
// execute command
// write response
}
}
}
Now you can see things happening. Okay so this is it ? Right ? RIGHT??? .
YES and no.
I want to leave you with a subtle warning that is really crucial when doing this.
func main(){
// ..... existing logic
// ......
// this here blocks until the context is cancelled
<- ctx.Done()
svr.Shutdown() // <- THIS GOES TO THROUGH CONNECTIONS AND CLOSES THEM AND
// NEW CONNECTIONS ARE NOT BEING ACCEPTED
// wait for all the goroutine to finish
wg.Wait()
}
Things i learned
- Closing connection listener only stops the new incoming connections. Existing connections require additional cleanup mechanism.
conn.Readcannot watch for the context closure. It only returns when data arrives , connection closes , or deadline expires (if setup properly).- Not every long-running function necessarily needs a
context.Context. It's better to ask questions on what components owns the lifecycle of a given routine and How ? - Cleanup and wait order matters a lot, If in wrong order it can lead to deadlock.
That's it . PEACE OFF.................
