New external site management. New external prober heartbeat feature.

internal/db/db.go

@@ -25,6 +25,12 @@ var fipSettleDelaySchema string
 //go:embed migrations/0004_inbound_checks_admin.sql
 var inboundChecksAdminSchema string
 
+//go:embed migrations/0005_unbounded_sites.sql
+var unboundedSitesSchema string
+
+//go:embed migrations/0006_prober_heartbeat.sql
+var proberHeartbeatSchema string
+
 // migrations is the ordered list of schema versions. Each entry's SQL is
 // applied, in order, for any version greater than the database's current
 // PRAGMA user_version — so a fresh database walks the whole list and an
@@ -37,6 +43,8 @@ var migrations = []struct {
 	{2, dynamicConfigSchema},
 	{3, fipSettleDelaySchema},
 	{4, inboundChecksAdminSchema},
+	{5, unboundedSitesSchema},
+	{6, proberHeartbeatSchema},
 }
 
 type DB struct {

internal/db/migrations/0005_unbounded_sites.sql

@@ -0,0 +1,19 @@
+-- Remove the hardcoded 3-site limit: ip_queue.site{1,2,3}_complete are
+-- replaced by a child table with no cap on site_idx, keyed by
+-- attempt_number the same way `checks` already is (see docs/USAGE.md).
+
+ALTER TABLE ip_queue DROP COLUMN site1_complete;
+ALTER TABLE ip_queue DROP COLUMN site2_complete;
+ALTER TABLE ip_queue DROP COLUMN site3_complete;
+
+-- Per-IP-per-site inbound-check completion. No fixed cap on site_idx, and
+-- no reset needed on retry — a new attempt_number simply has no rows yet,
+-- unlike the old fixed columns which had to be explicitly zeroed.
+CREATE TABLE ip_site_checks (
+  ip_id          INTEGER NOT NULL REFERENCES ip_queue(id),
+  attempt_number INTEGER NOT NULL,
+  site_idx       INTEGER NOT NULL,
+  complete       BOOLEAN NOT NULL DEFAULT 0,
+  completed_at   TIMESTAMP,
+  PRIMARY KEY (ip_id, attempt_number, site_idx)
+);

internal/db/migrations/0006_prober_heartbeat.sql

@@ -0,0 +1,8 @@
+-- Prober availability indication (see docs/USAGE.md, "Управление
+-- площадками"): sites gain the same hostname/state/heartbeat tracking
+-- validators already have, so the dashboard can show whether a prober is
+-- actually connected and polling.
+
+ALTER TABLE sites ADD COLUMN hostname          TEXT NOT NULL DEFAULT '';
+ALTER TABLE sites ADD COLUMN state             TEXT NOT NULL DEFAULT 'unregistered';
+ALTER TABLE sites ADD COLUMN last_heartbeat_at TIMESTAMP;

internal/db/models.go

@@ -12,6 +12,12 @@ const (
 	ValidatorChecking     = "checking"
 	ValidatorUnreachable  = "unreachable"
 
+	// Site (prober) availability states — no assigned/checking equivalent,
+	// since a prober isn't bound to one IP at a time.
+	SiteUnregistered = "unregistered"
+	SiteIdle         = "idle"
+	SiteUnreachable  = "unreachable"
+
 	IPQueued            = "queued"
 	IPAssigningFIP      = "assigning_fip"
 	IPAwaitingSelfCheck = "awaiting_self_check"
@@ -79,9 +85,6 @@ type IPQueueItem struct {
 	RetryCount       int
 	LeaseExpiresAt   *time.Time
 	EgressComplete   bool
-	Site1Complete    bool
-	Site2Complete    bool
-	Site3Complete    bool
 	OverallResult    string
 	AssignedAt       *time.Time
 	FIPAssociatedAt  *time.Time
@@ -119,10 +122,13 @@ type Event struct {
 }
 
 type Site struct {
-	Index     int
-	SiteID    string
-	CreatedAt time.Time
-	UpdatedAt time.Time
+	Index           int
+	SiteID          string
+	Hostname        string
+	State           string
+	LastHeartbeatAt *time.Time
+	CreatedAt       time.Time
+	UpdatedAt       time.Time
 }
 
 type TargetGroup struct {

internal/db/queries_dynconfig_test.go

@@ -28,13 +28,18 @@ func TestUpsertSiteValidation(t *testing.T) {
 	if err := d.UpsertSite(ctx, 0, "site-1"); !errors.Is(err, ErrValidation) {
 		t.Fatalf("expected ErrValidation for idx=0, got %v", err)
 	}
-	if err := d.UpsertSite(ctx, 4, "site-1"); !errors.Is(err, ErrValidation) {
-		t.Fatalf("expected ErrValidation for idx=4, got %v", err)
-	}
 	if err := d.UpsertSite(ctx, 1, ""); !errors.Is(err, ErrValidation) {
 		t.Fatalf("expected ErrValidation for empty site_id, got %v", err)
 	}
 
+	// No cap on the number of slots — idx=4 and idx=100 must both succeed.
+	if err := d.UpsertSite(ctx, 4, "site-4"); err != nil {
+		t.Fatalf("upsert site 4: %v", err)
+	}
+	if err := d.UpsertSite(ctx, 100, "site-100"); err != nil {
+		t.Fatalf("upsert site 100: %v", err)
+	}
+
 	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
 		t.Fatalf("upsert site 1: %v", err)
 	}
@@ -352,6 +357,9 @@ func TestDeleteIP(t *testing.T) {
 	if err := d.InsertEvent(ctx, Event{SourceType: "control-api", IPID: &ip.ID, EventType: "test_event", OccurredAt: Now()}); err != nil {
 		t.Fatalf("insert event: %v", err)
 	}
+	if err := d.SetSiteComplete(ctx, ip.ID, 1); err != nil {
+		t.Fatalf("set site complete: %v", err)
+	}
 
 	if err := d.DeleteIP(ctx, ip.ID); err != nil {
 		t.Fatalf("delete ip: %v", err)
@@ -366,6 +374,13 @@ func TestDeleteIP(t *testing.T) {
 	if len(checks) != 0 {
 		t.Fatalf("expected checks gone, got %+v", checks)
 	}
+	completed, err := d.ListCompletedSiteIndices(ctx, ip.ID, ip.AttemptNumber)
+	if err != nil {
+		t.Fatalf("list completed site indices: %v", err)
+	}
+	if len(completed) != 0 {
+		t.Fatalf("expected ip_site_checks rows gone, got %+v", completed)
+	}
 	events, err := d.ListEventsForIP(ctx, ip.ID)
 	if err != nil {
 		t.Fatalf("list events: %v", err)

internal/db/queries_ipqueue.go

@@ -196,7 +196,7 @@ func (d *DB) RequeueOrFail(ctx context.Context, ipID int64, validatorID string,
 		_, err = tx.ExecContext(ctx, `
 			UPDATE ip_queue SET
 				state=?, owner_validator_id=NULL, fip_id='', retry_count=?, attempt_number=attempt_number+1,
-				lease_expires_at=NULL, egress_complete=0, site1_complete=0, site2_complete=0, site3_complete=0,
+				lease_expires_at=NULL, egress_complete=0,
 				overall_result='', assigned_at=NULL, fip_associated_at=NULL, updated_at=?
 			WHERE id=?
 		`, nextState, retryCount, now, ipID)
@@ -282,7 +282,7 @@ func (d *DB) SubmitIPs(ctx context.Context, addresses []string) (SubmitIPsResult
 				UPDATE ip_queue SET
 					state=?, sequence=?, owner_validator_id=NULL, fip_id='', retry_count=0,
 					attempt_number=attempt_number+1, lease_expires_at=NULL, egress_complete=0,
-					site1_complete=0, site2_complete=0, site3_complete=0, overall_result='',
+					overall_result='',
 					assigned_at=NULL, fip_associated_at=NULL, aggregated_at=NULL, fip_released_at=NULL, updated_at=?
 				WHERE ip_address=?
 			`, IPQueued, seq, now, addr); err != nil {
@@ -409,6 +409,9 @@ func deleteIPTx(ctx context.Context, tx *sql.Tx, ipID int64) error {
 	if _, err := tx.ExecContext(ctx, `DELETE FROM events WHERE ip_id=?`, ipID); err != nil {
 		return fmt.Errorf("delete events: %w", err)
 	}
+	if _, err := tx.ExecContext(ctx, `DELETE FROM ip_site_checks WHERE ip_id=?`, ipID); err != nil {
+		return fmt.Errorf("delete ip_site_checks: %w", err)
+	}
 	res, err := tx.ExecContext(ctx, `DELETE FROM ip_queue WHERE id=?`, ipID)
 	if err != nil {
 		return fmt.Errorf("delete ip_queue row: %w", err)
@@ -424,16 +427,6 @@ func (d *DB) SetEgressComplete(ctx context.Context, ipID int64) error {
 	return err
 }
 
-// SetSiteComplete marks completion for prober site 1, 2, or 3.
-func (d *DB) SetSiteComplete(ctx context.Context, ipID int64, siteIndex int) error {
-	col := map[int]string{1: "site1_complete", 2: "site2_complete", 3: "site3_complete"}[siteIndex]
-	if col == "" {
-		return fmt.Errorf("invalid site index %d", siteIndex)
-	}
-	_, err := d.ExecContext(ctx, fmt.Sprintf(`UPDATE ip_queue SET %s=1, updated_at=? WHERE id=?`, col), timeToDB(Now()), ipID)
-	return err
-}
-
 func (d *DB) GetIP(ctx context.Context, ipID int64) (*IPQueueItem, error) {
 	row := d.QueryRowContext(ctx, ipQueueSelect+`WHERE id=?`, ipID)
 	return scanIPQueueItem(row)
@@ -479,7 +472,7 @@ func (d *DB) ListExpiredLeases(ctx context.Context, now time.Time) ([]IPQueueIte
 
 const ipQueueSelect = `
 	SELECT id, ip_address, sequence, state, owner_validator_id, fip_id, attempt_number, retry_count,
-		lease_expires_at, egress_complete, site1_complete, site2_complete, site3_complete, overall_result,
+		lease_expires_at, egress_complete, overall_result,
 		assigned_at, fip_associated_at, aggregated_at, fip_released_at, created_at, updated_at
 	FROM ip_queue
 `
@@ -504,7 +497,7 @@ func scanIPQueueItem(row rowScanner) (*IPQueueItem, error) {
 	if err := row.Scan(
 		&item.ID, &item.IPAddress, &item.Sequence, &item.State, &owner, &item.FIPID,
 		&item.AttemptNumber, &item.RetryCount, &leaseExpires,
-		&item.EgressComplete, &item.Site1Complete, &item.Site2Complete, &item.Site3Complete,
+		&item.EgressComplete,
 		&item.OverallResult, &assignedAt, &fipAssociatedAt, &aggregatedAt, &fipReleasedAt, &createdAt, &updatedAt,
 	); err != nil {
 		return nil, err

internal/db/queries_ipsitechecks.go

@@ -0,0 +1,39 @@
+package db
+
+import "context"
+
+// SetSiteComplete marks a given site's inbound checks complete for an IP's
+// *current* attempt. attempt_number is resolved via subquery from
+// ip_queue rather than taken as a parameter, so the call signature stays
+// the same as before the fixed-3-columns design was replaced by this
+// table. No cap on siteIndex.
+func (d *DB) SetSiteComplete(ctx context.Context, ipID int64, siteIndex int) error {
+	_, err := d.ExecContext(ctx, `
+		INSERT INTO ip_site_checks (ip_id, attempt_number, site_idx, complete, completed_at)
+		SELECT ?, attempt_number, ?, 1, ? FROM ip_queue WHERE id=?
+		ON CONFLICT(ip_id, attempt_number, site_idx) DO UPDATE SET complete=1, completed_at=excluded.completed_at
+	`, ipID, siteIndex, timeToDB(Now()), ipID)
+	return err
+}
+
+// ListCompletedSiteIndices returns the set of site indices that have
+// reported completion for the given IP's attempt.
+func (d *DB) ListCompletedSiteIndices(ctx context.Context, ipID int64, attemptNumber int) (map[int]bool, error) {
+	rows, err := d.QueryContext(ctx, `
+		SELECT site_idx FROM ip_site_checks WHERE ip_id=? AND attempt_number=? AND complete=1
+	`, ipID, attemptNumber)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	out := map[int]bool{}
+	for rows.Next() {
+		var idx int
+		if err := rows.Scan(&idx); err != nil {
+			return nil, err
+		}
+		out[idx] = true
+	}
+	return out, rows.Err()
+}

internal/db/queries_ipsitechecks_test.go

@@ -0,0 +1,92 @@
+package db
+
+import "testing"
+
+func TestSetSiteCompleteAndListCompletedSiteIndices(t *testing.T) {
+	d, ctx := newTestDB(t)
+
+	if err := d.SeedQueue(ctx, []string{"1.2.3.4"}); err != nil {
+		t.Fatalf("seed queue: %v", err)
+	}
+	ip, err := d.GetIPByAddress(ctx, "1.2.3.4")
+	if err != nil {
+		t.Fatalf("get ip: %v", err)
+	}
+
+	for _, idx := range []int{1, 2, 4} {
+		if err := d.SetSiteComplete(ctx, ip.ID, idx); err != nil {
+			t.Fatalf("set site %d complete: %v", idx, err)
+		}
+	}
+
+	completed, err := d.ListCompletedSiteIndices(ctx, ip.ID, ip.AttemptNumber)
+	if err != nil {
+		t.Fatalf("list completed site indices: %v", err)
+	}
+	for _, idx := range []int{1, 2, 4} {
+		if !completed[idx] {
+			t.Fatalf("expected site %d complete, got %+v", idx, completed)
+		}
+	}
+	if completed[3] {
+		t.Fatalf("expected site 3 not complete, got %+v", completed)
+	}
+
+	// Idempotent re-mark doesn't error or duplicate.
+	if err := d.SetSiteComplete(ctx, ip.ID, 1); err != nil {
+		t.Fatalf("re-mark site 1 complete: %v", err)
+	}
+}
+
+// TestNewAttemptStartsWithNoCompletedSites confirms a bumped
+// attempt_number (via SubmitIPs requeue) starts with an empty completed
+// set even though the previous attempt had rows — proving no explicit
+// reset is needed on retry, unlike the old fixed-3-column design.
+func TestNewAttemptStartsWithNoCompletedSites(t *testing.T) {
+	d, ctx := newTestDB(t)
+
+	if err := d.SeedQueue(ctx, []string{"1.2.3.4"}); err != nil {
+		t.Fatalf("seed queue: %v", err)
+	}
+	ip, err := d.GetIPByAddress(ctx, "1.2.3.4")
+	if err != nil {
+		t.Fatalf("get ip: %v", err)
+	}
+	if err := d.SetSiteComplete(ctx, ip.ID, 1); err != nil {
+		t.Fatalf("set site complete: %v", err)
+	}
+
+	// Force the IP to a terminal state, then requeue via SubmitIPs — the
+	// same path an admin-triggered recheck takes — to bump attempt_number.
+	if _, err := d.ExecContext(ctx, `UPDATE ip_queue SET state=? WHERE id=?`, IPFailed, ip.ID); err != nil {
+		t.Fatalf("force failed: %v", err)
+	}
+	if _, err := d.SubmitIPs(ctx, []string{"1.2.3.4"}); err != nil {
+		t.Fatalf("resubmit: %v", err)
+	}
+
+	requeued, err := d.GetIP(ctx, ip.ID)
+	if err != nil {
+		t.Fatalf("get ip after requeue: %v", err)
+	}
+	if requeued.AttemptNumber != ip.AttemptNumber+1 {
+		t.Fatalf("expected attempt_number bumped, got %d", requeued.AttemptNumber)
+	}
+
+	completed, err := d.ListCompletedSiteIndices(ctx, ip.ID, requeued.AttemptNumber)
+	if err != nil {
+		t.Fatalf("list completed site indices for new attempt: %v", err)
+	}
+	if len(completed) != 0 {
+		t.Fatalf("expected no completed sites for the new attempt, got %+v", completed)
+	}
+
+	// The old attempt's row is still there as history, not wiped.
+	old, err := d.ListCompletedSiteIndices(ctx, ip.ID, ip.AttemptNumber)
+	if err != nil {
+		t.Fatalf("list completed site indices for old attempt: %v", err)
+	}
+	if !old[1] {
+		t.Fatalf("expected old attempt's completion to remain as history, got %+v", old)
+	}
+}

internal/db/queries_sites.go

@@ -4,12 +4,13 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"time"
 )
 
 // ListSites returns all configured prober sites, ordered by slot index.
 func (d *DB) ListSites(ctx context.Context) ([]Site, error) {
 	rows, err := d.QueryContext(ctx, `
-		SELECT idx, site_id, created_at, updated_at FROM sites ORDER BY idx
+		SELECT idx, site_id, hostname, state, last_heartbeat_at, created_at, updated_at FROM sites ORDER BY idx
 	`)
 	if err != nil {
 		return nil, err
@@ -18,24 +19,37 @@ func (d *DB) ListSites(ctx context.Context) ([]Site, error) {
 
 	var out []Site
 	for rows.Next() {
-		var s Site
-		var createdAt, updatedAt string
-		if err := rows.Scan(&s.Index, &s.SiteID, &createdAt, &updatedAt); err != nil {
+		s, err := scanSite(rows)
+		if err != nil {
 			return nil, err
 		}
-		var err error
-		if s.CreatedAt, err = dbToTime(createdAt); err != nil {
-			return nil, err
-		}
-		if s.UpdatedAt, err = dbToTime(updatedAt); err != nil {
-			return nil, err
-		}
-		out = append(out, s)
+		out = append(out, *s)
 	}
 	return out, rows.Err()
 }
 
-// GetSiteIndex resolves a configured site_id to its 1/2/3 slot index. It
+func scanSite(row rowScanner) (*Site, error) {
+	var s Site
+	var lastHeartbeat sql.NullString
+	var createdAt, updatedAt string
+	if err := row.Scan(&s.Index, &s.SiteID, &s.Hostname, &s.State, &lastHeartbeat, &createdAt, &updatedAt); err != nil {
+		return nil, err
+	}
+	hb, err := nullStringToTimePtr(lastHeartbeat)
+	if err != nil {
+		return nil, err
+	}
+	s.LastHeartbeatAt = hb
+	if s.CreatedAt, err = dbToTime(createdAt); err != nil {
+		return nil, err
+	}
+	if s.UpdatedAt, err = dbToTime(updatedAt); err != nil {
+		return nil, err
+	}
+	return &s, nil
+}
+
+// GetSiteIndex resolves a configured site_id to its slot index. It
 // returns (0, nil) — not an error — if no site with that ID is configured,
 // matching the "0 means unconfigured" convention used throughout the
 // prober-facing handlers.
@@ -52,12 +66,21 @@ func (d *DB) GetSiteIndex(ctx context.Context, siteID string) (int, error) {
 }
 
 // UpsertSite assigns (or renames) the site occupying the given slot. idx
-// must be 1, 2, or 3 — the schema hard-caps the number of prober slots at
-// three (see ip_queue.site{1,2,3}_complete). site_id must be unique across
+// has no upper bound — an admin can configure as many prober sites as they
+// need — only idx >= 1 is enforced (slot numbering starts at 1, matching
+// the existing admin-facing convention). site_id must be unique across
 // slots.
+//
+// Re-pointing an existing slot to a (possibly different) site_id always
+// resets hostname/state/last_heartbeat_at to their "never connected"
+// defaults: a new site_id is a new prober identity, and even a same-value
+// PUT is treated as idempotent-intent rather than a heartbeat, so the old
+// prober's connection status must not silently linger. The prober
+// re-establishes state on its next register/heartbeat within
+// poll_interval_seconds.
 func (d *DB) UpsertSite(ctx context.Context, idx int, siteID string) error {
-	if idx < 1 || idx > 3 {
-		return fmt.Errorf("site index must be 1, 2, or 3, got %d: %w", idx, ErrValidation)
+	if idx < 1 {
+		return fmt.Errorf("site index must be >= 1, got %d: %w", idx, ErrValidation)
 	}
 	if siteID == "" {
 		return fmt.Errorf("site_id must not be empty: %w", ErrValidation)
@@ -76,8 +99,10 @@ func (d *DB) UpsertSite(ctx context.Context, idx int, siteID string) error {
 	_, err = d.ExecContext(ctx, `
 		INSERT INTO sites (idx, site_id, created_at, updated_at)
 		VALUES (?, ?, ?, ?)
-		ON CONFLICT(idx) DO UPDATE SET site_id=excluded.site_id, updated_at=excluded.updated_at
-	`, idx, siteID, now, now)
+		ON CONFLICT(idx) DO UPDATE SET
+			site_id=excluded.site_id, updated_at=excluded.updated_at,
+			hostname='', state=?, last_heartbeat_at=NULL
+	`, idx, siteID, now, now, SiteUnregistered)
 	if err != nil {
 		return fmt.Errorf("upsert site: %w", err)
 	}
@@ -95,3 +120,75 @@ func (d *DB) DeleteSite(ctx context.Context, idx int) error {
 	}
 	return nil
 }
+
+// RegisterSiteProber records that a prober has (re)connected for this
+// site: sets hostname, stamps last_heartbeat_at, and reactivates state to
+// idle if it was unregistered or unreachable — mirrors RegisterValidator's
+// reactivation UPDATE. Called once at prober startup. Returns
+// sql.ErrNoRows if site_id is unconfigured.
+func (d *DB) RegisterSiteProber(ctx context.Context, siteID, hostname string) error {
+	now := timeToDB(Now())
+	res, err := d.ExecContext(ctx, `
+		UPDATE sites SET hostname=?, last_heartbeat_at=?, updated_at=?,
+			state = CASE WHEN state IN (?, ?) THEN ? ELSE state END
+		WHERE site_id=?
+	`, hostname, now, now, SiteUnregistered, SiteUnreachable, SiteIdle, siteID)
+	if err != nil {
+		return fmt.Errorf("register site prober: %w", err)
+	}
+	n, _ := res.RowsAffected()
+	if n == 0 {
+		return sql.ErrNoRows
+	}
+	return nil
+}
+
+// SiteHeartbeat is the periodic "I'm alive" touch — mirrors Heartbeat for
+// validators exactly: stamps last_heartbeat_at and flips state from
+// unreachable back to idle (does not touch hostname).
+func (d *DB) SiteHeartbeat(ctx context.Context, siteID string) error {
+	now := timeToDB(Now())
+	res, err := d.ExecContext(ctx, `
+		UPDATE sites SET last_heartbeat_at=?, updated_at=?,
+			state = CASE WHEN state=? THEN ? ELSE state END
+		WHERE site_id=?
+	`, now, now, SiteUnreachable, SiteIdle, siteID)
+	if err != nil {
+		return fmt.Errorf("site heartbeat: %w", err)
+	}
+	n, _ := res.RowsAffected()
+	if n == 0 {
+		return sql.ErrNoRows
+	}
+	return nil
+}
+
+// ListStaleSiteHeartbeats returns sites whose last heartbeat predates the
+// given cutoff and that aren't already marked unreachable.
+func (d *DB) ListStaleSiteHeartbeats(ctx context.Context, cutoff time.Time) ([]Site, error) {
+	rows, err := d.QueryContext(ctx, `
+		SELECT idx, site_id, hostname, state, last_heartbeat_at, created_at, updated_at
+		FROM sites
+		WHERE state != ? AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ?
+	`, SiteUnreachable, timeToDB(cutoff))
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var out []Site
+	for rows.Next() {
+		s, err := scanSite(rows)
+		if err != nil {
+			return nil, err
+		}
+		out = append(out, *s)
+	}
+	return out, rows.Err()
+}
+
+func (d *DB) MarkSiteUnreachable(ctx context.Context, siteID string) error {
+	_, err := d.ExecContext(ctx, `UPDATE sites SET state=?, updated_at=? WHERE site_id=?`,
+		SiteUnreachable, timeToDB(Now()), siteID)
+	return err
+}

internal/db/queries_sites_heartbeat_test.go

@@ -0,0 +1,172 @@
+package db
+
+import (
+	"database/sql"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestRegisterSiteProberUnknownSiteReturnsErrNoRows(t *testing.T) {
+	d, ctx := newTestDB(t)
+	if err := d.RegisterSiteProber(ctx, "unknown-site", "host-1"); !errors.Is(err, sql.ErrNoRows) {
+		t.Fatalf("expected sql.ErrNoRows for unknown site, got %v", err)
+	}
+}
+
+func TestRegisterSiteProberSetsHostnameAndIdleState(t *testing.T) {
+	d, ctx := newTestDB(t)
+	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
+		t.Fatalf("upsert site: %v", err)
+	}
+
+	sites, err := d.ListSites(ctx)
+	if err != nil {
+		t.Fatalf("list sites: %v", err)
+	}
+	if sites[0].State != SiteUnregistered {
+		t.Fatalf("expected freshly-created slot unregistered, got %s", sites[0].State)
+	}
+
+	if err := d.RegisterSiteProber(ctx, "site-1", "probe-host-1"); err != nil {
+		t.Fatalf("register site prober: %v", err)
+	}
+	sites, err = d.ListSites(ctx)
+	if err != nil {
+		t.Fatalf("list sites after register: %v", err)
+	}
+	if sites[0].State != SiteIdle {
+		t.Fatalf("expected idle after register, got %s", sites[0].State)
+	}
+	if sites[0].Hostname != "probe-host-1" {
+		t.Fatalf("expected hostname set, got %q", sites[0].Hostname)
+	}
+	if sites[0].LastHeartbeatAt == nil {
+		t.Fatalf("expected last_heartbeat_at stamped")
+	}
+}
+
+func TestRegisterSiteProberReactivatesFromUnreachable(t *testing.T) {
+	d, ctx := newTestDB(t)
+	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
+		t.Fatalf("upsert site: %v", err)
+	}
+	if err := d.MarkSiteUnreachable(ctx, "site-1"); err != nil {
+		t.Fatalf("mark unreachable: %v", err)
+	}
+	if err := d.RegisterSiteProber(ctx, "site-1", "probe-host-1"); err != nil {
+		t.Fatalf("register site prober: %v", err)
+	}
+	sites, err := d.ListSites(ctx)
+	if err != nil {
+		t.Fatalf("list sites: %v", err)
+	}
+	if sites[0].State != SiteIdle {
+		t.Fatalf("expected reactivated to idle, got %s", sites[0].State)
+	}
+}
+
+func TestSiteHeartbeatRoundTrip(t *testing.T) {
+	d, ctx := newTestDB(t)
+	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
+		t.Fatalf("upsert site: %v", err)
+	}
+	if err := d.SiteHeartbeat(ctx, "unknown-site"); !errors.Is(err, sql.ErrNoRows) {
+		t.Fatalf("expected sql.ErrNoRows for unknown site, got %v", err)
+	}
+
+	if err := d.SiteHeartbeat(ctx, "site-1"); err != nil {
+		t.Fatalf("heartbeat: %v", err)
+	}
+	sites, err := d.ListSites(ctx)
+	if err != nil {
+		t.Fatalf("list sites: %v", err)
+	}
+	// A plain heartbeat (no prior register) doesn't flip an unregistered
+	// site to idle — only reactivates from unreachable, mirroring
+	// validators' Heartbeat.
+	if sites[0].State != SiteUnregistered {
+		t.Fatalf("expected still unregistered after bare heartbeat, got %s", sites[0].State)
+	}
+	if sites[0].LastHeartbeatAt == nil {
+		t.Fatalf("expected last_heartbeat_at stamped")
+	}
+	if sites[0].Hostname != "" {
+		t.Fatalf("expected heartbeat not to set hostname, got %q", sites[0].Hostname)
+	}
+
+	if err := d.MarkSiteUnreachable(ctx, "site-1"); err != nil {
+		t.Fatalf("mark unreachable: %v", err)
+	}
+	if err := d.SiteHeartbeat(ctx, "site-1"); err != nil {
+		t.Fatalf("heartbeat after unreachable: %v", err)
+	}
+	sites, err = d.ListSites(ctx)
+	if err != nil {
+		t.Fatalf("list sites: %v", err)
+	}
+	if sites[0].State != SiteIdle {
+		t.Fatalf("expected heartbeat to reactivate from unreachable to idle, got %s", sites[0].State)
+	}
+}
+
+func TestListStaleSiteHeartbeatsAndMarkUnreachable(t *testing.T) {
+	d, ctx := newTestDB(t)
+	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
+		t.Fatalf("upsert site: %v", err)
+	}
+	if err := d.UpsertSite(ctx, 2, "site-2"); err != nil {
+		t.Fatalf("upsert site: %v", err)
+	}
+	// site-3 never heartbeats at all — must not show up as "stale" (it's
+	// simply unregistered, which is a distinct, expected state).
+	if err := d.UpsertSite(ctx, 3, "site-3"); err != nil {
+		t.Fatalf("upsert site: %v", err)
+	}
+
+	if err := d.SiteHeartbeat(ctx, "site-1"); err != nil {
+		t.Fatalf("heartbeat site-1: %v", err)
+	}
+	if err := d.SiteHeartbeat(ctx, "site-2"); err != nil {
+		t.Fatalf("heartbeat site-2: %v", err)
+	}
+
+	cutoff := Now().Add(1 * time.Hour) // everything heartbeated is "before" this
+	stale, err := d.ListStaleSiteHeartbeats(ctx, cutoff)
+	if err != nil {
+		t.Fatalf("list stale site heartbeats: %v", err)
+	}
+	if len(stale) != 2 {
+		t.Fatalf("expected 2 stale sites (site-1, site-2), got %+v", stale)
+	}
+
+	for _, s := range stale {
+		if err := d.MarkSiteUnreachable(ctx, s.SiteID); err != nil {
+			t.Fatalf("mark unreachable %s: %v", s.SiteID, err)
+		}
+	}
+
+	// Now already-unreachable sites are excluded from a subsequent sweep.
+	stale, err = d.ListStaleSiteHeartbeats(ctx, cutoff)
+	if err != nil {
+		t.Fatalf("list stale site heartbeats after mark: %v", err)
+	}
+	if len(stale) != 0 {
+		t.Fatalf("expected no stale sites left, got %+v", stale)
+	}
+
+	sites, err := d.ListSites(ctx)
+	if err != nil {
+		t.Fatalf("list sites: %v", err)
+	}
+	byID := map[string]Site{}
+	for _, s := range sites {
+		byID[s.SiteID] = s
+	}
+	if byID["site-1"].State != SiteUnreachable || byID["site-2"].State != SiteUnreachable {
+		t.Fatalf("expected site-1/site-2 unreachable, got %+v", sites)
+	}
+	if byID["site-3"].State != SiteUnregistered {
+		t.Fatalf("expected site-3 still unregistered (never heartbeated), got %s", byID["site-3"].State)
+	}
+}