mirror of
https://github.com/multipleof4/ntfy.git
synced 2026-10-15 16:07:55 +00:00
WIP: Clsuter support (cross node delivery, leader election)
This commit is contained in:
binwiederhier
+103
@@ -0,0 +1,103 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"hash/fnv"
|
||||
"math"
|
||||
)
|
||||
|
||||
// BloomFilter is a fixed-size probabilistic set: Contains may return false positives (bounded by
|
||||
// the target rate the filter was sized for) but never false negatives. Elements cannot be
|
||||
// removed; rebuild the filter from scratch instead.
|
||||
type BloomFilter struct {
|
||||
bits []uint64
|
||||
k int // number of hash probes per element, derived via double hashing
|
||||
}
|
||||
|
||||
// NewBloomFilter creates a filter sized for n elements at the given false-positive rate.
|
||||
func NewBloomFilter(n int, fpRate float64) *BloomFilter {
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
if fpRate <= 0 || fpRate >= 1 {
|
||||
fpRate = 0.01
|
||||
}
|
||||
m := int(math.Ceil(-float64(n) * math.Log(fpRate) / (math.Ln2 * math.Ln2))) // bits
|
||||
k := int(math.Round(float64(m) / float64(n) * math.Ln2)) // probes
|
||||
if k < 1 {
|
||||
k = 1
|
||||
}
|
||||
return &BloomFilter{
|
||||
bits: make([]uint64, (m+63)/64),
|
||||
k: k,
|
||||
}
|
||||
}
|
||||
|
||||
// Add inserts an element into the filter.
|
||||
func (b *BloomFilter) Add(s string) {
|
||||
h1, h2 := hashPair(s)
|
||||
m := uint64(len(b.bits)) * 64
|
||||
for i := 0; i < b.k; i++ {
|
||||
bit := (h1 + uint64(i)*h2) % m
|
||||
b.bits[bit/64] |= 1 << (bit % 64)
|
||||
}
|
||||
}
|
||||
|
||||
// Contains reports whether the element may be in the set. A false result is definitive: the
|
||||
// element was never added.
|
||||
func (b *BloomFilter) Contains(s string) bool {
|
||||
h1, h2 := hashPair(s)
|
||||
m := uint64(len(b.bits)) * 64
|
||||
for i := 0; i < b.k; i++ {
|
||||
bit := (h1 + uint64(i)*h2) % m
|
||||
if b.bits[bit/64]&(1<<(bit%64)) == 0 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// MarshalBinary serializes the filter as [k][bits...], with 64-bit little-endian words.
|
||||
func (b *BloomFilter) MarshalBinary() ([]byte, error) {
|
||||
data := make([]byte, 1+len(b.bits)*8)
|
||||
data[0] = byte(b.k)
|
||||
for i, word := range b.bits {
|
||||
for j := 0; j < 8; j++ {
|
||||
data[1+i*8+j] = byte(word >> (8 * j))
|
||||
}
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// UnmarshalBloomFilter deserializes a filter produced by MarshalBinary.
|
||||
func UnmarshalBloomFilter(data []byte) (*BloomFilter, error) {
|
||||
if len(data) < 9 || (len(data)-1)%8 != 0 {
|
||||
return nil, errors.New("invalid bloom filter data")
|
||||
}
|
||||
b := &BloomFilter{
|
||||
bits: make([]uint64, (len(data)-1)/8),
|
||||
k: int(data[0]),
|
||||
}
|
||||
if b.k < 1 {
|
||||
return nil, errors.New("invalid bloom filter hash count")
|
||||
}
|
||||
for i := range b.bits {
|
||||
var word uint64
|
||||
for j := 0; j < 8; j++ {
|
||||
word |= uint64(data[1+i*8+j]) << (8 * j)
|
||||
}
|
||||
b.bits[i] = word
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// hashPair derives the two independent hash values used for double hashing (probe i uses
|
||||
// h1 + i*h2), from FNV-1a over the element and a domain-separated variant of it.
|
||||
func hashPair(s string) (uint64, uint64) {
|
||||
f := fnv.New64a()
|
||||
f.Write([]byte(s))
|
||||
h1 := f.Sum64()
|
||||
f.Write([]byte{0xff}) // Domain separation for the second hash
|
||||
h2 := f.Sum64() | 1 // Odd, so probes cycle through all bit positions
|
||||
return h1, h2
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package util_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"heckel.io/ntfy/v2/util"
|
||||
)
|
||||
|
||||
func TestBloomFilter_NoFalseNegatives(t *testing.T) {
|
||||
// The property routing correctness relies on: an added element is ALWAYS reported present
|
||||
b := util.NewBloomFilter(1000, 0.01)
|
||||
for i := 0; i < 1000; i++ {
|
||||
b.Add(fmt.Sprintf("topic-%d", i))
|
||||
}
|
||||
for i := 0; i < 1000; i++ {
|
||||
require.True(t, b.Contains(fmt.Sprintf("topic-%d", i)))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBloomFilter_FalsePositiveRate(t *testing.T) {
|
||||
b := util.NewBloomFilter(10000, 0.01)
|
||||
for i := 0; i < 10000; i++ {
|
||||
b.Add(fmt.Sprintf("added-%d", i))
|
||||
}
|
||||
falsePositives := 0
|
||||
for i := 0; i < 10000; i++ {
|
||||
if b.Contains(fmt.Sprintf("absent-%d", i)) {
|
||||
falsePositives++
|
||||
}
|
||||
}
|
||||
require.Less(t, falsePositives, 300, "expected ~1%% false positives, got %d/10000", falsePositives)
|
||||
}
|
||||
|
||||
func TestBloomFilter_EmptyContainsNothing(t *testing.T) {
|
||||
b := util.NewBloomFilter(100, 0.01)
|
||||
require.False(t, b.Contains("anything"))
|
||||
}
|
||||
|
||||
func TestBloomFilter_MarshalRoundTrip(t *testing.T) {
|
||||
b := util.NewBloomFilter(500, 0.01)
|
||||
for i := 0; i < 500; i++ {
|
||||
b.Add(fmt.Sprintf("topic-%d", i))
|
||||
}
|
||||
data, err := b.MarshalBinary()
|
||||
require.Nil(t, err)
|
||||
b2, err := util.UnmarshalBloomFilter(data)
|
||||
require.Nil(t, err)
|
||||
for i := 0; i < 500; i++ {
|
||||
require.True(t, b2.Contains(fmt.Sprintf("topic-%d", i)))
|
||||
}
|
||||
require.False(t, b2.Contains("never-added-topic"))
|
||||
}
|
||||
|
||||
func TestBloomFilter_UnmarshalGarbage(t *testing.T) {
|
||||
_, err := util.UnmarshalBloomFilter([]byte{})
|
||||
require.Error(t, err)
|
||||
_, err = util.UnmarshalBloomFilter([]byte{1, 2, 3})
|
||||
require.Error(t, err)
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// LingerQueue is a bounded, non-blocking batching queue: enqueued elements are emitted as
|
||||
// batches once a linger window expires, a batch count cap is reached, or a batch size cap is
|
||||
// reached, whichever comes first. Unlike BatchingQueue, producers never block: TryEnqueue drops
|
||||
// (returns false) when the queue is full, and Close flushes the remainder and terminates the
|
||||
// consumer channel, so per-entity queues can be created and destroyed dynamically.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// q := NewLingerQueue[int](64, 10, 0, nil, 500*time.Millisecond)
|
||||
// go func() {
|
||||
// for batch := range q.Dequeue() {
|
||||
// send(batch)
|
||||
// }
|
||||
// }()
|
||||
// q.TryEnqueue(1)
|
||||
// q.TryEnqueue(2) // emitted together as [1, 2] after <= 500ms
|
||||
type LingerQueue[T any] struct {
|
||||
in chan T
|
||||
out chan []T
|
||||
max int // max elements per batch
|
||||
maxSize int // max cumulative size per batch; 0 = no size cap
|
||||
size func(T) int // element size function; nil = no size cap
|
||||
linger time.Duration // max time the first element of a batch waits; 0 = emit immediately
|
||||
closed bool
|
||||
mu sync.Mutex // Protects closed, and guards TryEnqueue's send against Close's close(in)
|
||||
}
|
||||
|
||||
// NewLingerQueue creates a LingerQueue holding up to capacity queued elements, emitting batches
|
||||
// of up to max elements or maxSize cumulative size (as measured by size; pass 0/nil for no size
|
||||
// cap) after at most linger.
|
||||
func NewLingerQueue[T any](capacity, max, maxSize int, size func(T) int, linger time.Duration) *LingerQueue[T] {
|
||||
q := &LingerQueue[T]{
|
||||
in: make(chan T, capacity),
|
||||
out: make(chan []T),
|
||||
max: max,
|
||||
maxSize: maxSize,
|
||||
size: size,
|
||||
linger: linger,
|
||||
}
|
||||
go q.run()
|
||||
return q
|
||||
}
|
||||
|
||||
// TryEnqueue enqueues an element without blocking. It returns false if the queue is full or
|
||||
// closed; the caller decides how to account for the drop.
|
||||
func (q *LingerQueue[T]) TryEnqueue(t T) bool {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
if q.closed {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case q.in <- t:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Dequeue returns the channel emitting batches. It is closed after Close, once the remaining
|
||||
// elements have been flushed.
|
||||
func (q *LingerQueue[T]) Dequeue() <-chan []T {
|
||||
return q.out
|
||||
}
|
||||
|
||||
// Close stops the queue: remaining elements are flushed as final batches, then the Dequeue
|
||||
// channel is closed. TryEnqueue returns false after Close. Close is idempotent.
|
||||
func (q *LingerQueue[T]) Close() {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
if q.closed {
|
||||
return
|
||||
}
|
||||
q.closed = true
|
||||
close(q.in)
|
||||
}
|
||||
|
||||
// run is the batching loop: it blocks for the first element of a batch, then collects more until
|
||||
// the linger expires or a cap is hit, and emits the batch. It exits once the queue is closed and
|
||||
// drained. Note that receiving from the closed in channel still yields the buffered remainder
|
||||
// before reporting closed, which is what flushes on Close.
|
||||
func (q *LingerQueue[T]) run() {
|
||||
defer close(q.out)
|
||||
for {
|
||||
first, ok := <-q.in
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
batch := []T{first}
|
||||
bytes := q.sizeOf(first)
|
||||
var timeout <-chan time.Time
|
||||
if q.linger > 0 {
|
||||
timeout = time.After(q.linger)
|
||||
}
|
||||
closed := false
|
||||
collect:
|
||||
for len(batch) < q.max && (q.maxSize <= 0 || bytes < q.maxSize) {
|
||||
if timeout == nil {
|
||||
// Zero linger: greedily drain what is immediately available, never wait
|
||||
select {
|
||||
case t, ok := <-q.in:
|
||||
if !ok {
|
||||
closed = true
|
||||
break collect
|
||||
}
|
||||
batch = append(batch, t)
|
||||
bytes += q.sizeOf(t)
|
||||
default:
|
||||
break collect
|
||||
}
|
||||
} else {
|
||||
select {
|
||||
case t, ok := <-q.in:
|
||||
if !ok {
|
||||
closed = true
|
||||
break collect
|
||||
}
|
||||
batch = append(batch, t)
|
||||
bytes += q.sizeOf(t)
|
||||
case <-timeout:
|
||||
break collect
|
||||
}
|
||||
}
|
||||
}
|
||||
q.out <- batch
|
||||
if closed {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (q *LingerQueue[T]) sizeOf(t T) int {
|
||||
if q.size == nil {
|
||||
return 0
|
||||
}
|
||||
return q.size(t)
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package util_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"heckel.io/ntfy/v2/util"
|
||||
)
|
||||
|
||||
func TestLingerQueue_LingerFlush(t *testing.T) {
|
||||
// Items enqueued within the linger window are emitted as a single batch when it expires
|
||||
q := util.NewLingerQueue[int](16, 100, 0, nil, 50*time.Millisecond)
|
||||
defer q.Close()
|
||||
require.True(t, q.TryEnqueue(1))
|
||||
require.True(t, q.TryEnqueue(2))
|
||||
require.True(t, q.TryEnqueue(3))
|
||||
start := time.Now()
|
||||
batch := <-q.Dequeue()
|
||||
require.Equal(t, []int{1, 2, 3}, batch)
|
||||
require.GreaterOrEqual(t, time.Since(start), 30*time.Millisecond) // Waited out the linger
|
||||
}
|
||||
|
||||
func TestLingerQueue_MaxBatchFlush(t *testing.T) {
|
||||
// Hitting the count cap flushes early, before the linger expires
|
||||
q := util.NewLingerQueue[int](16, 5, 0, nil, time.Minute)
|
||||
defer q.Close()
|
||||
for i := 0; i < 12; i++ {
|
||||
require.True(t, q.TryEnqueue(i))
|
||||
}
|
||||
require.Len(t, <-q.Dequeue(), 5)
|
||||
require.Len(t, <-q.Dequeue(), 5)
|
||||
q.Close() // Flushes the remainder
|
||||
require.Len(t, <-q.Dequeue(), 2)
|
||||
}
|
||||
|
||||
func TestLingerQueue_SizeCapFlush(t *testing.T) {
|
||||
// Hitting the byte cap flushes early, before count cap or linger
|
||||
q := util.NewLingerQueue(16, 100, 10, func(s string) int { return len(s) }, time.Minute)
|
||||
defer q.Close()
|
||||
require.True(t, q.TryEnqueue("aaaa"))
|
||||
require.True(t, q.TryEnqueue("bbbb"))
|
||||
require.True(t, q.TryEnqueue("cccc")) // 12 bytes >= 10 -> flush
|
||||
batch := <-q.Dequeue()
|
||||
require.Equal(t, []string{"aaaa", "bbbb", "cccc"}, batch)
|
||||
}
|
||||
|
||||
func TestLingerQueue_TryEnqueueFull(t *testing.T) {
|
||||
// A full queue drops (returns false) instead of blocking the producer
|
||||
q := util.NewLingerQueue[int](1, 1, 0, nil, 0)
|
||||
defer q.Close()
|
||||
require.True(t, q.TryEnqueue(1)) // Taken by the batcher, blocks emitting (no consumer)
|
||||
waitForCond(t, func() bool { return q.TryEnqueue(2) }) // Fills the buffer once slot frees
|
||||
require.False(t, q.TryEnqueue(3)) // Buffer full, batcher blocked -> drop
|
||||
}
|
||||
|
||||
func TestLingerQueue_CloseFlushesAndCloses(t *testing.T) {
|
||||
q := util.NewLingerQueue[int](16, 100, 0, nil, time.Minute)
|
||||
require.True(t, q.TryEnqueue(1))
|
||||
require.True(t, q.TryEnqueue(2))
|
||||
q.Close()
|
||||
require.Equal(t, []int{1, 2}, <-q.Dequeue()) // Remainder flushed without waiting out the linger
|
||||
_, ok := <-q.Dequeue()
|
||||
require.False(t, ok) // Channel closed
|
||||
require.False(t, q.TryEnqueue(3))
|
||||
q.Close() // Idempotent
|
||||
}
|
||||
|
||||
func TestLingerQueue_LingerZeroImmediate(t *testing.T) {
|
||||
q := util.NewLingerQueue[int](16, 100, 0, nil, 0)
|
||||
defer q.Close()
|
||||
require.True(t, q.TryEnqueue(1))
|
||||
select {
|
||||
case batch := <-q.Dequeue():
|
||||
require.Contains(t, batch, 1)
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("expected immediate flush with zero linger")
|
||||
}
|
||||
}
|
||||
|
||||
func waitForCond(t *testing.T, f func() bool) {
|
||||
t.Helper()
|
||||
for i := 0; i < 100; i++ {
|
||||
if f() {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("timed out waiting for condition")
|
||||
}
|
||||
Reference in New Issue
Block a user