package db import ( "context" "database/sql" "fmt" "time" ) // RegisterValidator inserts a new validator or updates an existing one's // hostname/agent_version on re-registration. It never overwrites the state // of a validator that's mid-assignment, so an agent restarting while it // owns an IP doesn't silently lose that ownership. func (d *DB) RegisterValidator(ctx context.Context, validatorID, hostname, osPortID, agentVersion string) error { now := timeToDB(Now()) _, err := d.ExecContext(ctx, ` INSERT INTO validators (validator_id, hostname, os_port_id, agent_version, state, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(validator_id) DO UPDATE SET hostname=excluded.hostname, os_port_id=excluded.os_port_id, agent_version=excluded.agent_version, updated_at=excluded.updated_at `, validatorID, hostname, osPortID, agentVersion, ValidatorIdle, now, now) if err != nil { return fmt.Errorf("register validator: %w", err) } // A brand-new row already lands in ValidatorIdle via the INSERT branch; // a re-registering validator that was 'unregistered' or 'unreachable' // comes back: to idle when it holds no address, to assigned when it still // does (its current_ip_id is kept, so it must not be handed another). _, err = d.ExecContext(ctx, ` UPDATE validators SET state = CASE WHEN current_ip_id IS NULL THEN ? ELSE ? END, updated_at=? WHERE validator_id=? AND state IN (?, ?) `, ValidatorIdle, ValidatorAssigned, now, validatorID, ValidatorUnregistered, ValidatorUnreachable) if err != nil { return fmt.Errorf("register validator (reactivate): %w", err) } return nil } // Heartbeat records a sign of life. A validator that was marked unreachable // returns to idle if it holds no address, but to assigned if it still does: // it was only silent (for example busy with slow checks), and handing it a // second address while it works on the first would leave the second one // without an owner that can ever pick it up. func (d *DB) Heartbeat(ctx context.Context, validatorID string) error { now := timeToDB(Now()) res, err := d.ExecContext(ctx, ` UPDATE validators SET last_heartbeat_at=?, updated_at=?, state = CASE WHEN state=? THEN (CASE WHEN current_ip_id IS NULL THEN ? ELSE ? END) ELSE state END WHERE validator_id=? `, now, now, ValidatorUnreachable, ValidatorIdle, ValidatorAssigned, validatorID) if err != nil { return fmt.Errorf("heartbeat: %w", err) } n, _ := res.RowsAffected() if n == 0 { return sql.ErrNoRows } return nil } func (d *DB) GetValidator(ctx context.Context, validatorID string) (*Validator, error) { row := d.QueryRowContext(ctx, ` SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version, last_heartbeat_at, created_at, updated_at FROM validators WHERE validator_id=? `, validatorID) return scanValidator(row) } func (d *DB) ListValidators(ctx context.Context) ([]Validator, error) { rows, err := d.QueryContext(ctx, ` SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version, last_heartbeat_at, created_at, updated_at FROM validators ORDER BY validator_id `) if err != nil { return nil, err } defer rows.Close() var out []Validator for rows.Next() { v, err := scanValidator(rows) if err != nil { return nil, err } out = append(out, *v) } return out, rows.Err() } // ListIdleValidators returns validators currently eligible to be handed a // new IP to work on. func (d *DB) ListIdleValidators(ctx context.Context) ([]Validator, error) { rows, err := d.QueryContext(ctx, ` SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version, last_heartbeat_at, created_at, updated_at FROM validators WHERE state=? ORDER BY validator_id `, ValidatorIdle) if err != nil { return nil, err } defer rows.Close() var out []Validator for rows.Next() { v, err := scanValidator(rows) if err != nil { return nil, err } out = append(out, *v) } return out, rows.Err() } // ListStaleHeartbeats returns validators whose last heartbeat predates the // given cutoff and that aren't already marked unreachable. func (d *DB) ListStaleHeartbeats(ctx context.Context, cutoff time.Time) ([]Validator, error) { rows, err := d.QueryContext(ctx, ` SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version, last_heartbeat_at, created_at, updated_at FROM validators WHERE state != ? AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ? `, ValidatorUnreachable, timeToDB(cutoff)) if err != nil { return nil, err } defer rows.Close() var out []Validator for rows.Next() { v, err := scanValidator(rows) if err != nil { return nil, err } out = append(out, *v) } return out, rows.Err() } func (d *DB) MarkValidatorUnreachable(ctx context.Context, validatorID string) error { _, err := d.ExecContext(ctx, `UPDATE validators SET state=?, updated_at=? WHERE validator_id=?`, ValidatorUnreachable, timeToDB(Now()), validatorID) return err } // freeValidatorSQL releases a validator from the address it holds. It only // applies when the validator's current address is the one being released // (args: now, validator id, ip id): a late release of an old address must not // free a validator that has already moved on to another one. An unreachable // validator stays unreachable until its next heartbeat, so a dead validator // is not handed new addresses just because its lease was reclaimed. var freeValidatorSQL = fmt.Sprintf(` UPDATE validators SET current_ip_id=NULL, state = CASE WHEN state='%s' THEN state ELSE '%s' END, updated_at=? WHERE validator_id=? AND current_ip_id=?`, ValidatorUnreachable, ValidatorIdle) // FreeValidator releases a validator from the given address (see // freeValidatorSQL). Used after an IP finishes (success or failure) or is // reclaimed by the lease sweep. func (d *DB) FreeValidator(ctx context.Context, validatorID string, ipID int64) error { _, err := d.ExecContext(ctx, freeValidatorSQL, timeToDB(Now()), validatorID, ipID) return err } // ReconcileValidators repairs validators whose state disagrees with the // queue: a validator pointing at an address that no longer exists, is // finished, or belongs to another validator is released; an "assigned" // validator that holds nothing goes back to idle. The invariants normally // hold by construction (every change is one transaction); this heals what a // crash or an older version left behind. It returns the number of validators // repaired. func (d *DB) ReconcileValidators(ctx context.Context) (int64, error) { now := timeToDB(Now()) res, err := d.ExecContext(ctx, fmt.Sprintf(` UPDATE validators SET current_ip_id=NULL, state = CASE WHEN state='%s' THEN state ELSE '%s' END, updated_at=? WHERE current_ip_id IS NOT NULL AND NOT EXISTS ( SELECT 1 FROM ip_queue q WHERE q.id = validators.current_ip_id AND q.owner_validator_id = validators.validator_id AND q.state NOT IN ('%s','%s','%s'))`, ValidatorUnreachable, ValidatorIdle, IPDone, IPFailed, IPOccupied), now) if err != nil { return 0, fmt.Errorf("reconcile validators: %w", err) } n, _ := res.RowsAffected() res, err = d.ExecContext(ctx, ` UPDATE validators SET state=?, updated_at=? WHERE state=? AND current_ip_id IS NULL`, ValidatorIdle, now, ValidatorAssigned) if err != nil { return n, fmt.Errorf("reconcile validators: %w", err) } m, _ := res.RowsAffected() return n + m, nil } // AdminCreateValidator registers a brand-new validator via the admin API. // Unlike RegisterValidator (used by the agent's self-registration call), // this refuses to upsert over an existing row. func (d *DB) AdminCreateValidator(ctx context.Context, validatorID, osPortID string) error { var exists int err := d.QueryRowContext(ctx, `SELECT 1 FROM validators WHERE validator_id=?`, validatorID).Scan(&exists) if err != nil && err != sql.ErrNoRows { return err } if err == nil { return fmt.Errorf("validator %q: %w", validatorID, ErrConflict) } now := timeToDB(Now()) _, err = d.ExecContext(ctx, ` INSERT INTO validators (validator_id, hostname, os_port_id, agent_version, state, created_at, updated_at) VALUES (?, '', ?, '', ?, ?, ?) `, validatorID, osPortID, ValidatorIdle, now, now) if err != nil { return fmt.Errorf("create validator: %w", err) } return nil } // AdminUpdateValidatorPort updates an existing validator's Neutron port ID. func (d *DB) AdminUpdateValidatorPort(ctx context.Context, validatorID, osPortID string) error { res, err := d.ExecContext(ctx, ` UPDATE validators SET os_port_id=?, updated_at=? WHERE validator_id=? `, osPortID, timeToDB(Now()), validatorID) if err != nil { return fmt.Errorf("update validator port: %w", err) } if n, _ := res.RowsAffected(); n == 0 { return fmt.Errorf("validator %q: %w", validatorID, ErrNotFound) } return nil } // DeleteValidator removes a validator, refusing if it currently owns an IP. // ip_queue.owner_validator_id is a permanent historical record (set once an // IP is claimed, never cleared on completion — see RequeueOrFail/ // ReleaseFIP), so any validator that has ever processed an IP would // otherwise violate the FK constraint on delete; those historical // references are cleared in the same transaction once we've confirmed the // validator isn't currently busy. func (d *DB) DeleteValidator(ctx context.Context, validatorID string) error { tx, err := d.BeginTx(ctx, nil) if err != nil { return err } defer tx.Rollback() var currentIPID sql.NullInt64 err = tx.QueryRowContext(ctx, `SELECT current_ip_id FROM validators WHERE validator_id=?`, validatorID).Scan(¤tIPID) if err == sql.ErrNoRows { return fmt.Errorf("validator %q: %w", validatorID, ErrNotFound) } if err != nil { return err } if currentIPID.Valid { return fmt.Errorf("validator %q owns ip_id %d: %w", validatorID, currentIPID.Int64, ErrBusy) } if _, err := tx.ExecContext(ctx, `UPDATE ip_queue SET owner_validator_id=NULL WHERE owner_validator_id=?`, validatorID); err != nil { return fmt.Errorf("clear historical ip_queue references: %w", err) } if _, err := tx.ExecContext(ctx, `DELETE FROM validators WHERE validator_id=?`, validatorID); err != nil { return fmt.Errorf("delete validator: %w", err) } return tx.Commit() } type rowScanner interface { Scan(dest ...interface{}) error } func scanValidator(row rowScanner) (*Validator, error) { var v Validator var currentIPID sql.NullInt64 var lastHeartbeat sql.NullString var createdAt, updatedAt string if err := row.Scan(&v.ValidatorID, &v.Hostname, &v.OSPortID, &v.State, ¤tIPID, &v.AgentVersion, &lastHeartbeat, &createdAt, &updatedAt); err != nil { return nil, err } if currentIPID.Valid { v.CurrentIPID = ¤tIPID.Int64 } hb, err := nullStringToTimePtr(lastHeartbeat) if err != nil { return nil, err } v.LastHeartbeatAt = hb if v.CreatedAt, err = dbToTime(createdAt); err != nil { return nil, err } if v.UpdatedAt, err = dbToTime(updatedAt); err != nil { return nil, err } return &v, nil }