fix get status and get total size total cnt

Author: RenjieXiao

src/memray/_memray.pyx

@@ -1402,7 +1402,11 @@ cdef class SocketReader:
             (<AllocationRecord> alloc)._reader = self._reader
             yield alloc
 
+    def get_main_thread_id(self):
+        return self._reader.get().getMainThreadTid()
 
+    def get_skipped_frames_on_main_thread(self):
+        return self._reader.get().getSkippedFramesOnMainThread()
 
     def get_current_snapshot_raw_table(self, *, bool merge_threads):
         if self._impl is NULL:

src/memray/_memray/record_reader.cpp

@@ -820,7 +820,7 @@ RecordReader::Py_GetTraceInfo(std::vector<std::pair<FrameTree::index_t, size_t>>
     }
     for (const auto& kv : d_frame_map) {
         PyObject* pkey = PyLong_FromUnsignedLong(kv.first);
-        PyObject* pvalue = kv.second.toPythonObject(d_pystring_cache);
+        PyObject* pvalue = kv.second.toPythonObjectWithEntry(d_pystring_cache);
         PyDict_SetItem(py_d_frame_map, pkey, pvalue);
         Py_XDECREF(pkey);
         Py_XDECREF(pvalue);

src/memray/_memray/records.cpp

@@ -112,4 +112,38 @@ Frame::toPythonObject(python_helpers::PyUnicode_Cache& pystring_cache) const
     PyTuple_SET_ITEM(tuple, 2, pylineno);
     return tuple;
 }
+
+PyObject*
+Frame::toPythonObjectWithEntry(python_helpers::PyUnicode_Cache& pystring_cache) const
+{
+    PyObject* pyfunction_name = pystring_cache.getUnicodeObject(function_name);
+    if (pyfunction_name == nullptr) {
+        return nullptr;
+    }
+    PyObject* pyfilename = pystring_cache.getUnicodeObject(filename);
+    if (pyfilename == nullptr) {
+        return nullptr;
+    }
+    PyObject* pylineno = PyLong_FromLong(this->lineno);
+    if (pylineno == nullptr) {
+        return nullptr;
+    }
+    PyObject* pyisentry = PyBool_FromLong(this->is_entry_frame);
+    if (pyisentry == nullptr) {
+        return nullptr;
+    }
+    PyObject* tuple = PyTuple_New(4);
+    if (tuple == nullptr) {
+        Py_DECREF(pylineno);
+        Py_XDECREF(pyisentry);
+        return nullptr;
+    }
+    Py_INCREF(pyfunction_name);
+    Py_INCREF(pyfilename);
+    PyTuple_SET_ITEM(tuple, 0, pyfunction_name);
+    PyTuple_SET_ITEM(tuple, 1, pyfilename);
+    PyTuple_SET_ITEM(tuple, 2, pylineno);
+    PyTuple_SET_ITEM(tuple, 3, pyisentry);
+    return tuple;
+}
 }  // namespace memray::tracking_api

src/memray/_memray/records.h

@@ -246,6 +246,8 @@ struct Frame
 
     PyObject* toPythonObject(python_helpers::PyUnicode_Cache& pystring_cache) const;
 
