xk6-frostfs/internal/registry/obj_selector.go

package registry

import (
	"encoding/json"
	"errors"
	"sync"

	"go.etcd.io/bbolt"
)

type ObjSelector struct {
	boltDB    *bbolt.DB
	mu        sync.Mutex
	lastId    uint64
	objStatus string
}

func (o *ObjSelector) NextObject() (*ObjectInfo, error) {
	var foundObj *ObjectInfo
	err := o.boltDB.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return nil
		}

		c := b.Cursor()

		// We use mutex so that multiple VUs won't attempt to modify lastId simultaneously
		// TODO: consider singleton channel that will produce those ids on demand
		o.mu.Lock()
		defer o.mu.Unlock()

		// 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 key right after the lastId
		var keyBytes, objBytes []byte
		if o.lastId == 0 {
			keyBytes, objBytes = c.First()
		} else {
			c.Seek(encodeId(o.lastId))
			keyBytes, objBytes = c.Next()
		}

		// Iterate over objects to find the next object in the target status
		var obj ObjectInfo
		for ; keyBytes != nil; keyBytes, objBytes = c.Next() {
			if objBytes != nil {
				if err := json.Unmarshal(objBytes, &obj); err != nil {
					// Ignore malformed objects for now. Maybe it should be panic?
					continue
				}
				// If we reached an object in the target status, stop iterating
				if obj.Status == o.objStatus {
					foundObj = &obj
					break
				}
			}
		}

		// Update the last key
		if keyBytes != nil {
			o.lastId = decodeId(keyBytes)
			return nil
		}

		return errors.New("no objects are available")
	})
	return foundObj, err
}