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
|
package state
import (
"context"
"encoding/binary"
"fmt"
"os"
"path/filepath"
"sync"
"flux/core"
"github.com/cockroachdb/pebble/v2"
)
type PebbleStateBackend struct {
maxParallelism int
currentKey string
currentGroupKey int
currentNamespace any
db *pebble.DB
// for support statename storage optimize
// PebbleDB dont have column family as RockDB does
// manually convert name <-> id
mu sync.Mutex
stateMap map[string]uint8
// nextStateID prevent the hash collision
nextStateID uint8
}
func NewPebbleStateBackend(dir string, maxParalleism int) (*PebbleStateBackend, error) {
db, err := pebble.Open(dir, &pebble.Options{})
if err != nil {
return nil, fmt.Errorf("failed to open pebbledb %w", err)
}
return &PebbleStateBackend{
maxParallelism: maxParalleism,
db: db,
stateMap: make(map[string]uint8),
nextStateID: 1,
currentNamespace: &core.VoidNamespace{},
}, nil
}
func (p *PebbleStateBackend) SetCurrentKey(key string) error {
p.currentKey = key
p.currentGroupKey = core.AssignToKeyGroup(key, p.maxParallelism)
return nil
}
func (p *PebbleStateBackend) GetCurrentKey() string {
return p.currentKey
}
func (p *PebbleStateBackend) SetCurrentNamespace(namespace any) error {
p.currentNamespace = namespace
return nil
}
func (p *PebbleStateBackend) GetCurrentNamespace() any {
return p.currentNamespace
}
// GetValueState(ctx context.Context, key string) (ValueState[any], error)
func (p *PebbleStateBackend) GetValueState(ctx context.Context, stateName string) (core.ValueState[any], error) {
stateID := p.getStateID(stateName)
return &pebbleStateValue{
backend: p,
stateID: stateID,
}, nil
}
func (p *PebbleStateBackend) getStateID(stateName string) uint8 {
p.mu.Lock()
defer p.mu.Unlock()
if _, ok := p.stateMap[stateName]; !ok {
id := p.nextStateID
p.stateMap[stateName] = id
p.nextStateID++
}
return p.stateMap[stateName]
}
func (p *PebbleStateBackend) calculateNamespaceSize() int {
nsLen := 0
switch p.currentNamespace.(type) {
case core.TimeWindow:
nsLen = 16
case string:
nsLen = len(p.currentNamespace.(string))
default:
// void namespace
nsLen = 0
}
return nsLen
}
func (p *PebbleStateBackend) deserialize(stateID uint8) []byte {
// this will store in SSTable, order is important, keygroup MUST be first for scaling
// 2B(KeyGroup) + 1B(StateID) + 2B(KeyLen) + len(Key) + 16B(Window)
kl := len(p.currentKey)
b := make([]byte, 0, p.calculateNamespaceSize()+5+kl)
// add keygroup
b = binary.BigEndian.AppendUint16(b, uint16(p.currentGroupKey))
// add state id
b = append(b, stateID)
// add key
b = binary.BigEndian.AppendUint16(b, uint16(kl))
b = append(b, p.currentKey...)
// add window / namespace
switch ns := p.currentNamespace.(type) {
case core.TimeWindow:
b = binary.BigEndian.AppendUint64(b, uint64(ns.Start))
b = binary.BigEndian.AppendUint64(b, uint64(ns.End))
case string:
b = append(b, ns...)
default:
}
return b
}
func (p *PebbleStateBackend) Close() error {
return p.db.Close()
}
// Snapshot atomic
func (p *PebbleStateBackend) Snapshot(checkpointID uint64, targetDir string) error {
// get parent dir
parentDir := filepath.Dir(targetDir)
if err := os.MkdirAll(parentDir, 00755); err != nil {
return err
}
// create temp dir
tmpDir := targetDir + ".tmp"
_ = os.RemoveAll(tmpDir)
// rollback
success := false
defer func() {
if !success {
_ = os.RemoveAll(tmpDir)
}
}()
if err := p.db.Checkpoint(tmpDir, pebble.WithFlushedWAL()); err != nil {
return fmt.Errorf("pebble checkpoint failed for ID %d: %w", checkpointID, err)
}
_ = os.RemoveAll(targetDir)
if err := os.Rename(tmpDir, targetDir); err != nil {
return fmt.Errorf("failed to commit checkpoint atomically from %s to %s: %w", tmpDir, targetDir, err)
}
// fsync the parent dir
if parentFile, err := os.Open(parentDir); err == nil {
_ = parentFile.Sync()
_ = parentFile.Close()
}
success = true
return nil
}
|
