SciMesh/coordinator/internal/domain/task.go

// Package domain holds SciMesh's entities and the rules that govern them. It
// is the innermost layer: it imports nothing from this module and knows nothing
// about HTTP, SQL, or configuration. Every state transition a task can undergo
// is a method here, so the rules are unit-testable without a database.
package domain

import (
	"time"

	"github.com/google/uuid"
)

type TaskStatus string

const (
	TaskPending   TaskStatus = "pending"
	TaskLeased    TaskStatus = "leased"
	TaskRunning   TaskStatus = "running"
	TaskCompleted TaskStatus = "completed"
	TaskFailed    TaskStatus = "failed"
	TaskCancelled TaskStatus = "cancelled"
)

// ErrCodeLeaseExpired marks tasks failed by the reaper rather than by a worker.
const ErrCodeLeaseExpired = "lease_expired"

// Task is one independently executable chunk of a job.
//
// Nullable columns are pointers so "no lease" stays distinguishable from
// "lease owned by the empty string" — a plain string cannot express both.
type Task struct {
	ID               uuid.UUID
	JobID            uuid.UUID
	ChunkIndex       int
	Workload         string
	InputURI         string     // external input URI; empty for uploaded shards
	InputArtifactID  *uuid.UUID // coordinator-stored shard; nil for URI inputs
	InputSHA256      string
	Parameters       map[string]any
	Status           TaskStatus
	Attempt          int
	MaxAttempts      int
	LeaseOwner       *string
	LeaseExpiresAt   *time.Time
	ResultArtifactID *uuid.UUID
	Metrics          map[string]any
	ErrorCode        *string
	ErrorMessage     *string
	CreatedAt        time.Time
	StartedAt        *time.Time
	CompletedAt      *time.Time
	Version          int
}

// NewTask builds a pending task. maxAttempts <= 0 falls back to the default.
func NewTask(jobID uuid.UUID, chunkIndex int, workload, inputURI, inputSHA256 string,
	params map[string]any, maxAttempts int, now time.Time) (*Task, error) {

	if inputURI == "" {
		return nil, ErrInvalidInput
	}
	if inputSHA256 == "" {
		return nil, ErrInvalidInput // checksum is mandatory: workers verify inputs
	}
	if chunkIndex < 0 {
		return nil, ErrInvalidInput
	}
	if maxAttempts <= 0 {
		maxAttempts = DefaultMaxAttempts
	}
	return &Task{
		ID:          uuid.New(),
		JobID:       jobID,
		ChunkIndex:  chunkIndex,
		Workload:    workload,
		InputURI:    inputURI,
		InputSHA256: inputSHA256,
		Parameters:  params,
		Status:      TaskPending,
		Attempt:     0,
		MaxAttempts: maxAttempts,
		CreatedAt:   now,
	}, nil
}

// NewShardTask builds a pending task whose input is a coordinator-stored shard
// artifact rather than an external URI. The worker fetches it from the
// coordinator, so no InputURI is set — inputSHA256 is the shard's checksum.
func NewShardTask(jobID uuid.UUID, chunkIndex int, workload string, inputArtifactID uuid.UUID,
	inputSHA256 string, params map[string]any, maxAttempts int, now time.Time) (*Task, error) {

	if inputArtifactID == uuid.Nil || inputSHA256 == "" || chunkIndex < 0 {
		return nil, ErrInvalidInput
	}
	if maxAttempts <= 0 {
		maxAttempts = DefaultMaxAttempts
	}
	return &Task{
		ID:              uuid.New(),
		JobID:           jobID,
		ChunkIndex:      chunkIndex,
		Workload:        workload,
		InputArtifactID: &inputArtifactID,
		InputSHA256:     inputSHA256,
		Parameters:      params,
		Status:          TaskPending,
		Attempt:         0,
		MaxAttempts:     maxAttempts,
		CreatedAt:       now,
	}, nil
}

// DefaultMaxAttempts applies when a task does not specify its own ceiling.
const DefaultMaxAttempts = 3

// CanRetry reports whether any attempts remain.
func (t *Task) CanRetry() bool { return t.Attempt < t.MaxAttempts }

// IsLeaseHeldBy reports whether worker currently holds this task at attempt.
func (t *Task) IsLeaseHeldBy(worker string, attempt int, now time.Time) bool {
	return t.LeaseOwner != nil && t.LeaseExpiresAt != nil && now.Before(*t.LeaseExpiresAt) &&
		*t.LeaseOwner == worker && t.Attempt == attempt &&
		(t.Status == TaskLeased || t.Status == TaskRunning)
}

// AsClaimed projects the task into the trimmed view handed to a worker:
// everything needed to execute, nothing it has no business seeing.
func (t *Task) AsClaimed() ClaimedTask {
	ct := ClaimedTask{
		TaskID:          t.ID,
		JobID:           t.JobID,
		ChunkIndex:      t.ChunkIndex,
		Workload:        t.Workload,
		InputURI:        t.InputURI,
		InputArtifactID: t.InputArtifactID,
		InputSHA256:     t.InputSHA256,
		Parameters:      t.Parameters,
		Attempt:         t.Attempt,
	}
	if t.LeaseOwner != nil {
		ct.LeaseOwner = *t.LeaseOwner
	}
	if t.LeaseExpiresAt != nil {
		ct.LeaseExpiresAt = *t.LeaseExpiresAt
	}
	return ct
}

