System: allocation_tracker

Data structure

typedef struct {
  uint8_t allocator_id;
  void* ptr;
  size_t size;
  bool freed;
} allocation_t;

static const size_t canary_size = 8;
static char canary[canary_size];
static std::unordered_map<void*, allocation_t*> allocations;
static std::mutex tracker_lock;
static bool enabled = false;

// Memory allocation statistics
static size_t alloc_counter = 0;
static size_t free_counter = 0;
static size_t alloc_total_size = 0;
static size_t free_total_size = 0;

Initialize

void allocation_tracker_init(void) {
  std::unique_lock<std::mutex> lock(tracker_lock);
  if (enabled) return;

  // randomize the canary contents
  for (size_t i = 0; i < canary_size; i++) canary[i] = (char)osi_rand();

  LOG_INFO("canary initialized");

  enabled = true;
}

Head + User + Tail size

size_t allocation_tracker_resize_for_canary(size_t size) {
  return (!enabled) ? size : size + (2 * canary_size);
}

Alloc

void* allocation_tracker_notify_alloc(uint8_t allocator_id, void* ptr,
                                      size_t requested_size) {
  char* return_ptr;
  {
    std::unique_lock<std::mutex> lock(tracker_lock);
    if (!enabled || !ptr) return ptr;

    // Keep statistics
    alloc_counter++;
    alloc_total_size += allocation_tracker_resize_for_canary(requested_size);

    return_ptr = ((char*)ptr) + canary_size;

    auto map_entry = allocations.find(return_ptr);
    allocation_t* allocation;
    if (map_entry != allocations.end()) {
      allocation = map_entry->second;
      CHECK(allocation->freed);  // Must have been freed before
    } else {
      allocation = (allocation_t*)calloc(1, sizeof(allocation_t));
      allocations[return_ptr] = allocation;
    }

    allocation->allocator_id = allocator_id;
    allocation->freed = false;
    allocation->size = requested_size;
    allocation->ptr = return_ptr;
  }

  // Add the canary on both sides
  memcpy(return_ptr - canary_size, canary, canary_size);
  memcpy(return_ptr + requested_size, canary, canary_size);

  return return_ptr;
}

Free

void* allocation_tracker_notify_free(UNUSED_ATTR uint8_t allocator_id,
                                     void* ptr) {
  std::unique_lock<std::mutex> lock(tracker_lock);

  if (!enabled || !ptr) return ptr;

  auto map_entry = allocations.find(ptr);
  CHECK(map_entry != allocations.end());
  allocation_t* allocation = map_entry->second;
  CHECK(allocation);          // Must have been tracked before
  CHECK(!allocation->freed);  // Must not be a double free
  CHECK(allocation->allocator_id ==
        allocator_id);  // Must be from the same allocator

  // Keep statistics
  free_counter++;
  free_total_size += allocation_tracker_resize_for_canary(allocation->size);

  allocation->freed = true;

  UNUSED_ATTR const char* beginning_canary = ((char*)ptr) - canary_size;
  UNUSED_ATTR const char* end_canary = ((char*)ptr) + allocation->size;

  for (size_t i = 0; i < canary_size; i++) {
    CHECK(beginning_canary[i] == canary[i]);
    CHECK(end_canary[i] == canary[i]);
  }

  // Free the hash map entry to avoid unlimited memory usage growth.
  // Double-free of memory is detected with "assert(allocation)" above
  // as the allocation entry will not be present.
  allocations.erase(ptr);
  free(allocation);

  return ((char*)ptr) - canary_size;
}

Clear All Allocated Memory

void allocation_tracker_reset(void) {
  std::unique_lock<std::mutex> lock(tracker_lock);
  if (!enabled) return;

  allocations.clear();
}

Caculate All Allocated Memory Size

void allocation_tracker_reset(void) {
  std::unique_lock<std::mutex> lock(tracker_lock);
  if (!enabled) return;

  allocations.clear();
}

Deinitialize

void allocation_tracker_uninit(void) {
  std::unique_lock<std::mutex> lock(tracker_lock);
  if (!enabled) return;

  allocations.clear();
  enabled = false;
}