diff --git a/grace.go b/grace.go index cb8067b..5032998 100644 --- a/grace.go +++ b/grace.go @@ -59,11 +59,12 @@ type deadliner interface { // Allows for us to notice when the connection is closed. type conn struct { net.Conn - wg *sync.WaitGroup + wg *sync.WaitGroup + once sync.Once } -func (c conn) Close() error { - defer c.wg.Done() +func (c *conn) Close() error { + defer c.once.Do(c.wg.Done) return c.Conn.Close() } @@ -133,7 +134,7 @@ func (l *listener) Accept() (net.Conn, error) { } return nil, err } - return conn{Conn: c, wg: &l.wg}, nil + return &conn{Conn: c, wg: &l.wg}, nil } type Process struct {