63 lines
1.0 KiB
Go
63 lines
1.0 KiB
Go
|
package ingestion
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Batcher[T any] struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
items []T
|
||
|
|
size int
|
||
|
|
timeout time.Duration
|
||
|
|
createdAt time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewBatcher[T any](size int, timeout time.Duration) *Batcher[T] {
|
||
|
|
return &Batcher[T]{
|
||
|
|
items: make([]T, 0, size),
|
||
|
|
size: size,
|
||
|
|
timeout: timeout,
|
||
|
|
createdAt: time.Now(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (b *Batcher[T]) Add(item T) {
|
||
|
|
b.mu.Lock()
|
||
|
|
defer b.mu.Unlock()
|
||
|
|
|
||
|
|
if len(b.items) == 0 {
|
||
|
|
b.createdAt = time.Now()
|
||
|
|
}
|
||
|
|
b.items = append(b.items, item)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (b *Batcher[T]) IsFull() bool {
|
||
|
|
b.mu.Lock()
|
||
|
|
defer b.mu.Unlock()
|
||
|
|
return len(b.items) >= b.size
|
||
|
|
}
|
||
|
|
|
||
|
|
func (b *Batcher[T]) IsExpired() bool {
|
||
|
|
b.mu.Lock()
|
||
|
|
defer b.mu.Unlock()
|
||
|
|
if len(b.items) == 0 {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
return time.Since(b.createdAt) >= b.timeout
|
||
|
|
}
|
||
|
|
|
||
|
|
func (b *Batcher[T]) Get() []T {
|
||
|
|
b.mu.Lock()
|
||
|
|
defer b.mu.Unlock()
|
||
|
|
items := make([]T, len(b.items))
|
||
|
|
copy(items, b.items)
|
||
|
|
return items
|
||
|
|
}
|
||
|
|
|
||
|
|
func (b *Batcher[T]) Clear() {
|
||
|
|
b.mu.Lock()
|
||
|
|
defer b.mu.Unlock()
|
||
|
|
b.items = b.items[:0]
|
||
|
|
}
|