queuer.go

package utils

import (
	"encoding/binary"
	"fmt"

	"github.com/cosmos/cosmos-sdk/codec"
	sdk "github.com/cosmos/cosmos-sdk/types"
	gogoprototypes "github.com/gogo/protobuf/types"
	"github.com/tendermint/tendermint/libs/log"
)

//go:generate moq -out ./mock/queuer.go -pkg mock . KVQueue

// KVQueue represents a queue built with the KVStore
type KVQueue interface {
	Enqueue(key Key, value codec.ProtoMarshaler)
	Dequeue(value codec.ProtoMarshaler, filter ...func(value codec.ProtoMarshaler) bool) bool
	IsEmpty() bool

	//TODO: convert to iterator
	Keys() []Key
}

var _ KVQueue = GeneralKVQueue{}
var _ KVQueue = BlockHeightKVQueue{}

// GeneralKVQueue is a queue that orders items based on the given prioritizer function
type GeneralKVQueue struct {
	name        StringKey
	store       KVStore
	logger      log.Logger
	prioritizer func(value codec.ProtoMarshaler) Key
}

// NewGeneralKVQueue is the contructor for GeneralKVQueue
func NewGeneralKVQueue(name string, store KVStore, logger log.Logger, prioritizer func(value codec.ProtoMarshaler) Key) GeneralKVQueue {
	return GeneralKVQueue{
		name:        KeyFromStr(name),
		store:       store,
		logger:      logger,
		prioritizer: prioritizer,
	}
}

// Enqueue pushes the given value onto the top of the queue and stores the value at given key
func (q GeneralKVQueue) Enqueue(key Key, value codec.ProtoMarshaler) {
	q.store.Set(q.name.Append(q.prioritizer(value)).Append(key), &gogoprototypes.BytesValue{Value: key.AsKey()})
	q.store.Set(key, value)
}

// Dequeue pops the bottom of the queue and unmarshals it into the given object, and return true if anything
// in the queue is found and the value passes the optional filter function
func (q GeneralKVQueue) Dequeue(value codec.ProtoMarshaler, filter ...func(value codec.ProtoMarshaler) bool) bool {
	iter := sdk.KVStorePrefixIterator(q.store.KVStore, q.name.AsKey())
	defer CloseLogError(iter, q.logger)

	if !iter.Valid() {
		return false
	}

	var key gogoprototypes.BytesValue
	q.store.cdc.MustUnmarshalLengthPrefixed(iter.Value(), &key)

	if ok := q.store.Get(KeyFromBz(key.Value), value); !ok {
		return false
	}

	if len(filter) > 0 && !filter[0](value) {
		return false
	}

	q.store.Delete(KeyFromBz(iter.Key()))

	return true
}

// IsEmpty returns true if the queue is empty; otherwise, false
func (q GeneralKVQueue) IsEmpty() bool {
	iter := sdk.KVStorePrefixIterator(q.store.KVStore, q.name.AsKey())
	defer CloseLogError(iter, q.logger)

	return !iter.Valid()
}

// Keys returns a list with the keys for the values still enqueued
func (q GeneralKVQueue) Keys() []Key {
	iter := sdk.KVStorePrefixIterator(q.store.KVStore, q.name.AsKey())
	defer CloseLogError(iter, q.logger)

	var keys []Key
	for ; iter.Valid(); iter.Next() {
		var key gogoprototypes.BytesValue
		q.store.cdc.MustUnmarshalLengthPrefixed(iter.Value(), &key)

		keys = append(keys, KeyFromBz(key.Value))
	}

	return keys
}

// ImportState should only be used to populate state at genesis. Panics if the given state is invalid
func (q GeneralKVQueue) ImportState(state map[string]codec.ProtoMarshaler, validator ...func(map[string]codec.ProtoMarshaler) error) {
	if len(validator) > 0 {
		if err := validator[0](state); err != nil {
			panic(err)
		}
	}

	for key, value := range state {
		q.store.Set(q.name.AppendStr(key), value)
	}
}

// BlockHeightKVQueue is a queue that orders items with the block height at which the items are enqueued;
// the order of items that are enqueued at the same block height is deterministically based on their actual key
// in the KVStore
type BlockHeightKVQueue struct {
	GeneralKVQueue
}

