backport : dfdaemon to gracefully shutdown

client/daemon/daemon.go

@@ -789,6 +789,13 @@ func (cd *clientDaemon) Serve() error {
 func (cd *clientDaemon) Stop() {
 	cd.once.Do(func() {
 		close(cd.done)
+
+		if cd.ProxyManager.IsEnabled() {
+			if err := cd.ProxyManager.Stop(); err != nil {
+				logger.Errorf("proxy manager stop failed %s", err)
+			}
+		}
+
 		if cd.schedulerClient != nil {
 			if !cd.Option.KeepStorage {
 				logger.Info("leave host with scheduler client")
@@ -820,12 +827,6 @@ func (cd *clientDaemon) Stop() {
 			}
 		}
 
-		if cd.ProxyManager.IsEnabled() {
-			if err := cd.ProxyManager.Stop(); err != nil {
-				logger.Errorf("proxy manager stop failed %s", err)
-			}
-		}
-
 		if !cd.Option.KeepStorage {
 			logger.Infof("keep storage disabled")
 			cd.StorageManager.CleanUp()

client/daemon/proxy/proxy.go

@@ -116,6 +116,8 @@ type Proxy struct {
 	dumpHTTPContent bool
 
 	peerIDGenerator peer.IDGenerator
+
+	wg *sync.WaitGroup
 }
 
 // Option is a functional option for configuring the proxy
@@ -272,6 +274,7 @@ func NewProxyWithOptions(options ...Option) (*Proxy, error) {
 		directHandler: http.NewServeMux(),
 		tracer:        otel.Tracer("dfget-daemon-proxy"),
 		certCache:     lru.New(100),
+		wg:            &sync.WaitGroup{},
 	}
 
 	for _, opt := range options {
@@ -286,6 +289,8 @@ func NewProxyWithOptions(options ...Option) (*Proxy, error) {
 
 // ServeHTTP implements http.Handler.ServeHTTP
 func (proxy *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	proxy.wg.Add(1)
+	defer proxy.wg.Done()
 	metrics.ProxyRequestCount.WithLabelValues(r.Method).Add(1)
 	metrics.ProxyRequestRunningCount.WithLabelValues(r.Method).Add(1)
 	defer metrics.ProxyRequestRunningCount.WithLabelValues(r.Method).Sub(1)

client/daemon/proxy/proxy_manager.go

@@ -142,7 +142,12 @@ func (pm *proxyManager) ServeSNI(listener net.Listener) error {
 }
 
 func (pm *proxyManager) Stop() error {
-	return pm.Server.Shutdown(context.Background())
+	err := pm.Server.Shutdown(context.Background())
+	if err != nil {
+		logger.Errorf("proxy shut down error: %s", err)
+	}
+	pm.Proxy.wg.Wait()
+	return err
 }
 
 func (pm *proxyManager) IsEnabled() bool {

client/daemon/proxy/proxy_sni.go

@@ -66,6 +66,8 @@ func handshakeTLSConn(clientConn net.Conn, config *tls.Config) (net.Conn, error)
 }
 
 func (proxy *Proxy) handleTLSConn(clientConn net.Conn, port int) {
+	proxy.wg.Add(1)
+	defer proxy.wg.Done()
 	var serverName string
 	sConfig := new(tls.Config)
 	if !proxy.cert.Leaf.IsCA {