1#include "Python.h"
2#include "hashtable.h"
3#include "frameobject.h"
4#include "pythread.h"
5#include "osdefs.h"
6
7/* Trace memory blocks allocated by PyMem_RawMalloc() */
8#define TRACE_RAW_MALLOC
9
10/* Forward declaration */
11static void tracemalloc_stop(void);
12static void* raw_malloc(size_t size);
13static void raw_free(void *ptr);
14
15#ifdef Py_DEBUG
16#  define TRACE_DEBUG
17#endif
18
19/* Protected by the GIL */
20static struct {
21    PyMemAllocatorEx mem;
22    PyMemAllocatorEx raw;
23    PyMemAllocatorEx obj;
24} allocators;
25
26static struct {
27    /* Module initialized?
28       Variable protected by the GIL */
29    enum {
30        TRACEMALLOC_NOT_INITIALIZED,
31        TRACEMALLOC_INITIALIZED,
32        TRACEMALLOC_FINALIZED
33    } initialized;
34
35    /* Is tracemalloc tracing memory allocations?
36       Variable protected by the GIL */
37    int tracing;
38
39    /* limit of the number of frames in a traceback, 1 by default.
40       Variable protected by the GIL. */
41    int max_nframe;
42
43    /* use domain in trace key?
44       Variable protected by the GIL. */
45    int use_domain;
46} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 1, 0};
47
48#if defined(TRACE_RAW_MALLOC) && defined(WITH_THREAD)
49/* This lock is needed because tracemalloc_free() is called without
50   the GIL held from PyMem_RawFree(). It cannot acquire the lock because it
51   would introduce a deadlock in PyThreadState_DeleteCurrent(). */
52static PyThread_type_lock tables_lock;
53#  define TABLES_LOCK() PyThread_acquire_lock(tables_lock, 1)
54#  define TABLES_UNLOCK() PyThread_release_lock(tables_lock)
55#else
56   /* variables are protected by the GIL */
57#  define TABLES_LOCK()
58#  define TABLES_UNLOCK()
59#endif
60
61
62#define DEFAULT_DOMAIN 0
63
64/* Pack the frame_t structure to reduce the memory footprint. */
65typedef struct
66#ifdef __GNUC__
67__attribute__((packed))
68#endif
69{
70    uintptr_t ptr;
71    _PyTraceMalloc_domain_t domain;
72} pointer_t;
73
74/* Pack the frame_t structure to reduce the memory footprint on 64-bit
75   architectures: 12 bytes instead of 16. */
76typedef struct
77#ifdef __GNUC__
78__attribute__((packed))
79#elif defined(_MSC_VER)
80_declspec(align(4))
81#endif
82{
83    /* filename cannot be NULL: "<unknown>" is used if the Python frame
84       filename is NULL */
85    PyObject *filename;
86    unsigned int lineno;
87} frame_t;
88
89
90typedef struct {
91    Py_uhash_t hash;
92    int nframe;
93    frame_t frames[1];
94} traceback_t;
95
96#define TRACEBACK_SIZE(NFRAME) \
97        (sizeof(traceback_t) + sizeof(frame_t) * (NFRAME - 1))
98
99#define MAX_NFRAME \
100        ((INT_MAX - (int)sizeof(traceback_t)) / (int)sizeof(frame_t) + 1)
101
102
103static PyObject *unknown_filename = NULL;
104static traceback_t tracemalloc_empty_traceback;
105
106/* Trace of a memory block */
107typedef struct {
108    /* Size of the memory block in bytes */
109    size_t size;
110
111    /* Traceback where the memory block was allocated */
112    traceback_t *traceback;
113} trace_t;
114
115
116/* Size in bytes of currently traced memory.
117   Protected by TABLES_LOCK(). */
118static size_t tracemalloc_traced_memory = 0;
119
120/* Peak size in bytes of traced memory.
121   Protected by TABLES_LOCK(). */
122static size_t tracemalloc_peak_traced_memory = 0;
123
124/* Hash table used as a set to intern filenames:
125   PyObject* => PyObject*.
126   Protected by the GIL */
127static _Py_hashtable_t *tracemalloc_filenames = NULL;
128
129/* Buffer to store a new traceback in traceback_new().
130   Protected by the GIL. */
131static traceback_t *tracemalloc_traceback = NULL;
132
133/* Hash table used as a set to intern tracebacks:
134   traceback_t* => traceback_t*
135   Protected by the GIL */
136static _Py_hashtable_t *tracemalloc_tracebacks = NULL;
137
138/* pointer (void*) => trace (trace_t).
139   Protected by TABLES_LOCK(). */
140static _Py_hashtable_t *tracemalloc_traces = NULL;
141
142
143#ifdef TRACE_DEBUG
144static void
145tracemalloc_error(const char *format, ...)
146{
147    va_list ap;
148    fprintf(stderr, "tracemalloc: ");
149    va_start(ap, format);
150    vfprintf(stderr, format, ap);
151    va_end(ap);
152    fprintf(stderr, "\n");
153    fflush(stderr);
154}
155#endif
156
157
158#if defined(WITH_THREAD) && defined(TRACE_RAW_MALLOC)
159#define REENTRANT_THREADLOCAL
160
161/* If your OS does not provide native thread local storage, you can implement
162   it manually using a lock. Functions of thread.c cannot be used because
163   they use PyMem_RawMalloc() which leads to a reentrant call. */
164#if !(defined(_POSIX_THREADS) || defined(NT_THREADS))
165#  error "need native thread local storage (TLS)"
166#endif
167
168static int tracemalloc_reentrant_key = -1;
169
170/* Any non-NULL pointer can be used */
171#define REENTRANT Py_True
172
173static int
174get_reentrant(void)
175{
176    void *ptr;
177
178    assert(tracemalloc_reentrant_key != -1);
179    ptr = PyThread_get_key_value(tracemalloc_reentrant_key);
180    if (ptr != NULL) {
181        assert(ptr == REENTRANT);
182        return 1;
183    }
184    else
185        return 0;
186}
187
188static void
189set_reentrant(int reentrant)
190{
191    assert(reentrant == 0 || reentrant == 1);
192    assert(tracemalloc_reentrant_key != -1);
193
194    if (reentrant) {
195        assert(!get_reentrant());
196        PyThread_set_key_value(tracemalloc_reentrant_key, REENTRANT);
197    }
198    else {
199        assert(get_reentrant());
200        PyThread_set_key_value(tracemalloc_reentrant_key, NULL);
201    }
202}
203
204#else
205
206/* WITH_THREAD not defined: Python compiled without threads,
207   or TRACE_RAW_MALLOC not defined: variable protected by the GIL */
208static int tracemalloc_reentrant = 0;
209
210static int
211get_reentrant(void)
212{
213    return tracemalloc_reentrant;
214}
215
216static void
217set_reentrant(int reentrant)
218{
219    assert(reentrant != tracemalloc_reentrant);
220    tracemalloc_reentrant = reentrant;
221}
222#endif
223
224
225static Py_uhash_t
226hashtable_hash_pyobject(_Py_hashtable_t *ht, const void *pkey)
227{
228    PyObject *obj;
229
230    _Py_HASHTABLE_READ_KEY(ht, pkey, obj);
231    return PyObject_Hash(obj);
232}
233
234
235static int
236hashtable_compare_unicode(_Py_hashtable_t *ht, const void *pkey,
237                          const _Py_hashtable_entry_t *entry)
238{
239    PyObject *key1, *key2;
240
241    _Py_HASHTABLE_READ_KEY(ht, pkey, key1);
242    _Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, key2);
243
244    if (key1 != NULL && key2 != NULL)
245        return (PyUnicode_Compare(key1, key2) == 0);
246    else
247        return key1 == key2;
248}
249
250
251static Py_uhash_t
252hashtable_hash_pointer_t(_Py_hashtable_t *ht, const void *pkey)
253{
254    pointer_t ptr;
255    Py_uhash_t hash;
256
257    _Py_HASHTABLE_READ_KEY(ht, pkey, ptr);
258
259    hash = (Py_uhash_t)_Py_HashPointer((void*)ptr.ptr);
260    hash ^= ptr.domain;
261    return hash;
262}
263
264
265static int
266hashtable_compare_pointer_t(_Py_hashtable_t *ht, const void *pkey,
267                            const _Py_hashtable_entry_t *entry)
268{
269    pointer_t ptr1, ptr2;
270
271    _Py_HASHTABLE_READ_KEY(ht, pkey, ptr1);
272    _Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, ptr2);
273
274    /* compare pointer before domain, because pointer is more likely to be
275       different */
276    return (ptr1.ptr == ptr2.ptr && ptr1.domain == ptr2.domain);
277
278}
279
280
281static _Py_hashtable_t *
282hashtable_new(size_t key_size, size_t data_size,
283              _Py_hashtable_hash_func hash_func,
284              _Py_hashtable_compare_func compare_func)
285{
286    _Py_hashtable_allocator_t hashtable_alloc = {malloc, free};
287    return _Py_hashtable_new_full(key_size, data_size, 0,
288                                  hash_func, compare_func,
289                                  &hashtable_alloc);
290}
291
292
293static void*
294raw_malloc(size_t size)
295{
296    return allocators.raw.malloc(allocators.raw.ctx, size);
297}
298
299static void
300raw_free(void *ptr)
301{
302    allocators.raw.free(allocators.raw.ctx, ptr);
303}
304
305
306static Py_uhash_t
307hashtable_hash_traceback(_Py_hashtable_t *ht, const void *pkey)
308{
309    traceback_t *traceback;
310
311    _Py_HASHTABLE_READ_KEY(ht, pkey, traceback);
312    return traceback->hash;
313}
314
315
316static int
317hashtable_compare_traceback(_Py_hashtable_t *ht, const void *pkey,
318                            const _Py_hashtable_entry_t *entry)
319{
320    traceback_t *traceback1, *traceback2;
321    const frame_t *frame1, *frame2;
322    int i;
323
324    _Py_HASHTABLE_READ_KEY(ht, pkey, traceback1);
325    _Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, traceback2);
326
327    if (traceback1->nframe != traceback2->nframe)
328        return 0;
329
330    for (i=0; i < traceback1->nframe; i++) {
331        frame1 = &traceback1->frames[i];
332        frame2 = &traceback2->frames[i];
333
334        if (frame1->lineno != frame2->lineno)
335            return 0;
336
337        if (frame1->filename != frame2->filename) {
338            assert(PyUnicode_Compare(frame1->filename, frame2->filename) != 0);
339            return 0;
340        }
341    }
342    return 1;
343}
344
345
346static void
347tracemalloc_get_frame(PyFrameObject *pyframe, frame_t *frame)
348{
349    PyCodeObject *code;
350    PyObject *filename;
351    _Py_hashtable_entry_t *entry;
352    int lineno;
353
354    frame->filename = unknown_filename;
355    lineno = PyFrame_GetLineNumber(pyframe);
356    if (lineno < 0)
357        lineno = 0;
358    frame->lineno = (unsigned int)lineno;
359
360    code = pyframe->f_code;
361    if (code == NULL) {
362#ifdef TRACE_DEBUG
363        tracemalloc_error("failed to get the code object of the frame");
364#endif
365        return;
366    }
367
368    if (code->co_filename == NULL) {
369#ifdef TRACE_DEBUG
370        tracemalloc_error("failed to get the filename of the code object");
371#endif
372        return;
373    }
374
375    filename = code->co_filename;
376    assert(filename != NULL);
377    if (filename == NULL)
378        return;
379
380    if (!PyUnicode_Check(filename)) {
381#ifdef TRACE_DEBUG
382        tracemalloc_error("filename is not a unicode string");
383#endif
384        return;
385    }
386    if (!PyUnicode_IS_READY(filename)) {
387        /* Don't make a Unicode string ready to avoid reentrant calls
388           to tracemalloc_malloc() or tracemalloc_realloc() */
389#ifdef TRACE_DEBUG
390        tracemalloc_error("filename is not a ready unicode string");
391#endif
392        return;
393    }
394
395    /* intern the filename */
396    entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_filenames, filename);
397    if (entry != NULL) {
398        _Py_HASHTABLE_ENTRY_READ_KEY(tracemalloc_filenames, entry, filename);
399    }
400    else {
401        /* tracemalloc_filenames is responsible to keep a reference
402           to the filename */
403        Py_INCREF(filename);
404        if (_Py_HASHTABLE_SET_NODATA(tracemalloc_filenames, filename) < 0) {
405            Py_DECREF(filename);
406#ifdef TRACE_DEBUG
407            tracemalloc_error("failed to intern the filename");
408#endif
409            return;
410        }
411    }
412
413    /* the tracemalloc_filenames table keeps a reference to the filename */
414    frame->filename = filename;
415}
416
417
418static Py_uhash_t
419traceback_hash(traceback_t *traceback)
420{
421    /* code based on tuplehash() of Objects/tupleobject.c */
422    Py_uhash_t x, y;  /* Unsigned for defined overflow behavior. */
423    int len = traceback->nframe;
424    Py_uhash_t mult = _PyHASH_MULTIPLIER;
425    frame_t *frame;
426
427    x = 0x345678UL;
428    frame = traceback->frames;
429    while (--len >= 0) {
430        y = (Py_uhash_t)PyObject_Hash(frame->filename);
431        y ^= (Py_uhash_t)frame->lineno;
432        frame++;
433
434        x = (x ^ y) * mult;
435        /* the cast might truncate len; that doesn't change hash stability */
436        mult += (Py_uhash_t)(82520UL + len + len);
437    }
438    x += 97531UL;
439    return x;
440}
441
442
443static void
444traceback_get_frames(traceback_t *traceback)
445{
446    PyThreadState *tstate;
447    PyFrameObject *pyframe;
448
449#ifdef WITH_THREAD
450    tstate = PyGILState_GetThisThreadState();
451#else
452    tstate = PyThreadState_Get();
453#endif
454    if (tstate == NULL) {
455#ifdef TRACE_DEBUG
456        tracemalloc_error("failed to get the current thread state");
457#endif
458        return;
459    }
460
461    for (pyframe = tstate->frame; pyframe != NULL; pyframe = pyframe->f_back) {
462        tracemalloc_get_frame(pyframe, &traceback->frames[traceback->nframe]);
463        assert(traceback->frames[traceback->nframe].filename != NULL);
464        traceback->nframe++;
465        if (traceback->nframe == tracemalloc_config.max_nframe)
466            break;
467    }
468}
469
470
471static traceback_t *
472traceback_new(void)
473{
474    traceback_t *traceback;
475    _Py_hashtable_entry_t *entry;
476
477#ifdef WITH_THREAD
478    assert(PyGILState_Check());
479#endif
480
481    /* get frames */
482    traceback = tracemalloc_traceback;
483    traceback->nframe = 0;
484    traceback_get_frames(traceback);
485    if (traceback->nframe == 0)
486        return &tracemalloc_empty_traceback;
487    traceback->hash = traceback_hash(traceback);
488
489    /* intern the traceback */
490    entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_tracebacks, traceback);
491    if (entry != NULL) {
492        _Py_HASHTABLE_ENTRY_READ_KEY(tracemalloc_tracebacks, entry, traceback);
493    }
494    else {
495        traceback_t *copy;
496        size_t traceback_size;
497
498        traceback_size = TRACEBACK_SIZE(traceback->nframe);
499
500        copy = raw_malloc(traceback_size);
501        if (copy == NULL) {
502#ifdef TRACE_DEBUG
503            tracemalloc_error("failed to intern the traceback: malloc failed");
504#endif
505            return NULL;
506        }
507        memcpy(copy, traceback, traceback_size);
508
509        if (_Py_HASHTABLE_SET_NODATA(tracemalloc_tracebacks, copy) < 0) {
510            raw_free(copy);
511#ifdef TRACE_DEBUG
512            tracemalloc_error("failed to intern the traceback: putdata failed");
513#endif
514            return NULL;
515        }
516        traceback = copy;
517    }
518    return traceback;
519}
520
521
522static int
523tracemalloc_use_domain_cb(_Py_hashtable_t *old_traces,
524                           _Py_hashtable_entry_t *entry, void *user_data)
525{
526    uintptr_t ptr;
527    pointer_t key;
528    _Py_hashtable_t *new_traces = (_Py_hashtable_t *)user_data;
529    const void *pdata = _Py_HASHTABLE_ENTRY_PDATA(old_traces, entry);
530
531    _Py_HASHTABLE_ENTRY_READ_KEY(old_traces, entry, ptr);
532    key.ptr = ptr;
533    key.domain = DEFAULT_DOMAIN;
534
535    return _Py_hashtable_set(new_traces,
536                             sizeof(key), &key,
537                             old_traces->data_size, pdata);
538}
539
540
541/* Convert tracemalloc_traces from compact key (uintptr_t) to pointer_t key.
542 * Return 0 on success, -1 on error. */
543static int
544tracemalloc_use_domain(void)
545{
546    _Py_hashtable_t *new_traces = NULL;
547
548    assert(!tracemalloc_config.use_domain);
549
550    new_traces = hashtable_new(sizeof(pointer_t),
551                               sizeof(trace_t),
552                               hashtable_hash_pointer_t,
553                               hashtable_compare_pointer_t);
554    if (new_traces == NULL) {
555        return -1;
556    }
557
558    if (_Py_hashtable_foreach(tracemalloc_traces, tracemalloc_use_domain_cb,
559                              new_traces) < 0)
560    {
561        _Py_hashtable_destroy(new_traces);
562        return -1;
563    }
564
565    _Py_hashtable_destroy(tracemalloc_traces);
566    tracemalloc_traces = new_traces;
567
568    tracemalloc_config.use_domain = 1;
569
570    return 0;
571}
572
573
574static void
575tracemalloc_remove_trace(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
576{
577    trace_t trace;
578    int removed;
579
580    assert(tracemalloc_config.tracing);
581
582    if (tracemalloc_config.use_domain) {
583        pointer_t key = {ptr, domain};
584        removed = _Py_HASHTABLE_POP(tracemalloc_traces, key, trace);
585    }
586    else {
587        removed = _Py_HASHTABLE_POP(tracemalloc_traces, ptr, trace);
588    }
589    if (!removed) {
590        return;
591    }
592
593    assert(tracemalloc_traced_memory >= trace.size);
594    tracemalloc_traced_memory -= trace.size;
595}
596
597#define REMOVE_TRACE(ptr) \
598            tracemalloc_remove_trace(DEFAULT_DOMAIN, (uintptr_t)(ptr))
599
600
601static int
602tracemalloc_add_trace(_PyTraceMalloc_domain_t domain, uintptr_t ptr,
603                      size_t size)
604{
605    pointer_t key = {ptr, domain};
606    traceback_t *traceback;
607    trace_t trace;
608    _Py_hashtable_entry_t* entry;
609    int res;
610
611    assert(tracemalloc_config.tracing);
612
613    traceback = traceback_new();
614    if (traceback == NULL) {
615        return -1;
616    }
617
618    if (!tracemalloc_config.use_domain && domain != DEFAULT_DOMAIN) {
619        /* first trace using a non-zero domain whereas traces use compact
620           (uintptr_t) keys: switch to pointer_t keys. */
621        if (tracemalloc_use_domain() < 0) {
622            return -1;
623        }
624    }
625
626    if (tracemalloc_config.use_domain) {
627        entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, key);
628    }
629    else {
630        entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, ptr);
631    }
632
633    if (entry != NULL) {
634        /* the memory block is already tracked */
635        _Py_HASHTABLE_ENTRY_READ_DATA(tracemalloc_traces, entry, trace);
636        assert(tracemalloc_traced_memory >= trace.size);
637        tracemalloc_traced_memory -= trace.size;
638
639        trace.size = size;
640        trace.traceback = traceback;
641        _Py_HASHTABLE_ENTRY_WRITE_DATA(tracemalloc_traces, entry, trace);
642    }
643    else {
644        trace.size = size;
645        trace.traceback = traceback;
646
647        if (tracemalloc_config.use_domain) {
648            res = _Py_HASHTABLE_SET(tracemalloc_traces, key, trace);
649        }
650        else {
651            res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace);
652        }
653        if (res != 0) {
654            return res;
655        }
656    }
657
658    assert(tracemalloc_traced_memory <= SIZE_MAX - size);
659    tracemalloc_traced_memory += size;
660    if (tracemalloc_traced_memory > tracemalloc_peak_traced_memory)
661        tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
662    return 0;
663}
664
665#define ADD_TRACE(ptr, size) \
666            tracemalloc_add_trace(DEFAULT_DOMAIN, (uintptr_t)(ptr), size)
667
668
669static void*
670tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
671{
672    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
673    void *ptr;
674
675    assert(elsize == 0 || nelem <= SIZE_MAX / elsize);
676
677    if (use_calloc)
678        ptr = alloc->calloc(alloc->ctx, nelem, elsize);
679    else
680        ptr = alloc->malloc(alloc->ctx, nelem * elsize);
681    if (ptr == NULL)
682        return NULL;
683
684    TABLES_LOCK();
685    if (ADD_TRACE(ptr, nelem * elsize) < 0) {
686        /* Failed to allocate a trace for the new memory block */
687        TABLES_UNLOCK();
688        alloc->free(alloc->ctx, ptr);
689        return NULL;
690    }
691    TABLES_UNLOCK();
692    return ptr;
693}
694
695
696static void*
697tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
698{
699    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
700    void *ptr2;
701
702    ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
703    if (ptr2 == NULL)
704        return NULL;
705
706    if (ptr != NULL) {
707        /* an existing memory block has been resized */
708
709        TABLES_LOCK();
710
711        /* tracemalloc_add_trace() updates the trace if there is already
712           a trace at address (domain, ptr2) */
713        if (ptr2 != ptr) {
714            REMOVE_TRACE(ptr);
715        }
716
717        if (ADD_TRACE(ptr2, new_size) < 0) {
718            /* Memory allocation failed. The error cannot be reported to
719               the caller, because realloc() may already have shrunk the
720               memory block and so removed bytes.
721
722               This case is very unlikely: a hash entry has just been
723               released, so the hash table should have at least one free entry.
724
725               The GIL and the table lock ensures that only one thread is
726               allocating memory. */
727            assert(0 && "should never happen");
728        }
729        TABLES_UNLOCK();
730    }
731    else {
732        /* new allocation */
733
734        TABLES_LOCK();
735        if (ADD_TRACE(ptr2, new_size) < 0) {
736            /* Failed to allocate a trace for the new memory block */
737            TABLES_UNLOCK();
738            alloc->free(alloc->ctx, ptr2);
739            return NULL;
740        }
741        TABLES_UNLOCK();
742    }
743    return ptr2;
744}
745
746
747static void
748tracemalloc_free(void *ctx, void *ptr)
749{
750    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
751
752    if (ptr == NULL)
753        return;
754
755     /* GIL cannot be locked in PyMem_RawFree() because it would introduce
756        a deadlock in PyThreadState_DeleteCurrent(). */
757
758    alloc->free(alloc->ctx, ptr);
759
760    TABLES_LOCK();
761    REMOVE_TRACE(ptr);
762    TABLES_UNLOCK();
763}
764
765
766static void*
767tracemalloc_alloc_gil(int use_calloc, void *ctx, size_t nelem, size_t elsize)
768{
769    void *ptr;
770
771    if (get_reentrant()) {
772        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
773        if (use_calloc)
774            return alloc->calloc(alloc->ctx, nelem, elsize);
775        else
776            return alloc->malloc(alloc->ctx, nelem * elsize);
777    }
778
779    /* Ignore reentrant call. PyObjet_Malloc() calls PyMem_Malloc() for
780       allocations larger than 512 bytes, don't trace the same memory
781       allocation twice. */
782    set_reentrant(1);
783
784    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
785
786    set_reentrant(0);
787    return ptr;
788}
789
790
791static void*
792tracemalloc_malloc_gil(void *ctx, size_t size)
793{
794    return tracemalloc_alloc_gil(0, ctx, 1, size);
795}
796
797
798static void*
799tracemalloc_calloc_gil(void *ctx, size_t nelem, size_t elsize)
800{
801    return tracemalloc_alloc_gil(1, ctx, nelem, elsize);
802}
803
804
805static void*
806tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
807{
808    void *ptr2;
809
810    if (get_reentrant()) {
811        /* Reentrant call to PyMem_Realloc() and PyMem_RawRealloc().
812           Example: PyMem_RawRealloc() is called internally by pymalloc
813           (_PyObject_Malloc() and  _PyObject_Realloc()) to allocate a new
814           arena (new_arena()). */
815        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
816
817        ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
818        if (ptr2 != NULL && ptr != NULL) {
819            TABLES_LOCK();
820            REMOVE_TRACE(ptr);
821            TABLES_UNLOCK();
822        }
823        return ptr2;
824    }
825
826    /* Ignore reentrant call. PyObjet_Realloc() calls PyMem_Realloc() for
827       allocations larger than 512 bytes. Don't trace the same memory
828       allocation twice. */
829    set_reentrant(1);
830
831    ptr2 = tracemalloc_realloc(ctx, ptr, new_size);
832
833    set_reentrant(0);
834    return ptr2;
835}
836
837
838#ifdef TRACE_RAW_MALLOC
839static void*
840tracemalloc_raw_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
841{
842#ifdef WITH_THREAD
843    PyGILState_STATE gil_state;
844#endif
845    void *ptr;
846
847    if (get_reentrant()) {
848        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
849        if (use_calloc)
850            return alloc->calloc(alloc->ctx, nelem, elsize);
851        else
852            return alloc->malloc(alloc->ctx, nelem * elsize);
853    }
854
855    /* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()
856       indirectly which would call PyGILState_Ensure() if reentrant are not
857       disabled. */
858    set_reentrant(1);
859
860#ifdef WITH_THREAD
861    gil_state = PyGILState_Ensure();
862    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
863    PyGILState_Release(gil_state);
864#else
865    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
866#endif
867
868    set_reentrant(0);
869    return ptr;
870}
871
872
873static void*
874tracemalloc_raw_malloc(void *ctx, size_t size)
875{
876    return tracemalloc_raw_alloc(0, ctx, 1, size);
877}
878
879
880static void*
881tracemalloc_raw_calloc(void *ctx, size_t nelem, size_t elsize)
882{
883    return tracemalloc_raw_alloc(1, ctx, nelem, elsize);
884}
885
886
887static void*
888tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)
889{
890#ifdef WITH_THREAD
891    PyGILState_STATE gil_state;
892#endif
893    void *ptr2;
894
895    if (get_reentrant()) {
896        /* Reentrant call to PyMem_RawRealloc(). */
897        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
898
899        ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
900
901        if (ptr2 != NULL && ptr != NULL) {
902            TABLES_LOCK();
903            REMOVE_TRACE(ptr);
904            TABLES_UNLOCK();
905        }
906        return ptr2;
907    }
908
909    /* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()
910       indirectly which would call PyGILState_Ensure() if reentrant calls are
911       not disabled. */
912    set_reentrant(1);
913
914#ifdef WITH_THREAD
915    gil_state = PyGILState_Ensure();
916    ptr2 = tracemalloc_realloc(ctx, ptr, new_size);
917    PyGILState_Release(gil_state);
918#else
919    ptr2 = tracemalloc_realloc(ctx, ptr, new_size);
920#endif
921
922    set_reentrant(0);
923    return ptr2;
924}
925#endif   /* TRACE_RAW_MALLOC */
926
927
928static int
929tracemalloc_clear_filename(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
930                           void *user_data)
931{
932    PyObject *filename;
933
934    _Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, filename);
935    Py_DECREF(filename);
936    return 0;
937}
938
939
940static int
941traceback_free_traceback(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
942                         void *user_data)
943{
944    traceback_t *traceback;
945
946    _Py_HASHTABLE_ENTRY_READ_KEY(ht, entry, traceback);
947    raw_free(traceback);
948    return 0;
949}
950
951
952/* reentrant flag must be set to call this function and GIL must be held */
953static void
954tracemalloc_clear_traces(void)
955{
956#ifdef WITH_THREAD
957    /* The GIL protects variables againt concurrent access */
958    assert(PyGILState_Check());
959#endif
960
961    TABLES_LOCK();
962    _Py_hashtable_clear(tracemalloc_traces);
963    tracemalloc_traced_memory = 0;
964    tracemalloc_peak_traced_memory = 0;
965    TABLES_UNLOCK();
966
967    _Py_hashtable_foreach(tracemalloc_tracebacks, traceback_free_traceback, NULL);
968    _Py_hashtable_clear(tracemalloc_tracebacks);
969
970    _Py_hashtable_foreach(tracemalloc_filenames, tracemalloc_clear_filename, NULL);
971    _Py_hashtable_clear(tracemalloc_filenames);
972}
973
974
975static int
976tracemalloc_init(void)
977{
978    if (tracemalloc_config.initialized == TRACEMALLOC_FINALIZED) {
979        PyErr_SetString(PyExc_RuntimeError,
980                        "the tracemalloc module has been unloaded");
981        return -1;
982    }
983
984    if (tracemalloc_config.initialized == TRACEMALLOC_INITIALIZED)
985        return 0;
986
987    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);
988
989#ifdef REENTRANT_THREADLOCAL
990    tracemalloc_reentrant_key = PyThread_create_key();
991    if (tracemalloc_reentrant_key == -1) {
992#ifdef MS_WINDOWS
993        PyErr_SetFromWindowsErr(0);
994#else
995        PyErr_SetFromErrno(PyExc_OSError);
996#endif
997        return -1;
998    }
999#endif
1000
1001#if defined(WITH_THREAD) && defined(TRACE_RAW_MALLOC)
1002    if (tables_lock == NULL) {
1003        tables_lock = PyThread_allocate_lock();
1004        if (tables_lock == NULL) {
1005            PyErr_SetString(PyExc_RuntimeError, "cannot allocate lock");
1006            return -1;
1007        }
1008    }
1009#endif
1010
1011    tracemalloc_filenames = hashtable_new(sizeof(PyObject *), 0,
1012                                          hashtable_hash_pyobject,
1013                                          hashtable_compare_unicode);
1014
1015    tracemalloc_tracebacks = hashtable_new(sizeof(traceback_t *), 0,
1016                                           hashtable_hash_traceback,
1017                                           hashtable_compare_traceback);
1018
1019    if (tracemalloc_config.use_domain) {
1020        tracemalloc_traces = hashtable_new(sizeof(pointer_t),
1021                                           sizeof(trace_t),
1022                                           hashtable_hash_pointer_t,
1023                                           hashtable_compare_pointer_t);
1024    }
1025    else {
1026        tracemalloc_traces = hashtable_new(sizeof(uintptr_t),
1027                                           sizeof(trace_t),
1028                                           _Py_hashtable_hash_ptr,
1029                                           _Py_hashtable_compare_direct);
1030    }
1031
1032    if (tracemalloc_filenames == NULL || tracemalloc_tracebacks == NULL
1033       || tracemalloc_traces == NULL) {
1034        PyErr_NoMemory();
1035        return -1;
1036    }
1037
1038    unknown_filename = PyUnicode_FromString("<unknown>");
1039    if (unknown_filename == NULL)
1040        return -1;
1041    PyUnicode_InternInPlace(&unknown_filename);
1042
1043    tracemalloc_empty_traceback.nframe = 1;
1044    /* borrowed reference */
1045    tracemalloc_empty_traceback.frames[0].filename = unknown_filename;
1046    tracemalloc_empty_traceback.frames[0].lineno = 0;
1047    tracemalloc_empty_traceback.hash = traceback_hash(&tracemalloc_empty_traceback);
1048
1049    tracemalloc_config.initialized = TRACEMALLOC_INITIALIZED;
1050    return 0;
1051}
1052
1053
1054static void
1055tracemalloc_deinit(void)
1056{
1057    if (tracemalloc_config.initialized != TRACEMALLOC_INITIALIZED)
1058        return;
1059    tracemalloc_config.initialized = TRACEMALLOC_FINALIZED;
1060
1061    tracemalloc_stop();
1062
1063    /* destroy hash tables */
1064    _Py_hashtable_destroy(tracemalloc_tracebacks);
1065    _Py_hashtable_destroy(tracemalloc_filenames);
1066    _Py_hashtable_destroy(tracemalloc_traces);
1067
1068#if defined(WITH_THREAD) && defined(TRACE_RAW_MALLOC)
1069    if (tables_lock != NULL) {
1070        PyThread_free_lock(tables_lock);
1071        tables_lock = NULL;
1072    }
1073#endif
1074
1075#ifdef REENTRANT_THREADLOCAL
1076    PyThread_delete_key(tracemalloc_reentrant_key);
1077    tracemalloc_reentrant_key = -1;
1078#endif
1079
1080    Py_XDECREF(unknown_filename);
1081}
1082
1083
1084static int
1085tracemalloc_start(int max_nframe)
1086{
1087    PyMemAllocatorEx alloc;
1088    size_t size;
1089
1090    if (tracemalloc_init() < 0)
1091        return -1;
1092
1093    if (tracemalloc_config.tracing) {
1094        /* hook already installed: do nothing */
1095        return 0;
1096    }
1097
1098    assert(1 <= max_nframe && max_nframe <= MAX_NFRAME);
1099    tracemalloc_config.max_nframe = max_nframe;
1100
1101    /* allocate a buffer to store a new traceback */
1102    size = TRACEBACK_SIZE(max_nframe);
1103    assert(tracemalloc_traceback == NULL);
1104    tracemalloc_traceback = raw_malloc(size);
1105    if (tracemalloc_traceback == NULL) {
1106        PyErr_NoMemory();
1107        return -1;
1108    }
1109
1110#ifdef TRACE_RAW_MALLOC
1111    alloc.malloc = tracemalloc_raw_malloc;
1112    alloc.calloc = tracemalloc_raw_calloc;
1113    alloc.realloc = tracemalloc_raw_realloc;
1114    alloc.free = tracemalloc_free;
1115
1116    alloc.ctx = &allocators.raw;
1117    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);
1118    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);
1119#endif
1120
1121    alloc.malloc = tracemalloc_malloc_gil;
1122    alloc.calloc = tracemalloc_calloc_gil;
1123    alloc.realloc = tracemalloc_realloc_gil;
1124    alloc.free = tracemalloc_free;
1125
1126    alloc.ctx = &allocators.mem;
1127    PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &allocators.mem);
1128    PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);
1129
1130    alloc.ctx = &allocators.obj;
1131    PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &allocators.obj);
1132    PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
1133
1134    /* everything is ready: start tracing Python memory allocations */
1135    tracemalloc_config.tracing = 1;
1136
1137    return 0;
1138}
1139
1140
1141static void
1142tracemalloc_stop(void)
1143{
1144    if (!tracemalloc_config.tracing)
1145        return;
1146
1147    /* stop tracing Python memory allocations */
1148    tracemalloc_config.tracing = 0;
1149
1150    /* unregister the hook on memory allocators */
1151#ifdef TRACE_RAW_MALLOC
1152    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);
1153#endif
1154    PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &allocators.mem);
1155    PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &allocators.obj);
1156
1157    tracemalloc_clear_traces();
1158
1159    /* release memory */
1160    raw_free(tracemalloc_traceback);
1161    tracemalloc_traceback = NULL;
1162}
1163
1164PyDoc_STRVAR(tracemalloc_is_tracing_doc,
1165    "is_tracing()->bool\n"
1166    "\n"
1167    "True if the tracemalloc module is tracing Python memory allocations,\n"
1168    "False otherwise.");
1169
1170
1171static PyObject*
1172py_tracemalloc_is_tracing(PyObject *self)
1173{
1174    return PyBool_FromLong(tracemalloc_config.tracing);
1175}
1176
1177PyDoc_STRVAR(tracemalloc_clear_traces_doc,
1178    "clear_traces()\n"
1179    "\n"
1180    "Clear traces of memory blocks allocated by Python.");
1181
1182
1183static PyObject*
1184py_tracemalloc_clear_traces(PyObject *self)
1185{
1186    if (!tracemalloc_config.tracing)
1187        Py_RETURN_NONE;
1188
1189    set_reentrant(1);
1190    tracemalloc_clear_traces();
1191    set_reentrant(0);
1192
1193    Py_RETURN_NONE;
1194}
1195
1196
1197static PyObject*
1198frame_to_pyobject(frame_t *frame)
1199{
1200    PyObject *frame_obj, *lineno_obj;
1201
1202    frame_obj = PyTuple_New(2);
1203    if (frame_obj == NULL)
1204        return NULL;
1205
1206    Py_INCREF(frame->filename);
1207    PyTuple_SET_ITEM(frame_obj, 0, frame->filename);
1208
1209    lineno_obj = PyLong_FromUnsignedLong(frame->lineno);
1210    if (lineno_obj == NULL) {
1211        Py_DECREF(frame_obj);
1212        return NULL;
1213    }
1214    PyTuple_SET_ITEM(frame_obj, 1, lineno_obj);
1215
1216    return frame_obj;
1217}
1218
1219
1220static PyObject*
1221traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
1222{
1223    int i;
1224    PyObject *frames, *frame;
1225
1226    if (intern_table != NULL) {
1227        if (_Py_HASHTABLE_GET(intern_table, traceback, frames)) {
1228            Py_INCREF(frames);
1229            return frames;
1230        }
1231    }
1232
1233    frames = PyTuple_New(traceback->nframe);
1234    if (frames == NULL)
1235        return NULL;
1236
1237    for (i=0; i < traceback->nframe; i++) {
1238        frame = frame_to_pyobject(&traceback->frames[i]);
1239        if (frame == NULL) {
1240            Py_DECREF(frames);
1241            return NULL;
1242        }
1243        PyTuple_SET_ITEM(frames, i, frame);
1244    }
1245
1246    if (intern_table != NULL) {
1247        if (_Py_HASHTABLE_SET(intern_table, traceback, frames) < 0) {
1248            Py_DECREF(frames);
1249            PyErr_NoMemory();
1250            return NULL;
1251        }
1252        /* intern_table keeps a new reference to frames */
1253        Py_INCREF(frames);
1254    }
1255    return frames;
1256}
1257
1258
1259static PyObject*
1260trace_to_pyobject(_PyTraceMalloc_domain_t domain, trace_t *trace,
1261                  _Py_hashtable_t *intern_tracebacks)
1262{
1263    PyObject *trace_obj = NULL;
1264    PyObject *obj;
1265
1266    trace_obj = PyTuple_New(3);
1267    if (trace_obj == NULL)
1268        return NULL;
1269
1270    obj = PyLong_FromSize_t(domain);
1271    if (obj == NULL) {
1272        Py_DECREF(trace_obj);
1273        return NULL;
1274    }
1275    PyTuple_SET_ITEM(trace_obj, 0, obj);
1276
1277    obj = PyLong_FromSize_t(trace->size);
1278    if (obj == NULL) {
1279        Py_DECREF(trace_obj);
1280        return NULL;
1281    }
1282    PyTuple_SET_ITEM(trace_obj, 1, obj);
1283
1284    obj = traceback_to_pyobject(trace->traceback, intern_tracebacks);
1285    if (obj == NULL) {
1286        Py_DECREF(trace_obj);
1287        return NULL;
1288    }
1289    PyTuple_SET_ITEM(trace_obj, 2, obj);
1290
1291    return trace_obj;
1292}
1293
1294
1295typedef struct {
1296    _Py_hashtable_t *traces;
1297    _Py_hashtable_t *tracebacks;
1298    PyObject *list;
1299} get_traces_t;
1300
1301static int
1302tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entry,
1303                            void *user_data)
1304{
1305    get_traces_t *get_traces = user_data;
1306    _PyTraceMalloc_domain_t domain;
1307    trace_t trace;
1308    PyObject *tracemalloc_obj;
1309    int res;
1310
1311    if (tracemalloc_config.use_domain) {
1312        pointer_t key;
1313        _Py_HASHTABLE_ENTRY_READ_KEY(traces, entry, key);
1314        domain = key.domain;
1315    }
1316    else {
1317        domain = DEFAULT_DOMAIN;
1318    }
1319    _Py_HASHTABLE_ENTRY_READ_DATA(traces, entry, trace);
1320
1321    tracemalloc_obj = trace_to_pyobject(domain, &trace, get_traces->tracebacks);
1322    if (tracemalloc_obj == NULL)
1323        return 1;
1324
1325    res = PyList_Append(get_traces->list, tracemalloc_obj);
1326    Py_DECREF(tracemalloc_obj);
1327    if (res < 0)
1328        return 1;
1329
1330    return 0;
1331}
1332
1333
1334static int
1335tracemalloc_pyobject_decref_cb(_Py_hashtable_t *tracebacks,
1336                               _Py_hashtable_entry_t *entry,
1337                               void *user_data)
1338{
1339    PyObject *obj;
1340    _Py_HASHTABLE_ENTRY_READ_DATA(tracebacks, entry, obj);
1341    Py_DECREF(obj);
1342    return 0;
1343}
1344
1345
1346PyDoc_STRVAR(tracemalloc_get_traces_doc,
1347    "_get_traces() -> list\n"
1348    "\n"
1349    "Get traces of all memory blocks allocated by Python.\n"
1350    "Return a list of (size: int, traceback: tuple) tuples.\n"
1351    "traceback is a tuple of (filename: str, lineno: int) tuples.\n"
1352    "\n"
1353    "Return an empty list if the tracemalloc module is disabled.");
1354
1355static PyObject*
1356py_tracemalloc_get_traces(PyObject *self, PyObject *obj)
1357{
1358    get_traces_t get_traces;
1359    int err;
1360
1361    get_traces.traces = NULL;
1362    get_traces.tracebacks = NULL;
1363    get_traces.list = PyList_New(0);
1364    if (get_traces.list == NULL)
1365        goto error;
1366
1367    if (!tracemalloc_config.tracing)
1368        return get_traces.list;
1369
1370    /* the traceback hash table is used temporarily to intern traceback tuple
1371       of (filename, lineno) tuples */
1372    get_traces.tracebacks = hashtable_new(sizeof(traceback_t *),
1373                                          sizeof(PyObject *),
1374                                          _Py_hashtable_hash_ptr,
1375                                          _Py_hashtable_compare_direct);
1376    if (get_traces.tracebacks == NULL) {
1377        PyErr_NoMemory();
1378        goto error;
1379    }
1380
1381    TABLES_LOCK();
1382    get_traces.traces = _Py_hashtable_copy(tracemalloc_traces);
1383    TABLES_UNLOCK();
1384
1385    if (get_traces.traces == NULL) {
1386        PyErr_NoMemory();
1387        goto error;
1388    }
1389
1390    set_reentrant(1);
1391    err = _Py_hashtable_foreach(get_traces.traces,
1392                                tracemalloc_get_traces_fill, &get_traces);
1393    set_reentrant(0);
1394    if (err)
1395        goto error;
1396
1397    goto finally;
1398
1399error:
1400    Py_CLEAR(get_traces.list);
1401
1402finally:
1403    if (get_traces.tracebacks != NULL) {
1404        _Py_hashtable_foreach(get_traces.tracebacks,
1405                              tracemalloc_pyobject_decref_cb, NULL);
1406        _Py_hashtable_destroy(get_traces.tracebacks);
1407    }
1408    if (get_traces.traces != NULL) {
1409        _Py_hashtable_destroy(get_traces.traces);
1410    }
1411
1412    return get_traces.list;
1413}
1414
1415
1416static traceback_t*
1417tracemalloc_get_traceback(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
1418{
1419    trace_t trace;
1420    int found;
1421
1422    if (!tracemalloc_config.tracing)
1423        return NULL;
1424
1425    TABLES_LOCK();
1426    if (tracemalloc_config.use_domain) {
1427        pointer_t key = {ptr, domain};
1428        found = _Py_HASHTABLE_GET(tracemalloc_traces, key, trace);
1429    }
1430    else {
1431        found = _Py_HASHTABLE_GET(tracemalloc_traces, ptr, trace);
1432    }
1433    TABLES_UNLOCK();
1434
1435    if (!found)
1436        return NULL;
1437
1438    return trace.traceback;
1439}
1440
1441
1442PyDoc_STRVAR(tracemalloc_get_object_traceback_doc,
1443    "_get_object_traceback(obj)\n"
1444    "\n"
1445    "Get the traceback where the Python object obj was allocated.\n"
1446    "Return a tuple of (filename: str, lineno: int) tuples.\n"
1447    "\n"
1448    "Return None if the tracemalloc module is disabled or did not\n"
1449    "trace the allocation of the object.");
1450
1451static PyObject*
1452py_tracemalloc_get_object_traceback(PyObject *self, PyObject *obj)
1453{
1454    PyTypeObject *type;
1455    void *ptr;
1456    traceback_t *traceback;
1457
1458    type = Py_TYPE(obj);
1459    if (PyType_IS_GC(type))
1460        ptr = (void *)((char *)obj - sizeof(PyGC_Head));
1461    else
1462        ptr = (void *)obj;
1463
1464    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, (uintptr_t)ptr);
1465    if (traceback == NULL)
1466        Py_RETURN_NONE;
1467
1468    return traceback_to_pyobject(traceback, NULL);
1469}
1470
1471
1472#define PUTS(fd, str) _Py_write_noraise(fd, str, (int)strlen(str))
1473
1474static void
1475_PyMem_DumpFrame(int fd, frame_t * frame)
1476{
1477    PUTS(fd, "  File \"");
1478    _Py_DumpASCII(fd, frame->filename);
1479    PUTS(fd, "\", line ");
1480    _Py_DumpDecimal(fd, frame->lineno);
1481    PUTS(fd, "\n");
1482}
1483
1484/* Dump the traceback where a memory block was allocated into file descriptor
1485   fd. The function may block on TABLES_LOCK() but it is unlikely. */
1486void
1487_PyMem_DumpTraceback(int fd, const void *ptr)
1488{
1489    traceback_t *traceback;
1490    int i;
1491
1492    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, (uintptr_t)ptr);
1493    if (traceback == NULL)
1494        return;
1495
1496    PUTS(fd, "Memory block allocated at (most recent call first):\n");
1497    for (i=0; i < traceback->nframe; i++) {
1498        _PyMem_DumpFrame(fd, &traceback->frames[i]);
1499    }
1500    PUTS(fd, "\n");
1501}
1502
1503#undef PUTS
1504
1505
1506PyDoc_STRVAR(tracemalloc_start_doc,
1507    "start(nframe: int=1)\n"
1508    "\n"
1509    "Start tracing Python memory allocations. Set also the maximum number \n"
1510    "of frames stored in the traceback of a trace to nframe.");
1511
1512static PyObject*
1513py_tracemalloc_start(PyObject *self, PyObject *args)
1514{
1515    Py_ssize_t nframe = 1;
1516    int nframe_int;
1517
1518    if (!PyArg_ParseTuple(args, "|n:start", &nframe))
1519        return NULL;
1520
1521    if (nframe < 1 || nframe > MAX_NFRAME) {
1522        PyErr_Format(PyExc_ValueError,
1523                     "the number of frames must be in range [1; %i]",
1524                     MAX_NFRAME);
1525        return NULL;
1526    }
1527    nframe_int = Py_SAFE_DOWNCAST(nframe, Py_ssize_t, int);
1528
1529    if (tracemalloc_start(nframe_int) < 0)
1530        return NULL;
1531
1532    Py_RETURN_NONE;
1533}
1534
1535PyDoc_STRVAR(tracemalloc_stop_doc,
1536    "stop()\n"
1537    "\n"
1538    "Stop tracing Python memory allocations and clear traces\n"
1539    "of memory blocks allocated by Python.");
1540
1541
1542static PyObject*
1543py_tracemalloc_stop(PyObject *self)
1544{
1545    tracemalloc_stop();
1546    Py_RETURN_NONE;
1547}
1548
1549
1550PyDoc_STRVAR(tracemalloc_get_traceback_limit_doc,
1551    "get_traceback_limit() -> int\n"
1552    "\n"
1553    "Get the maximum number of frames stored in the traceback\n"
1554    "of a trace.\n"
1555    "\n"
1556    "By default, a trace of an allocated memory block only stores\n"
1557    "the most recent frame: the limit is 1.");
1558
1559static PyObject*
1560py_tracemalloc_get_traceback_limit(PyObject *self)
1561{
1562    return PyLong_FromLong(tracemalloc_config.max_nframe);
1563}
1564
1565
1566PyDoc_STRVAR(tracemalloc_get_tracemalloc_memory_doc,
1567    "get_tracemalloc_memory() -> int\n"
1568    "\n"
1569    "Get the memory usage in bytes of the tracemalloc module\n"
1570    "used internally to trace memory allocations.");
1571
1572static PyObject*
1573tracemalloc_get_tracemalloc_memory(PyObject *self)
1574{
1575    size_t size;
1576    PyObject *size_obj;
1577
1578    size = _Py_hashtable_size(tracemalloc_tracebacks);
1579    size += _Py_hashtable_size(tracemalloc_filenames);
1580
1581    TABLES_LOCK();
1582    size += _Py_hashtable_size(tracemalloc_traces);
1583    TABLES_UNLOCK();
1584
1585    size_obj = PyLong_FromSize_t(size);
1586    return Py_BuildValue("N", size_obj);
1587}
1588
1589
1590PyDoc_STRVAR(tracemalloc_get_traced_memory_doc,
1591    "get_traced_memory() -> (int, int)\n"
1592    "\n"
1593    "Get the current size and peak size of memory blocks traced\n"
1594    "by the tracemalloc module as a tuple: (current: int, peak: int).");
1595
1596static PyObject*
1597tracemalloc_get_traced_memory(PyObject *self)
1598{
1599    Py_ssize_t size, peak_size;
1600    PyObject *size_obj, *peak_size_obj;
1601
1602    if (!tracemalloc_config.tracing)
1603        return Py_BuildValue("ii", 0, 0);
1604
1605    TABLES_LOCK();
1606    size = tracemalloc_traced_memory;
1607    peak_size = tracemalloc_peak_traced_memory;
1608    TABLES_UNLOCK();
1609
1610    size_obj = PyLong_FromSize_t(size);
1611    peak_size_obj = PyLong_FromSize_t(peak_size);
1612    return Py_BuildValue("NN", size_obj, peak_size_obj);
1613}
1614
1615
1616static PyMethodDef module_methods[] = {
1617    {"is_tracing", (PyCFunction)py_tracemalloc_is_tracing,
1618     METH_NOARGS, tracemalloc_is_tracing_doc},
1619    {"clear_traces", (PyCFunction)py_tracemalloc_clear_traces,
1620     METH_NOARGS, tracemalloc_clear_traces_doc},
1621    {"_get_traces", (PyCFunction)py_tracemalloc_get_traces,
1622     METH_NOARGS, tracemalloc_get_traces_doc},
1623    {"_get_object_traceback", (PyCFunction)py_tracemalloc_get_object_traceback,
1624     METH_O, tracemalloc_get_object_traceback_doc},
1625    {"start", (PyCFunction)py_tracemalloc_start,
1626      METH_VARARGS, tracemalloc_start_doc},
1627    {"stop", (PyCFunction)py_tracemalloc_stop,
1628      METH_NOARGS, tracemalloc_stop_doc},
1629    {"get_traceback_limit", (PyCFunction)py_tracemalloc_get_traceback_limit,
1630     METH_NOARGS, tracemalloc_get_traceback_limit_doc},
1631    {"get_tracemalloc_memory", (PyCFunction)tracemalloc_get_tracemalloc_memory,
1632     METH_NOARGS, tracemalloc_get_tracemalloc_memory_doc},
1633    {"get_traced_memory", (PyCFunction)tracemalloc_get_traced_memory,
1634     METH_NOARGS, tracemalloc_get_traced_memory_doc},
1635
1636    /* sentinel */
1637    {NULL, NULL}
1638};
1639
1640PyDoc_STRVAR(module_doc,
1641"Debug module to trace memory blocks allocated by Python.");
1642
1643static struct PyModuleDef module_def = {
1644    PyModuleDef_HEAD_INIT,
1645    "_tracemalloc",
1646    module_doc,
1647    0, /* non-negative size to be able to unload the module */
1648    module_methods,
1649    NULL,
1650};
1651
1652PyMODINIT_FUNC
1653PyInit__tracemalloc(void)
1654{
1655    PyObject *m;
1656    m = PyModule_Create(&module_def);
1657    if (m == NULL)
1658        return NULL;
1659
1660    if (tracemalloc_init() < 0)
1661        return NULL;
1662
1663    return m;
1664}
1665
1666
1667static int
1668parse_sys_xoptions(PyObject *value)
1669{
1670    PyObject *valuelong;
1671    long nframe;
1672
1673    if (value == Py_True)
1674        return 1;
1675
1676    assert(PyUnicode_Check(value));
1677    if (PyUnicode_GetLength(value) == 0)
1678        return -1;
1679
1680    valuelong = PyLong_FromUnicodeObject(value, 10);
1681    if (valuelong == NULL)
1682        return -1;
1683
1684    nframe = PyLong_AsLong(valuelong);
1685    Py_DECREF(valuelong);
1686    if (nframe == -1 && PyErr_Occurred())
1687        return -1;
1688
1689    if (nframe < 1 || nframe > MAX_NFRAME)
1690        return -1;
1691
1692    return Py_SAFE_DOWNCAST(nframe, long, int);
1693}
1694
1695
1696int
1697_PyTraceMalloc_Init(void)
1698{
1699    char *p;
1700    int nframe;
1701
1702#ifdef WITH_THREAD
1703    assert(PyGILState_Check());
1704#endif
1705
1706    if ((p = Py_GETENV("PYTHONTRACEMALLOC")) && *p != '\0') {
1707        char *endptr = p;
1708        long value;
1709
1710        errno = 0;
1711        value = strtol(p, &endptr, 10);
1712        if (*endptr != '\0'
1713            || value < 1
1714            || value > MAX_NFRAME
1715            || errno == ERANGE)
1716        {
1717            Py_FatalError("PYTHONTRACEMALLOC: invalid number of frames");
1718            return -1;
1719        }
1720
1721        nframe = (int)value;
1722    }
1723    else {
1724        PyObject *xoptions, *key, *value;
1725
1726        xoptions = PySys_GetXOptions();
1727        if (xoptions == NULL)
1728            return -1;
1729
1730        key = PyUnicode_FromString("tracemalloc");
1731        if (key == NULL)
1732            return -1;
1733
1734        value = PyDict_GetItemWithError(xoptions, key);
1735        Py_DECREF(key);
1736        if (value == NULL) {
1737            if (PyErr_Occurred())
1738                return -1;
1739
1740            /* -X tracemalloc is not used */
1741            return 0;
1742        }
1743
1744        nframe = parse_sys_xoptions(value);
1745        Py_DECREF(value);
1746        if (nframe < 0) {
1747            Py_FatalError("-X tracemalloc=NFRAME: invalid number of frames");
1748        }
1749    }
1750
1751    return tracemalloc_start(nframe);
1752}
1753
1754
1755void
1756_PyTraceMalloc_Fini(void)
1757{
1758#ifdef WITH_THREAD
1759    assert(PyGILState_Check());
1760#endif
1761    tracemalloc_deinit();
1762}
1763
1764int
1765_PyTraceMalloc_Track(_PyTraceMalloc_domain_t domain, uintptr_t ptr,
1766                     size_t size)
1767{
1768    int res;
1769#ifdef WITH_THREAD
1770    PyGILState_STATE gil_state;
1771#endif
1772
1773    if (!tracemalloc_config.tracing) {
1774        /* tracemalloc is not tracing: do nothing */
1775        return -2;
1776    }
1777
1778#ifdef WITH_THREAD
1779    gil_state = PyGILState_Ensure();
1780#endif
1781
1782    TABLES_LOCK();
1783    res = tracemalloc_add_trace(domain, ptr, size);
1784    TABLES_UNLOCK();
1785
1786#ifdef WITH_THREAD
1787    PyGILState_Release(gil_state);
1788#endif
1789    return res;
1790}
1791
1792
1793int
1794_PyTraceMalloc_Untrack(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
1795{
1796    if (!tracemalloc_config.tracing) {
1797        /* tracemalloc is not tracing: do nothing */
1798        return -2;
1799    }
1800
1801    TABLES_LOCK();
1802    tracemalloc_remove_trace(domain, ptr);
1803    TABLES_UNLOCK();
1804
1805    return 0;
1806}
1807
1808
1809PyObject*
1810_PyTraceMalloc_GetTraceback(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
1811{
1812    traceback_t *traceback;
1813
1814    traceback = tracemalloc_get_traceback(domain, ptr);
1815    if (traceback == NULL)
1816        Py_RETURN_NONE;
1817
1818    return traceback_to_pyobject(traceback, NULL);
1819}
1820