github.com/searKing/golang/go@v1.2.117/net/tcp/shutdown.go (about)

     1  // Copyright 2020 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package tcp
     6  
     7  import (
     8  	"context"
     9  	"time"
    10  )
    11  
    12  var shutdownPollInterval = 500 * time.Millisecond
    13  
    14  func (srv *Server) Shutdown(ctx context.Context) error {
    15  	srv.inShutdown.Store(true)
    16  
    17  	srv.mu.Lock()
    18  	lnerr := srv.closeListenersLocked()
    19  	srv.closeDoneChanLocked()
    20  	for _, f := range srv.onShutdown {
    21  		go f()
    22  	}
    23  	srv.mu.Unlock()
    24  
    25  	ticker := time.NewTicker(shutdownPollInterval)
    26  	defer ticker.Stop()
    27  	for {
    28  		if srv.closeIdleConns() {
    29  			return lnerr
    30  		}
    31  		select {
    32  		case <-ctx.Done():
    33  			return ctx.Err()
    34  		case <-ticker.C:
    35  		}
    36  	}
    37  }
    38  
    39  func (srv *Server) RegisterOnShutdown(f func()) {
    40  	srv.mu.Lock()
    41  	srv.onShutdown = append(srv.onShutdown, f)
    42  	srv.mu.Unlock()
    43  }
    44  
    45  // closeIdleConns closes all idle connections and reports whether the
    46  // server is quiescent.
    47  func (srv *Server) closeIdleConns() bool {
    48  	srv.mu.Lock()
    49  	defer srv.mu.Unlock()
    50  	quiescent := true
    51  	for c := range srv.activeConn {
    52  		st, unixSec := c.getState()
    53  		// Issue 22682: treat StateNew connections as if
    54  		// they're idle if we haven't read the first request's
    55  		// header in over 5 seconds.
    56  		if st == StateNew && unixSec < time.Now().Unix()-5 {
    57  			st = StateIdle
    58  		}
    59  		if st != StateIdle || unixSec == 0 {
    60  			// Assume unixSec == 0 means it's a very new
    61  			// connection, without state set yet.
    62  			quiescent = false
    63  			continue
    64  		}
    65  		c.rwc.Close()
    66  		delete(srv.activeConn, c)
    67  	}
    68  	return quiescent
    69  }
    70  
    71  func (srv *Server) closeListenersLocked() error {
    72  	var err error
    73  	for ln := range srv.listeners {
    74  		if cerr := (*ln).Close(); cerr != nil && err == nil {
    75  			err = cerr
    76  		}
    77  		delete(srv.listeners, ln)
    78  	}
    79  	return err
    80  }