Skip to content

Commit 8197d5f

Browse files
committed
Visit threads for getting allocated blocks
Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:
1 parent 385db5b commit 8197d5f

File tree

1 file changed

+32
-0
lines changed

1 file changed

+32
-0
lines changed

Objects/obmalloc.c

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1023,6 +1023,16 @@ get_state(void)
10231023
#define narenas_highwater (state->mgmt.narenas_highwater)
10241024
#define raw_allocated_blocks (state->mgmt.raw_allocated_blocks)
10251025

1026+
#ifdef WITH_MIMALLOC
1027+
static bool count_blocks(
1028+
const mi_heap_t* heap, const mi_heap_area_t* area,
1029+
void* block, size_t block_size, void* allocated_blocks)
1030+
{
1031+
*(size_t *)allocated_blocks += area->used;
1032+
return 1;
1033+
}
1034+
#endif
1035+
10261036
Py_ssize_t
10271037
_PyInterpreterState_GetAllocatedBlocks(PyInterpreterState *interp)
10281038
{
@@ -1034,6 +1044,28 @@ _PyInterpreterState_GetAllocatedBlocks(PyInterpreterState *interp)
10341044
"the interpreter doesn't have its own allocator");
10351045
}
10361046
#endif
1047+
1048+
#ifdef WITH_MIMALLOC
1049+
// TODO(sgross): this only counts the current thread's blocks.
1050+
if (_PyMem_MimallocEnabled()) {
1051+
size_t allocated_blocks = 0;
1052+
1053+
PyThreadState *tstate = PyInterpreterState_ThreadHead(interp);
1054+
while (tstate != NULL) {
1055+
mi_heap_tag_t tags[] = {
1056+
mi_heap_tag_default,
1057+
mi_heap_tag_obj,
1058+
};
1059+
for (size_t i = 0; i != sizeof(tags)/sizeof(*tags); i++) {
1060+
mi_heap_t *heap = tstate->heaps[i];
1061+
mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
1062+
}
1063+
tstate = PyThreadState_Next(tstate);
1064+
}
1065+
return allocated_blocks;
1066+
}
1067+
#endif
1068+
10371069
OMState *state = &interp->obmalloc;
10381070

10391071
Py_ssize_t n = raw_allocated_blocks;

0 commit comments

Comments
 (0)