150 lines
2.6 KiB
Go
150 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"io/ioutil"
|
|
"path"
|
|
"sync"
|
|
"time"
|
|
|
|
"go.etcd.io/bbolt"
|
|
)
|
|
|
|
type Driver interface {
|
|
Close() error
|
|
ForEach(context.Context, string, func(string, []byte) error) error
|
|
Get(context.Context, string, string) ([]byte, error)
|
|
Set(context.Context, string, string, []byte) error
|
|
}
|
|
|
|
type RAM struct {
|
|
m map[string]map[string][]byte
|
|
lock *sync.RWMutex
|
|
}
|
|
|
|
func NewRAM() RAM {
|
|
return RAM{
|
|
m: make(map[string]map[string][]byte),
|
|
lock: &sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
func (ram RAM) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func (ram RAM) ForEach(ctx context.Context, ns string, cb func(string, []byte) error) error {
|
|
ram.lock.RLock()
|
|
defer ram.lock.RUnlock()
|
|
|
|
for k, v := range ram.m[ns] {
|
|
if ctx.Err() != nil {
|
|
break
|
|
}
|
|
if err := cb(k, v); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return ctx.Err()
|
|
}
|
|
|
|
func (ram RAM) Get(_ context.Context, ns, id string) ([]byte, error) {
|
|
ram.lock.RLock()
|
|
defer ram.lock.RUnlock()
|
|
|
|
if _, ok := ram.m[ns]; !ok {
|
|
return nil, nil
|
|
}
|
|
|
|
return ram.m[ns][id], nil
|
|
}
|
|
|
|
func (ram RAM) Set(_ context.Context, ns, id string, v []byte) error {
|
|
ram.lock.Lock()
|
|
defer ram.lock.Unlock()
|
|
|
|
if _, ok := ram.m[ns]; !ok {
|
|
ram.m[ns] = map[string][]byte{}
|
|
}
|
|
ram.m[ns][id] = v
|
|
|
|
return nil
|
|
}
|
|
|
|
type BBolt struct {
|
|
db *bbolt.DB
|
|
}
|
|
|
|
func NewTestDBIn(d string) BBolt {
|
|
d, err := ioutil.TempDir(d, "test-db-*")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
db, err := NewDB(path.Join(d, "bb"))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return db
|
|
}
|
|
|
|
func NewDB(p string) (BBolt, error) {
|
|
db, err := bbolt.Open(p, 0600, &bbolt.Options{
|
|
Timeout: time.Second,
|
|
})
|
|
return BBolt{db: db}, err
|
|
}
|
|
|
|
func (bb BBolt) Close() error {
|
|
return bb.db.Close()
|
|
}
|
|
|
|
func (bb BBolt) ForEach(ctx context.Context, db string, cb func(string, []byte) error) error {
|
|
return bb.db.View(func(tx *bbolt.Tx) error {
|
|
bkt := tx.Bucket([]byte(db))
|
|
if bkt == nil {
|
|
return nil
|
|
}
|
|
|
|
c := bkt.Cursor()
|
|
for k, v := c.First(); k != nil && ctx.Err() == nil; k, v = c.Next() {
|
|
if err := cb(string(k), v); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return ctx.Err()
|
|
})
|
|
}
|
|
|
|
func (bb BBolt) Get(_ context.Context, db, id string) ([]byte, error) {
|
|
var b []byte
|
|
err := bb.db.View(func(tx *bbolt.Tx) error {
|
|
bkt := tx.Bucket([]byte(db))
|
|
if bkt == nil {
|
|
return nil
|
|
}
|
|
|
|
b = bkt.Get([]byte(id))
|
|
return nil
|
|
})
|
|
return b, err
|
|
}
|
|
|
|
func (bb BBolt) Set(_ context.Context, db, id string, value []byte) error {
|
|
return bb.db.Update(func(tx *bbolt.Tx) error {
|
|
bkt := tx.Bucket([]byte(db))
|
|
if bkt == nil {
|
|
var err error
|
|
bkt, err = tx.CreateBucket([]byte(db))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if value == nil {
|
|
return bkt.Delete([]byte(id))
|
|
}
|
|
return bkt.Put([]byte(id), value)
|
|
})
|
|
}
|