250 lines
6.2 KiB
Go
250 lines
6.2 KiB
Go
package graphvent
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"reflect"
|
|
"sync"
|
|
|
|
badger "github.com/dgraph-io/badger/v3"
|
|
)
|
|
|
|
type Database interface {
|
|
WriteNodeInit(*Context, *Node) error
|
|
WriteNodeChanges(*Context, *Node, map[ExtType]Changes) error
|
|
LoadNode(*Context, NodeID) (*Node, error)
|
|
}
|
|
|
|
const WRITE_BUFFER_SIZE = 1000000
|
|
type BadgerDB struct {
|
|
*badger.DB
|
|
sync.Mutex
|
|
buffer [WRITE_BUFFER_SIZE]byte
|
|
}
|
|
|
|
func (db *BadgerDB) WriteNodeInit(ctx *Context, node *Node) error {
|
|
if node == nil {
|
|
return fmt.Errorf("Cannot serialize nil *Node")
|
|
}
|
|
|
|
return db.Update(func(tx *badger.Txn) error {
|
|
db.Lock()
|
|
defer db.Unlock()
|
|
|
|
// Get the base key bytes
|
|
id_ser, err := node.ID.MarshalBinary()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
cur := 0
|
|
|
|
// Write Node value
|
|
written, err := Serialize(ctx, node, db.buffer[cur:])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = tx.Set(id_ser, db.buffer[cur:cur+written])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
cur += written
|
|
|
|
// Write empty signal queue
|
|
sigqueue_id := append(id_ser, []byte(" - SIGQUEUE")...)
|
|
written, err = Serialize(ctx, node.SignalQueue, db.buffer[cur:])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = tx.Set(sigqueue_id, db.buffer[cur:cur+written])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
cur += written
|
|
|
|
// Write node extension list
|
|
ext_list := []ExtType{}
|
|
for ext_type := range(node.Extensions) {
|
|
ext_list = append(ext_list, ext_type)
|
|
}
|
|
written, err = Serialize(ctx, ext_list, db.buffer[cur:])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ext_list_id := append(id_ser, []byte(" - EXTLIST")...)
|
|
err = tx.Set(ext_list_id, db.buffer[cur:cur+written])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cur += written
|
|
|
|
// For each extension:
|
|
for ext_type, ext := range(node.Extensions) {
|
|
// Write each extension's current value
|
|
ext_id := binary.BigEndian.AppendUint64(id_ser, uint64(ext_type))
|
|
written, err := SerializeValue(ctx, reflect.ValueOf(ext).Elem(), db.buffer[cur:])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = tx.Set(ext_id, db.buffer[cur:cur+written])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cur += written
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (db *BadgerDB) WriteNodeChanges(ctx *Context, node *Node, changes map[ExtType]Changes) error {
|
|
return db.Update(func(tx *badger.Txn) error {
|
|
db.Lock()
|
|
defer db.Unlock()
|
|
|
|
// Get the base key bytes
|
|
id_bytes := ([16]byte)(node.ID)
|
|
|
|
cur := 0
|
|
|
|
// Write the signal queue if it needs to be written
|
|
if node.writeSignalQueue {
|
|
node.writeSignalQueue = false
|
|
|
|
sigqueue_id := append(id_bytes[:], []byte(" - SIGQUEUE")...)
|
|
written, err := Serialize(ctx, node.SignalQueue, db.buffer[cur:])
|
|
if err != nil {
|
|
return fmt.Errorf("SignalQueue Serialize Error: %+v, %w", node.SignalQueue, err)
|
|
}
|
|
err = tx.Set(sigqueue_id, db.buffer[cur:cur+written])
|
|
if err != nil {
|
|
return fmt.Errorf("SignalQueue set error: %+v, %w", node.SignalQueue, err)
|
|
}
|
|
cur += written
|
|
}
|
|
|
|
// For each ext in changes
|
|
for ext_type := range(changes) {
|
|
// Write each ext
|
|
ext, exists := node.Extensions[ext_type]
|
|
if exists == false {
|
|
return fmt.Errorf("%s is not an extension in %s", ext_type, node.ID)
|
|
}
|
|
ext_id := binary.BigEndian.AppendUint64(id_bytes[:], uint64(ext_type))
|
|
written, err := SerializeValue(ctx, reflect.ValueOf(ext).Elem(), db.buffer[cur:])
|
|
if err != nil {
|
|
return fmt.Errorf("Extension serialize err: %s, %w", reflect.TypeOf(ext), err)
|
|
}
|
|
|
|
err = tx.Set(ext_id, db.buffer[cur:cur+written])
|
|
if err != nil {
|
|
return fmt.Errorf("Extension set err: %s, %w", reflect.TypeOf(ext), err)
|
|
}
|
|
cur += written
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (db *BadgerDB) LoadNode(ctx *Context, id NodeID) (*Node, error) {
|
|
var node *Node = nil
|
|
|
|
err := db.View(func(tx *badger.Txn) error {
|
|
// Get the base key bytes
|
|
id_ser, err := id.MarshalBinary()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Get the node value
|
|
node_item, err := tx.Get(id_ser)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = node_item.Value(func(val []byte) error {
|
|
ctx.Log.Logf("db", "DESERIALIZE_NODE(%d bytes): %+v", len(val), val)
|
|
node, err = Deserialize[*Node](ctx, val)
|
|
return err
|
|
})
|
|
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to deserialize Node %s - %w", id, err)
|
|
}
|
|
|
|
// Get the signal queue
|
|
sigqueue_id := append(id_ser, []byte(" - SIGQUEUE")...)
|
|
sigqueue_item, err := tx.Get(sigqueue_id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = sigqueue_item.Value(func(val []byte) error {
|
|
node.SignalQueue, err = Deserialize[[]QueuedSignal](ctx, val)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Get the extension list
|
|
ext_list_id := append(id_ser, []byte(" - EXTLIST")...)
|
|
ext_list_item, err := tx.Get(ext_list_id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
var ext_list []ExtType
|
|
ext_list_item.Value(func(val []byte) error {
|
|
ext_list, err = Deserialize[[]ExtType](ctx, val)
|
|
return err
|
|
})
|
|
|
|
// Get the extensions
|
|
for _, ext_type := range(ext_list) {
|
|
ext_id := binary.BigEndian.AppendUint64(id_ser, uint64(ext_type))
|
|
ext_item, err := tx.Get(ext_id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
ext_info, exists := ctx.Extensions[ext_type]
|
|
if exists == false {
|
|
return fmt.Errorf("Extension %s not in context", ext_type)
|
|
}
|
|
|
|
var ext Extension
|
|
var ok bool
|
|
err = ext_item.Value(func(val []byte) error {
|
|
value, _, err := DeserializeValue(ctx, val, ext_info.Type)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
ext, ok = value.Addr().Interface().(Extension)
|
|
if ok == false {
|
|
return fmt.Errorf("Parsed value %+v is not extension", value.Type())
|
|
}
|
|
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
node.Extensions[ext_type] = ext
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
} else if node == nil {
|
|
return nil, fmt.Errorf("Tried to return nil *Node from BadgerDB.LoadNode without error")
|
|
}
|
|
|
|
return node, nil
|
|
}
|