diff --git a/device/channels.go b/device/channels.go index bdba7dee2..711f033a3 100644 --- a/device/channels.go +++ b/device/channels.go @@ -83,10 +83,18 @@ func newAutodrainingInboundQueue(device *Device) *autodrainingInboundQueue { q := &autodrainingInboundQueue{ c: make(chan *QueueInboundElementsContainer, QueueInboundSize), } - runtime.SetFinalizer(q, device.flushInboundQueue) + if device.needsInboundQueueFinalizer() { + runtime.SetFinalizer(q, device.flushInboundQueue) + } return q } +func (device *Device) needsInboundQueueFinalizer() bool { + return device.pool.messageBuffers.hasAccounting() || + device.pool.inboundElements.hasAccounting() || + device.pool.inboundElementsContainer.hasAccounting() +} + func (device *Device) flushInboundQueue(q *autodrainingInboundQueue) { for { select { @@ -116,10 +124,18 @@ func newAutodrainingOutboundQueue(device *Device) *autodrainingOutboundQueue { q := &autodrainingOutboundQueue{ c: make(chan *QueueOutboundElementsContainer, QueueOutboundSize), } - runtime.SetFinalizer(q, device.flushOutboundQueue) + if device.needsOutboundQueueFinalizer() { + runtime.SetFinalizer(q, device.flushOutboundQueue) + } return q } +func (device *Device) needsOutboundQueueFinalizer() bool { + return device.pool.messageBuffers.hasAccounting() || + device.pool.outboundElements.hasAccounting() || + device.pool.outboundElementsContainer.hasAccounting() +} + func (device *Device) flushOutboundQueue(q *autodrainingOutboundQueue) { for { select { diff --git a/device/channels_test.go b/device/channels_test.go new file mode 100644 index 000000000..990497f64 --- /dev/null +++ b/device/channels_test.go @@ -0,0 +1,53 @@ +/* SPDX-License-Identifier: MIT + * + * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved. + */ + +package device + +import "testing" + +func TestAutodrainingQueueFinalizerNeedTracksPoolAccounting(t *testing.T) { + unbounded := func() *WaitPool { return NewWaitPool(0, func() any { return nil }) } + bounded := func() *WaitPool { return NewWaitPool(1, func() any { return nil }) } + + device := &Device{} + device.pool.messageBuffers = unbounded() + device.pool.inboundElements = unbounded() + device.pool.inboundElementsContainer = unbounded() + device.pool.outboundElements = unbounded() + device.pool.outboundElementsContainer = unbounded() + + if device.needsInboundQueueFinalizer() { + t.Fatal("unbounded inbound pools should not need queue finalizer") + } + if device.needsOutboundQueueFinalizer() { + t.Fatal("unbounded outbound pools should not need queue finalizer") + } + + device.pool.inboundElementsContainer = bounded() + if !device.needsInboundQueueFinalizer() { + t.Fatal("bounded inbound pool should need queue finalizer") + } + if device.needsOutboundQueueFinalizer() { + t.Fatal("bounded inbound pool should not affect outbound queue finalizer") + } + + device.pool.inboundElementsContainer = unbounded() + device.pool.outboundElementsContainer = bounded() + if device.needsInboundQueueFinalizer() { + t.Fatal("bounded outbound pool should not affect inbound queue finalizer") + } + if !device.needsOutboundQueueFinalizer() { + t.Fatal("bounded outbound pool should need queue finalizer") + } + + device.pool.outboundElementsContainer = unbounded() + device.pool.messageBuffers = bounded() + if !device.needsInboundQueueFinalizer() { + t.Fatal("bounded message buffer pool should need inbound queue finalizer") + } + if !device.needsOutboundQueueFinalizer() { + t.Fatal("bounded message buffer pool should need outbound queue finalizer") + } +} diff --git a/device/pools.go b/device/pools.go index 47f952b3c..68714fa2a 100644 --- a/device/pools.go +++ b/device/pools.go @@ -23,6 +23,10 @@ func NewWaitPool(max uint32, new func() any) *WaitPool { return p } +func (p *WaitPool) hasAccounting() bool { + return p != nil && p.max != 0 +} + func (p *WaitPool) Get() any { if p.max != 0 { p.lock.Lock()