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

     1  // Copyright 2022 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/edwingeng/deque"
    21  	"github.com/pingcap/tiflow/pkg/errors"
    22  )
    23  
    24  // MockMessageSender defines a mock message sender
    25  type MockMessageSender struct {
    26  	mu            sync.Mutex
    27  	msgBox        map[msgBoxIndex]deque.Deque
    28  	nodeBlackList map[NodeID]struct{}
    29  	isBlocked     bool
    30  
    31  	injectedErrCh chan error
    32  }
    33  
    34  // NewMockMessageSender creates a new MockMessageSender instance
    35  func NewMockMessageSender() *MockMessageSender {
    36  	return &MockMessageSender{
    37  		msgBox:        make(map[msgBoxIndex]deque.Deque),
    38  		nodeBlackList: make(map[NodeID]struct{}),
    39  		injectedErrCh: make(chan error, 1),
    40  	}
    41  }
    42  
    43  type msgBoxIndex struct {
    44  	topic  Topic
    45  	target NodeID
    46  }
    47  
    48  // SendToNodeB implements pkg/p2p.MessageSender.SendToNodeB
    49  func (m *MockMessageSender) SendToNodeB(
    50  	ctx context.Context,
    51  	targetNodeID NodeID,
    52  	topic Topic,
    53  	message interface{},
    54  ) error {
    55  	m.mu.Lock()
    56  	defer m.mu.Unlock()
    57  
    58  	// Handle mock offline nodes.
    59  	if _, exists := m.nodeBlackList[targetNodeID]; exists {
    60  		return errors.ErrExecutorNotFoundForMessage.GenWithStackByArgs()
    61  	}
    62  
    63  	select {
    64  	case err := <-m.injectedErrCh:
    65  		return err
    66  	default:
    67  	}
    68  
    69  	// TODO Handle the `m.isBlocked == true` case
    70  	q := m.getQueue(targetNodeID, topic)
    71  	q.PushBack(message)
    72  	return nil
    73  }
    74  
    75  // SendToNode implements pkg/p2p.MessageSender.SendToNode
    76  func (m *MockMessageSender) SendToNode(
    77  	_ context.Context,
    78  	targetNodeID NodeID,
    79  	topic Topic,
    80  	message interface{},
    81  ) (bool, error) {
    82  	m.mu.Lock()
    83  	defer m.mu.Unlock()
    84  
    85  	// Handle mock offline nodes.
    86  	if _, exists := m.nodeBlackList[targetNodeID]; exists {
    87  		return false, nil
    88  	}
    89  
    90  	select {
    91  	case err := <-m.injectedErrCh:
    92  		return false, err
    93  	default:
    94  	}
    95  
    96  	if m.isBlocked {
    97  		return false, nil
    98  	}
    99  
   100  	q := m.getQueue(targetNodeID, topic)
   101  	q.PushBack(message)
   102  
   103  	return true, nil
   104  }
   105  
   106  // TryPop tries to get a message from message sender
   107  func (m *MockMessageSender) TryPop(targetNodeID NodeID, topic Topic) (interface{}, bool) {
   108  	m.mu.Lock()
   109  	defer m.mu.Unlock()
   110  
   111  	q := m.getQueue(targetNodeID, topic)
   112  	if q.Empty() {
   113  		return nil, false
   114  	}
   115  
   116  	return q.PopFront(), true
   117  }
   118  
   119  func (m *MockMessageSender) getQueue(target NodeID, topic Topic) deque.Deque {
   120  	mapKey := msgBoxIndex{
   121  		topic:  topic,
   122  		target: target,
   123  	}
   124  
   125  	q, ok := m.msgBox[mapKey]
   126  	if !ok {
   127  		q = deque.NewDeque()
   128  		m.msgBox[mapKey] = q
   129  	}
   130  
   131  	return q
   132  }
   133  
   134  // SetBlocked makes the message send blocking
   135  func (m *MockMessageSender) SetBlocked(isBlocked bool) {
   136  	m.mu.Lock()
   137  	defer m.mu.Unlock()
   138  
   139  	m.isBlocked = isBlocked
   140  }
   141  
   142  // InjectError injects error to simulate error scenario
   143  func (m *MockMessageSender) InjectError(err error) {
   144  	m.injectedErrCh <- err
   145  }
   146  
   147  // MarkNodeOffline marks a node as offline.
   148  func (m *MockMessageSender) MarkNodeOffline(nodeID NodeID) {
   149  	m.mu.Lock()
   150  	defer m.mu.Unlock()
   151  
   152  	m.nodeBlackList[nodeID] = struct{}{}
   153  }
   154  
   155  // MarkNodeOnline marks a node as online.
   156  // Note that by default all nodes are treated as online
   157  // to facilitate testing.
   158  func (m *MockMessageSender) MarkNodeOnline(nodeID NodeID) {
   159  	m.mu.Lock()
   160  	defer m.mu.Unlock()
   161  
   162  	delete(m.nodeBlackList, nodeID)
   163  }