task_config_manager.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/chaosdaemon/tasks

     1  // Copyright 2021 Chaos Mesh Authors.
     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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  
    16  package tasks
    17  
    18  import (
    19  	"github.com/pkg/errors"
    20  
    21  	"github.com/chaos-mesh/chaos-mesh/pkg/cerr"
    22  )
    23  
    24  var ErrNotFoundID = cerr.NotFound("ID")
    25  var ErrNotFoundTaskID = cerr.NotFound("TaskID")
    26  var ErrNotFoundTypeTaskConfig = cerr.NotFoundType[TaskConfig]()
    27  
    28  var ErrDiffID = cerr.FromErr(errors.New("different IsID"))
    29  
    30  var ErrTaskConfigMapNotInit = cerr.NotInit[map[TaskID]TaskConfig]().WrapName("TaskConfigMap").Err()
    31  
    32  type IsID interface {
    33  	ToID() string
    34  }
    35  
    36  // Object ensure the outer config change will not change
    37  // the data inside the TaskManager.
    38  type Object interface {
    39  	DeepCopy() Object
    40  }
    41  
    42  // Mergeable introduces the data gathering ability.
    43  type Mergeable interface {
    44  	Merge(a Mergeable) error
    45  }
    46  
    47  // TaskConfig defines a composite of flexible config with an immutable target.
    48  // TaskConfig.Id is the ID of task.
    49  // TaskConfig.Data is the config provided by developer.
    50  type TaskConfig struct {
    51  	Id   IsID
    52  	Data Object
    53  }
    54  
    55  func NewTaskConfig(id IsID, data Object) TaskConfig {
    56  	return TaskConfig{
    57  		id,
    58  		data.DeepCopy(),
    59  	}
    60  }
    61  
    62  type TaskID = string
    63  
    64  // TaskConfigManager provides some basic methods on TaskConfig.
    65  // If developers wants to use MergeTaskConfig, they must implement Mergeable for the TaskConfig.
    66  type TaskConfigManager struct {
    67  	TaskConfigMap map[TaskID]TaskConfig
    68  }
    69  
    70  func NewTaskConfigManager() TaskConfigManager {
    71  	return TaskConfigManager{make(map[TaskID]TaskConfig)}
    72  }
    73  
    74  func (m TaskConfigManager) AddTaskConfig(id TaskID, task TaskConfig) error {
    75  	if m.TaskConfigMap == nil {
    76  		return ErrTaskConfigMapNotInit
    77  	}
    78  	if _, ok := m.TaskConfigMap[id]; ok {
    79  		return errors.Wrapf(cerr.ErrDuplicateEntity, "uid: %s, task: %v", id, task)
    80  	}
    81  	m.TaskConfigMap[id] = task
    82  	return nil
    83  }
    84  
    85  func (m TaskConfigManager) UpdateTaskConfig(id TaskID, task TaskConfig) (TaskConfig, error) {
    86  	if m.TaskConfigMap == nil {
    87  		return TaskConfig{}, ErrTaskConfigMapNotInit
    88  	}
    89  	taskOld, ok := m.TaskConfigMap[id]
    90  	if !ok {
    91  		return TaskConfig{}, ErrNotFoundTaskID.WrapInput(id).WrapInput(task).Err()
    92  	}
    93  	if taskOld.Id != task.Id {
    94  		return TaskConfig{}, ErrDiffID.Wrapf("expect: %v, input: %v", taskOld.Id, task.Id).Err()
    95  	}
    96  	m.TaskConfigMap[id] = task
    97  	return taskOld, nil
    98  }
    99  
   100  // DeleteTaskConfig Delete task inside the TaskConfigManager
   101  func (m TaskConfigManager) DeleteTaskConfig(id TaskID) error {
   102  	if m.TaskConfigMap == nil {
   103  		return ErrTaskConfigMapNotInit
   104  	}
   105  	_, ok := m.TaskConfigMap[id]
   106  	if !ok {
   107  		return ErrNotFoundTypeTaskConfig.WrapInput(id).Err()
   108  	}
   109  	delete(m.TaskConfigMap, id)
   110  	return nil
   111  }
   112  
   113  func (m TaskConfigManager) GetConfigWithUID(id TaskID) (TaskConfig, error) {
   114  	t, ok := m.TaskConfigMap[id]
   115  	if !ok {
   116  		return TaskConfig{}, ErrNotFoundTaskID.WrapInput(id).Err()
   117  	}
   118  	return t, nil
   119  }
   120  
   121  func (m TaskConfigManager) GetUIDsWithPID(id IsID) []TaskID {
   122  	uIds := make([]TaskID, 0)
   123  	for uid, task := range m.TaskConfigMap {
   124  		if task.Id == id {
   125  			uIds = append(uIds, uid)
   126  		}
   127  	}
   128  	return uIds
   129  }
   130  
   131  func (m TaskConfigManager) CheckTask(uid TaskID, pid IsID) error {
   132  	t, ok := m.TaskConfigMap[uid]
   133  	if !ok {
   134  		return ErrNotFoundTaskID.WrapInput(uid).Err()
   135  	}
   136  	if t.Id != pid {
   137  		return ErrDiffID.Wrapf("expect: %v, input: %v", t.Id, pid).Err()
   138  	}
   139  	return nil
   140  }
   141  
   142  // MergeTaskConfig will sum the TaskConfig with a same TaskConfig.Id.
   143  // If developers want to use it with type T, they must implement Mergeable for *T.
   144  // IMPORTANT: Just here , we do not assume A.Merge(B) == B.Merge(A).
   145  // What MergeTaskConfig do : A := new(TaskConfig), A.Merge(B).Merge(C).Merge(D)... , A marked as uid.
   146  func (m TaskConfigManager) MergeTaskConfig(uid TaskID) (TaskConfig, error) {
   147  	if m.TaskConfigMap == nil {
   148  		return TaskConfig{}, ErrTaskConfigMapNotInit
   149  	}
   150  	taskRaw, ok := m.TaskConfigMap[uid]
   151  	if !ok {
   152  		return TaskConfig{}, ErrNotFoundTaskID.WrapInput(uid).Err()
   153  	}
   154  
   155  	task := TaskConfig{
   156  		Id:   taskRaw.Id,
   157  		Data: taskRaw.Data.DeepCopy(),
   158  	}
   159  	uids := m.GetUIDsWithPID(task.Id)
   160  
   161  	for _, uidTemp := range uids {
   162  		if uid == uidTemp {
   163  			continue
   164  		}
   165  		taskTemp, ok := m.TaskConfigMap[uidTemp]
   166  		if !ok {
   167  			return TaskConfig{}, ErrNotFoundTypeTaskConfig.WrapInput(uidTemp).Err()
   168  		}
   169  		AddableData, ok := task.Data.(Mergeable)
   170  		if !ok {
   171  			return TaskConfig{}, cerr.NotImpl[Mergeable]().WrapInput(task.Data).Err()
   172  		}
   173  		AddableTempData, ok := taskTemp.Data.(Mergeable)
   174  		if !ok {
   175  			return TaskConfig{}, cerr.NotImpl[Mergeable]().WrapInput(taskTemp.Data).Err()
   176  		}
   177  		err := AddableData.Merge(AddableTempData)
   178  		if err != nil {
   179  			return TaskConfig{}, err
   180  		}
   181  	}
   182  	return task, nil
   183  }
   184