+    PyObject* toPythonObjectWithEntry(python_helpers::PyUnicode_Cache& pystring_cache) const;
+
     auto operator==(const Frame& other) const -> bool
     {
         return (function_name == other.function_name && filename == other.filename

src/memray/_memray/socket_reader_thread.cpp

@@ -93,16 +93,21 @@ BackgroundSocketReader::Py_GetSnapshotAllocationRecordsAndStatsData(bool merge_t
     std::unordered_map<int, uint64_t> cnt_by_alloc;
     std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>> top_size;
     std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>> top_cnt;
+    std::uint64_t total_size;
+    std::uint64_t total_cnt;
     {
         std::lock_guard<std::mutex> lock(d_mutex);
         stack_to_allocation = d_aggregator.getSnapshotAllocations(merge_threads);
         cnt_by_size = stats_aggregator.allocationCountBySize();
         cnt_by_alloc = stats_aggregator.allocationCountByAllocator();
         top_size = stats_aggregator.topLocationsBySize(largest_num);
         top_cnt = stats_aggregator.topLocationsByCount(largest_num);
+        total_cnt = stats_aggregator.totalAllocations();
+        total_size = stats_aggregator.totalBytesAllocated();
     }
     PyObject* snaps = api::Py_ListFromSnapshotAllocationRecords(stack_to_allocation);
-    PyObject* stats = Py_GetStatsData(cnt_by_size, cnt_by_alloc, top_size, top_cnt);
+    PyObject* stats =
+            Py_GetStatsData(cnt_by_size, cnt_by_alloc, top_size, top_cnt, total_size, total_cnt);
     PyObject* result = PyTuple_Pack(2, snaps, stats);
     Py_XDECREF(snaps);
     Py_XDECREF(stats);
@@ -120,7 +125,9 @@ BackgroundSocketReader::Py_GetStatsData(
         const std::unordered_map<size_t, uint64_t>& cnt_by_size,
         const std::unordered_map<int, uint64_t>& cnt_by_alloc,
         std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>>& top_size,
-        std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>>& top_cnt)
+        std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>>& top_cnt,
+        std::uint64_t total_size,
+        std::uint64_t total_cnt)
 {
     PyObject* result = PyList_New(0);
     if (result == nullptr) {
@@ -163,8 +170,21 @@ BackgroundSocketReader::Py_GetStatsData(
         return nullptr;
     }
     for (const auto& it : top_size) {
-        PyObject* pk = PyLong_FromSize_t(it.first);
-        PyObject* pv = PyLong_FromSize_t(it.second.value_or(-1));
+        //        PyObject* pk = PyLong_FromSize_t(it.second.value_or(0));
+        //        PyObject * pk = d_record_reader ->Py_GetFrame(it.second.value_or(0));
+        PyObject* pk;
+        try {  // todo: optimize
+            pk = d_record_reader->Py_GetFrame(it.second.value_or(0));
+        } catch (std::exception& e) {
+            PyObject* function = PyUnicode_FromString("");
+            PyObject* file = PyUnicode_FromString("");
+            PyObject* line = PyLong_FromLong(0);
+            pk = PyTuple_Pack(3, function, file, line);
+            Py_XDECREF(function);
+            Py_XDECREF(file);
+            Py_XDECREF(line);
+        };
+        PyObject* pv = PyLong_FromSize_t(it.first);
         PyObject* pair = PyTuple_Pack(2, pk, pv);
         PyList_Append(py_top_size, pair);
         Py_XDECREF(pk);
@@ -180,8 +200,21 @@ BackgroundSocketReader::Py_GetStatsData(
         return nullptr;
     }
     for (const auto& it : top_cnt) {
-        PyObject* pk = PyLong_FromSize_t(it.first);
-        PyObject* pv = PyLong_FromSize_t(it.second.value_or(-1));
+        //        PyObject* pk = PyLong_FromSize_t(it.second.value_or(0));
+        //        PyObject * pk = d_record_reader ->Py_GetFrame(it.second.value_or(0));
+        PyObject* pk;
+        try {  // todo: optimize
+            pk = d_record_reader->Py_GetFrame(it.second.value_or(0));
+        } catch (std::exception& e) {
+            PyObject* function = PyUnicode_FromString("");
+            PyObject* file = PyUnicode_FromString("");
+            PyObject* line = PyLong_FromLong(0);
+            pk = PyTuple_Pack(3, function, file, line);
+            Py_XDECREF(function);
+            Py_XDECREF(file);
+            Py_XDECREF(line);
+        };
+        PyObject* pv = PyLong_FromSize_t(it.first);
         PyObject* pair = PyTuple_Pack(2, pk, pv);
         PyList_Append(py_top_cnt, pair);
         Py_XDECREF(pk);
@@ -191,6 +224,13 @@ BackgroundSocketReader::Py_GetStatsData(
     PyList_Append(result, py_top_cnt);
     Py_XDECREF(py_top_cnt);
 
+    PyObject* py_total_size = PyLong_FromUnsignedLong(total_size);
+    PyObject* py_total_cnt = PyLong_FromUnsignedLong(total_cnt);
+    PyList_Append(result, py_total_size);
+    PyList_Append(result, py_total_cnt);
+    Py_XDECREF(py_total_size);
+    Py_XDECREF(py_total_cnt);
+
     return result;
 }
 

src/memray/_memray/socket_reader_thread.h

@@ -45,7 +45,9 @@ class BackgroundSocketReader
             const std::unordered_map<size_t, uint64_t>& cnt_by_size,
             const std::unordered_map<int, uint64_t>& cnt_by_alloc,
             std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>>& top_size,
-            std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>>& top_cnt);
+            std::vector<std::pair<uint64_t, std::optional<memray::tracking_api::frame_id_t>>>& top_cnt,
+            std::uint64_t total_size,
+            std::uint64_t total_cnt);
     PyObject* Py_GetSnapshotAllocationRecordsAndStatsData(bool merge_threads, int largest_num);
 };