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
|
mu sync.Mutex
|
||||||
listeners map[net.Listener]struct{}
|
listeners map[net.Listener]struct{}
|
||||||
conns map[*gossh.ServerConn]struct{}
|
conns map[*gossh.ServerConn]struct{}
|
||||||
|
connWg sync.WaitGroup
|
||||||
doneChan chan struct{}
|
doneChan chan struct{}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -122,16 +123,17 @@ func (srv *Server) Shutdown(ctx context.Context) error {
|
|||||||
srv.closeDoneChanLocked()
|
srv.closeDoneChanLocked()
|
||||||
srv.mu.Unlock()
|
srv.mu.Unlock()
|
||||||
|
|
||||||
listenerWgChan := make(chan struct{}, 1)
|
finished := make(chan struct{}, 1)
|
||||||
go func() {
|
go func() {
|
||||||
srv.listenerWg.Wait()
|
srv.listenerWg.Wait()
|
||||||
listenerWgChan <- struct{}{}
|
srv.connWg.Wait()
|
||||||
|
finished <- struct{}{}
|
||||||
}()
|
}()
|
||||||
|
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
return ctx.Err()
|
return ctx.Err()
|
||||||
case <-listenerWgChan:
|
case <-finished:
|
||||||
return lnerr
|
return lnerr
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -319,7 +321,9 @@ func (srv *Server) trackConn(c *gossh.ServerConn, add bool) {
|
|||||||
}
|
}
|
||||||
if add {
|
if add {
|
||||||
srv.conns[c] = struct{}{}
|
srv.conns[c] = struct{}{}
|
||||||
|
srv.connWg.Add(1)
|
||||||
} else {
|
} else {
|
||||||
delete(srv.conns, c)
|
delete(srv.conns, c)
|
||||||
|
srv.connWg.Done()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user