Add agent restart lifecycle command
Extends the existing CheckIn RPC with a one-shot AgentCommand (restart only -- stop/uninstall need real per-platform OS service-manager integration and stay deliberately out of scope), delivered at-most-once: cleared the instant it's handed to the agent in a response, since a restarting agent's process is gone before it could ever confirm receipt. On restart, the agent flushes whatever's buffered, aborts its source task, and exits cleanly, relying entirely on the host's own service manager to bring it back up. Issuing a command is gated at RoleAdmin (stricter than config editing's RoleEditor) and logged into the same audit_log table Phase 7's AI interactions use, via a new agent_command event type. A real bug was found and fixed during live verification: the first implementation tried to atomically read-and-clear pending_command in a single INSERT...ON CONFLICT statement using a sibling CTE referenced only from RETURNING, on the assumption that Postgres evaluates every part of a WITH query against one pre-statement snapshot. That's wrong specifically for FOR UPDATE, which always reads the latest row version including one written earlier in the same statement -- confirmed empirically (a restart command was always coming back empty even when genuinely pending, so the agent never received it). Fixed by splitting into two real, ordered statements inside one explicit transaction. See /docs/agent-management-design.md's "Lifecycle commands" section.
This commit is contained in:
@@ -10,9 +10,11 @@ package agentregistry
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/sentry/sentry/ingest/internal/grpcserver"
|
||||
@@ -49,16 +51,47 @@ func New(pool *pgxpool.Pool) *Registry {
|
||||
|
||||
var _ grpcserver.AgentRegistry = (*Registry)(nil)
|
||||
|
||||
func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver.AgentCheckIn) (grpcserver.AgentOverride, error) {
|
||||
func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver.AgentCheckIn) (grpcserver.CheckInResult, error) {
|
||||
if tenantID == "" {
|
||||
tenantID = defaultTenantID
|
||||
}
|
||||
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: beginning transaction: %w", err)
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
// Read (and lock) whatever pending_command exists BEFORE the
|
||||
// upsert below clears it -- as two real, ordered statements in one
|
||||
// transaction, not one clever statement. An earlier version tried
|
||||
// to do this with a single INSERT...ON CONFLICT plus a sibling CTE
|
||||
// referenced only from RETURNING, on the assumption that Postgres
|
||||
// evaluates every part of a WITH query against the same pre-
|
||||
// statement snapshot; that assumption is wrong specifically for
|
||||
// FOR UPDATE, which always locks (and therefore reads) the latest
|
||||
// row version to do its job, including versions written earlier in
|
||||
// the SAME statement -- confirmed empirically against a live
|
||||
// Postgres (the CTE's FOR UPDATE was reading its own sibling
|
||||
// UPDATE's just-cleared NULL, always reporting "no command" even
|
||||
// when one was genuinely pending). Two statements in one
|
||||
// transaction has no such ambiguity: the SELECT strictly
|
||||
// happens-before the UPDATE, full stop. FOR UPDATE against zero
|
||||
// rows (a brand-new agent's first-ever check-in) is a harmless
|
||||
// no-op -- there's nothing to lock or have a pending command yet.
|
||||
var pendingCommand *string
|
||||
err = tx.QueryRow(ctx, `SELECT pending_command FROM agents WHERE tenant_id = $1 AND host = $2 FOR UPDATE`,
|
||||
tenantID, info.Host,
|
||||
).Scan(&pendingCommand)
|
||||
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
||||
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: reading pending command: %w", err)
|
||||
}
|
||||
|
||||
var (
|
||||
desiredOverride []byte
|
||||
desiredVersion *string
|
||||
)
|
||||
err := r.pool.QueryRow(ctx, `
|
||||
err = tx.QueryRow(ctx, `
|
||||
INSERT INTO agents (
|
||||
id, tenant_id, host, service,
|
||||
reported_agent_version, reported_source_kind, reported_source_detail,
|
||||
@@ -76,7 +109,8 @@ func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver
|
||||
reported_heartbeat_on = EXCLUDED.reported_heartbeat_on,
|
||||
reported_heartbeat_ms = EXCLUDED.reported_heartbeat_ms,
|
||||
last_seen_at = now(),
|
||||
applied_override_version = EXCLUDED.applied_override_version
|
||||
applied_override_version = EXCLUDED.applied_override_version,
|
||||
pending_command = NULL
|
||||
RETURNING desired_override, desired_override_version`,
|
||||
uuid.NewString(), tenantID, info.Host, info.Service,
|
||||
info.AgentVersion, info.SourceKind, info.SourceDetail,
|
||||
@@ -85,18 +119,27 @@ func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver
|
||||
info.AppliedOverrideVersion,
|
||||
).Scan(&desiredOverride, &desiredVersion)
|
||||
if err != nil {
|
||||
return grpcserver.AgentOverride{}, fmt.Errorf("agentregistry: upserting check-in: %w", err)
|
||||
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: upserting check-in: %w", err)
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: committing check-in: %w", err)
|
||||
}
|
||||
|
||||
result := grpcserver.CheckInResult{}
|
||||
if pendingCommand != nil {
|
||||
result.Command = *pendingCommand
|
||||
}
|
||||
|
||||
if desiredVersion == nil || len(desiredOverride) == 0 {
|
||||
return grpcserver.AgentOverride{HasOverride: false}, nil
|
||||
return result, nil
|
||||
}
|
||||
|
||||
var fields overrideFields
|
||||
if err := json.Unmarshal(desiredOverride, &fields); err != nil {
|
||||
return grpcserver.AgentOverride{}, fmt.Errorf("agentregistry: parsing stored override: %w", err)
|
||||
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: parsing stored override: %w", err)
|
||||
}
|
||||
return grpcserver.AgentOverride{
|
||||
result.Override = grpcserver.AgentOverride{
|
||||
HasOverride: true,
|
||||
BatchMaxSize: fields.BatchMaxSize,
|
||||
BatchFlushIntervalMS: fields.BatchFlushIntervalMS,
|
||||
@@ -104,5 +147,6 @@ func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver
|
||||
HeartbeatIntervalMS: fields.HeartbeatIntervalMS,
|
||||
JournaldUnit: fields.JournaldUnit,
|
||||
Version: *desiredVersion,
|
||||
}, nil
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user