repo init
This commit is contained in:
ayurishchev committed
commit
7e44db87b2
48 files changed
+5346
No files matched your search
@@ -0,0 +1,85 @@
|
||||
// Package db owns the SQLite connection, schema migrations, and all queries
|
||||
// used by the Control API. It is the only package in the system that talks
|
||||
// to the database directly — agents and probers never connect to it.
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
//go:embed migrations/0001_init.sql
|
||||
var initSchema string
|
||||
|
||||
type DB struct {
|
||||
*sql.DB
|
||||
}
|
||||
|
||||
// Open opens (creating if necessary) the SQLite database at path, applies
|
||||
// pragmas suited to a single-writer WAL workload, and runs any pending
|
||||
// schema migrations.
|
||||
func Open(ctx context.Context, path string) (*DB, error) {
|
||||
sqlDB, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open sqlite: %w", err)
|
||||
}
|
||||
// Control API is the sole writer; one connection avoids SQLITE_BUSY
|
||||
// entirely for writes while still allowing concurrent reads via WAL.
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
|
||||
for _, pragma := range []string{
|
||||
"PRAGMA journal_mode=WAL",
|
||||
"PRAGMA synchronous=NORMAL",
|
||||
"PRAGMA foreign_keys=ON",
|
||||
"PRAGMA busy_timeout=5000",
|
||||
} {
|
||||
if _, err := sqlDB.ExecContext(ctx, pragma); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("apply pragma %q: %w", pragma, err)
|
||||
}
|
||||
}
|
||||
|
||||
d := &DB{DB: sqlDB}
|
||||
if err := d.migrate(ctx); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("migrate: %w", err)
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
|
||||
// migrate applies the embedded schema exactly once, tracked via
|
||||
// PRAGMA user_version so repeated startups are no-ops.
|
||||
func (d *DB) migrate(ctx context.Context) error {
|
||||
var version int
|
||||
if err := d.QueryRowContext(ctx, "PRAGMA user_version").Scan(&version); err != nil {
|
||||
return fmt.Errorf("read user_version: %w", err)
|
||||
}
|
||||
if version >= 1 {
|
||||
return nil
|
||||
}
|
||||
|
||||
tx, err := d.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
if _, err := tx.ExecContext(ctx, initSchema); err != nil {
|
||||
return fmt.Errorf("apply 0001_init.sql: %w", err)
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, "PRAGMA user_version=1"); err != nil {
|
||||
return fmt.Errorf("set user_version: %w", err)
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// Now returns the current time truncated to millisecond precision, the
|
||||
// granularity used consistently for all timestamp columns.
|
||||
func Now() time.Time {
|
||||
return time.Now().UTC().Truncate(time.Millisecond)
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
CREATE TABLE validators (
|
||||
validator_id TEXT PRIMARY KEY,
|
||||
hostname TEXT NOT NULL DEFAULT '',
|
||||
os_port_id TEXT NOT NULL DEFAULT '',
|
||||
state TEXT NOT NULL DEFAULT 'unregistered',
|
||||
current_ip_id INTEGER REFERENCES ip_queue(id),
|
||||
agent_version TEXT NOT NULL DEFAULT '',
|
||||
last_heartbeat_at TIMESTAMP,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE TABLE ip_queue (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ip_address TEXT NOT NULL UNIQUE,
|
||||
sequence INTEGER NOT NULL,
|
||||
state TEXT NOT NULL DEFAULT 'queued',
|
||||
owner_validator_id TEXT REFERENCES validators(validator_id),
|
||||
fip_id TEXT NOT NULL DEFAULT '',
|
||||
attempt_number INTEGER NOT NULL DEFAULT 1,
|
||||
retry_count INTEGER NOT NULL DEFAULT 0,
|
||||
lease_expires_at TIMESTAMP,
|
||||
egress_complete BOOLEAN NOT NULL DEFAULT 0,
|
||||
site1_complete BOOLEAN NOT NULL DEFAULT 0,
|
||||
site2_complete BOOLEAN NOT NULL DEFAULT 0,
|
||||
site3_complete BOOLEAN NOT NULL DEFAULT 0,
|
||||
overall_result TEXT NOT NULL DEFAULT '',
|
||||
assigned_at TIMESTAMP,
|
||||
aggregated_at TIMESTAMP,
|
||||
fip_released_at TIMESTAMP,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
CREATE INDEX idx_ip_queue_state_seq ON ip_queue(state, sequence);
|
||||
|
||||
CREATE TABLE checks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ip_id INTEGER NOT NULL REFERENCES ip_queue(id),
|
||||
ip_address TEXT NOT NULL,
|
||||
attempt_number INTEGER NOT NULL,
|
||||
validator_id TEXT NOT NULL DEFAULT '',
|
||||
source TEXT NOT NULL,
|
||||
check_type TEXT NOT NULL,
|
||||
target TEXT NOT NULL DEFAULT '',
|
||||
success BOOLEAN NOT NULL,
|
||||
latency_ms INTEGER NOT NULL DEFAULT 0,
|
||||
detail TEXT NOT NULL DEFAULT '',
|
||||
checked_at TIMESTAMP NOT NULL,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE(ip_id, attempt_number, source, check_type, target)
|
||||
);
|
||||
CREATE INDEX idx_checks_ip_attempt ON checks(ip_id, attempt_number);
|
||||
|
||||
CREATE TABLE events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
source_type TEXT NOT NULL,
|
||||
source_id TEXT NOT NULL DEFAULT '',
|
||||
ip_id INTEGER REFERENCES ip_queue(id),
|
||||
event_type TEXT NOT NULL,
|
||||
payload TEXT NOT NULL DEFAULT '',
|
||||
occurred_at TIMESTAMP NOT NULL,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
CREATE INDEX idx_events_ip ON events(ip_id);
|
||||
CREATE INDEX idx_events_source ON events(source_type, source_id);
|
||||
@@ -0,0 +1,117 @@
|
||||
package db
|
||||
|
||||
import "time"
|
||||
|
||||
// Validator and IP lifecycle states. Kept as typed string constants rather
|
||||
// than a Go enum type so they round-trip through SQLite TEXT columns and
|
||||
// JSON without conversion.
|
||||
const (
|
||||
ValidatorUnregistered = "unregistered"
|
||||
ValidatorIdle = "idle"
|
||||
ValidatorAssigned = "assigned"
|
||||
ValidatorChecking = "checking"
|
||||
ValidatorUnreachable = "unreachable"
|
||||
|
||||
IPQueued = "queued"
|
||||
IPAssigningFIP = "assigning_fip"
|
||||
IPAwaitingSelfCheck = "awaiting_self_check"
|
||||
IPChecking = "checking"
|
||||
IPAggregating = "aggregating"
|
||||
IPDone = "done"
|
||||
IPFailed = "failed"
|
||||
|
||||
ResultPass = "pass"
|
||||
ResultPartial = "partial"
|
||||
ResultFail = "fail"
|
||||
|
||||
SourceEgress = "egress"
|
||||
)
|
||||
|
||||
// InboundSource returns the checks.source value for the given prober site
|
||||
// index (1-based), e.g. InboundSource(1) == "inbound-site-1".
|
||||
func InboundSource(siteIndex int) string {
|
||||
return "inbound-site-" + itoa(siteIndex)
|
||||
}
|
||||
|
||||
func itoa(n int) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
neg := n < 0
|
||||
if neg {
|
||||
n = -n
|
||||
}
|
||||
var buf [20]byte
|
||||
i := len(buf)
|
||||
for n > 0 {
|
||||
i--
|
||||
buf[i] = byte('0' + n%10)
|
||||
n /= 10
|
||||
}
|
||||
if neg {
|
||||
i--
|
||||
buf[i] = '-'
|
||||
}
|
||||
return string(buf[i:])
|
||||
}
|
||||
|
||||
type Validator struct {
|
||||
ValidatorID string
|
||||
Hostname string
|
||||
OSPortID string
|
||||
State string
|
||||
CurrentIPID *int64
|
||||
AgentVersion string
|
||||
LastHeartbeatAt *time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type IPQueueItem struct {
|
||||
ID int64
|
||||
IPAddress string
|
||||
Sequence int
|
||||
State string
|
||||
OwnerValidatorID *string
|
||||
FIPID string
|
||||
AttemptNumber int
|
||||
RetryCount int
|
||||
LeaseExpiresAt *time.Time
|
||||
EgressComplete bool
|
||||
Site1Complete bool
|
||||
Site2Complete bool
|
||||
Site3Complete bool
|
||||
OverallResult string
|
||||
AssignedAt *time.Time
|
||||
AggregatedAt *time.Time
|
||||
FIPReleasedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type Check struct {
|
||||
ID int64
|
||||
IPID int64
|
||||
IPAddress string
|
||||
AttemptNumber int
|
||||
ValidatorID string
|
||||
Source string
|
||||
CheckType string
|
||||
Target string
|
||||
Success bool
|
||||
LatencyMS int64
|
||||
Detail string
|
||||
CheckedAt time.Time
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
ID int64
|
||||
SourceType string
|
||||
SourceID string
|
||||
IPID *int64
|
||||
EventType string
|
||||
Payload string
|
||||
OccurredAt time.Time
|
||||
CreatedAt time.Time
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
// UpsertCheck records (or, on retry, overwrites) a single check result. The
|
||||
// UNIQUE(ip_id, attempt_number, source, check_type, target) constraint plus
|
||||
// this upsert is what makes agent/prober result submission safely
|
||||
// retryable without producing duplicate rows.
|
||||
func (d *DB) UpsertCheck(ctx context.Context, c Check) error {
|
||||
_, err := d.ExecContext(ctx, `
|
||||
INSERT INTO checks (ip_id, ip_address, attempt_number, validator_id, source, check_type, target,
|
||||
success, latency_ms, detail, checked_at, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(ip_id, attempt_number, source, check_type, target) DO UPDATE SET
|
||||
validator_id=excluded.validator_id,
|
||||
success=excluded.success,
|
||||
latency_ms=excluded.latency_ms,
|
||||
detail=excluded.detail,
|
||||
checked_at=excluded.checked_at
|
||||
`, c.IPID, c.IPAddress, c.AttemptNumber, c.ValidatorID, c.Source, c.CheckType, c.Target,
|
||||
c.Success, c.LatencyMS, c.Detail, timeToDB(c.CheckedAt), timeToDB(Now()))
|
||||
return err
|
||||
}
|
||||
|
||||
// ListChecksForAttempt returns every check recorded for an IP's current
|
||||
// attempt — the input to overall-result aggregation.
|
||||
func (d *DB) ListChecksForAttempt(ctx context.Context, ipID int64, attemptNumber int) ([]Check, error) {
|
||||
rows, err := d.QueryContext(ctx, `
|
||||
SELECT id, ip_id, ip_address, attempt_number, validator_id, source, check_type, target,
|
||||
success, latency_ms, detail, checked_at, created_at
|
||||
FROM checks WHERE ip_id=? AND attempt_number=?
|
||||
ORDER BY source, check_type, target
|
||||
`, ipID, attemptNumber)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []Check
|
||||
for rows.Next() {
|
||||
var c Check
|
||||
var checkedAt, createdAt string
|
||||
if err := rows.Scan(&c.ID, &c.IPID, &c.IPAddress, &c.AttemptNumber, &c.ValidatorID, &c.Source,
|
||||
&c.CheckType, &c.Target, &c.Success, &c.LatencyMS, &c.Detail, &checkedAt, &createdAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if c.CheckedAt, err = dbToTime(checkedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if c.CreatedAt, err = dbToTime(createdAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package db
|
||||
|
||||
import "context"
|
||||
|
||||
func (d *DB) InsertEvent(ctx context.Context, e Event) error {
|
||||
_, err := d.ExecContext(ctx, `
|
||||
INSERT INTO events (source_type, source_id, ip_id, event_type, payload, occurred_at, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
`, e.SourceType, e.SourceID, e.IPID, e.EventType, e.Payload, timeToDB(e.OccurredAt), timeToDB(Now()))
|
||||
return err
|
||||
}
|
||||
|
||||
// ListEventsForIP returns the audit trail for a single IP, most recent
|
||||
// first — used by the admin detail endpoint.
|
||||
func (d *DB) ListEventsForIP(ctx context.Context, ipID int64) ([]Event, error) {
|
||||
rows, err := d.QueryContext(ctx, `
|
||||
SELECT id, source_type, source_id, ip_id, event_type, payload, occurred_at, created_at
|
||||
FROM events WHERE ip_id=? ORDER BY occurred_at DESC
|
||||
`, ipID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanEvents(rows)
|
||||
}
|
||||
|
||||
func (d *DB) ListRecentEvents(ctx context.Context, limit int) ([]Event, error) {
|
||||
rows, err := d.QueryContext(ctx, `
|
||||
SELECT id, source_type, source_id, ip_id, event_type, payload, occurred_at, created_at
|
||||
FROM events ORDER BY id DESC LIMIT ?
|
||||
`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanEvents(rows)
|
||||
}
|
||||
|
||||
func scanEvents(rows interface {
|
||||
Next() bool
|
||||
Scan(...interface{}) error
|
||||
Err() error
|
||||
}) ([]Event, error) {
|
||||
var out []Event
|
||||
for rows.Next() {
|
||||
var e Event
|
||||
var ipID *int64
|
||||
var occurredAt, createdAt string
|
||||
if err := rows.Scan(&e.ID, &e.SourceType, &e.SourceID, &ipID, &e.EventType, &e.Payload, &occurredAt, &createdAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
e.IPID = ipID
|
||||
var err error
|
||||
if e.OccurredAt, err = dbToTime(occurredAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if e.CreatedAt, err = dbToTime(createdAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, e)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
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'
|
||||
// (but not mid-assignment) should also come back to idle.
|
||||
_, err = d.ExecContext(ctx, `
|
||||
UPDATE validators SET state=?, updated_at=?
|
||||
WHERE validator_id=? AND state IN (?, ?)
|
||||
`, ValidatorIdle, now, validatorID, ValidatorUnregistered, ValidatorUnreachable)
|
||||
if err != nil {
|
||||
return fmt.Errorf("register validator (reactivate): %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
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 ? ELSE state END
|
||||
WHERE validator_id=?
|
||||
`, now, now, ValidatorUnreachable, ValidatorIdle, 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
|
||||
}
|
||||
|
||||
// FreeValidator returns a validator to idle with no assigned IP. Used after
|
||||
// an IP finishes (success or failure) or is reclaimed by the lease sweep.
|
||||
func (d *DB) FreeValidator(ctx context.Context, validatorID string) error {
|
||||
_, err := d.ExecContext(ctx, `
|
||||
UPDATE validators SET state=?, current_ip_id=NULL, updated_at=?
|
||||
WHERE validator_id=?
|
||||
`, ValidatorIdle, timeToDB(Now()), validatorID)
|
||||
return err
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Timestamps are stored as RFC3339Nano TEXT explicitly (rather than relying
|
||||
// on driver-specific time.Time marshaling) so the on-disk format is stable
|
||||
// and easy to inspect with the sqlite3 CLI.
|
||||
|
||||
const timeLayout = time.RFC3339Nano
|
||||
|
||||
func timeToDB(t time.Time) string {
|
||||
return t.UTC().Format(timeLayout)
|
||||
}
|
||||
|
||||
func timePtrToDB(t *time.Time) interface{} {
|
||||
if t == nil {
|
||||
return nil
|
||||
}
|
||||
return timeToDB(*t)
|
||||
}
|
||||
|
||||
func dbToTime(s string) (time.Time, error) {
|
||||
return time.Parse(timeLayout, s)
|
||||
}
|
||||
|
||||
func nullStringToTimePtr(ns sql.NullString) (*time.Time, error) {
|
||||
if !ns.Valid || ns.String == "" {
|
||||
return nil, nil
|
||||
}
|
||||
t, err := dbToTime(ns.String)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
Reference in new issue
Block a user