ladill-hosting/app/Services/Hosting/ServerTaskDispatchService.php

<?php

namespace App\Services\Hosting;

use App\Models\CustomerHostingOrder;
use App\Models\ServerAgent;
use App\Models\ServerTask;
use Illuminate\Support\Collection;
use Illuminate\Support\Facades\DB;
use Illuminate\Support\Str;

class ServerTaskDispatchService
{
    public function __construct(
        private ServerTaskResultService $resultService,
    ) {}

    /**
     * @return Collection<int, ServerTask>
     */
    public function claimTasks(ServerAgent $agent, int $limit = 10): Collection
    {
        $this->recoverStaleTasks($agent->customer_hosting_order_id);

        $claimedIds = DB::transaction(function () use ($agent, $limit): array {
            $taskIds = ServerTask::query()
                ->where('customer_hosting_order_id', $agent->customer_hosting_order_id)
                ->where('status', ServerTask::STATUS_QUEUED)
                ->where(function ($query): void {
                    $query->whereNull('available_at')
                        ->orWhere('available_at', '<=', now());
                })
                ->whereNull('lock_token')
                ->orderBy('id')
                ->limit($limit)
                ->lockForUpdate()
                ->pluck('id')
                ->all();

            $claimed = [];

            foreach ($taskIds as $taskId) {
                $lockToken = Str::random(48);

                $updated = ServerTask::query()
                    ->whereKey($taskId)
                    ->whereNull('lock_token')
                    ->update([
                        'server_agent_id' => $agent->id,
                        'status' => ServerTask::STATUS_DISPATCHED,
                        'lock_token' => $lockToken,
                        'locked_at' => now(),
                        'attempt_count' => DB::raw('attempt_count + 1'),
                        'started_at' => DB::raw('COALESCE(started_at, CURRENT_TIMESTAMP)'),
                        'last_heartbeat_at' => now(),
                    ]);

                if ($updated) {
                    $claimed[] = $taskId;
                }
            }

            return $claimed;
        });

        return ServerTask::query()->whereIn('id', $claimedIds)->orderBy('id')->get();
    }

    public function markProgress(ServerTask $task, string $lockToken, int $progress, ?string $message = null, array $result = []): ServerTask
    {
        $task = $this->assertLock($task, $lockToken);

        $task->update([
            'status' => ServerTask::STATUS_RUNNING,
            'progress' => $progress,
            'result' => $this->mergedTaskResult($task, $message, $result),
            'started_at' => $task->started_at ?? now(),
            'last_heartbeat_at' => now(),
            'locked_at' => now(),
        ]);

        return $task->fresh();
    }

    public function markCompleted(ServerTask $task, string $lockToken, ?string $message = null, array $result = []): ServerTask
    {
        $task = $this->assertLock($task, $lockToken);

        $task->update([
            'status' => ServerTask::STATUS_COMPLETED,
            'progress' => 100,
            'result' => $this->mergedTaskResult($task, $message, $result),
            'completed_at' => now(),
            'last_heartbeat_at' => now(),
            'error_message' => null,
            'lock_token' => null,
            'locked_at' => null,
        ]);

        return $this->resultService->apply($task->fresh());
    }

    /**
     * @return array{task: ServerTask, requeued: bool}
     */
    public function markFailed(ServerTask $task, string $lockToken, string $message, array $result = [], bool $retryable = true): array
    {
        $task = $this->assertLock($task, $lockToken);
        $mergedResult = $this->mergedTaskResult($task, $message, $result);

        if ($task->canRetry($retryable)) {
            $task->update([
                'status' => ServerTask::STATUS_QUEUED,
                'progress' => 0,
                'available_at' => now()->addSeconds(max(15, (int) $task->retry_backoff_seconds)),
                'failed_at' => null,
                'last_heartbeat_at' => now(),
                'error_message' => $message,
                'result' => array_merge($mergedResult, [
                    'retrying' => true,
                    'retry_scheduled_at' => now()->addSeconds(max(15, (int) $task->retry_backoff_seconds))->toIso8601String(),
                ]),
                'lock_token' => null,
                'locked_at' => null,
            ]);

            return [
                'task' => $task->fresh(),
                'requeued' => true,
            ];
        }

        $task->update([
            'status' => ServerTask::STATUS_FAILED,
            'failed_at' => now(),
            'last_heartbeat_at' => now(),
            'error_message' => $message,
            'result' => array_merge($mergedResult, ['retrying' => false]),
            'lock_token' => null,
            'locked_at' => null,
        ]);

        return [
            'task' => $task->fresh(),
            'requeued' => false,
        ];
    }

    public function recoverStaleTasks(int|CustomerHostingOrder $order): int
    {
        $orderId = $order instanceof CustomerHostingOrder ? $order->id : $order;
        $tasks = ServerTask::query()
            ->where('customer_hosting_order_id', $orderId)
            ->whereIn('status', [ServerTask::STATUS_DISPATCHED, ServerTask::STATUS_RUNNING])
            ->get();

        $recovered = 0;

        foreach ($tasks as $task) {
            $referenceTime = $task->last_heartbeat_at ?? $task->locked_at ?? $task->updated_at;
            if ($referenceTime === null) {
                continue;
            }

            if ($referenceTime->gt(now()->subSeconds(max(30, (int) $task->stale_after_seconds)))) {
                continue;
            }

            if ($task->canRetry()) {
                $task->update([
                    'status' => ServerTask::STATUS_QUEUED,
                    'progress' => 0,
                    'server_agent_id' => null,
                    'available_at' => now()->addSeconds(max(15, (int) $task->retry_backoff_seconds)),
                    'error_message' => 'Recovered after agent heartbeat timeout.',
                    'result' => array_merge((array) ($task->result ?? []), [
                        'recovered_from_stale_lock' => true,
                        'stale_recovered_at' => now()->toIso8601String(),
                    ]),
                    'lock_token' => null,
                    'locked_at' => null,
                    'last_heartbeat_at' => null,
                ]);
            } else {
                $task->update([
                    'status' => ServerTask::STATUS_FAILED,
                    'failed_at' => now(),
                    'error_message' => 'Task failed after stale lock recovery exhausted its retry budget.',
                    'result' => array_merge((array) ($task->result ?? []), [
                        'recovered_from_stale_lock' => true,
                        'stale_recovered_at' => now()->toIso8601String(),
                        'retrying' => false,
                    ]),
                    'lock_token' => null,
                    'locked_at' => null,
                ]);
            }

            $recovered++;
        }

        return $recovered;
    }

    private function assertLock(ServerTask $task, string $lockToken): ServerTask
    {
        $task->refresh();
        abort_if($task->lock_token === null, 409, 'Task lock has expired.');
        abort_unless(hash_equals((string) $task->lock_token, $lockToken), 409, 'Task lock token mismatch.');

        return $task;
    }

    /**
     * @param  array<string, mixed>  $result
     * @return array<string, mixed>
     */
    private function mergedTaskResult(ServerTask $task, ?string $message, array $result): array
    {
        return array_filter(array_merge((array) ($task->result ?? []), $result, [
            'message' => $message,
        ]), static fn ($value): bool => $value !== null);
    }
}