device: add options for queue configuration

device/channels.go

@@ -22,9 +22,9 @@ type outboundQueue struct {
 	wg sync.WaitGroup
 }
 
-func newOutboundQueue() *outboundQueue {
+func newOutboundQueue(capacity int) *outboundQueue {
 	q := &outboundQueue{
-		c: make(chan *QueueOutboundElementsContainer, QueueOutboundSize),
+		c: make(chan *QueueOutboundElementsContainer, capacity),
 	}
 	q.wg.Add(1)
 	go func() {
@@ -40,9 +40,9 @@ type inboundQueue struct {
 	wg sync.WaitGroup
 }
 
-func newInboundQueue() *inboundQueue {
+func newInboundQueue(capacity int) *inboundQueue {
 	q := &inboundQueue{
-		c: make(chan *QueueInboundElementsContainer, QueueInboundSize),
+		c: make(chan *QueueInboundElementsContainer, capacity),
 	}
 	q.wg.Add(1)
 	go func() {
@@ -58,9 +58,9 @@ type handshakeQueue struct {
 	wg sync.WaitGroup
 }
 
-func newHandshakeQueue() *handshakeQueue {
+func newHandshakeQueue(capacity int) *handshakeQueue {
 	q := &handshakeQueue{
-		c: make(chan QueueHandshakeElement, QueueHandshakeSize),
+		c: make(chan QueueHandshakeElement, capacity),
 	}
 	q.wg.Add(1)
 	go func() {

device/device.go

@@ -20,6 +20,8 @@ import (
 )
 
 type Device struct {
+	config config
+
 	state struct {
 		// state holds the device's state. It is accessed atomically.
 		// Use the device.deviceState method to read it.
@@ -96,6 +98,79 @@ type Device struct {
 	log      *Logger
 }
 
+type config struct {
+	queueStagedSize            int
+	queueOutboundSize          int
+	queueInboundSize           int
+	queueHandshakeSize         int
+	preallocatedBuffersPerPool uint32
+}
+
+func defaultConfig() config {
+	return config{
+		queueStagedSize:            DefaultQueueStagedSize,
+		queueOutboundSize:          DefaultQueueOutboundSize,
+		queueInboundSize:           DefaultQueueInboundSize,
+		queueHandshakeSize:         DefaultQueueHandshakeSize,
+		preallocatedBuffersPerPool: DefaultPreallocatedBuffersPerPool,
+	}
+}
+
+// An Option configures a [Device].
+type Option interface {
+	apply(*config)
+}
+
+type optionFunc func(*config)
+
+func (f optionFunc) apply(config *config) {
+	f(config)
+}
+
+// WithQueueStagedSize sets the capacity of each peer's staged packet queue.
+// Staged packet queues must be buffered, so max(1, size) is applied to the
+// user-supplied value. [DefaultQueueStagedSize] is the default.
+func WithQueueStagedSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueStagedSize = max(1, size)
+	})
+}
+
+// WithQueueOutboundSize sets the capacity of the device and per-peer outbound
+// packet queues. [DefaultQueueOutboundSize] is the default.
+func WithQueueOutboundSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueOutboundSize = size
+	})
+}
+
+// WithQueueInboundSize sets the capacity of the device and per-peer inbound
+// packet queues. [DefaultQueueInboundSize] is the default.
+func WithQueueInboundSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueInboundSize = size
+	})
+}
+
+// WithQueueHandshakeSize sets the capacity of the device's handshake queue.
+// [DefaultQueueHandshakeSize] is the default.
+func WithQueueHandshakeSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueHandshakeSize = size
+	})
+}
+
+// WithPreallocatedBuffersPerPool sets the maximum number of packet memory
+// pool outstanding items. A value of zero is unlimited. Care must be taken to
+// not set a size that is too small, which can lead to immediate deadlock, or
+// increase the probability of deadlock once packets start flowing.
+// See tailscale/corp#46396. [DefaultPreallocatedBuffersPerPool] is the default.
+func WithPreallocatedBuffersPerPool(size uint32) Option {
+	return optionFunc(func(config *config) {
+		config.preallocatedBuffersPerPool = size
+	})
+}
+
 // deviceState represents the state of a Device.
 // There are three states: down, up, closed.
 // Transitions:
@@ -230,7 +305,9 @@ func (device *Device) Down() error {
 func (device *Device) IsUnderLoad() bool {
 	// check if currently under load
 	now := time.Now()
-	underLoad := len(device.queue.handshake.c) >= QueueHandshakeSize/8
+	// max(1, ...) is required on the right hand side, otherwise underLoad would
+	// always be true when queueHandshakeSize < 8.
+	underLoad := len(device.queue.handshake.c) >= max(1, device.config.queueHandshakeSize/8)
 	if underLoad {
 		device.rate.underLoadUntil.Store(now.Add(UnderLoadAfterTime).UnixNano())
 		return true
@@ -294,8 +371,11 @@ func (device *Device) SetPrivateKey(sk NoisePrivateKey) error {
 	return nil
 }
 
-func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger) *Device {
-	device := new(Device)
+func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger, opts ...Option) *Device {
+	device := &Device{config: defaultConfig()}
+	for _, opt := range opts {
+		opt.apply(&device.config)
+	}
 	device.state.state.Store(uint32(deviceStateDown))
 	device.closed = make(chan struct{})
 	device.log = logger
@@ -315,9 +395,9 @@ func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger) *Device {
 
 	// create queues
 
-	device.queue.handshake = newHandshakeQueue()
-	device.queue.encryption = newOutboundQueue()
-	device.queue.decryption = newInboundQueue()
+	device.queue.handshake = newHandshakeQueue(device.config.queueHandshakeSize)
+	device.queue.encryption = newOutboundQueue(device.config.queueOutboundSize)
+	device.queue.decryption = newInboundQueue(device.config.queueInboundSize)
 
 	// start workers
 

device/device_test.go

@@ -652,3 +652,36 @@ func TestPeerIfRunningWaitQueueActors(t *testing.T) {
 		}
 	})
 }
+
+func TestDeviceConfig(t *testing.T) {
+	c := defaultConfig()
+	if c.queueStagedSize != DefaultQueueStagedSize ||
+		c.queueOutboundSize != DefaultQueueOutboundSize ||
+		c.queueInboundSize != DefaultQueueInboundSize ||
+		c.queueHandshakeSize != DefaultQueueHandshakeSize ||
+		c.preallocatedBuffersPerPool != DefaultPreallocatedBuffersPerPool {
+		t.Fatalf("default device config: %+v", c)
+	}
+
+}
+
+func TestDeviceOptions(t *testing.T) {
+	c := defaultConfig()
+	opts := []Option{
+		WithQueueStagedSize(1),
+		WithQueueOutboundSize(2),
+		WithQueueInboundSize(3),
+		WithQueueHandshakeSize(4),
+		WithPreallocatedBuffersPerPool(5),
+	}
+	for _, opt := range opts {
+		opt.apply(&c)
+	}
+	if c.queueStagedSize != 1 ||
+		c.queueOutboundSize != 2 ||
+		c.queueInboundSize != 3 ||
+		c.queueHandshakeSize != 4 ||
+		c.preallocatedBuffersPerPool != 5 {
+		t.Fatalf("configured device config: %+v", c)
+	}
+}

device/peer.go

@@ -121,7 +121,7 @@ func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
 
 	// staged is never closed, and it can be accessed concurrent to [Peer.Start],
 	// so we init here instead of [Peer.Start].
-	peer.queue.staged = make(chan *QueueOutboundElementsContainer, QueueStagedSize)
+	peer.queue.staged = make(chan *QueueOutboundElementsContainer, device.config.queueStagedSize)
 
 	// map public key
 	_, ok := device.peers.keyMap[pk]
