#include "Python.h"

#include "hashtable.h"

#include "frameobject.h"

#include "pythread.h"

#include "osdefs.h"

#include "clinic/_tracemalloc.c.h"

#define TRACE_RAW_MALLOC

static void tracemalloc_stop(void);

static void* raw_malloc(size_t size);

static void raw_free(void *ptr);

#ifdef Py_DEBUG

# define TRACE_DEBUG

static struct {

PyMemAllocatorEx mem;

PyMemAllocatorEx raw;

PyMemAllocatorEx obj;

} allocators;

static struct {

/* Module initialized?

enum {

TRACEMALLOC_NOT_INITIALIZED,

TRACEMALLOC_INITIALIZED,

TRACEMALLOC_FINALIZED

} initialized;

/* Is tracemalloc tracing memory allocations?

int tracing;

/* limit of the number of frames in a traceback, 1 by default.

int max_nframe;

/* use domain in trace key?

int use_domain;

} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 1, 0};

#if defined(TRACE_RAW_MALLOC) && defined(WITH_THREAD)

static PyThread_type_lock tables_lock;

# define TABLES_LOCK() PyThread_acquire_lock(tables_lock, 1)

# define TABLES_UNLOCK() PyThread_release_lock(tables_lock)

/* variables are protected by the GIL */

# define TABLES_LOCK()

# define TABLES_UNLOCK()

#define DEFAULT_DOMAIN 0

typedef struct

#ifdef __GNUC__

__attribute__((packed))

{

uintptr_t ptr;

unsigned int domain;

} pointer_t;

typedef struct

#ifdef __GNUC__

__attribute__((packed))

#elif defined(_MSC_VER)

#pragma pack(push, 4)

{

/* filename cannot be NULL: "<unknown>" is used if the Python frame

PyObject *filename;

unsigned int lineno;

} frame_t;

#ifdef _MSC_VER

#pragma pack(pop)

typedef struct {

Py_uhash_t hash;

int nframe;

frame_t frames[1];

} traceback_t;

#define TRACEBACK_SIZE(NFRAME) \

(sizeof(traceback_t) + sizeof(frame_t) * (NFRAME - 1))

#define MAX_NFRAME \

((INT_MAX - (int)sizeof(traceback_t)) / (int)sizeof(frame_t) + 1)

static PyObject *unknown_filename = NULL;

static traceback_t tracemalloc_empty_traceback;

typedef struct {

/* Size of the memory block in bytes */

size_t size;

/* Traceback where the memory block was allocated */

traceback_t *traceback;

} trace_t;

static size_t tracemalloc_traced_memory = 0;

static size_t tracemalloc_peak_traced_memory = 0;

static _Py_hashtable_t *tracemalloc_filenames = NULL;

static traceback_t *tracemalloc_traceback = NULL;

static _Py_hashtable_t *tracemalloc_tracebacks = NULL;

static _Py_hashtable_t *tracemalloc_traces = NULL;

#ifdef TRACE_DEBUG

static void

tracemalloc_error(const char *format, ...)

{

va_list ap;

fprintf(stderr, "tracemalloc: ");

va_start(ap, format);

vfprintf(stderr, format, ap);

va_end(ap);

fprintf(stderr, "\n");

fflush(stderr);

}

#if defined(WITH_THREAD) && defined(TRACE_RAW_MALLOC)

#define REENTRANT_THREADLOCAL

#if !(defined(_POSIX_THREADS) || defined(NT_THREADS))

# error "need native thread local storage (TLS)"

static int tracemalloc_reentrant_key = -1;

#define REENTRANT Py_True

static int

get_reentrant(void)

{

void *ptr;

assert(tracemalloc_reentrant_key != -1);

ptr = PyThread_get_key_value(tracemalloc_reentrant_key);

if (ptr != NULL) {

assert(ptr == REENTRANT);

return 1;

}

else

return 0;

}

static void

set_reentrant(int reentrant)

{

assert(reentrant == 0 || reentrant == 1);

assert(tracemalloc_reentrant_key != -1);

if (reentrant) {

assert(!get_reentrant());

PyThread_set_key_value(tracemalloc_reentrant_key, REENTRANT);

}

else {

assert(get_reentrant());

PyThread_set_key_value(tracemalloc_reentrant_key, NULL);

}

}

static int tracemalloc_reentrant = 0;

static int

get_reentrant(void)

{

return tracemalloc_reentrant;

}

static void

set_reentrant(int reentrant)

{

assert(reentrant != tracemalloc_reentrant);

tracemalloc_reentrant = reentrant;

}

static Py_uhash_t

hashtable_hash_pyobject(_Py_hashtable_t *ht, const void *pkey)

{

PyObject *obj;

_Py_HASHTABLE_READ_KEY(ht, pkey, obj);

return PyObject_Hash(obj);

}

static int

hashtable_compare_unicode(_Py_hashtable_t *ht, const void *pkey,

const _Py_hashtable_entry_t *entry)

{

PyObject *key1, *key2;

_Py_HASHTABLE_READ_KEY(ht, pkey, key1);

_Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, key2);

if (key1 != NULL && key2 != NULL)

return (PyUnicode_Compare(key1, key2) == 0);

else

return key1 == key2;

}

