github.com/MetalBlockchain/metalgo@v1.11.9/database/linkeddb/linkeddb.go

// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package linkeddb

import (
	"slices"
	"sync"

	"github.com/MetalBlockchain/metalgo/cache"
	"github.com/MetalBlockchain/metalgo/database"
)

const (
	defaultCacheSize = 1024
)

var (
	headKey = []byte{0x01}

	_ LinkedDB          = (*linkedDB)(nil)
	_ database.Iterator = (*iterator)(nil)
)

// LinkedDB provides a key value interface while allowing iteration.
type LinkedDB interface {
	database.KeyValueReaderWriterDeleter

	IsEmpty() (bool, error)
	HeadKey() ([]byte, error)
	Head() (key []byte, value []byte, err error)

	NewIterator() database.Iterator
	NewIteratorWithStart(start []byte) database.Iterator
}

type linkedDB struct {
	// lock ensure that this data structure handles its thread safety correctly.
	lock sync.RWMutex

	cacheLock sync.Mutex
	// these variables provide caching for the head key.
	headKeyIsSynced, headKeyExists, headKeyIsUpdated, updatedHeadKeyExists bool
	headKey, updatedHeadKey                                                []byte
	// these variables provide caching for the nodes.
	nodeCache    cache.Cacher[string, *node] // key -> *node
	updatedNodes map[string]*node

	// db is the underlying database that this list is stored in.
	db database.Database
	// batch writes to [db] atomically.
	batch database.Batch
}

type node struct {
	Value       []byte `serialize:"true"`
	HasNext     bool   `serialize:"true"`
	Next        []byte `serialize:"true"`
	HasPrevious bool   `serialize:"true"`
	Previous    []byte `serialize:"true"`
}

func New(db database.Database, cacheSize int) LinkedDB {
	return &linkedDB{
		nodeCache:    &cache.LRU[string, *node]{Size: cacheSize},
		updatedNodes: make(map[string]*node),
		db:           db,
		batch:        db.NewBatch(),
	}
}

func NewDefault(db database.Database) LinkedDB {
	return New(db, defaultCacheSize)
}

func (ldb *linkedDB) Has(key []byte) (bool, error) {
	ldb.lock.RLock()
	defer ldb.lock.RUnlock()

	return ldb.db.Has(nodeKey(key))
}

func (ldb *linkedDB) Get(key []byte) ([]byte, error) {
	ldb.lock.RLock()
	defer ldb.lock.RUnlock()

	node, err := ldb.getNode(key)
	return node.Value, err
}

func (ldb *linkedDB) Put(key, value []byte) error {
	ldb.lock.Lock()
	defer ldb.lock.Unlock()

	ldb.resetBatch()

	// If the key already has a node in the list, update that node.
	existingNode, err := ldb.getNode(key)
	if err == nil {
		existingNode.Value = slices.Clone(value)
		if err := ldb.putNode(key, existingNode); err != nil {
			return err
		}
		return ldb.writeBatch()
	}
	if err != database.ErrNotFound {
		return err
	}

	// The key isn't currently in the list, so we should add it as the head.
	newHead := node{Value: slices.Clone(value)}
	if headKey, err := ldb.getHeadKey(); err == nil {
		// The list currently has a head, so we need to update the old head.
		oldHead, err := ldb.getNode(headKey)
		if err != nil {
			return err
		}
		oldHead.HasPrevious = true
		oldHead.Previous = key
		if err := ldb.putNode(headKey, oldHead); err != nil {
			return err
		}

		newHead.HasNext = true
		newHead.Next = headKey
	} else if err != database.ErrNotFound {
		return err
	}
	if err := ldb.putNode(key, newHead); err != nil {
		return err
	}
	if err := ldb.putHeadKey(key); err != nil {
		return err
	}
	return ldb.writeBatch()
}

