github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/client/mock_executor_group.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 client
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  	"time"
    20  
    21  	"github.com/pingcap/tiflow/engine/model"
    22  	"github.com/pingcap/tiflow/pkg/errors"
    23  )
    24  
    25  // MockExecutorGroup is a stub implementation for ExecutorGroup.
    26  type MockExecutorGroup struct {
    27  	mu      sync.RWMutex
    28  	clients map[model.ExecutorID]ExecutorClient
    29  
    30  	updateNotifyCh chan struct{}
    31  }
    32  
    33  // NewMockExecutorGroup returns a new MockExecutorGroup.
    34  func NewMockExecutorGroup() *MockExecutorGroup {
    35  	return &MockExecutorGroup{
    36  		clients:        make(map[model.ExecutorID]ExecutorClient),
    37  		updateNotifyCh: make(chan struct{}, 1),
    38  	}
    39  }
    40  
    41  // GetExecutorClient returns the ExecutorClient associated with id.
    42  func (g *MockExecutorGroup) GetExecutorClient(id model.ExecutorID) (ExecutorClient, bool) {
    43  	g.mu.RLock()
    44  	defer g.mu.RUnlock()
    45  
    46  	if client, ok := g.clients[id]; ok {
    47  		return client, true
    48  	}
    49  	return nil, false
    50  }
    51  
    52  // GetExecutorClientB tries to get the ExecutorClient blockingly.
    53  func (g *MockExecutorGroup) GetExecutorClientB(ctx context.Context, id model.ExecutorID) (ExecutorClient, error) {
    54  	ticker := time.NewTicker(100 * time.Millisecond)
    55  	defer ticker.Stop()
    56  
    57  	for {
    58  		select {
    59  		case <-ctx.Done():
    60  			return nil, errors.Trace(ctx.Err())
    61  		case <-ticker.C:
    62  		case <-g.updateNotifyCh:
    63  		}
    64  
    65  		if client, ok := g.GetExecutorClient(id); ok {
    66  			return client, nil
    67  		}
    68  	}
    69  }
    70  
    71  // AddClient adds a client to the client map.
    72  func (g *MockExecutorGroup) AddClient(id model.ExecutorID, client ExecutorClient) {
    73  	g.mu.Lock()
    74  	defer g.mu.Unlock()
    75  
    76  	g.clients[id] = client
    77  
    78  	select {
    79  	case g.updateNotifyCh <- struct{}{}:
    80  	default:
    81  	}
    82  }
    83  
    84  // RemoveClient removes a client from the client map.
    85  func (g *MockExecutorGroup) RemoveClient(id model.ExecutorID) bool {
    86  	g.mu.Lock()
    87  	defer g.mu.Unlock()
    88  
    89  	client, exists := g.clients[id]
    90  	if !exists {
    91  		return false
    92  	}
    93  
    94  	client.Close()
    95  	delete(g.clients, id)
    96  
    97  	return true
    98  }