github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/p2p/server_wrapper.go (about)

     1  // Copyright 2021 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package p2p
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  
    20  	"github.com/modern-go/reflect2"
    21  	"github.com/pingcap/failpoint"
    22  	"github.com/pingcap/log"
    23  	"github.com/pingcap/tiflow/proto/p2p"
    24  	"go.uber.org/zap"
    25  	"google.golang.org/grpc"
    26  	"google.golang.org/grpc/codes"
    27  	"google.golang.org/grpc/keepalive"
    28  	gRPCPeer "google.golang.org/grpc/peer"
    29  	"google.golang.org/grpc/status"
    30  )
    31  
    32  type streamWrapper struct {
    33  	MessageServerStream
    34  	ctx    context.Context
    35  	cancel context.CancelFunc
    36  }
    37  
    38  func wrapStream(stream MessageServerStream) *streamWrapper {
    39  	ctx, cancel := context.WithCancel(stream.Context())
    40  	return &streamWrapper{
    41  		MessageServerStream: stream,
    42  		ctx:                 ctx,
    43  		cancel:              cancel,
    44  	}
    45  }
    46  
    47  func (w *streamWrapper) Context() context.Context {
    48  	return w.ctx
    49  }
    50  
    51  // ServerWrapper implements a CDCPeerToPeerServer, and it
    52  // maintains an inner CDCPeerToPeerServer instance that can
    53  // be replaced as needed.
    54  type ServerWrapper struct {
    55  	rwMu        sync.RWMutex
    56  	innerServer p2p.CDCPeerToPeerServer
    57  	cfg         *MessageServerConfig
    58  
    59  	wrappedStreamsMu sync.Mutex
    60  	wrappedStreams   map[*streamWrapper]struct{}
    61  }
    62  
    63  // NewServerWrapper creates a new ServerWrapper
    64  func NewServerWrapper(cfg *MessageServerConfig) *ServerWrapper {
    65  	return &ServerWrapper{
    66  		wrappedStreams: map[*streamWrapper]struct{}{},
    67  		cfg:            cfg,
    68  	}
    69  }
    70  
    71  // ServerOptions returns server option for creating grpc servers.
    72  func (s *ServerWrapper) ServerOptions() []grpc.ServerOption {
    73  	keepaliveParams := keepalive.ServerParameters{
    74  		Time:    s.cfg.KeepAliveTime,
    75  		Timeout: s.cfg.KeepAliveTimeout,
    76  	}
    77  	return []grpc.ServerOption{
    78  		grpc.MaxRecvMsgSize(s.cfg.MaxRecvMsgSize),
    79  		grpc.KeepaliveParams(keepaliveParams),
    80  	}
    81  }
    82  
    83  // SendMessage implements p2p.CDCPeerToPeerServer
    84  func (s *ServerWrapper) SendMessage(stream p2p.CDCPeerToPeer_SendMessageServer) error {
    85  	s.rwMu.RLock()
    86  	innerServer := s.innerServer
    87  	s.rwMu.RUnlock()
    88  
    89  	if innerServer == nil {
    90  		var addr string
    91  		peer, ok := gRPCPeer.FromContext(stream.Context())
    92  		if ok {
    93  			addr = peer.Addr.String()
    94  		}
    95  		log.Debug("gRPC server received request while message server is not running.", zap.String("addr", addr))
    96  		return status.New(codes.Unavailable, "message server is not running").Err()
    97  	}
    98  
    99  	wrappedStream := wrapStream(stream)
   100  	s.wrappedStreamsMu.Lock()
   101  	s.wrappedStreams[wrappedStream] = struct{}{}
   102  	s.wrappedStreamsMu.Unlock()
   103  	defer func() {
   104  		s.wrappedStreamsMu.Lock()
   105  		delete(s.wrappedStreams, wrappedStream)
   106  		s.wrappedStreamsMu.Unlock()
   107  		wrappedStream.cancel()
   108  	}()
   109  
   110  	// Used in unit tests to simulate a race situation between `SendMessage` and `Reset`.
   111  	// TODO think of another way to make tests parallelizable.
   112  	failpoint.Inject("ServerWrapperSendMessageDelay", func() {})
   113  	return innerServer.SendMessage(wrappedStream)
   114  }
   115  
   116  // Reset resets the inner server object in the ServerWrapper
   117  func (s *ServerWrapper) Reset(inner p2p.CDCPeerToPeerServer) {
   118  	s.rwMu.Lock()
   119  	defer s.rwMu.Unlock()
   120  
   121  	s.wrappedStreamsMu.Lock()
   122  	defer s.wrappedStreamsMu.Unlock()
   123  
   124  	for wrappedStream := range s.wrappedStreams {
   125  		wrappedStream.cancel()
   126  	}
   127  
   128  	// reflect2.IsNil handles two cases for us:
   129  	// 1) null value
   130  	// 2) an interface with a null value but a not-null type info.
   131  	if reflect2.IsNil(inner) {
   132  		s.innerServer = nil
   133  		return
   134  	}
   135  	s.innerServer = inner
   136  }