diff --git a/server.go b/server.go index 26419831d..515547b58 100644 --- a/server.go +++ b/server.go @@ -74,9 +74,18 @@ func (s *Server) RegisterService(name string, desc *ServiceDesc) { } func (s *Server) Serve(ctx context.Context, l net.Listener) error { - s.addListener(l) + s.mu.Lock() + s.addListenerLocked(l) defer s.closeListener(l) + select { + case <-s.done: + s.mu.Unlock() + return ErrServerClosed + default: + } + s.mu.Unlock() + var ( backoff time.Duration handshaker = s.config.handshaker @@ -191,6 +200,10 @@ func (s *Server) Close() error { func (s *Server) addListener(l net.Listener) { s.mu.Lock() defer s.mu.Unlock() + s.addListenerLocked(l) +} + +func (s *Server) addListenerLocked(l net.Listener) { s.listeners[l] = struct{}{} }