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

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
}