graphvent/db.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
}