static Py_uhash_t

hashtable_hash_pointer_t(_Py_hashtable_t *ht, const void *pkey)

{

pointer_t ptr;

Py_uhash_t hash;

_Py_HASHTABLE_READ_KEY(ht, pkey, ptr);

hash = (Py_uhash_t)_Py_HashPointer((void*)ptr.ptr);

hash ^= ptr.domain;

return hash;

}

static int

hashtable_compare_pointer_t(_Py_hashtable_t *ht, const void *pkey,

const _Py_hashtable_entry_t *entry)

{

pointer_t ptr1, ptr2;

_Py_HASHTABLE_READ_KEY(ht, pkey, ptr1);

_Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, ptr2);

/* compare pointer before domain, because pointer is more likely to be

return (ptr1.ptr == ptr2.ptr && ptr1.domain == ptr2.domain);

}

static _Py_hashtable_t *

hashtable_new(size_t key_size, size_t data_size,

_Py_hashtable_hash_func hash_func,

_Py_hashtable_compare_func compare_func)

{

_Py_hashtable_allocator_t hashtable_alloc = {malloc, free};

return _Py_hashtable_new_full(key_size, data_size, 0,

hash_func, compare_func,

&hashtable_alloc);

}

static void*

raw_malloc(size_t size)

{

return allocators.raw.malloc(allocators.raw.ctx, size);

}

static void

raw_free(void *ptr)

{

allocators.raw.free(allocators.raw.ctx, ptr);

}

static Py_uhash_t

hashtable_hash_traceback(_Py_hashtable_t *ht, const void *pkey)

{

traceback_t *traceback;

_Py_HASHTABLE_READ_KEY(ht, pkey, traceback);

return traceback->hash;

}

static int

hashtable_compare_traceback(_Py_hashtable_t *ht, const void *pkey,

const _Py_hashtable_entry_t *entry)

{

traceback_t *traceback1, *traceback2;

const frame_t *frame1, *frame2;

int i;

_Py_HASHTABLE_READ_KEY(ht, pkey, traceback1);

_Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, traceback2);

if (traceback1->nframe != traceback2->nframe)

return 0;

for (i=0; i < traceback1->nframe; i++) {

frame1 = &traceback1->frames[i];

frame2 = &traceback2->frames[i];

if (frame1->lineno != frame2->lineno)

return 0;

if (frame1->filename != frame2->filename) {

assert(PyUnicode_Compare(frame1->filename, frame2->filename) != 0);

return 0;

}

}

return 1;

}

static void

tracemalloc_get_frame(PyFrameObject *pyframe, frame_t *frame)

{

PyCodeObject *code;

PyObject *filename;

_Py_hashtable_entry_t *entry;

int lineno;

frame->filename = unknown_filename;

lineno = PyFrame_GetLineNumber(pyframe);

if (lineno < 0)

lineno = 0;

frame->lineno = (unsigned int)lineno;

code = pyframe->f_code;

if (code == NULL) {

#ifdef TRACE_DEBUG

tracemalloc_error("failed to get the code object of the frame");

return;

}

if (code->co_filename == NULL) {

#ifdef TRACE_DEBUG

tracemalloc_error("failed to get the filename of the code object");

return;

}

filename = code->co_filename;

assert(filename != NULL);

if (filename == NULL)

return;

if (!PyUnicode_Check(filename)) {

#ifdef TRACE_DEBUG

tracemalloc_error("filename is not a unicode string");

return;

}

if (!PyUnicode_IS_READY(filename)) {

/* Don't make a Unicode string ready to avoid reentrant calls

#ifdef TRACE_DEBUG

tracemalloc_error("filename is not a ready unicode string");

return;

}

/* intern the filename */

entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_filenames, filename);

if (entry != NULL) {

_Py_HASHTABLE_ENTRY_READ_KEY(tracemalloc_filenames, entry, filename);

}

