cloud-ip-validator/internal/db/queries_scale_test.go

package db

import (
	"context"
	"fmt"
	"reflect"
	"sort"
	"testing"
	"time"
)

// scaleAddrs returns n distinct addresses 10.<a>.<b>.<c> in ascending order.
func scaleAddrs(n int) []string {
	out := make([]string, n)
	for i := 0; i < n; i++ {
		out[i] = fmt.Sprintf("10.%d.%d.%d", (i/65536)%256, (i/256)%256, i%256)
	}
	return out
}

// finishWithChecks submits the address, records ok successes and bad
// failures as the current cycle's checks and, when finish != "", finishes the
// row with that overall result.
func finishWithChecks(t *testing.T, d *DB, addr string, ok, bad int, finish string) {
	t.Helper()
	ctx := context.Background()
	ip, err := d.GetIPByAddress(ctx, addr)
	if err != nil {
		t.Fatalf("get %s: %v", addr, err)
	}
	for i := 0; i < ok+bad; i++ {
		if err := d.UpsertCheck(ctx, Check{
			IPID: ip.ID, IPAddress: addr, AttemptNumber: ip.AttemptNumber,
			Source: SourceEgress, CheckType: "https", Target: fmt.Sprintf("https://t%d.test", i),
			Success: i < ok, CheckedAt: Now(),
		}); err != nil {
			t.Fatalf("upsert check: %v", err)
		}
	}
	if finish != "" {
		if err := d.FinishIP(ctx, ip.ID, finish); err != nil {
			t.Fatalf("finish: %v", err)
		}
	}
}

func TestListIPsPageFiltersTotalOrder(t *testing.T) {
	d, ctx := newTestDB(t)
	addrs := []string{"10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.1.1", "10.0.1.2", "192.168.0.10"}
	if _, err := d.SubmitIPs(ctx, addrs); err != nil {
		t.Fatalf("submit: %v", err)
	}
	// 10.0.0.1 pass, 10.0.0.2 fail, 10.0.0.3 checking, 10.0.1.1 occupied.
	finishWithChecks(t, d, "10.0.0.1", 1, 0, ResultPass)
	time.Sleep(3 * time.Millisecond)
	finishWithChecks(t, d, "10.0.0.2", 0, 1, ResultFail)
	ip3, _ := d.GetIPByAddress(ctx, "10.0.0.3")
	if err := d.SetChecking(ctx, ip3.ID, time.Minute); err != nil {
		t.Fatal(err)
	}
	ip4, _ := d.GetIPByAddress(ctx, "10.0.1.1")
	if err := d.MarkFIPOccupied(ctx, ip4.ID, ""); err != nil {
		t.Fatal(err)
	}

	addrsOf := func(items []IPQueueItem) []string {
		out := []string{}
		for _, it := range items {
			out = append(out, it.IPAddress)
		}
		return out
	}

	items, total, err := d.ListIPsPage(ctx, IPFilter{}, 4, 0)
	if err != nil || total != 6 || !reflect.DeepEqual(addrsOf(items), addrs[:4]) {
		t.Fatalf("page 1: total=%d items=%v err=%v", total, addrsOf(items), err)
	}
	items, total, _ = d.ListIPsPage(ctx, IPFilter{}, 4, 4)
	if total != 6 || !reflect.DeepEqual(addrsOf(items), addrs[4:]) {
		t.Fatalf("page 2: total=%d items=%v", total, addrsOf(items))
	}
	items, total, _ = d.ListIPsPage(ctx, IPFilter{}, 4, 100)
	if total != 6 || len(items) != 0 || items == nil {
		t.Fatalf("offset past end: total=%d items=%v", total, items)
	}

	items, total, _ = d.ListIPsPage(ctx, IPFilter{States: []string{IPDone, IPFailed}}, 50, 0)
	if total != 2 || !reflect.DeepEqual(addrsOf(items), []string{"10.0.0.1", "10.0.0.2"}) {
		t.Fatalf("states filter: total=%d items=%v", total, addrsOf(items))
	}
	items, total, _ = d.ListIPsPage(ctx, IPFilter{States: []string{IPQueued}, Query: "10.0.1."}, 50, 0)
	if total != 1 || addrsOf(items)[0] != "10.0.1.2" {
		t.Fatalf("state+q filter: total=%d items=%v", total, addrsOf(items))
	}
	items, total, _ = d.ListIPsPage(ctx, IPFilter{Result: ResultFail}, 50, 0)
	if total != 1 || addrsOf(items)[0] != "10.0.0.2" {
		t.Fatalf("result filter: total=%d items=%v", total, addrsOf(items))
	}
	// q is a plain substring, not a LIKE pattern: % and _ match literally.
	if _, total, _ = d.ListIPsPage(ctx, IPFilter{Query: "%"}, 50, 0); total != 0 {
		t.Fatalf("expected literal substring match, total=%d", total)
	}

	// Newest aggregated first; never-aggregated rows last.
	items, _, _ = d.ListIPsPage(ctx, IPFilter{Order: IPOrderAggregatedAtDesc}, 50, 0)
	got := addrsOf(items)
	if got[0] != "10.0.1.1" && got[0] != "10.0.0.2" {
		t.Fatalf("expected a finished row first, got %v", got)
	}
	if items[0].AggregatedAt == nil || items[len(items)-1].AggregatedAt != nil {
		t.Fatalf("expected aggregated rows first and unaggregated last: %v", got)
	}
	for i := 1; i < len(items); i++ {
		a, b := items[i-1].AggregatedAt, items[i].AggregatedAt
		if a != nil && b != nil && a.Before(*b) {
			t.Fatalf("not sorted by aggregated_at desc at %d: %v", i, got)
		}
	}
}

