|
1 | 1 | package batch_task |
2 | 2 |
|
3 | 3 | import ( |
4 | | - "maps" |
5 | 4 | "sync" |
| 5 | + |
| 6 | + "github.com/sirupsen/logrus" |
| 7 | +) |
| 8 | + |
| 9 | +type taskKey int |
| 10 | + |
| 11 | +const ( |
| 12 | + _ taskKey = iota |
| 13 | + refreshPath |
| 14 | + MoveSrcPath |
| 15 | + MoveDstPath |
6 | 16 | ) |
7 | 17 |
|
8 | | -type FinishHook func(allTaskArgs map[string]TaskMap) |
9 | | -type BatchTaskHook struct { |
10 | | - name string |
11 | | - mu sync.Mutex |
12 | | - tasks map[string]struct{} |
13 | | - taskArgs map[string]TaskMap |
| 18 | +type TaskPayload map[taskKey]any |
| 19 | +type FinishHook func(payloads []TaskPayload) |
| 20 | +type BatchTasCoordinator struct { |
| 21 | + name string |
| 22 | + mu sync.Mutex |
| 23 | + |
| 24 | + pendingTasks map[string][]TaskPayload |
| 25 | + finishCounts map[string]int |
14 | 26 | allFinishHook FinishHook |
15 | 27 | } |
16 | 28 |
|
17 | | -func NewBatchTaskHook(name string) *BatchTaskHook { |
18 | | - return &BatchTaskHook{ |
19 | | - name: name, |
20 | | - tasks: map[string]struct{}{}, |
21 | | - taskArgs: map[string]TaskMap{}, |
| 29 | +func NewBatchTasCoordinator(name string) *BatchTasCoordinator { |
| 30 | + return &BatchTasCoordinator{ |
| 31 | + name: name, |
| 32 | + pendingTasks: map[string][]TaskPayload{}, |
| 33 | + finishCounts: map[string]int{}, |
22 | 34 | } |
23 | 35 | } |
24 | 36 |
|
25 | | -func (bt *BatchTaskHook) SetAllFinishHook(f FinishHook) *BatchTaskHook { |
| 37 | +func (bt *BatchTasCoordinator) SetAllFinishHook(f FinishHook) *BatchTasCoordinator { |
26 | 38 | bt.allFinishHook = f |
27 | 39 | return bt |
28 | 40 | } |
29 | 41 |
|
30 | | -func (bt *BatchTaskHook) AddTask(taskID string, taskMap TaskMap) { |
| 42 | +// 自动添加refreshPath在第1个任务 |
| 43 | +func (bt *BatchTasCoordinator) AddTask(targetPath string, playload TaskPayload) { |
31 | 44 | bt.mu.Lock() |
32 | 45 | defer bt.mu.Unlock() |
33 | | - bt.tasks[taskID] = struct{}{} |
34 | | - if existingMap, ok := bt.taskArgs[taskID]; ok { |
35 | | - maps.Copy(existingMap, taskMap) |
| 46 | + defer func() { |
| 47 | + logrus.Debugf("AddTask:%s ,%d", targetPath, len(bt.pendingTasks[targetPath])) |
| 48 | + }() |
| 49 | + if playloads, ok := bt.pendingTasks[targetPath]; ok { |
| 50 | + if _, ok := playload[refreshPath]; ok { |
| 51 | + if len(playload) == 1 { |
| 52 | + t := playloads[0] |
| 53 | + t[refreshPath] = targetPath |
| 54 | + playloads[0] = t |
| 55 | + return |
| 56 | + } |
| 57 | + delete(playload, refreshPath) |
| 58 | + } |
| 59 | + bt.pendingTasks[targetPath] = append(playloads, playload) |
36 | 60 | } else { |
37 | | - bt.taskArgs[taskID] = taskMap |
| 61 | + playload[refreshPath] = targetPath |
| 62 | + bt.pendingTasks[targetPath] = []TaskPayload{playload} |
38 | 63 | } |
39 | 64 | } |
40 | 65 |
|
41 | | -func (bt *BatchTaskHook) RemoveTask(taskID string, allFinish bool) { |
| 66 | +func (bt *BatchTasCoordinator) MarkTaskFinish(targetPath string) { |
42 | 67 | bt.mu.Lock() |
43 | 68 | defer bt.mu.Unlock() |
44 | | - delete(bt.tasks, taskID) |
45 | | - if len(bt.tasks) == 0 && allFinish { |
46 | | - if bt.allFinishHook != nil { |
47 | | - bt.allFinishHook(bt.taskArgs) |
| 69 | + finishCount := bt.finishCounts[targetPath] |
| 70 | + finishCount++ |
| 71 | + logrus.Debugf("MarkTaskFinish:%s ,%v", targetPath, finishCount) |
| 72 | + if playloads, ok := bt.pendingTasks[targetPath]; ok { |
| 73 | + if len(playloads) == finishCount { |
| 74 | + delete(bt.pendingTasks, targetPath) |
| 75 | + delete(bt.finishCounts, targetPath) |
| 76 | + if bt.allFinishHook != nil { |
| 77 | + logrus.Debugf("allFinishHook:%s", targetPath) |
| 78 | + bt.allFinishHook(playloads) |
| 79 | + } |
| 80 | + return |
48 | 81 | } |
49 | | - clear(bt.taskArgs) |
50 | | - } |
51 | | -} |
52 | | - |
53 | | -func (bt *BatchTaskHook) GetAllTaskArgs() map[string]TaskMap { |
54 | | - bt.mu.Lock() |
55 | | - defer bt.mu.Unlock() |
56 | | - |
57 | | - result := map[string]TaskMap{} |
58 | | - for taskID, args := range bt.taskArgs { |
59 | | - copyArgs := TaskMap{} |
60 | | - maps.Copy(copyArgs, args) |
61 | | - result[taskID] = copyArgs |
62 | 82 | } |
63 | | - return result |
| 83 | + bt.finishCounts[targetPath] = finishCount |
64 | 84 | } |
0 commit comments