From 4b1e3ebc0054a8fe9c022175d1a049b070e2c96e Mon Sep 17 00:00:00 2001
From: Fabian Holler <fabian.holler@simplesurance.de>
Date: Fri, 20 Oct 2023 14:40:00 +0200
Subject: [PATCH] move shared Stop() and GracefulStop() code into
 stop(graceful)

---
 server.go | 37 ++++++++++++-------------------------
 1 file changed, 12 insertions(+), 25 deletions(-)

diff --git a/server.go b/server.go
index 8534c907cca6..816e9e77930d 100644
--- a/server.go
+++ b/server.go
@@ -1817,36 +1817,17 @@ func ServerTransportStreamFromContext(ctx context.Context) ServerTransportStream
 // pending RPCs on the client side will get notified by connection
 // errors.
 func (s *Server) Stop() {
-	s.quit.Fire()
-	defer s.done.Fire()
-
-	s.channelzRemoveOnce.Do(func() { channelz.RemoveEntry(s.channelzID) })
-
-	s.mu.Lock()
-	s.closeListeners()
-	s.mu.Unlock()
-
-	// Wait for serving threads to be ready to exit.  Only then can we be sure no
-	// new conns will be created.
-	s.serveWG.Wait()
-
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	s.closeServerTransports()
-
-	if s.opts.numServerWorkers > 0 {
-		s.stopServerWorkers()
-	}
-
-	s.waitForServerConnRemoval()
-	s.finishEventLog()
+	s.stop(false)
 }
 
 // GracefulStop stops the gRPC server gracefully. It stops the server from
 // accepting new connections and RPCs and blocks until all the pending RPCs are
 // finished.
 func (s *Server) GracefulStop() {
+	s.stop(true)
+}
+
+func (s *Server) stop(graceful bool) {
 	s.quit.Fire()
 	defer s.done.Fire()
 
@@ -1854,7 +1835,9 @@ func (s *Server) GracefulStop() {
 
 	s.mu.Lock()
 	s.closeListeners()
-	s.drainAllServerTransports()
+	if graceful {
+		s.drainAllServerTransports()
+	}
 	s.mu.Unlock()
 
 	// Wait for serving threads to be ready to exit.  Only then can we be sure no
@@ -1864,6 +1847,10 @@ func (s *Server) GracefulStop() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	if !graceful {
+		s.closeServerTransports()
+	}
+
 	if s.opts.numServerWorkers > 0 {
 		s.stopServerWorkers()
 	}