116 lines
2.2 KiB
C++
116 lines
2.2 KiB
C++
#include <stdlib.h>
|
|
#include <map>
|
|
#include <set>
|
|
#include "bandit/bandit.h"
|
|
|
|
using std::map;
|
|
using std::set;
|
|
|
|
bool _enabled = false;
|
|
static size_t _allocation_count = 0;
|
|
static map<void *, size_t> _outstanding_allocations;
|
|
static size_t _allocation_failure_index = -1;
|
|
|
|
namespace record_alloc {
|
|
|
|
void start() {
|
|
_enabled = true;
|
|
_allocation_count = 0;
|
|
_outstanding_allocations.clear();
|
|
_allocation_failure_index = -1;
|
|
}
|
|
|
|
void stop() {
|
|
_enabled = false;
|
|
}
|
|
|
|
void fail_at_allocation_index(size_t failure_index) {
|
|
_allocation_failure_index = failure_index;
|
|
}
|
|
|
|
set<size_t> outstanding_allocation_indices() {
|
|
set<size_t> result;
|
|
for (const auto &entry : _outstanding_allocations) {
|
|
result.insert(entry.second);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
size_t allocation_count() {
|
|
return _allocation_count;
|
|
}
|
|
|
|
} // namespace record_alloc
|
|
|
|
static bool can_allocate() {
|
|
if (!_enabled)
|
|
return true;
|
|
|
|
return _allocation_count < _allocation_failure_index;
|
|
}
|
|
|
|
static void *record_allocation(void *result) {
|
|
if (!_enabled)
|
|
return result;
|
|
|
|
_outstanding_allocations[result] = _allocation_count;
|
|
_allocation_count++;
|
|
return result;
|
|
}
|
|
|
|
static void record_allocation_failure() {
|
|
_allocation_count++;
|
|
}
|
|
|
|
static void record_deallocation(void *pointer) {
|
|
if (!_enabled)
|
|
return;
|
|
|
|
auto entry = _outstanding_allocations.find(pointer);
|
|
if (entry != _outstanding_allocations.end()) {
|
|
_outstanding_allocations.erase(entry);
|
|
}
|
|
}
|
|
|
|
extern "C" {
|
|
|
|
void *ts_record_malloc(size_t size) {
|
|
if (can_allocate()) {
|
|
return record_allocation(malloc(size));
|
|
} else {
|
|
record_allocation_failure();
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
void *ts_record_realloc(void *pointer, size_t size) {
|
|
if (can_allocate()) {
|
|
record_deallocation(pointer);
|
|
return record_allocation(realloc(pointer, size));
|
|
} else {
|
|
record_allocation_failure();
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
void *ts_record_calloc(size_t count, size_t size) {
|
|
if (can_allocate()) {
|
|
return record_allocation(calloc(count, size));
|
|
} else {
|
|
record_allocation_failure();
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
void ts_record_free(void *pointer) {
|
|
free(pointer);
|
|
record_deallocation(pointer);
|
|
}
|
|
|
|
bool ts_record_allocations_toggle(bool value) {
|
|
bool previous_value = _enabled;
|
|
_enabled = value;
|
|
return previous_value;
|
|
}
|
|
|
|
}
|