From 3f1cb09e0d367eac4bcb36b96fb2075ececf8c00 Mon Sep 17 00:00:00 2001 From: Jakub Novak Date: Tue, 7 Jul 2026 15:16:06 +0000 Subject: [PATCH 1/3] fix(api): prevent uint64 underflow in node allocated metrics --- .../internal/orchestrator/nodemanager/mock.go | 7 +++ .../internal/orchestrator/nodemanager/node.go | 14 ++++-- .../orchestrator/nodemanager/node_test.go | 46 ++++++++++++++++++- 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/packages/api/internal/orchestrator/nodemanager/mock.go b/packages/api/internal/orchestrator/nodemanager/mock.go index a7483a6869..a3f581e182 100644 --- a/packages/api/internal/orchestrator/nodemanager/mock.go +++ b/packages/api/internal/orchestrator/nodemanager/mock.go @@ -127,6 +127,13 @@ func WithFeatureFlags(ff *featureflags.Client) TestOptions { } } +// WithAllocatedMemoryBytes sets the initial allocated memory metric for the test node +func WithAllocatedMemoryBytes(bytes uint64) TestOptions { + return func(node *TestNode) { + node.metrics.MemoryAllocatedBytes = bytes + } +} + // MockSandboxClientCustom allows custom error logic per call type MockSandboxClientCustom struct { orchestrator.SandboxServiceClient diff --git a/packages/api/internal/orchestrator/nodemanager/node.go b/packages/api/internal/orchestrator/nodemanager/node.go index 5b412d9abf..a377742fa3 100644 --- a/packages/api/internal/orchestrator/nodemanager/node.go +++ b/packages/api/internal/orchestrator/nodemanager/node.go @@ -215,7 +215,15 @@ func (n *Node) OptimisticRemove(ctx context.Context, res SandboxResources) { n.metricsMu.Lock() defer n.metricsMu.Unlock() - // Directly subtract from the current metrics view - n.metrics.CpuAllocated -= uint32(res.CPUs) - n.metrics.MemoryAllocatedBytes -= uint64(res.MiBMemory) * 1024 * 1024 + // Directly subtract from the current metrics view. + cpu := uint32(res.CPUs) + memory := uint64(res.MiBMemory) * 1024 * 1024 + + // Prevent underflow due to race condition (the sandbox was most likely already removed by the node sync) + if cpu <= n.metrics.CpuAllocated { + n.metrics.CpuAllocated -= cpu + } + if memory <= n.metrics.MemoryAllocatedBytes { + n.metrics.MemoryAllocatedBytes -= memory + } } diff --git a/packages/api/internal/orchestrator/nodemanager/node_test.go b/packages/api/internal/orchestrator/nodemanager/node_test.go index adc749dc7d..f8ae828ea1 100644 --- a/packages/api/internal/orchestrator/nodemanager/node_test.go +++ b/packages/api/internal/orchestrator/nodemanager/node_test.go @@ -85,7 +85,7 @@ func TestNode_OptimisticRemove_FlagEnabled(t *testing.T) { require.NoError(t, err) // 4. Initialize Node with the injected ffClient - some resources are already allocated at initialization - node := NewTestNode("test-node", api.NodeStatusReady, 4, 8192, WithFeatureFlags(ffClient)) + node := NewTestNode("test-node", api.NodeStatusReady, 4, 8192, WithFeatureFlags(ffClient), WithAllocatedMemoryBytes(8192*1024*1024)) initialMetrics := node.Metrics() // 5. Call the method @@ -101,6 +101,50 @@ func TestNode_OptimisticRemove_FlagEnabled(t *testing.T) { assert.Equal(t, initialMetrics.MemoryAllocatedBytes-uint64(res.MiBMemory)*1024*1024, newMetrics.MemoryAllocatedBytes) } +func TestNode_OptimisticRemove_SkipsWhenItWouldUnderflow(t *testing.T) { + t.Parallel() + + td := ldtestdata.DataSource() + td.Update(td.Flag(featureflags.OptimisticResourceAccountingFlag.Key()).VariationForAll(true)) + + ffClient, err := featureflags.NewClientWithDatasource(td) + require.NoError(t, err) + + // Node with less allocated than what will be removed: 1 CPU, 512 MiB + node := NewTestNode("test-node", api.NodeStatusReady, 1, 8192, WithFeatureFlags(ffClient), WithAllocatedMemoryBytes(512*1024*1024)) + initialMetrics := node.Metrics() + + res := SandboxResources{ + CPUs: 2, + MiBMemory: 1024, + } + node.OptimisticRemove(t.Context(), res) + + // Counters must never wrap to ~2^32/2^64; subtraction is skipped instead + newMetrics := node.Metrics() + assert.Equal(t, initialMetrics.CpuAllocated, newMetrics.CpuAllocated) + assert.Equal(t, initialMetrics.MemoryAllocatedBytes, newMetrics.MemoryAllocatedBytes) +} + +func TestNode_OptimisticRemove_FreshNodeDoesNotUnderflow(t *testing.T) { + t.Parallel() + + td := ldtestdata.DataSource() + td.Update(td.Flag(featureflags.OptimisticResourceAccountingFlag.Key()).VariationForAll(true)) + + ffClient, err := featureflags.NewClientWithDatasource(td) + require.NoError(t, err) + + // Fresh node: nothing allocated yet (e.g. poll overwrote counters after sandbox already left the orchestrator) + node := NewTestNode("test-node", api.NodeStatusReady, 0, 8192, WithFeatureFlags(ffClient)) + + node.OptimisticRemove(t.Context(), SandboxResources{CPUs: 2, MiBMemory: 1024}) + + newMetrics := node.Metrics() + assert.Equal(t, uint32(0), newMetrics.CpuAllocated) + assert.Equal(t, uint64(0), newMetrics.MemoryAllocatedBytes) +} + func TestNode_OptimisticRemove_FlagDisabled(t *testing.T) { t.Parallel() From 8e9f6d83360d6fc8f86d2fdbf78fd45c2a5d5e2b Mon Sep 17 00:00:00 2001 From: Jakub Novak Date: Tue, 7 Jul 2026 15:18:04 +0000 Subject: [PATCH 2/3] chore: simplify --- .../api/internal/orchestrator/nodemanager/node.go | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/packages/api/internal/orchestrator/nodemanager/node.go b/packages/api/internal/orchestrator/nodemanager/node.go index a377742fa3..3d7aa25d63 100644 --- a/packages/api/internal/orchestrator/nodemanager/node.go +++ b/packages/api/internal/orchestrator/nodemanager/node.go @@ -220,10 +220,12 @@ func (n *Node) OptimisticRemove(ctx context.Context, res SandboxResources) { memory := uint64(res.MiBMemory) * 1024 * 1024 // Prevent underflow due to race condition (the sandbox was most likely already removed by the node sync) - if cpu <= n.metrics.CpuAllocated { - n.metrics.CpuAllocated -= cpu - } - if memory <= n.metrics.MemoryAllocatedBytes { - n.metrics.MemoryAllocatedBytes -= memory + if cpu > n.metrics.CpuAllocated || memory > n.metrics.MemoryAllocatedBytes { + logger.L().Warn(ctx, "OptimisticRemove would cause underflow, skipping", logger.WithNodeID(n.ID), zap.Uint32("cpuAllocated", n.metrics.CpuAllocated), zap.Uint64("memoryAllocatedBytes", n.metrics.MemoryAllocatedBytes), zap.Uint32("cpuToRemove", cpu), zap.Uint64("memoryToRemove", memory)) + + return } + + n.metrics.CpuAllocated -= cpu + n.metrics.MemoryAllocatedBytes -= memory } From f1f7fe22725e6fc0ddb435e9372eb338726921d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Nov=C3=A1k?= Date: Tue, 7 Jul 2026 08:55:45 -0700 Subject: [PATCH 3/3] Apply suggestion from @arkamar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Petr Vaněk --- packages/api/internal/orchestrator/nodemanager/node.go | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/api/internal/orchestrator/nodemanager/node.go b/packages/api/internal/orchestrator/nodemanager/node.go index 3d7aa25d63..6f8ad11bdd 100644 --- a/packages/api/internal/orchestrator/nodemanager/node.go +++ b/packages/api/internal/orchestrator/nodemanager/node.go @@ -215,7 +215,6 @@ func (n *Node) OptimisticRemove(ctx context.Context, res SandboxResources) { n.metricsMu.Lock() defer n.metricsMu.Unlock() - // Directly subtract from the current metrics view. cpu := uint32(res.CPUs) memory := uint64(res.MiBMemory) * 1024 * 1024