handle race condition that may arise between accept and add

This commit is contained in:
Naitik Shah 2013-09-10 12:42:11 -07:00
parent 4817bd8812
commit f41bb0a2be

View File

@ -91,7 +91,19 @@ func (l *listener) Close() error {
return err return err
} }
func (l *listener) Accept() (net.Conn, error) { func (l *listener) Accept() (c net.Conn, err error) {
// Presume we'll accept and decrement in defer if we don't. If we did this
// after a successful accept we would have a race condition where we may end
// up incorrectly shutting down between the time we do a successful accept
// and the increment.
l.wg.Add(1)
defer func() {
// If we didn't accept, we decrement our presumptuous count above.
if c == nil {
l.wg.Done()
}
}()
l.closedMutex.RLock() l.closedMutex.RLock()
if l.closed { if l.closed {
l.closedMutex.RUnlock() l.closedMutex.RUnlock()
@ -99,7 +111,7 @@ func (l *listener) Accept() (net.Conn, error) {
} }
l.closedMutex.RUnlock() l.closedMutex.RUnlock()
c, err := l.Listener.Accept() c, err = l.Listener.Accept()
if err != nil { if err != nil {
if strings.HasSuffix(err.Error(), errClosed) { if strings.HasSuffix(err.Error(), errClosed) {
return nil, ErrAlreadyClosed return nil, ErrAlreadyClosed
@ -118,7 +130,6 @@ func (l *listener) Accept() (net.Conn, error) {
} }
return nil, err return nil, err
} }
l.wg.Add(1)
return conn{Conn: c, wg: &l.wg}, nil return conn{Conn: c, wg: &l.wg}, nil
} }