-
Notifications
You must be signed in to change notification settings - Fork 672
/
state.go
243 lines (201 loc) · 7.42 KB
/
state.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
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
// Copyright (C) 2019-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package state
import (
"errors"
"fmt"
"time"
"github.com/ava-labs/avalanchego/cache"
"github.com/ava-labs/avalanchego/database"
"github.com/ava-labs/avalanchego/ids"
"github.com/ava-labs/avalanchego/snow/choices"
"github.com/ava-labs/avalanchego/utils/wrappers"
)
const cacheSize = 1000
var (
errWrongType = errors.New("value in the database was the wrong type")
_ State = &state{}
)
// State is a key-value store where every value is associated with a "type ID".
// Every different type of value must have its own type ID.
//
// For example, if you're storing blocks, accounts and addresses, each of those types
// must have their own type ID.
//
// Each type ID is associated with a function that specifies how to unmarshal bytes
// to a struct/value of a given type.
//
// State has built-in support for putting and getting choices.Status and ids.ID
// To put/get any other type, you must first register that type using RegisterType
type State interface {
// In [db], add a key-value pair.
// [value] will be converted to bytes by calling Bytes() on it.
// [typeID] must have already been registered using RegisterType.
// If [value] is nil, the value associated with [key] and [typeID] is deleted (if it exists).
Put(db database.Database, typeID uint64, key ids.ID, value interface{}) error
// From [db], get the value of type [typeID] whose key is [key]
// Returns database.ErrNotFound if the entry doesn't exist
Get(db database.Database, typeID uint64, key ids.ID) (interface{}, error)
// Return whether [key] exists in [db] for type [typeID]
Has(db database.Database, typeID uint64, key ids.ID) (bool, error)
// PutStatus associates [key] with [status] in [db]
PutStatus(db database.Database, key ids.ID, status choices.Status) error
// GetStatus gets the status associated with [key] in [db]
GetStatus(db database.Database, key ids.ID) choices.Status
// PutID associates [key] with [ID] in [db]
PutID(db database.Database, key ids.ID, ID ids.ID) error
// GetID gets the ID associated with [key] in [db]
GetID(db database.Database, key ids.ID) (ids.ID, error)
// PutTime associates [key] with [time] in [db]
PutTime(db database.Database, key ids.ID, time time.Time) error
// GetTime gets the time associated with [key] in [db]
GetTime(db database.Database, key ids.ID) (time.Time, error)
// Register a new type.
// When values that were Put with [typeID] are retrieved from the database,
// they will be unmarshaled from bytes using [unmarshal].
// Returns an error if there is already a type with ID [typeID]
RegisterType(typeID uint64,
marshal func(interface{}) ([]byte, error),
unmarshal func([]byte) (interface{}, error),
) error
}
type state struct {
// Keys: Type ID
// Values: Function that unmarshals values
// that were Put with that type ID
unmarshallers map[uint64]func([]byte) (interface{}, error)
// Keys: Type ID
// Values: Function that marshals values
// that were Put with that type ID
marshallers map[uint64]func(interface{}) ([]byte, error)
// Keys: Type ID
// Values: Cache that stores uniqueIDs for values that were put with that type ID
// (Saves us from having to re-compute uniqueIDs)
uniqueIDCaches map[uint64]*cache.LRU
}
func (s *state) RegisterType(
typeID uint64,
marshal func(interface{}) ([]byte, error),
unmarshal func([]byte) (interface{}, error),
) error {
if _, exists := s.unmarshallers[typeID]; exists {
return fmt.Errorf("there is already a type with ID %d", typeID)
}
s.marshallers[typeID] = marshal
s.unmarshallers[typeID] = unmarshal
return nil
}
func (s *state) Put(db database.Database, typeID uint64, key ids.ID, value interface{}) error {
marshaller, exists := s.marshallers[typeID]
if !exists {
return fmt.Errorf("typeID %d has not been registered", typeID)
}
// Get the unique ID of thie key/typeID pair
uID := s.uniqueID(key, typeID)
if value == nil {
return db.Delete(uID[:])
}
// Put the byte repr. of the value in the database
valueBytes, err := marshaller(value)
if err != nil {
return err
}
return db.Put(uID[:], valueBytes)
}
func (s *state) Has(db database.Database, typeID uint64, key ids.ID) (bool, error) {
key = s.uniqueID(key, typeID)
return db.Has(key[:])
}
func (s *state) Get(db database.Database, typeID uint64, key ids.ID) (interface{}, error) {
unmarshal, exists := s.unmarshallers[typeID]
if !exists {
return nil, fmt.Errorf("typeID %d has not been registered", typeID)
}
// The unique ID of this key/typeID pair
uID := s.uniqueID(key, typeID)
// Get the value from the database
valueBytes, err := db.Get(uID[:])
if err != nil {
return nil, err
}
// Unmarshal the value from bytes and return it
return unmarshal(valueBytes)
}
// PutStatus associates [key] with [status] in [db]
func (s *state) PutStatus(db database.Database, key ids.ID, status choices.Status) error {
return s.Put(db, StatusTypeID, key, status)
}
// GetStatus gets the status associated with [key] in [db]
// Return choices.Processing if can't get the status from database
func (s *state) GetStatus(db database.Database, key ids.ID) choices.Status {
statusInterface, err := s.Get(db, StatusTypeID, key)
if err != nil {
return choices.Processing
}
status, ok := statusInterface.(choices.Status)
if !ok || status.Valid() != nil {
return choices.Processing
}
return status
}
// PutID associates [key] with [ID] in [db]
func (s *state) PutID(db database.Database, key ids.ID, id ids.ID) error {
return s.Put(db, IDTypeID, key, id)
}
// GetID gets the ID associated with [key] in [db]
func (s *state) GetID(db database.Database, key ids.ID) (ids.ID, error) {
IDInterface, err := s.Get(db, IDTypeID, key)
if err != nil {
return ids.ID{}, err
}
if ID, ok := IDInterface.(ids.ID); ok {
return ID, nil
}
return ids.ID{}, errWrongType
}
// PutTime associates [key] with [time] in [db]
func (s *state) PutTime(db database.Database, key ids.ID, time time.Time) error {
return s.Put(db, TimeTypeID, key, time)
}
// GetTime gets the time associated with [key] in [db]
func (s *state) GetTime(db database.Database, key ids.ID) (time.Time, error) {
timeInterface, err := s.Get(db, TimeTypeID, key)
if err != nil {
return time.Time{}, err
}
if time, ok := timeInterface.(time.Time); ok {
return time, nil
}
return time.Time{}, errWrongType
}
// Prefix [ID] with [typeID] to prevent key collisions in the database
func (s *state) uniqueID(id ids.ID, typeID uint64) ids.ID {
uIDCache, cacheExists := s.uniqueIDCaches[typeID]
if cacheExists {
if uID, uIDExists := uIDCache.Get(id); uIDExists { // Get the uniqueID associated with [typeID] and [ID]
return uID.(ids.ID)
}
} else {
s.uniqueIDCaches[typeID] = &cache.LRU{Size: cacheSize}
}
uID := id.Prefix(typeID)
s.uniqueIDCaches[typeID].Put(id, uID)
return uID
}
// NewState returns a new State
func NewState() (State, error) {
state := &state{
marshallers: make(map[uint64]func(interface{}) ([]byte, error)),
unmarshallers: make(map[uint64]func([]byte) (interface{}, error)),
uniqueIDCaches: make(map[uint64]*cache.LRU),
}
// Register ID, Status and time.Time so they can be put/get without client code
// having to register them
errs := wrappers.Errs{}
errs.Add(
state.RegisterType(IDTypeID, marshalID, unmarshalID),
state.RegisterType(StatusTypeID, marshalStatus, unmarshalStatus),
state.RegisterType(TimeTypeID, marshalTime, unmarshalTime),
)
return state, errs.Err
}