Add treapNode pool. Reduce cloneTreapNode() allocations.

database/ffldb/dbcache.go

@@ -285,9 +285,12 @@ func (iter *dbCacheIterator) Error() error {
 // dbCacheSnapshot defines a snapshot of the database cache and underlying
 // database at a particular point in time.
 type dbCacheSnapshot struct {
-	dbSnapshot    *leveldb.Snapshot
-	pendingKeys   *treap.Immutable
-	pendingRemove *treap.Immutable
+	dbSnapshot        *leveldb.Snapshot
+	pendingKeys       *treap.Immutable
+	pendingRemove     *treap.Immutable
+	pendingKeysSnap   *treap.SnapRecord
+	pendingRemoveSnap *treap.SnapRecord
+	cacheFlushed      int
 }
 
 // Has returns whether or not the passed key exists.
@@ -327,6 +330,18 @@ func (snap *dbCacheSnapshot) Get(key []byte) []byte {
 // Release releases the snapshot.
 func (snap *dbCacheSnapshot) Release() {
 	snap.dbSnapshot.Release()
+	if snap.cacheFlushed > 0 && snap.pendingKeys != nil {
+		snap.pendingKeys.Recycle(snap.pendingKeysSnap)
+	}
+	if snap.cacheFlushed > 0 && snap.pendingRemove != nil {
+		snap.pendingRemove.Recycle(snap.pendingRemoveSnap)
+	}
+	if snap.pendingKeysSnap != nil {
+		snap.pendingKeysSnap.Release()
+	}
+	if snap.pendingRemoveSnap != nil {
+		snap.pendingRemoveSnap.Release()
+	}
 	snap.pendingKeys = nil
 	snap.pendingRemove = nil
 }
@@ -411,6 +426,12 @@ func (c *dbCache) Snapshot() (*dbCacheSnapshot, error) {
 		pendingKeys:   c.cachedKeys,
 		pendingRemove: c.cachedRemove,
 	}
+	if cacheSnapshot.pendingKeys != nil {
+		cacheSnapshot.pendingKeysSnap = cacheSnapshot.pendingKeys.Snapshot()
+	}
+	if cacheSnapshot.pendingRemove != nil {
+		cacheSnapshot.pendingRemoveSnap = cacheSnapshot.pendingRemove.Snapshot()
+	}
 	c.cacheLock.RUnlock()
 	return cacheSnapshot, nil
 }
@@ -485,7 +506,7 @@ func (c *dbCache) commitTreaps(pendingKeys, pendingRemove TreapForEacher) error
 // cache to the underlying database.
 //
 // This function MUST be called with the database write lock held.
-func (c *dbCache) flush() error {
+func (c *dbCache) flush(tx *transaction) error {
 	c.lastFlush = time.Now()
 
 	// Sync the current write file associated with the block store.  This is
@@ -499,12 +520,10 @@ func (c *dbCache) flush() error {
 	// Since the cached keys to be added and removed use an immutable treap,
 	// a snapshot is simply obtaining the root of the tree under the lock
 	// which is used to atomically swap the root.
-	c.cacheLock.Lock()
+	c.cacheLock.RLock()
 	cachedKeys := c.cachedKeys
 	cachedRemove := c.cachedRemove
-	c.cachedKeys = treap.NewImmutable()
-	c.cachedRemove = treap.NewImmutable()
-	c.cacheLock.Unlock()
+	c.cacheLock.RUnlock()
 
 	// Nothing to do if there is no data to flush.
 	if cachedKeys.Len() == 0 && cachedRemove.Len() == 0 {
@@ -516,6 +535,20 @@ func (c *dbCache) flush() error {
 		return err
 	}
 
+	c.cacheLock.Lock()
+	c.cachedKeys = treap.NewImmutable()
+	c.cachedRemove = treap.NewImmutable()
+	c.cacheLock.Unlock()
+
+	cachedKeys.Recycle(nil)
+	cachedRemove.Recycle(nil)
+
+	// Make a note that cache was flushed so tx.snapshot.Release()
+	// can also call Recycle() to free more nodes.
+	if tx != nil && tx.snapshot != nil {
+		tx.snapshot.cacheFlushed++
+	}
+
 	return nil
 }
 
@@ -564,7 +597,7 @@ func (c *dbCache) commitTx(tx *transaction) error {
 	// Flush the cache and write the current transaction directly to the
 	// database if a flush is needed.
 	if c.needsFlush(tx) {
-		if err := c.flush(); err != nil {
+		if err := c.flush(tx); err != nil {
 			return err
 		}
 
@@ -574,9 +607,16 @@ func (c *dbCache) commitTx(tx *transaction) error {
 			return err
 		}
 
+		pk := tx.pendingKeys
+		pr := tx.pendingRemove
+
 		// Clear the transaction entries since they have been committed.
 		tx.pendingKeys = nil
 		tx.pendingRemove = nil
+
+		pk.Recycle()
+		pr.Recycle()
+
 		return nil
 	}
 
@@ -593,26 +633,31 @@ func (c *dbCache) commitTx(tx *transaction) error {
 
 	// Apply every key to add in the database transaction to the cache.
 	tx.pendingKeys.ForEach(func(k, v []byte) bool {
-		newCachedRemove = newCachedRemove.Delete(k)
-		newCachedKeys = newCachedKeys.Put(k, v)
+		treap.DeleteM(&newCachedRemove, k, tx.snapshot.pendingRemoveSnap)
+		treap.PutM(&newCachedKeys, k, v, tx.snapshot.pendingKeysSnap)
 		return true
 	})
+	pk := tx.pendingKeys
 	tx.pendingKeys = nil
+	pk.Recycle()
 
 	// Apply every key to remove in the database transaction to the cache.
 	tx.pendingRemove.ForEach(func(k, v []byte) bool {
-		newCachedKeys = newCachedKeys.Delete(k)
-		newCachedRemove = newCachedRemove.Put(k, nil)
+		treap.DeleteM(&newCachedKeys, k, tx.snapshot.pendingKeysSnap)
+		treap.PutM(&newCachedRemove, k, nil, tx.snapshot.pendingRemoveSnap)
 		return true
 	})
+	pr := tx.pendingRemove
 	tx.pendingRemove = nil
+	pr.Recycle()
 
 	// Atomically replace the immutable treaps which hold the cached keys to
 	// add and delete.
 	c.cacheLock.Lock()
 	c.cachedKeys = newCachedKeys
 	c.cachedRemove = newCachedRemove
 	c.cacheLock.Unlock()
+
 	return nil
 }
 
@@ -622,7 +667,7 @@ func (c *dbCache) commitTx(tx *transaction) error {
 // This function MUST be called with the database write lock held.
 func (c *dbCache) Close() error {
 	// Flush any outstanding cached entries to disk.
-	if err := c.flush(); err != nil {
+	if err := c.flush(nil); err != nil {
 		// Even if there is an error while flushing, attempt to close
 		// the underlying database.  The error is ignored since it would
 		// mask the flush error.

database/ffldb/whitebox_test.go

@@ -319,7 +319,7 @@ func testWriteFailures(tc *testContext) bool {
 		file: &mockFile{forceSyncErr: true, maxSize: -1},
 	}
 	store.writeCursor.Unlock()
-	err := tc.db.(*db).cache.flush()
+	err := tc.db.(*db).cache.flush(nil)
 	if !checkDbError(tc.t, testName, err, database.ErrDriverSpecific) {
 		return false
 	}

database/internal/treap/common.go

@@ -6,6 +6,7 @@ package treap
 
 import (
 	"math/rand"
+	"sync"
 	"time"
 )
 
@@ -23,7 +24,7 @@ const (
 	// size in that case is acceptable since it avoids the need to import
 	// unsafe.  It consists of 24-bytes for each key and value + 8 bytes for
 	// each of the priority, left, and right fields (24*2 + 8*3).
-	nodeFieldsSize = 72
+	nodeFieldsSize = 96
 )
 
 var (
@@ -33,13 +34,24 @@ var (
 	emptySlice = make([]byte, 0)
 )
 
+const (
+	// mutableGeneration is the generation number for nodes in a Mutable treap.
+	mutableGeneration int = -1
+	// recycleGeneration indicates node is scheduled for recycling back to nodePool.
+	recycleGeneration int = -2
+	// poolGeneration is the generation number for free nodes in the nodePool.
+	poolGeneration int = -3
+)
+
 // treapNode represents a node in the treap.
 type treapNode struct {
-	key      []byte
-	value    []byte
-	priority int
-	left     *treapNode
-	right    *treapNode
+	key        []byte
+	value      []byte
+	priority   int
+	left       *treapNode
+	right      *treapNode
+	generation int
+	next       *treapNode
 }
 
 // nodeSize returns the number of bytes the specified node occupies including
@@ -48,10 +60,40 @@ func nodeSize(node *treapNode) uint64 {
 	return nodeFieldsSize + uint64(len(node.key)+len(node.value))
 }
 
-// newTreapNode returns a new node from the given key, value, and priority.  The
-// node is not initially linked to any others.
-func newTreapNode(key, value []byte, priority int) *treapNode {
-	return &treapNode{key: key, value: value, priority: priority}
+func poolNewTreapNode() interface{} {
+	return &treapNode{key: nil, value: nil, priority: 0, generation: poolGeneration}
+}
+
+// Pool of treapNode available for reuse.
+var nodePool sync.Pool = sync.Pool{New: poolNewTreapNode}
+
+// getTreapNode returns a node from nodePool with the given key, value, priority,
+// and generation. The node is not initially linked to any others.
+func getTreapNode(key, value []byte, priority int, generation int) *treapNode {
+	n := nodePool.Get().(*treapNode)
+	n.key = key
+	n.value = value
+	n.priority = priority
+	n.left = nil
+	n.right = nil
+	n.generation = generation
+	n.next = nil
+	return n
+}
+
+// putTreapNode returns a node back to nodePool for reuse.
+func putTreapNode(n *treapNode) {
+	if n.generation <= poolGeneration {
+		panic("double free of treapNode detected")
+	}
+	n.key = nil
+	n.value = nil
+	n.priority = 0
+	n.left = nil
+	n.right = nil
+	n.generation = poolGeneration
+	n.next = nil
+	nodePool.Put(n)
 }
 
 // parentStack represents a stack of parent treap nodes that are used during

database/internal/treap/common_test.go

@@ -49,7 +49,7 @@ testLoop:
 		for j := 0; j < test.numNodes; j++ {
 			var key [4]byte
 			binary.BigEndian.PutUint32(key[:], uint32(j))
-			node := newTreapNode(key[:], key[:], 0)
+			node := getTreapNode(key[:], key[:], 0, 0)
 			nodes = append(nodes, node)
 		}