Add reference counting to packet buffers.

PiperOrigin-RevId: 408426639

2 files changed, 5 insertions, 7 deletions

diff --git a/pkg/tcpip/link/qdisc/fifo/endpoint.go b/pkg/tcpip/link/qdisc/fifo/endpoint.go
index c15cbf81b..a68b274b2 100644
--- a/pkg/tcpip/link/qdisc/fifo/endpoint.go
+++ b/pkg/tcpip/link/qdisc/fifo/endpoint.go
@@ -94,9 +94,7 @@ func (q *queueDispatcher) dispatchLoop() {
 			// We pass a protocol of zero here because each packet carries its
 			// NetworkProtocol.
 			q.lower.WritePackets(stack.RouteInfo{}, batch, 0 /* protocol */)
-			for pkt := batch.Front(); pkt != nil; pkt = pkt.Next() {
-				batch.Remove(pkt)
-			}
+			batch.DecRef()
 			batch.Reset()
 		}
 	}
diff --git a/pkg/tcpip/link/qdisc/fifo/packet_buffer_queue.go b/pkg/tcpip/link/qdisc/fifo/packet_buffer_queue.go
index eb5abb906..09e5b8314 100644
--- a/pkg/tcpip/link/qdisc/fifo/packet_buffer_queue.go
+++ b/pkg/tcpip/link/qdisc/fifo/packet_buffer_queue.go
@@ -54,13 +54,13 @@ func (q *packetBufferQueue) setLimit(limit int) {
 // enqueue adds the given packet to the queue.
 //
 // Returns true when the PacketBuffer is successfully added to the queue, in
-// which case ownership of the reference is transferred to the queue. And
-// returns false if the queue is full, in which case ownership is retained by
-// the caller.
+// which case the queue acquires a reference to the PacketBuffer, and
+// returns false if the queue is full.
 func (q *packetBufferQueue) enqueue(s *stack.PacketBuffer) bool {
 	q.mu.Lock()
 	r := q.used < q.limit
 	if r {
+		s.IncRef()
 		q.list.PushBack(s)
 		q.used++
 	}
@@ -70,7 +70,7 @@ func (q *packetBufferQueue) enqueue(s *stack.PacketBuffer) bool {
 }
 
 // dequeue removes and returns the next PacketBuffer from queue, if one exists.
-// Ownership is transferred to the caller.
+// Caller is responsible for calling DecRef on the PacketBuffer.
 func (q *packetBufferQueue) dequeue() *stack.PacketBuffer {
 	q.mu.Lock()
 	s := q.list.Front()