-
Notifications
You must be signed in to change notification settings - Fork 375
/
snapshot.go
84 lines (73 loc) · 2.59 KB
/
snapshot.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
// Copyright (c) 2021 Terminus, Inc.
//
// This program is free software: you can use, redistribute, and/or modify
// it under the terms of the GNU Affero General Public License, version 3
// or later ("AGPL"), as published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
package throttler
import (
"encoding/json"
"fmt"
"github.com/erda-project/erda/modules/pipeline/pipengine/queue/enhancedqueue"
)
type SnapshotObj struct {
Name string `json:"name"`
QueueByName map[string]json.RawMessage `json:"queueByName"` // queueByName 无法从 keyRelatedQueues 中还原,因为可能存在空队列
KeyRelatedQueues map[string][]string `json:"keyRelatedQueues"` // 只关心加入了哪些队列,无需重复引用队列
}
func (t *throttler) Export() json.RawMessage {
t.lock.Lock()
defer t.lock.Unlock()
obj := SnapshotObj{
Name: t.name,
QueueByName: make(map[string]json.RawMessage),
KeyRelatedQueues: make(map[string][]string),
}
// queueByName
for qName, queue := range t.queueByName {
obj.QueueByName[qName] = queue.Export()
}
// keyRelatedQueues
for key, queueByName := range t.keyRelatedQueues {
for qName := range queueByName {
obj.KeyRelatedQueues[key] = append(obj.KeyRelatedQueues[key], qName)
}
}
b, _ := json.Marshal(&obj)
return b
}
func (t *throttler) Import(rawMsg json.RawMessage) error {
t.lock.Lock()
defer t.lock.Unlock()
var obj SnapshotObj
if err := json.Unmarshal(rawMsg, &obj); err != nil {
return err
}
_nt := NewNamedThrottler(obj.Name, nil)
nt := _nt.(*throttler)
t.name = nt.name
// queueByName
t.queueByName = make(map[string]*enhancedqueue.EnhancedQueue)
for qName, queueRawJSON := range obj.QueueByName {
eq := enhancedqueue.NewEnhancedQueue(0)
if err := eq.Import(queueRawJSON); err != nil {
return fmt.Errorf("failed to import enhanced queue, queue name: %s, err: %v", qName, err)
}
t.queueByName[qName] = eq
}
// keyRelatedQueues
t.keyRelatedQueues = make(map[string]map[string]*enhancedqueue.EnhancedQueue)
for key, relatedQueues := range obj.KeyRelatedQueues {
t.keyRelatedQueues[key] = make(map[string]*enhancedqueue.EnhancedQueue)
for _, qName := range relatedQueues {
t.keyRelatedQueues[key][qName] = t.queueByName[qName]
}
}
return nil
}