// NewBlockHeightKVQueue is the constructor of BlockHeightKVQueue
func NewBlockHeightKVQueue(name string, store KVStore, blockHeight int64, logger log.Logger) BlockHeightKVQueue {
	bz := make([]byte, 8)
	binary.BigEndian.PutUint64(bz, uint64(blockHeight))

	return BlockHeightKVQueue{
		GeneralKVQueue: NewGeneralKVQueue(name, store, logger, func(_ codec.ProtoMarshaler) Key {
			return KeyFromBz(bz)
		}),
	}
}

// SequenceKVQueue is a queue that orders items with the sequence number at which the items are enqueued;
type SequenceKVQueue struct {
	store   KVStore
	maxSize uint64
	seqNo   uint64
	size    uint64
	logger  log.Logger
}

var (
	sizeKey     = KeyFromStr("size")
	sequenceKey = KeyFromStr("sequence")
	queueKey    = KeyFromStr("queue")
)

// NewSequenceKVQueue is the constructor of SequenceKVQueue. The prefixStore should isolate the namespace of the queue
func NewSequenceKVQueue(prefixStore KVStore, maxSize uint64, logger log.Logger) SequenceKVQueue {
	var seqNo uint64
	bz := prefixStore.GetRaw(sequenceKey)
	if bz != nil {
		seqNo = binary.BigEndian.Uint64(bz)
	}

	var size uint64
	bz = prefixStore.GetRaw(sizeKey)
	if bz != nil {
		size = binary.BigEndian.Uint64(bz)
	}

	return SequenceKVQueue{store: prefixStore, maxSize: maxSize, seqNo: seqNo, size: size, logger: logger}
}

// Enqueue pushes the given value onto the top of the queue and stores the value at given key. Returns queue position and status
func (q *SequenceKVQueue) Enqueue(value codec.ProtoMarshaler) error {
	if q.size >= q.maxSize {
		return fmt.Errorf("queue is full")
	}

	bz := make([]byte, 8)
	binary.BigEndian.PutUint64(bz, q.seqNo)

	q.store.Set(queueKey.Append(KeyFromBz(bz)), value)
	q.incrSize()
	q.incrSeqNo()

	return nil
}

// Dequeue pops the nth value off and unmarshals it into the given object, returns false if n is out of range
// n is 0-based
func (q *SequenceKVQueue) Dequeue(n uint64, value codec.ProtoMarshaler) bool {
	key, ok := q.peek(n, value)
	if ok {
		q.store.Delete(key)
		q.decrSize()
	}
	return ok
}

// Size returns the given current size of the queue
func (q SequenceKVQueue) Size() uint64 {
	return q.size
}

// Peek unmarshals the value into the given object at the given index and returns its sequence number, returns false if n is out of range
// n is 0-based
func (q SequenceKVQueue) Peek(n uint64, value codec.ProtoMarshaler) bool {
	_, ok := q.peek(n, value)
	return ok
}

func (q SequenceKVQueue) peek(n uint64, value codec.ProtoMarshaler) (Key, bool) {
	if n >= q.size {
		return nil, false
	}

	var i uint64
	iter := q.store.Iterator(queueKey)
	defer CloseLogError(iter, q.logger)

	for ; iter.Valid() && i < n; iter.Next() {
		i++
	}
	iter.UnmarshalValue(value)
	return iter.GetKey(), true
}

func (q *SequenceKVQueue) incrSize() {
	q.size++
	bz := make([]byte, 8)
	binary.BigEndian.PutUint64(bz, q.size)
	q.store.SetRaw(sizeKey, bz)
}

func (q *SequenceKVQueue) decrSize() {
	if q.size > 0 {
		q.size--
		bz := make([]byte, 8)
		binary.BigEndian.PutUint64(bz, q.size)
		q.store.SetRaw(sizeKey, bz)
	}
}

func (q *SequenceKVQueue) incrSeqNo() {
	q.seqNo++
	bz := make([]byte, 8)
	binary.BigEndian.PutUint64(bz, q.seqNo)
	q.store.SetRaw(sequenceKey, bz)
}