-
Notifications
You must be signed in to change notification settings - Fork 175
/
identifier_map.go
176 lines (146 loc) · 4.53 KB
/
identifier_map.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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
package stdmap
import (
"fmt"
"github.com/onflow/flow-go/model/flow"
"github.com/onflow/flow-go/module/mempool"
"github.com/onflow/flow-go/module/mempool/model"
)
// IdentifierMap represents a concurrency-safe memory pool for IdMapEntity.
type IdentifierMap struct {
*Backend
}
// NewIdentifierMap creates a new memory pool for IdMapEntity.
func NewIdentifierMap(limit uint) (*IdentifierMap, error) {
i := &IdentifierMap{
Backend: NewBackend(WithLimit(limit)),
}
return i, nil
}
// Append will append the id to the list of identifiers associated with key.
func (i *IdentifierMap) Append(key, id flow.Identifier) error {
return i.Backend.Run(func(backdata mempool.BackData) error {
var ids map[flow.Identifier]struct{}
entity, ok := backdata.ByID(key)
if !ok {
// no record with key is available in the mempool,
// initializes ids.
ids = make(map[flow.Identifier]struct{})
} else {
idMapEntity, ok := entity.(model.IdMapEntity)
if !ok {
return fmt.Errorf("could not assert entity to IdMapEntity")
}
ids = idMapEntity.IDs
if _, ok := ids[id]; ok {
// id is already associated with the key
// no need to append
return nil
}
// removes map entry associated with key for update
if _, removed := backdata.Remove(key); !removed {
return fmt.Errorf("potential race condition on removing from identifier map")
}
}
// appends id to the ids list
ids[id] = struct{}{}
// adds the new ids list associated with key to mempool
idMapEntity := model.IdMapEntity{
Key: key,
IDs: ids,
}
if added := backdata.Add(key, idMapEntity); !added {
return fmt.Errorf("potential race condition on adding to identifier map")
}
return nil
})
}
// Get returns list of all identifiers associated with key and true, if the key exists in the mempool.
// Otherwise it returns nil and false.
func (i *IdentifierMap) Get(key flow.Identifier) ([]flow.Identifier, bool) {
ids := make([]flow.Identifier, 0)
err := i.Run(func(backdata mempool.BackData) error {
entity, ok := backdata.ByID(key)
if !ok {
return fmt.Errorf("could not retrieve key from backend")
}
mapEntity, ok := entity.(model.IdMapEntity)
if !ok {
return fmt.Errorf("could not assert entity as IdMapEntity")
}
for id := range mapEntity.IDs {
ids = append(ids, id)
}
return nil
})
if err != nil {
return nil, false
}
return ids, true
}
// Has returns true if the key exists in the map, i.e., there is at least an id
// attached to it.
func (i *IdentifierMap) Has(key flow.Identifier) bool {
return i.Backend.Has(key)
}
// Remove removes the given key with all associated identifiers.
func (i *IdentifierMap) Remove(key flow.Identifier) bool {
return i.Backend.Remove(key)
}
// RemoveIdFromKey removes the id from the list of identifiers associated with key.
// If the list becomes empty, it also removes the key from the map.
func (i *IdentifierMap) RemoveIdFromKey(key, id flow.Identifier) error {
err := i.Backend.Run(func(backdata mempool.BackData) error {
// var ids map[flow.Identifier]struct{}
entity, ok := backdata.ByID(key)
if !ok {
// entity key has already been removed
return nil
}
idMapEntity, ok := entity.(model.IdMapEntity)
if !ok {
return fmt.Errorf("could not assert entity to IdMapEntity")
}
if _, ok := idMapEntity.IDs[id]; !ok {
// id has already been removed from the key map
return nil
}
// removes map entry associated with key for update
if _, removed := backdata.Remove(key); !removed {
return fmt.Errorf("potential race condition on removing from identifier map")
}
// removes id from the secondary map of the key
delete(idMapEntity.IDs, id)
if len(idMapEntity.IDs) == 0 {
// all ids related to key are removed, so there is no need
// to add key back to the idMapEntity
return nil
}
// adds the new ids list associated with key to mempool
idMapEntity = model.IdMapEntity{
Key: key,
IDs: idMapEntity.IDs,
}
if added := backdata.Add(key, idMapEntity); !added {
return fmt.Errorf("potential race condition on adding to identifier map")
}
return nil
})
return err
}
// Size returns number of IdMapEntities in mempool
func (i *IdentifierMap) Size() uint {
return i.Backend.Size()
}
// Keys returns a list of all keys in the mempool
func (i *IdentifierMap) Keys() ([]flow.Identifier, bool) {
entities := i.Backend.All()
keys := make([]flow.Identifier, 0)
for _, entity := range entities {
idMapEntity, ok := entity.(model.IdMapEntity)
if !ok {
return nil, false
}
keys = append(keys, idMapEntity.Key)
}
return keys, true
}