github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/worker/subtask_holder.go (about)

     1  // Copyright 2019 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  	"github.com/pingcap/tiflow/dm/relay"
    21  )
    22  
    23  // subTaskHolder holds subtask instances.
    24  type subTaskHolder struct {
    25  	// protect concurrent access of `subTasks`
    26  	// the internal operation of the subtask should be protected by any mechanism in the subtask itself.
    27  	mu       sync.RWMutex
    28  	subTasks map[string]*SubTask
    29  }
    30  
    31  // newSubTaskHolder creates a new subTaskHolder instance.
    32  func newSubTaskHolder() *subTaskHolder {
    33  	return &subTaskHolder{
    34  		subTasks: make(map[string]*SubTask, 5),
    35  	}
    36  }
    37  
    38  // recordSubTask records subtask instance.
    39  func (h *subTaskHolder) recordSubTask(st *SubTask) {
    40  	h.mu.Lock()
    41  	defer h.mu.Unlock()
    42  	h.subTasks[st.cfg.Name] = st
    43  }
    44  
    45  // removeSubTask removes subtask instance.
    46  func (h *subTaskHolder) removeSubTask(name string) {
    47  	h.mu.Lock()
    48  	defer h.mu.Unlock()
    49  	delete(h.subTasks, name)
    50  }
    51  
    52  // resetAllSubTasks does Close, change cfg.UseRelay then Init the subtasks.
    53  func (h *subTaskHolder) resetAllSubTasks(relay relay.Process) {
    54  	h.mu.Lock()
    55  	defer h.mu.Unlock()
    56  	for _, st := range h.subTasks {
    57  		stage := st.Stage()
    58  		st.Close()
    59  		// TODO: make a st.reset
    60  		st.ctx, st.cancel = context.WithCancel(context.Background())
    61  		st.cfg.UseRelay = relay != nil
    62  		st.Run(stage, st.getValidatorStage(), relay)
    63  	}
    64  }
    65  
    66  // closeAllSubTasks closes all subtask instances.
    67  func (h *subTaskHolder) closeAllSubTasks() {
    68  	h.mu.Lock()
    69  	defer h.mu.Unlock()
    70  	for _, st := range h.subTasks {
    71  		st.Close()
    72  	}
    73  	h.subTasks = make(map[string]*SubTask)
    74  }
    75  
    76  // killAllSubTasks kill and stop all subtask instances.
    77  func (h *subTaskHolder) killAllSubTasks() {
    78  	h.mu.Lock()
    79  	defer h.mu.Unlock()
    80  	for _, st := range h.subTasks {
    81  		st.Kill()
    82  	}
    83  	h.subTasks = make(map[string]*SubTask)
    84  }
    85  
    86  // findSubTask finds subtask instance by name.
    87  func (h *subTaskHolder) findSubTask(name string) *SubTask {
    88  	h.mu.RLock()
    89  	defer h.mu.RUnlock()
    90  	return h.subTasks[name]
    91  }
    92  
    93  // getAllSubTasks returns all subtask instances.
    94  func (h *subTaskHolder) getAllSubTasks() map[string]*SubTask {
    95  	h.mu.RLock()
    96  	defer h.mu.RUnlock()
    97  	result := make(map[string]*SubTask, len(h.subTasks))
    98  	for name, st := range h.subTasks {
    99  		result[name] = st
   100  	}
   101  	return result
   102  }