163 lines
4.4 KiB
Go
163 lines
4.4 KiB
Go
// Copyright 2024 The Cockroach Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
|
// implied. See the License for the specific language governing
|
|
// permissions and limitations under the License.
|
|
|
|
package fifo
|
|
|
|
import "sync"
|
|
|
|
// Queue implements an allocation efficient FIFO queue. It is not safe for
|
|
// concurrent access.
|
|
//
|
|
// Note that the queue provides pointer access to the internal storage (via
|
|
// PeekFront and PushBack) so it must be used with care. These pointers must not
|
|
// be used once the respective element is popped out of the queue.
|
|
//
|
|
// -- Implementation --
|
|
//
|
|
// The queue is implemented as a linked list of nodes, where each node is a
|
|
// small ring buffer. The nodes are allocated using a sync.Pool (a single pool
|
|
// is created for any given type and is used for all queues of that type).
|
|
type Queue[T any] struct {
|
|
len int
|
|
head, tail *queueNode[T]
|
|
|
|
pool *QueueBackingPool[T]
|
|
}
|
|
|
|
// MakeQueue constructs a new Queue.
|
|
//
|
|
// The pool should be a singleton object initialized with MakeQueueBackingPool.
|
|
// A single pool can and should be used by all queues of that type.
|
|
func MakeQueue[T any](pool *QueueBackingPool[T]) Queue[T] {
|
|
return Queue[T]{
|
|
pool: pool,
|
|
}
|
|
}
|
|
|
|
// Len returns the current length of the queue.
|
|
func (q *Queue[T]) Len() int {
|
|
return q.len
|
|
}
|
|
|
|
// PushBack adds t to the end of the queue.
|
|
//
|
|
// The returned pointer can be used to modify the element while it is in the
|
|
// queue; it is valid until the element is removed from the queue.
|
|
func (q *Queue[T]) PushBack(t T) *T {
|
|
if q.head == nil {
|
|
q.head = q.pool.get()
|
|
q.tail = q.head
|
|
} else if q.tail.IsFull() {
|
|
newTail := q.pool.get()
|
|
q.tail.next = newTail
|
|
q.tail = newTail
|
|
}
|
|
q.len++
|
|
return q.tail.PushBack(t)
|
|
}
|
|
|
|
// PeekFront returns the current head of the queue, or nil if the queue is
|
|
// empty.
|
|
//
|
|
// The result is only valid until the next call to PopFront.
|
|
func (q *Queue[T]) PeekFront() *T {
|
|
if q.len == 0 {
|
|
return nil
|
|
}
|
|
return q.head.PeekFront()
|
|
}
|
|
|
|
// PopFront removes the current head of the queue.
|
|
//
|
|
// It is illegal to call PopFront on an empty queue.
|
|
func (q *Queue[T]) PopFront() {
|
|
q.head.PopFront()
|
|
// If this is the only node, we don't want to release it; otherwise we would
|
|
// allocate/free a node every time we transition between the queue being empty
|
|
// and non-empty.
|
|
if q.head.len == 0 && q.head.next != nil {
|
|
oldHead := q.head
|
|
q.head = oldHead.next
|
|
q.pool.put(oldHead)
|
|
}
|
|
q.len--
|
|
}
|
|
|
|
// QueueBackingPool is a sync.Pool that used to allocate internal nodes
|
|
// for Queue[T].
|
|
type QueueBackingPool[T any] struct {
|
|
pool sync.Pool
|
|
}
|
|
|
|
// MakeQueueBackingPool makes a queue backing pool. It is intented to be used to
|
|
// initialize a singleton (global) variable. A single pool can and should be
|
|
// used by all queues of that type.
|
|
func MakeQueueBackingPool[T any]() QueueBackingPool[T] {
|
|
return QueueBackingPool[T]{
|
|
pool: sync.Pool{
|
|
New: func() interface{} { return &queueNode[T]{} },
|
|
},
|
|
}
|
|
}
|
|
|
|
func (qp *QueueBackingPool[T]) get() *queueNode[T] {
|
|
return qp.pool.Get().(*queueNode[T])
|
|
}
|
|
|
|
func (qp *QueueBackingPool[T]) put(n *queueNode[T]) {
|
|
*n = queueNode[T]{}
|
|
qp.pool.Put(n)
|
|
}
|
|
|
|
// We batch the allocation of this many queue objects. The value was chosen
|
|
// without experimentation - it provides a reasonable amount of amortization
|
|
// without a very large increase in memory overhead if T is large.
|
|
const queueNodeSize = 8
|
|
|
|
type queueNode[T any] struct {
|
|
buf [queueNodeSize]T
|
|
head, len int32
|
|
next *queueNode[T]
|
|
}
|
|
|
|
func (qn *queueNode[T]) IsFull() bool {
|
|
return qn.len == queueNodeSize
|
|
}
|
|
|
|
func (qn *queueNode[T]) PushBack(t T) *T {
|
|
if invariants && qn.len >= queueNodeSize {
|
|
panic("cannot push back into a full node")
|
|
}
|
|
i := (qn.head + qn.len) % queueNodeSize
|
|
qn.buf[i] = t
|
|
qn.len++
|
|
return &qn.buf[i]
|
|
}
|
|
|
|
func (qn *queueNode[T]) PeekFront() *T {
|
|
return &qn.buf[qn.head]
|
|
}
|
|
|
|
func (qn *queueNode[T]) PopFront() T {
|
|
if invariants && qn.len == 0 {
|
|
panic("cannot pop from empty queue")
|
|
}
|
|
t := qn.buf[qn.head]
|
|
var zero T
|
|
qn.buf[qn.head] = zero
|
|
qn.head = (qn.head + 1) % queueNodeSize
|
|
qn.len--
|
|
return t
|
|
}
|