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

package db

import (
	"context"
	"database/sql"
	"fmt"
	"time"
)

// SeedQueue inserts the configured IP address list in order, assigning each
// a stable sequence number. Re-running with the same list is a no-op for
// addresses already present (ON CONFLICT DO NOTHING keyed by the UNIQUE
// ip_address column), so restarting control-api against the same config
// never re-queues already-processed addresses.
func (d *DB) SeedQueue(ctx context.Context, addresses []string) error {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	now := timeToDB(Now())
	for i, addr := range addresses {
		_, err := tx.ExecContext(ctx, `
			INSERT INTO ip_queue (ip_address, sequence, state, created_at, updated_at)
			VALUES (?, ?, ?, ?, ?)
			ON CONFLICT(ip_address) DO NOTHING
		`, addr, i, IPQueued, now, now)
		if err != nil {
			return fmt.Errorf("seed %s: %w", addr, err)
		}
	}
	return tx.Commit()
}

// ClaimNextQueued atomically hands the next queued IP (lowest sequence) to
// the given idle validator. It returns (nil, nil) if the validator isn't
// idle or no IP is queued. The DB connection pool is capped at one physical
// connection (see Open), so this transaction already has exclusive access
// to the database for its duration — no other claim, requeue, or update can
// interleave — which combined with the conditional UPDATEs (checked via
// RowsAffected) guarantees a single IP is never claimed by two validators.
func (d *DB) ClaimNextQueued(ctx context.Context, validatorID string, leaseTTL time.Duration) (*IPQueueItem, error) {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	var state string
	err = tx.QueryRowContext(ctx, `SELECT state FROM validators WHERE validator_id=?`, validatorID).Scan(&state)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if state != ValidatorIdle {
		return nil, nil
	}

	var item IPQueueItem
	err = tx.QueryRowContext(ctx, `
		SELECT id, ip_address, sequence, attempt_number, retry_count
		FROM ip_queue WHERE state=? ORDER BY sequence LIMIT 1
	`, IPQueued).Scan(&item.ID, &item.IPAddress, &item.Sequence, &item.AttemptNumber, &item.RetryCount)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	now := Now()
	lease := now.Add(leaseTTL)
	res, err := tx.ExecContext(ctx, `
		UPDATE ip_queue SET state=?, owner_validator_id=?, assigned_at=?, lease_expires_at=?, updated_at=?
		WHERE id=? AND state=?
	`, IPAssigningFIP, validatorID, timeToDB(now), timeToDB(lease), timeToDB(now), item.ID, IPQueued)
	if err != nil {
		return nil, err
	}
	if n, _ := res.RowsAffected(); n != 1 {
		return nil, nil
	}

	res, err = tx.ExecContext(ctx, `
		UPDATE validators SET state=?, current_ip_id=?, updated_at=?
		WHERE validator_id=? AND state=?
	`, ValidatorAssigned, item.ID, timeToDB(now), validatorID, ValidatorIdle)
	if err != nil {
		return nil, err
	}
	if n, _ := res.RowsAffected(); n != 1 {
		return nil, nil
	}

	if err := tx.Commit(); err != nil {
		return nil, err
	}

	item.State = IPAssigningFIP
	ownerID := validatorID
	item.OwnerValidatorID = &ownerID
	item.AssignedAt = &now
	item.LeaseExpiresAt = &lease
	return &item, nil
}

func (d *DB) SetFIPAssociated(ctx context.Context, ipID int64, fipID string, leaseTTL time.Duration) error {
	now := Now()
	_, err := d.ExecContext(ctx, `
		UPDATE ip_queue SET state=?, fip_id=?, lease_expires_at=?, updated_at=?
		WHERE id=?
	`, IPAwaitingSelfCheck, fipID, timeToDB(now.Add(leaseTTL)), timeToDB(now), ipID)
	return err
}

func (d *DB) SetChecking(ctx context.Context, ipID int64, leaseTTL time.Duration) error {
	now := Now()
	_, err := d.ExecContext(ctx, `
		UPDATE ip_queue SET state=?, lease_expires_at=?, updated_at=?
		WHERE id=?
	`, IPChecking, timeToDB(now.Add(leaseTTL)), timeToDB(now), ipID)
	return err
}

func (d *DB) SetAggregating(ctx context.Context, ipID int64) error {
	_, err := d.ExecContext(ctx, `UPDATE ip_queue SET state=?, updated_at=? WHERE id=?`,
		IPAggregating, timeToDB(Now()), ipID)
	return err
}

// FinishIP records the aggregated result and marks the IP done or failed.
func (d *DB) FinishIP(ctx context.Context, ipID int64, result string) error {
	state := IPDone
	if result == ResultFail {
		state = IPFailed
	}
	now := timeToDB(Now())
	_, err := d.ExecContext(ctx, `
		UPDATE ip_queue SET state=?, overall_result=?, aggregated_at=?, updated_at=?
		WHERE id=?
	`, state, result, now, now, ipID)
	return err
}

// ReleaseFIP records that the floating IP has been disassociated and frees
// the owning validator back to idle, in one transaction.
func (d *DB) ReleaseFIP(ctx context.Context, ipID int64, validatorID string) error {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	now := timeToDB(Now())
	if _, err := tx.ExecContext(ctx, `UPDATE ip_queue SET fip_released_at=?, updated_at=? WHERE id=?`, now, now, ipID); err != nil {
		return err
	}
	if _, err := tx.ExecContext(ctx, `
		UPDATE validators SET state=?, current_ip_id=NULL, updated_at=?
		WHERE validator_id=?
	`, ValidatorIdle, now, validatorID); err != nil {
		return err
	}
	return tx.Commit()
}

