diff --git a/device/channels.go b/device/channels.go index be15d1c..1eaec56 100644 --- a/device/channels.go +++ b/device/channels.go @@ -126,7 +126,7 @@ func (device *Device) flushOutboundQueue(q *autodrainingOutboundQueue) { case elemsContainer := <-q.c: elemsContainer.Lock() for _, elem := range elemsContainer.elems { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) } device.PutOutboundElementsContainer(elemsContainer) diff --git a/device/device.go b/device/device.go index 8afa6fd..529951a 100644 --- a/device/device.go +++ b/device/device.go @@ -71,11 +71,11 @@ type Device struct { cookieChecker CookieChecker pool struct { - inboundElementsContainer *WaitPool - outboundElementsContainer *WaitPool + inboundElementsContainer *sync.Pool + outboundElementsContainer *sync.Pool messageBuffers *WaitPool - inboundElements *WaitPool - outboundElements *WaitPool + inboundElements *sync.Pool + outboundElements *sync.Pool } queue struct { diff --git a/device/peer.go b/device/peer.go index bca121d..a2703bd 100644 --- a/device/peer.go +++ b/device/peer.go @@ -25,6 +25,8 @@ type Peer struct { rxBytes atomic.Uint64 // bytes received from peer lastHandshakeNano atomic.Int64 // nano seconds since epoch + queuedOutboundPackets atomic.Int32 // packets in staged+outbound queues, for input backpressure + endpoint struct { sync.Mutex val conn.Endpoint @@ -193,6 +195,7 @@ func (peer *Peer) Start() { // reset routine state peer.stopping.Wait() peer.stopping.Add(2) + peer.queuedOutboundPackets.Store(0) peer.handshake.mutex.Lock() peer.handshake.lastSentHandshake = time.Now().Add(-(RekeyTimeout + time.Second)) diff --git a/device/pools.go b/device/pools.go index 2c18f41..b7536a3 100644 --- a/device/pools.go +++ b/device/pools.go @@ -7,6 +7,8 @@ package device import ( "sync" + + "github.com/sagernet/sing/common/buf" ) type WaitPool struct { @@ -47,23 +49,23 @@ func (p *WaitPool) Put(x any) { } func (device *Device) PopulatePools() { - device.pool.inboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any { + device.pool.inboundElementsContainer = &sync.Pool{New: func() any { s := make([]*QueueInboundElement, 0, device.BatchSize()) return &QueueInboundElementsContainer{elems: s} - }) - device.pool.outboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any { + }} + device.pool.outboundElementsContainer = &sync.Pool{New: 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 = &sync.Pool{New: func() any { return new(QueueInboundElement) - }) - device.pool.outboundElements = NewWaitPool(PreallocatedBuffersPerPool, func() any { + }} + device.pool.outboundElements = &sync.Pool{New: func() any { return new(QueueOutboundElement) - }) + }} } func (device *Device) GetInboundElementsContainer() *QueueInboundElementsContainer { @@ -102,6 +104,20 @@ func (device *Device) PutMessageBuffer(msg *[MaxMessageSize]byte) { device.pool.messageBuffers.Put(msg) } +// Outbound buffers come from the sing allocator instead of the bounded +// messageBuffers pool: the injection paths (InputPacket/InputPackets) run on +// the caller's shared read loop, which must never block on pool exhaustion, +// and their packets are far smaller than MaxMessageSize, so they are allocated +// by actual size. This also keeps the bounded pool exclusively for the receive +// path, so outbound backlog can no longer starve it. +func (device *Device) GetOutboundBuffer(size int) []byte { + return buf.Get(size) +} + +func (device *Device) PutOutboundBuffer(buffer []byte) { + _ = buf.Put(buffer) +} + func (device *Device) GetInboundElement() *QueueInboundElement { return device.pool.inboundElements.Get().(*QueueInboundElement) } diff --git a/device/send.go b/device/send.go index 7a94068..bf4b85c 100644 --- a/device/send.go +++ b/device/send.go @@ -45,7 +45,7 @@ import ( */ type QueueOutboundElement struct { - buffer *[MaxMessageSize]byte // slice holding the packet data + buffer []byte // sing-allocated buffer holding the packet data // packet is always a slice of "buffer". The starting offset in buffer // is either: // a) MessageEncapsulatingTransportSize+MessageTransportHeaderSize (plaintext) @@ -63,7 +63,7 @@ type QueueOutboundElementsContainer struct { func (device *Device) NewOutboundElement() *QueueOutboundElement { elem := device.GetOutboundElement() - elem.buffer = device.GetMessageBuffer() + elem.buffer = device.GetOutboundBuffer(MaxMessageSize) elem.nonce = 0 // keypair and peer were cleared (if necessary) by clearPointers. return elem @@ -89,9 +89,10 @@ func (peer *Peer) SendKeepalive() { elemsContainer.elems = append(elemsContainer.elems, elem) select { case peer.queue.staged <- elemsContainer: + peer.queuedOutboundPackets.Add(1) peer.device.log.Verbosef("%v - Sending keepalive packet", peer) default: - peer.device.PutMessageBuffer(elem.buffer) + peer.device.PutOutboundBuffer(elem.buffer) peer.device.PutOutboundElement(elem) peer.device.PutOutboundElementsContainer(elemsContainer) } @@ -238,7 +239,7 @@ func (device *Device) RoutineReadFromTUN() { defer func() { for _, elem := range elems { if elem != nil { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) } } @@ -295,7 +296,7 @@ func (device *Device) RoutineReadFromTUN() { peer.SendStagedPackets() } else { for _, elem := range elemsForPeer.elems { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) } device.PutOutboundElementsContainer(elemsForPeer) @@ -322,22 +323,33 @@ func (device *Device) RoutineReadFromTUN() { } } +// maxQueuedInputPackets bounds the staged+outbound backlog of a peer fed via +// InputPacket/InputPackets. Injected packets beyond it are dropped before they +// are copied into pooled message buffers, like a full qdisc: injection has no +// flow control, and the queues are bounded in containers (up to a full batch +// each), so without this cap a flood is buffered instead of dropped. +const maxQueuedInputPackets = 2048 + func (device *Device) InputPacket(destination []byte, packetSlices [][]byte) { peer := device.allowedips.Lookup(destination) if peer == nil { return } - elem := device.NewOutboundElement() - packet := elem.buffer[MessageEncapsulatingTransportSize+MessageTransportHeaderSize:] + if peer.queuedOutboundPackets.Load() >= maxQueuedInputPackets { + return + } var totalLength int for _, packetSlice := range packetSlices { totalLength += len(packetSlice) } - if totalLength > len(packet) { - device.PutMessageBuffer(elem.buffer) - device.PutOutboundElement(elem) + allocLength := MessageEncapsulatingTransportSize + MessageTransportHeaderSize + totalLength + PaddingMultiple + chacha20poly1305.Overhead + if allocLength > MaxMessageSize { return } + elem := device.GetOutboundElement() + elem.buffer = device.GetOutboundBuffer(allocLength) + elem.nonce = 0 + packet := elem.buffer[MessageEncapsulatingTransportSize+MessageTransportHeaderSize:] var n int for _, packetSlice := range packetSlices { n += copy(packet[n:], packetSlice) @@ -349,7 +361,7 @@ func (device *Device) InputPacket(destination []byte, packetSlices [][]byte) { peer.StagePackets(elemsForPeer) peer.SendStagedPackets() } else { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) device.PutOutboundElementsContainer(elemsForPeer) } @@ -369,17 +381,21 @@ func (device *Device) InputPackets(packets []*InputPacketRef) []*InputPacketRef unmatched = append(unmatched, packetRef) continue } - elem := device.NewOutboundElement() - packet := elem.buffer[MessageEncapsulatingTransportSize+MessageTransportHeaderSize:] + if peer.queuedOutboundPackets.Load() >= maxQueuedInputPackets { + continue + } var totalLength int for _, packetSlice := range packetRef.PacketSlices { totalLength += len(packetSlice) } - if totalLength > len(packet) { - device.PutMessageBuffer(elem.buffer) - device.PutOutboundElement(elem) + allocLength := MessageEncapsulatingTransportSize + MessageTransportHeaderSize + totalLength + PaddingMultiple + chacha20poly1305.Overhead + if allocLength > MaxMessageSize { continue } + elem := device.GetOutboundElement() + elem.buffer = device.GetOutboundBuffer(allocLength) + elem.nonce = 0 + packet := elem.buffer[MessageEncapsulatingTransportSize+MessageTransportHeaderSize:] var n int for _, packetSlice := range packetRef.PacketSlices { n += copy(packet[n:], packetSlice) @@ -398,7 +414,7 @@ func (device *Device) InputPackets(packets []*InputPacketRef) []*InputPacketRef peer.SendStagedPackets() } else { for _, elem := range elemsForPeer.elems { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) } device.PutOutboundElementsContainer(elemsForPeer) @@ -408,6 +424,7 @@ func (device *Device) InputPackets(packets []*InputPacketRef) []*InputPacketRef } func (peer *Peer) StagePackets(elems *QueueOutboundElementsContainer) { + peer.queuedOutboundPackets.Add(int32(len(elems.elems))) for { select { case peer.queue.staged <- elems: @@ -416,8 +433,9 @@ func (peer *Peer) StagePackets(elems *QueueOutboundElementsContainer) { } select { case tooOld := <-peer.queue.staged: + peer.queuedOutboundPackets.Add(-int32(len(tooOld.elems))) for _, elem := range tooOld.elems { - peer.device.PutMessageBuffer(elem.buffer) + peer.device.PutOutboundBuffer(elem.buffer) peer.device.PutOutboundElement(elem) } peer.device.PutOutboundElementsContainer(tooOld) @@ -464,6 +482,8 @@ top: elemsContainer.elems = elemsContainer.elems[:i] if elemsContainerOOO != nil { + // Already counted at their original staging; StagePackets will count them again. + peer.queuedOutboundPackets.Add(-int32(len(elemsContainerOOO.elems))) peer.StagePackets(elemsContainerOOO) // XXX: Out of order, but we can't front-load go chans } @@ -477,8 +497,9 @@ top: peer.queue.outbound.c <- elemsContainer peer.device.queue.encryption.c <- elemsContainer } else { + peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems))) for _, elem := range elemsContainer.elems { - peer.device.PutMessageBuffer(elem.buffer) + peer.device.PutOutboundBuffer(elem.buffer) peer.device.PutOutboundElement(elem) } peer.device.PutOutboundElementsContainer(elemsContainer) @@ -497,8 +518,9 @@ func (peer *Peer) FlushStagedPackets() { for { select { case elemsContainer := <-peer.queue.staged: + peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems))) for _, elem := range elemsContainer.elems { - peer.device.PutMessageBuffer(elem.buffer) + peer.device.PutOutboundBuffer(elem.buffer) peer.device.PutOutboundElement(elem) } peer.device.PutOutboundElementsContainer(elemsContainer) @@ -592,8 +614,9 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) { // TODO: rework peer shutdown order to ensure // that we never accidentally keep timers alive longer than necessary. elemsContainer.Lock() + peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems))) for _, elem := range elemsContainer.elems { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) } device.PutOutboundElementsContainer(elemsContainer) @@ -615,8 +638,9 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) { if dataSent { peer.timersDataSent() } + peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems))) for _, elem := range elemsContainer.elems { - device.PutMessageBuffer(elem.buffer) + device.PutOutboundBuffer(elem.buffer) device.PutOutboundElement(elem) } device.PutOutboundElementsContainer(elemsContainer)