else {

/* tracemalloc_filenames is responsible to keep a reference

Py_INCREF(filename);

if (_Py_HASHTABLE_SET_NODATA(tracemalloc_filenames, filename) < 0) {

Py_DECREF(filename);

#ifdef TRACE_DEBUG

tracemalloc_error("failed to intern the filename");

return;

}

}

/* the tracemalloc_filenames table keeps a reference to the filename */

frame->filename = filename;

}

static Py_uhash_t

traceback_hash(traceback_t *traceback)

{

/* code based on tuplehash() of Objects/tupleobject.c */

Py_uhash_t x, y; /* Unsigned for defined overflow behavior. */

int len = traceback->nframe;

Py_uhash_t mult = _PyHASH_MULTIPLIER;

frame_t *frame;

x = 0x345678UL;

frame = traceback->frames;

while (--len >= 0) {

y = (Py_uhash_t)PyObject_Hash(frame->filename);

y ^= (Py_uhash_t)frame->lineno;

frame++;

x = (x ^ y) * mult;

/* the cast might truncate len; that doesn't change hash stability */

mult += (Py_uhash_t)(82520UL + len + len);

}

x += 97531UL;

return x;

}

static void

traceback_get_frames(traceback_t *traceback)

{

PyThreadState *tstate;

PyFrameObject *pyframe;

#ifdef WITH_THREAD

tstate = PyGILState_GetThisThreadState();

tstate = PyThreadState_Get();

if (tstate == NULL) {

#ifdef TRACE_DEBUG

tracemalloc_error("failed to get the current thread state");

return;

}

for (pyframe = tstate->frame; pyframe != NULL; pyframe = pyframe->f_back) {

tracemalloc_get_frame(pyframe, &traceback->frames[traceback->nframe]);

assert(traceback->frames[traceback->nframe].filename != NULL);

traceback->nframe++;

if (traceback->nframe == tracemalloc_config.max_nframe)

break;

}

}

static traceback_t *

traceback_new(void)

{

traceback_t *traceback;

_Py_hashtable_entry_t *entry;

#ifdef WITH_THREAD

assert(PyGILState_Check());

/* get frames */

traceback = tracemalloc_traceback;

traceback->nframe = 0;

traceback_get_frames(traceback);

if (traceback->nframe == 0)

return &tracemalloc_empty_traceback;

traceback->hash = traceback_hash(traceback);

/* intern the traceback */

entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_tracebacks, traceback);

if (entry != NULL) {

_Py_HASHTABLE_ENTRY_READ_KEY(tracemalloc_tracebacks, entry, traceback);

}

else {

traceback_t *copy;

size_t traceback_size;

traceback_size = TRACEBACK_SIZE(traceback->nframe);

copy = raw_malloc(traceback_size);

if (copy == NULL) {

#ifdef TRACE_DEBUG

tracemalloc_error("failed to intern the traceback: malloc failed");

return NULL;

}

memcpy(copy, traceback, traceback_size);

if (_Py_HASHTABLE_SET_NODATA(tracemalloc_tracebacks, copy) < 0) {

raw_free(copy);

#ifdef TRACE_DEBUG

tracemalloc_error("failed to intern the traceback: putdata failed");

return NULL;

}

traceback = copy;

}

return traceback;

}

static int

tracemalloc_use_domain_cb(_Py_hashtable_t *old_traces,

_Py_hashtable_entry_t *entry, void *user_data)

{

uintptr_t ptr;

pointer_t key;

_Py_hashtable_t *new_traces = (_Py_hashtable_t *)user_data;

const void *pdata = _Py_HASHTABLE_ENTRY_PDATA(old_traces, entry);

_Py_HASHTABLE_ENTRY_READ_KEY(old_traces, entry, ptr);

key.ptr = ptr;

key.domain = DEFAULT_DOMAIN;

return _Py_hashtable_set(new_traces,

sizeof(key), &key,

old_traces->data_size, pdata);

}

static int

tracemalloc_use_domain(void)

{

_Py_hashtable_t *new_traces = NULL;

assert(!tracemalloc_config.use_domain);

new_traces = hashtable_new(sizeof(pointer_t),

sizeof(trace_t),

hashtable_hash_pointer_t,

hashtable_compare_pointer_t);

if (new_traces == NULL) {

return -1;

}

if (_Py_hashtable_foreach(tracemalloc_traces, tracemalloc_use_domain_cb,

new_traces) < 0)

{

_Py_hashtable_destroy(new_traces);

return -1;

}

_Py_hashtable_destroy(tracemalloc_traces);

tracemalloc_traces = new_traces;

tracemalloc_config.use_domain = 1;

return 0;

}

