Memory allocation tracking using anonymous lambda wrapper

Question

I had this idea to exploit anonymous lambdas that are immediately executed that wrap around an allocation function to track those allocations in a statically created Stats object.

This is the code I came up with after some time

  #define TW_INLINE __attribute__((always_inline))

  #include <iostream>
  #include <cstdint>

  using u32 = std::uint32_t;
  using cstring = const char*;

  using namespace std;

  struct MemoryBlock {};

  template<typename T>
  TW_INLINE MemoryBlock twalloc(T& alloc_, u32 size_) {
    // normally an alloc_.Allocate() would allocaty memory somehow, might be
    // malloc for HeapAllocator I just stub it here 
    malloc(128);    
    return MemoryBlock{}; // this would hold allocation info 
  }

  struct AllocationInfo
  {
    static AllocationInfo* s_head;
    static AllocationInfo* s_last;

    u32 allocationCount;
    cstring file;
    u32 line;
    u32 allocationSize;
    AllocationInfo* next;

    AllocationInfo(cstring file_, u32 line_, u32 size_);
    ~AllocationInfo() = default;

    static void DumpAllocationInfo();
  };

  template<typename T>
  TW_INLINE MemoryBlock _twalloc(T& alloc_, u32 size_) {
    return twalloc(alloc_, size_);
  }


  #define twalloc(alloc, size)                                                \
    [](auto& alloc_, u32 size_, cstring file_, u32 line_) {                   \
      static AllocationInfo info{ file_, line_, size_ };                      \
      ++info.allocationCount;                                                 \
      return _twalloc(alloc_, size_);                                         \
    }(alloc, size, __FILE__, __LINE__)                                      




  AllocationInfo* AllocationInfo::s_head = nullptr;

  AllocationInfo* AllocationInfo::s_last = nullptr;

  AllocationInfo::AllocationInfo(cstring file_, u32 line_, u32 size_)
    : file(file_)
    , line(line_)
    , allocationSize(size_)
    , allocationCount(0)
    , next(nullptr)
  {
    if (!s_head) {
      s_head = this;
      s_last = s_head;
    }
    else {
      s_last = s_last->next = this;
    }
  }

  void AllocationInfo::DumpAllocationInfo() {

    for (auto n = s_head; n; n = n->next) {
      printf("Allocation %s (%d) size: %d, count: %d\n",
        n->file, n->line, n->allocationSize, n->allocationCount);
    }
  }


  template<typename TAllocator>
  void SomeFunc(TAllocator& alloc, u32 i) {

    auto block1 = twalloc(alloc, 32);
    if (i % 2) {
      auto block2 = twalloc(alloc, 32);
    }
    if (i > 8) {
      auto block3 = twalloc(alloc, 32);
    }
  }

  // some allocator type that has Allocate Deallocate functions 
  // not needed here for the tracking demo
  struct AnAllocator { 
  };

  int main() {
    AnAllocator alloc;

    for (int i = 0; i < 10; ++i) {
      SomeFunc(alloc, i);
    }
    AllocationInfo::DumpAllocationInfo();
    return 0;
  }

This actually requires C++14 but only because I use templated twalloc function so the lambda wrapper needs this auto to work. This is not required though if Allocator is some opaque base interface

This code essentialy works on every compiler I checked (GCC 4.9, Clang 3.9, latest MSVC) but I am worried about eventual code bloat.

I am actually lacking assembly knowledge and I tried to understand what gets generated but it was hard for me to follow.

This is GCC 4.9 output using -Os:

AllocationInfo::AllocationInfo(char const*, unsigned int, unsigned int):
        cmpq    $0, AllocationInfo::s_head(%rip)
        movl    $0, (%rdi)
        movq    %rsi, 8(%rdi)
        movl    %edx, 16(%rdi)
        movl    %ecx, 20(%rdi)
        movq    $0, 24(%rdi)
        jne     .L2
        movq    %rdi, AllocationInfo::s_head(%rip)
        jmp     .L4
.L2:
        movq    AllocationInfo::s_last(%rip), %rax
        movq    %rdi, 24(%rax)
.L4:
        movq    %rdi, AllocationInfo::s_last(%rip)
        ret
.LC1:
        .string "Allocation %s (%d) size: %d, count: %d\n"
AllocationInfo::DumpAllocationInfo():
        pushq   %rbx
        movq    AllocationInfo::s_head(%rip), %rbx
.L7:
        testq   %rbx, %rbx
        je      .L10
        movl    20(%rbx), %ecx
        movl    16(%rbx), %edx
        movl    $.LC1, %edi
        movq    8(%rbx), %rsi
        movl    (%rbx), %r8d
        xorl    %eax, %eax
        call    printf
        movq    24(%rbx), %rbx
        jmp     .L7
.L10:
        popq    %rbx
        ret
.LC3:
        .string "/tmp/gcc-explorer-compiler116818-58-17dtlch/example.cpp"
void SomeFunc<AnAllocator>(AnAllocator&, unsigned int):
        cmpb    $0, _ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info(%rip)
        pushq   %rbx
        movl    %esi, %ebx
        jne     .L13
        movl    $_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info, %edi
        call    __cxa_guard_acquire
        testl   %eax, %eax
        je      .L13
        movl    $_ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info, %edi
        movl    $32, %ecx
        movl    $86, %edx
        movl    $.LC3, %esi
        call    AllocationInfo::AllocationInfo(char const*, unsigned int, unsigned int)
        movl    $_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info, %edi
        call    __cxa_guard_release
