Verified feat: Add on pebble state backend support pebbledb

1 files changed, 126 insertions, 0 deletions

diff --git a/state/pebble_backend.go b/state/pebble_backend.go
new file mode 100644
index 0000000..e428988
--- /dev/null
+++ b/state/pebble_backend.go
@@ -0,0 +1,126 @@
+package state
+
+import (
+        "context"
+        "encoding/binary"
+        "fmt"
+        "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,
+        }, 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 nil, 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(statename string) []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
+        stateID := p.getStateID(statename)
+        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
+}