From cd7ac13b86fde959a222441a6a206e9e3cb9e135 Mon Sep 17 00:00:00 2001
From: Simon Law <sfllaw@tailscale.com>
Date: Thu, 18 Jun 2026 18:02:42 -0700
Subject: [PATCH] device: convert runtime.SetFinalizer to AddCleanup (#71)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

In PR #66, we tried to address a memory leak by avoiding
runtime.SetFinalizer for autodrainingInboundQueue and
autodrainingOutboundQueue unless there was something to do.

However, when there is work to be done, these finalizers still leak
memory because they’re still holding on to a cyclical reference to q.
This applies to any platform that relies on a bounded device.WaitPool,
like Android and iOS which both declare PreallocatedBuffersPerPool.

This patch converts this logic to runtime.AddCleanup which is designed
to avoid this problem.

Updates tailscale/corp#42776

Signed-off-by: Simon Law <sfllaw@tailscale.com>
---
 device/channels.go | 12 ++++++------
 device/peer.go     |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/device/channels.go b/device/channels.go
index 9ac767f..9af6e3d 100644
--- a/device/channels.go
+++ b/device/channels.go
@@ -84,7 +84,7 @@ func newAutodrainingInboundQueue(device *Device) *autodrainingInboundQueue {
 		c: make(chan *QueueInboundElementsContainer, QueueInboundSize),
 	}
 	if device.needsInboundQueueFinalizer() {
-		runtime.SetFinalizer(q, device.flushInboundQueue)
+		runtime.AddCleanup(q, device.flushInboundQueue, q.c)
 	}
 	return q
 }
@@ -93,10 +93,10 @@ func (device *Device) needsInboundQueueFinalizer() bool {
 	return device.pool.messageBuffers.hasAccounting()
 }
 
-func (device *Device) flushInboundQueue(q *autodrainingInboundQueue) {
+func (device *Device) flushInboundQueue(c <-chan *QueueInboundElementsContainer) {
 	for {
 		select {
-		case elemsContainer := <-q.c:
+		case elemsContainer := <-c:
 			elemsContainer.filling.Wait()
 			for _, elem := range elemsContainer.elems {
 				device.PutMessageBuffer(elem.buffer)
@@ -123,7 +123,7 @@ func newAutodrainingOutboundQueue(device *Device) *autodrainingOutboundQueue {
 		c: make(chan *QueueOutboundElementsContainer, QueueOutboundSize),
 	}
 	if device.needsOutboundQueueFinalizer() {
-		runtime.SetFinalizer(q, device.flushOutboundQueue)
+		runtime.AddCleanup(q, device.flushOutboundQueue, q.c)
 	}
 	return q
 }
@@ -132,10 +132,10 @@ func (device *Device) needsOutboundQueueFinalizer() bool {
 	return device.pool.messageBuffers.hasAccounting()
 }
 
-func (device *Device) flushOutboundQueue(q *autodrainingOutboundQueue) {
+func (device *Device) flushOutboundQueue(c <-chan *QueueOutboundElementsContainer) {
 	for {
 		select {
-		case elemsContainer := <-q.c:
+		case elemsContainer := <-c:
 			elemsContainer.filling.Wait()
 			for _, elem := range elemsContainer.elems {
 				device.PutOutboundBuffer(elem.buffer)
diff --git a/device/peer.go b/device/peer.go
index 4f90144..c0ca59a 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -244,8 +244,8 @@ func (peer *Peer) Start() {
 
 	peer.timersStart()
 
-	device.flushInboundQueue(peer.queue.inbound)
-	device.flushOutboundQueue(peer.queue.outbound)
+	device.flushInboundQueue(peer.queue.inbound.c)
+	device.flushOutboundQueue(peer.queue.outbound.c)
 
 	// Use the device batch size, not the bind batch size, as the device size is
 	// the size of the batch pools.