From 4178d916c2a2e61e7380f2e1c051e8ef890df31e Mon Sep 17 00:00:00 2001 From: Mikers Date: Mon, 6 Jul 2026 06:34:14 -1000 Subject: [PATCH] treedb: pool append-only key arena chunks --- TreeDB/caching/db.go | 35 +++ TreeDB/caching/memory_stats_test.go | 28 +++ TreeDB/caching/pool_pressure_test.go | 21 +- TreeDB/caching/retained_arena_trim_test.go | 4 + TreeDB/internal/memtable/append_only.go | 208 +++++++++++++++++- .../memtable/append_only_pool_test.go | 182 +++++++++++++++ 6 files changed, 472 insertions(+), 6 deletions(-) diff --git a/TreeDB/caching/db.go b/TreeDB/caching/db.go index 904d83abf9..a5b84308cf 100644 --- a/TreeDB/caching/db.go +++ b/TreeDB/caching/db.go @@ -401,6 +401,7 @@ const ( postFlushAppendOnlyMemLeaseKeep = 24 postCheckpointAppendOnlyMemLeaseKeep = 8 appendOnlyEntryPoolHighPressureDropBytes = uint64(32 << 20) + appendOnlyKeyArenaPoolHighPressureDropBytes = uint64(32 << 20) appendOnlyValueArenaPoolHighPressureDropBytes = uint64(32 << 20) ) @@ -760,6 +761,7 @@ func maybeTrimEntrySliceLeasesUnderPressure(level poolPressureLevel, sampledAt t entrySlicePoolTrimDropBytesTotal.Add(uint64(droppedBytes)) } maybeDropAppendOnlyEntryPoolsUnderPressure(level) + maybeDropAppendOnlyKeyArenaPoolsUnderPressure(level) maybeDropAppendOnlyValueArenaPoolsUnderPressure(level) } @@ -778,6 +780,21 @@ func maybeDropAppendOnlyEntryPoolsUnderPressure(level poolPressureLevel) { memtable.DropAppendOnlyEntryPools() } +func maybeDropAppendOnlyKeyArenaPoolsUnderPressure(level poolPressureLevel) { + if level == poolPressureNormal { + return + } + stats := memtable.AppendOnlyKeyArenaPoolStatsSnapshot() + retainedBytes := stats.RetainedBytesEstimate + if retainedBytes == 0 { + return + } + if level != poolPressureCritical && retainedBytes < appendOnlyKeyArenaPoolHighPressureDropBytes { + return + } + memtable.DropAppendOnlyKeyArenaPools() +} + func maybeDropAppendOnlyValueArenaPoolsUnderPressure(level poolPressureLevel) { if level == poolPressureNormal { return @@ -10719,6 +10736,7 @@ func (db *DB) dropColdAppendOnlyPools() { db.appendOnlyMemPool.Store(&sync.Pool{}) db.appendOnlyMemPoolDropTotal.Add(1) memtable.DropAppendOnlyEntryPools() + memtable.DropAppendOnlyKeyArenaPools() memtable.DropAppendOnlyValueArenaPools() } @@ -30089,6 +30107,7 @@ func (db *DB) Stats() map[string]string { appendOnlyMemNewAllocQueueBytes := db.appendOnlyMemNewAllocQueueBytes.Load() appendOnlyMemPoolDropTotal := db.appendOnlyMemPoolDropTotal.Load() appendOnlyEntryPoolStats := memtable.AppendOnlyEntryPoolStatsSnapshot() + appendOnlyKeyArenaPoolStats := memtable.AppendOnlyKeyArenaPoolStatsSnapshot() appendOnlyValueArenaPoolStats := memtable.AppendOnlyValueArenaPoolStatsSnapshot() appendOnlyEntryReserveStats := memtable.AppendOnlyEntryReserveStatsSnapshot() appendOnlyMemLeaseCount, appendOnlyMemLeaseEntryCapacity, appendOnlyMemLeaseEntryBackingBytes, appendOnlyMemLeaseValueArena := db.appendOnlyMemLeaseStats() @@ -30108,6 +30127,14 @@ func (db *DB) Stats() map[string]string { stats["treedb.cache.append_only.entry_pool_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyEntryPoolStats.DropBytesTotal) stats["treedb.cache.append_only.entry_pool_admission_drops_total"] = fmt.Sprintf("%d", appendOnlyEntryPoolStats.AdmissionDropsTotal) stats["treedb.cache.append_only.entry_pool_admission_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyEntryPoolStats.AdmissionDropBytesTotal) + stats["treedb.cache.append_only.key_arena_pool_retained_bytes_estimate"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.RetainedBytesEstimate) + stats["treedb.cache.append_only.key_arena_pool_retained_bytes_max_estimate"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.RetainedBytesMaxEstimate) + stats["treedb.cache.append_only.key_arena_pool_gets_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.GetsTotal) + stats["treedb.cache.append_only.key_arena_pool_puts_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.PutsTotal) + stats["treedb.cache.append_only.key_arena_pool_drops_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.DropsTotal) + stats["treedb.cache.append_only.key_arena_pool_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.DropBytesTotal) + stats["treedb.cache.append_only.key_arena_pool_admission_drops_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.AdmissionDropsTotal) + stats["treedb.cache.append_only.key_arena_pool_admission_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.AdmissionDropBytesTotal) stats["treedb.cache.append_only.value_arena_pool_retained_bytes_estimate"] = fmt.Sprintf("%d", appendOnlyValueArenaPoolStats.RetainedBytesEstimate) stats["treedb.cache.append_only.value_arena_pool_retained_bytes_max_estimate"] = fmt.Sprintf("%d", appendOnlyValueArenaPoolStats.RetainedBytesMaxEstimate) stats["treedb.cache.append_only.value_arena_pool_gets_total"] = fmt.Sprintf("%d", appendOnlyValueArenaPoolStats.GetsTotal) @@ -30158,6 +30185,14 @@ func (db *DB) Stats() map[string]string { stats["treedb.process.append_only.entry_pool_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyEntryPoolStats.DropBytesTotal) stats["treedb.process.append_only.entry_pool_admission_drops_total"] = fmt.Sprintf("%d", appendOnlyEntryPoolStats.AdmissionDropsTotal) stats["treedb.process.append_only.entry_pool_admission_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyEntryPoolStats.AdmissionDropBytesTotal) + stats["treedb.process.append_only.key_arena_pool_retained_bytes_estimate"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.RetainedBytesEstimate) + stats["treedb.process.append_only.key_arena_pool_retained_bytes_max_estimate"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.RetainedBytesMaxEstimate) + stats["treedb.process.append_only.key_arena_pool_gets_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.GetsTotal) + stats["treedb.process.append_only.key_arena_pool_puts_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.PutsTotal) + stats["treedb.process.append_only.key_arena_pool_drops_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.DropsTotal) + stats["treedb.process.append_only.key_arena_pool_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.DropBytesTotal) + stats["treedb.process.append_only.key_arena_pool_admission_drops_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.AdmissionDropsTotal) + stats["treedb.process.append_only.key_arena_pool_admission_drop_bytes_total"] = fmt.Sprintf("%d", appendOnlyKeyArenaPoolStats.AdmissionDropBytesTotal) stats["treedb.process.append_only.value_arena_pool_retained_bytes_estimate"] = fmt.Sprintf("%d", appendOnlyValueArenaPoolStats.RetainedBytesEstimate) stats["treedb.process.append_only.value_arena_pool_retained_bytes_max_estimate"] = fmt.Sprintf("%d", appendOnlyValueArenaPoolStats.RetainedBytesMaxEstimate) stats["treedb.process.append_only.value_arena_pool_gets_total"] = fmt.Sprintf("%d", appendOnlyValueArenaPoolStats.GetsTotal) diff --git a/TreeDB/caching/memory_stats_test.go b/TreeDB/caching/memory_stats_test.go index a58b02b6f7..607dcedd23 100644 --- a/TreeDB/caching/memory_stats_test.go +++ b/TreeDB/caching/memory_stats_test.go @@ -164,6 +164,14 @@ func TestProcessMemoryStatsIncludeRuntimeBreakdown(t *testing.T) { "treedb.cache.append_only.entry_pool_drop_bytes_total", "treedb.cache.append_only.entry_pool_admission_drops_total", "treedb.cache.append_only.entry_pool_admission_drop_bytes_total", + "treedb.cache.append_only.key_arena_pool_retained_bytes_estimate", + "treedb.cache.append_only.key_arena_pool_retained_bytes_max_estimate", + "treedb.cache.append_only.key_arena_pool_gets_total", + "treedb.cache.append_only.key_arena_pool_puts_total", + "treedb.cache.append_only.key_arena_pool_drops_total", + "treedb.cache.append_only.key_arena_pool_drop_bytes_total", + "treedb.cache.append_only.key_arena_pool_admission_drops_total", + "treedb.cache.append_only.key_arena_pool_admission_drop_bytes_total", "treedb.cache.append_only.value_arena_pool_retained_bytes_estimate", "treedb.cache.append_only.value_arena_pool_retained_bytes_max_estimate", "treedb.cache.append_only.value_arena_pool_gets_total", @@ -195,6 +203,14 @@ func TestProcessMemoryStatsIncludeRuntimeBreakdown(t *testing.T) { "treedb.process.append_only.entry_pool_drop_bytes_total", "treedb.process.append_only.entry_pool_admission_drops_total", "treedb.process.append_only.entry_pool_admission_drop_bytes_total", + "treedb.process.append_only.key_arena_pool_retained_bytes_estimate", + "treedb.process.append_only.key_arena_pool_retained_bytes_max_estimate", + "treedb.process.append_only.key_arena_pool_gets_total", + "treedb.process.append_only.key_arena_pool_puts_total", + "treedb.process.append_only.key_arena_pool_drops_total", + "treedb.process.append_only.key_arena_pool_drop_bytes_total", + "treedb.process.append_only.key_arena_pool_admission_drops_total", + "treedb.process.append_only.key_arena_pool_admission_drop_bytes_total", "treedb.process.append_only.value_arena_pool_retained_bytes_estimate", "treedb.process.append_only.value_arena_pool_retained_bytes_max_estimate", "treedb.process.append_only.value_arena_pool_gets_total", @@ -275,6 +291,18 @@ func TestProcessMemoryStatsIncludeRuntimeBreakdown(t *testing.T) { if got := mustStatInt64(t, stats, "treedb.process.append_only.entry_pool_admission_drop_bytes_total"); got != mustStatInt64(t, stats, "treedb.cache.append_only.entry_pool_admission_drop_bytes_total") { t.Fatalf("append_only entry pool admission drop bytes mismatch process=%d cache=%d", got, mustStatInt64(t, stats, "treedb.cache.append_only.entry_pool_admission_drop_bytes_total")) } + if got := mustStatInt64(t, stats, "treedb.process.append_only.key_arena_pool_retained_bytes_estimate"); got != mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_retained_bytes_estimate") { + t.Fatalf("append_only key arena pool retained mismatch process=%d cache=%d", got, mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_retained_bytes_estimate")) + } + if got := mustStatInt64(t, stats, "treedb.process.append_only.key_arena_pool_drop_bytes_total"); got != mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_drop_bytes_total") { + t.Fatalf("append_only key arena pool drop bytes mismatch process=%d cache=%d", got, mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_drop_bytes_total")) + } + if got := mustStatInt64(t, stats, "treedb.process.append_only.key_arena_pool_admission_drops_total"); got != mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_admission_drops_total") { + t.Fatalf("append_only key arena pool admission drops mismatch process=%d cache=%d", got, mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_admission_drops_total")) + } + if got := mustStatInt64(t, stats, "treedb.process.append_only.key_arena_pool_admission_drop_bytes_total"); got != mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_admission_drop_bytes_total") { + t.Fatalf("append_only key arena pool admission drop bytes mismatch process=%d cache=%d", got, mustStatInt64(t, stats, "treedb.cache.append_only.key_arena_pool_admission_drop_bytes_total")) + } if got := mustStatInt64(t, stats, "treedb.process.append_only.value_arena_pool_retained_bytes_estimate"); got != mustStatInt64(t, stats, "treedb.cache.append_only.value_arena_pool_retained_bytes_estimate") { t.Fatalf("append_only value arena pool retained mismatch process=%d cache=%d", got, mustStatInt64(t, stats, "treedb.cache.append_only.value_arena_pool_retained_bytes_estimate")) } diff --git a/TreeDB/caching/pool_pressure_test.go b/TreeDB/caching/pool_pressure_test.go index 0245277eb9..7cd9e04685 100644 --- a/TreeDB/caching/pool_pressure_test.go +++ b/TreeDB/caching/pool_pressure_test.go @@ -157,9 +157,11 @@ func TestPoolPressureDropsAppendOnlyEntryPools(t *testing.T) { resetPoolPressureStateForTest() memtable.DropAppendOnlyEntryPools() + memtable.DropAppendOnlyKeyArenaPools() memtable.DropAppendOnlyValueArenaPools() t.Cleanup(func() { memtable.DropAppendOnlyEntryPools() + memtable.DropAppendOnlyKeyArenaPools() memtable.DropAppendOnlyValueArenaPools() resetPoolPressureStateForTest() }) @@ -175,8 +177,15 @@ func TestPoolPressureDropsAppendOnlyEntryPools(t *testing.T) { if before.RetainedBytesEstimate == 0 { t.Fatal("test setup did not retain append-only entry pool bytes") } + keyArena := memtable.NewAppendOnlyWithCapacity(0) + keyArena.Set([]byte("key-arena-retained-chunk"), nil) + keyArena.Reset() + beforeKeyArena := memtable.AppendOnlyKeyArenaPoolStatsSnapshot() + if beforeKeyArena.RetainedBytesEstimate == 0 { + t.Fatal("test setup did not retain append-only key arena pool bytes") + } valueArena := memtable.NewAppendOnlyWithCapacity(0) - valueArena.Set([]byte("value-arena-key"), make([]byte, 4096)) + valueArena.Set([]byte("valuekey"), make([]byte, 4096)) valueArena.ResetDropEntries() beforeValueArena := memtable.AppendOnlyValueArenaPoolStatsSnapshot() if beforeValueArena.RetainedBytesEstimate == 0 { @@ -195,6 +204,16 @@ func TestPoolPressureDropsAppendOnlyEntryPools(t *testing.T) { if got := after.DropBytesTotal; got < before.DropBytesTotal+before.RetainedBytesEstimate { t.Fatalf("append-only entry pool drop bytes=%d want at least %d", got, before.DropBytesTotal+before.RetainedBytesEstimate) } + afterKeyArena := memtable.AppendOnlyKeyArenaPoolStatsSnapshot() + if got := afterKeyArena.RetainedBytesEstimate; got != 0 { + t.Fatalf("append-only key arena pool retained bytes=%d want 0", got) + } + if got := afterKeyArena.DropsTotal; got != beforeKeyArena.DropsTotal+1 { + t.Fatalf("append-only key arena pool drops=%d want %d", got, beforeKeyArena.DropsTotal+1) + } + if got := afterKeyArena.DropBytesTotal; got < beforeKeyArena.DropBytesTotal+beforeKeyArena.RetainedBytesEstimate { + t.Fatalf("append-only key arena pool drop bytes=%d want at least %d", got, beforeKeyArena.DropBytesTotal+beforeKeyArena.RetainedBytesEstimate) + } afterValueArena := memtable.AppendOnlyValueArenaPoolStatsSnapshot() if got := afterValueArena.RetainedBytesEstimate; got != 0 { t.Fatalf("append-only value arena pool retained bytes=%d want 0", got) diff --git a/TreeDB/caching/retained_arena_trim_test.go b/TreeDB/caching/retained_arena_trim_test.go index ae4e2c62a6..2ab98ab778 100644 --- a/TreeDB/caching/retained_arena_trim_test.go +++ b/TreeDB/caching/retained_arena_trim_test.go @@ -162,6 +162,7 @@ func TestStoreMemtableMode_DropsAppendOnlyPoolsOnColdTransition(t *testing.T) { t.Fatal("expected append-only memtable pool") } beforeEntryPoolDrops := memtable.AppendOnlyEntryPoolDropTotal() + beforeKeyArenaPoolDrops := memtable.AppendOnlyKeyArenaPoolDropTotal() beforeValueArenaPoolDrops := memtable.AppendOnlyValueArenaPoolDropTotal() db.storeMemtableMode(memtable.ModeBTree) @@ -179,6 +180,9 @@ func TestStoreMemtableMode_DropsAppendOnlyPoolsOnColdTransition(t *testing.T) { if got := memtable.AppendOnlyEntryPoolDropTotal(); got != beforeEntryPoolDrops+1 { t.Fatalf("append-only entry pool drops=%d want %d", got, beforeEntryPoolDrops+1) } + if got := memtable.AppendOnlyKeyArenaPoolDropTotal(); got != beforeKeyArenaPoolDrops+1 { + t.Fatalf("append-only key arena pool drops=%d want %d", got, beforeKeyArenaPoolDrops+1) + } if got := memtable.AppendOnlyValueArenaPoolDropTotal(); got != beforeValueArenaPoolDrops+1 { t.Fatalf("append-only value arena pool drops=%d want %d", got, beforeValueArenaPoolDrops+1) } diff --git a/TreeDB/internal/memtable/append_only.go b/TreeDB/internal/memtable/append_only.go index 8284908eae..207ea04d8e 100644 --- a/TreeDB/internal/memtable/append_only.go +++ b/TreeDB/internal/memtable/append_only.go @@ -34,7 +34,11 @@ const ( appendOnlyIteratorPoolMaxCap = 1 << 20 appendOnlyIteratorPtrPoolMaxCap = 1 << 20 appendOnlyReusableKeyMaxCap = 1 << 10 - appendOnlyKeyArenaDefaultChunk = 2 << 10 + appendOnlyKeyArenaMinShift = 11 + appendOnlyKeyArenaMaxShift = 20 + appendOnlyKeyArenaClassCount = appendOnlyKeyArenaMaxShift - appendOnlyKeyArenaMinShift + 1 + appendOnlyKeyArenaDefaultChunk = 1 << appendOnlyKeyArenaMinShift + appendOnlyKeyArenaPoolMaxCap = 1 << appendOnlyKeyArenaMaxShift appendOnlyValueArenaMinShift = 12 appendOnlyValueArenaMaxShift = 20 appendOnlyValueArenaClassCount = appendOnlyValueArenaMaxShift - appendOnlyValueArenaMinShift + 1 @@ -58,6 +62,7 @@ const ( ) var appendOnlyEntryPoolRetainBudgetBytes = uint64(256 << 20) +var appendOnlyKeyArenaPoolRetainBudgetBytes = uint64(64 << 20) var appendOnlyValueArenaPoolRetainBudgetBytes = uint64(64 << 20) // Serializes package-pool replacement and retained entry backing accounting. @@ -81,6 +86,16 @@ var appendOnlyEntryReserveGrowCallsTotal atomic.Uint64 var appendOnlyEntryReserveGrowBytesTotal atomic.Uint64 var appendOnlyEntryReserveSkippedGrowthAllocsTotal atomic.Uint64 var appendOnlyEntryReserveSkippedGrowthBytesTotal atomic.Uint64 +var appendOnlyKeyArenaPoolMu sync.Mutex +var appendOnlyKeyArenaPoolBins [appendOnlyKeyArenaClassCount][][]byte +var appendOnlyKeyArenaPoolRetainedBytes atomic.Uint64 +var appendOnlyKeyArenaPoolRetainedBytesMax atomic.Uint64 +var appendOnlyKeyArenaPoolGetTotal atomic.Uint64 +var appendOnlyKeyArenaPoolPutTotal atomic.Uint64 +var appendOnlyKeyArenaPoolDropTotal atomic.Uint64 +var appendOnlyKeyArenaPoolDropBytesTotal atomic.Uint64 +var appendOnlyKeyArenaPoolAdmissionDropTotal atomic.Uint64 +var appendOnlyKeyArenaPoolAdmissionDropBytesTotal atomic.Uint64 var appendOnlyValueArenaPoolPtrs [appendOnlyValueArenaClassCount]atomic.Pointer[sync.Pool] var appendOnlyValueArenaPoolRetainedBytes atomic.Uint64 var appendOnlyValueArenaPoolRetainedBytesMax atomic.Uint64 @@ -189,6 +204,17 @@ type AppendOnlyValueArenaStats struct { RetainedBytes int64 } +type AppendOnlyKeyArenaPoolStats struct { + RetainedBytesEstimate uint64 + RetainedBytesMaxEstimate uint64 + GetsTotal uint64 + PutsTotal uint64 + DropsTotal uint64 + DropBytesTotal uint64 + AdmissionDropsTotal uint64 + AdmissionDropBytesTotal uint64 +} + type AppendOnlyValueArenaPoolStats struct { RetainedBytesEstimate uint64 RetainedBytesMaxEstimate uint64 @@ -407,6 +433,160 @@ func appendOnlyValueArenaPoolForClass(class int) *sync.Pool { return appendOnlyValueArenaPoolPtrs[class].Load() } +func appendOnlyKeyArenaPoolBytes(capacity int) uint64 { + if capacity <= 0 { + return 0 + } + return uint64(capacity) +} + +func appendOnlyKeyArenaClassForLen(length int) (idx int, classCap int, ok bool) { + if length <= 0 || length > appendOnlyKeyArenaPoolMaxCap { + return 0, 0, false + } + classCap = 1 << uint(bits.Len(uint(length-1))) + if classCap < appendOnlyKeyArenaDefaultChunk { + classCap = appendOnlyKeyArenaDefaultChunk + } + if classCap > appendOnlyKeyArenaPoolMaxCap { + return 0, 0, false + } + shift := bits.Len(uint(classCap)) - 1 + idx = shift - appendOnlyKeyArenaMinShift + if idx < 0 || idx >= appendOnlyKeyArenaClassCount { + return 0, 0, false + } + return idx, classCap, true +} + +func appendOnlyKeyArenaClassForCap(capacity int) (idx int, ok bool) { + if capacity < appendOnlyKeyArenaDefaultChunk || capacity > appendOnlyKeyArenaPoolMaxCap { + return 0, false + } + if capacity&(capacity-1) != 0 { + return 0, false + } + shift := bits.TrailingZeros(uint(capacity)) + idx = shift - appendOnlyKeyArenaMinShift + if idx < 0 || idx >= appendOnlyKeyArenaClassCount { + return 0, false + } + return idx, true +} + +func updateAppendOnlyKeyArenaPoolRetainedBytesMax(value uint64) { + for { + prev := appendOnlyKeyArenaPoolRetainedBytesMax.Load() + if value <= prev || appendOnlyKeyArenaPoolRetainedBytesMax.CompareAndSwap(prev, value) { + return + } + } +} + +func getAppendOnlyKeyArenaChunkFromPool(length int) ([]byte, bool) { + idx, classCap, ok := appendOnlyKeyArenaClassForLen(length) + if !ok { + return nil, false + } + appendOnlyKeyArenaPoolMu.Lock() + bin := appendOnlyKeyArenaPoolBins[idx] + if len(bin) == 0 { + appendOnlyKeyArenaPoolMu.Unlock() + return nil, false + } + last := len(bin) - 1 + chunk := bin[last] + bin[last] = nil + appendOnlyKeyArenaPoolBins[idx] = bin[:last] + bytes := appendOnlyKeyArenaPoolBytes(cap(chunk)) + if bytes > 0 { + current := appendOnlyKeyArenaPoolRetainedBytes.Load() + if current > bytes { + appendOnlyKeyArenaPoolRetainedBytes.Store(current - bytes) + } else { + appendOnlyKeyArenaPoolRetainedBytes.Store(0) + } + } + appendOnlyKeyArenaPoolGetTotal.Add(1) + appendOnlyKeyArenaPoolMu.Unlock() + if cap(chunk) < classCap { + return nil, false + } + return chunk[:0], true +} + +func recordAppendOnlyKeyArenaPoolAdmissionDrop(bytes uint64) { + appendOnlyKeyArenaPoolAdmissionDropTotal.Add(1) + if bytes > 0 { + appendOnlyKeyArenaPoolAdmissionDropBytesTotal.Add(bytes) + } +} + +func putAppendOnlyKeyArenaChunk(chunk []byte) { + if chunk == nil { + return + } + bytes := appendOnlyKeyArenaPoolBytes(cap(chunk)) + idx, ok := appendOnlyKeyArenaClassForCap(cap(chunk)) + if !ok { + if bytes > 0 { + recordAppendOnlyKeyArenaPoolAdmissionDrop(bytes) + } + return + } + appendOnlyKeyArenaPoolMu.Lock() + budget := appendOnlyKeyArenaPoolRetainBudgetBytes + retained := appendOnlyKeyArenaPoolRetainedBytes.Load() + if budget == 0 || retained >= budget || bytes > budget-retained { + appendOnlyKeyArenaPoolMu.Unlock() + recordAppendOnlyKeyArenaPoolAdmissionDrop(bytes) + return + } + next := retained + bytes + appendOnlyKeyArenaPoolRetainedBytes.Store(next) + updateAppendOnlyKeyArenaPoolRetainedBytesMax(next) + appendOnlyKeyArenaPoolPutTotal.Add(1) + appendOnlyKeyArenaPoolBins[idx] = append(appendOnlyKeyArenaPoolBins[idx], chunk[:0]) + appendOnlyKeyArenaPoolMu.Unlock() +} + +func dropAppendOnlyKeyArenaPools() { + appendOnlyKeyArenaPoolMu.Lock() + for i := range appendOnlyKeyArenaPoolBins { + appendOnlyKeyArenaPoolBins[i] = nil + } + dropped := appendOnlyKeyArenaPoolRetainedBytes.Swap(0) + appendOnlyKeyArenaPoolMu.Unlock() + if dropped > 0 { + appendOnlyKeyArenaPoolDropBytesTotal.Add(dropped) + } + appendOnlyKeyArenaPoolDropTotal.Add(1) +} + +// DropAppendOnlyKeyArenaPools abandons package-level append-only key-arena +// chunks. It is intended for cold transitions away from append-only mutable +// memtables and for tests that need deterministic pool state. +func DropAppendOnlyKeyArenaPools() { + dropAppendOnlyKeyArenaPools() +} + +func AppendOnlyKeyArenaPoolDropTotal() uint64 { + return appendOnlyKeyArenaPoolDropTotal.Load() +} + +func AppendOnlyKeyArenaPoolStatsSnapshot() AppendOnlyKeyArenaPoolStats { + return AppendOnlyKeyArenaPoolStats{ + RetainedBytesEstimate: appendOnlyKeyArenaPoolRetainedBytes.Load(), + RetainedBytesMaxEstimate: appendOnlyKeyArenaPoolRetainedBytesMax.Load(), + GetsTotal: appendOnlyKeyArenaPoolGetTotal.Load(), + PutsTotal: appendOnlyKeyArenaPoolPutTotal.Load(), + DropsTotal: appendOnlyKeyArenaPoolDropTotal.Load(), + DropBytesTotal: appendOnlyKeyArenaPoolDropBytesTotal.Load(), + AdmissionDropsTotal: appendOnlyKeyArenaPoolAdmissionDropTotal.Load(), + AdmissionDropBytesTotal: appendOnlyKeyArenaPoolAdmissionDropBytesTotal.Load(), + } +} + func appendOnlyValueArenaPoolBytes(capacity int) uint64 { if capacity <= 0 { return 0 @@ -947,9 +1127,13 @@ func (a *appendOnlyKeyArena) alloc(length int) []byte { if length > chunkCap { chunkCap = 1 << uint(bits.Len(uint(length-1))) } - chunk := make([]byte, chunkCap) - a.chunks = append(a.chunks, chunk) - a.cur = chunk + chunk, ok := getAppendOnlyKeyArenaChunkFromPool(chunkCap) + if !ok { + chunk = make([]byte, chunkCap) + } + full := chunk[:cap(chunk)] + a.chunks = append(a.chunks, full) + a.cur = full a.curPos = 0 } out := a.cur[a.curPos : a.curPos+length : a.curPos+length] @@ -959,6 +1143,7 @@ func (a *appendOnlyKeyArena) alloc(length int) []byte { func (a *appendOnlyKeyArena) reset() { for i := range a.chunks { + putAppendOnlyKeyArenaChunk(a.chunks[i]) a.chunks[i] = nil } a.chunks = a.chunks[:0] @@ -966,6 +1151,15 @@ func (a *appendOnlyKeyArena) reset() { a.curPos = 0 } +func (a *appendOnlyKeyArena) drop() { + for i := range a.chunks { + a.chunks[i] = nil + } + a.chunks = nil + a.cur = nil + a.curPos = 0 +} + func appendOnlyNextCapacity(current int) int { if current < appendOnlyMinInitialEntries { return appendOnlyMinInitialEntries @@ -2220,7 +2414,11 @@ func (m *AppendOnly) release(poolEntries bool) { m.runMergeBuf = nil m.snapshot = nil m.indexBuf = nil - m.keyArena.reset() + if poolEntries { + m.keyArena.reset() + } else { + m.keyArena.drop() + } m.valueArena.reset() m.valueArena.dropRetained() m.count = 0 diff --git a/TreeDB/internal/memtable/append_only_pool_test.go b/TreeDB/internal/memtable/append_only_pool_test.go index 6cb7e34238..c3e74b383c 100644 --- a/TreeDB/internal/memtable/append_only_pool_test.go +++ b/TreeDB/internal/memtable/append_only_pool_test.go @@ -663,6 +663,188 @@ func TestAppendOnlyLargeNonInlineKeyUsesReusableSlice(t *testing.T) { } } +func TestAppendOnlyKeyArenaClassForLen(t *testing.T) { + tests := []struct { + name string + length int + wantCap int + wantPool bool + }{ + {name: "zero", length: 0, wantPool: false}, + {name: "single-byte", length: 1, wantCap: appendOnlyKeyArenaDefaultChunk, wantPool: true}, + {name: "default", length: appendOnlyKeyArenaDefaultChunk, wantCap: appendOnlyKeyArenaDefaultChunk, wantPool: true}, + {name: "round-up", length: appendOnlyKeyArenaDefaultChunk + 1, wantCap: appendOnlyKeyArenaDefaultChunk * 2, wantPool: true}, + {name: "retention-maximum", length: appendOnlyKeyArenaPoolMaxCap, wantCap: appendOnlyKeyArenaPoolMaxCap, wantPool: true}, + {name: "above-retention-maximum", length: appendOnlyKeyArenaPoolMaxCap + 1, wantPool: false}, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + _, gotCap, gotPool := appendOnlyKeyArenaClassForLen(tc.length) + if gotPool != tc.wantPool { + t.Fatalf("pooled=%v want %v", gotPool, tc.wantPool) + } + if gotPool && gotCap != tc.wantCap { + t.Fatalf("cap=%d want %d", gotCap, tc.wantCap) + } + }) + } +} + +func TestAppendOnlyKeyArenaPoolStatsTrackRetainedChunks(t *testing.T) { + DropAppendOnlyKeyArenaPools() + t.Cleanup(DropAppendOnlyKeyArenaPools) + + before := AppendOnlyKeyArenaPoolStatsSnapshot() + chunk := make([]byte, appendOnlyKeyArenaDefaultChunk) + wantBytes := appendOnlyKeyArenaPoolBytes(cap(chunk)) + + putAppendOnlyKeyArenaChunk(chunk) + afterPut := AppendOnlyKeyArenaPoolStatsSnapshot() + if got := afterPut.RetainedBytesEstimate; got != wantBytes { + t.Fatalf("retained bytes after put=%d want %d", got, wantBytes) + } + if got := afterPut.PutsTotal; got != before.PutsTotal+1 { + t.Fatalf("puts total=%d want %d", got, before.PutsTotal+1) + } + + got, ok := getAppendOnlyKeyArenaChunkFromPool(1) + if !ok { + t.Fatal("expected pooled key arena chunk") + } + afterGet := AppendOnlyKeyArenaPoolStatsSnapshot() + if cap(got) != appendOnlyKeyArenaDefaultChunk { + t.Fatalf("pooled cap=%d want %d", cap(got), appendOnlyKeyArenaDefaultChunk) + } + if gotRetained := afterGet.RetainedBytesEstimate; gotRetained != 0 { + t.Fatalf("retained bytes after pooled get=%d want 0", gotRetained) + } + if gotGets := afterGet.GetsTotal; gotGets != before.GetsTotal+1 { + t.Fatalf("gets total=%d want %d", gotGets, before.GetsTotal+1) + } + + putAppendOnlyKeyArenaChunk(got) + beforeDrop := AppendOnlyKeyArenaPoolStatsSnapshot() + DropAppendOnlyKeyArenaPools() + afterDrop := AppendOnlyKeyArenaPoolStatsSnapshot() + if gotRetained := afterDrop.RetainedBytesEstimate; gotRetained != 0 { + t.Fatalf("retained bytes after drop=%d want 0", gotRetained) + } + if gotDrops := afterDrop.DropsTotal; gotDrops != beforeDrop.DropsTotal+1 { + t.Fatalf("drops total=%d want %d", gotDrops, beforeDrop.DropsTotal+1) + } + if gotDropBytes := afterDrop.DropBytesTotal; gotDropBytes < beforeDrop.DropBytesTotal+beforeDrop.RetainedBytesEstimate { + t.Fatalf("drop bytes total=%d want at least %d", gotDropBytes, beforeDrop.DropBytesTotal+beforeDrop.RetainedBytesEstimate) + } +} + +func TestAppendOnlyKeyArenaPoolRetainsBackingAcrossGC(t *testing.T) { + DropAppendOnlyKeyArenaPools() + t.Cleanup(DropAppendOnlyKeyArenaPools) + + chunk := make([]byte, appendOnlyKeyArenaDefaultChunk) + wantBytes := appendOnlyKeyArenaPoolBytes(cap(chunk)) + putAppendOnlyKeyArenaChunk(chunk) + + runtime.GC() + + afterGC := AppendOnlyKeyArenaPoolStatsSnapshot() + if got := afterGC.RetainedBytesEstimate; got != wantBytes { + t.Fatalf("retained bytes after GC=%d want %d", got, wantBytes) + } + + reused, ok := getAppendOnlyKeyArenaChunkFromPool(1) + if !ok { + t.Fatal("expected retained key arena chunk after GC") + } + if cap(reused) != cap(chunk) { + t.Fatalf("reused cap after GC=%d want %d", cap(reused), cap(chunk)) + } + if got := AppendOnlyKeyArenaPoolStatsSnapshot().RetainedBytesEstimate; got != 0 { + t.Fatalf("retained bytes after get=%d want 0", got) + } + putAppendOnlyKeyArenaChunk(reused) +} + +func TestAppendOnlyKeyArenaReusesPooledChunkAfterHardReset(t *testing.T) { + DropAppendOnlyKeyArenaPools() + t.Cleanup(DropAppendOnlyKeyArenaPools) + + m := NewAppendOnlyWithCapacity(0) + key := bytes.Repeat([]byte("k"), 32) + m.Set(key, nil) + if m.count != 1 || len(m.entries[0].key) == 0 { + t.Fatalf("test setup did not store a key-arena key") + } + firstPtr := uintptr(unsafe.Pointer(&m.entries[0].key[0])) + + m.ResetWithCapacityHard(0, 0) + if got := AppendOnlyKeyArenaPoolStatsSnapshot().RetainedBytesEstimate; got == 0 { + t.Fatalf("retained key arena pool bytes=0 after hard reset") + } + beforeGet := AppendOnlyKeyArenaPoolStatsSnapshot().GetsTotal + + key[0] = 'z' + m.Set(key, nil) + if m.count != 1 || len(m.entries[0].key) == 0 { + t.Fatalf("post-reset write did not store a key-arena key") + } + reusedPtr := uintptr(unsafe.Pointer(&m.entries[0].key[0])) + if reusedPtr != firstPtr { + t.Fatalf("post-reset key arena pointer=%#x want reused %#x", reusedPtr, firstPtr) + } + if got := AppendOnlyKeyArenaPoolStatsSnapshot().GetsTotal; got != beforeGet+1 { + t.Fatalf("key arena pool gets=%d want %d", got, beforeGet+1) + } +} + +func TestAppendOnlyReleaseDropEntriesDropsKeyArenaWithoutPooling(t *testing.T) { + DropAppendOnlyKeyArenaPools() + t.Cleanup(DropAppendOnlyKeyArenaPools) + + m := NewAppendOnlyWithCapacity(0) + m.Set(bytes.Repeat([]byte("k"), 32), nil) + if len(m.keyArena.chunks) == 0 { + t.Fatal("test setup did not populate key arena") + } + + m.ReleaseDropEntries() + + if got := AppendOnlyKeyArenaPoolStatsSnapshot().RetainedBytesEstimate; got != 0 { + t.Fatalf("retained key arena pool bytes after cold release=%d want 0", got) + } + if len(m.keyArena.chunks) != 0 || m.keyArena.cur != nil || m.keyArena.curPos != 0 { + t.Fatalf("key arena not dropped after cold release: chunks=%d cur=%v pos=%d", len(m.keyArena.chunks), m.keyArena.cur != nil, m.keyArena.curPos) + } +} + +func TestPutAppendOnlyKeyArenaChunkAdmissionDropsOnRetainBudgetPressure(t *testing.T) { + DropAppendOnlyKeyArenaPools() + t.Cleanup(DropAppendOnlyKeyArenaPools) + + savedBudget := appendOnlyKeyArenaPoolRetainBudgetBytes + appendOnlyKeyArenaPoolRetainBudgetBytes = appendOnlyKeyArenaPoolBytes(appendOnlyKeyArenaDefaultChunk) + t.Cleanup(func() { + appendOnlyKeyArenaPoolRetainBudgetBytes = savedBudget + }) + + before := AppendOnlyKeyArenaPoolStatsSnapshot() + putAppendOnlyKeyArenaChunk(make([]byte, appendOnlyKeyArenaDefaultChunk)) + putAppendOnlyKeyArenaChunk(make([]byte, appendOnlyKeyArenaDefaultChunk)) + + after := AppendOnlyKeyArenaPoolStatsSnapshot() + if got, want := after.RetainedBytesEstimate, appendOnlyKeyArenaPoolBytes(appendOnlyKeyArenaDefaultChunk); got != want { + t.Fatalf("retained bytes=%d want %d", got, want) + } + if got := after.PutsTotal - before.PutsTotal; got != 1 { + t.Fatalf("puts delta=%d want 1", got) + } + if got := after.AdmissionDropsTotal - before.AdmissionDropsTotal; got != 1 { + t.Fatalf("admission drops delta=%d want 1", got) + } +} + func TestAppendOnlyResetClearsValueArenaState(t *testing.T) { m := NewAppendOnlyWithCapacity(0) key := []byte("k")