func (ldb *linkedDB) Delete(key []byte) error {
	ldb.lock.Lock()
	defer ldb.lock.Unlock()

	currentNode, err := ldb.getNode(key)
	if err == database.ErrNotFound {
		return nil
	}
	if err != nil {
		return err
	}

	ldb.resetBatch()

	// We're trying to delete this node.
	if err := ldb.deleteNode(key); err != nil {
		return err
	}

	switch {
	case currentNode.HasPrevious:
		// We aren't modifying the head.
		previousNode, err := ldb.getNode(currentNode.Previous)
		if err != nil {
			return err
		}
		previousNode.HasNext = currentNode.HasNext
		previousNode.Next = currentNode.Next
		if err := ldb.putNode(currentNode.Previous, previousNode); err != nil {
			return err
		}
		if currentNode.HasNext {
			// We aren't modifying the tail.
			nextNode, err := ldb.getNode(currentNode.Next)
			if err != nil {
				return err
			}
			nextNode.HasPrevious = true
			nextNode.Previous = currentNode.Previous
			if err := ldb.putNode(currentNode.Next, nextNode); err != nil {
				return err
			}
		}
	case !currentNode.HasNext:
		// This is the only node, so we don't have a head anymore.
		if err := ldb.deleteHeadKey(); err != nil {
			return err
		}
	default:
		// The next node will be the new head.
		if err := ldb.putHeadKey(currentNode.Next); err != nil {
			return err
		}
		nextNode, err := ldb.getNode(currentNode.Next)
		if err != nil {
			return err
		}
		nextNode.HasPrevious = false
		nextNode.Previous = nil
		if err := ldb.putNode(currentNode.Next, nextNode); err != nil {
			return err
		}
	}
	return ldb.writeBatch()
}

func (ldb *linkedDB) IsEmpty() (bool, error) {
	_, err := ldb.HeadKey()
	if err == database.ErrNotFound {
		return true, nil
	}
	return false, err
}

func (ldb *linkedDB) HeadKey() ([]byte, error) {
	ldb.lock.RLock()
	defer ldb.lock.RUnlock()

	return ldb.getHeadKey()
}

func (ldb *linkedDB) Head() ([]byte, []byte, error) {
	ldb.lock.RLock()
	defer ldb.lock.RUnlock()

	headKey, err := ldb.getHeadKey()
	if err != nil {
		return nil, nil, err
	}
	head, err := ldb.getNode(headKey)
	return headKey, head.Value, err
}

// This iterator does not guarantee that keys are returned in lexicographic
// order.
func (ldb *linkedDB) NewIterator() database.Iterator {
	return &iterator{ldb: ldb}
}

// NewIteratorWithStart returns an iterator that starts at [start].
// This iterator does not guarantee that keys are returned in lexicographic
// order.
// If [start] is not in the list, starts iterating from the list head.
func (ldb *linkedDB) NewIteratorWithStart(start []byte) database.Iterator {
	hasStartKey, err := ldb.Has(start)
	if err == nil && hasStartKey {
		return &iterator{
			ldb:         ldb,
			initialized: true,
			nextKey:     start,
		}
	}
	// If the start key isn't present, start from the head
	return ldb.NewIterator()
}

func (ldb *linkedDB) getHeadKey() ([]byte, error) {
	// If the ldb read lock is held, then there needs to be additional
	// synchronization here to avoid racy behavior.
	ldb.cacheLock.Lock()
	defer ldb.cacheLock.Unlock()

	if ldb.headKeyIsSynced {
		if ldb.headKeyExists {
			return ldb.headKey, nil
		}
		return nil, database.ErrNotFound
	}
	headKey, err := ldb.db.Get(headKey)
	if err == nil {
		ldb.headKeyIsSynced = true
		ldb.headKeyExists = true
		ldb.headKey = headKey
		return headKey, nil
	}
	if err == database.ErrNotFound {
		ldb.headKeyIsSynced = true
		ldb.headKeyExists = false
		return nil, database.ErrNotFound
	}
	return headKey, err
}

func (ldb *linkedDB) putHeadKey(key []byte) error {
	ldb.headKeyIsUpdated = true
	ldb.updatedHeadKeyExists = true
	ldb.updatedHeadKey = key
	return ldb.batch.Put(headKey, key)
}