static void

tracemalloc_remove_trace(unsigned int domain, uintptr_t ptr)

{

trace_t trace;

int removed;

assert(tracemalloc_config.tracing);

if (tracemalloc_config.use_domain) {

pointer_t key = {ptr, domain};

removed = _Py_HASHTABLE_POP(tracemalloc_traces, key, trace);

}

else {

removed = _Py_HASHTABLE_POP(tracemalloc_traces, ptr, trace);

}

if (!removed) {

return;

}

assert(tracemalloc_traced_memory >= trace.size);

tracemalloc_traced_memory -= trace.size;

}

#define REMOVE_TRACE(ptr) \

tracemalloc_remove_trace(DEFAULT_DOMAIN, (uintptr_t)(ptr))

static int

tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,

size_t size)

{

pointer_t key = {ptr, domain};

traceback_t *traceback;

trace_t trace;

_Py_hashtable_entry_t* entry;

int res;

assert(tracemalloc_config.tracing);

traceback = traceback_new();

if (traceback == NULL) {

return -1;

}

if (!tracemalloc_config.use_domain && domain != DEFAULT_DOMAIN) {

/* first trace using a non-zero domain whereas traces use compact

if (tracemalloc_use_domain() < 0) {

return -1;

}

}

if (tracemalloc_config.use_domain) {

entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, key);

}

else {

entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, ptr);

}

if (entry != NULL) {

/* the memory block is already tracked */

_Py_HASHTABLE_ENTRY_READ_DATA(tracemalloc_traces, entry, trace);

assert(tracemalloc_traced_memory >= trace.size);

tracemalloc_traced_memory -= trace.size;

trace.size = size;

trace.traceback = traceback;

_Py_HASHTABLE_ENTRY_WRITE_DATA(tracemalloc_traces, entry, trace);

}

else {

trace.size = size;

trace.traceback = traceback;

if (tracemalloc_config.use_domain) {

res = _Py_HASHTABLE_SET(tracemalloc_traces, key, trace);

}

else {

res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace);

}

if (res != 0) {

return res;

}

}

assert(tracemalloc_traced_memory <= SIZE_MAX - size);

tracemalloc_traced_memory += size;

if (tracemalloc_traced_memory > tracemalloc_peak_traced_memory)

tracemalloc_peak_traced_memory = tracemalloc_traced_memory;

return 0;

}

#define ADD_TRACE(ptr, size) \

tracemalloc_add_trace(DEFAULT_DOMAIN, (uintptr_t)(ptr), size)

static void*

tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)

{

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

void *ptr;

assert(elsize == 0 || nelem <= SIZE_MAX / elsize);

if (use_calloc)

ptr = alloc->calloc(alloc->ctx, nelem, elsize);

else

ptr = alloc->malloc(alloc->ctx, nelem * elsize);

if (ptr == NULL)

return NULL;

TABLES_LOCK();

if (ADD_TRACE(ptr, nelem * elsize) < 0) {

/* Failed to allocate a trace for the new memory block */

TABLES_UNLOCK();

alloc->free(alloc->ctx, ptr);

return NULL;

}

TABLES_UNLOCK();

return ptr;

}

static void*

tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)

{

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

void *ptr2;

ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);

if (ptr2 == NULL)

return NULL;

if (ptr != NULL) {

/* an existing memory block has been resized */

TABLES_LOCK();

/* tracemalloc_add_trace() updates the trace if there is already

if (ptr2 != ptr) {

REMOVE_TRACE(ptr);

}

if (ADD_TRACE(ptr2, new_size) < 0) {

/* Memory allocation failed. The error cannot be reported to

assert(0 && "should never happen");

}

TABLES_UNLOCK();

}

else {

/* new allocation */

TABLES_LOCK();

if (ADD_TRACE(ptr2, new_size) < 0) {

/* Failed to allocate a trace for the new memory block */

TABLES_UNLOCK();

alloc->free(alloc->ctx, ptr2);

return NULL;

}

TABLES_UNLOCK();

}

return ptr2;

}

static void

tracemalloc_free(void *ctx, void *ptr)