// verifyLease is the guard every worker-driven transition shares: the caller
// must own the lease and reference the attempt it was granted.
func (t *Task) verifyLease(worker string, attempt int, now time.Time) error {
	// A task is worker-owned while leased or running: the first heartbeat moves
	// it from leased to running, but ownership rules are identical for both.
	if t.Status != TaskLeased && t.Status != TaskRunning {
		return ErrTaskNotLeased
	}
	if t.LeaseOwner == nil || *t.LeaseOwner != worker {
		return ErrLeaseConflict
	}
	if t.Attempt != attempt {
		return ErrStaleAttempt
	}
	if t.LeaseExpiresAt == nil || !now.Before(*t.LeaseExpiresAt) {
		return ErrLeaseConflict
	}
	return nil
}

// RenewLease extends the lease of the worker that holds it. The first heartbeat
// also acknowledges start, moving the task from leased to running.
func (t *Task) RenewLease(worker string, attempt int, now, until time.Time) error {
	if err := t.verifyLease(worker, attempt, now); err != nil {
		return err
	}
	t.LeaseExpiresAt = &until
	if t.Status == TaskLeased {
		t.Status = TaskRunning
	}
	t.Version++
	return nil
}

// CompleteWith records a successful result.
//
// Idempotency comes first deliberately: a worker whose network dropped will
// retry the same manifest, and that must succeed rather than trip the lease
// check on a task the coordinator already finished. A *different* manifest for
// an already-completed task is a genuine conflict.
func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any,
	worker string, attempt int, now time.Time) error {

	if resultArtifactID == uuid.Nil {
		return ErrInvalidInput
	}

	if t.Status == TaskCompleted {
		if t.Attempt == attempt && t.ResultArtifactID != nil && *t.ResultArtifactID == resultArtifactID {
			return nil // same attempt, same artifact — replay of a successful call
		}
		return ErrResultConflict
	}

	if err := t.verifyLease(worker, attempt, now); err != nil {
		return err
	}

	t.Status = TaskCompleted
	t.ResultArtifactID = &resultArtifactID
	t.Metrics = metrics
	t.CompletedAt = &now
	t.LeaseOwner = nil
	t.LeaseExpiresAt = nil
	t.ErrorCode = nil
	t.ErrorMessage = nil
	t.Version++
	return nil
}

// Fail records a worker-reported failure. A retryable failure with attempts
// left returns the task to the queue; otherwise it terminates as failed.
func (t *Task) Fail(worker string, attempt int, code, message string, retryable bool, now time.Time) error {
	if err := t.verifyLease(worker, attempt, now); err != nil {
		return err
	}
	t.ErrorCode = &code
	t.ErrorMessage = &message
	t.LeaseOwner = nil
	t.LeaseExpiresAt = nil
	t.Version++

	if retryable && t.CanRetry() {
		t.Status = TaskPending
		return nil
	}
	t.Status = TaskFailed
	t.CompletedAt = &now
	return nil
}

// ExpireLease is applied by the reaper when a lease elapses without a
// heartbeat: requeue while attempts remain, otherwise fail terminally.
func (t *Task) ExpireLease(now time.Time) {
	// Both a leased and a running task can go silent and must be reclaimed.
	if t.Status != TaskLeased && t.Status != TaskRunning {
		return
	}
	t.LeaseOwner = nil
	t.LeaseExpiresAt = nil
	t.Version++

	if t.CanRetry() {
		t.Status = TaskPending
		return
	}
	code, msg := ErrCodeLeaseExpired, "lease expired after the final attempt"
	t.ErrorCode = &code
	t.ErrorMessage = &msg
	t.Status = TaskFailed
	t.CompletedAt = &now
}

// Cancel prevents any further worker transition for a task that has not
// reached a terminal result. A cancelled lease deliberately becomes invalid:
// a worker still running locally must not upload or complete after its job was
// stopped by the operator.
func (t *Task) Cancel(now time.Time) bool {
	if t.Status == TaskCompleted || t.Status == TaskFailed || t.Status == TaskCancelled {
		return false
	}
	t.Status = TaskCancelled
	t.LeaseOwner = nil
	t.LeaseExpiresAt = nil
	t.ErrorCode = nil
	t.ErrorMessage = nil
	t.CompletedAt = &now
	t.Version++
	return true
}

// ClaimedTask is the worker-facing projection of a leased task. Input is either
// an external URI or a coordinator-stored shard (InputArtifactID set); the
// transport turns the latter into a coordinator download URL.
type ClaimedTask struct {
	TaskID          uuid.UUID
	JobID           uuid.UUID
	ChunkIndex      int
	Workload        string
	InputURI        string
	InputArtifactID *uuid.UUID
	InputSHA256     string
	Parameters      map[string]any
	Attempt         int
	LeaseOwner      string
	LeaseExpiresAt  time.Time
}

// ResultManifest is a completed task's output, ordered for the stitcher. It
// points at the coordinator-owned result artifact rather than a worker URI.
type ResultManifest struct {
	TaskID           uuid.UUID
	ChunkIndex       int
	ResultArtifactID uuid.UUID
	Metrics          map[string]any
}