[johng]

Obtain a reference to lz->long_cnt, lz->long_step
such that we don't hit a use after free.

• Issue: itertoolsmodule: free-threading use after free bug in counting slow mode #153981

[johng]

This appears to add some very slight overhead to repr in slow mode, no impact to the likely more heavily used fast mode

[picnixz]

Performances are not really inportant for __repr__. You usually don't print a count in a hot path.

[johng]

@picnixz would you be able to kindly review again, thanks!

[johng]

@picnixz kindly pining for a re-review, thank you!

[picnixz]

It may just seem dumb but since __repr__ is fairly used just for debugging, I don't think we really need to worry about performances; so why not just having everything in a critical section? would that work?

[johng]

Would that not risk locking other threads executing more latency sensitive functions that are holding the lock?

[picnixz]

We don't care here IMO. If you print the counter, just expect it to possbly be slow. And can you simply benchmark this please?

[johng]

If I add a lock around the whole of repr I can significantly slow down the actual counting module if another thread is printing. I'm sure this example is a bit contrived however I think a real possibility another thread printing could slow down the counting module.

Build	next() thread time for 1M calls (median of 5)
PR as-is (lock only while taking references)	108 / 116 / 112 ms
Whole repr in a critical section	1479 / 2095 / 1783 ms

simple_bench.py (see below)

Details

import statistics
import threading
import time
from itertools import count

N = 1_000_000


def bench():
    it = count(10**18, 2)  # slow mode: next() takes the critical section
    done = threading.Event()
    elapsed = []

    def counter():
        start = time.perf_counter()
        for _ in range(N):
            next(it)
        elapsed.append(time.perf_counter() - start)

    def reprer():
        while not done.is_set():
            repr(it)

    r = threading.Thread(target=reprer)
    c = threading.Thread(target=counter)
    r.start()
    c.start()
    c.join()
    done.set()
    r.join()
    return elapsed[0]


times = [bench() for _ in range(5)]
print(f"next() thread: {statistics.median(times) * 1000:.1f} ms "
      f"for {N:,} calls while another thread calls repr() (median of 5)")

And just to confirm a lock held across repr would look something like this

static PyObject *
count_repr_lock_held(countobject *lz)
{
    if (lz->long_cnt == NULL) {
        Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->cnt);
        return PyUnicode_FromFormat("%s(%zd)",
                                    _PyType_Name(Py_TYPE(lz)), cnt);
    }

    if (PyLong_Check(lz->long_step)) {
        long step = PyLong_AsLong(lz->long_step);
        if (step == -1 && PyErr_Occurred()) {
            PyErr_Clear();
        }
        if (step == 1) {
            /* Don't display step when it is an integer equal to 1 */
            return PyUnicode_FromFormat("%s(%R)",
                                        _PyType_Name(Py_TYPE(lz)),
                                        lz->long_cnt);
        }
    }
    return PyUnicode_FromFormat("%s(%R, %R)",
                                _PyType_Name(Py_TYPE(lz)),
                                lz->long_cnt, lz->long_step);
}

static PyObject *
count_repr(PyObject *op)
{
    PyObject *result;
    Py_BEGIN_CRITICAL_SECTION(op);
    result = count_repr_lock_held(countobject_CAST(op));
    Py_END_CRITICAL_SECTION();
    return result;
}

[picnixz]

Two possibilities:

• I think it's expected that printing continuously may slow down the rest of the access. So there should be NO hot paths being affected (if you decide to repr() the count, you are no longer in a hot path).
• OTOH, the reason why I suggested the whole critical section is because the code either has a lots of ifs or a UnicodeWriter. Would it be possible to actually have a minimal change with neither (maybe just storing the message and then returning it, like msg = PyUnicode... as we have, and then goto finally where finally would free the resources?)

[johng]

I'm not sure if this was what you were after but I think this reduces the branches considerably if the if/else branches were the main sticking point. Agree looking back it was rather messy for all the approaches!

[picnixz]

You can put PyObject *msg = here instead.

[picnixz]

Can we have the if check out of the format call? I don't exactly remember but I think passing more va args than what the format expects is UB>