func TestCountAndNonTerminal(t *testing.T) {
	d, ctx := newTestDB(t)

	byState, total, err := d.CountIPsByState(ctx)
	if err != nil || total != 0 || len(byState) != 0 {
		t.Fatalf("empty: %v %d %v", byState, total, err)
	}
	if any, err := d.AnyNonTerminalIP(ctx); err != nil || any {
		t.Fatalf("empty queue must not report non-terminal: %v %v", any, err)
	}

	if _, err := d.SubmitIPs(ctx, []string{"1.1.1.1", "1.1.1.2", "1.1.1.3", "1.1.1.4"}); err != nil {
		t.Fatal(err)
	}
	finishWithChecks(t, d, "1.1.1.1", 1, 0, ResultPass)
	finishWithChecks(t, d, "1.1.1.2", 1, 1, ResultPartial)
	ip3, _ := d.GetIPByAddress(ctx, "1.1.1.3")
	if err := d.CancelIP(ctx, ip3.ID); err != nil {
		t.Fatal(err)
	}

	byState, total, err = d.CountIPsByState(ctx)
	if err != nil || total != 4 || byState[IPDone] != 2 || byState[IPFailed] != 1 || byState[IPQueued] != 1 {
		t.Fatalf("by state: %v total=%d err=%v", byState, total, err)
	}
	byResult, err := d.CountIPsByResult(ctx)
	if err != nil || len(byResult) != 3 || byResult[ResultPass] != 1 || byResult[ResultPartial] != 1 || byResult[ResultCancelled] != 1 {
		t.Fatalf("by result: %v err=%v", byResult, err)
	}
	if any, _ := d.AnyNonTerminalIP(ctx); !any {
		t.Fatalf("a queued row is non-terminal")
	}
	ip4, _ := d.GetIPByAddress(ctx, "1.1.1.4")
	if err := d.MarkFIPOccupied(ctx, ip4.ID, ""); err != nil {
		t.Fatal(err)
	}
	if any, _ := d.AnyNonTerminalIP(ctx); any {
		t.Fatalf("done/failed/occupied only: expected terminal")
	}
}

