diff --git a/apps/api/internal/store/execution_pools.go b/apps/api/internal/store/execution_pools.go index 52f79af..65b4f72 100644 --- a/apps/api/internal/store/execution_pools.go +++ b/apps/api/internal/store/execution_pools.go @@ -101,14 +101,15 @@ func (s *Store) ListWorkers(ctx context.Context, now time.Time) ([]executionpool if now.IsZero() { now = time.Now() } + cutoff := now.Add(-workerHeartbeatStaleAfter) rows, err := s.pool.Query(ctx, ` SELECT worker_id, instance_id, pool_id, endpoint, protocol_version, revision, capabilities, allocated_capacity, safe_capacity, heavy_capacity, active_tasks, pressure_state, heartbeat_at, load_sampled_at FROM gateway_worker_instances WHERE status = 'active' - AND heartbeat_at > $1 - $2::interval -ORDER BY pool_id, instance_id`, now, workerHeartbeatStaleAfter.String()) + AND heartbeat_at > $1::timestamptz +ORDER BY pool_id, instance_id`, cutoff) if err != nil { return nil, err } @@ -346,6 +347,7 @@ func (s *Store) ListCapacity(ctx context.Context, now time.Time) ([]executionpoo if now.IsZero() { now = time.Now() } + cutoff := now.Add(-workerHeartbeatStaleAfter) rows, err := s.pool.Query(ctx, ` SELECT pool.pool_id, COUNT(worker.instance_id)::int, @@ -358,12 +360,12 @@ SELECT pool.pool_id, COALESCE(MAX(worker.load_sampled_at), $1) FROM gateway_execution_pools pool LEFT JOIN gateway_worker_instances worker - ON worker.pool_id = pool.pool_id + ON worker.pool_id = pool.pool_id AND worker.status = 'active' - AND worker.heartbeat_at > $1 - $2::interval + AND worker.heartbeat_at > $2::timestamptz WHERE pool.state = 'active' GROUP BY pool.pool_id -ORDER BY pool.pool_id`, now, workerHeartbeatStaleAfter.String()) +ORDER BY pool.pool_id`, now, cutoff) if err != nil { return nil, err } diff --git a/apps/api/internal/store/execution_pools_integration_test.go b/apps/api/internal/store/execution_pools_integration_test.go new file mode 100644 index 0000000..303364d --- /dev/null +++ b/apps/api/internal/store/execution_pools_integration_test.go @@ -0,0 +1,20 @@ +package store + +import ( + "context" + "testing" + "time" +) + +func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if _, err := db.ListWorkers(ctx, time.Now()); err != nil { + t.Fatalf("list workers with timestamp cutoff: %v", err) + } + if _, err := db.ListCapacity(ctx, time.Now()); err != nil { + t.Fatalf("list capacity with timestamp cutoff: %v", err) + } +}