feat(gRPC): support gRPC graceful shutdown

pkg/remote/trans/nphttp2/conn_pool.go

@@ -121,7 +121,10 @@ func (p *connPool) newTransport(ctx context.Context, dialer remote.Dialer, netwo
 		opts,
 		p.remoteService,
 		func(grpc.GoAwayReason) {
-			// do nothing
+			// remove connection from the pool.
+			// we do not need to close this grpc transport manually
+			// since grpc client is responsible for doing this.
+			p.conns.Delete(address)
 		},
 		func() {
 			// do nothing

pkg/remote/trans/nphttp2/grpc/controlbuf.go

@@ -22,6 +22,7 @@ package grpc
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"runtime"
 	"sync"
@@ -145,10 +146,11 @@ func (h *headerFrame) isTransportResponseFrame() bool {
 }
 
 type cleanupStream struct {
-	streamID uint32
-	rst      bool
-	rstCode  http2.ErrCode
-	onWrite  func()
+	streamID      uint32
+	rst           bool
+	rstCode       http2.ErrCode
+	onWrite       func()
+	onFinishWrite func()
 }
 
 func (c *cleanupStream) isTransportResponseFrame() bool { return c.rst } // Results in a RST_STREAM
@@ -451,19 +453,20 @@ func (c *controlBuffer) get(block bool) (interface{}, error) {
 		select {
 		case <-c.ch:
 		case <-c.done:
-			c.finish()
-			return nil, ErrConnClosing
+			return nil, c.finish(ErrConnClosing)
 		}
 	}
 }
 
-func (c *controlBuffer) finish() {
+func (c *controlBuffer) finish(err error) (rErr error) {
 	c.mu.Lock()
 	if c.err != nil {
+		rErr = c.err
 		c.mu.Unlock()
 		return
 	}
-	c.err = ErrConnClosing
+	c.err = err
+	rErr = err
 	// There may be headers for streams in the control buffer.
 	// These streams need to be cleaned out since the transport
 	// is still not aware of these yet.
@@ -473,10 +476,11 @@ func (c *controlBuffer) finish() {
 			continue
 		}
 		if hdr.onOrphaned != nil { // It will be nil on the server-side.
-			hdr.onOrphaned(ErrConnClosing)
+			hdr.onOrphaned(err)
 		}
 	}
 	c.mu.Unlock()
+	return
 }
 
 type side int
@@ -564,6 +568,10 @@ func (l *loopyWriter) run(remoteAddr string) (err error) {
 			klog.Debugf("KITEX: grpc transport loopyWriter.run returning, error=%v, remoteAddr=%s", err, remoteAddr)
 			err = nil
 		}
+		// make sure the Graceful Shutdown behaviour triggered
+		if errors.Is(err, errGracefulShutdown) {
+			l.framer.writer.Flush()
+		}
 	}()
 	for {
 		it, err := l.cbuf.get(true)
@@ -778,6 +786,11 @@ func (l *loopyWriter) outFlowControlSizeRequestHandler(o *outFlowControlSizeRequ
 }
 
 func (l *loopyWriter) cleanupStreamHandler(c *cleanupStream) error {
+	if c.onFinishWrite != nil {
+		defer func() {
+			c.onFinishWrite()
+		}()
+	}
 	c.onWrite()
 	if str, ok := l.estdStreams[c.streamID]; ok {
 		// On the server side it could be a trailers-only response or

pkg/remote/trans/nphttp2/grpc/controlbuf_test.go

@@ -71,5 +71,5 @@ func TestControlBuf(t *testing.T) {
 	cb.throttle()
 
 	// test finish()
-	cb.finish()
+	cb.finish(ErrConnClosing)
 }

pkg/remote/trans/nphttp2/grpc/graceful_shutdown_test.go

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2024 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package grpc
+
+import (
+	"context"
+	"errors"
+	"math"
+	"strings"
+	"testing"
+
+	"github.com/cloudwego/kitex/internal/test"
+	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/codes"
+)
+
+func TestGracefulShutdown(t *testing.T) {
+	onGoAwayCh := make(chan struct{})
+	srv, cli := setUpWithOnGoAway(t, 10000, &ServerConfig{MaxStreams: math.MaxUint32}, gracefulShutdown, ConnectOptions{}, func(reason GoAwayReason) {
+		close(onGoAwayCh)
+	})
+	defer cli.Close(errSelfCloseForTest)
+
+	stream, err := cli.NewStream(context.Background(), &CallHdr{})
+	test.Assert(t, err == nil, err)
+	<-srv.srvReady
+	finishCh := make(chan struct{})
+	go srv.gracefulShutdown(finishCh)
+	err = cli.Write(stream, nil, []byte("hello"), &Options{})
+	test.Assert(t, err == nil, err)
+	msg := make([]byte, 5)
+	num, err := stream.Read(msg)
+	test.Assert(t, err == nil, err)
+	test.Assert(t, num == 5, num)
+	// waiting onGoAway triggered
+	<-onGoAwayCh
+	// this transport could not create Stream anymore
+	_, err = cli.NewStream(context.Background(), &CallHdr{})
+	test.Assert(t, errors.Is(err, errStreamDrain), err)
+	// wait for the server transport to be closed
+	<-finishCh
+	_, err = stream.Read(msg)
+	test.Assert(t, err != nil, err)
+	st := stream.Status()
+	test.Assert(t, strings.Contains(st.Message(), gracefulShutdownMsg), st.Message())
+	test.Assert(t, st.Code() == codes.Unavailable, st.Code())
+}

pkg/remote/trans/nphttp2/grpc/http2_client.go

@@ -22,6 +22,8 @@ package grpc
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"io"
 	"math"
 	"net"
@@ -468,9 +470,10 @@ func (t *http2Client) NewStream(ctx context.Context, callHdr *CallHdr) (_ *Strea
 		s.id = h.streamID
 		s.fc = &inFlow{limit: uint32(t.initialWindowSize)}
 		t.mu.Lock()
-		if t.activeStreams == nil { // Can be niled from Close().
+		// Don't create a stream if the transport is in a state of graceful shutdown or already closed
+		if t.state == draining || t.activeStreams == nil { // Can be niled from Close().
 			t.mu.Unlock()
-			return false // Don't create a stream if the transport is already closed.
+			return false
 		}
 		t.activeStreams[s.id] = s
 		t.mu.Unlock()
@@ -533,11 +536,17 @@ func (t *http2Client) CloseStream(s *Stream, err error) {
 	)
 	if err != nil {
 		rst = true
-		rstCode = http2.ErrCodeCancel
+		if errors.Is(err, errGracefulShutdown) {
+			rstCode = gracefulShutdownCode
+		} else {
+			rstCode = http2.ErrCodeCancel
+		}
 	}
 	t.closeStream(s, err, rst, rstCode, status.Convert(err), nil, false)
 }
 
+// before invoking closeStream, pls do not hold the t.mu
+// because accessing the controlbuf while holding t.mu will cause a deadlock.
 func (t *http2Client) closeStream(s *Stream, err error, rst bool, rstCode http2.ErrCode, st *status.Status, mdata map[string][]string, eosReceived bool) {
 	// Set stream status to done.
 	if s.swapState(streamDone) == streamDone {
@@ -617,7 +626,7 @@ func (t *http2Client) Close(err error) error {
 		t.kpDormancyCond.Signal()
 	}
 	t.mu.Unlock()
-	t.controlBuf.finish()
+	t.controlBuf.finish(err)
 	t.cancel()
 	cErr := t.conn.Close()
 
@@ -812,7 +821,13 @@ func (t *http2Client) handleRSTStream(f *http2.RSTStreamFrame) {
 			statusCode = codes.DeadlineExceeded
 		}
 	}
-	t.closeStream(s, io.EOF, false, http2.ErrCodeNo, status.Newf(statusCode, "stream terminated by RST_STREAM with error code: %v", f.ErrCode), nil, false)
+	var msg string
+	if f.ErrCode == gracefulShutdownCode {
+		msg = gracefulShutdownMsg
+	} else {
+		msg = fmt.Sprintf("stream terminated by RST_STREAM with error code: %v", f.ErrCode)
+	}
+	t.closeStream(s, io.EOF, false, http2.ErrCodeNo, status.New(statusCode, msg), nil, false)
 }
 
 func (t *http2Client) handleSettings(f *grpcframe.SettingsFrame, isFirst bool) {
@@ -917,29 +932,42 @@ func (t *http2Client) handleGoAway(f *grpcframe.GoAwayFrame) {
 		// Notify the clientconn about the GOAWAY before we set the state to
 		// draining, to allow the client to stop attempting to create streams
 		// before disallowing new streams on this connection.
-		if t.onGoAway != nil {
-			t.onGoAway(t.goAwayReason)
+		if t.state != draining {
+			if t.onGoAway != nil {
+				t.onGoAway(t.goAwayReason)
+			}
+			t.state = draining
 		}
-		t.state = draining
 	}
 	// All streams with IDs greater than the GoAwayId
 	// and smaller than the previous GoAway ID should be killed.
 	upperLimit := t.prevGoAwayID
 	if upperLimit == 0 { // This is the first GoAway Frame.
 		upperLimit = math.MaxUint32 // Kill all streams after the GoAway ID.
 	}
+	t.prevGoAwayID = id
+	active := len(t.activeStreams)
+	if active <= 0 {
+		t.mu.Unlock()
+		t.Close(connectionErrorf(true, nil, "received goaway and there are no active streams"))
+		return
+	}
+
+	var unprocessedStream []*Stream
 	for streamID, stream := range t.activeStreams {
 		if streamID > id && streamID <= upperLimit {
 			// The stream was unprocessed by the server.
 			atomic.StoreUint32(&stream.unprocessed, 1)
-			t.closeStream(stream, errStreamDrain, false, http2.ErrCodeNo, statusGoAway, nil, false)
+			unprocessedStream = append(unprocessedStream, stream)
 		}
 	}
-	t.prevGoAwayID = id
-	active := len(t.activeStreams)
 	t.mu.Unlock()
-	if active == 0 {
-		t.Close(connectionErrorf(true, nil, "received goaway and there are no active streams"))
+
+	// we should not access controlBuf with t.mu held since it will cause deadlock.
+	// Pls refer to checkForStreamQuota in NewStream, it gets the controlbuf.mu and
+	// wants to get the t.mu.
+	for _, stream := range unprocessedStream {
+		t.closeStream(stream, errStreamDrain, false, http2.ErrCodeNo, statusGoAway, nil, false)
 	}
 }