57 lines
989 B
Go
57 lines
989 B
Go
package util
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type BatchingQueue[T any] struct {
|
|
batchSize int
|
|
timeout time.Duration
|
|
in []T
|
|
out chan []T
|
|
mu sync.Mutex
|
|
}
|
|
|
|
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),
|
|
}
|
|
ticker := time.NewTicker(timeout)
|
|
go func() {
|
|
for range ticker.C {
|
|
elements := q.popAll()
|
|
if len(elements) > 0 {
|
|
q.out <- elements
|
|
}
|
|
}
|
|
}()
|
|
return q
|
|
}
|
|
|
|
func (c *BatchingQueue[T]) Push(element T) {
|
|
c.mu.Lock()
|
|
c.in = append(c.in, element)
|
|
limitReached := len(c.in) == c.batchSize
|
|
c.mu.Unlock()
|
|
if limitReached {
|
|
c.out <- c.popAll()
|
|
}
|
|
}
|
|
|
|
func (c *BatchingQueue[T]) Pop() <-chan []T {
|
|
return c.out
|
|
}
|
|
|
|
func (c *BatchingQueue[T]) popAll() []T {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
elements := make([]T, len(c.in))
|
|
copy(elements, c.in)
|
|
c.in = c.in[:0]
|
|
return elements
|
|
}
|