grpc: Wait until resources finish cleaning up in Stop() and GracefulStop()

internal/transport/http2_server.go

@@ -68,15 +68,15 @@ var serverConnectionCounter uint64
 
 // http2Server implements the ServerTransport interface with HTTP2.
 type http2Server struct {
-	lastRead    int64 // Keep this field 64-bit aligned. Accessed atomically.
-	done        chan struct{}
-	conn        net.Conn
-	loopy       *loopyWriter
-	readerDone  chan struct{} // sync point to enable testing.
-	writerDone  chan struct{} // sync point to enable testing.
-	peer        peer.Peer
-	inTapHandle tap.ServerInHandle
-	framer      *framer
+	lastRead        int64 // Keep this field 64-bit aligned. Accessed atomically.
+	done            chan struct{}
+	conn            net.Conn
+	loopy           *loopyWriter
+	readerDone      chan struct{} // sync point to enable testing.
+	loopyWriterDone chan struct{}
+	peer            peer.Peer
+	inTapHandle     tap.ServerInHandle
+	framer          *framer
 	// The max number of concurrent streams.
 	maxStreams uint32
 	// controlBuf delivers all the control related tasks (e.g., window
@@ -251,7 +251,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 		peer:              peer,
 		framer:            framer,
 		readerDone:        make(chan struct{}),
-		writerDone:        make(chan struct{}),
+		loopyWriterDone:   make(chan struct{}),
 		maxStreams:        config.MaxStreams,
 		inTapHandle:       config.InTapHandle,
 		fc:                &trInFlow{limit: uint32(icwz)},
@@ -323,7 +323,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 		t.loopy = newLoopyWriter(serverSide, t.framer, t.controlBuf, t.bdpEst, t.conn, t.logger)
 		t.loopy.ssGoAwayHandler = t.outgoingGoAwayHandler
 		t.loopy.run()
-		close(t.writerDone)
+		close(t.loopyWriterDone)
 	}()
 	go t.keepalive()
 	return t, nil
@@ -608,7 +608,10 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 // typically run in a separate goroutine.
 // traceCtx attaches trace to ctx and returns the new context.
 func (t *http2Server) HandleStreams(ctx context.Context, handle func(*Stream)) {
-	defer close(t.readerDone)
+	defer func() {
+		<-t.loopyWriterDone
+		close(t.readerDone)
+	}()
 	for {
 		t.controlBuf.throttle()
 		frame, err := t.framer.fr.ReadFrame()

internal/transport/transport_test.go

@@ -1685,7 +1685,7 @@ func testFlowControlAccountCheck(t *testing.T, msgSize int, wc windowSizeConfig)
 	client.Close(errors.New("closed manually by test"))
 	st.Close(errors.New("closed manually by test"))
 	<-st.readerDone
-	<-st.writerDone
+	<-st.loopyWriterDone
 	<-client.readerDone
 	<-client.writerDone
 	for _, cstream := range clientStreams {

server.go

@@ -992,15 +992,11 @@
 		}
 	}()
 
-	var wg sync.WaitGroup
 	streamQuota := newHandlerQuota(s.opts.maxConcurrentStreams)
 	st.HandleStreams(ctx, func(stream *transport.Stream) {
-		wg.Add(1)
-
 		streamQuota.acquire()
 		f := func() {
 			defer streamQuota.release()
-			defer wg.Done()
 			s.handleStream(st, stream)
 		}
 
@@ -1014,7 +1010,6 @@
 		}
 		go f()
 	})
-	wg.Wait()
 }
 
 var _ http.Handler = (*Server)(nil)
@@ -1854,62 +1849,66 @@
 // pending RPCs on the client side will get notified by connection
 // errors.
 func (s *Server) Stop() {
-	s.quit.Fire()
+	s.stop(false)
+}
 
-	defer func() {
-		s.serveWG.Wait()
-		s.done.Fire()
-	}()
+// 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()
 
 	s.channelzRemoveOnce.Do(func() { channelz.RemoveEntry(s.channelzID) })
 
 	s.mu.Lock()
-	listeners := s.lis
-	s.lis = nil
-	conns := s.conns
-	s.conns = nil
-	// interrupt GracefulStop if Stop and GracefulStop are called concurrently.
-	s.cv.Broadcast()
+	s.closeListeners()
+	// Wait for serving threads to be ready to exit.  Only then can we be sure no
+	// new conns will be created.
 	s.mu.Unlock()
+	s.serveWG.Wait()
 
-	for lis := range listeners {
-		lis.Close()
-	}
-	for _, cs := range conns {
-		for st := range cs {
-			st.Close(errors.New("Server.Stop called"))
-		}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if graceful {
+		s.drainAllServerTransports()
+	} else {
+		s.closeServerTransports()
 	}
+
 	if s.opts.numServerWorkers > 0 {
 		s.stopServerWorkers()
 	}
 
-	s.mu.Lock()
-	if s.events != nil {
-		s.events.Finish()
-		s.events = nil
-	}
-	s.mu.Unlock()
+	s.waitForServerConnRemoval()
+	s.finishEventLog()
 }
 
