Skip to content

Commit 990e118

Browse files
authored
Handle more memray allocation types (#578)
1 parent 50181e0 commit 990e118

File tree

1 file changed

+13
-2
lines changed

1 file changed

+13
-2
lines changed

cubed/diagnostics/memray.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,13 @@
1212

1313

1414
class AllocationType(Enum):
15+
# integer values match memray AllocatorType
1516
MALLOC = 1
1617
FREE = 2
1718
CALLOC = 3
1819
REALLOC = 4
20+
MMAP = 10
21+
MUNMAP = 11
1922

2023

2124
@dataclass()
@@ -75,6 +78,8 @@ def get_allocations_over_threshold(result_file, mem_threshold):
7578
allocation_type = AllocationType.CALLOC
7679
elif a.allocator == memray.AllocatorType.REALLOC:
7780
allocation_type = AllocationType.REALLOC
81+
elif a.allocator == memray.AllocatorType.MMAP:
82+
allocation_type = AllocationType.MMAP
7883
else:
7984
raise ValueError(f"Unsupported memray.AllocatorType {a.allocator}")
8085
allocation = Allocation(
@@ -88,13 +93,19 @@ def get_allocations_over_threshold(result_file, mem_threshold):
8893
address_to_allocation[a.address] = allocation
8994
yield allocation
9095
elif (
91-
a.allocator == memray.AllocatorType.FREE
96+
a.allocator in (memray.AllocatorType.FREE, memray.AllocatorType.MUNMAP)
9297
and a.address in address_to_allocation
9398
):
99+
if a.allocator == memray.AllocatorType.FREE:
100+
allocation_type = AllocationType.FREE
101+
elif a.allocator == memray.AllocatorType.MUNMAP:
102+
allocation_type = AllocationType.MUNMAP
103+
else:
104+
raise ValueError(f"Unsupported memray.AllocatorType {a.allocator}")
94105
allocation = address_to_allocation.pop(a.address)
95106
yield Allocation(
96107
allocation.object_id,
97-
AllocationType.FREE,
108+
allocation_type,
98109
allocation.memory,
99110
address=a.address,
100111
)

0 commit comments

Comments
 (0)