github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/internal/worker/master_info.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 worker
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  
    20  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    21  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    22  	"go.uber.org/atomic"
    23  )
    24  
    25  // MasterInfoProvider is an object that can provide the caller
    26  // information on the master.
    27  type MasterInfoProvider interface {
    28  	MasterID() frameModel.MasterID
    29  	MasterNode() p2p.NodeID
    30  	Epoch() frameModel.Epoch
    31  	SyncRefreshMasterInfo(ctx context.Context) error
    32  	IsMasterSideClosed() bool
    33  }
    34  
    35  // MasterClient must implement MasterInfoProvider.
    36  var _ MasterInfoProvider = (*MasterClient)(nil)
    37  
    38  // MockMasterInfoProvider defines a mock provider that implements MasterInfoProvider
    39  type MockMasterInfoProvider struct {
    40  	mu         sync.RWMutex
    41  	masterID   frameModel.MasterID
    42  	masterNode p2p.NodeID
    43  	epoch      frameModel.Epoch
    44  
    45  	refreshCount     atomic.Int64
    46  	masterSideClosed atomic.Bool
    47  }
    48  
    49  // NewMockMasterInfoProvider creates a new MockMasterInfoProvider
    50  func NewMockMasterInfoProvider(
    51  	masterID frameModel.MasterID,
    52  	masterNode p2p.NodeID,
    53  	epoch frameModel.Epoch,
    54  ) *MockMasterInfoProvider {
    55  	return &MockMasterInfoProvider{
    56  		masterID:   masterID,
    57  		masterNode: masterNode,
    58  		epoch:      epoch,
    59  	}
    60  }
    61  
    62  // MasterID implements MasterInfoProvider.MasterID
    63  func (p *MockMasterInfoProvider) MasterID() frameModel.MasterID {
    64  	p.mu.RLock()
    65  	defer p.mu.RUnlock()
    66  
    67  	return p.masterID
    68  }
    69  
    70  // MasterNode implements MasterInfoProvider.MasterNode
    71  func (p *MockMasterInfoProvider) MasterNode() p2p.NodeID {
    72  	p.mu.RLock()
    73  	defer p.mu.RUnlock()
    74  
    75  	return p.masterNode
    76  }
    77  
    78  // Epoch implements MasterInfoProvider.Epoch
    79  func (p *MockMasterInfoProvider) Epoch() frameModel.Epoch {
    80  	p.mu.RLock()
    81  	defer p.mu.RUnlock()
    82  
    83  	return p.epoch
    84  }
    85  
    86  // SyncRefreshMasterInfo implements MasterInfoProvider.RefreshMasterInfo
    87  func (p *MockMasterInfoProvider) SyncRefreshMasterInfo(ctx context.Context) error {
    88  	p.refreshCount.Add(1)
    89  	return nil
    90  }
    91  
    92  // Set sets given information to the MockMasterInfoProvider
    93  func (p *MockMasterInfoProvider) Set(masterID frameModel.MasterID, masterNode p2p.NodeID, epoch frameModel.Epoch) {
    94  	p.mu.Lock()
    95  	defer p.mu.Unlock()
    96  
    97  	p.masterID = masterID
    98  	p.masterNode = masterNode
    99  	p.epoch = epoch
   100  }
   101  
   102  // SetMasterClosed marks the mock master has having marked the current
   103  // worker closed.
   104  func (p *MockMasterInfoProvider) SetMasterClosed() {
   105  	p.masterSideClosed.Store(true)
   106  }
   107  
   108  // RefreshCount returns refresh time, it is used in unit test only
   109  func (p *MockMasterInfoProvider) RefreshCount() int {
   110  	return int(p.refreshCount.Load())
   111  }
   112  
   113  // IsMasterSideClosed implements MasterInfoProvider.IsMasterSideClosed
   114  func (p *MockMasterInfoProvider) IsMasterSideClosed() bool {
   115  	return p.masterSideClosed.Load()
   116  }