From e06861fbec3a3c5e3d28a993691811db0be78a12 Mon Sep 17 00:00:00 2001
From: WeidiDeng <weidi_deng@icloud.com>
Date: Mon, 28 Oct 2024 11:13:07 +0800
Subject: [PATCH] stop the loop when listener is closed

---
 layer4/listener.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/layer4/listener.go b/layer4/listener.go
index 70c2733d..b64067e4 100644
--- a/layer4/listener.go
+++ b/layer4/listener.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"runtime"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/caddyserver/caddy/v2"
@@ -114,6 +115,7 @@ type listener struct {
 	logger        *zap.Logger
 	compiledRoute Handler
 
+	closed atomic.Bool
 	// closed when there is a non-recoverable error and all handle goroutines are done
 	connChan chan net.Conn
 	err      error
@@ -122,10 +124,21 @@ type listener struct {
 	wg *sync.WaitGroup
 }
 
+func (l *listener) Close() error {
+	err := l.Listener.Close()
+	l.closed.Store(true)
+	return err
+}
+
 // loop accept connection from underlying listener and pipe the connection if there are any
 func (l *listener) loop() {
 	for {
 		conn, err := l.Listener.Accept()
+		// listener closed
+		if l.closed.Load() {
+			break
+		}
+
 		var nerr net.Error
 		if errors.As(err, &nerr) && nerr.Temporary() {
 			l.logger.Error("temporary error accepting connection", zap.Error(err))