@@ -250,8 +250,8 @@ func (peer *Peer) Start() {
 	device.log.Verbosef("%v - Starting", peer)
 
 	// init inbound & outbound packet queues
-	peer.queue.outbound = make(chan *QueueOutboundElementsContainer, QueueOutboundSize)
-	peer.queue.inbound = make(chan *QueueInboundElementsContainer, QueueInboundSize)
+	peer.queue.outbound = make(chan *QueueOutboundElementsContainer, device.config.queueOutboundSize)
+	peer.queue.inbound = make(chan *QueueInboundElementsContainer, device.config.queueInboundSize)
 
 	// reset routine state
 	peer.runningState.queueReaders.Wait()

device/pools.go

@@ -47,21 +47,21 @@ func (p *WaitPool) Put(x any) {
 }
 
 func (device *Device) PopulatePools() {
-	device.pool.inboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.inboundElementsContainer = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		s := make([]*QueueInboundElement, 0, device.BatchSize())
 		return &QueueInboundElementsContainer{elems: s}
 	})
-	device.pool.outboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.outboundElementsContainer = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		s := make([]*QueueOutboundElement, 0, device.BatchSize())
 		return &QueueOutboundElementsContainer{elems: s}
 	})
-	device.pool.messageBuffers = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.messageBuffers = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		return new([MaxMessageSize]byte)
 	})
-	device.pool.inboundElements = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.inboundElements = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		return new(QueueInboundElement)
 	})
-	device.pool.outboundElements = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.outboundElements = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		return new(QueueOutboundElement)
 	})
 }

device/queueconstants_android.go

@@ -10,10 +10,10 @@ import "github.com/tailscale/wireguard-go/conn"
 /* Reduce memory consumption for Android */
 
 const (
-	QueueStagedSize            = conn.IdealBatchSize
-	QueueOutboundSize          = 1024
-	QueueInboundSize           = 1024
-	QueueHandshakeSize         = 1024
-	MaxSegmentSize             = 2200
-	PreallocatedBuffersPerPool = 4096
+	DefaultQueueStagedSize            = conn.IdealBatchSize
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	MaxSegmentSize                    = 2200
+	DefaultPreallocatedBuffersPerPool = 4096
 )

device/queueconstants_default.go

@@ -10,10 +10,10 @@ package device
 import "github.com/tailscale/wireguard-go/conn"
 
 const (
-	QueueStagedSize            = conn.IdealBatchSize
-	QueueOutboundSize          = 1024
-	QueueInboundSize           = 1024
-	QueueHandshakeSize         = 1024
-	MaxSegmentSize             = (1 << 16) - 1 // largest possible UDP datagram
-	PreallocatedBuffersPerPool = 0             // Disable and allow for infinite memory growth
+	DefaultQueueStagedSize            = conn.IdealBatchSize
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	MaxSegmentSize                    = (1 << 16) - 1 // largest possible UDP datagram
+	DefaultPreallocatedBuffersPerPool = 0             // Disable and allow for infinite memory growth
 )

device/queueconstants_ios.go

@@ -7,15 +7,15 @@
 
 package device
 
-// Fit within memory limits for iOS's Network Extension API, which has stricter requirements.
-// These are vars instead of consts, because heavier network extensions might want to reduce
-// them further.
-var (
-	QueueStagedSize                   = 128
-	QueueOutboundSize                 = 1024
-	QueueInboundSize                  = 1024
-	QueueHandshakeSize                = 1024
-	PreallocatedBuffersPerPool uint32 = 1024
+// Fit within memory limits for iOS's Network Extension API, which has stricter
+// requirements. Heavier network extensions can reduce these further using
+// [Device.Option]'s.
+const (
+	DefaultQueueStagedSize            = 128
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	DefaultPreallocatedBuffersPerPool = 1024
 )
 
 const MaxSegmentSize = 1700

device/queueconstants_windows.go

@@ -6,10 +6,10 @@
 package device
 
 const (
-	QueueStagedSize            = 128
-	QueueOutboundSize          = 1024
-	QueueInboundSize           = 1024
-	QueueHandshakeSize         = 1024
-	MaxSegmentSize             = 2048 - 32 // largest possible UDP datagram
-	PreallocatedBuffersPerPool = 0         // Disable and allow for infinite memory growth
+	DefaultQueueStagedSize            = 128
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	MaxSegmentSize                    = 2048 - 32 // largest possible UDP datagram
+	DefaultPreallocatedBuffersPerPool = 0         // Disable and allow for infinite memory growth
 )