Fix rare double-free issue

[arhamchopra]

Potential, but very rare chance of error flagged by claude

def test_pynode_init_exception_no_double_free(self):                                                                                                                                                                                                                             [50/2121]
    # Regression test for the double-free fixed in PyNode::init
    # (cpp/csp/python/PyNode.cpp).  The buggy code decref'd the generator
    # frame-slot tuple *before* computing and assigning its replacement.
    # These frame-slot tuples are shared entries in the generator code
    # object's co_consts (refcount > 1), so if init raised after that early
    # decref but before the slot was reassigned, the slot was left pointing
    # at the const tuple, which got decref'd a *second* time when the
    # half-built node's generator was torn down.  That over-decref'd the
    # shared co_consts tuple and eventually double-freed it.
    #
    # We deterministically force an exception inside the init window by
    # inflating a (deliberately small, so graph wiring doesn't choke) output
    # basket's shape past the C++ basket-size limit right before the engine
    # constructs the node, so PyNode::init raises its "output basket size
    # exceeds limit" ValueError while iterating the frame locals.  We then
    # assert the shared const tuple's refcount is unchanged -- with the bug
    # it drops by one every time init raises.
    from csp.impl.wiring.node import NodeDef

    @csp.node
    def basket_out(x: ts[int]) -> csp.OutputBasket(List[ts[int]], shape=2):
        if csp.ticked(x):
            csp.output([x, x])

    # the ('output_proxy', N) tuple is the shared co_consts entry that the
    # buggy code over-decref'd; grab and hold a reference to it
    output_proxy_const = None
    for const in basket_out._impl.__code__.co_consts:
        if isinstance(const, tuple) and len(const) == 2 and const[0] == "output_proxy":
            output_proxy_const = const
            break
    self.assertIsNotNone(output_proxy_const, "could not locate output_proxy const tuple in generator")

    @csp.graph
    def g():
        csp.add_graph_output("o", basket_out(csp.const(1))[0])

    huge_shape = 2**40  # > OutputId::maxBasketElements() (2**31)
    orig_create = NodeDef._create

    def inflate_shape_create(self, engine, memo):
        # after graph build, before C++ node construction, inflate the basket
        # shape so PyNode::init raises inside its (buggy) init window
        if self._signature._name == "basket_out" and self._output_types is not None:
            self._output_types = tuple(
                (huge_shape, ot[1]) if (isinstance(ot, tuple) and len(ot) == 2 and isinstance(ot[0], int)) else ot
                for ot in self._output_types
            )
        return orig_create(self, engine, memo)

    gc.collect()
    before = sys.getrefcount(output_proxy_const)
    NodeDef._create = inflate_shape_create
    try:
        with self.assertRaises(ValueError):
            csp.run(g, starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1))
    finally:
        NodeDef._create = orig_create
    gc.collect()
    after = sys.getrefcount(output_proxy_const)

    self.assertEqual(
        before,
        after,
        f"PyNode::init raising an exception over-decref'd the shared generator const tuple "
        f"(refcount {before} -> {after}), which leads to a double-free",
    )
    # AssertionError: 6 != 5 : PyNode::init raising an exception over-decref'd the shared generator const tuple (refcount 6 -> 5), which leads to a double-free

[timkpaine]

is there a test case or repro that we can include?

[arhamchopra]

is there a test case or repro that we can include?

Updated description, no point adding this as a test since this is not possible in a real-life scenario.