-// 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.quit.Fire()
-	defer s.done.Fire()
-
-	s.channelzRemoveOnce.Do(func() { channelz.RemoveEntry(s.channelzID) })
-	s.mu.Lock()
-	if s.conns == nil {
-		s.mu.Unlock()
-		return
+// s.mu must be held by the caller.
+// waitForServerConnRemoval blocks until s.conns is empty.
+func (s *Server) waitForServerConnRemoval() {
+	for len(s.conns) != 0 {
+		s.cv.Wait()
 	}
+	s.conns = nil
+}
 
-	for lis := range s.lis {
-		lis.Close()
+// s.mu must be held by the caller.
+func (s *Server) closeServerTransports() {
+	for _, conns := range s.conns {
+		for st := range conns {
+			st.Close(errors.New("Server.Stop called"))
+		}
 	}
-	s.lis = nil
+}
+
+// s.mu must be held by the caller.
+func (s *Server) drainAllServerTransports() {
 	if !s.drain {
 		for _, conns := range s.conns {
 			for st := range conns {
@@ -1918,22 +1917,22 @@
 		}
 		s.drain = true
 	}
+}
 
-	// Wait for serving threads to be ready to exit.  Only then can we be sure no
-	// new conns will be created.
-	s.mu.Unlock()
-	s.serveWG.Wait()
-	s.mu.Lock()
-
-	for len(s.conns) != 0 {
-		s.cv.Wait()
-	}
-	s.conns = nil
+// s.mu must be held by the caller.
+func (s *Server) finishEventLog() {
 	if s.events != nil {
 		s.events.Finish()
 		s.events = nil
 	}
-	s.mu.Unlock()
+}
+
+// s.mu must be held by the caller.
+func (s *Server) closeListeners() {
+	for lis := range s.lis {
+		lis.Close()
+	}
+	s.lis = nil
 }
 
 // contentSubtype must be lowercase

test/gracefulstop_test.go

@@ -24,6 +24,7 @@ import (
 	"net"
 	"sync"
 	"testing"
+	"time"
 
 	"golang.org/x/net/http2"
 	"google.golang.org/grpc"
@@ -216,3 +217,89 @@ func (s) TestGracefulStopClosesConnAfterLastStream(t *testing.T) {
 		<-gracefulStopDone // Wait for GracefulStop to return.
 	})
 }
+
+func (s) TestGracefulStopBlocksUntilGRPCConnectionsTerminate(t *testing.T) {
+	// This tests ensures that GracefulStop() blocks until all ongoing
+	// client GRPC calls finished.
+	unblockGRPCCall := make(chan struct{})
+	grpcCallExecuting := make(chan struct{})
+	ss := &stubserver.StubServer{
+		UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+			close(grpcCallExecuting)
+			<-unblockGRPCCall
+			return &testpb.SimpleResponse{}, nil
+		},
+	}
+
+	err := ss.Start(nil)
+	if err != nil {
+		t.Fatalf("StubServer.start failed: %s", err)
+	}
+	t.Cleanup(ss.Stop)
+
+	grpcClientCallReturned := make(chan struct{})
+	go func() {
+		clt := ss.Client
+		_, err := clt.UnaryCall(context.Background(), &testpb.SimpleRequest{})
+		if err != nil {
+			t.Errorf("grpc call failed with error: %s", err)
+		}
+		close(grpcClientCallReturned)
+	}()
+
+	gracefulStopReturned := make(chan struct{})
+	<-grpcCallExecuting
+	go func() {
+		ss.S.GracefulStop()
+		close(gracefulStopReturned)
+	}()
+
+	select {
+	case <-gracefulStopReturned:
+		t.Error("GracefulStop returned before GRPC method call ended")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	unblockGRPCCall <- struct{}{}
+	<-grpcClientCallReturned
+	<-gracefulStopReturned
+}
+
+func (s) TestStopAbortsBlockingGRPCCall(t *testing.T) {
+	// This tests ensures that when Stop() is called while an ongoing grpc
+	// call is blocking that:
+	// - Stop() returns
+	// - and the GRPC call on the client side fails with an connection
+	// closed error
+	unblockGRPCCall := make(chan struct{})
+	grpcCallExecuting := make(chan struct{})
+	ss := &stubserver.StubServer{
+		UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+			close(grpcCallExecuting)
+			<-unblockGRPCCall
+			return &testpb.SimpleResponse{}, nil
+		},
+	}
+
+	err := ss.Start(nil)
+	if err != nil {
+		t.Fatalf("StubServer.start failed: %s", err)
+	}
+	t.Cleanup(ss.Stop)
+
+	grpcClientCallReturned := make(chan struct{})
+	go func() {
+		clt := ss.Client
+		_, err := clt.UnaryCall(context.Background(), &testpb.SimpleRequest{})
+		if err == nil || !isConnClosedErr(err) {
+			t.Errorf("expected grpc call to fail with connection closed error, got: %v", err)
+		}
+		close(grpcClientCallReturned)
+	}()
+
+	<-grpcCallExecuting
+	ss.S.Stop()
+
+	unblockGRPCCall <- struct{}{}
+	<-grpcClientCallReturned
+}