feat: implement graceful shutdown for first child listener and enhance connection handling

Author: 0xJacky

firstchild_listener.go

@@ -10,15 +10,37 @@ type firstChildListener struct {
 	ch        chan net.Conn
 	closeOnce sync.Once
 	wg        sync.WaitGroup
+	closed    chan struct{}
 }
 
 func (sl *firstChildListener) Forward(conn net.Conn) {
 	sl.wg.Add(1)
+	
+	// Check if listener is closed before sending
+	select {
+	case <-sl.closed:
+		// Listener is closed, don't forward the connection
+		sl.wg.Done()
+		conn.Close()
+		return
+	default:
+	}
+	
 	select {
 	case sl.ch <- conn:
+	case <-sl.closed:
+		// Listener was closed while trying to send
+		sl.wg.Done()
+		conn.Close()
 	default:
 		go func() {
-			sl.ch <- conn
+			select {
+			case sl.ch <- conn:
+			case <-sl.closed:
+				// Listener was closed while waiting
+				sl.wg.Done()
+				conn.Close()
+			}
 		}()
 	}
 }
@@ -37,8 +59,9 @@ func (sl *firstChildListener) Addr() net.Addr {
 }
 
 func (sl *firstChildListener) Close() error {
-	sl.wg.Wait()
 	sl.closeOnce.Do(func() {
+		close(sl.closed)
+		sl.wg.Wait()
 		close(sl.ch)
 	})
 	return nil

risefront.go

@@ -38,6 +38,7 @@ type Dialer interface {
 type Config struct {
 	Addresses []string                   // Addresses to listen to.
 	Run       func([]net.Listener) error // Handle the connections. All open connections should be properly closed before returning (srv.Shutdown for http.Server for instance).
+	Shutdown  func()                     // A callback to trigger graceful shutdown of the Run function.
 
 	Name          string                                                    // Name of the socket file
 	Dialer        Dialer                                                    // Dialer for child-parent communication. Let empty for default dialer (PrefixDialer{}).
@@ -220,8 +221,9 @@ func (cfg Config) runParent(ctx context.Context, socket string) error {
 
 		// first child listener
 		sl := firstChildListener{
-			addr: ln.Addr(),
-			ch:   make(chan net.Conn),
+			addr:   ln.Addr(),
+			ch:     make(chan net.Conn),
+			closed: make(chan struct{}),
 		}
 		firstChildListeners = append(firstChildListeners, &sl)
 
@@ -564,6 +566,11 @@ func Restart() {
 	err := globalConfig.createChild()
 	if err != nil {
 		globalConfig.LogHandler(ErrorLevel, "restart.createChild", err)
+		return
+	}
+
+	if globalConfig.Shutdown != nil {
+		globalConfig.Shutdown()
 	}
 }