// TestListRegistryPageMatchesListRegistry builds a mixed dataset (live rows
// with every overall result, an in-progress live row, deleted rows whose
// last cycle classifies as pass/partial/fail, and an address without checks)
// and verifies that ListRegistryPage's SQL filter selects exactly the rows
// ListRegistry+fillRegistrySummary labels with the same LastResult.
func TestListRegistryPageMatchesListRegistry(t *testing.T) {
	d, ctx := newTestDB(t)
	addrs := scaleAddrs(14)
	if _, err := d.SubmitIPs(ctx, addrs); err != nil {
		t.Fatal(err)
	}
	// live rows with an aggregated result
	finishWithChecks(t, d, addrs[0], 2, 0, ResultPass)
	finishWithChecks(t, d, addrs[1], 1, 1, ResultPartial)
	finishWithChecks(t, d, addrs[2], 0, 2, ResultFail)
	ip3, _ := d.GetIPByAddress(ctx, addrs[3])
	if err := d.CancelIP(ctx, ip3.ID); err != nil {
		t.Fatal(err)
	}
	// live row, passing checks recorded, but cycle unfinished -> no verdict
	finishWithChecks(t, d, addrs[4], 2, 0, "")
	// rows to be deleted: result derived from the checks of the last cycle
	finishWithChecks(t, d, addrs[5], 2, 0, ResultPass)    // -> pass
	finishWithChecks(t, d, addrs[6], 1, 2, ResultPartial) // -> partial
	finishWithChecks(t, d, addrs[7], 0, 2, ResultFail)    // -> fail
	finishWithChecks(t, d, addrs[8], 0, 0, "")            // deleted, no checks -> ""
	// a deleted row whose first cycle failed but whose last passed
	finishWithChecks(t, d, addrs[9], 0, 1, ResultFail)
	if _, err := d.DeleteIPs(ctx, []string{addrs[5], addrs[6], addrs[7], addrs[8], addrs[9]}); err != nil {
		t.Fatal(err)
	}
	if _, err := d.SubmitIPs(ctx, []string{addrs[9]}); err != nil {
		t.Fatal(err)
	}
	finishWithChecks(t, d, addrs[9], 1, 0, ResultPass)
	if _, err := d.DeleteIPs(ctx, []string{addrs[9]}); err != nil {
		t.Fatal(err)
	}

	all, err := d.ListRegistry(ctx)
	if err != nil || len(all) != 14 {
		t.Fatalf("list registry: n=%d err=%v", len(all), err)
	}

	// No filter: same rows in the same order as ListRegistry, paged.
	var paged []RegistrySummary
	for off := 0; ; off += 5 {
		page, total, err := d.ListRegistryPage(ctx, RegistryFilter{}, 5, off)
		if err != nil || total != 14 {
			t.Fatalf("page off=%d total=%d err=%v", off, total, err)
		}
		if len(page) == 0 {
			break
		}
		paged = append(paged, page...)
	}
	if len(paged) != len(all) {
		t.Fatalf("paged %d rows, want %d", len(paged), len(all))
	}
	for i := range all {
		if all[i].IPAddress != paged[i].IPAddress || all[i].LastResult != paged[i].LastResult {
			t.Fatalf("row %d differs: %+v vs %+v", i, all[i], paged[i])
		}
	}

	for _, res := range []string{ResultPass, ResultPartial, ResultFail, ResultCancelled} {
		var want []string
		for _, s := range all {
			if s.LastResult == res {
				want = append(want, s.IPAddress)
			}
		}
		page, total, err := d.ListRegistryPage(ctx, RegistryFilter{LastResult: res}, 100, 0)
		if err != nil {
			t.Fatal(err)
		}
		var got []string
		for _, s := range page {
			got = append(got, s.IPAddress)
			if s.LastResult != res {
				t.Fatalf("%s: row %s has LastResult %q", res, s.IPAddress, s.LastResult)
			}
		}
		sort.Strings(want)
		sort.Strings(got)
		if total != len(want) || !reflect.DeepEqual(got, want) || len(want) == 0 {
			t.Fatalf("last_result=%s: total=%d got=%v want=%v", res, total, got, want)
		}
	}

	// pass: addrs[0] (live) + addrs[5] + addrs[9] (deleted, latest cycle).
	if _, total, _ := d.ListRegistryPage(ctx, RegistryFilter{LastResult: ResultPass}, 100, 0); total != 3 {
		t.Fatalf("expected 3 pass rows, got %d", total)
	}

	// q filter + LIMIT applies after the filter, total is the filtered count.
	page, total, err := d.ListRegistryPage(ctx, RegistryFilter{Query: "10.0.0.1"}, 2, 0)
	if err != nil || total != 5 || len(page) != 2 { // 10.0.0.1, .10-.13
		t.Fatalf("q filter: total=%d len=%d err=%v", total, len(page), err)
	}
	page, total, _ = d.ListRegistryPage(ctx, RegistryFilter{Query: "10.0.0.1", LastResult: ResultPass}, 10, 0)
	if total != 0 || len(page) != 0 {
		// 10.0.0.1 is partial; 10.0.0.10-13 have no verdict or fail.
		t.Fatalf("q+last_result: total=%d page=%+v", total, page)
	}
}

