140 lines
2.4 KiB
Go
140 lines
2.4 KiB
Go
package badger
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"git.markbailey.dev/cerbervs/joist/internal/errors"
|
|
"github.com/dgraph-io/badger/v4"
|
|
)
|
|
|
|
const defaultBaseDir = "./data/badger"
|
|
|
|
type BadgerOpts func(*badgerStore)
|
|
|
|
func WithBaseDir(baseDir string) BadgerOpts {
|
|
return func(b *badgerStore) {
|
|
b.baseDir = baseDir
|
|
}
|
|
}
|
|
|
|
func WithSubDir(subDir string) BadgerOpts {
|
|
return func(b *badgerStore) {
|
|
b.subDir = subDir
|
|
}
|
|
}
|
|
|
|
type badgerStore struct {
|
|
baseDir string
|
|
subDir string
|
|
db *badger.DB
|
|
}
|
|
|
|
func (b *badgerStore) Put(ctx context.Context, key string, val []byte, ttl time.Duration) error {
|
|
done := make(chan error, 1)
|
|
|
|
go func() {
|
|
err := b.db.Update(func(txn *badger.Txn) error {
|
|
e := badger.NewEntry([]byte(key), val).WithTTL(ttl)
|
|
return txn.SetEntry(e)
|
|
})
|
|
done <- err
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case err := <-done:
|
|
return err
|
|
}
|
|
}
|
|
|
|
func (b *badgerStore) Get(ctx context.Context, key string) ([]byte, error) {
|
|
done := make(chan struct {
|
|
val []byte
|
|
err error
|
|
}, 1)
|
|
|
|
go func() {
|
|
var valCopy []byte
|
|
err := b.db.View(func(txn *badger.Txn) error {
|
|
item, err := txn.Get([]byte(key))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
valCopy, err = item.ValueCopy(nil)
|
|
return err
|
|
})
|
|
done <- struct {
|
|
val []byte
|
|
err error
|
|
}{valCopy, err}
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case result := <-done:
|
|
return result.val, result.err
|
|
}
|
|
}
|
|
|
|
func (b *badgerStore) Delete(ctx context.Context, key string) error {
|
|
done := make(chan error, 1)
|
|
|
|
go func() {
|
|
err := b.db.Update(func(txn *badger.Txn) error {
|
|
return txn.Delete([]byte(key))
|
|
})
|
|
done <- err
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case err := <-done:
|
|
return err
|
|
}
|
|
}
|
|
|
|
func (b *badgerStore) Close() error {
|
|
return b.db.Close()
|
|
}
|
|
|
|
func (b *badgerStore) dir() string {
|
|
if b.baseDir == "" {
|
|
b.baseDir = defaultBaseDir
|
|
}
|
|
|
|
if b.subDir == "" {
|
|
return b.baseDir
|
|
}
|
|
|
|
return filepath.Join(b.baseDir, b.subDir)
|
|
}
|
|
|
|
func NewBadgerStore(opts ...BadgerOpts) (*badgerStore, error) {
|
|
s := &badgerStore{}
|
|
|
|
for _, opt := range opts {
|
|
opt(s)
|
|
}
|
|
|
|
fi, err := os.Stat(s.dir())
|
|
if err == nil && !fi.IsDir() {
|
|
if err := os.MkdirAll(s.dir(), 0755); err != nil {
|
|
return nil, errors.NewErrBadgerInit(s.baseDir, s.subDir)
|
|
}
|
|
}
|
|
|
|
bopts := badger.DefaultOptions(s.dir())
|
|
s.db, err = badger.Open(bopts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return s, nil
|
|
}
|