Assign distinct thread IDs to greenlets

Greenlets aren't threads, but they are distinct threads of execution
within a single process, with distinct stacks. Reflect this by
generating distinct thread IDs for them.

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>

src/memray/_memray/tracking_api.cpp

@@ -375,7 +375,7 @@ PythonStackTracker::installGreenletTraceFunctionIfNeeded()
 }
 
 void
-PythonStackTracker::handleGreenletSwitch(PyObject*, PyObject*)
+PythonStackTracker::handleGreenletSwitch(PyObject* from, PyObject* to)
 {
     RecursionGuard guard;
 
@@ -388,6 +388,25 @@ PythonStackTracker::handleGreenletSwitch(PyObject*, PyObject*)
         emitPendingPushesAndPops();
     }
 
+    // Save current TID on old greenlet. Print errors but otherwise ignore them.
+    PyObject* tid = PyLong_FromUnsignedLong(t_tid);
+    if (!tid || 0 != PyObject_SetAttrString(from, "_memray_tid", tid)) {
+        PyErr_Print();
+    }
+    Py_XDECREF(tid);
+
+    // Restore TID from new greenlet, or generate a new one. Ignore errors:
+    // maybe we haven't seen this TID before, or maybe someone overwrote our
+    // attribute, but either way we can recover by generating a new one.
+    tid = PyObject_GetAttrString(to, "_memray_tid");
+    if (!tid || !PyLong_CheckExact(tid)) {
+        PyErr_Clear();
+        t_tid = generate_next_tid();
+    } else {
+        t_tid = PyLong_AsUnsignedLong(tid);
+    }
+    Py_XDECREF(tid);
+
     // Re-create our TLS stack from our Python frames, most recent last.
     // Note: `frame` may be null; the new greenlet may not have a Python stack.
     PyFrameObject* frame = PyEval_GetFrame();

tests/integration/test_greenlet.py

@@ -196,3 +196,8 @@ def stack(alloc):
 
     assert stack(vallocs[6]) == ["valloc", "test", "<module>"]
     assert vallocs[6].size == 70 * 1024
+
+    # 0 and 2 are fruit, 1 and 3 and 4 and 5 are animal, 6 is main.
+    assert vallocs[0].tid != vallocs[1].tid != vallocs[6].tid
+    assert vallocs[0].tid == vallocs[2].tid
+    assert vallocs[1].tid == vallocs[3].tid == vallocs[4].tid == vallocs[5].tid