func TestClearAllIPsKeepsHistoryAndFreesValidators(t *testing.T) {
	d, ctx := newTestDB(t)
	if err := d.AdminCreateValidator(ctx, "validator-1", "port-1"); err != nil {
		t.Fatal(err)
	}
	addrs := []string{"5.5.5.1", "5.5.5.2", "5.5.5.3"}
	if _, err := d.SubmitIPs(ctx, addrs); err != nil {
		t.Fatal(err)
	}
	finishWithChecks(t, d, "5.5.5.2", 1, 1, ResultPartial)
	claimed, err := d.ClaimNextQueued(ctx, "validator-1", time.Minute)
	if err != nil || claimed == nil {
		t.Fatalf("claim: %v %v", claimed, err)
	}
	if err := d.SetFIPAssociated(ctx, claimed.ID, "fip-9", time.Minute); err != nil {
		t.Fatal(err)
	}
	if err := d.SetEgressComplete(ctx, claimed.ID); err != nil {
		t.Fatal(err)
	}
	if err := d.SetSiteComplete(ctx, claimed.ID, 1); err != nil {
		t.Fatal(err)
	}

	refs, err := d.ListFIPRefs(ctx)
	if err != nil || len(refs) != 1 || refs[0].FIPID != "fip-9" || refs[0].IPAddress != claimed.IPAddress {
		t.Fatalf("fip refs: %+v err=%v", refs, err)
	}
	refs, err = d.ListFIPRefsByAddresses(ctx, []string{claimed.IPAddress, "nope"})
	if err != nil || len(refs) != 1 {
		t.Fatalf("fip refs by address: %+v err=%v", refs, err)
	}
	if refs, _ := d.ListFIPRefsByAddresses(ctx, []string{"5.5.5.3"}); len(refs) != 0 {
		t.Fatalf("row without fip must not be listed: %+v", refs)
	}

	deleted, err := d.ClearAllIPs(ctx)
	if err != nil {
		t.Fatalf("clear: %v", err)
	}
	sort.Strings(deleted)
	if !reflect.DeepEqual(deleted, addrs) {
		t.Fatalf("deleted %v, want %v", deleted, addrs)
	}
	if items, _ := d.ListIPs(ctx); len(items) != 0 {
		t.Fatalf("queue not empty: %+v", items)
	}
	v, err := d.GetValidator(ctx, "validator-1")
	if err != nil || v.State != ValidatorIdle || v.CurrentIPID != nil {
		t.Fatalf("validator not freed: %+v err=%v", v, err)
	}
	// History and registry rows survive, detached from the queue.
	reg, err := d.GetRegistryByAddress(ctx, "5.5.5.2")
	if err != nil || reg.TotalCycles != 1 || reg.LastResult != ResultPartial || reg.InQueue {
		t.Fatalf("registry after clear: %+v err=%v", reg, err)
	}
	checks, err := d.ListChecksForRegistry(ctx, reg.ID, nil)
	if err != nil || len(checks) != 2 || checks[0].IPID != 0 {
		t.Fatalf("checks after clear: %+v err=%v", checks, err)
	}
	// Clearing an empty queue is fine and returns an empty (non-nil) list.
	if deleted, err := d.ClearAllIPs(ctx); err != nil || deleted == nil || len(deleted) != 0 {
		t.Fatalf("second clear: %v %v", deleted, err)
	}
	// The same addresses can be re-submitted afterwards.
	if res, err := d.SubmitIPs(ctx, addrs); err != nil || len(res.Added) != 3 {
		t.Fatalf("resubmit: %+v err=%v", res, err)
	}
}

