github.com/MetalBlockchain/metalgo@v1.11.9/x/merkledb/intermediate_node_db.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package merkledb
     5  
     6  import (
     7  	"github.com/MetalBlockchain/metalgo/cache"
     8  	"github.com/MetalBlockchain/metalgo/database"
     9  	"github.com/MetalBlockchain/metalgo/utils"
    10  )
    11  
    12  // Holds intermediate nodes. That is, those without values.
    13  // Changes to this database aren't written to [baseDB] until
    14  // they're evicted from the [nodeCache] or Flush is called.
    15  type intermediateNodeDB struct {
    16  	bufferPool *utils.BytesPool
    17  
    18  	// The underlying storage.
    19  	// Keys written to [baseDB] are prefixed with [intermediateNodePrefix].
    20  	baseDB database.Database
    21  
    22  	// The write buffer contains nodes that have been changed but have not been written to disk.
    23  	// Note that a call to Put may cause a node to be evicted
    24  	// from the cache, which will call [OnEviction].
    25  	// A non-nil error returned from Put is considered fatal.
    26  	// Keys in [nodeCache] aren't prefixed with [intermediateNodePrefix].
    27  	writeBuffer onEvictCache[Key, *node]
    28  
    29  	// If a value is nil, the corresponding key isn't in the trie.
    30  	nodeCache cache.Cacher[Key, *node]
    31  
    32  	// the number of bytes to evict during an eviction batch
    33  	evictionBatchSize int
    34  	metrics           metrics
    35  	tokenSize         int
    36  	hasher            Hasher
    37  }
    38  
    39  func newIntermediateNodeDB(
    40  	db database.Database,
    41  	bufferPool *utils.BytesPool,
    42  	metrics metrics,
    43  	cacheSize int,
    44  	writeBufferSize int,
    45  	evictionBatchSize int,
    46  	tokenSize int,
    47  	hasher Hasher,
    48  ) *intermediateNodeDB {
    49  	result := &intermediateNodeDB{
    50  		metrics:           metrics,
    51  		baseDB:            db,
    52  		bufferPool:        bufferPool,
    53  		evictionBatchSize: evictionBatchSize,
    54  		tokenSize:         tokenSize,
    55  		hasher:            hasher,
    56  		nodeCache:         cache.NewSizedLRU(cacheSize, cacheEntrySize),
    57  	}
    58  	result.writeBuffer = newOnEvictCache(
    59  		writeBufferSize,
    60  		cacheEntrySize,
    61  		result.onEviction,
    62  	)
    63  
    64  	return result
    65  }
    66  
    67  // A non-nil error is considered fatal and closes [db.baseDB].
    68  func (db *intermediateNodeDB) onEviction(key Key, n *node) error {
    69  	writeBatch := db.baseDB.NewBatch()
    70  	totalSize := cacheEntrySize(key, n)
    71  	if err := db.addToBatch(writeBatch, key, n); err != nil {
    72  		_ = db.baseDB.Close()
    73  		return err
    74  	}
    75  
    76  	// Evict the oldest [evictionBatchSize] nodes from the cache
    77  	// and write them to disk. We write a batch of them, rather than
    78  	// just [n], so that we don't immediately evict and write another
    79  	// node, because each time this method is called we do a disk write.
    80  	// Evicts a total number of bytes, rather than a number of nodes
    81  	for totalSize < db.evictionBatchSize {
    82  		key, n, exists := db.writeBuffer.removeOldest()
    83  		if !exists {
    84  			// The cache is empty.
    85  			break
    86  		}
    87  		totalSize += cacheEntrySize(key, n)
    88  		if err := db.addToBatch(writeBatch, key, n); err != nil {
    89  			_ = db.baseDB.Close()
    90  			return err
    91  		}
    92  	}
    93  	if err := writeBatch.Write(); err != nil {
    94  		_ = db.baseDB.Close()
    95  		return err
    96  	}
    97  	return nil
    98  }
    99  
   100  func (db *intermediateNodeDB) addToBatch(b database.KeyValueWriterDeleter, key Key, n *node) error {
   101  	dbKey := db.constructDBKey(key)
   102  	defer db.bufferPool.Put(dbKey)
   103  
   104  	db.metrics.DatabaseNodeWrite()
   105  	if n == nil {
   106  		return b.Delete(*dbKey)
   107  	}
   108  	return b.Put(*dbKey, n.bytes())
   109  }
   110  
   111  func (db *intermediateNodeDB) Get(key Key) (*node, error) {
   112  	if cachedValue, isCached := db.nodeCache.Get(key); isCached {
   113  		db.metrics.IntermediateNodeCacheHit()
   114  		if cachedValue == nil {
   115  			return nil, database.ErrNotFound
   116  		}
   117  		return cachedValue, nil
   118  	}
   119  	if cachedValue, isCached := db.writeBuffer.Get(key); isCached {
   120  		db.metrics.IntermediateNodeCacheHit()
   121  		if cachedValue == nil {
   122  			return nil, database.ErrNotFound
   123  		}
   124  		return cachedValue, nil
   125  	}
   126  	db.metrics.IntermediateNodeCacheMiss()
   127  
   128  	dbKey := db.constructDBKey(key)
   129  	defer db.bufferPool.Put(dbKey)
   130  
   131  	db.metrics.DatabaseNodeRead()
   132  	nodeBytes, err := db.baseDB.Get(*dbKey)
   133  	if err != nil {
   134  		return nil, err
   135  	}
   136  
   137  	return parseNode(db.hasher, key, nodeBytes)
   138  }
   139  
   140  // constructDBKey returns a key that can be used in [db.baseDB].
   141  // We need to be able to differentiate between two keys of equal
   142  // byte length but different bit length, so we add padding to differentiate.
   143  // Additionally, we add a prefix indicating it is part of the intermediateNodeDB.
   144  func (db *intermediateNodeDB) constructDBKey(key Key) *[]byte {
   145  	if db.tokenSize == 8 {
   146  		// For tokens of size byte, no padding is needed since byte
   147  		// length == token length
   148  		return addPrefixToKey(db.bufferPool, intermediateNodePrefix, key.Bytes())
   149  	}
   150  
   151  	var (
   152  		prefixLen              = len(intermediateNodePrefix)
   153  		prefixBitLen           = 8 * prefixLen
   154  		dualIndex              = dualBitIndex(db.tokenSize)
   155  		paddingByteValue  byte = 1 << dualIndex
   156  		paddingSliceValue      = []byte{paddingByteValue}
   157  		paddingKey             = Key{
   158  			value:  byteSliceToString(paddingSliceValue),
   159  			length: db.tokenSize,
   160  		}
   161  	)
   162  
   163  	bufferPtr := db.bufferPool.Get(bytesNeeded(prefixBitLen + key.length + db.tokenSize))
   164  	copy(*bufferPtr, intermediateNodePrefix)                          // add prefix
   165  	copy((*bufferPtr)[prefixLen:], key.Bytes())                       // add key
   166  	extendIntoBuffer(*bufferPtr, paddingKey, prefixBitLen+key.length) // add padding
   167  	return bufferPtr
   168  }
   169  
   170  func (db *intermediateNodeDB) Put(key Key, n *node) error {
   171  	db.nodeCache.Put(key, n)
   172  	return db.writeBuffer.Put(key, n)
   173  }
   174  
   175  func (db *intermediateNodeDB) Flush() error {
   176  	db.nodeCache.Flush()
   177  	return db.writeBuffer.Flush()
   178  }
   179  
   180  func (db *intermediateNodeDB) Delete(key Key) error {
   181  	db.nodeCache.Put(key, nil)
   182  	return db.writeBuffer.Put(key, nil)
   183  }
   184  
   185  func (db *intermediateNodeDB) Clear() error {
   186  	db.nodeCache.Flush()
   187  
   188  	// Reset the buffer. Note we don't flush because that would cause us to
   189  	// persist intermediate nodes we're about to delete.
   190  	db.writeBuffer = newOnEvictCache(
   191  		db.writeBuffer.maxSize,
   192  		db.writeBuffer.size,
   193  		db.writeBuffer.onEviction,
   194  	)
   195  	return database.AtomicClearPrefix(db.baseDB, db.baseDB, intermediateNodePrefix)
   196  }