Cancelling a net.Listener via Context in Golang

user_658736 0
Viewed 2824

I'm implementing a TCP server application that accepts incoming TCP connections in an infinite loop.

I'm trying to use Context throughout the application to allow shutting down, which is generally working great.

The one thing I'm struggling with is cancelling a net.Listener that is waiting on Accept(). I'm using a ListenConfig which, I believe, has the advantage of taking a Context when then creating a Listener. However, cancelling this Context does not have the intended effect of aborting the Accept call.

Here's a small app that demonstrates the same problem:

package main

import (
    "context"
    "fmt"
    "net"
    "time"
)

func main() {
    lc := net.ListenConfig{}

    ctx, cancel := context.WithCancel(context.Background())

    go func() {
        time.Sleep(2*time.Second)
        fmt.Println("cancelling context...")
        cancel()
    }()
    ln, err := lc.Listen(ctx, "tcp", ":9801")
    if err != nil {
        fmt.Println("error creating listener:", err)
    } else {
        fmt.Println("listen returned without error")
        defer ln.Close()
    }
    conn, err := ln.Accept()
    if err != nil {
        fmt.Println("accept returned error:", err)
    } else {
        fmt.Println("accept returned without error")
        defer conn.Close()
    }
}

I expect that, if no clients connect, when the Context is cancelled 2 seconds after startup, the Accept() should abort. However, it just sits there until you Ctrl-C out.

Is my expectation wrong? If so, what is the point of the Context passed to ListenConfig.Listen()?

Is there another way to achieve the same goal?

go tcp
1 Answers

I believe you should be closing the listener when your timeout runs out. Then, when Accept returns an error, check that it's intentional (e.g. the timeout elapsed).

This blog post shows how to do a safe shutdown of a TCP server without a context. The interesting part of the code is:

type Server struct {
    listener net.Listener
    quit     chan interface{}
    wg       sync.WaitGroup
}

func NewServer(addr string) *Server {
    s := &Server{
        quit: make(chan interface{}),
    }
    l, err := net.Listen("tcp", addr)
    if err != nil {
        log.Fatal(err)
    }
    s.listener = l
    s.wg.Add(1)
    go s.serve()
    return s
}

func (s *Server) Stop() {
    close(s.quit)
    s.listener.Close()
    s.wg.Wait()
}

func (s *Server) serve() {
    defer s.wg.Done()

    for {
        conn, err := s.listener.Accept()
        if err != nil {
            select {
            case <-s.quit:
                return
            default:
                log.Println("accept error", err)
            }
        } else {
            s.wg.Add(1)
            go func() {
                s.handleConection(conn)
                s.wg.Done()
            }()
        }
    }
}

func (s *Server) handleConection(conn net.Conn) {
    defer conn.Close()
    buf := make([]byte, 2048)
    for {
        n, err := conn.Read(buf)
        if err != nil && err != io.EOF {
            log.Println("read error", err)
            return
        }
        if n == 0 {
            return
        }
        log.Printf("received from %v: %s", conn.RemoteAddr(), string(buf[:n]))
    }
}

In your case you should call Stop when the context runs out.

If you look at the source code of TCPConn.Accept, you'll see it basically calls the underlying socket accept, and the context is not piped through there. But Accept is simple to cancel by closing the listener, so piping the context all the way isn't strictly necessary.

eliben_82060
Related