From 084910d7cccde5b27b08703cb71922f6160edca2 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Fri, 12 Jun 2026 12:17:46 -0400
Subject: [PATCH] fix(diff): release graphs from memory as they are completed

Because the `as_completed()` call was outside of the with scope, we were
waiting for all futures to complete before processing any of them. Also
we were storing futures (which contain the full output) in a dict.

This meant we were holding the output of all graph generations in memory
simultaneously. In Gecko, this will eat through your available RAM in no
time.
---
 src/taskgraph/main.py | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py
index d1ab84e12..e74da8a7d 100644
--- a/src/taskgraph/main.py
+++ b/src/taskgraph/main.py
@@ -222,6 +222,7 @@ def logfile(spec):
         return 0
 
     futures = {}
+    returncode = 0
     with ProcessPoolExecutor(max_workers=options["max_workers"]) as executor:
         for spec in parameters:
             f = executor.submit(
@@ -229,25 +230,24 @@ def logfile(spec):
             )
             futures[f] = spec
 
-    returncode = 0
-    for future in as_completed(futures):
-        output_file = options["output_file"]
-        spec = futures[future]
-        e = future.exception()
-        if e:
-            returncode = 1
-            out = "".join(traceback.format_exception(type(e), e, e.__traceback__))
-            if options["diff"]:
-                # Dump to console so we don't accidentally diff the tracebacks.
-                output_file = None
-        else:
-            out = future.result()
+        for future in as_completed(futures):
+            output_file = options["output_file"]
+            spec = futures.pop(future)
+            e = future.exception()
+            if e:
+                returncode = 1
+                out = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+                if options["diff"]:
+                    # Dump to console so we don't accidentally diff the tracebacks.
+                    output_file = None
+            else:
+                out = future.result()
 
-        dump_output(
-            out,
-            path=output_file,
-            params_spec=spec if len(parameters) > 1 else None,
-        )
+            dump_output(
+                out,
+                path=output_file,
+                params_spec=spec if len(parameters) > 1 else None,
+            )
 
     return returncode