feat(gRPC): support gRPC graceful shutdown

pkg/remote/trans/nphttp2/conn_pool.go

@@ -121,7 +121,7 @@ func (p *connPool) newTransport(ctx context.Context, dialer remote.Dialer, netwo
  opts,
  p.remoteService,
  func(grpc.GoAwayReason) {
- // do nothing
+ p.Clean(network, address)
  },
  func() {
  // do nothing

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

@@ -0,0 +1,52 @@
+/*
+ * 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"
+ "math"
+ "testing"
+ "time"
+
+ "github.com/cloudwego/kitex/internal/test"
+)
+
+func TestGracefulShutdown(t *testing.T) {
+ srv, cli := setUp(t, 0, math.MaxUint32, gracefulShutdown)
+ defer cli.Close(errSelfCloseForTest)
+
+ stream, err := cli.NewStream(context.Background(), &CallHdr{})
+ test.Assert(t, err == nil, err)
+ <-srv.srvReady
+ go srv.gracefulShutdown()
+ 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)
+ _, err = cli.NewStream(context.Background(), &CallHdr{})
+ test.Assert(t, err != nil, err)
+ t.Logf("NewStream err: %v", err)
+ time.Sleep(1 * time.Second)
+ err = cli.Write(stream, nil, []byte("hello"), &Options{})
+ test.Assert(t, err != nil, err)
+ t.Logf("After timeout, Write err: %v", err)
+ _, err = stream.Read(msg)
+ test.Assert(t, err != nil, err)
+ t.Logf("After timeout, Read err: %v", err)
+}

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

@@ -468,9 +468,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()
@@ -917,29 +918,40 @@ 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)
+ unprocessedStream = append(unprocessedStream, stream)
  t.closeStream(stream, errStreamDrain, false, http2.ErrCodeNo, statusGoAway, nil, false)
  }
  }
- 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"))
+
+ for _, stream := range unprocessedStream {
+ t.closeStream(stream, errStreamDrain, false, http2.ErrCodeNo, statusGoAway, nil, false)
  }
 }
 

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

@@ -1081,10 +1081,8 @@ func (t *http2Server) outgoingGoAwayHandler(g *goAway) (bool, error) {
  if err := t.framer.WriteGoAway(sid, g.code, g.debugData); err != nil {
  return false, err
  }
+ t.framer.writer.Flush()
  if g.closeConn {
- // Abruptly close the connection following the GoAway (via
- // loopywriter). But flush out what's inside the buffer first.
- t.framer.writer.Flush()
  return false, fmt.Errorf("transport: Connection closing")
  }
  return true, nil
@@ -1096,15 +1094,15 @@ func (t *http2Server) outgoingGoAwayHandler(g *goAway) (bool, error) {
  // originated before the GoAway reaches the client.
  // After getting the ack or timer expiration send out another GoAway this
  // time with an ID of the max stream server intends to process.
- if err := t.framer.WriteGoAway(math.MaxUint32, http2.ErrCodeNo, []byte{}); err != nil {
+ if err := t.framer.WriteGoAway(math.MaxUint32, http2.ErrCodeNo, g.debugData); err != nil {
  return false, err
  }
  if err := t.framer.WritePing(false, goAwayPing.data); err != nil {
  return false, err
  }
 
  gofunc.RecoverGoFuncWithInfo(context.Background(), func() {
- timer := time.NewTimer(time.Minute)
+ timer := time.NewTimer(10 * time.Second)
  defer timer.Stop()
  select {
  case <-t.drainChan:

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

@@ -57,6 +57,10 @@ type server struct {
  conns map[ServerTransport]bool
  h *testStreamHandler
  ready chan struct{}
+ hdlWG sync.WaitGroup
+ transWG sync.WaitGroup
+
+ srvReady chan struct{}
 }
 
 var (
@@ -77,6 +81,7 @@ func init() {
 
 type testStreamHandler struct {
  t *http2Server
+ srv *server
  notify chan struct{}
  getNotified chan struct{}
 }
@@ -92,6 +97,8 @@ const (
  invalidHeaderField
  delayRead
  pingpong
+
+ gracefulShutdown
 )
 
 func (h *testStreamHandler) handleStreamAndNotify(s *Stream) {
@@ -292,6 +299,20 @@ func (h *testStreamHandler) handleStreamDelayRead(t *testing.T, s *Stream) {
  }
 }
 
+func (h *testStreamHandler) gracefulShutdown(t *testing.T, s *Stream) {
+ close(h.srv.srvReady)
+ msg := make([]byte, 5)
+ num, err := s.Read(msg)
+ test.Assert(t, err == nil, err)
+ test.Assert(t, num == 5, num)
+ test.Assert(t, string(msg) == "hello", string(msg))
+ err = h.t.Write(s, nil, msg, &Options{})
+ test.Assert(t, err == nil, err)
+ _, err = s.Read(msg)
+ test.Assert(t, err != nil, err)
+ t.Logf("Server-side after timeout err: %v", err)
+}
+
 // start starts server. Other goroutines should block on s.readyChan for further operations.
 func (s *server) start(t *testing.T, port int, serverConfig *ServerConfig, ht hType) {
  // 创建 listener
@@ -329,6 +350,7 @@ func (s *server) start(t *testing.T, port int, serverConfig *ServerConfig, ht hT
  s.conns[transport] = true
  h := &testStreamHandler{t: transport.(*http2Server)}
  s.h = h
+ h.srv = s
  s.mu.Unlock()
  switch ht {
  case notifyCall:
@@ -379,12 +401,26 @@ func (s *server) start(t *testing.T, port int, serverConfig *ServerConfig, ht hT
  }, func(ctx context.Context, method string) context.Context {
  return ctx
  })
+ case gracefulShutdown:
+ s.transWG.Add(1)
+ go func() {
+ defer s.transWG.Done()
+ transport.HandleStreams(func(stream *Stream) {
+ s.hdlWG.Add(1)
+ go func() {
+ defer s.hdlWG.Done()
+ h.gracefulShutdown(t, stream)
+ }()
+ }, func(ctx context.Context, method string) context.Context { return ctx })
+ }()
  default:
- go transport.HandleStreams(func(s *Stream) {
- go h.handleStream(t, s)
- }, func(ctx context.Context, method string) context.Context {
- return ctx
- })
+ go func() {
+ transport.HandleStreams(func(s *Stream) {
+ go h.handleStream(t, s)
+ }, func(ctx context.Context, method string) context.Context {
+ return ctx
+ })
+ }()
  }
  return ctx
  }
@@ -434,6 +470,40 @@ func (s *server) stop() {
  s.mu.Unlock()
 }
 
+func (s *server) gracefulShutdown() {
+ ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+ defer cancel()
+ s.lis.Close()
+ s.mu.Lock()
+ for trans := range s.conns {
+ trans.Drain()
+ }
+ s.mu.Unlock()
+ timeout, _ := ctx.Deadline()
+ graceTimer := time.NewTimer(time.Until(timeout))
+ exitCh := make(chan struct{})
+ go func() {
+ select {
+ case <-graceTimer.C:
+ s.mu.Lock()
+ for trans := range s.conns {
+ trans.Close()
+ }
+ s.mu.Unlock()
+ return
+ case <-exitCh:
+ return
+ }
+ }()
+ s.hdlWG.Wait()
+ s.transWG.Wait()
+ close(exitCh)
+ s.conns = nil
+ if err := s.eventLoop.Shutdown(ctx); err != nil {
+ fmt.Printf("netpoll server exit failed, err=%v", err)
+ }
+}
+
 func (s *server) addr() string {
  if s.lis == nil {
  return ""
@@ -442,7 +512,7 @@ func (s *server) addr() string {
 }
 
 func setUpServerOnly(t *testing.T, port int, serverConfig *ServerConfig, ht hType) *server {
- server := &server{startedErr: make(chan error, 1), ready: make(chan struct{})}
+ server := &server{startedErr: make(chan error, 1), ready: make(chan struct{}), srvReady: make(chan struct{})}
  go server.start(t, port, serverConfig, ht)
  server.wait(t, time.Second)
  return server

pkg/remote/trans/nphttp2/server_handler.go

@@ -62,6 +62,7 @@ func newSvrTransHandler(opt *remote.ServerOption) (*svrTransHandler, error) {
  opt: opt,
  svcSearcher: opt.SvcSearcher,
  codec: grpc.NewGRPCCodec(grpc.WithThriftCodec(opt.PayloadCodec)),
+ transports: make(map[grpcTransport.ServerTransport]struct{}),
  }, nil
 }
 
@@ -72,6 +73,11 @@ type svrTransHandler struct {
  svcSearcher remote.ServiceSearcher
  inkHdlFunc endpoint.Endpoint
  codec remote.Codec
+ mu sync.Mutex
+ transports map[grpcTransport.ServerTransport]struct{}
+
+ hdlWG sync.WaitGroup
+ transWG sync.WaitGroup
 }
 
 var prefaceReadAtMost = func() int {
@@ -119,9 +125,11 @@ func (t *svrTransHandler) Read(ctx context.Context, conn net.Conn, msg remote.Me
 func (t *svrTransHandler) OnRead(ctx context.Context, conn net.Conn) error {
  svrTrans := ctx.Value(ctxKeySvrTransport).(*SvrTrans)
  tr := svrTrans.tr
-
+ defer t.transWG.Done()
  tr.HandleStreams(func(s *grpcTransport.Stream) {
+ t.hdlWG.Add(1)
  gofunc.GoFunc(ctx, func() {
+ defer t.hdlWG.Done()
  t.handleFunc(s, svrTrans, conn)
  })
  }, func(ctx context.Context, method string) context.Context {
@@ -315,6 +323,10 @@ func (t *svrTransHandler) OnActive(ctx context.Context, conn net.Conn) (context.
  if err != nil {
  return nil, err
  }
+ t.transWG.Add(1)
+ t.mu.Lock()
+ t.transports[tr] = struct{}{}
+ t.mu.Unlock()
  pool := &sync.Pool{
  New: func() interface{} {
  // init rpcinfo
@@ -329,6 +341,9 @@ func (t *svrTransHandler) OnActive(ctx context.Context, conn net.Conn) (context.
 // 连接关闭时回调
 func (t *svrTransHandler) OnInactive(ctx context.Context, conn net.Conn) {
  tr := ctx.Value(ctxKeySvrTransport).(*SvrTrans).tr
+ t.mu.Lock()
+ delete(t.transports, tr)
+ t.mu.Unlock()
  tr.Close()
 }
 
@@ -349,6 +364,43 @@ func (t *svrTransHandler) SetInvokeHandleFunc(inkHdlFunc endpoint.Endpoint) {
 func (t *svrTransHandler) SetPipeline(p *remote.TransPipeline) {
 }
 
+func (t *svrTransHandler) GracefulShutdown(ctx context.Context) error {
+ t.mu.Lock()
+ for trans := range t.transports {
+ trans.Drain()
+ }
+ t.mu.Unlock()
+
+ exitCh := make(chan struct{})
+ // todo: think about a better grace time duration
+ graceTime := time.Minute * 3
+ exitTimeout, ok := ctx.Deadline()
+ if ok {
+ graceTime = time.Until(exitTimeout)
+ }
+ graceTimer := time.NewTimer(graceTime)
+ gofunc.GoFunc(ctx, func() {
+ select {
+ case <-graceTimer.C:
+ t.mu.Lock()
+ for trans := range t.transports {
+ // use CloseWithErr
+ trans.Close()
+ }
+ t.mu.Unlock()
+ return
+ case <-exitCh:
+ return
+ }
+ })
+
+ t.hdlWG.Wait()
+ t.transWG.Wait()
+ close(exitCh)
+
+ return nil
+}
+
 func (t *svrTransHandler) startTracer(ctx context.Context, ri rpcinfo.RPCInfo) context.Context {
  c := t.opt.TracerCtl.DoStart(ctx, ri)
  return c