github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/internal/master/mock_handle.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 master
    15  
    16  import (
    17  	"context"
    18  
    19  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    20  	"github.com/pingcap/tiflow/engine/model"
    21  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    22  	"github.com/pingcap/tiflow/pkg/errors"
    23  	"go.uber.org/atomic"
    24  )
    25  
    26  // MockHandle implements WorkerHandle, it can work as either a RunningHandle or
    27  // a TombstoneHandle.
    28  type MockHandle struct {
    29  	WorkerID      frameModel.WorkerID
    30  	WorkerStatus  *frameModel.WorkerStatus
    31  	ExecutorID    model.ExecutorID
    32  	MessageSender *p2p.MockMessageSender
    33  	IsTombstone   bool
    34  
    35  	sendMessageCount atomic.Int64
    36  }
    37  
    38  // GetTombstone implements WorkerHandle.GetTombstone
    39  func (h *MockHandle) GetTombstone() TombstoneHandle {
    40  	if h.IsTombstone {
    41  		return h
    42  	}
    43  	return nil
    44  }
    45  
    46  // Unwrap implements WorkerHandle.Unwrap
    47  func (h *MockHandle) Unwrap() RunningHandle {
    48  	if !h.IsTombstone {
    49  		return mockRunningHandle{
    50  			BaseHandle: h,
    51  			handler:    h,
    52  		}
    53  	}
    54  	return nil
    55  }
    56  
    57  // Status implements WorkerHandle.Status
    58  func (h *MockHandle) Status() *frameModel.WorkerStatus {
    59  	return h.WorkerStatus
    60  }
    61  
    62  // ID implements WorkerHandle.ID
    63  func (h *MockHandle) ID() frameModel.WorkerID {
    64  	return h.WorkerID
    65  }
    66  
    67  // CleanTombstone implements TombstoneHandle.CleanTombstone
    68  func (h *MockHandle) CleanTombstone(ctx context.Context) error {
    69  	// TODO implement me
    70  	panic("implement me")
    71  }
    72  
    73  // SendMessageCount returns the send message count, used in unit test only.
    74  func (h *MockHandle) SendMessageCount() int {
    75  	return int(h.sendMessageCount.Load())
    76  }
    77  
    78  type mockRunningHandle struct {
    79  	BaseHandle
    80  	handler *MockHandle
    81  }
    82  
    83  // SendMessage implements RunningHandle.SendMessage
    84  func (rh mockRunningHandle) SendMessage(ctx context.Context, topic p2p.Topic, message interface{}, nonblocking bool) error {
    85  	h := rh.handler
    86  	if h.IsTombstone {
    87  		return errors.ErrSendingMessageToTombstone.GenWithStackByCause(h.WorkerID)
    88  	}
    89  
    90  	h.sendMessageCount.Add(1)
    91  	if h.MessageSender == nil {
    92  		return nil
    93  	}
    94  
    95  	var err error
    96  	if nonblocking {
    97  		_, err = h.MessageSender.SendToNode(ctx, p2p.NodeID(h.ExecutorID), topic, message)
    98  	} else {
    99  		err = h.MessageSender.SendToNodeB(ctx, p2p.NodeID(h.ExecutorID), topic, message)
   100  	}
   101  	return err
   102  }