server: Return ErrServerClosed

1 files changed, 6 insertions, 0 deletions

diff --git a/server.go b/server.go
index 69a8b22..c2983bb 100644
--- a/server.go
+++ b/server.go
@@ -219,6 +219,9 @@ func (srv *Server) Serve(ctx context.Context, l net.Listener) error {
 
 	select {
 	case <-lnctx.Done():
+		if atomic.LoadInt32(&srv.closed) == 1 {
+			return ErrServerClosed
+		}
 		return lnctx.Err()
 	case err := <-errch:
 		return err
@@ -234,6 +237,9 @@ func (srv *Server) serve(ctx context.Context, l net.Listener) error {
 		if err != nil {
 			select {
 			case <-ctx.Done():
+				if atomic.LoadInt32(&srv.closed) == 1 {
+					return ErrServerClosed
+				}
 				return ctx.Err()
 			default:
 			}