Add LZ4-based allocator hibernation

Signed-off-by: Vadim Markovtsev <vadim@sourced.tech>

internal/rbtree/lz4.go

@@ -0,0 +1,32 @@
+package rbtree
+
+/*
+#cgo CFLAGS: -std=c99
+#include "lz4hc.c"
+*/
+import "C"
+import "unsafe"
+
+// CompressUInt32Slice compresses a slice of uint32-s with LZ4.
+func CompressUInt32Slice(data []uint32) []byte {
+	dstSize := C.LZ4_compressBound(C.int(len(data) * 4))
+	dst := make([]byte, dstSize)
+	dstSize = C.LZ4_compress_HC(
+		unsafe.Pointer(&data[0]),
+		unsafe.Pointer(&dst[0]),
+		C.int(len(data) * 4),
+		dstSize,
+		12)
+	finalDst := make([]byte, dstSize)
+	copy(finalDst, dst[:dstSize])
+	return finalDst
+}
+
+// DecompressUInt32Slice decompresses a slice of uint32-s previously compressed with LZ4.
+// `result` must be preallocated.
+func DecompressUInt32Slice(data []byte, result []uint32) {
+	C.LZ4_decompress_fast(
+		unsafe.Pointer(&data[0]),
+		unsafe.Pointer(&result[0]),
+		C.int(len(result) * 4))
+}

internal/rbtree/lz4_test.go

@@ -0,0 +1,23 @@
+package rbtree
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCompressDecompressUInt32Slice(t *testing.T) {
+	data := make([]uint32, 1000)
+	for i := range data {
+		data[i] = 7
+	}
+	packed := CompressUInt32Slice(data)
+	assert.Len(t, packed, 29)
+	for i := range data {
+		data[i] = 0
+	}
+	DecompressUInt32Slice(packed, data)
+	for i := range data {
+		assert.Equal(t, uint32(7), data[i], i)
+	}
+}

internal/rbtree/rbtree.go

@@ -2,6 +2,7 @@ package rbtree
 
 import (
 	"math"
+	"sync"
 )
 
 //
@@ -16,8 +17,12 @@ type Item struct {
 
 // Allocator is the allocator for nodes in a RBTree.
 type Allocator struct {
+	HibernationThreshold int
+
 	storage []node
 	gaps map[uint32]bool
+	hibernatedStorage [6][]byte
+	hibernatedLen int
 }
 
 // NewAllocator creates a new allocator for RBTree's nodes.
@@ -51,6 +56,71 @@ func (allocator *Allocator) Clone() *Allocator {
 	return newAllocator
 }
 
+// Hibernate compresses the allocated memory.
+func (allocator *Allocator) Hibernate() {
+	if allocator.HibernationThreshold == 0 || len(allocator.storage) < allocator.HibernationThreshold {
+		return
+	}
+	allocator.hibernatedLen = len(allocator.storage)
+	buffers := [6][]uint32{}
+	for i := 0; i < len(buffers); i++ {
+		buffers[i] = make([]uint32, len(allocator.storage))
+	}
+	// we deinterleave to achieve a better compression ratio
+	for i, n := range allocator.storage {
+		buffers[0][i] = n.item.Key
+		buffers[1][i] = n.item.Value
+		buffers[2][i] = n.left
+		buffers[3][i] = n.parent
+		buffers[4][i] = n.right
+		if n.color {
+			buffers[5][i] = 1
+		}
+	}
+	allocator.storage = nil
+	wg := &sync.WaitGroup{}
+	wg.Add(len(buffers))
+	for i, buffer := range buffers {
+		go func(i int, buffer []uint32) {
+			allocator.hibernatedStorage[i] = CompressUInt32Slice(buffer)
+			buffers[i] = nil
+			wg.Done()
+		}(i, buffer)
+	}
+	wg.Wait()
+}
+
+// Boot performs the opposite of Hibernate() - decompresses and restores the allocated memory.
+func (allocator *Allocator) Boot() {
+	if allocator.hibernatedLen == 0 {
+		// not hibernated
+		return
+	}
+	buffers := [6][]uint32{}
+	wg := &sync.WaitGroup{}
+	wg.Add(len(buffers))
+	for i := 0; i < len(buffers); i++ {
+		go func(i int) {
+			buffers[i] = make([]uint32, allocator.hibernatedLen)
+			DecompressUInt32Slice(allocator.hibernatedStorage[i], buffers[i])
+			allocator.hibernatedStorage[i] = nil
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+	allocator.storage = make([]node, allocator.hibernatedLen, (allocator.hibernatedLen * 3) / 2)
+	for i := range allocator.storage {
+		n := &allocator.storage[i]
+		n.item.Key = buffers[0][i]
+		n.item.Value = buffers[1][i]
+		n.left = buffers[2][i]
+		n.parent = buffers[3][i]
+		n.right = buffers[4][i]
+		n.color = buffers[5][i] > 0
+	}
+	allocator.hibernatedLen = 0
+}
+
 func (allocator *Allocator) malloc() uint32 {
 	if len(allocator.gaps) > 0 {
 		var key uint32