77 lines
1.5 KiB
Go
77 lines
1.5 KiB
Go
|
// Package sample provides a bounded queue that keeps a random sample of recent
|
||
|
|
// items instead of a FIFO. A FIFO fills with the first N items and drops the
|
||
|
|
// rest, which starves later work; random replacement gives representative
|
||
|
|
// coverage at bounded memory. Not for hot paths.
|
||
|
|
package sample
|
||
|
|
|
||
|
|
import (
|
||
|
|
mathrand "math/rand/v2"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Queue[T any] struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
items []T
|
||
|
|
cap int
|
||
|
|
notify chan struct{}
|
||
|
|
stop chan struct{}
|
||
|
|
once sync.Once
|
||
|
|
}
|
||
|
|
|
||
|
|
func New[T any](capacity int) *Queue[T] {
|
||
|
|
if capacity < 1 {
|
||
|
|
capacity = 1
|
||
|
|
}
|
||
|
|
return &Queue[T]{
|
||
|
|
cap: capacity,
|
||
|
|
notify: make(chan struct{}, 1),
|
||
|
|
stop: make(chan struct{}),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *Queue[T]) Push(v T) {
|
||
|
|
q.mu.Lock()
|
||
|
|
if len(q.items) < q.cap {
|
||
|
|
q.items = append(q.items, v)
|
||
|
|
} else {
|
||
|
|
q.items[mathrand.IntN(len(q.items))] = v
|
||
|
|
}
|
||
|
|
q.mu.Unlock()
|
||
|
|
select {
|
||
|
|
case q.notify <- struct{}{}:
|
||
|
|
default:
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *Queue[T]) Pop() (T, bool) {
|
||
|
|
q.mu.Lock()
|
||
|
|
n := len(q.items)
|
||
|
|
if n == 0 {
|
||
|
|
q.mu.Unlock()
|
||
|
|
var zero T
|
||
|
|
return zero, false
|
||
|
|
}
|
||
|
|
i := mathrand.IntN(n)
|
||
|
|
v := q.items[i]
|
||
|
|
q.items[i] = q.items[n-1]
|
||
|
|
q.items = q.items[:n-1]
|
||
|
|
q.mu.Unlock()
|
||
|
|
return v, true
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *Queue[T]) Len() int {
|
||
|
|
q.mu.Lock()
|
||
|
|
defer q.mu.Unlock()
|
||
|
|
return len(q.items)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *Queue[T]) Notify() <-chan struct{} { return q.notify }
|
||
|
|
func (q *Queue[T]) Stop() <-chan struct{} { return q.stop }
|
||
|
|
|
||
|
|
func (q *Queue[T]) Close() {
|
||
|
|
q.once.Do(func() { close(q.stop) })
|
||
|
|
select {
|
||
|
|
case q.notify <- struct{}{}:
|
||
|
|
default:
|
||
|
|
}
|
||
|
|
}
|