knowledge-governance/internal/platform/postgres/v1_scheduler_integration_test.go

package postgres

import (
	"context"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

// TestV1SourceRefreshScheduler proves the dispatcher ACTUALLY consumes the
// per-source schedule: a due, enabled source with an active revision interval
// gets exactly one catalog refresh job, its next horizon advances by that
// interval, and the disabled / in-flight / not-due / unconfigured states are
// each honoured.
func TestV1SourceRefreshScheduler(t *testing.T) {
	databaseURL := isolatedIntegrationDatabaseURL(t, os.Getenv("TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	if err := ApplyMigrations(ctx, pool); err != nil {
		t.Fatal(err)
	}
	now := time.Date(2026, 9, 12, 8, 0, 0, 0, time.UTC)

	seed := func(id, sourceType, name string, enabled bool, schedule string, nextRefresh *time.Time, revisionInterval string) {
		var nextArg any
		if nextRefresh != nil {
			nextArg = nextRefresh.UTC()
		}
		if _, err := pool.Exec(ctx, `INSERT INTO source_connections(id,source_type,name,enabled,schedule,next_refresh_at) VALUES($1,$2,$3,$4,$5::jsonb,$6)`,
			id, sourceType, name, enabled, schedule, nextArg); err != nil {
			t.Fatal(err)
		}
		revisionSchedule := "{}"
		if revisionInterval != "" {
			revisionSchedule = `{"refresh_interval_seconds":` + revisionInterval + `}`
		}
		if _, err := pool.Exec(ctx, `INSERT INTO source_connection_revisions(id,connection_id,revision_number,state,source_type,name,enabled,schedule,created_by,validated_at,activated_at)
VALUES('active-'||$1,$1,1,'active',$2,$3,$4,$5::jsonb,'tester',$6,$6)`, id, sourceType, name, enabled, revisionSchedule, now.UTC()); err != nil {
			t.Fatal(err)
		}
		if _, err := pool.Exec(ctx, `UPDATE source_connections SET current_revision_id='active-'||$1 WHERE id=$1`, id); err != nil {
			t.Fatal(err)
		}
	}
	past := now.Add(-time.Minute)

	// due hourly source
	seed("conn-due", "alist", "到期源", true, "{}", &past, "3600")
	// disabled even though due
	seed("conn-off", "alist", "停用源", false, "{}", &past, "3600")
	// due but the active revision has no interval (defensive floor 300s advances)
	seed("conn-nointerval", "alist", "无周期源", true, "{}", &past, "")
	// an import-only manual source is never dispatchable even if a horizon is
	// somehow present (activation rejects schedules for it)
	seed("conn-offline", "manual", "离线手工源", true, "{}", &past, "")
	// not configured: next_refresh_at NULL
	seed("conn-uncfg", "alist", "未配置源", true, "{}", nil, "3600")
	// due and has a pending refresh job already (in-flight window)
	seed("conn-busy", "alist", "在途源", true, "{}", &past, "3600")
	if _, err := pool.Exec(ctx, `INSERT INTO jobs(id,type,payload,state,priority,available_at,max_attempts,idempotency_key,created_at,updated_at)
VALUES('job-busy','source.catalog.refresh',jsonb_build_object('connection_id','conn-busy'),'pending',20,$1,3,'manual-busy',$1,$1)`, now.UTC()); err != nil {
		t.Fatal(err)
	}

	store := NewStore(pool)
	claimed, enqueued, err := store.ClaimDueSourceRefreshes(ctx, now)
	if err != nil {
		t.Fatal(err)
	}
	if claimed != 3 || enqueued != 2 {
		t.Fatalf("claimed=%d enqueued=%d want 3/2", claimed, enqueued)
	}

	// exactly one new job for conn-due and conn-manual, none for conn-busy
	var dueJobs, busyJobs int
	if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM jobs WHERE type='source.catalog.refresh' AND payload->>'connection_id'='conn-due'`).Scan(&dueJobs); err != nil {
		t.Fatal(err)
	}
	if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM jobs WHERE type='source.catalog.refresh' AND payload->>'connection_id'='conn-busy'`).Scan(&busyJobs); err != nil {
		t.Fatal(err)
	}
	if dueJobs != 1 || busyJobs != 1 {
		t.Fatalf("dueJobs=%d busyJobs=%d want 1/1", dueJobs, busyJobs)
	}

	// horizons advanced by the ACTIVE revision interval (hourly source -> +1h)
	var horizon time.Time
	if err := pool.QueryRow(ctx, `SELECT next_refresh_at FROM source_connections WHERE id='conn-due'`).Scan(&horizon); err != nil {
		t.Fatal(err)
	}
	if !horizon.Equal(now.Add(time.Hour)) {
		t.Fatalf("conn-due horizon=%v want %v", horizon, now.Add(time.Hour))
	}
	if err := pool.QueryRow(ctx, `SELECT next_refresh_at FROM source_connections WHERE id='conn-nointerval'`).Scan(&horizon); err != nil {
		t.Fatal(err)
	}
	if !horizon.Equal(now.Add(5 * time.Minute)) {
		t.Fatalf("conn-nointerval horizon=%v want five minute floor", horizon)
	}

	// manual sources are never dispatchable: horizon frozen at the seeded past
	var offlineHorizon time.Time
	if err := pool.QueryRow(ctx, `SELECT next_refresh_at FROM source_connections WHERE id='conn-offline'`).Scan(&offlineHorizon); err != nil {
		t.Fatal(err)
	}
	if !offlineHorizon.Equal(past) {
		t.Fatalf("manual source was claimed: horizon moved to %v", offlineHorizon)
	}

	// disabled and unconfigured sources untouched
	// neither was claimed: no last_scheduled_at stamp; conn-uncfg keeps NULL
	// horizon, conn-off keeps its seeded past value untouched
	var offScheduled, uncfgScheduled bool
	if err := pool.QueryRow(ctx, `SELECT last_scheduled_at IS NOT NULL FROM source_connections WHERE id='conn-off'`).Scan(&offScheduled); err != nil {
		t.Fatal(err)
	}
	if err := pool.QueryRow(ctx, `SELECT last_scheduled_at IS NOT NULL FROM source_connections WHERE id='conn-uncfg'`).Scan(&uncfgScheduled); err != nil {
		t.Fatal(err)
	}
	if offScheduled || uncfgScheduled {
		t.Fatalf("offScheduled=%v uncfgScheduled=%v want both false", offScheduled, uncfgScheduled)
	}

	// second tick at the same instant claims nothing (horizons already future)
	claimed2, enqueued2, err := store.ClaimDueSourceRefreshes(ctx, now)
	if err != nil {
		t.Fatal(err)
	}
	if claimed2 != 0 || enqueued2 != 0 {
		t.Fatalf("replay claimed=%d enqueued=%d want 0/0", claimed2, enqueued2)
	}

	// simulate the worker finishing the two scheduled jobs (conn-busy's stays
	// pending); after one hour all three are due again, but only the two whose
	// jobs completed get a NEW job with a NEW idempotency window.
	if _, err := pool.Exec(ctx, `UPDATE jobs SET state='completed', updated_at=$1::timestamptz WHERE type='source.catalog.refresh' AND payload->>'connection_id' IN ('conn-due','conn-nointerval') AND state='pending'`, now.Add(30*time.Minute)); err != nil {
		t.Fatal(err)
	}
	claimed3, enqueued3, err := store.ClaimDueSourceRefreshes(ctx, now.Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if claimed3 != 3 || enqueued3 != 2 {
		t.Fatalf("hour later claimed=%d enqueued=%d want 3/2", claimed3, enqueued3)
	}
}

// TestV1ConcurrentDispatchersClaimOnce proves the SKIP LOCKED lease: two
// dispatcher goroutines ticking at the same instant enqueue exactly ONE refresh
// job for the same due source, never two.
func TestV1ConcurrentDispatchersClaimOnce(t *testing.T) {
	databaseURL := isolatedIntegrationDatabaseURL(t, os.Getenv("TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	if err := ApplyMigrations(ctx, pool); err != nil {
		t.Fatal(err)
	}
	now := time.Date(2026, 9, 12, 10, 0, 0, 0, time.UTC)
	past := now.Add(-time.Minute)
	if _, err := pool.Exec(ctx, `INSERT INTO source_connections(id,source_type,name,enabled,next_refresh_at)
VALUES('conn-race','alist','并发源',true,$1)`, past); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `INSERT INTO source_connection_revisions(id,connection_id,revision_number,state,source_type,name,enabled,schedule,created_by,validated_at,activated_at)
VALUES('active-race','conn-race',1,'active','alist','并发源',true,'{"refresh_interval_seconds":3600}','tester',$1,$1)`, now); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `UPDATE source_connections SET current_revision_id='active-race' WHERE id='conn-race'`); err != nil {
		t.Fatal(err)
	}
	store := NewStore(pool)

	var wg sync.WaitGroup
	totalEnqueued := 0
	var mu sync.Mutex
	for i := 0; i < 4; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, enq, err := store.ClaimDueSourceRefreshes(ctx, now)
			if err != nil {
				t.Errorf("dispatch: %v", err)
				return
			}
			mu.Lock()
			totalEnqueued += enq
			mu.Unlock()
		}()
	}
	wg.Wait()
	if totalEnqueued != 1 {
		t.Fatalf("total enqueued=%d want exactly 1 across 4 concurrent dispatchers", totalEnqueued)
	}
	var jobs int
	if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM jobs WHERE type='source.catalog.refresh' AND payload->>'connection_id'='conn-race'`).Scan(&jobs); err != nil {
		t.Fatal(err)
	}
	if jobs != 1 {
		t.Fatalf("stored jobs=%d want 1", jobs)
	}
}