Fasthttp: Handling connections after closing net.Listener

[mtojek]

Hi @valyala ,

I am curious if it is possible (or will be in the future) to trace the worker pool in the context of connection pool size. After closing the net.Listener, the HTTP server does not accept any new connections and it allows earlier accepted ones to be finished (to be served until disconnection).
I need to find a way to monitor the connection pool for the case of graceful shutdown, just to be notified when the number of active connections decreases to zero.

The reference implementation of HTTP server supporting such feature is available at:
https://github.com/tylerb/graceful
(it stores net.Conn in a hash map)

Kind regards,
Marcin

[valyala]

Hi Marcin,
Yo can implement custom net.Listener, which will wrap accepted connections into a custom net.Conn with overridden Close method, which decrements the number of open connections. This way you may keep track on the number of open connections through custom listener and wait until this number drops to zero before returning from overridden net.Listener.Close. I'd suggest using non-zero timeout when waiting for pending connections, since certain clients may keep connections open for extended periods of time, so the graceful shutdown may never finish.

[mtojek]

Thank you @valyala for the proposed solution, but it may be more difficult if I had already used another listener like e.g. https://github.com/kavu/go_reuseport.

[valyala]

@mtojek , just wrap any listener into 'graceful shutdown' listener. Something like the following (I didn't test it, but you should get the idea):

type GracefulListener struct {
    // inner listener
    ln net.Listener

    // maximum wait time for graceful shutdown
    maxWaitTime time.Duration

    // this channel is closed during graceful shutdown on zero open connections.
    done chan struct{}

    // the number of open connections
    connsCount uint64

    // becomes non-zero when graceful shutdown starts
    shutdown uint64
}

// NewGracefulListener wraps the given listener into 'graceful shutdown' listener.
func NewGracefulListener(ln net.Listener, maxWaitTime time.Duration) net.Listener {
    return &GracefulListener{
        ln: ln,
        maxWaitTime: maxWaitTime,
        done: make(chan struct{}),
    }
}

func (ln *GracefulListener) Accept() (net.Conn, error) {
    c, err := ln.ln.Accept()
    if err != nil {
        return nil, err
    }
    atomic.AddUint64(&ln.connsCount, 1)
    return &gracefulConn{
        Conn: c,
        ln: ln,
    }, nil
}

// Close closes the inner listener and waits until all the pending open connections
// are closed before returning.
func (ln *GracefulListener) Close() error {
    err := ln.ln.Close()
    if err != nil {
        return nil
    }
    return ln.waitForZeroConns()
}

func (ln *GracefulListener) waitForZeroConns() error {
    atomic.AddUint64(&ln.shutdown, 1)
    select {
    case <-ln.done:
        return nil
    case <-time.After(ln.maxWaitTime):
        return fmt.Errorf("cannot complete graceful shutdown in %s", ln.maxWaitTime)
    }
}

func (ln *GracefulListener) closeConn() {
    connsCount := atomic.AddUint64(&ln.connsCount, ^uint64(0))
    if atomic.LoadUint64(&ln.shutdown) != 0 && connsCount == 0 {
        close(ln.done)
    }
}

type gracefulConn struct {
    net.Conn
    ln *GracefulListener
}

func (c *gracefulConn) Close() error {
    err := c.Conn.Close()
    if err != nil {
        return err
    }
    c.ln.closeConn()
}

[derekstavis]

Here's @valyala snippet with a couple of fixes that were need to make it work:

type GracefulListener struct {
    // inner listener
    ln net.Listener

    // maximum wait time for graceful shutdown
    maxWaitTime time.Duration

    // this channel is closed during graceful shutdown on zero open connections.
    done chan struct{}

    // the number of open connections
    connsCount uint64

    // becomes non-zero when graceful shutdown starts
    shutdown uint64
}

// NewGracefulListener wraps the given listener into 'graceful shutdown' listener.
func newGracefulListener(ln net.Listener, maxWaitTime time.Duration) net.Listener {
    return &GracefulListener{
        ln:          ln,
        maxWaitTime: maxWaitTime,
        done:        make(chan struct{}),
    }
}

func (ln *GracefulListener) Accept() (net.Conn, error) {
    c, err := ln.ln.Accept()

    if err != nil {
        return nil, err
    }

    atomic.AddUint64(&ln.connsCount, 1)

    return &gracefulConn{
        Conn: c,
        ln:   ln,
    }, nil
}

func (ln *GracefulListener) Addr() net.Addr {
    return ln.ln.Addr()
}

// Close closes the inner listener and waits until all the pending open connections
// are closed before returning.
func (ln *GracefulListener) Close() error {
    err := ln.ln.Close()

    if err != nil {
        return nil
    }

    return ln.waitForZeroConns()
}

func (ln *GracefulListener) waitForZeroConns() error {
    atomic.AddUint64(&ln.shutdown, 1)

    if atomic.LoadUint64(&ln.connsCount) == 0 {
        close(ln.done)
        return nil
    }

    select {
    case <-ln.done:
        return nil
    case <-time.After(ln.maxWaitTime):
        return fmt.Errorf("cannot complete graceful shutdown in %s", ln.maxWaitTime)
    }

    return nil
}

func (ln *GracefulListener) closeConn() {
    connsCount := atomic.AddUint64(&ln.connsCount, ^uint64(0))

    if atomic.LoadUint64(&ln.shutdown) != 0 && connsCount == 0 {
        close(ln.done)
    }
}

type gracefulConn struct {
    net.Conn
    ln *GracefulListener
}

func (c *gracefulConn) Close() error {
    err := c.Conn.Close()

    if err != nil {
        return err
    }

    c.ln.closeConn()

    return nil
}

[mtojek]

Even though I already solved the problem, Big Up for responding after one year, thanks! :)

[jeffrand]

Thanks for writing this! I added it to a repo so we can use it in a few projects

https://github.com/seatgeek/graceful_listener

[dstroot]

I wrote a full server example. I wanted a fully working sample for future reference. It is pretty limited but I like the structure and layout. I didn't write any tests yet. Feedback welcome and encouraged. Cheers!

https://github.com/dstroot/fasthttp-graceful-shutdown