{

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

if (ptr == NULL)

return;

/* GIL cannot be locked in PyMem_RawFree() because it would introduce

alloc->free(alloc->ctx, ptr);

TABLES_LOCK();

REMOVE_TRACE(ptr);

TABLES_UNLOCK();

}

static void*

tracemalloc_alloc_gil(int use_calloc, void *ctx, size_t nelem, size_t elsize)

{

void *ptr;

if (get_reentrant()) {

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

if (use_calloc)

return alloc->calloc(alloc->ctx, nelem, elsize);

else

return alloc->malloc(alloc->ctx, nelem * elsize);

}

/* Ignore reentrant call. PyObjet_Malloc() calls PyMem_Malloc() for

set_reentrant(1);

ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);

set_reentrant(0);

return ptr;

}

static void*

tracemalloc_malloc_gil(void *ctx, size_t size)

{

return tracemalloc_alloc_gil(0, ctx, 1, size);

}

static void*

tracemalloc_calloc_gil(void *ctx, size_t nelem, size_t elsize)

{

return tracemalloc_alloc_gil(1, ctx, nelem, elsize);

}

static void*

tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)

{

void *ptr2;

if (get_reentrant()) {

/* Reentrant call to PyMem_Realloc() and PyMem_RawRealloc().

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);

if (ptr2 != NULL && ptr != NULL) {

TABLES_LOCK();

REMOVE_TRACE(ptr);

TABLES_UNLOCK();

}

return ptr2;

}

/* Ignore reentrant call. PyObjet_Realloc() calls PyMem_Realloc() for

set_reentrant(1);

ptr2 = tracemalloc_realloc(ctx, ptr, new_size);

set_reentrant(0);

return ptr2;

}

#ifdef TRACE_RAW_MALLOC

static void*

tracemalloc_raw_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)

{

#ifdef WITH_THREAD

PyGILState_STATE gil_state;

void *ptr;

if (get_reentrant()) {

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

if (use_calloc)

return alloc->calloc(alloc->ctx, nelem, elsize);

else

return alloc->malloc(alloc->ctx, nelem * elsize);

}

/* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()

set_reentrant(1);

#ifdef WITH_THREAD

gil_state = PyGILState_Ensure();

ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);

PyGILState_Release(gil_state);

ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);

set_reentrant(0);

return ptr;

}

static void*

tracemalloc_raw_malloc(void *ctx, size_t size)

{

return tracemalloc_raw_alloc(0, ctx, 1, size);

}

static void*

tracemalloc_raw_calloc(void *ctx, size_t nelem, size_t elsize)

{

return tracemalloc_raw_alloc(1, ctx, nelem, elsize);

}

static void*

tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)

{

#ifdef WITH_THREAD

PyGILState_STATE gil_state;

void *ptr2;

if (get_reentrant()) {

/* Reentrant call to PyMem_RawRealloc(). */

PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;

ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);

if (ptr2 != NULL && ptr != NULL) {

TABLES_LOCK();

REMOVE_TRACE(ptr);

TABLES_UNLOCK();

}

return ptr2;

}

/* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()

set_reentrant(1);

#ifdef WITH_THREAD

gil_state = PyGILState_Ensure();

ptr2 = tracemalloc_realloc(ctx, ptr, new_size);

PyGILState_Release(gil_state);

ptr2 = tracemalloc_realloc(ctx, ptr, new_size);

set_reentrant(0);

return ptr2;

}

#endif /* TRACE_RAW_MALLOC */

static int

tracemalloc_clear_filename(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,

void *user_data)

{

PyObject *filename;

_Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, filename);

Py_DECREF(filename);

return 0;

}

static int

traceback_free_traceback(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,

void *user_data)

{

traceback_t *traceback;

_Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, traceback);

raw_free(traceback);

return 0;

}

static void

tracemalloc_clear_traces(void)

{

#ifdef WITH_THREAD

/* The GIL protects variables againt concurrent access */

assert(PyGILState_Check());

TABLES_LOCK();

_Py_hashtable_clear(tracemalloc_traces);

tracemalloc_traced_memory = 0;

tracemalloc_peak_traced_memory = 0;

TABLES_UNLOCK();

_Py_hashtable_foreach(tracemalloc_tracebacks, traceback_free_traceback, NULL);

_Py_hashtable_clear(tracemalloc_tracebacks);

_Py_hashtable_foreach(tracemalloc_filenames, tracemalloc_clear_filename, NULL);

_Py_hashtable_clear(tracemalloc_filenames);

}

static int

tracemalloc_init(void)

{

if (tracemalloc_config.initialized == TRACEMALLOC_FINALIZED) {

PyErr_SetString(PyExc_RuntimeError,

"the tracemalloc module has been unloaded");

return -1;

}

if (tracemalloc_config.initialized == TRACEMALLOC_INITIALIZED)

return 0;

PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);

#ifdef REENTRANT_THREADLOCAL

tracemalloc_reentrant_key = PyThread_create_key();

if (tracemalloc_reentrant_key == -1) {