From 536acd4cf00c875a86cdcb802b7e7f3d22b1ca1b Mon Sep 17 00:00:00 2001
From: Viktor Liu <viktor@netbird.io>
Date: Wed, 22 Apr 2026 10:18:25 +0200
Subject: [PATCH] device: support live-tuning pool cap via SetMax /
 Device.SetPreallocatedBuffersPerPool

---
 device/pool_size.go | 16 ++++++++++++++--
 device/pools.go     | 32 +++++++++++++++++++++++---------
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/device/pool_size.go b/device/pool_size.go
index 3f5dca4..8910a4a 100644
--- a/device/pool_size.go
+++ b/device/pool_size.go
@@ -7,8 +7,20 @@ package device
 
 // SetPreallocatedBuffersPerPool sets the cap on the number of buffers held by
 // each per-Device pool. Zero disables the cap (upstream default on
-// non-mobile platforms). Must be called before NewDevice; changes take effect
-// only for Devices created after this call.
+// non-mobile platforms). Changes affect Devices created after this call.
+// To retune a live Device, use Device.SetPreallocatedBuffersPerPool.
 func SetPreallocatedBuffersPerPool(n uint32) {
 	PreallocatedBuffersPerPool = n
 }
+
+// SetPreallocatedBuffersPerPool updates the cap on this Device's pools in
+// place. Takes effect immediately; goroutines blocked in Get are unblocked if
+// the cap was raised. Has no effect if the Device was created with
+// PreallocatedBuffersPerPool == 0.
+func (device *Device) SetPreallocatedBuffersPerPool(n uint32) {
+	device.pool.messageBuffers.SetMax(n)
+	device.pool.inboundElements.SetMax(n)
+	device.pool.outboundElements.SetMax(n)
+	device.pool.inboundElementsContainer.SetMax(n)
+	device.pool.outboundElementsContainer.SetMax(n)
+}
diff --git a/device/pools.go b/device/pools.go
index 2c18f41..179f83d 100644
--- a/device/pools.go
+++ b/device/pools.go
@@ -10,23 +10,24 @@ import (
 )
 
 type WaitPool struct {
-	pool  sync.Pool
-	cond  sync.Cond
-	lock  sync.Mutex
-	count uint32 // Get calls not yet Put back
-	max   uint32
+	pool    sync.Pool
+	cond    sync.Cond
+	lock    sync.Mutex
+	count   uint32 // Get calls not yet Put back
+	max     uint32
+	tracked bool // true if max was non-zero at construction; enables SetMax
 }
 
 func NewWaitPool(max uint32, new func() any) *WaitPool {
-	p := &WaitPool{pool: sync.Pool{New: new}, max: max}
+	p := &WaitPool{pool: sync.Pool{New: new}, max: max, tracked: max != 0}
 	p.cond = sync.Cond{L: &p.lock}
 	return p
 }
 
 func (p *WaitPool) Get() any {
-	if p.max != 0 {
+	if p.tracked {
 		p.lock.Lock()
-		for p.count >= p.max {
+		for p.max != 0 && p.count >= p.max {
 			p.cond.Wait()
 		}
 		p.count++
@@ -37,7 +38,7 @@ func (p *WaitPool) Get() any {
 
 func (p *WaitPool) Put(x any) {
 	p.pool.Put(x)
-	if p.max == 0 {
+	if !p.tracked {
 		return
 	}
 	p.lock.Lock()
@@ -46,6 +47,19 @@ func (p *WaitPool) Put(x any) {
 	p.cond.Signal()
 }
 
+// SetMax updates the pool cap. Takes effect immediately; waiters are
+// broadcast so they re-check against the new value. Has no effect if the
+// pool was constructed with max == 0 (unbounded, fast-path Get/Put).
+func (p *WaitPool) SetMax(n uint32) {
+	if !p.tracked {
+		return
+	}
+	p.lock.Lock()
+	p.max = n
+	p.cond.Broadcast()
+	p.lock.Unlock()
+}
+
 func (device *Device) PopulatePools() {
 	device.pool.inboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any {
 		s := make([]*QueueInboundElement, 0, device.BatchSize())