func (ldb *linkedDB) deleteHeadKey() error {
	ldb.headKeyIsUpdated = true
	ldb.updatedHeadKeyExists = false
	return ldb.batch.Delete(headKey)
}

func (ldb *linkedDB) getNode(key []byte) (node, error) {
	// If the ldb read lock is held, then there needs to be additional
	// synchronization here to avoid racy behavior.
	ldb.cacheLock.Lock()
	defer ldb.cacheLock.Unlock()

	keyStr := string(key)
	if n, exists := ldb.nodeCache.Get(keyStr); exists {
		if n == nil {
			return node{}, database.ErrNotFound
		}
		return *n, nil
	}

	nodeBytes, err := ldb.db.Get(nodeKey(key))
	if err == database.ErrNotFound {
		ldb.nodeCache.Put(keyStr, nil)
		return node{}, err
	}
	if err != nil {
		return node{}, err
	}
	n := node{}
	_, err = Codec.Unmarshal(nodeBytes, &n)
	if err == nil {
		ldb.nodeCache.Put(keyStr, &n)
	}
	return n, err
}

func (ldb *linkedDB) putNode(key []byte, n node) error {
	ldb.updatedNodes[string(key)] = &n
	nodeBytes, err := Codec.Marshal(CodecVersion, n)
	if err != nil {
		return err
	}
	return ldb.batch.Put(nodeKey(key), nodeBytes)
}

func (ldb *linkedDB) deleteNode(key []byte) error {
	ldb.updatedNodes[string(key)] = nil
	return ldb.batch.Delete(nodeKey(key))
}

func (ldb *linkedDB) resetBatch() {
	ldb.headKeyIsUpdated = false
	clear(ldb.updatedNodes)
	ldb.batch.Reset()
}

func (ldb *linkedDB) writeBatch() error {
	if err := ldb.batch.Write(); err != nil {
		return err
	}
	if ldb.headKeyIsUpdated {
		ldb.headKeyIsSynced = true
		ldb.headKeyExists = ldb.updatedHeadKeyExists
		ldb.headKey = ldb.updatedHeadKey
	}
	for key, n := range ldb.updatedNodes {
		ldb.nodeCache.Put(key, n)
	}
	return nil
}

type iterator struct {
	ldb                    *linkedDB
	initialized, exhausted bool
	key, value, nextKey    []byte
	err                    error
}

func (it *iterator) Next() bool {
	// If the iterator has been exhausted, there is no next value.
	if it.exhausted {
		it.key = nil
		it.value = nil
		return false
	}

	it.ldb.lock.RLock()
	defer it.ldb.lock.RUnlock()

	// If the iterator was not yet initialized, do it now.
	if !it.initialized {
		it.initialized = true
		headKey, err := it.ldb.getHeadKey()
		if err == database.ErrNotFound {
			it.exhausted = true
			it.key = nil
			it.value = nil
			return false
		}
		if err != nil {
			it.exhausted = true
			it.key = nil
			it.value = nil
			it.err = err
			return false
		}
		it.nextKey = headKey
	}

	nextNode, err := it.ldb.getNode(it.nextKey)
	if err == database.ErrNotFound {
		it.exhausted = true
		it.key = nil
		it.value = nil
		return false
	}
	if err != nil {
		it.exhausted = true
		it.key = nil
		it.value = nil
		it.err = err
		return false
	}
	it.key = it.nextKey
	it.value = nextNode.Value
	it.nextKey = nextNode.Next
	it.exhausted = !nextNode.HasNext
	return true
}

func (it *iterator) Error() error {
	return it.err
}

func (it *iterator) Key() []byte {
	return it.key
}

func (it *iterator) Value() []byte {
	return it.value
}

func (*iterator) Release() {}

func nodeKey(key []byte) []byte {
	newKey := make([]byte, len(key)+1)
	copy(newKey[1:], key)
	return newKey
}