forked from wrenn/wrenn
Implement host registration, JWT refresh tokens, and multi-host scheduling
Replaces the hardcoded CP_HOST_AGENT_ADDR single-agent setup with a DB-driven registration system supporting multiple host agents (BYOC). Key changes: - Host agents register via one-time token, receive a 7-day JWT + 60-day refresh token; heartbeat loop auto-refreshes on 401/403 and pauses all sandboxes if refresh fails - HostClientPool: lazy Connect RPC client cache keyed by host ID, replacing the single static agent client throughout the API and service layers - RoundRobinScheduler: picks an online host for each new sandbox via ListActiveHosts; extensible for future scheduling strategies - HostMonitor (replaces Reconciler): passive heartbeat staleness check marks hosts unreachable and sandboxes missing after 90s; active reconciliation per online host restores missing-but-alive sandboxes and stops orphans - Graceful host delete: returns 409 with affected sandbox list without ?force=true; force-delete destroys sandboxes then evicts pool client - Snapshot delete broadcasts to all online hosts (templates have no host_id) - sandbox.Manager.PauseAll: pauses all running VMs on CP connectivity loss - New migration: host_refresh_tokens table with token rotation (issue-then- revoke ordering to prevent lockout on mid-rotation crash) - New sandbox status 'missing' (reversible, unlike 'stopped') and host status 'unreachable'; both reflected in OpenAPI spec - Fix: refresh token auth failure now returns 401 (was 400 via generic 'invalid' substring match in serviceErrToHTTP)
This commit is contained in:
@ -234,6 +234,50 @@ func (q *Queries) InsertHostToken(ctx context.Context, arg InsertHostTokenParams
|
||||
return i, err
|
||||
}
|
||||
|
||||
const listActiveHosts = `-- name: ListActiveHosts :many
|
||||
SELECT id, type, team_id, provider, availability_zone, arch, cpu_cores, memory_mb, disk_gb, address, status, last_heartbeat_at, metadata, created_by, created_at, updated_at, cert_fingerprint, mtls_enabled FROM hosts WHERE status NOT IN ('pending', 'offline') ORDER BY created_at
|
||||
`
|
||||
|
||||
// Returns all hosts that have completed registration (not pending/offline).
|
||||
func (q *Queries) ListActiveHosts(ctx context.Context) ([]Host, error) {
|
||||
rows, err := q.db.Query(ctx, listActiveHosts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []Host
|
||||
for rows.Next() {
|
||||
var i Host
|
||||
if err := rows.Scan(
|
||||
&i.ID,
|
||||
&i.Type,
|
||||
&i.TeamID,
|
||||
&i.Provider,
|
||||
&i.AvailabilityZone,
|
||||
&i.Arch,
|
||||
&i.CpuCores,
|
||||
&i.MemoryMb,
|
||||
&i.DiskGb,
|
||||
&i.Address,
|
||||
&i.Status,
|
||||
&i.LastHeartbeatAt,
|
||||
&i.Metadata,
|
||||
&i.CreatedBy,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
&i.CertFingerprint,
|
||||
&i.MtlsEnabled,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, i)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const listHosts = `-- name: ListHosts :many
|
||||
SELECT id, type, team_id, provider, availability_zone, arch, cpu_cores, memory_mb, disk_gb, address, status, last_heartbeat_at, metadata, created_by, created_at, updated_at, cert_fingerprint, mtls_enabled FROM hosts ORDER BY created_at DESC
|
||||
`
|
||||
@ -461,6 +505,15 @@ func (q *Queries) MarkHostTokenUsed(ctx context.Context, id string) error {
|
||||
return err
|
||||
}
|
||||
|
||||
const markHostUnreachable = `-- name: MarkHostUnreachable :exec
|
||||
UPDATE hosts SET status = 'unreachable', updated_at = NOW() WHERE id = $1
|
||||
`
|
||||
|
||||
func (q *Queries) MarkHostUnreachable(ctx context.Context, id string) error {
|
||||
_, err := q.db.Exec(ctx, markHostUnreachable, id)
|
||||
return err
|
||||
}
|
||||
|
||||
const registerHost = `-- name: RegisterHost :execrows
|
||||
UPDATE hosts
|
||||
SET arch = $2,
|
||||
@ -521,6 +574,20 @@ func (q *Queries) UpdateHostHeartbeat(ctx context.Context, id string) error {
|
||||
return err
|
||||
}
|
||||
|
||||
const updateHostHeartbeatAndStatus = `-- name: UpdateHostHeartbeatAndStatus :exec
|
||||
UPDATE hosts
|
||||
SET last_heartbeat_at = NOW(),
|
||||
status = CASE WHEN status = 'unreachable' THEN 'online' ELSE status END,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
`
|
||||
|
||||
// Updates last_heartbeat_at and transitions unreachable hosts back to online.
|
||||
func (q *Queries) UpdateHostHeartbeatAndStatus(ctx context.Context, id string) error {
|
||||
_, err := q.db.Exec(ctx, updateHostHeartbeatAndStatus, id)
|
||||
return err
|
||||
}
|
||||
|
||||
const updateHostStatus = `-- name: UpdateHostStatus :exec
|
||||
UPDATE hosts SET status = $2, updated_at = NOW() WHERE id = $1
|
||||
`
|
||||
|
||||
Reference in New Issue
Block a user