github.com/godaddy-x/freego@v1.0.156/node/listener.go (about)

     1  package node
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"sync/atomic"
     7  	"time"
     8  )
     9  
    10  type GracefulListener struct {
    11  	// inner listener
    12  	ln net.Listener
    13  	// maximum wait time for graceful shutdown
    14  	maxWaitTime time.Duration
    15  	// this channel is closed during graceful shutdown on zero open connections.
    16  	done chan struct{}
    17  	// the number of open connections
    18  	connsCount uint64
    19  	// becomes non-zero when graceful shutdown starts
    20  	shutdown uint64
    21  }
    22  
    23  // NewGracefulListener wraps the given listener into 'graceful shutdown' listener.
    24  func NewGracefulListener(address string, maxWaitTime time.Duration) net.Listener {
    25  	ln, err := net.Listen("tcp4", address)
    26  	if err != nil {
    27  		panic(err)
    28  	}
    29  	return &GracefulListener{
    30  		ln:          ln,
    31  		maxWaitTime: maxWaitTime,
    32  		done:        make(chan struct{}),
    33  	}
    34  }
    35  
    36  func (ln *GracefulListener) Accept() (net.Conn, error) {
    37  	c, err := ln.ln.Accept()
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	atomic.AddUint64(&ln.connsCount, 1)
    42  	return &gracefulConn{Conn: c, ln: ln}, nil
    43  }
    44  
    45  func (ln *GracefulListener) Addr() net.Addr {
    46  	return ln.ln.Addr()
    47  }
    48  
    49  // Close closes the inner listener and waits until all the pending open connections
    50  // are closed before returning.
    51  func (ln *GracefulListener) Close() error {
    52  	if err := ln.ln.Close(); err != nil {
    53  		return nil
    54  	}
    55  	return ln.waitForZeroConns()
    56  }
    57  
    58  func (ln *GracefulListener) waitForZeroConns() error {
    59  	atomic.AddUint64(&ln.shutdown, 1)
    60  	fmt.Println("waitForZeroConns", atomic.LoadUint64(&ln.connsCount))
    61  	if atomic.LoadUint64(&ln.connsCount) == 0 {
    62  		close(ln.done)
    63  		return nil
    64  	}
    65  	select {
    66  	case <-ln.done:
    67  		return nil
    68  	case <-time.After(ln.maxWaitTime):
    69  		return fmt.Errorf("cannot complete graceful shutdown in %s", ln.maxWaitTime)
    70  	}
    71  }
    72  
    73  func (ln *GracefulListener) closeConn() {
    74  	// 相当于减1
    75  	connsCount := atomic.AddUint64(&ln.connsCount, ^uint64(0))
    76  	if atomic.LoadUint64(&ln.shutdown) != 0 && connsCount == 0 {
    77  		close(ln.done)
    78  	}
    79  }
    80  
    81  type gracefulConn struct {
    82  	net.Conn
    83  	ln *GracefulListener
    84  }
    85  
    86  func (c *gracefulConn) Close() error {
    87  	if err := c.Conn.Close(); err != nil {
    88  		return err
    89  	}
    90  	c.ln.closeConn()
    91  	return nil
    92  }