func TestMigration0009Indexes(t *testing.T) {
	d, ctx := newTestDB(t)
	for _, name := range []string{"idx_ip_queue_registry", "idx_ip_queue_state_aggregated"} {
		var n int
		if err := d.QueryRowContext(ctx, `SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name=?`, name).Scan(&n); err != nil || n != 1 {
			t.Fatalf("index %s missing (n=%d err=%v)", name, n, err)
		}
	}
	var ver int
	if err := d.QueryRowContext(ctx, `PRAGMA user_version`).Scan(&ver); err != nil || ver < 9 {
		t.Fatalf("user_version=%d err=%v", ver, err)
	}
}

// TestScaleSmoke6440 pushes a realistic project size through the hot paths
// with a loose time bound: the point is the absence of O(n^2) / N+1 work, not
// a benchmark.
func TestScaleSmoke6440(t *testing.T) {
	if testing.Short() {
		t.Skip("scale smoke test skipped in -short mode")
	}
	d, ctx := newTestDB(t)
	addrs := scaleAddrs(6440)

	start := time.Now()
	for off := 0; off < len(addrs); off += 500 {
		end := min(off+500, len(addrs))
		if _, err := d.SubmitIPs(ctx, addrs[off:end]); err != nil {
			t.Fatalf("submit chunk: %v", err)
		}
	}
	submitDur := time.Since(start)

	start = time.Now()
	page, total, err := d.ListRegistryPage(ctx, RegistryFilter{}, 100, 3000)
	if err != nil || total != 6440 || len(page) != 100 {
		t.Fatalf("registry page: total=%d len=%d err=%v", total, len(page), err)
	}
	if _, total, err = d.ListRegistryPage(ctx, RegistryFilter{LastResult: ResultPass}, 100, 0); err != nil || total != 0 {
		t.Fatalf("registry last_result filter: total=%d err=%v", total, err)
	}
	registryDur := time.Since(start)

	start = time.Now()
	if items, total, err := d.ListIPsPage(ctx, IPFilter{States: []string{IPQueued}, Query: "10.0.1."}, 50, 0); err != nil || total != 256 || len(items) != 50 {
		t.Fatalf("ips page: total=%d len=%d err=%v", total, len(items), err)
	}
	if by, total, err := d.CountIPsByState(ctx); err != nil || total != 6440 || by[IPQueued] != 6440 {
		t.Fatalf("count: %v %d %v", by, total, err)
	}
	if any, err := d.AnyNonTerminalIP(ctx); err != nil || !any {
		t.Fatalf("any non terminal: %v %v", any, err)
	}
	queryDur := time.Since(start)

	start = time.Now()
	deleted, err := d.ClearAllIPs(ctx)
	if err != nil || len(deleted) != 6440 {
		t.Fatalf("clear: n=%d err=%v", len(deleted), err)
	}
	clearDur := time.Since(start)

	t.Logf("submit=%v registry=%v queries=%v clear=%v", submitDur, registryDur, queryDur, clearDur)
	if registryDur > 10*time.Second || queryDur > 5*time.Second || clearDur > 10*time.Second {
		t.Fatalf("too slow: registry=%v queries=%v clear=%v", registryDur, queryDur, clearDur)
	}
}