Wait for connections to finish when shutting down
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.
This commit is contained in:
parent
2a96aa1cdd
commit
d3a6756290
10
server.go
10
server.go
@ -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{}
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
@ -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()
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user