From 489a7628114d1dfa652d3eece03e9f38358b5a7a Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Wed, 8 Jan 2025 23:12:06 +0800
Subject: [PATCH] service graceful shutdown

---
 config/parsing/tls.go | 17 ++++++++++++-----
 logger/logger.go      |  7 ++++++-
 recorder/http.go      |  4 ++--
 service/service.go    | 14 +++++++++++++-
 4 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/config/parsing/tls.go b/config/parsing/tls.go
index ac6d4741..7df2b1ad 100644
--- a/config/parsing/tls.go
+++ b/config/parsing/tls.go
@@ -12,6 +12,7 @@ import (
 	"crypto/x509/pkix"
 	"encoding/pem"
 	"math/big"
+	"sync/atomic"
 	"time"
 
 	"github.com/go-gost/core/logger"
@@ -20,14 +21,19 @@ import (
 )
 
 var (
-	defaultTLSConfig *tls.Config
+	defaultTLSConfig atomic.Value
 )
 
 func DefaultTLSConfig() *tls.Config {
-	return defaultTLSConfig
+	v, _ := defaultTLSConfig.Load().(*tls.Config)
+	return v
 }
 
-func BuildDefaultTLSConfig(cfg *config.TLSConfig) {
+func SetDefaultTLSConfig(cfg *tls.Config) {
+	defaultTLSConfig.Store(cfg)
+}
+
+func BuildDefaultTLSConfig(cfg *config.TLSConfig) (*tls.Config, error) {
 	log := logger.Default()
 
 	if cfg == nil {
@@ -43,7 +49,7 @@ func BuildDefaultTLSConfig(cfg *config.TLSConfig) {
 		// generate random self-signed certificate.
 		cert, err := genCertificate(cfg.Validity, cfg.Organization, cfg.CommonName)
 		if err != nil {
-			log.Fatal(err)
+			return nil, err
 		}
 		tlsConfig = &tls.Config{
 			Certificates: []tls.Certificate{cert},
@@ -52,7 +58,8 @@ func BuildDefaultTLSConfig(cfg *config.TLSConfig) {
 	} else {
 		log.Debug("load global TLS certificate files OK")
 	}
-	defaultTLSConfig = tlsConfig
+
+	return tlsConfig, nil
 }
 
 func genCertificate(validity time.Duration, org string, cn string) (cert tls.Certificate, err error) {
diff --git a/logger/logger.go b/logger/logger.go
index effc4652..029301ca 100644
--- a/logger/logger.go
+++ b/logger/logger.go
@@ -84,9 +84,14 @@ func NewLogger(opts ...Option) logger.Logger {
 		log.SetLevel(logrus.InfoLevel)
 	}
 
-	return &logrusLogger{
+	l := &logrusLogger{
 		logger: logrus.NewEntry(log),
 	}
+	if options.Name != "" {
+		l.logger = l.logger.WithField("logger", options.Name)
+	}
+
+	return l
 }
 
 // WithFields adds new fields to log.
diff --git a/recorder/http.go b/recorder/http.go
index 98c1c593..fbf95776 100644
--- a/recorder/http.go
+++ b/recorder/http.go
@@ -3,7 +3,7 @@ package recorder
 import (
 	"bytes"
 	"context"
-	"fmt"
+	"errors"
 	"net/http"
 	"strings"
 	"time"
@@ -80,7 +80,7 @@ func (r *httpRecorder) Record(ctx context.Context, b []byte, opts ...recorder.Re
 	defer resp.Body.Close()
 
 	if resp.StatusCode >= http.StatusBadRequest {
-		return fmt.Errorf(resp.Status)
+		return errors.New(resp.Status)
 	}
 
 	return nil
diff --git a/service/service.go b/service/service.go
index bbb74a28..10807dfa 100644
--- a/service/service.go
+++ b/service/service.go
@@ -2,11 +2,13 @@ package service
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
 	"os/exec"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/go-gost/core/admission"
@@ -157,6 +159,9 @@ func (s *defaultService) Serve() error {
 		defer v.Dec()
 	}
 
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
 	var tempDelay time.Duration
 	for {
 		conn, e := s.listener.Accept()
@@ -179,7 +184,10 @@ func (s *defaultService) Serve() error {
 				continue
 			}
 			s.setState(StateClosed)
-			s.options.logger.Errorf("accept: %v", e)
+
+			if !errors.Is(e, net.ErrClosed) {
+				s.options.logger.Errorf("accept: %v", e)
+			}
 
 			return e
 		}
@@ -224,7 +232,11 @@ func (s *defaultService) Serve() error {
 			continue
 		}
 
+		wg.Add(1)
+
 		go func() {
+			defer wg.Done()
+
 			if v := xmetrics.GetCounter(xmetrics.MetricServiceRequestsCounter,
 				metrics.Labels{"service": s.name, "client": clientIP}); v != nil {
 				v.Inc()