worker/app/models/task_selector.go

package models

import (
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// DefaultCheckTaskLeaseTTL is how long a leased check remains owned before the reaper
// returns it to the queue. It is deliberately larger than the worker's heartbeat
// (10s) so a healthy worker can finish a 30s check without the reaper stealing
// the lease, and deliberately smaller than the heartbeat timeout (2m) so a
// crashed worker sees its tasks reaped within one watchdog cycle.
const DefaultCheckTaskLeaseTTL = 60 * time.Second

// DefaultTaskLeaseTTL remains the check-task default for existing confirmation
// callers. Generic task selection must use TaskLeaseTTL so notification work is
// not reclaimed during its longer execution window.
const DefaultTaskLeaseTTL = DefaultCheckTaskLeaseTTL

// Notification execution is bounded by the worker runner at five minutes. The
// additional minute covers result serialization, websocket/HTTP transport, and
// a control-plane scheduling delay before the persisted lease may be reaped.
const (
	DefaultNotificationExecutionTimeout = 5 * time.Minute
	NotificationTaskReportMargin        = time.Minute
	DefaultNotificationTaskLeaseTTL     = DefaultNotificationExecutionTimeout + NotificationTaskReportMargin
)

// TaskLeaseTTL returns the persisted lease lifetime for a task kind.
func TaskLeaseTTL(kind string) time.Duration {
	if kind == TaskKindNotification {
		return DefaultNotificationTaskLeaseTTL
	}
	return DefaultCheckTaskLeaseTTL
}

// DefaultTaskMaxAttempts is the retry budget for a task before it moves to dead.
const DefaultTaskMaxAttempts = 5

// DefaultNotificationTaskDeadline is assigned to manually replayed notification
// dead letters. Normal producer tasks may be deadline-free, but a replay must
// never inherit an already-expired deadline.
const DefaultNotificationTaskDeadline = 15 * time.Minute

// ErrNotificationMethodNotAuthorized is returned by EnqueueNotificationTask
// when the producer can prove no worker in the pool is authorized for the
// (method, account) pair. The caller may skip the enqueue or log + continue.
var ErrNotificationMethodNotAuthorized = errors.New("no worker authorized for method/account")

// EnqueueNotificationTaskInput is the pre-rendered envelope produced by the
// notifier producer. All slices are required; the selector never reads them.
type EnqueueNotificationTaskInput struct {
	AccountID      int64
	NotificationID int64
	ContactID      int64
	MessageID      *int64
	MonitorID      *int64
	CheckID        *int64
	EventIDs       []int64
	Method         string // "email", "telegram", "webhook", "mattermost", "sms", "voice"
	Subject        string
	BodyText       string
	BodyHTML       string
	BodyMarkdown   string
	Language       string
	MessageKind    string // "down", "up", "exp", "test"
	NotBefore      time.Time
	Deadline       *time.Time
	MaxAttempts    int
	Payload        []byte // marshaled task-specific data
	IdempotencyKey string // optional; manual test tasks use a unique key and do not have event IDs
}

// EnqueueNotificationTask writes one Task row keyed by a stable idempotency key.
// A second call with the same key (same notification/contact/event triple) is a
// no-op so the producer is safe to call more than once per pass.
//
// The capability precheck uses the same NotificationMethods + NotificationAccounts
// rule the selector does, so the producer can skip enqueueing work that no
// operated worker could ever pick up (sms/voice until phase 4).
func EnqueueNotificationTask(input *EnqueueNotificationTaskInput) (*Task, error) {
	return EnqueueNotificationTaskTx(DB(), input)
}

// EnqueueNotificationTaskTx is the transactional form used by state machines
// that must commit their transition, audit event, message, and task together.
func EnqueueNotificationTaskTx(tx *gorm.DB, input *EnqueueNotificationTaskInput) (*Task, error) {
	if tx == nil {
		return nil, errors.New("enqueue: nil transaction")
	}
	if input.AccountID == 0 || input.ContactID == 0 {
		return nil, errors.New("enqueue: account_id and contact_id are required")
	}

	idempotencyKey := input.IdempotencyKey
	if idempotencyKey == "" {
		if input.NotificationID == 0 || len(input.EventIDs) == 0 {
			return nil, errors.New("enqueue: notification_id and event_ids are required without explicit idempotency_key")
		}
		idempotencyKey = notificationIdempotencyKey(input.NotificationID, input.ContactID, input.EventIDs[0])
	}

	// Fast path: row already exists from a previous producer tick. Returning
	// the existing row is the idempotency guarantee — second calls return the
	// same id, no second INSERT.
	var existing Task
	if err := tx.Where("idempotency_key = ?", idempotencyKey).First(&existing).Error; err == nil {
		return &existing, nil
	} else if !errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, err
	}

	maxAttempts := input.MaxAttempts
	if maxAttempts <= 0 {
		maxAttempts = DefaultTaskMaxAttempts
	}
	if !anyWorkerCanDeliver(input.Method, input.AccountID) {
		return nil, fmt.Errorf("%w: method=%s account=%d", ErrNotificationMethodNotAuthorized, input.Method, input.AccountID)
	}

	notBefore := input.NotBefore
	if notBefore.IsZero() {
		notBefore = time.Now()
	}

	now := time.Now()
	task := &Task{
		JobID:          uuid.New().String(),
		Kind:           TaskKindNotification,
		State:          TaskStateQueued,
		AccountID:      input.AccountID,
		MessageID:      input.MessageID,
		ContactID:      &input.ContactID,
		MonitorID:      input.MonitorID,
		CheckID:        input.CheckID,
		NotBefore:      notBefore,
		Deadline:       input.Deadline,
		Attempts:       0,
		MaxAttempts:    maxAttempts,
		IdempotencyKey: idempotencyKey,
	}
	if len(input.Payload) > 0 {
		task.Payload = input.Payload
	}
	task.CreatedAt = now
	task.UpdatedAt = now

	// ON CONFLICT DO NOTHING so a concurrent producer tick racing with us on
	// the same idempotency_key loses the race but does not duplicate the row.
	if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(task).Error; err != nil {
		return nil, err
	}
	if task.ID == 0 {
		// Lost the race. Re-read and return the winner.
		if err := tx.Where("idempotency_key = ?", idempotencyKey).First(&existing).Error; err != nil {
			return nil, err
		}
		return &existing, nil
	}
	return task, nil
}

