make gRPC connections synced

Signed-off-by: Yu Juncen <yu745514916@live.com>

br/pkg/backup/prepare_snap/env.go

@@ -17,6 +17,7 @@ package preparesnap
 import (
  "context"
  "slices"
+ "sync"
  "time"
 
  "github.com/docker/go-units"
@@ -110,6 +111,34 @@ func (c CliEnv) GetAllLiveStores(ctx context.Context) ([]*metapb.Store, error) {
  return withoutTiFlash, err
 }
 
+func AdaptForGRPCInTest(p PrepareClient) PrepareClient {
+ return &gRPCGoAdapter{
+ inner: p,
+ }
+}
+
+// GrpcGoAdapter makes the `Send` call synchronous.
+// grpc-go doesn't guarantee concurrency call to `Send` or `Recv` is safe.
+// But concurrency call to `send` and `recv` is safe.
+// This type is exported for testing.
+type gRPCGoAdapter struct {
+ inner PrepareClient
+ sendMu sync.Mutex
+ recvMu sync.Mutex
+}
+
+func (s *gRPCGoAdapter) Send(req *brpb.PrepareSnapshotBackupRequest) error {
+ s.sendMu.Lock()
+ defer s.sendMu.Unlock()
+ return s.inner.Send(req)
+}
+
+func (s *gRPCGoAdapter) Recv() (*brpb.PrepareSnapshotBackupResponse, error) {
+ s.recvMu.Lock()
+ defer s.recvMu.Unlock()
+ return s.inner.Recv()
+}
+
 func (c CliEnv) ConnectToStore(ctx context.Context, storeID uint64) (PrepareClient, error) {
  var cli brpb.Backup_PrepareSnapshotBackupClient
  err := c.Mgr.TryWithConn(ctx, storeID, func(cc *grpc.ClientConn) error {
@@ -124,7 +153,7 @@ func (c CliEnv) ConnectToStore(ctx context.Context, storeID uint64) (PrepareClie
  if err != nil {
  return nil, err
  }
- return cli, nil
+ return &gRPCGoAdapter{inner: cli}, nil
 }
 
 func (c CliEnv) LoadRegionsInKeyRange(ctx context.Context, startKey []byte, endKey []byte) (regions []Region, err error) {

br/pkg/backup/prepare_snap/prepare_test.go

@@ -161,7 +161,7 @@ func (m *mockStores) ConnectToStore(ctx context.Context, storeID uint64) (Prepar
  }
  m.onCreateStore(m.stores[storeID])
  }
- return m.stores[storeID], nil
+ return AdaptForGRPCInTest(m.stores[storeID]), nil
 }
 
 func (m *mockStores) LoadRegionsInKeyRange(ctx context.Context, startKey []byte, endKey []byte) (regions []Region, err error) {