Keep subfield memoization keys alive during execution

src/graphql/execution/executor.py

@@ -399,7 +399,9 @@ def __init__(  # noqa: PLR0913, PLR0917
         self.pending_incremental_futures = set()
         self.background_futures = set()
         self.async_work_finished_hook_task = None
-        self._relevant_sub_fields: dict[tuple, CollectedFields] = {}
+        self._relevant_sub_fields: dict[
+            tuple[GraphQLObjectType, int], tuple[FieldDetailsList, CollectedFields]
+        ] = {}
         self._stream_usages: RefMap[FieldDetailsList, StreamUsage] = RefMap()
 
     @classmethod
@@ -2117,30 +2119,23 @@ def collect_subfields(
         :meta private:
         """
         relevant_sub_fields = self._relevant_sub_fields
-        # We cannot use the field_details_list itself as key for the cache, since it
-        # is not hashable as a list. We also do not want to use the field_details_list
-        # itself (converted to a tuple) as keys, since hashing them is slow.
-        # Therefore, we use the ids of the field_details_list items as keys. Note that
-        # we do not use the id of the list, since we want to hit the cache for all
-        # lists of the same nodes, not only for the same list of nodes. Also, the list
-        # id may even be reused, in which case we would get wrong results from cache.
-        key = (
-            (return_type, id(field_details_list[0]))
-            if len(field_details_list) == 1  # optimize most frequent case
-            else (return_type, *map(id, field_details_list))
+        # Like GraphQL.js, we memoize on the identity of the field_details_list.
+        # Since lists cannot be weakly referenced, we keep the list alive together
+        # with the cached value, so that its id cannot be reused during execution.
+        key = (return_type, id(field_details_list))
+        cached = relevant_sub_fields.get(key)
+        if cached is not None:
+            return cached[1]
+        collected_fields = collect_subfields(
+            self.schema,
+            self.fragments,
+            self.variable_values,
+            self.operation,
+            return_type,
+            field_details_list,
+            self.hide_suggestions,
         )
-        collected_fields: CollectedFields | None = relevant_sub_fields.get(key)
-        if collected_fields is None:
-            collected_fields = collect_subfields(
-                self.schema,
-                self.fragments,
-                self.variable_values,
-                self.operation,
-                return_type,
-                field_details_list,
-                self.hide_suggestions,
-            )
-            relevant_sub_fields[key] = collected_fields
+        relevant_sub_fields[key] = (field_details_list, collected_fields)
         return collected_fields
 
 

tests/execution/test_executor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import gc
 from collections.abc import Awaitable
 from typing import Any, cast
 
@@ -1367,6 +1368,48 @@ def memoizes_collect_subfields_results():
 
         assert third is not first
 
+    def retains_memoization_keys_until_execution_finishes():
+        released: list[str] = []
+
+        class TrackedFieldDetails(FieldDetails):
+            def __del__(self) -> None:
+                released.append(self.node.name.value)
+
+        item_type = GraphQLObjectType(
+            "Item",
+            {"key": GraphQLField(GraphQLString), "name": GraphQLField(GraphQLString)},
+        )
+        schema = GraphQLSchema(
+            GraphQLObjectType(
+                "Query", {"a": GraphQLField(item_type), "b": GraphQLField(item_type)}
+            )
+        )
+        document = parse("{ a { key } b { name } }")
+        executor = Executor.build(schema, document)
+        assert isinstance(executor, Executor)
+        a_node, b_node = executor.operation.selection_set.selections
+        assert isinstance(a_node, FieldNode)
+        assert isinstance(b_node, FieldNode)
+
+        first_fields: list[FieldDetails] = [TrackedFieldDetails(a_node, None)]
+        first = executor.collect_subfields(item_type, first_fields)
+        del first_fields
+        gc.collect()
+        assert released == []
+
+        second_fields: list[FieldDetails] = [TrackedFieldDetails(b_node, None)]
+        second = executor.collect_subfields(item_type, second_fields)
+        assert second is not first
+        assert list(first.grouped_field_set) == ["key"]
+        assert list(second.grouped_field_set) == ["name"]
+        del second_fields
+        gc.collect()
+        assert released == []
+
+        del executor
+        gc.collect()
+        assert sorted(released) == ["a", "b"]
+
 
 def describe_base_executor_without_incremental_delivery():
     def deduplicates_errors_at_nulled_positions():