From db553983b5e146cd0166bcf1ff59aae4d87a8a0d Mon Sep 17 00:00:00 2001
From: Toon Verwerft <toonverwerft@gmail.com>
Date: Fri, 2 Aug 2024 14:32:22 +0200
Subject: [PATCH] Cleanup serialized context coming back from a task worker

---
 src/Runner/FixableTaskResult.php               |  8 ++++++++
 .../ParallelProcessingMiddleware.php           | 18 +++++++++++++++++-
 src/Runner/TaskResult.php                      |  8 ++++++++
 src/Runner/TaskResultInterface.php             |  2 ++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/Runner/FixableTaskResult.php b/src/Runner/FixableTaskResult.php
index d94079735..6068aeb0d 100644
--- a/src/Runner/FixableTaskResult.php
+++ b/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;
diff --git a/src/Runner/TaskHandler/Middleware/ParallelProcessingMiddleware.php b/src/Runner/TaskHandler/Middleware/ParallelProcessingMiddleware.php
index 648a5804f..a4b58f4d2 100644
--- a/src/Runner/TaskHandler/Middleware/ParallelProcessingMiddleware.php
+++ b/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());
+    }
 }
diff --git a/src/Runner/TaskResult.php b/src/Runner/TaskResult.php
index 35ccc8f94..c1b0c2702 100644
--- a/src/Runner/TaskResult.php
+++ b/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;
diff --git a/src/Runner/TaskResultInterface.php b/src/Runner/TaskResultInterface.php
index 3a3a3e76b..8196dc21d 100644
--- a/src/Runner/TaskResultInterface.php
+++ b/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;
 }