151 lines
3.4 KiB
Go
151 lines
3.4 KiB
Go
package usage
|
|
|
|
import (
|
|
"RenewCMS/internal/domain/image"
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
const FlushKeyLimit = 500
|
|
|
|
func FlushInterval() time.Duration {
|
|
return 30 * time.Second
|
|
}
|
|
|
|
type hitKey struct {
|
|
variantName string
|
|
sourceName string
|
|
}
|
|
|
|
type Tracker struct {
|
|
repository image.VariantUsageRepository
|
|
now func() time.Time
|
|
interval time.Duration
|
|
keyLimit int
|
|
flushing sync.Mutex
|
|
pending sync.Mutex
|
|
hits map[hitKey]time.Time
|
|
rejections map[time.Time]int
|
|
wake chan struct{}
|
|
}
|
|
|
|
var _ image.VariantUsageTracker = &Tracker{}
|
|
|
|
func NewTracker(repository image.VariantUsageRepository, now func() time.Time) *Tracker {
|
|
return NewTrackerWith(repository, now, FlushInterval(), FlushKeyLimit)
|
|
}
|
|
|
|
func NewTrackerWith(repository image.VariantUsageRepository, now func() time.Time, interval time.Duration, keyLimit int) *Tracker {
|
|
return &Tracker{
|
|
repository: repository,
|
|
now: now,
|
|
interval: interval,
|
|
keyLimit: keyLimit,
|
|
hits: map[hitKey]time.Time{},
|
|
rejections: map[time.Time]int{},
|
|
wake: make(chan struct{}, 1),
|
|
}
|
|
}
|
|
|
|
func (t *Tracker) RecordHit(variantName string, sourceName string) {
|
|
at := t.now().UTC()
|
|
t.pending.Lock()
|
|
defer t.pending.Unlock()
|
|
keepLatest(t.hits, hitKey{variantName, sourceName}, at)
|
|
if len(t.hits) >= t.keyLimit {
|
|
t.wakeLoop()
|
|
}
|
|
}
|
|
|
|
func (t *Tracker) RecordRejection() {
|
|
day := image.RejectionDay(t.now())
|
|
t.pending.Lock()
|
|
defer t.pending.Unlock()
|
|
t.rejections[day]++
|
|
}
|
|
|
|
func (t *Tracker) wakeLoop() {
|
|
select {
|
|
case t.wake <- struct{}{}:
|
|
default:
|
|
}
|
|
}
|
|
|
|
func keepLatest(hits map[hitKey]time.Time, key hitKey, at time.Time) {
|
|
if current, found := hits[key]; !found || at.After(current) {
|
|
hits[key] = at
|
|
}
|
|
}
|
|
|
|
func (t *Tracker) Flush() error {
|
|
t.flushing.Lock()
|
|
defer t.flushing.Unlock()
|
|
|
|
hits, rejections := t.take()
|
|
if len(hits) == 0 && len(rejections) == 0 {
|
|
return nil
|
|
}
|
|
batch := image.VariantUsageBatch{
|
|
Hits: hitsOf(hits),
|
|
Rejections: rejectionsOf(rejections),
|
|
PruneBefore: image.RejectionPruneBefore(t.now()),
|
|
}
|
|
if err := t.repository.Record(batch); err != nil {
|
|
t.restore(hits, rejections)
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (t *Tracker) take() (map[hitKey]time.Time, map[time.Time]int) {
|
|
t.pending.Lock()
|
|
defer t.pending.Unlock()
|
|
hits, rejections := t.hits, t.rejections
|
|
t.hits, t.rejections = map[hitKey]time.Time{}, map[time.Time]int{}
|
|
return hits, rejections
|
|
}
|
|
|
|
func (t *Tracker) restore(hits map[hitKey]time.Time, rejections map[time.Time]int) {
|
|
t.pending.Lock()
|
|
defer t.pending.Unlock()
|
|
for key, at := range hits {
|
|
keepLatest(t.hits, key, at)
|
|
}
|
|
for day, count := range rejections {
|
|
t.rejections[day] += count
|
|
}
|
|
}
|
|
|
|
func hitsOf(hits map[hitKey]time.Time) []image.VariantHit {
|
|
listed := make([]image.VariantHit, 0, len(hits))
|
|
for key, at := range hits {
|
|
listed = append(listed, image.VariantHit{VariantName: key.variantName, SourceName: key.sourceName, At: at})
|
|
}
|
|
return listed
|
|
}
|
|
|
|
func rejectionsOf(rejections map[time.Time]int) []image.RejectionCount {
|
|
listed := make([]image.RejectionCount, 0, len(rejections))
|
|
for day, count := range rejections {
|
|
listed = append(listed, image.RejectionCount{Day: day, Count: count})
|
|
}
|
|
return listed
|
|
}
|
|
|
|
func (t *Tracker) Run(ctx context.Context) {
|
|
ticker := time.NewTicker(t.interval)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
_ = t.Flush()
|
|
return
|
|
case <-ticker.C:
|
|
_ = t.Flush()
|
|
case <-t.wake:
|
|
_ = t.Flush()
|
|
}
|
|
}
|
|
}
|