.L13:
        incl    _ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info(%rip)
        testb   $1, %bl
        je      .L15
        cmpb    $0, _ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info(%rip)
        jne     .L17
        movl    $_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info, %edi
        call    __cxa_guard_acquire
        testl   %eax, %eax
        je      .L17
        movl    $_ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info, %edi
        movl    $32, %ecx
        movl    $88, %edx
        movl    $.LC3, %esi
        call    AllocationInfo::AllocationInfo(char const*, unsigned int, unsigned int)
        movl    $_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info, %edi
        call    __cxa_guard_release
.L17:
        incl    _ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info(%rip)
.L15:
        cmpl    $8, %ebx
        jbe     .L11
        cmpb    $0, _ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info(%rip)
        jne     .L21
        movl    $_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info, %edi
        call    __cxa_guard_acquire
        testl   %eax, %eax
        je      .L21
        movl    $_ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info, %edi
        movl    $32, %ecx
        movl    $91, %edx
        movl    $.LC3, %esi
        call    AllocationInfo::AllocationInfo(char const*, unsigned int, unsigned int)
        movl    $_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info, %edi
        call    __cxa_guard_release
.L21:
        incl    _ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info(%rip)
.L11:
        popq    %rbx
        ret
main:
        pushq   %rbx
        xorl    %ebx, %ebx
        subq    $16, %rsp
.L37:
        leaq    15(%rsp), %rdi
        movl    %ebx, %esi
        incl    %ebx
        call    void SomeFunc<AnAllocator>(AnAllocator&, unsigned int)
        cmpl    $10, %ebx
        jne     .L37
        call    AllocationInfo::DumpAllocationInfo()
        addq    $16, %rsp
        xorl    %eax, %eax
        popq    %rbx
        ret
        pushq   %rax
        movl    std::__ioinit, %edi
        call    std::ios_base::Init::Init()
        popq    %rcx
        movl    $__dso_handle, %edx
        movl    std::__ioinit, %esi
        movl    std::ios_base::Init::~Init(), %edi
        jmp     __cxa_atexit
_ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info:
        .zero   32
_ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info:
        .zero   32
_ZZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info:
        .zero   32
_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE1_clIS0_EEDaRT_jS4_jE4info:
        .zero   8
_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE0_clIS0_EEDaRT_jS4_jE4info:
        .zero   8
_ZGVZZ8SomeFuncI11AnAllocatorEvRT_jENKUlS2_jPKcjE_clIS0_EEDaRT_jS4_jE4info:
        .zero   8
AllocationInfo::s_last:
        .zero   8
AllocationInfo::s_head:
        .zero   8

Am I guessing correctly that all those lambda invocations got inlined by the compiler inside the SomeFunc? In this case is this solution feasible or you wouldn't recommend it?

This is something I would like to use during development not in production executable and I actually find this to be more elegant than prepending allocations with some header and tracking all allocations elsewhere. This seems dead simple solution that seems to work.

Edit: What about using atomic bool flag to manually initialize the memory region once? So instead of static object, which causes the compiler to emit those __cxa_acquire/release that eventually lead to complex mutex lock, use a raw uninitialized static buffer that after the atomic check passes does the initialization once.

The __cxa_acquire/release are gone and I see this instead

lock cmpxchgb %dl, lambda()::initialized(%rip)

This should be faster? But it can't be that simple.

Answer 1

Your approach avoids any additional heap memory allocations for the allocation record structure, but I'm not convinced it is the best approach for the following reasons:

• It will definitely bloat the executable image since each and every allocation call (executed or not) adds a static AllocationInfo instance to the prog data.

• This is not thread safe, unless you make each AllocationInfo thread-local, which would further degrade performance. But actually, C++11 requires static variable initialization to be thread-safe, and we can see in the assembly output that the compiler generates calls to mutex locks before the constructor calls (__cxa_guard_acquire, __cxa_guard_release), but not around the allocationCount increment, so you're paying a performance cost already for partial thread-safety.

• There is a subtle problem with this code. If a client of your API calls twalloc from a DLL and that DLL gets unloaded before the main process exits, then pointers to the static instances of AllocationInfo created inside the DLL will no longer be valid and the program would crash when trying to walk the memory blocks at the end of main.

Alternatives

Other ways of keeping the allocation records that you should explore, with pros and cons:

1. Append or prepend the records to the user block:

   • Takes no additional heap allocations and will only consume space once allocations are actually ran, adding only a small constant overhead per allocation.
   • Record data might get corrupted is the user stomps over your header/footer. This is actually quite common and the main reason to avoid this approach.

2. Store the records in a separate pool:

   • No risk of having the user stomping over your data.
   • Will only consume space once allocations are actually ran.
   • Records can be kept together in a linear pool, which might improve search and iteration due to better CPU cache usage.
   • Might require dynamically allocating a new chunk of memory, but this can be mitigated with the use of memory pools and simple recycling schemes.
   • Requires some form of lookup to associate a pointer to the corresponding allocation record (think of a hash-table), so there might be some overhead when deallocating a user block.

Thread safety and contention

Furthermore, you should try be avoid thread contention (locking) as much as possible. Since your allocation function already expects and allocator object, then you should keep the allocation records inside that allocator. This way each thread can declare is own private allocator and avoid the necessity of mutex locks altogether.