// NotificationIdempotencyKey exposes the producer's idempotency key so the
// result handler and admin tooling can match a Task row back to the logical
// (notification, contact, event) tuple without re-deriving the format.
func NotificationIdempotencyKey(notificationID, contactID, eventID int64) string {
	return notificationIdempotencyKey(notificationID, contactID, eventID)
}

func notificationIdempotencyKey(notificationID, contactID, eventID int64) string {
	return fmt.Sprintf("notif:%d:contact:%d:event:%d", notificationID, contactID, eventID)
}

// anyWorkerCanDeliver returns true if at least one active worker in the pool is
// authorized to deliver the given (method, account) pair. Used by the producer
// to skip enqueues no worker could ever pick up.
func anyWorkerCanDeliver(method string, accountID int64) bool {
	var nodes []WorkerNode
	if err := DB().Where("status = ? AND last_seen > ?", "active", time.Now().Add(-WorkerHeartbeatFreshness)).Find(&nodes).Error; err != nil {
		log.Printf("task_selector: cannot list workers: %v", err)
		// Be permissive on lookup failure: the selector's own filter would still
		// hold the lease back, so the worst case is a queued task nobody picks
		// up — which the reaper eventually dead-letters.
		return true
	}
	for i := range nodes {
		if nodes[i].SupportsTaskEnvelope() && nodes[i].CanDeliverNotification(method, accountID) {
			return true
		}
	}
	return false
}

