github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/internal/master/worker_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  	"github.com/pingcap/log"
    20  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    21  	"github.com/pingcap/tiflow/engine/model"
    22  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    23  	"go.uber.org/zap"
    24  )
    25  
    26  // BaseHandle provides some common api of a worker, no matter it is running or dead.
    27  type BaseHandle interface {
    28  	Status() *frameModel.WorkerStatus
    29  	ID() frameModel.WorkerID
    30  }
    31  
    32  // WorkerHandle defines the interface of a worker, businiss logic can use this
    33  // handler to get RunningHandle or TombstoneHandle
    34  type WorkerHandle interface {
    35  	BaseHandle
    36  
    37  	GetTombstone() TombstoneHandle
    38  	Unwrap() RunningHandle
    39  }
    40  
    41  // RunningHandle represents a running worker
    42  type RunningHandle interface {
    43  	BaseHandle
    44  
    45  	SendMessage(
    46  		ctx context.Context,
    47  		topic p2p.Topic,
    48  		message interface{},
    49  		nonblocking bool,
    50  	) error
    51  }
    52  
    53  // TombstoneHandle represents a dead worker.
    54  type TombstoneHandle interface {
    55  	BaseHandle
    56  
    57  	// CleanTombstone cleans the metadata from the metastore,
    58  	// and cleans the state managed by the framework.
    59  	// Do not call any other methods on this handle after
    60  	// CleanTombstone is called.
    61  	CleanTombstone(ctx context.Context) error
    62  }
    63  
    64  type runningHandleImpl struct {
    65  	workerID   frameModel.WorkerID
    66  	executorID model.ExecutorID
    67  	manager    *WorkerManager
    68  }
    69  
    70  func (h *runningHandleImpl) Status() *frameModel.WorkerStatus {
    71  	h.manager.mu.Lock()
    72  	defer h.manager.mu.Unlock()
    73  
    74  	entry, exists := h.manager.workerEntries[h.workerID]
    75  	if !exists {
    76  		log.Panic("Using a stale handle", zap.String("worker-id", h.workerID))
    77  	}
    78  
    79  	return entry.Status()
    80  }
    81  
    82  func (h *runningHandleImpl) ID() frameModel.WorkerID {
    83  	return h.workerID
    84  }
    85  
    86  func (h *runningHandleImpl) GetTombstone() TombstoneHandle {
    87  	return nil
    88  }
    89  
    90  func (h *runningHandleImpl) Unwrap() RunningHandle {
    91  	return h
    92  }
    93  
    94  func (h *runningHandleImpl) SendMessage(
    95  	ctx context.Context,
    96  	topic p2p.Topic,
    97  	message interface{},
    98  	nonblocking bool,
    99  ) error {
   100  	var err error
   101  	if nonblocking {
   102  		_, err = h.manager.messageSender.SendToNode(ctx, p2p.NodeID(h.executorID), topic, message)
   103  	} else {
   104  		err = h.manager.messageSender.SendToNodeB(ctx, p2p.NodeID(h.executorID), topic, message)
   105  	}
   106  
   107  	return err
   108  }
   109  
   110  type tombstoneHandleImpl struct {
   111  	workerID frameModel.WorkerID
   112  	manager  *WorkerManager
   113  }
   114  
   115  func (h *tombstoneHandleImpl) Status() *frameModel.WorkerStatus {
   116  	h.manager.mu.Lock()
   117  	defer h.manager.mu.Unlock()
   118  
   119  	entry, exists := h.manager.workerEntries[h.workerID]
   120  	if !exists {
   121  		log.Panic("Using a stale handle", zap.String("worker-id", h.workerID))
   122  	}
   123  
   124  	return entry.Status()
   125  }
   126  
   127  func (h *tombstoneHandleImpl) ID() frameModel.WorkerID {
   128  	return h.workerID
   129  }
   130  
   131  func (h *tombstoneHandleImpl) GetTombstone() TombstoneHandle {
   132  	return h
   133  }
   134  
   135  func (h *tombstoneHandleImpl) Unwrap() RunningHandle {
   136  	return nil
   137  }
   138  
   139  func (h *tombstoneHandleImpl) CleanTombstone(ctx context.Context) error {
   140  	ok, err := h.manager.workerMetaClient.Remove(ctx, h.workerID)
   141  	if err != nil {
   142  		return err
   143  	}
   144  	if !ok {
   145  		log.Info("Tombstone already cleaned", zap.String("worker-id", h.workerID))
   146  		// Idempotent for robustness.
   147  		return nil
   148  	}
   149  	log.Info("Worker tombstone is cleaned", zap.String("worker-id", h.workerID))
   150  	h.manager.removeTombstoneEntry(h.workerID)
   151  
   152  	return nil
   153  }