diff --git a/backend/internal/handler/admin/grok_import_probe.go b/backend/internal/handler/admin/grok_import_probe.go index d301776419..f1df15bba9 100644 --- a/backend/internal/handler/admin/grok_import_probe.go +++ b/backend/internal/handler/admin/grok_import_probe.go @@ -33,12 +33,6 @@ type grokImportProbeScheduler struct { timeout time.Duration } -type grokImportProbeSchedulerSnapshot struct { - Queued int - Workers int - MaxWorkers int -} - var defaultGrokImportProbeScheduler = newGrokImportProbeScheduler( grokImportProbeConcurrency, grokImportProbeTimeout, @@ -103,19 +97,6 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) { return task, true } -func (s *grokImportProbeScheduler) snapshot() grokImportProbeSchedulerSnapshot { - if s == nil { - return grokImportProbeSchedulerSnapshot{} - } - s.mu.Lock() - defer s.mu.Unlock() - return grokImportProbeSchedulerSnapshot{ - Queued: len(s.queue), - Workers: s.workers, - MaxWorkers: s.maxWorkers, - } -} - func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) { defer func() { if recovered := recover(); recovered != nil { diff --git a/backend/internal/handler/admin/grok_import_probe_test.go b/backend/internal/handler/admin/grok_import_probe_test.go index 27d0b744e2..3b8fc0ca6e 100644 --- a/backend/internal/handler/admin/grok_import_probe_test.go +++ b/backend/internal/handler/admin/grok_import_probe_test.go @@ -86,6 +86,25 @@ func (s *grokImportProbeStub) snapshot() (map[int64]int, int, bool) { return calls, s.maxActive, s.deadlineSeen } +type grokImportProbeSchedulerTestSnapshot struct { + queued int + workers int + maxWorkers int +} + +func snapshotGrokImportProbeScheduler(s *grokImportProbeScheduler) grokImportProbeSchedulerTestSnapshot { + if s == nil { + return grokImportProbeSchedulerTestSnapshot{} + } + s.mu.Lock() + defer s.mu.Unlock() + return grokImportProbeSchedulerTestSnapshot{ + queued: len(s.queue), + workers: s.workers, + maxWorkers: s.maxWorkers, + } +} + func newGrokOAuthImportAccount(id int64) *service.Account { return &service.Account{ ID: id, @@ -117,8 +136,8 @@ func TestGrokImportProbeSchedulerProbesSingleAccountOnce(t *testing.T) { require.Equal(t, 1, maxActive) require.True(t, deadlineSeen) require.Eventually(t, func() bool { - snapshot := scheduler.snapshot() - return snapshot.Queued == 0 && snapshot.Workers == 0 + snapshot := snapshotGrokImportProbeScheduler(scheduler) + return snapshot.queued == 0 && snapshot.workers == 0 }, time.Second, 10*time.Millisecond) } @@ -136,10 +155,10 @@ func TestGrokImportProbeSchedulerQueuesBatchWithoutPerTaskGoroutines(t *testing. for i := 0; i < 3; i++ { awaitGrokProbeSignal(t, prober.started) } - snapshot := scheduler.snapshot() - require.Equal(t, 97, snapshot.Queued) - require.Equal(t, 3, snapshot.Workers) - require.Equal(t, 3, snapshot.MaxWorkers) + snapshot := snapshotGrokImportProbeScheduler(scheduler) + require.Equal(t, 97, snapshot.queued) + require.Equal(t, 3, snapshot.workers) + require.Equal(t, 3, snapshot.maxWorkers) select { case id := <-prober.started: t.Fatalf("probe %d started before a concurrency slot was released", id) @@ -157,10 +176,10 @@ func TestGrokImportProbeSchedulerQueuesBatchWithoutPerTaskGoroutines(t *testing. } require.Equal(t, 3, maxActive) require.Eventually(t, func() bool { - snapshot = scheduler.snapshot() - return snapshot.Queued == 0 && snapshot.Workers == 0 + snapshot = snapshotGrokImportProbeScheduler(scheduler) + return snapshot.queued == 0 && snapshot.workers == 0 }, time.Second, 10*time.Millisecond) - require.Equal(t, 3, snapshot.MaxWorkers) + require.Equal(t, 3, snapshot.maxWorkers) } func TestGrokImportProbeSchedulerTimeoutCancelsProbe(t *testing.T) {