package discovery import ( "context" "crypto/rand" "crypto/sha256" "encoding/hex" "encoding/json" "errors" "fmt" "regexp" "strings" "sync" "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) // DefaultJobType is the job_runs.job_type discovery claims are recorded under. const DefaultJobType = "discovery" const ( defaultWindow = time.Minute defaultLeaseTTL = 5 * time.Minute maxEventType = 160 maxSummary = 500 maxDedupKey = 255 ) var ( // ErrLeaseLost reports that the job run this process claimed was taken over // or completed elsewhere, so its result must not overwrite the newer one. ErrLeaseLost = errors.New("discovery job lease was lost") // ErrInvalidEvent reports an event that cannot be persisted safely. ErrInvalidEvent = errors.New("invalid discovery event") // ErrUnavailable reports a store without a usable database pool. ErrUnavailable = errors.New("discovery store is unavailable") uuidPattern = regexp.MustCompile("^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-8][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$") allowedSeverity = map[string]struct{}{"info": {}, "attention": {}, "warning": {}, "critical": {}} ) // PostgresStore persists discovery job runs and discovered events. // // It deliberately mirrors alertworker.PostgresLeaseStore: a claim is a row in // job_runs guarded by (job_type, job_key, scheduled_at) with a bounded // lease_owner/lease_until pair, so two workers never run the same discovery // window twice and a worker that crashes mid-run has its lease reclaimed once // it expires instead of blocking discovery forever. Events are inserted with // ON CONFLICT DO NOTHING against the natural (source_id, dedup_key, // occurred_at) key, which is what makes a repeated pass idempotent. type PostgresStore struct { Pool *pgxpool.Pool // JobType overrides the job_runs.job_type value; DefaultJobType is used when empty. JobType string // Owner identifies this worker in job_runs.lease_owner. Owner string // Window is the schedule granularity a claim is truncated to. Two claims of // the same key inside one window are the same unit of work. Window time.Duration // LeaseTTL bounds how long a claim blocks another worker after a crash. LeaseTTL time.Duration // Now is injectable for tests; time.Now is used when nil. Now func() time.Time mu sync.Mutex claims map[string]time.Time } // NewPostgresStore validates the configuration and returns a ready store. func NewPostgresStore(pool *pgxpool.Pool, owner string) (*PostgresStore, error) { if pool == nil { return nil, ErrUnavailable } if strings.TrimSpace(owner) == "" || len(owner) > 120 { return nil, errors.New("discovery store requires a bounded owner") } return &PostgresStore{Pool: pool, JobType: DefaultJobType, Owner: owner, Window: defaultWindow, LeaseTTL: defaultLeaseTTL}, nil } func (s *PostgresStore) jobType() string { if strings.TrimSpace(s.JobType) == "" { return DefaultJobType } return s.JobType } func (s *PostgresStore) window() time.Duration { if s.Window <= 0 { return defaultWindow } return s.Window } func (s *PostgresStore) leaseTTL() time.Duration { if s.LeaseTTL <= 0 { return defaultLeaseTTL } return s.LeaseTTL } func (s *PostgresStore) now() time.Time { if s.Now != nil { return s.Now().UTC() } return time.Now().UTC() } // Claim reserves the discovery window that contains at for this worker. It // returns false without an error when another worker holds a live lease or the // window already completed, which is the normal "nothing to do" outcome. func (s *PostgresStore) Claim(ctx context.Context, jobKey string, at time.Time) (bool, error) { if s == nil || s.Pool == nil { return false, ErrUnavailable } if strings.TrimSpace(jobKey) == "" || len(jobKey) > 255 { return false, errors.New("discovery job key is invalid") } if at.IsZero() { at = s.now() } now := s.now() scheduledAt := at.UTC().Truncate(s.window()) leaseUntil := now.Add(s.leaseTTL()) tx, err := s.Pool.BeginTx(ctx, pgx.TxOptions{}) if err != nil { return false, fmt.Errorf("begin discovery claim: %w", err) } defer func() { _ = tx.Rollback(ctx) }() var insertedID string err = tx.QueryRow(ctx, `INSERT INTO job_runs (id,job_type,job_key,scheduled_at,started_at,status,lease_owner,lease_until) VALUES ($1,$2,$3,$4,$5,'running',$6,$7) ON CONFLICT (job_type,job_key,scheduled_at) DO NOTHING RETURNING id`, newID(), s.jobType(), jobKey, scheduledAt, now, s.Owner, leaseUntil).Scan(&insertedID) if err == nil { if err := tx.Commit(ctx); err != nil { return false, fmt.Errorf("commit discovery claim: %w", err) } s.rememberClaim(jobKey, scheduledAt) return true, nil } if !errors.Is(err, pgx.ErrNoRows) { return false, fmt.Errorf("insert discovery claim: %w", err) } var status string var existingUntil *time.Time if err := tx.QueryRow(ctx, `SELECT status,lease_until FROM job_runs WHERE job_type=$1 AND job_key=$2 AND scheduled_at=$3 FOR UPDATE`, s.jobType(), jobKey, scheduledAt).Scan(&status, &existingUntil); err != nil { return false, fmt.Errorf("read discovery claim: %w", err) } if status != "running" || (existingUntil != nil && existingUntil.After(now)) { if err := tx.Commit(ctx); err != nil { return false, fmt.Errorf("commit discovery claim contention: %w", err) } return false, nil } tag, err := tx.Exec(ctx, `UPDATE job_runs SET started_at=$1,completed_at=NULL,error_code=NULL,lease_owner=$2,lease_until=$3 WHERE job_type=$4 AND job_key=$5 AND scheduled_at=$6 AND status='running' AND (lease_until IS NULL OR lease_until <= $7)`, now, s.Owner, leaseUntil, s.jobType(), jobKey, scheduledAt, now) if err != nil { return false, fmt.Errorf("reclaim discovery lease: %w", err) } if tag.RowsAffected() != 1 { if err := tx.Commit(ctx); err != nil { return false, fmt.Errorf("commit discovery reclaim contention: %w", err) } return false, nil } if err := tx.Commit(ctx); err != nil { return false, fmt.Errorf("commit discovery reclaim: %w", err) } s.rememberClaim(jobKey, scheduledAt) return true, nil } // Finish records the outcome of a claimed run. It only updates the row this // worker still owns, so a run whose lease expired cannot overwrite the result // of the worker that took over. func (s *PostgresStore) Finish(ctx context.Context, run JobRun) error { if s == nil || s.Pool == nil { return ErrUnavailable } if strings.TrimSpace(run.Key) == "" { return errors.New("discovery job key is required") } scheduledAt, ok := s.takeClaim(run.Key) if !ok { reference := run.StartedAt if reference.IsZero() { reference = s.now() } scheduledAt = reference.UTC().Truncate(s.window()) } status := "completed" if run.Status != "succeeded" && run.Status != "completed" { status = "failed" } errorCode := boundedText(run.ErrorCode, 160) counts, err := json.Marshal(map[string]any{"status": status, "attempts": run.Attempts}) if err != nil { return fmt.Errorf("encode discovery counts: %w", err) } completedAt := run.CompletedAt if completedAt.IsZero() { completedAt = s.now() } tag, err := s.Pool.Exec(ctx, `UPDATE job_runs SET status=$1,completed_at=$2,counts=$3::jsonb,error_code=NULLIF($4,''),lease_owner=NULL,lease_until=NULL WHERE job_type=$5 AND job_key=$6 AND scheduled_at=$7 AND lease_owner=$8 AND status='running'`, status, completedAt.UTC(), counts, errorCode, s.jobType(), run.Key, scheduledAt, s.Owner) if err != nil { return fmt.Errorf("complete discovery job: %w", err) } if tag.RowsAffected() != 1 { return ErrLeaseLost } return nil } // Emit persists one discovered event and reports whether it was new. A repeated // pass emitting the same (source, dedup key, occurrence time) is a no-op. func (s *PostgresStore) Emit(ctx context.Context, event Event) (bool, error) { if s == nil || s.Pool == nil { return false, ErrUnavailable } event, err := normalizeEvent(event) if err != nil { return false, err } attributes, err := json.Marshal(map[string]any{}) if err != nil { return false, fmt.Errorf("encode discovery attributes: %w", err) } var inserted string err = s.Pool.QueryRow(ctx, `INSERT INTO events (id,event_type,severity,entity_id,source_id,occurred_at,received_at,dedup_key,summary,attributes) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10::jsonb) ON CONFLICT (source_id,dedup_key,occurred_at) DO NOTHING RETURNING id`, DeterministicEventID(event), event.Type, event.Severity, nullableUUID(event.EntityID), event.SourceID, event.OccurredAt, s.now(), event.DedupKey, event.Summary, attributes).Scan(&inserted) if errors.Is(err, pgx.ErrNoRows) { return false, nil } if err != nil { return false, fmt.Errorf("emit discovery event: %w", err) } return true, nil } // DeterministicEventID derives a stable UUID from the event identity so a // retried emit reuses the same primary key instead of racing on a new one. func DeterministicEventID(event Event) string { sum := sha256.Sum256([]byte("itworx-pulse/discovery-event/v1/" + event.SourceID + "\x00" + event.DedupKey + "\x00" + event.OccurredAt.UTC().Format(time.RFC3339Nano))) return formatUUID(sum[:16], 0x50) } func normalizeEvent(event Event) (Event, error) { event.SourceID = strings.TrimSpace(event.SourceID) event.DedupKey = strings.TrimSpace(event.DedupKey) event.Type = strings.TrimSpace(event.Type) event.Summary = strings.TrimSpace(event.Summary) event.EntityID = strings.TrimSpace(event.EntityID) event.Severity = strings.ToLower(strings.TrimSpace(event.Severity)) if event.Severity == "" { event.Severity = "info" } if !uuidPattern.MatchString(event.SourceID) { return Event{}, fmt.Errorf("%w: source id must be a registered data source UUID", ErrInvalidEvent) } if event.EntityID != "" && !uuidPattern.MatchString(event.EntityID) { return Event{}, fmt.Errorf("%w: entity id must be a UUID", ErrInvalidEvent) } if _, ok := allowedSeverity[event.Severity]; !ok { return Event{}, fmt.Errorf("%w: severity is unsupported", ErrInvalidEvent) } if event.DedupKey == "" || len(event.DedupKey) > maxDedupKey { return Event{}, fmt.Errorf("%w: dedup key must be 1-%d characters", ErrInvalidEvent, maxDedupKey) } if event.Type == "" || len(event.Type) > maxEventType { return Event{}, fmt.Errorf("%w: type must be 1-%d characters", ErrInvalidEvent, maxEventType) } if event.Summary == "" || len(event.Summary) > maxSummary { return Event{}, fmt.Errorf("%w: summary must be 1-%d characters", ErrInvalidEvent, maxSummary) } if event.OccurredAt.IsZero() { return Event{}, fmt.Errorf("%w: occurrence time is required", ErrInvalidEvent) } event.OccurredAt = event.OccurredAt.UTC() return event, nil } func (s *PostgresStore) rememberClaim(jobKey string, scheduledAt time.Time) { s.mu.Lock() defer s.mu.Unlock() if s.claims == nil { s.claims = make(map[string]time.Time) } s.claims[jobKey] = scheduledAt } func (s *PostgresStore) takeClaim(jobKey string) (time.Time, bool) { s.mu.Lock() defer s.mu.Unlock() scheduledAt, ok := s.claims[jobKey] if ok { delete(s.claims, jobKey) } return scheduledAt, ok } func nullableUUID(value string) any { if value == "" { return nil } return value } func boundedText(value string, max int) string { value = strings.TrimSpace(value) if len(value) > max { return value[:max] } return value } func newID() string { raw := make([]byte, 16) if _, err := rand.Read(raw); err != nil { sum := sha256.Sum256([]byte(fmt.Sprintf("itworx-pulse/discovery-run/v1/%d", time.Now().UnixNano()))) copy(raw, sum[:16]) } return formatUUID(raw, 0x40) } func formatUUID(raw []byte, version byte) string { b := make([]byte, 16) copy(b, raw) b[6] = (b[6] & 0x0f) | version b[8] = (b[8] & 0x3f) | 0x80 return fmt.Sprintf("%s-%s-%s-%s-%s", hex.EncodeToString(b[0:4]), hex.EncodeToString(b[4:6]), hex.EncodeToString(b[6:8]), hex.EncodeToString(b[8:10]), hex.EncodeToString(b[10:16])) }