From e068b65c78b24db468bb57e36bcb3662d472ae86 Mon Sep 17 00:00:00 2001
From: rubyist <scott.barron@github.com>
Date: Sun, 3 Nov 2013 13:03:01 -0500
Subject: [PATCH] Ensure that the WaitGroup isn't marked done twice when
 edgecase errors occurr

---
 grace.go | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/grace.go b/grace.go
index cb8067b..9fd97ef 100644
--- a/grace.go
+++ b/grace.go
@@ -59,11 +59,15 @@ type deadliner interface {
 // Allows for us to notice when the connection is closed.
 type conn struct {
 	net.Conn
-	wg *sync.WaitGroup
+	wg     *sync.WaitGroup
+	closed bool
 }
 
-func (c conn) Close() error {
-	defer c.wg.Done()
+func (c *conn) Close() error {
+	if !c.closed {
+		c.closed = true
+		defer c.wg.Done()
+	}
 	return c.Conn.Close()
 }
 
@@ -133,7 +137,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 {