dispatch

type ActiveDistributedRun struct {
	AttemptKey string       `json:"attemptKey"`
	DAGRun     ir.DAGRunRef `json:"dagRun"`
	Root       ir.DAGRunRef `json:"root,omitzero"`
	AttemptID  string       `json:"attemptId"`
	WorkerID   string       `json:"workerId"`
	Status     ir.Status    `json:"status"`
	UpdatedAt  int64        `json:"updatedAt"`
}

type ActiveDistributedRunStore interface {
	Upsert(ctx context.Context, record ActiveDistributedRun) error
	Delete(ctx context.Context, attemptKey string) error
	Get(ctx context.Context, attemptKey string) (*ActiveDistributedRun, error)
	ListAll(ctx context.Context) ([]ActiveDistributedRun, error)
}

type ClaimedDispatchTask struct {
	Task       *DispatchTask
	ClaimToken string
	ClaimedAt  time.Time
	WorkerID   string
	PollerID   string
	Owner      CoordinatorEndpoint
}

type CoordinatorEndpoint struct {
	ID   string `json:"id,omitempty"`
	Host string `json:"host,omitempty"`
	Port int    `json:"port,omitempty"`
}

type DAGRunLease struct {
	AttemptKey            string              `json:"attemptKey"`
	DAGRun                ir.DAGRunRef        `json:"dagRun"`
	Root                  ir.DAGRunRef        `json:"root,omitzero"`
	AttemptID             string              `json:"attemptId"`
	ProfileName           string              `json:"profileName,omitempty"`
	QueueName             string              `json:"queueName"`
	WorkerID              string              `json:"workerId"`
	Owner                 CoordinatorEndpoint `json:"owner"`
	ClaimToken            string              `json:"claimToken,omitempty"`
	WorkspaceBundleDigest string              `json:"workspaceBundleDigest,omitempty"`
	ClaimedAt             int64               `json:"claimedAt"`
	LastHeartbeatAt       int64               `json:"lastHeartbeatAt"`
}

type DAGRunLeaseStore interface {
	Upsert(ctx context.Context, lease DAGRunLease) error
	Touch(ctx context.Context, attemptKey string, observedAt time.Time) error
	Delete(ctx context.Context, attemptKey string) error
	Get(ctx context.Context, attemptKey string) (*DAGRunLease, error)
	ListByQueue(ctx context.Context, queueName string) ([]DAGRunLease, error)
	ListAll(ctx context.Context) ([]DAGRunLease, error)
}

type DAGRunStatusResult struct {
	Found  bool
	Status *ir.DAGRunStatus
}

type DefinitionError struct {
	Err error
}

type DispatchAdmissionBindRequest struct {
	ReservationToken string
	Task             *DispatchTask
}

type DispatchAdmissionDecision struct {
	Reserved         bool
	Reason           DispatchAdmissionRejectReason
	ReservationToken string
}

type DispatchAdmissionRejectReason string

type DispatchAdmissionRequest struct {
	QueueName             string
	MaxConcurrency        int
	NonAdmissionOccupancy int
	AttemptKey            string
	AttemptID             string
	DAGRun                ir.DAGRunRef
	StaleThreshold        time.Duration
}

type DispatchAdmissionStore interface {
	ReserveAdmission(ctx context.Context, req DispatchAdmissionRequest) (*DispatchAdmissionDecision, error)
	BindAdmission(ctx context.Context, req DispatchAdmissionBindRequest) error
	ReleaseAdmissionToken(ctx context.Context, reservationToken string) error
	FinalizeAdmissionAttempt(ctx context.Context, attemptKey string) error
	CleanupAdmissions(ctx context.Context, staleThreshold time.Duration) error
}

type DispatchOperation int32

type DispatchRequest struct {
	Task                      *DispatchTask
	AdmissionReservationToken string
}

