Allow bounded leased completion uploads during storage draining

.github/workflows/phpunit-feature.yml

@@ -215,7 +215,9 @@ jobs:
                 -e DW_TEST_S3_ACCESS_KEY_ID=dw-minio-access \
                 -e DW_TEST_S3_SECRET_ACCESS_KEY=dw-minio-secret-key \
                 -e DW_TEST_BACKUP_MYSQL_HOST="dw-server-backup-mysql-${{ github.run_id }}-${{ github.run_attempt }}-${{ github.job }}" \
+                -e DW_TEST_COMPLETION_MYSQL_HOST="dw-server-backup-mysql-${{ github.run_id }}-${{ github.run_attempt }}-${{ github.job }}" \
                 -e DW_TEST_BACKUP_PGSQL_HOST="dw-server-backup-pgsql-${{ github.run_id }}-${{ github.run_attempt }}-${{ github.job }}" \
+                -e DW_TEST_COMPLETION_PGSQL_HOST="dw-server-backup-pgsql-${{ github.run_id }}-${{ github.run_attempt }}-${{ github.job }}" \
                 -w /app \
                 --entrypoint sh \
                 durable-workflow-server-corpus-validator \

app/Http/Controllers/Api/RuntimeExternalPayloadController.php

@@ -7,6 +7,10 @@
 use App\Support\RuntimeExternalPayloadReference;
 use App\Support\RuntimeExternalPayloadRegistry;
 use App\Support\RuntimeExternalPayloadUploadBody;
+use App\Support\RuntimePayloadCompletionContext;
+use App\Support\RuntimePayloadCompletionUploads;
+use App\Support\StorageAdmissionPaused;
+use App\Support\StoragePressure;
 use Illuminate\Http\JsonResponse;
 use Illuminate\Http\Request;
 use Symfony\Component\HttpFoundation\Response;
@@ -62,12 +66,24 @@ public function store(Request $request): JsonResponse
         }
 
         $namespace = (string) $request->attributes->get('namespace', config('server.default_namespace'));
