Skip to content

Commit 14f5df9

Browse files
p-sawickihauntsaninja
p-sawicki
authored and
hauntsaninja
committed
[mypyc] Clear coroutine env on coroutine completion (#21734)
The `env_class` object associated with a mypyc coroutine is not immediately cleared when the coroutine completes. This can significantly increase memory usage, since the env class may hold references to captured locals and values spilled across suspension points. The objects are eventually collectible by the GC but the collection might be delayed in cases where a nested coroutine is awaited, eg. ```python async def allocate(size: int) -> None: payload = bytearray(size) async def nested() -> int: return payload[-1] assert await nested() == 0 ``` With nesting mypyc creates env classes for both `allocate` and `nested` that may reference each other and form a cycle. To fix this, clear the `env_class` immediately after the coroutine completes. This matches behavior of cpython, which clears frames of completed coroutines immediately in the eval loop.
1 parent 6dfa06d commit 14f5df9

7 files changed

Lines changed: 318 additions & 19 deletions

File tree

‎mypyc/codegen/emitclass.py‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
252252
getseters_name = f"{name_prefix}_getseters"
253253
vtable_name = f"{name_prefix}_vtable"
254254
traverse_name = f"{name_prefix}_traverse"
255-
clear_name = f"{name_prefix}_clear"
255+
clear_name = emitter.native_function_name(cl.clear)
256256
dealloc_name = f"{name_prefix}_dealloc"
257257
methods_name = f"{name_prefix}_methods"
258258
vtable_setup_name = f"{name_prefix}_trait_vtable_setup"
@@ -271,7 +271,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
271271
fields["tp_dealloc"] = f"(destructor){name_prefix}_dealloc"
272272
if not cl.is_acyclic:
273273
fields["tp_traverse"] = f"(traverseproc){name_prefix}_traverse"
274-
fields["tp_clear"] = f"(inquiry){name_prefix}_clear"
274+
fields["tp_clear"] = f"(inquiry){clear_name}"
275275
# Populate .tp_finalize and generate a finalize method only if __del__ is defined for this class.
276276
del_method = next((e.method for e in cl.vtable_entries if e.name == "__del__"), None)
277277
if del_method:
@@ -344,7 +344,11 @@ def emit_line() -> None:
344344
if not cl.is_acyclic:
345345
generate_traverse_for_class(cl, traverse_name, emitter)
346346
emit_line()
347-
generate_clear_for_class(cl, clear_name, emitter)
347+
generate_clear_for_class(cl, cl.clear, emitter)
348+
emit_line()
349+
generate_clear_for_class(
350+
cl, cl.clear_on_completion, emitter, skip_attrs=cl.attrs_to_keep_alive_on_completion
351+
)
348352
emit_line()
349353
generate_dealloc_for_class(cl, dealloc_name, clear_name, bool(del_method), emitter)
350354
emit_line()
@@ -895,12 +899,21 @@ def generate_traverse_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -
895899
emitter.emit_line("}")
896900

897901

898-
def generate_clear_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -> None:
899-
emitter.emit_line("static int")
900-
emitter.emit_line(f"{func_name}({cl.struct_name(emitter.names)} *self)")
902+
def generate_clear_for_class(
903+
cl: ClassIR, func_decl: FuncDecl, emitter: Emitter, skip_attrs: set[str] | None = None
904+
) -> None:
905+
if skip_attrs is None:
906+
skip_attrs = set()
907+
emitter.emit_line("static " + native_function_header(func_decl, emitter))
901908
emitter.emit_line("{")
909+
emitter.emit_line(
910+
f"{cl.struct_name(emitter.names)} *self = "
911+
f"({cl.struct_name(emitter.names)} *)cpy_r_self;"
912+
)
902913
for base in reversed(cl.base_mro):
903914
for attr, rtype in base.attributes.items():
915+
if attr in skip_attrs:
916+
continue
904917
emitter.emit_gc_clear(f"self->{emitter.attr(attr)}", rtype)
905918
base_args = "(PyObject *)self"
906919
if cl.builtin_base:
@@ -951,7 +964,7 @@ def generate_dealloc_for_class(
951964
if not cl.is_acyclic:
952965
emitter.emit_line("PyObject_GC_UnTrack(self);")
953966
if cl.builtin_base:
954-
emitter.emit_line(f"{clear_func_name}(self);")
967+
emitter.emit_line(f"{clear_func_name}((PyObject *)self);")
955968
# For native subclasses of builtins such as dict, the base deallocator
956969
# is responsible for tearing down base-owned storage and freeing memory.
957970
# Re-track self if base is GC-aware to match cpython's subtype_dealloc.
@@ -966,7 +979,7 @@ def generate_dealloc_for_class(
966979
emit_reuse_dealloc(cl, emitter)
967980
# The trashcan is needed to handle deep recursive deallocations
968981
emitter.emit_line(f"CPy_TRASHCAN_BEGIN(self, {dealloc_func_name})")
969-
emitter.emit_line(f"{clear_func_name}(self);")
982+
emitter.emit_line(f"{clear_func_name}((PyObject *)self);")
970983
emitter.emit_line("Py_TYPE(self)->tp_free((PyObject *)self);")
971984
emitter.emit_line("CPy_TRASHCAN_END(self)")
972985
emitter.emit_line("done: ;")

‎mypyc/ir/class_ir.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from mypyc.common import PROPSET_PREFIX, JsonDict
88
from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
99
from mypyc.ir.ops import DeserMaps, Value
10-
from mypyc.ir.rtypes import RInstance, RType, deserialize_type, object_rprimitive
10+
from mypyc.ir.rtypes import RInstance, RType, c_int_rprimitive, deserialize_type, object_rprimitive
1111
from mypyc.namegen import NameGenerator, exported_name
1212

1313
# Some notes on the vtable layout: Each concrete class has a vtable
@@ -143,8 +143,25 @@ def __init__(
143143
module_name,
144144
FuncSignature([RuntimeArg("type", object_rprimitive)], RInstance(self)),
145145
)
146+
self.clear = FuncDecl(
147+
name + "_clear",
148+
None,
149+
module_name,
150+
FuncSignature([RuntimeArg("self", RInstance(self))], c_int_rprimitive),
151+
internal=True,
152+
)
153+
self.clear_on_completion = FuncDecl(
154+
name + "_clear_on_completion",
155+
None,
156+
module_name,
157+
FuncSignature([RuntimeArg("self", RInstance(self))], c_int_rprimitive),
158+
internal=True,
159+
)
146160
# Attributes defined in the class (not inherited)
147161
self.attributes: dict[str, RType] = {}
162+
# Attributes that must survive generator/coroutine completion because
163+
# escaped nested functions may still read them as closure variables.
164+
self.attrs_to_keep_alive_on_completion: set[str] = set()
148165
# Final attributes defined in the class (not inherited)
149166
self.final_attributes: set[str] = set()
150167
# Deletable attributes
@@ -412,8 +429,11 @@ def serialize(self) -> JsonDict:
412429
"_serializable": self._serializable,
413430
"builtin_base": self.builtin_base,
414431
"ctor": self.ctor.serialize(),
432+
"clear": self.clear.serialize(),
433+
"clear_on_completion": self.clear_on_completion.serialize(),
415434
# We serialize dicts as lists to ensure order is preserved
416435
"attributes": [(k, t.serialize()) for k, t in self.attributes.items()],
436+
"attrs_to_keep_alive_on_completion": sorted(self.attrs_to_keep_alive_on_completion),
417437
"final_attributes": sorted(self.final_attributes),
418438
# We try to serialize a name reference, but if the decl isn't in methods
419439
# then we can't be sure that will work so we serialize the whole decl.
@@ -474,7 +494,10 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
474494
ir._serializable = data["_serializable"]
475495
ir.builtin_base = data["builtin_base"]
476496
ir.ctor = FuncDecl.deserialize(data["ctor"], ctx)
497+
ir.clear = FuncDecl.deserialize(data["clear"], ctx)
498+
ir.clear_on_completion = FuncDecl.deserialize(data["clear_on_completion"], ctx)
477499
ir.attributes = {k: deserialize_type(t, ctx) for k, t in data["attributes"]}
500+
ir.attrs_to_keep_alive_on_completion = set(data["attrs_to_keep_alive_on_completion"])
478501
ir.final_attributes = set(data["final_attributes"])
479502
ir.method_decls = {
480503
k: ctx.functions[v].decl if isinstance(v, str) else FuncDecl.deserialize(v, ctx)

‎mypyc/irbuild/builder.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1568,12 +1568,15 @@ def add_var_to_env_class(
15681568
base: FuncInfo | ImplicitClass,
15691569
reassign: bool = False,
15701570
always_defined: bool = False,
1571+
keep_alive_on_completion: bool = False,
15711572
prefix: str = "",
15721573
) -> AssignmentTarget:
15731574
# First, define the variable name as an attribute of the environment class, and then
15741575
# construct a target for that attribute.
15751576
name = prefix + remangle_redefinition_name(var.name)
15761577
self.fn_info.env_class.attributes[name] = rtype
1578+
if keep_alive_on_completion:
1579+
self.fn_info.env_class.attrs_to_keep_alive_on_completion.add(name)
15771580
if always_defined:
15781581
self.fn_info.env_class.attrs_with_defaults.add(name)
15791582
attr_target = AssignmentTargetAttr(base.curr_env_reg, name)

‎mypyc/irbuild/env_class.py‎

Lines changed: 39 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ def g() -> int:
1717

1818
from __future__ import annotations
1919

20-
from mypy.nodes import Argument, FuncDef, SymbolNode, Var
20+
from mypy.nodes import Argument, FuncDef, FuncItem, SymbolNode, Var
2121
from mypyc.common import (
2222
BITMAP_BITS,
2323
ENV_ATTR_NAME,
@@ -60,6 +60,8 @@ class is generated, the function environment has not yet been
6060
# If the function is nested, its environment class must contain an environment
6161
# attribute pointing to its encapsulating functions' environment class.
6262
env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class)
63+
if builder.fn_info.contains_nested:
64+
env_class.attrs_to_keep_alive_on_completion.add(ENV_ATTR_NAME)
6365
env_class.mro = [env_class]
6466
builder.fn_info.env_class = env_class
6567
builder.classes.append(env_class)
@@ -229,7 +231,17 @@ def add_args_to_env(
229231
rtype = builder.type_to_rtype(arg.variable.type)
230232
assert base is not None, "base cannot be None for adding nonlocal args"
231233
builder.add_var_to_env_class(
232-
arg.variable, rtype, base, reassign=reassign, prefix=prefix
234+
arg.variable,
235+
rtype,
236+
base,
237+
reassign=reassign,
238+
keep_alive_on_completion=(
239+
is_free_variable(builder, arg.variable)
240+
or is_free_variable_in_nested_func(
241+
builder, builder.fn_info.fitem, arg.variable
242+
)
243+
),
244+
prefix=prefix,
233245
)
234246

235247

@@ -256,7 +268,12 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None:
256268
if isinstance(var, Var):
257269
rtype = builder.type_to_rtype(var.type)
258270
builder.add_var_to_env_class(
259-
var, rtype, env_for_func, reassign=False, prefix=prefix
271+
var,
272+
rtype,
273+
env_for_func,
274+
reassign=False,
275+
keep_alive_on_completion=True,
276+
prefix=prefix,
260277
)
261278

262279
if builder.fn_info.fitem in builder.encapsulating_funcs:
@@ -267,10 +284,16 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None:
267284
# the same name and signature across conditional blocks
268285
# will generate different callable classes, so the callable
269286
# class that gets instantiated must be generic.
287+
nested_prefix = prefix
270288
if nested_fn.is_generator or nested_fn.is_coroutine:
271-
prefix = GENERATOR_ATTRIBUTE_PREFIX
289+
nested_prefix = GENERATOR_ATTRIBUTE_PREFIX
272290
builder.add_var_to_env_class(
273-
nested_fn, object_rprimitive, env_for_func, reassign=False, prefix=prefix
291+
nested_fn,
292+
object_rprimitive,
293+
env_for_func,
294+
reassign=False,
295+
keep_alive_on_completion=is_free_variable(builder, nested_fn),
296+
prefix=nested_prefix,
274297
)
275298

276299

@@ -308,3 +331,14 @@ def setup_func_for_recursive_call(
308331
def is_free_variable(builder: IRBuilder, symbol: SymbolNode) -> bool:
309332
fitem = builder.fn_info.fitem
310333
return fitem in builder.free_variables and symbol in builder.free_variables[fitem]
334+
335+
336+
def is_free_variable_in_nested_func(
337+
builder: IRBuilder, fitem: FuncItem, symbol: SymbolNode
338+
) -> bool:
339+
for nested in builder.encapsulating_funcs.get(fitem, []):
340+
if symbol in builder.free_variables.get(nested, set()):
341+
return True
342+
if is_free_variable_in_nested_func(builder, nested, symbol):
343+
return True
344+
return False

‎mypyc/irbuild/generator.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
Call,
2424
Goto,
2525
Integer,
26+
LoadErrorValue,
2627
MethodCall,
2728
RaiseStandardError,
2829
Register,
@@ -49,7 +50,7 @@
4950
load_outer_envs,
5051
setup_func_for_recursive_call,
5152
)
52-
from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl
53+
from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl, gen_generator_func_cleanup
5354
from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
5455
from mypyc.primitives.exc_ops import (
5556
error_catch_op,
@@ -107,13 +108,20 @@ class that implements the function (each function gets a separate class).
107108
setup_func_for_recursive_call(
108109
builder, fitem, builder.fn_info.generator_class, prefix=GENERATOR_ATTRIBUTE_PREFIX
109110
)
111+
cleanup_on_error = BasicBlock()
112+
builder.builder.push_error_handler(cleanup_on_error)
110113
create_switch_for_generator_class(builder)
111114
add_raise_exception_blocks_to_generator_class(builder, fitem.line)
112115

113116
add_vars_to_env(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
114117

115118
builder.accept(fitem.body)
116119
builder.maybe_add_implicit_return()
120+
builder.builder.pop_error_handler()
121+
122+
builder.activate_block(cleanup_on_error)
123+
gen_generator_func_cleanup(builder, fitem.line)
124+
builder.add(Return(builder.add(LoadErrorValue(object_rprimitive))))
117125

118126
populate_switch_for_generator_class(builder)
119127

‎mypyc/irbuild/nonlocalcontrol.py‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@
88
from abc import abstractmethod
99
from typing import TYPE_CHECKING
1010

11+
from mypyc.ir.class_ir import ClassIR
1112
from mypyc.ir.ops import (
13+
ERR_NEVER,
1214
NO_TRACEBACK_LINE_NO,
1315
BasicBlock,
1416
Branch,
17+
Call,
1518
Goto,
1619
Integer,
1720
Register,
@@ -90,10 +93,7 @@ class GeneratorNonlocalControl(BaseNonlocalControl):
9093
"""Default nonlocal control in a generator function outside statements."""
9194

9295
def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
93-
# Assign an invalid next label number so that the next time
94-
# __next__ is called, we jump to the case in which
95-
# StopIteration is raised.
96-
builder.assign(builder.fn_info.generator_class.next_label_target, Integer(-1), line)
96+
gen_generator_func_cleanup(builder, line)
9797

9898
# Raise a StopIteration containing a field for the value that
9999
# should be returned. Before doing so, create a new block
@@ -132,6 +132,25 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
132132
builder.add(Return(Integer(0, object_rprimitive)))
133133

134134

135+
def gen_class_clear_on_completion(builder: IRBuilder, obj: Value, cl: ClassIR, line: int) -> None:
136+
call = Call(cl.clear_on_completion, [obj], line)
137+
call.error_kind = ERR_NEVER
138+
builder.add(call)
139+
140+
141+
def gen_generator_func_cleanup(builder: IRBuilder, line: int) -> None:
142+
"""Clear references held by a completed generator or coroutine."""
143+
cls = builder.fn_info.generator_class
144+
145+
# Assign an invalid next label number so that the next time __next__ is
146+
# called, we jump to the case in which StopIteration is raised.
147+
builder.assign(cls.next_label_target, Integer(-1), line)
148+
149+
gen_class_clear_on_completion(builder, cls.curr_env_reg, builder.fn_info.env_class, line)
150+
if builder.fn_info.env_class is not cls.ir:
151+
gen_class_clear_on_completion(builder, cls.self_reg, cls.ir, line)
152+
153+
135154
class CleanupNonlocalControl(NonlocalControl):
136155
"""Abstract nonlocal control that runs some cleanup code."""
137156

0 commit comments

Comments
 (0)