Prevent double waitgroup decrement (#10170)

* Prevent double waitgroup decrement

1 files changed, 7 insertions, 3 deletions

diff --git a/modules/graceful/server.go b/modules/graceful/server.go
index 19ce8a866f..09be933999 100644
--- a/modules/graceful/server.go
+++ b/modules/graceful/server.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -215,9 +216,12 @@ func (wl *wrappedListener) Accept() (net.Conn, error) {
 		}
 	}
 
+	closed := int32(0)
+
 	c = wrappedConn{
 		Conn:   c,
 		server: wl.server,
+		closed: &closed,
 	}
 
 	wl.server.wg.Add(1)
@@ -241,12 +245,12 @@ func (wl *wrappedListener) File() (*os.File, error) {
 type wrappedConn struct {
 	net.Conn
 	server *Server
+	closed *int32
 }
 
 func (w wrappedConn) Close() error {
-	err := w.Conn.Close()
-	if err == nil {
+	if atomic.CompareAndSwapInt32(w.closed, 0, 1) {
 		w.server.wg.Done()
 	}
-	return err
+	return w.Conn.Close()
 }