github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/p2p/mock_cluster.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  	"fmt"
    19  	"net"
    20  	"sync"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/phayes/freeport"
    25  	"github.com/pingcap/tiflow/pkg/security"
    26  	proto "github.com/pingcap/tiflow/proto/p2p"
    27  	"github.com/stretchr/testify/require"
    28  	"google.golang.org/grpc"
    29  )
    30  
    31  // MockCluster mocks the whole peer-messaging cluster.
    32  type MockCluster struct {
    33  	Nodes map[NodeID]*MockNode
    34  }
    35  
    36  // MockNode represents one mock node.
    37  type MockNode struct {
    38  	Addr string
    39  	ID   NodeID
    40  
    41  	Server *MessageServer
    42  	Router MessageRouter
    43  
    44  	cancel func()
    45  	wg     sync.WaitGroup
    46  }
    47  
    48  // read only
    49  var serverConfig4MockCluster = &MessageServerConfig{
    50  	MaxPendingMessageCountPerTopic:       256,
    51  	MaxPendingTaskCount:                  102400,
    52  	SendChannelSize:                      1,
    53  	AckInterval:                          time.Millisecond * 100,
    54  	WorkerPoolSize:                       4,
    55  	MaxPeerCount:                         1024,
    56  	WaitUnregisterHandleTimeoutThreshold: time.Millisecond * 100,
    57  }
    58  
    59  // read only
    60  var clientConfig4MockCluster = &MessageClientConfig{
    61  	SendChannelSize:         1,
    62  	BatchSendInterval:       time.Millisecond * 100,
    63  	MaxBatchCount:           128,
    64  	MaxBatchBytes:           8192,
    65  	RetryRateLimitPerSecond: 10.0, // using 10.0 instead of 1.0 to accelerate testing
    66  	DialTimeout:             time.Second * 3,
    67  	MaxRecvMsgSize:          4 * 1024 * 1024, // 4MB
    68  }
    69  
    70  func newMockNode(t *testing.T, id NodeID) *MockNode {
    71  	port := freeport.GetPort()
    72  	addr := fmt.Sprintf("127.0.0.1:%d", port)
    73  
    74  	ctx, cancel := context.WithCancel(context.Background())
    75  
    76  	ret := &MockNode{
    77  		Addr:   addr,
    78  		ID:     id,
    79  		Server: NewMessageServer(id, serverConfig4MockCluster),
    80  		// Note that TLS is disabled.
    81  		Router: NewMessageRouter(id, &security.Credential{}, clientConfig4MockCluster),
    82  		cancel: cancel,
    83  	}
    84  
    85  	grpcServer := grpc.NewServer()
    86  	proto.RegisterCDCPeerToPeerServer(grpcServer, ret.Server)
    87  
    88  	ret.wg.Add(1)
    89  	go func() {
    90  		defer ret.wg.Done()
    91  		lis, err := net.Listen("tcp", addr)
    92  		require.NoError(t, err)
    93  		_ = grpcServer.Serve(lis)
    94  	}()
    95  
    96  	ret.wg.Add(1)
    97  	go func() {
    98  		defer ret.wg.Done()
    99  		err := ret.Server.Run(ctx, nil)
   100  		require.Error(t, err)
   101  		require.Regexp(t, ".*context canceled.*", err.Error())
   102  	}()
   103  
   104  	ret.wg.Add(1)
   105  	go func() {
   106  		defer ret.wg.Done()
   107  		<-ctx.Done()
   108  		grpcServer.Stop()
   109  	}()
   110  
   111  	return ret
   112  }
   113  
   114  // Close closes the mock node.
   115  func (n *MockNode) Close() {
   116  	n.Router.Close()
   117  	n.cancel()
   118  	n.wg.Wait()
   119  }
   120  
   121  // NewMockCluster creates a mock cluster.
   122  func NewMockCluster(t *testing.T, nodeCount int) *MockCluster {
   123  	ret := &MockCluster{
   124  		Nodes: make(map[NodeID]*MockNode),
   125  	}
   126  
   127  	for i := 0; i < nodeCount; i++ {
   128  		id := fmt.Sprintf("capture-%d", i)
   129  		ret.Nodes[id] = newMockNode(t, id)
   130  	}
   131  
   132  	for _, sourceNode := range ret.Nodes {
   133  		for _, targetNode := range ret.Nodes {
   134  			sourceNode.Router.AddPeer(targetNode.ID, targetNode.Addr)
   135  		}
   136  	}
   137  
   138  	return ret
   139  }
   140  
   141  // Close closes the mock cluster.
   142  func (c *MockCluster) Close() {
   143  	for _, node := range c.Nodes {
   144  		node.Close()
   145  	}
   146  }