frostfs-node/cmd/frostfs-lens/internal/tui/db.go

161 lines
3.2 KiB
Go
Raw Normal View History

package tui
import (
"context"
"errors"
"fmt"
"go.etcd.io/bbolt"
)
type Item[T any] struct {
val T
err error
}
func resolvePath(tx *bbolt.Tx, path [][]byte) (*bbolt.Bucket, error) {
if len(path) == 0 {
return nil, errors.New("can't find bucket without path")
}
name := path[0]
bucket := tx.Bucket(name)
if bucket == nil {
return nil, fmt.Errorf("no bucket with name %s", name)
}
for _, name := range path[1:] {
bucket = bucket.Bucket(name)
if bucket == nil {
return nil, fmt.Errorf("no bucket with name %s", name)
}
}
return bucket, nil
}
func load[T any](
ctx context.Context, db *bbolt.DB, path [][]byte, bufferSize int,
filter func(key, value []byte) bool, transform func(key, value []byte) T,
) (<-chan Item[T], error) {
buffer := make(chan Item[T], bufferSize)
go func() {
defer close(buffer)
err := db.View(func(tx *bbolt.Tx) error {
var cursor *bbolt.Cursor
if len(path) == 0 {
cursor = tx.Cursor()
} else {
bucket, err := resolvePath(tx, path)
if err != nil {
buffer <- Item[T]{err: fmt.Errorf("can't find bucket: %w", err)}
return nil
}
cursor = bucket.Cursor()
}
key, value := cursor.First()
for {
if key == nil {
return nil
}
if filter != nil && !filter(key, value) {
key, value = cursor.Next()
continue
}
select {
case <-ctx.Done():
return nil
case buffer <- Item[T]{val: transform(key, value)}:
key, value = cursor.Next()
}
}
})
if err != nil {
buffer <- Item[T]{err: err}
}
}()
return buffer, nil
}
func LoadBuckets(
ctx context.Context, db *bbolt.DB, path [][]byte, bufferSize int,
) (<-chan Item[*Bucket], error) {
buffer, err := load(
ctx, db, path, bufferSize,
func(_, value []byte) bool {
return value == nil
},
func(key, _ []byte) *Bucket {
base := make([][]byte, 0, len(path))
base = append(base, path...)
return &Bucket{
Name: key,
Path: append(base, key),
}
},
)
if err != nil {
return nil, fmt.Errorf("can't start iterating bucket: %w", err)
}
return buffer, nil
}
func LoadRecords(
ctx context.Context, db *bbolt.DB, path [][]byte, bufferSize int,
) (<-chan Item[*Record], error) {
buffer, err := load(
ctx, db, path, bufferSize,
func(_, value []byte) bool {
return value != nil
},
func(key, value []byte) *Record {
base := make([][]byte, 0, len(path))
base = append(base, path...)
return &Record{
Key: key,
Value: value,
Path: append(base, key),
}
},
)
if err != nil {
return nil, fmt.Errorf("can't start iterating bucket: %w", err)
}
return buffer, nil
}
// HasBuckets checks if a bucket has nested buckets. It relies on assumption
// that a bucket can have either nested buckets or records but not both.
func HasBuckets(ctx context.Context, db *bbolt.DB, path [][]byte) (bool, error) {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
buffer, err := load(
ctx, db, path, 1,
nil,
func(_, value []byte) []byte { return value },
)
if err != nil {
return false, err
}
x, ok := <-buffer
if !ok {
return false, nil
}
if x.err != nil {
return false, err
}
if x.val != nil {
return false, err
}
return true, nil
}