// TasksForWorker leases up to `limit` due tasks for the worker. The selection
// is one transaction so the FOR UPDATE SKIP LOCKED + UPDATE that flips state
// from queued to leased is atomic. Notification tasks are filtered by the worker's
// notification_methods + notification_accounts capability set; check tasks are
// filtered by check_types in their payload.
//
// The function is safe to call from multiple goroutines for different workers.
// Two workers that hit the DB at the same time will see disjoint task sets.
func TasksForWorker(worker *WorkerNode, limit int) ([]Task, error) {
	if worker == nil {
		return nil, errors.New("TasksForWorker: worker is nil")
	}
	if !worker.SupportsTaskEnvelope() {
		return nil, nil
	}
	if limit <= 0 {
		limit = 1
	}

	notifMethods := worker.NotificationMethods()
	notifAccounts := worker.AccessibleAccountIDs()
	hasNotif := len(notifMethods) > 0

	tx := DB().Begin()
	if tx.Error != nil {
		return nil, tx.Error
	}
	defer func() {
		if r := recover(); r != nil {
			_ = tx.Rollback().Error
			panic(r)
		}
	}()

	now := time.Now()
	var out []Task

	// First pass: notification tasks the worker is authorized to deliver. We
	// also bump attempts and flip state to leased in the same row so the
	// outer selector+lease is atomic. The method filter is a JSONB extract on
	// payload->>'method' so a single worker query can target one method list.
	if hasNotif {
		notifQuery := tx.Clauses(SkipLockedClause).
			Where("state = ? AND kind = ?", TaskStateQueued, TaskKindNotification).
			Where("not_before <= ?", now).
			Where("(deadline IS NULL OR deadline > ?)", now).
			Where("payload->>'method' IN (?)", notifMethods).
			Where("payload->>'method' <> ''")
		if len(notifAccounts) > 0 {
			notifQuery = notifQuery.Where("account_id IN (?)", notifAccounts)
		}

		var picked []Task
		if err := notifQuery.Limit(limit).Find(&picked).Error; err != nil {
			_ = tx.Rollback().Error
			return nil, err
		}

		for i := range picked {
			row := picked[i]
			newAttempts := row.Attempts + 1
			leaseToken := uuid.NewString()
			leaseUntil := now.Add(TaskLeaseTTL(row.Kind))
			if err := tx.Model(&row).Updates(map[string]interface{}{
				colState:           TaskStateLeased,
				"lease_owner":      worker.WorkerID,
				"lease_expires_at": leaseUntil,
				"attempts":         newAttempts,
				"lease_token":      leaseToken,
				"updated_at":       now,
			}).Error; err != nil {
				_ = tx.Rollback().Error
				return nil, err
			}
			row.State = TaskStateLeased
			row.LeaseOwner = worker.WorkerID
			row.LeaseExpiresAt = &leaseUntil
			row.Attempts = newAttempts
			row.LeaseToken = leaseToken
			out = append(out, row)
		}
	}

	remaining := limit - len(out)
	if checkTypes := worker.CheckTypes(); remaining > 0 && len(checkTypes) > 0 {
		checkQuery := tx.Clauses(SkipLockedClause).
			Where("state = ? AND kind = ?", TaskStateQueued, TaskKindCheck).
			Where("not_before <= ?", now).
			Where("(deadline IS NULL OR deadline > ?)", now).
			Where("payload->>'kind' IN (?)", checkTypes)
		if accounts := worker.AccessibleAccountIDs(); len(accounts) > 0 {
			checkQuery = checkQuery.Where("account_id IN (?)", accounts)
		}
		var picked []Task
		if err := checkQuery.Limit(remaining).Find(&picked).Error; err != nil {
			_ = tx.Rollback().Error
			return nil, err
		}
		for i := range picked {
			row := picked[i]
			newAttempts := row.Attempts + 1
			leaseToken := uuid.NewString()
			leaseUntil := now.Add(TaskLeaseTTL(row.Kind))
			if err := tx.Model(&row).Updates(map[string]interface{}{
				colState: TaskStateLeased, "lease_owner": worker.WorkerID,
				"lease_expires_at": leaseUntil, "attempts": newAttempts, "lease_token": leaseToken, "updated_at": now,
			}).Error; err != nil {
				_ = tx.Rollback().Error
				return nil, err
			}
			row.State, row.LeaseOwner, row.LeaseExpiresAt, row.Attempts, row.LeaseToken = TaskStateLeased, worker.WorkerID, &leaseUntil, newAttempts, leaseToken
			out = append(out, row)
		}
	}

	if err := tx.Commit().Error; err != nil {
		return nil, err
	}
	return out, nil
}

