Skip to content

Commit

Permalink
Wait for connections to finish when shutting down
Browse files Browse the repository at this point in the history
PR #74 introduced a WaitGroup for listeners, but it doesn't wait for
open connections before closing the server. This patch waits until all
conns are closed before returning from Shutdown.
  • Loading branch information
jbarnette authored and belak committed May 7, 2018
1 parent 2a96aa1 commit d3a6756
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions server.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ type Server struct {
mu sync.Mutex
listeners map[net.Listener]struct{}
conns map[*gossh.ServerConn]struct{}
connWg sync.WaitGroup
doneChan chan struct{}
}

Expand Down Expand Up @@ -122,16 +123,17 @@ func (srv *Server) Shutdown(ctx context.Context) error {
srv.closeDoneChanLocked()
srv.mu.Unlock()

listenerWgChan := make(chan struct{}, 1)
finished := make(chan struct{}, 1)
go func() {
srv.listenerWg.Wait()
listenerWgChan <- struct{}{}
srv.connWg.Wait()
finished <- struct{}{}
}()

select {
case <-ctx.Done():
return ctx.Err()
case <-listenerWgChan:
case <-finished:
return lnerr
}
}
Expand Down Expand Up @@ -319,7 +321,9 @@ func (srv *Server) trackConn(c *gossh.ServerConn, add bool) {
}
if add {
srv.conns[c] = struct{}{}
srv.connWg.Add(1)
} else {
delete(srv.conns, c)
srv.connWg.Done()
}
}

0 comments on commit d3a6756

Please sign in to comment.