ITWorx-Pulse-Public/internal/workerruntime/probejob.go

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
}