package workerruntime import ( "context" "encoding/json" "errors" "fmt" "log/slog" "time" "github.com/itworx/pulse/internal/probe" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) const ( // MaxProbeBatch bounds one probe scan. It matches the documented scale // target of 300 service probes (SYSTEM_ARCHITECTURE section 7). MaxProbeBatch = 300 // DefaultProbeConcurrency bounds simultaneous probe executions. DefaultProbeConcurrency = 16 ) // ProbeStore is the persistence surface of the probe execution job. type ProbeStore interface { // ListDue returns enabled, unarchived probes whose interval has elapsed. ListDue(ctx context.Context, now time.Time, limit int) ([]probe.Definition, error) // CountEnabled reports how many probes are configured and enabled. It // separates "nothing was due" from "nothing is configured", so an empty // installation reports Disabled instead of a green probe subsystem. CountEnabled(ctx context.Context) (int, error) // SaveResults records probe results and reports how many were new. SaveResults(ctx context.Context, results []probe.Result) (int, error) } // ProbeJob executes due probes and records their results. // // Execution goes through probe.Scheduler and probe.ProbeExecutor with the // configured network policy. HTTP transport and dialing are deliberately left // to the executor's own safe client, so the SSRF protections in // internal/probe/policy.go (blocked link-local/metadata/loopback ranges, // redirect re-validation at dial time, response caps) stay fully in force. type ProbeJob struct { Store ProbeStore Scheduler *probe.Scheduler Logger *slog.Logger Now func() time.Time } // NewProbeJob builds a probe job with a bounded scheduler around the policy. func NewProbeJob(store ProbeStore, policy probe.NetworkPolicy, logger *slog.Logger) (*ProbeJob, error) { if store == nil { return nil, fmt.Errorf("%w: probe job requires a store", ErrInvalidConfig) } if err := policy.Validate(); err != nil { return nil, err } scheduler, err := probe.NewScheduler(probe.ProbeExecutor{Policy: policy}, probe.SchedulerConfig{ MaxConcurrent: DefaultProbeConcurrency, MaxAttempts: 2, AttemptTimeout: 10 * time.Second, RetryBackoff: 250 * time.Millisecond, }) if err != nil { return nil, err } return &ProbeJob{Store: store, Scheduler: scheduler, Logger: logger}, nil } func (j *ProbeJob) now() time.Time { if j.Now != nil { return j.Now().UTC() } return time.Now().UTC() } func (j *ProbeJob) logger() *slog.Logger { if j.Logger != nil { return j.Logger } return slog.Default() } // Run executes every due probe once and persists the results. func (j *ProbeJob) Run(ctx context.Context) (Outcome, error) { if j == nil || j.Store == nil || j.Scheduler == nil { return Outcome{Disabled: true, Reason: "probe_execution_not_configured"}, nil } now := j.now() definitions, err := j.Store.ListDue(ctx, now, MaxProbeBatch) if err != nil { return Outcome{}, fmt.Errorf("list due probes: %w", err) } counts := map[string]int64{"due": int64(len(definitions))} if len(definitions) == 0 { configured, err := j.Store.CountEnabled(ctx) if err != nil { return Outcome{Counts: counts}, fmt.Errorf("count enabled probes: %w", err) } counts["enabled"] = int64(configured) if configured == 0 { return Outcome{Disabled: true, Reason: "no_probes_configured", Counts: counts}, nil } return Outcome{Counts: counts}, nil } // An individually invalid definition is skipped and counted rather than // failing the whole batch, so one bad probe cannot stop the other 299. valid := make([]probe.Definition, 0, len(definitions)) for _, definition := range definitions { if err := definition.Validate(); err != nil { counts["invalid"]++ j.logger().Error("probe definition is invalid", "probe", definition.ID, "error", boundedError(err)) continue } valid = append(valid, definition) } if len(valid) == 0 { return Outcome{Counts: counts}, fmt.Errorf("all %d due probes are invalid", len(definitions)) } report, err := j.Scheduler.Run(ctx, valid) if err != nil { return Outcome{Counts: counts}, fmt.Errorf("execute probes: %w", err) } for _, result := range report.Results { counts["state_"+result.State]++ } counts["executed"] = int64(len(report.Results)) // A batch that ran into the job timeout or a shutdown still produced real // observations. They are persisted under a short detached deadline so the // work is not silently thrown away, and the job still reports the failure. saveCtx := ctx if ctx.Err() != nil { detached, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() saveCtx = detached } saved, err := j.Store.SaveResults(saveCtx, report.Results) if err != nil { return Outcome{Counts: counts}, fmt.Errorf("save probe results: %w", err) } counts["saved"] = int64(saved) return Outcome{Counts: counts}, nil } // Shutdown cancels in-flight probes and waits for them, bounded by ctx. func (j *ProbeJob) Shutdown(ctx context.Context) error { if j == nil || j.Scheduler == nil { return nil } return j.Scheduler.Shutdown(ctx) } // PostgresProbeStore reads probe definitions and writes probe_results. type PostgresProbeStore struct { Pool *pgxpool.Pool // SourceID optionally attributes results to a data source. SourceID string } // ListDue returns probes whose interval has elapsed since their last result. func (s PostgresProbeStore) ListDue(ctx context.Context, now time.Time, limit int) ([]probe.Definition, error) { if s.Pool == nil { return nil, fmt.Errorf("%w: probe store has no database pool", ErrInvalidConfig) } if limit < 1 || limit > MaxProbeBatch { return nil, fmt.Errorf("%w: probe batch must be between 1 and %d", ErrInvalidConfig, MaxProbeBatch) } rows, err := s.Pool.Query(ctx, ` SELECT p.id::text,p.service_id::text,COALESCE(p.endpoint_id::text,''),p.name,p.probe_type,p.target,p.interval_seconds,p.timeout_seconds, p.expected_status_codes,p.follow_redirects,p.verify_tls,p.content_assertion,COALESCE(p.secret_reference,''), COALESCE(p.network_policy_id::text,''),p.revision,p.created_at,p.updated_at FROM probes p LEFT JOIN LATERAL (SELECT max(r.observed_at) AS last_observed_at FROM probe_results r WHERE r.probe_id = p.id) latest ON true WHERE p.archived_at IS NULL AND p.enabled = true AND (latest.last_observed_at IS NULL OR latest.last_observed_at <= $1::timestamptz - make_interval(secs => p.interval_seconds)) ORDER BY p.id ASC LIMIT $2`, now.UTC(), limit) if err != nil { return nil, fmt.Errorf("query due probes: %w", err) } defer rows.Close() definitions := make([]probe.Definition, 0, limit) for rows.Next() { var definition probe.Definition var target, expected, assertion []byte var intervalSeconds, timeoutSeconds int if err := rows.Scan(&definition.ID, &definition.ServiceID, &definition.EndpointID, &definition.Name, &definition.Type, &target, &intervalSeconds, &timeoutSeconds, &expected, &definition.FollowRedirects, &definition.VerifyTLS, &assertion, &definition.SecretReference, &definition.NetworkPolicyID, &definition.Revision, &definition.CreatedAt, &definition.UpdatedAt); err != nil { return nil, fmt.Errorf("scan probe definition: %w", err) } definition.Enabled = true definition.Interval = time.Duration(intervalSeconds) * time.Second definition.Timeout = time.Duration(timeoutSeconds) * time.Second if err := json.Unmarshal(target, &definition.Target); err != nil { return nil, fmt.Errorf("decode probe target: %w", err) } if len(expected) > 0 { if err := json.Unmarshal(expected, &definition.ExpectedStatusCodes); err != nil { return nil, fmt.Errorf("decode probe expected status codes: %w", err) } } if len(assertion) > 0 { if err := json.Unmarshal(assertion, &definition.ContentAssertion); err != nil { return nil, fmt.Errorf("decode probe content assertion: %w", err) } } definitions = append(definitions, definition) } if err := rows.Err(); err != nil { return nil, fmt.Errorf("read probe definition rows: %w", err) } return definitions, nil } // CountEnabled counts configured, enabled, unarchived probes. func (s PostgresProbeStore) CountEnabled(ctx context.Context) (int, error) { if s.Pool == nil { return 0, fmt.Errorf("%w: probe store has no database pool", ErrInvalidConfig) } var count int if err := s.Pool.QueryRow(ctx, `SELECT count(*) FROM probes WHERE archived_at IS NULL AND enabled = true`).Scan(&count); err != nil { return 0, fmt.Errorf("count enabled probes: %w", err) } return count, nil } // SaveResults writes probe results and any observed certificate. Both inserts // use their natural key with ON CONFLICT DO NOTHING, so replaying a batch after // a crash cannot create duplicate history. func (s PostgresProbeStore) SaveResults(ctx context.Context, results []probe.Result) (int, error) { if s.Pool == nil { return 0, fmt.Errorf("%w: probe store has no database pool", ErrInvalidConfig) } if len(results) > MaxProbeBatch { return 0, fmt.Errorf("probe result batch of %d exceeds the %d bound", len(results), MaxProbeBatch) } if len(results) == 0 { return 0, nil } tx, err := s.Pool.BeginTx(ctx, pgx.TxOptions{}) if err != nil { return 0, fmt.Errorf("begin probe result save: %w", err) } defer func() { _ = tx.Rollback(ctx) }() saved := 0 for _, result := range results { attributes, err := json.Marshal(nonNilAttributes(result.Attributes)) if err != nil { return saved, fmt.Errorf("encode probe attributes: %w", err) } observedAt := result.ObservedAt.UTC() completedAt := result.CompletedAt.UTC() if completedAt.IsZero() { completedAt = observedAt } var inserted string err = tx.QueryRow(ctx, `INSERT INTO probe_results (id,probe_id,source_id,observed_at,completed_at,state,response_time_ms,status_code,error_class,error_message,attributes) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,NULLIF($9,''),NULLIF($10,''),$11::jsonb) ON CONFLICT (probe_id,observed_at) DO NOTHING RETURNING id`, newID(), result.ProbeID, nullableUUID(s.SourceID), observedAt, completedAt, result.State, result.ResponseTimeMS, result.StatusCode, boundedCode(result.ErrorClass), boundedText(result.ErrorMessage, 500), attributes).Scan(&inserted) switch { case err == nil: saved++ case errors.Is(err, pgx.ErrNoRows): // The result for this probe and observation time already exists. default: return saved, fmt.Errorf("insert probe result: %w", err) } if result.Certificate == nil { continue } certificate := result.Certificate if _, err := tx.Exec(ctx, `INSERT INTO service_certificates (id,service_id,endpoint_id,source_id,observed_at,expires_at,issuer,subject,hostname_valid,verification_state) VALUES ($1,$2,$3,$4,$5,$6,NULLIF($7,''),NULLIF($8,''),$9,$10) ON CONFLICT DO NOTHING`, newID(), certificate.ServiceID, nullableUUID(certificate.EndpointID), nullableUUID(s.SourceID), certificate.ObservedAt.UTC(), certificate.ExpiresAt, boundedText(certificate.Issuer, 500), boundedText(certificate.Subject, 500), certificate.HostnameValid, certificate.VerificationState); err != nil { return saved, fmt.Errorf("insert service certificate: %w", err) } } if err := tx.Commit(ctx); err != nil { return saved, fmt.Errorf("commit probe results: %w", err) } return saved, nil } func nonNilAttributes(value map[string]any) map[string]any { if value == nil { return map[string]any{} } return value } func nullableUUID(value string) any { if value == "" { return nil } return value } func boundedText(value string, max int) string { if len(value) > max { return value[:max] } return value }