ntfy/util/batching_queue.go

package util

import (
	"sync"
	"time"
)

// BatchingQueue is a queue that creates batches of the enqueued elements based on a
// max batch size and a batch timeout.
//
// Example:
//
//	q := NewBatchingQueue[int](2, 500 * time.Millisecond)
//	go func() {
//	  for batch := range q.Dequeue() {
//	    fmt.Println(batch)
//	  }
//	}()
//	q.Enqueue(1)
//	q.Enqueue(2)
//	q.Enqueue(3)
//	time.Sleep(time.Second)
//
// This example will emit batch [1, 2] immediately (because the batch size is 2), and
// a batch [3] after 500ms.
type BatchingQueue[T any] struct {
	batchSize int
	timeout   time.Duration
	in        []T
	out       chan []T
	mu        sync.Mutex
}

// NewBatchingQueue creates a new BatchingQueue
func NewBatchingQueue[T any](batchSize int, timeout time.Duration) *BatchingQueue[T] {
	q := &BatchingQueue[T]{
		batchSize: batchSize,
		timeout:   timeout,
		in:        make([]T, 0),
		out:       make(chan []T),
	}
	go q.timeoutTicker()
	return q
}

// Enqueue enqueues an element to the queue. If the configured batch size is reached,
// the batch will be emitted immediately.
func (q *BatchingQueue[T]) Enqueue(element T) {
	q.mu.Lock()
	q.in = append(q.in, element)
	limitReached := len(q.in) == q.batchSize
	q.mu.Unlock()
	if limitReached {
		q.out <- q.dequeueAll()
	}
}

// Dequeue returns a channel emitting batches of elements
func (q *BatchingQueue[T]) Dequeue() <-chan []T {
	return q.out
}

func (q *BatchingQueue[T]) dequeueAll() []T {
	q.mu.Lock()
	defer q.mu.Unlock()
	elements := make([]T, len(q.in))
	copy(elements, q.in)
	q.in = q.in[:0]
	return elements
}

func (q *BatchingQueue[T]) timeoutTicker() {
	if q.timeout == 0 {
		return
	}
	ticker := time.NewTicker(q.timeout)
	for range ticker.C {
		elements := q.dequeueAll()
		if len(elements) > 0 {
			q.out <- elements
		}
	}
}