// AvailableWorkerTaskCapacity returns unoccupied local worker slots. Durable
// unexpired leases and the worker's heartbeat-reported active/queued workload
// describe the same work from different sides, so the larger value is used to
// avoid both over-dispatching and double-counting a healthy worker.
func AvailableWorkerTaskCapacity(worker *WorkerNode) (int, error) {
	if worker == nil {
		return 0, errors.New("worker capacity: worker is nil")
	}
	concurrency := worker.Concurrency
	if concurrency < 1 {
		concurrency = 1
	}
	var leased int64
	if err := DB().Model(&Task{}).Where("state = ? AND lease_owner = ? AND lease_expires_at > ?", TaskStateLeased, worker.WorkerID, time.Now()).Count(&leased).Error; err != nil {
		return 0, err
	}
	var confirmationLeases int64
	if err := DB().Model(&CheckAttempt{}).Where("worker_node_id = ? AND kind = ? AND state = ? AND lease_expires_at > ?", worker.ID, AttemptKindConfirm, AttemptStateLeased, time.Now()).Count(&confirmationLeases).Error; err != nil {
		return 0, err
	}
	used := int(leased + confirmationLeases)
	if reported := worker.ReportedWorkload(); reported > used {
		used = reported
	}
	if used >= concurrency {
		return 0, nil
	}
	return concurrency - used, nil
}

// GetTaskByJobID returns one task row keyed by its unique job_id. The result
// handler uses this to validate that the incoming JobID exists and matches the
// calling worker before it mutates state.
func GetTaskByJobID(jobID string) (*Task, error) {
	if jobID == "" {
		return nil, errors.New("GetTaskByJobID: empty job_id")
	}
	var task Task
	if err := DB().Where("job_id = ?", jobID).First(&task).Error; err != nil {
		return nil, err
	}
	return &task, nil
}

// TasksForWorkerTx is the variant exposed for tests so a single SELECT inside a
// caller-provided transaction can be inspected without the auto-commit wrapper.
// Production code should use TasksForWorker.
func TasksForWorkerTx(tx *gorm.DB, worker *WorkerNode, limit int) ([]Task, error) {
	if tx == nil {
		return nil, errors.New("TasksForWorkerTx: nil tx")
	}
	notifMethods := worker.NotificationMethods()
	if len(notifMethods) == 0 {
		return nil, nil
	}
	now := time.Now()
	var out []Task
	// Use SKIP LOCKED to avoid contention between workers (mirror of ChecksForWorker).
	q := tx.Clauses(SkipLockedClause).
		Where("state = ? AND kind = ?", TaskStateQueued, TaskKindNotification).
		Where("not_before <= ?", now).
		Where("(deadline IS NULL OR deadline > ?)", now).
		Where("payload->>'method' IN (?)", notifMethods)
	if accounts := worker.AccessibleAccountIDs(); len(accounts) > 0 {
		q = q.Where("account_id IN (?)", accounts)
	}
	if err := q.Limit(limit).Find(&out).Error; err != nil {
		return nil, err
	}
	return out, nil
}