-        $reference = $this->registry->upload(
-            $namespace,
-            $data,
-            (string) $request->header('X-Durable-Workflow-Payload-Codec'),
-            strtolower((string) $request->header('X-Durable-Workflow-Payload-SHA256')),
-        );
+        $context = $request->attributes->get(RuntimePayloadCompletionContext::class);
+        $codec = (string) $request->header('X-Durable-Workflow-Payload-Codec');
+        $sha256 = strtolower((string) $request->header('X-Durable-Workflow-Payload-SHA256'));
+        try {
+            if ($context instanceof RuntimePayloadCompletionContext) {
+                $reference = app(RuntimePayloadCompletionUploads::class)->upload($namespace, $context, $data, $codec, $sha256, $this->registry);
+            } else {
+                // Pressure may change while an admitted HTTP upload is being read.
+                $snapshot = app(StoragePressure::class)->snapshot();
+                if ($snapshot['state'] !== 'normal') {
+                    throw new StorageAdmissionPaused($snapshot);
+                }
+                $reference = $this->registry->upload($namespace, $data, $codec, $sha256,
+                    fn () => app(StoragePressure::class)->requireNewWork());
+            }
+        } finally {
+            $data->close();
+        }
 
         $this->audit->record($request, 'external_payload.uploaded', [
             'reference_identity_sha256' => hash('sha256', $reference['reference_id']),

app/Http/Middleware/EnforceStorageAdmission.php

@@ -3,6 +3,8 @@
 namespace App\Http\Middleware;
 
 use App\Support\ControlPlaneProtocol;
+use App\Support\RuntimePayloadCompletionContext;
+use App\Support\RuntimePayloadCompletionUploads;
 use App\Support\StoragePressure;
 use App\Support\WorkerProtocol;
 use Closure;
@@ -39,6 +41,21 @@ public function handle(Request $request, Closure $next): Response
     {
         $snapshot = $this->pressure->snapshot();
         $action = class_basename((string) $request->route()?->getActionName());
+        if ($action === 'RuntimeExternalPayloadController@store'
+            && $request->headers->has(RuntimePayloadCompletionContext::HEADER)
+            && in_array($snapshot['state'], ['normal', 'draining'], true)) {
+            return app(RequireRole::class)->handle($request, function (Request $request) use ($next): Response {
+                $context = RuntimePayloadCompletionContext::parse((string) $request->header(RuntimePayloadCompletionContext::HEADER));
+                app(RuntimePayloadCompletionUploads::class)->authorize(
+                    (string) $request->attributes->get('namespace'), $context,
+                    strtolower((string) $request->header('X-Durable-Workflow-Payload-SHA256')),
+                    (int) $request->header('X-Durable-Workflow-Payload-Size'),
+                );
+                $request->attributes->set(RuntimePayloadCompletionContext::class, $context);
+
+                return $next($request);
+            }, 'worker');
+        }
         if ($snapshot['state'] === 'normal' || $request->isMethodSafe()
             || $action === 'WorkerController@workflowTaskHistory'
             || ($snapshot['state'] === 'draining' && in_array($action, self::DRAIN_ACTIONS, true))) {

app/Models/RuntimePayloadCompletionBudget.php

@@ -0,0 +1,18 @@
+<?php
+
+namespace App\Models;
+
+use Illuminate\Database\Eloquent\Model;
+
+class RuntimePayloadCompletionBudget extends Model
+{
+    public $incrementing = false;
+
+    protected $keyType = 'string';
+
+    protected $guarded = [];
+
+    protected $casts = [
+        'context' => 'array', 'slots' => 'array', 'objects' => 'array', 'expires_at' => 'immutable_datetime',
+    ];
+}

app/Support/RuntimeExternalPayloadCleanup.php

@@ -44,6 +44,8 @@ public function runPass(
         $limit = min(self::MAX_BATCH_SIZE, max(1, $limit));
         $cutoff ??= now();
 
+        app(RuntimePayloadCompletionUploads::class)->cleanup($namespace, $limit, $cutoff);
+
         $candidates = $this->expiredQuery($namespace, $cutoff)
             ->orderBy('expires_at')
             ->orderBy('id')

app/Support/RuntimeExternalPayloadObjectLock.php

@@ -22,11 +22,14 @@ final class RuntimeExternalPayloadObjectLock
      * @param  Closure(): TReturn  $callback
      * @return TReturn
      */
-    public function transaction(string $uri, Closure $callback): mixed
+    public function transaction(string $uri, Closure $callback, ?Closure $beforeLock = null): mixed
     {
         $bucket = hexdec(substr(hash('sha256', $uri), 0, 4)) % self::BUCKETS;
 
-        return DB::transaction(function () use ($bucket, $callback): mixed {
+        return DB::transaction(function () use ($bucket, $callback, $beforeLock): mixed {
+            // Completion ownership locks must precede payload locks, as they do
+            // when the engine commits an activity result.
+            $beforeLock?->__invoke();
             $lock = DB::table('runtime_external_payload_object_locks')
                 ->where('bucket', $bucket)
                 ->lockForUpdate()

app/Support/RuntimeExternalPayloadReference.php

@@ -72,6 +72,14 @@ public static function transportManifest(): array
                     'X-Durable-Workflow-Payload-SHA256',
                 ],
                 'idempotency' => 'content_addressed_per_namespace',
+                'completion_context' => [
+                    'schema' => RuntimePayloadCompletionContext::SCHEMA,
+                    'header' => RuntimePayloadCompletionContext::HEADER,
+                    'use' => 'retry_draining_refusal_for_current_worker_completion_only',
+                    'max_bytes_per_lease' => RuntimePayloadCompletionUploads::maxBytes(),
+                    'max_slots_per_lease' => RuntimePayloadCompletionUploads::MAX_SLOTS,
+                    'fenced_or_stale_admission' => 'refuse',
+                ],
             ],
             'fetch' => [
                 'method' => 'GET',
@@ -121,6 +129,10 @@ public static function transportManifest(): array
                 'external_payload_namespace_bytes_exhausted' => ['status' => 429, 'retryable' => true],
                 'external_payload_namespace_objects_exhausted' => ['status' => 429, 'retryable' => true],
                 'external_payload_namespace_quota_unavailable' => ['status' => 503, 'retryable' => true],
+                'external_payload_completion_invalid' => ['status' => 422, 'retryable' => false],
+                'external_payload_completion_lease_rejected' => ['status' => 409, 'retryable' => false],
+                'external_payload_completion_conflict' => ['status' => 409, 'retryable' => false],
+                'external_payload_completion_budget_exhausted' => ['status' => 429, 'retryable' => true],
             ],
             'audit_events' => [
                 'external_payload.uploaded',

app/Support/RuntimeExternalPayloadRegistry.php

@@ -19,7 +19,7 @@ public function __construct(
     /**
      * @return array{schema: string, reference_id: string, codec: string, size_bytes: int, sha256: string}
      */
-    public function upload(string $namespace, string|ExternalPayloadStream $data, string $codec, string $sha256): array
+    public function upload(string $namespace, string|ExternalPayloadStream $data, string $codec, string $sha256, ?\Closure $beforeWrite = null): array
     {
         $namespace = $this->namespace($namespace);
         try {
@@ -56,6 +56,7 @@ public function upload(string $namespace, string|ExternalPayloadStream $data, st
             sizeBytes: $sizeBytes,
             retained: false,
             expiresAt: $expiresAt,
+            beforeWrite: $beforeWrite,
         ));
     }
 
@@ -465,6 +466,7 @@ private function store(
         int $sizeBytes,
         bool $retained,
         mixed $expiresAt,
+        ?\Closure $beforeWrite = null,
     ): RuntimeExternalPayload {
         $observedSize = $data instanceof ExternalPayloadStream ? $data->sizeBytes : strlen($data);
         $observedHash = $data instanceof ExternalPayloadStream ? $data->sha256 : hash('sha256', $data);
@@ -479,17 +481,12 @@ private function store(
         }
 
         try {
-            $this->objectLock->transaction($uri, fn (): RuntimeExternalPayload => $this->track(
-                $namespace,
-                $uri,
-                $codec,
-                $sha256,
-                $sizeBytes,
-                false,
-                $expiresAt,
-                RuntimeExternalPayload::UPLOAD_WRITING,
-            ));
-        } catch (RuntimeExternalPayloadException $exception) {
+            $this->objectLock->transaction($uri, function () use ($namespace, $uri, $codec, $sha256, $sizeBytes, $expiresAt, $beforeWrite): RuntimeExternalPayload {
+                $beforeWrite?->__invoke();
+
+                return $this->track($namespace, $uri, $codec, $sha256, $sizeBytes, false, $expiresAt, RuntimeExternalPayload::UPLOAD_WRITING);
+            }, $beforeWrite);
+        } catch (RuntimeExternalPayloadException|StorageAdmissionPaused $exception) {
             throw $exception;
         } catch (Throwable $exception) {
             throw $this->unavailable('External payload reference registration failed before storage commit.', $exception);
@@ -506,7 +503,9 @@ private function store(
                 $sizeBytes,
                 $retained,
                 $expiresAt,
+                $beforeWrite,
             ): RuntimeExternalPayload {
+                $beforeWrite?->__invoke();
                 $row = RuntimeExternalPayload::query()
                     ->where('namespace', $namespace)
                     ->where('storage_uri_sha256', hash('sha256', $uri))
@@ -546,8 +545,8 @@ private function store(
                     $expiresAt,
                     RuntimeExternalPayload::UPLOAD_READY,
                 );
-            });
-        } catch (RuntimeExternalPayloadException $exception) {
+            }, $beforeWrite);
+        } catch (RuntimeExternalPayloadException|StorageAdmissionPaused $exception) {
             throw $exception;
         } catch (Throwable $exception) {
             // The writing row was committed before the backing write. If this

app/Support/RuntimePayloadCompletionContext.php

@@ -0,0 +1,109 @@
+<?php
+
+namespace App\Support;
+
+use JsonException;
+
+final readonly class RuntimePayloadCompletionContext
+{
+    public const HEADER = 'X-Durable-Workflow-Payload-Completion';
+
+    public const SCHEMA = 'durable-workflow.v2.payload-completion-context.v1';
+
+    /** @param list<int|string> $slot */
+    private function __construct(
+        public string $kind,
+        public string $taskId,
+        public int|string $attempt,
+        public string $leaseOwner,
+        public string $operation,
+        public array $slot,
+    ) {}
+
+    public static function parse(string $header): self
+    {
+        if (strlen($header) > 4096) {
+            throw self::invalid();
+        }
+        try {
+            $value = json_decode($header, true, 8, JSON_THROW_ON_ERROR);
+        } catch (JsonException) {
+            throw self::invalid();
+        }
+        $keys = ['schema', 'kind', 'task_id', 'attempt', 'lease_owner', 'operation', 'slot'];
+        if (! is_array($value) || count($value) !== count($keys)
+            || array_diff($keys, array_keys($value)) !== []
+            || ($value['schema'] ?? null) !== self::SCHEMA
+            || ! in_array($value['kind'] ?? null, ['activity', 'workflow', 'query'], true)
+            || ! is_string($value['operation'])
+            || ! self::identifier($value['task_id']) || ! self::identifier($value['lease_owner'])
+            || ! is_array($value['slot']) || ! array_is_list($value['slot'])) {
+            throw self::invalid();
+        }
+        if ($value['kind'] === 'activity'
+            ? ! self::identifier($value['attempt'])
+            : (! is_int($value['attempt']) || $value['attempt'] < 1)) {
+            throw self::invalid();
+        }
+        $validSlot = match ($value['kind'].'.'.($value['operation'] ?? '')) {
+            'activity.complete' => $value['slot'] === ['result'],
+            'activity.fail' => $value['slot'] === ['failure', 'details'],
+            'query.complete' => $value['slot'] === ['result_envelope'],
+            'workflow.complete' => self::workflowSlot($value['slot']),
+            default => false,
+        };
+        if (! $validSlot) {
+            throw self::invalid();
+        }
+
+        return new self($value['kind'], $value['task_id'], $value['attempt'],
+            $value['lease_owner'], $value['operation'], $value['slot']);
+    }
+
+    public function scope(string $namespace): string
+    {
+        // Complete/fail and every payload slot share one allowance for this lease.
+        return hash('sha256', json_encode([$namespace, $this->kind, $this->taskId,
+            $this->attempt, $this->leaseOwner], JSON_THROW_ON_ERROR));
+    }
+
+    public function slotIdentity(): string
+    {
+        return hash('sha256', json_encode([$this->operation, $this->slot], JSON_THROW_ON_ERROR));
+    }
+
+    public function toArray(): array
+    {
+        return ['schema' => self::SCHEMA, 'kind' => $this->kind, 'task_id' => $this->taskId,
+            'attempt' => $this->attempt, 'lease_owner' => $this->leaseOwner,
+            'operation' => $this->operation, 'slot' => $this->slot];
+    }
+
+    private static function identifier(mixed $value): bool
+    {
+        return is_string($value) && $value !== '' && strlen($value) <= 255
+            && preg_match('/[\x00-\x1f\x7f]/', $value) === 0;
+    }
+
+    private static function workflowSlot(array $slot): bool
+    {
+        if (($slot[0] ?? null) !== 'commands' || ! is_int($slot[1] ?? null) || $slot[1] < 0) {
+            return false;
+        }
+        if (count($slot) === 3) {
+            return in_array($slot[2], ['arguments', 'entries', 'request_payload', 'result'], true);
+        }
+        if (count($slot) === 4) {
+            return array_slice($slot, 2) === ['exception', 'details'];
+        }
+
+        return count($slot) === 6 && $slot[2] === 'workflow_stream' && $slot[3] === 'items'
+            && is_int($slot[4]) && $slot[4] >= 0 && $slot[5] === 'payload';
+    }
+
+    private static function invalid(): RuntimeExternalPayloadException
+    {
+        return new RuntimeExternalPayloadException('external_payload_completion_invalid', 422, false,
+            'A canonical worker completion context is required.');
+    }
+}