Prevent double waitgroup decrement (#10170) (#10175)

* Prevent double waitgroup decrement

Co-authored-by: zeripath <art27@cantab.net>

diff --git a/modules/graceful/server.go b/modules/graceful/server.go
index 19ce8a866f3d6ecb21ef797ab52161f777a8713d..09be9339997bf55799cc4e003af6aea65efdca8b 100644 (file)
--- 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()
 }