1
2
3
4
5
6
7
8
9
10
11
12
13
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
37
38 type Object interface {
39 DeepCopy() Object
40 }
41
42
43 type Mergeable interface {
44 Merge(a Mergeable) error
45 }
46
47
48
49
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
65
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
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
143
144
145
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