GH-148726: Add heap_size to generational GC

Include/internal/pycore_gc.h

@@ -223,12 +223,14 @@ static inline void _PyObject_GC_TRACK(
                           "object is in generation which is garbage collected",
                           filename, lineno, __func__);
 
-    PyGC_Head *generation0 = _PyInterpreterState_GET()->gc.generation0;
+    struct _gc_runtime_state *gcstate = &_PyInterpreterState_GET()->gc;
+    PyGC_Head *generation0 = gcstate->generation0;
     PyGC_Head *last = (PyGC_Head*)(generation0->_gc_prev);
     _PyGCHead_SET_NEXT(last, gc);
     _PyGCHead_SET_PREV(gc, last);
     _PyGCHead_SET_NEXT(gc, generation0);
     generation0->_gc_prev = (uintptr_t)gc;
+    gcstate->heap_size++;
 #endif
 }
 
@@ -263,6 +265,8 @@ static inline void _PyObject_GC_UNTRACK(
     _PyGCHead_SET_PREV(next, prev);
     gc->_gc_next = 0;
     gc->_gc_prev &= _PyGC_PREV_MASK_FINALIZED;
+    struct _gc_runtime_state *gcstate = &_PyInterpreterState_GET()->gc;
+    gcstate->heap_size--;
 #endif
 }
 

Include/internal/pycore_interp_structs.h

@@ -181,6 +181,8 @@ struct gc_generation {
 struct gc_generation_stats {
     PyTime_t ts_start;
     PyTime_t ts_stop;
+    /* heap_size on the start of the collection */
+    Py_ssize_t heap_size;
     /* total number of collections */
     Py_ssize_t collections;
     /* total number of collected objects */
@@ -244,6 +246,9 @@ struct _gc_runtime_state {
     /* a list of callbacks to be invoked when collection is performed */
     PyObject *callbacks;
 
+    /* The number of live objects. */
+    Py_ssize_t heap_size;
+
     /* This is the number of objects that survived the last full
        collection. It approximates the number of long lived objects
        tracked by the GC.
@@ -278,7 +283,8 @@ struct _gc_runtime_state {
         { .threshold = 2000, }, \
         { .threshold = 10, }, \
         { .threshold = 10, }, \
-    },
+    }, \
+    .heap_size = 0,
 #else
 #define GC_GENERATION_INIT \
     .young = { .threshold = 2000, }, \

Lib/test/test_gc.py

@@ -1288,6 +1288,15 @@ def test_tuple_untrack_counts(self):
         # Use n // 2 just in case some other objects were collected.
         self.assertTrue(new_count - count > (n // 2))
 
+    @requires_gil_enabled('need generational GC')
+    @unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
+    def test_heap_size(self):
+        count = _testinternalcapi.get_tracked_heap_size()
+        l = []
+        self.assertEqual(count + 1, _testinternalcapi.get_tracked_heap_size())
+        del l
+        self.assertEqual(count, _testinternalcapi.get_tracked_heap_size())
+
 
 class GCCallbackTests(unittest.TestCase):
     def setUp(self):

Modules/_testinternalcapi.c

@@ -2725,8 +2725,7 @@ has_deferred_refcount(PyObject *self, PyObject *op)
 static PyObject *
 get_tracked_heap_size(PyObject *self, PyObject *Py_UNUSED(ignored))
 {
-    // Generational GC doesn't track heap_size, return -1.
-    return PyLong_FromInt64(-1);
+    return PyLong_FromInt64(PyInterpreterState_Get()->gc.heap_size);
 }
 
 static PyObject *

Python/gc.c

@@ -1405,14 +1405,15 @@ add_stats(GCState *gcstate, int gen, struct gc_generation_stats *stats)
     memcpy(cur_stats, prev_stats, sizeof(struct gc_generation_stats));
 
     cur_stats->ts_start = stats->ts_start;
-    cur_stats->ts_stop = stats->ts_stop;
+    cur_stats->heap_size = stats->heap_size;
 
     cur_stats->collections += 1;
     cur_stats->collected += stats->collected;
     cur_stats->uncollectable += stats->uncollectable;
     cur_stats->candidates += stats->candidates;
 
     cur_stats->duration += stats->duration;
+    cur_stats->ts_stop = stats->ts_stop;
 }
 
 /* This is the main function.  Read this to understand how the
@@ -1465,6 +1466,7 @@ gc_collect_main(PyThreadState *tstate, int generation, _PyGC_Reason reason)
         invoke_gc_callback(tstate, "start", generation, &stats);
     }
 
+    stats.heap_size = gcstate->heap_size;
     // ignore error: don't interrupt the GC if reading the clock fails
     (void)PyTime_PerfCounterRaw(&stats.ts_start);
     if (gcstate->debug & _PyGC_DEBUG_STATS) {
@@ -2087,6 +2089,8 @@ PyObject_GC_Del(void *op)
     PyGC_Head *g = AS_GC(op);
     if (_PyObject_GC_IS_TRACKED(op)) {
         gc_list_remove(g);
+        GCState *gcstate = get_gc_state();
+        gcstate->heap_size--;
 #ifdef Py_DEBUG
         PyObject *exc = PyErr_GetRaisedException();
         if (PyErr_WarnExplicitFormat(PyExc_ResourceWarning, "gc", 0,