Update

vendor/github.com/tailscale/wireguard-go/device/pools.go

@@ -7,6 +7,8 @@ package device
 
 import (
 	"sync"
+
+	"github.com/tailscale/wireguard-go/tun"
 )
 
 type WaitPool struct {
@@ -46,29 +48,77 @@ func (p *WaitPool) Put(x any) {
 	p.cond.Signal()
 }
 
+const (
+	// smallPacketBufSize is the size to be used for [packetBuf.slab]'s pooled in
+	// [Device.pool.smallPacketBufs]. It's large enough to fit WireGuard keepalive
+	// messages, WireGuard-encrypted [MaxPriorityMessageContentSize], and
+	// WireGuard-encrypted internet MTU messages.
+	smallPacketBufSize = 2048
+
+	internetMTU = 1500
+	// assert that [smallPacketBufSize] is sufficient to hold [internetMTU]-sized
+	// plaintext read via [tun.Device.Read], for eventual encryption and
+	// transmission as a WireGuard transport message.
+	_ uint = smallPacketBufSize -
+		tun.ReadPacketSpacing -
+		internetMTU -
+		outboundPlaintextTailroom
+	// assert that [smallPacketBufSize] is sufficient to hold [MaxPriorityMessageContentSize]-sized
+	// plaintext for eventual encryption and transmission as a WireGuard transport
+	// message. This assertion also covers keepalive messages, which only have
+	// the auth tag, accounted for in [outboundPlaintextTailroom], following
+	// the WireGuard header.
+	_ uint = smallPacketBufSize -
+		outboundPlaintextHeadroom -
+		MaxPriorityMessageContentSize -
+		outboundPlaintextTailroom
+
+	// minPacketBufSizeForDistinctSmallPool represents the minimum [packetBuf.slab]
+	// size at which a distinct ([Device.pool.smallPacketBufs]) is created.
+	minPacketBufSizeForDistinctSmallPool = 2 * smallPacketBufSize
+)
+
 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 {
-		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)
 	})
+	packetBufSize := singlePacketSlabSize
+	if device.BatchSize() > 1 {
+		packetBufSize = batchingSlabSize
+	}
+	device.pool.packetBufs = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
+		return newPacketBuf(packetBufSize, func(buf *packetBuf) {
+			device.pool.packetBufs.Put(buf)
+		})
+	})
+	// A distinct smallPacketBufs pool makes TUN readers copy packets that fit
+	// into [smallPacketBufSize] in the interest of memory savings. Alias
+	// smallPacketBufs to the "full-sized" pool if there is little to be gained
+	// in memory savings, or if the device is configured to limit outstanding
+	// packet memory pool items, which we don't want to artificially inflate.
+	device.pool.smallPacketBufs = device.pool.packetBufs
+	if packetBufSize >= minPacketBufSizeForDistinctSmallPool && device.config.preallocatedBuffersPerPool == 0 {
+		device.pool.smallPacketBufs = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
+			return newPacketBuf(smallPacketBufSize, func(buf *packetBuf) {
+				device.pool.smallPacketBufs.Put(buf)
+			})
+		})
+	}
 }
 
 func (device *Device) GetInboundElementsContainer() *QueueInboundElementsContainer {
 	c := device.pool.inboundElementsContainer.Get().(*QueueInboundElementsContainer)
-	c.Mutex = sync.Mutex{}
 	return c
 }
 
@@ -82,7 +132,6 @@ func (device *Device) PutInboundElementsContainer(c *QueueInboundElementsContain
 
 func (device *Device) GetOutboundElementsContainer() *QueueOutboundElementsContainer {
 	c := device.pool.outboundElementsContainer.Get().(*QueueOutboundElementsContainer)
-	c.Mutex = sync.Mutex{}
 	return c
 }
 
@@ -94,12 +143,25 @@ func (device *Device) PutOutboundElementsContainer(c *QueueOutboundElementsConta
 	device.pool.outboundElementsContainer.Put(c)
 }
 
-func (device *Device) GetMessageBuffer() *[MaxMessageSize]byte {
-	return device.pool.messageBuffers.Get().(*[MaxMessageSize]byte)
+// hasDistinctSmallPacketBufPool returns true if [Device.getSmallPacketBuf]
+// will fetch from a pool independent of [Device.getPacketBuf].
+func (device *Device) hasDistinctSmallPacketBufPool() bool {
+	return device.pool.smallPacketBufs != device.pool.packetBufs
 }
 
-func (device *Device) PutMessageBuffer(msg *[MaxMessageSize]byte) {
-	device.pool.messageBuffers.Put(msg)
+// getSmallPacketBuf returns a [*packetBuf] smaller or equal in size to what
+// [Device.getPacketBuf] returns depending on [Device] batching and [WaitPool]
+// configuration.
+func (device *Device) getSmallPacketBuf() *packetBuf {
+	b := device.pool.smallPacketBufs.Get().(*packetBuf)
+	b.incRef()
+	return b
+}
+
+func (device *Device) getPacketBuf() *packetBuf {
+	b := device.pool.packetBufs.Get().(*packetBuf)
+	b.incRef()
+	return b
 }
 
 func (device *Device) GetInboundElement() *QueueInboundElement {
@@ -107,15 +169,25 @@ func (device *Device) GetInboundElement() *QueueInboundElement {
 }
 
 func (device *Device) PutInboundElement(elem *QueueInboundElement) {
+	elem.buffer.decRef()
 	elem.clearPointers()
 	device.pool.inboundElements.Put(elem)
 }
 
+// GetOutboundElement returns a [*QueueOutboundElement] with all its fields
+// set to their respective zero values.
 func (device *Device) GetOutboundElement() *QueueOutboundElement {
-	return device.pool.outboundElements.Get().(*QueueOutboundElement)
+	elem := device.pool.outboundElements.Get().(*QueueOutboundElement)
+	elem.plaintextOffset = 0
+	elem.nonce = 0
+	// buffer, packet, keypair, and peer were cleared (if necessary) by [QueueOutboundElement.clearPointers].
+	return elem
 }
 
 func (device *Device) PutOutboundElement(elem *QueueOutboundElement) {
+	if elem.buffer != nil {
+		elem.buffer.decRef()
+	}
 	elem.clearPointers()
 	device.pool.outboundElements.Put(elem)
 }