xk6-frostfs/internal/registry/obj_selector.go

package registry

import (
	"context"
	"fmt"
	"time"

	"github.com/nspcc-dev/neo-go/pkg/io"
	"go.etcd.io/bbolt"
)

const nextObjectTimeout = 10 * time.Second

type ObjFilter struct {
	Status string
	Age    int
}

type ObjSelector struct {
	ctx       context.Context
	objChan   chan *ObjectInfo
	boltDB    *bbolt.DB
	filter    *ObjFilter
	cacheSize int
	kind      SelectorKind
}

// objectSelectCache is the default maximum size of a batch to select from DB.
const objectSelectCache = 1000

// NewObjSelector creates a new instance of object selector that can iterate over
// objects in the specified registry.
func NewObjSelector(registry *ObjRegistry, selectionSize int, kind SelectorKind, filter *ObjFilter) *ObjSelector {
	if selectionSize <= 0 {
		selectionSize = objectSelectCache
	}
	if filter == nil || filter.Status == "" {
		panic("filtering without status is not supported")
	}
	objSelector := &ObjSelector{
		ctx:       registry.ctx,
		boltDB:    registry.boltDB,
		filter:    filter,
		objChan:   make(chan *ObjectInfo, selectionSize*2),
		cacheSize: selectionSize,
		kind:      kind,
	}

	go objSelector.selectLoop()

	return objSelector
}

// NextObject returns the next object from the registry that matches filter of
// the selector. NextObject only roams forward from the current position of the
// selector. If there are no objects that match the filter, blocks until one of
// the following happens:
//   - a "new" next object is available;
//   - underlying registry context is done, nil objects will be returned on the
//     currently blocked and every further NextObject calls.
func (o *ObjSelector) NextObject() *ObjectInfo {
	if o.kind == SelectorOneshot {
		return <-o.objChan
	}

	select {
	case <-time.After(nextObjectTimeout):
		return nil
	case obj := <-o.objChan:
		return obj
	}
}

// Count returns total number of objects that match filter of the selector.
func (o *ObjSelector) Count() (int, error) {
	count := 0
	err := o.boltDB.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(o.filter.Status))
		if b == nil {
			return nil
		}

		if o.filter.Age == 0 {
			count = b.Stats().KeyN
			return nil
		}

		return b.ForEach(func(_, objBytes []byte) error {
			if objBytes != nil {
				r := io.NewBinReaderFromBuf(objBytes)

				var obj ObjectInfo
				obj.decodeFilterableFields(r)
				if r.Err != nil {
					// Ignore malformed objects
					return nil
				}
				if o.filter.match(obj) {
					count++
				}
			}
			return nil
		})
	})
	return count, err
}

func (o *ObjSelector) selectLoop() {
	cache := make([]*ObjectInfo, 0, o.cacheSize)
	var lastID uint64
	defer close(o.objChan)

	for {
		select {
		case <-o.ctx.Done():
			return
		default:
		}

		// cache the objects
		err := o.boltDB.View(func(tx *bbolt.Tx) error {
			b := tx.Bucket([]byte(o.filter.Status))
			if b == nil {
				return nil
			}

			c := b.Cursor()

			// Establish the start position for searching the next object:
			// If we should go from the beginning (lastID=0), then we start
			// from the first element. Otherwise, we start from the last
			// handled ID + 1.
			var keyBytes, objBytes []byte
			if lastID == 0 {
				keyBytes, objBytes = c.First()
			} else {
				keyBytes, objBytes = c.Seek(encodeId(lastID))
				if keyBytes != nil && decodeId(keyBytes) == lastID {
					keyBytes, objBytes = c.Next()
				}
			}

			// Iterate over objects to find the next object matching the filter.
			for ; keyBytes != nil && len(cache) != o.cacheSize; keyBytes, objBytes = c.Next() {
				if objBytes != nil {
					var obj ObjectInfo
					if err := obj.Unmarshal(objBytes); err != nil {
						// Ignore malformed objects for now. Maybe it should be panic?
						continue
					}

					if o.filter.match(obj) {
						cache = append(cache, &obj)
					}
				}
			}

			if len(cache) > 0 {
				lastID = cache[len(cache)-1].Id
			}

			return nil
		})
		if err != nil {
			panic(fmt.Errorf("fetching objects failed: %w", err))
		}

		for _, obj := range cache {
			select {
			case <-o.ctx.Done():
				return
			case o.objChan <- obj:
			}
		}

		if o.kind == SelectorOneshot && len(cache) != o.cacheSize {
			return
		}

		if o.kind != SelectorLooped && len(cache) != o.cacheSize {
			// no more objects, wait a little; the logic could be improved.
			select {
			case <-time.After(time.Second):
			case <-o.ctx.Done():
				return
			}
		}

		if o.kind == SelectorLooped && len(cache) != o.cacheSize {
			lastID = 0
		}

		// clean handled objects
		cache = cache[:0]
	}
}

func (f *ObjFilter) match(o ObjectInfo) bool {
	if f.Status != "" && f.Status != o.Status {
		return false
	}
	if f.Age != 0 {
		objAge := time.Now().UTC().Unix() - o.CreatedAt
		if objAge < int64(f.Age) {
			return false
		}
	}
	return true
}