Cleanup serialized context coming back from a task worker

Author: veewee

src/Runner/FixableTaskResult.php

@@ -69,6 +69,14 @@ public function getContext(): ContextInterface
         return $this->result->getContext();
     }
 
+    public function withContext(ContextInterface $context): static
+    {
+        $new = clone $this;
+        $new->result = $this->result->withContext($context);
+
+        return $new;
+    }
+
     public function withAppendedMessage(string $message): TaskResultInterface
     {
         $new = clone $this;

src/Runner/TaskHandler/Middleware/ParallelProcessingMiddleware.php

@@ -58,7 +58,9 @@ static function () use ($task, $runnerContext, $next, $currentEnv): TaskResultIn
 
         return async(function () use ($task, $runnerContext, $execution): TaskResultInterface {
             try {
-                return $execution->getFuture()->await();
+                $result = $execution->getFuture()->await();
+
+                return $this->swapSerializedContext($result, $runnerContext);
             } catch (\Throwable $exception) {
                 return TaskResult::createFailed(
                     $task,
@@ -75,4 +77,18 @@ private function wrapException(\Throwable $error): ParallelException
             ? ParallelException::fromVerboseThrowable($error)
             : ParallelException::fromThrowable($error);
     }
+
+    /**
+     * The results coming back from a parallel worker contains a serialized context.
+     * This context can be rather big, which can cause memory issues.
+     *
+     * This method will swap the serialized context with the original
+     * context so that the serialized one gets garbage collected.
+     */
+    private function swapSerializedContext(
+        TaskResultInterface $result,
+        TaskRunnerContext $runnerContext
+    ): TaskResultInterface {
+        return $result->withContext($runnerContext->getTaskContext());
+    }
 }

src/Runner/TaskResult.php

@@ -95,6 +95,14 @@ public function getContext(): ContextInterface
         return $this->context;
     }
 
+    public function withContext(ContextInterface $context): static
+    {
+        $new = clone $this;
+        $new->context = $context;
+
+        return $new;
+    }
+
     public function withAppendedMessage(string $message): TaskResultInterface
     {
         $new = clone $this;

src/Runner/TaskResultInterface.php

@@ -33,5 +33,7 @@ public function getMessage(): string;
 
     public function getContext(): ContextInterface;
 
+    public function withContext(ContextInterface $context): static;
+
     public function withAppendedMessage(string $message): self;
 }