type DispatchTask struct {
	RootDAGRunName string
	RootDAGRunID   string

	ParentDAGRunName string
	ParentDAGRunID   string

	Operation    DispatchOperation
	DAGRunID     string
	Target       string
	Definition   string
	AttemptID    string
	AttemptKey   string
	Step         string
	Params       string
	ParallelItem string
	// PassedEnv carries resolved "KEY=value" pairs the parent opted to share
	// with the child run via the step's pass_env field.
	PassedEnv      []string
	QueueName      string
	WorkerID       string
	TargetWorkerID string
	ProfileName    string
	DefinitionID   string
	TriggerActor   string

	PreviousStatus *ir.DAGRunStatus

	BaseConfig          string
	BaseConfigWorkspace *string
	Labels              string
	ScheduleTime        string
	SourceFile          string
	SourceWorkDir       string

	WorkerSelector map[string]string

	// HoldUntil is set when the task was accepted with no healthy worker
	// matching it. Until then the task waits for one to register; after it
	// the coordinator fails the run instead.
	HoldUntil time.Time

	ExternalStepRetry   bool
	IncludeDownstream   bool
	BypassPreconditions bool
	RetryPath           string

	WorkspaceBundleDigest      string
	WorkspaceBundleSize        int64
	WorkspaceBundleDAGPath     string
	WorkspaceBundleOriginalRef string
	WorkspaceBundleResolvedRef string

	Owner      CoordinatorEndpoint
	ClaimToken string
}

type DispatchTaskClaim struct {
	WorkerID     string
	PollerID     string
	Labels       map[string]string
	Owner        CoordinatorEndpoint
	ClaimTimeout time.Duration
}

type DispatchTaskStore interface {
	Enqueue(ctx context.Context, task *DispatchTask) error
	ClaimNext(ctx context.Context, claim DispatchTaskClaim) (*ClaimedDispatchTask, error)
	GetClaim(ctx context.Context, claimToken string) (*ClaimedDispatchTask, error)
	ReleaseClaim(ctx context.Context, claimToken string) error
	DeleteClaim(ctx context.Context, claimToken string) error
	ListBundleDigests(ctx context.Context) ([]string, error)
	CountOutstandingByQueue(ctx context.Context, queueName string, claimTimeout time.Duration) (int, error)
	HasOutstandingAttempt(ctx context.Context, attemptKey string, claimTimeout time.Duration) (bool, error)
	// ListExpiredHeld returns the pending tasks whose HoldUntil has passed at
	// now. They are no longer claimable; the coordinator fails their runs and
	// removes them with DeletePending.
	ListExpiredHeld(ctx context.Context, now time.Time) ([]*DispatchTask, error)
	// ListStalePending returns the pending tasks, held ones aside, enqueued
	// before cutoff. They stay claimable; the coordinator removes, with
	// DeletePending, the ones whose runs no longer await them.
	ListStalePending(ctx context.Context, cutoff time.Time) ([]*DispatchTask, error)
	// DeletePending removes the pending task of an attempt, if any. A task
	// already claimed is left alone.
	DeletePending(ctx context.Context, attemptKey string) error
}

type Dispatcher interface {
	Dispatch(ctx context.Context, req DispatchRequest) error
	Cleanup(ctx context.Context) error
	GetDAGRunStatus(ctx context.Context, dagName, dagRunID string, rootRef *ir.DAGRunRef) (*DAGRunStatusResult, error)
	RequestCancel(ctx context.Context, dagName, dagRunID string, rootRef *ir.DAGRunRef) error
}

type RunningTask struct {
	DAGRunID         string `json:"dagRunId,omitempty"`
	DAGName          string `json:"dagName,omitempty"`
	StartedAt        int64  `json:"startedAt,omitempty"`
	RootDAGRunName   string `json:"rootDagRunName,omitempty"`
	RootDAGRunID     string `json:"rootDagRunId,omitempty"`
	ParentDAGRunName string `json:"parentDagRunName,omitempty"`
	ParentDAGRunID   string `json:"parentDagRunId,omitempty"`
	AttemptKey       string `json:"attemptKey,omitempty"`
}

type WorkerHeartbeatRecord struct {
	WorkerID        string            `json:"workerId"`
	Labels          map[string]string `json:"labels,omitempty"`
	Stats           *WorkerStats      `json:"stats,omitempty"`
	LastHeartbeatAt int64             `json:"lastHeartbeatAt"`
}

type WorkerHeartbeatStore interface {
	Upsert(ctx context.Context, record WorkerHeartbeatRecord) error
	Get(ctx context.Context, workerID string) (*WorkerHeartbeatRecord, error)
	List(ctx context.Context) ([]WorkerHeartbeatRecord, error)
	DeleteStale(ctx context.Context, before time.Time) (int, error)
}

type WorkerStats struct {
	TotalPollers int32          `json:"totalPollers,omitempty"`
	BusyPollers  int32          `json:"busyPollers,omitempty"`
	RunningTasks []*RunningTask `json:"runningTasks,omitempty"`
}