// RequeueOrFail is used by both the retry path (association/self-check
// failure) and the lease-sweep reclaim path. It clears ownership and
// per-attempt progress, bumps attempt_number and retry_count, and either
// sends the IP back to the queue or marks it permanently failed once
// maxRetries is exceeded. The owning validator (if any) is freed in the
// same transaction.
func (d *DB) RequeueOrFail(ctx context.Context, ipID int64, validatorID string, maxRetries int) error {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	var retryCount int
	if err := tx.QueryRowContext(ctx, `SELECT retry_count FROM ip_queue WHERE id=?`, ipID).Scan(&retryCount); err != nil {
		return err
	}
	retryCount++

	now := timeToDB(Now())
	nextState := IPQueued
	if retryCount > maxRetries {
		nextState = IPFailed
	}

	if nextState == IPQueued {
		_, 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,
				overall_result='', assigned_at=NULL, updated_at=?
			WHERE id=?
		`, nextState, retryCount, now, ipID)
	} else {
		_, err = tx.ExecContext(ctx, `
			UPDATE ip_queue SET
				state=?, retry_count=?, overall_result=?, aggregated_at=?, updated_at=?
			WHERE id=?
		`, nextState, retryCount, ResultFail, now, now, ipID)
	}
	if err != nil {
		return err
	}

	if validatorID != "" {
		if _, err := tx.ExecContext(ctx, `
			UPDATE validators SET state=?, current_ip_id=NULL, updated_at=?
			WHERE validator_id=?
		`, ValidatorIdle, now, validatorID); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (d *DB) SetEgressComplete(ctx context.Context, ipID int64) error {
	_, err := d.ExecContext(ctx, `UPDATE ip_queue SET egress_complete=1, updated_at=? WHERE id=?`, timeToDB(Now()), ipID)
	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)
}

func (d *DB) GetIPByAddress(ctx context.Context, address string) (*IPQueueItem, error) {
	row := d.QueryRowContext(ctx, ipQueueSelect+`WHERE ip_address=?`, address)
	return scanIPQueueItem(row)
}

func (d *DB) ListIPs(ctx context.Context) ([]IPQueueItem, error) {
	rows, err := d.QueryContext(ctx, ipQueueSelect+`ORDER BY sequence`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanIPQueueItems(rows)
}

// ListChecking returns all IPs currently in the checking state — the set a
// prober should be actively probing.
func (d *DB) ListChecking(ctx context.Context) ([]IPQueueItem, error) {
	rows, err := d.QueryContext(ctx, ipQueueSelect+`WHERE state=? ORDER BY sequence`, IPChecking)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanIPQueueItems(rows)
}

// ListReadyToAggregate returns checking-state IPs where every source has
// reported completion, or whose checking window has expired.
func (d *DB) ListReadyToAggregate(ctx context.Context, windowDeadline time.Time) ([]IPQueueItem, error) {
	rows, err := d.QueryContext(ctx, ipQueueSelect+`
		WHERE state=? AND (
			(egress_complete=1 AND site1_complete=1 AND site2_complete=1 AND site3_complete=1)
			OR assigned_at < ?
		)`, IPChecking, timeToDB(windowDeadline))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanIPQueueItems(rows)
}

// ListExpiredLeases returns non-terminal IPs whose lease has expired —
// candidates for the lease sweep (crash recovery + stuck-validator reclaim).
func (d *DB) ListExpiredLeases(ctx context.Context, now time.Time) ([]IPQueueItem, error) {
	rows, err := d.QueryContext(ctx, ipQueueSelect+`
		WHERE state NOT IN (?, ?) AND lease_expires_at IS NOT NULL AND lease_expires_at < ?
	`, IPDone, IPFailed, timeToDB(now))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanIPQueueItems(rows)
}

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,
		assigned_at, aggregated_at, fip_released_at, created_at, updated_at
	FROM ip_queue
`

func scanIPQueueItems(rows *sql.Rows) ([]IPQueueItem, error) {
	var out []IPQueueItem
	for rows.Next() {
		item, err := scanIPQueueItem(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *item)
	}
	return out, rows.Err()
}

func scanIPQueueItem(row rowScanner) (*IPQueueItem, error) {
	var item IPQueueItem
	var owner sql.NullString
	var leaseExpires, assignedAt, aggregatedAt, fipReleasedAt sql.NullString
	var createdAt, updatedAt string
	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.OverallResult, &assignedAt, &aggregatedAt, &fipReleasedAt, &createdAt, &updatedAt,
	); err != nil {
		return nil, err
	}
	if owner.Valid {
		item.OwnerValidatorID = &owner.String
	}
	var err error
	if item.LeaseExpiresAt, err = nullStringToTimePtr(leaseExpires); err != nil {
		return nil, err
	}
	if item.AssignedAt, err = nullStringToTimePtr(assignedAt); err != nil {
		return nil, err
	}
	if item.AggregatedAt, err = nullStringToTimePtr(aggregatedAt); err != nil {
		return nil, err
	}
	if item.FIPReleasedAt, err = nullStringToTimePtr(fipReleasedAt); err != nil {
		return nil, err
	}
	if item.CreatedAt, err = dbToTime(createdAt); err != nil {
		return nil, err
	}
	if item.UpdatedAt, err = dbToTime(updatedAt); err != nil {
		return nil, err
	}
	return &item, nil
}