grpc: eliminate panics in server worker implementation

server.go

@@ -144,7 +144,8 @@ type Server struct {
 	channelzID *channelz.Identifier
 	czData     *channelzData
 
-	serverWorkerChannel chan func()
+	serverWorkerChannel      chan func()
+	serverWorkerChannelClose func()
 }
 
 type serverOptions struct {
@@ -623,15 +624,14 @@ func (s *Server) serverWorker() {
 // connections to reduce the time spent overall on runtime.morestack.
 func (s *Server) initServerWorkers() {
 	s.serverWorkerChannel = make(chan func())
+	s.serverWorkerChannelClose = grpcsync.OnceFunc(func() {
+		close(s.serverWorkerChannel)
+	})
 	for i := uint32(0); i < s.opts.numServerWorkers; i++ {
 		go s.serverWorker()
 	}
 }
 
-func (s *Server) stopServerWorkers() {
-	close(s.serverWorkerChannel)
-}
-
 // NewServer creates a gRPC server which has no service registered and has not
 // started to accept requests yet.
 func NewServer(opt ...ServerOption) *Server {
@@ -1898,15 +1898,19 @@ func (s *Server) stop(graceful bool) {
 		s.closeServerTransportsLocked()
 	}
 
-	if s.opts.numServerWorkers > 0 {
-		s.stopServerWorkers()
-	}
-
 	for len(s.conns) != 0 {
 		s.cv.Wait()
 	}
 	s.conns = nil
 
+	if s.opts.numServerWorkers > 0 {
+		// Closing the channel (only once, via grpcsync.OnceFunc) after all the
+		// connections have been closed above ensures that there are no
+		// goroutines executing the callback passed to st.HandleStreams (where
+		// the channel is written to).
+		s.serverWorkerChannelClose()
+	}
+
 	if s.events != nil {
 		s.events.Finish()
 		s.events = nil

test/server_test.go

@@ -21,17 +21,81 @@ package test
 import (
 	"context"
 	"io"
+	"runtime"
+	"sync"
 	"testing"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/status"
 
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
 	testpb "google.golang.org/grpc/interop/grpc_testing"
 )
 
+// Tests the case where the server worker goroutine option is enabled, and a
+// number of RPCs are initiated around the same time that Stop() is called. This
+// used to result in a write to a closed channel. This test verifies that there
+// is no panic.
+func (s) TestServerWorkers_RPCsAndStop(t *testing.T) {
+	ss := stubserver.StartTestService(t, nil, grpc.NumStreamWorkers(uint32(runtime.NumCPU())))
+	// This deferred stop takes care of stopping the server when one of the
+	// below grpc.Dials fail, and the test exits early.
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	const numChannels = 20
+	const numRPCLoops = 20
+	var wg sync.WaitGroup
+	for i := 0; i < numChannels; i++ {
+		cc, err := grpc.Dial(ss.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
+		if err != nil {
+			t.Fatalf("[iteration: %d] grpc.Dial(%s) failed: %v", i, ss.Address, err)
+		}
+		defer cc.Close()
+
+		client := testgrpc.NewTestServiceClient(cc)
+		for j := 0; j < numRPCLoops; j++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				for {
+					if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+						t.Logf("EmptyCall() failed: %v", err)
+						return
+					}
+				}
+			}()
+		}
+	}
+	// Call Stop() concurrently with the above RPC attempts.
+	ss.Stop()
+	wg.Wait()
+}
+
+// Tests the case where the server worker goroutine option is enabled, and both
+// Stop() and GracefulStop() care called. This used to result in a close of a
+// closed channel. This test verifies that there is no panic.
+func (s) TestServerGracefulStopAndStop(t *testing.T) {
+	ss := stubserver.StartTestService(t, nil, grpc.NumStreamWorkers(uint32(runtime.NumCPU())))
+	defer ss.Stop()
+
+	if err := ss.StartClient(grpc.WithTransportCredentials(insecure.NewCredentials())); err != nil {
+		t.Fatalf("Failed to create client to stub server: %v", err)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	client := testgrpc.NewTestServiceClient(ss.CC)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall() failed: %v", err)
+	}
+
+	ss.S.GracefulStop()
+}
+
 type ctxKey string
 
 // TestServerReturningContextError verifies that if a context error is returned