diff --git a/Modules/_sqlite/module.c b/Modules/_sqlite/module.c index 6bfb1b73f8..a5bc61cf59 100644 --- a/Modules/_sqlite/module.c +++ b/Modules/_sqlite/module.c @@ -361,6 +361,79 @@ do { \ } \ } while (0) +static const int MAX_ALLOC = 10 * 1024; +static struct { + int sz; + void *ptr; +} pool[MAX_ALLOC] = { 0, }; +static int n_allocated = 0; + +static void * +pysqlite_mem_malloc(int size) +{ + if (n_allocated >= MAX_ALLOC) { + return NULL; + } + void *ptr = PyMem_RawMalloc(size); + if (ptr) { + pool[n_allocated].sz = size; + pool[n_allocated].ptr = ptr; + n_allocated++; + } + return ptr; +} + +static int +pysqlite_mem_size(void *ptr) +{ + for (int i = 0; i < n_allocated; i++) { + if (pool[i].ptr == ptr) { + return pool[i].sz; + } + } + return 0; +} + +static void +pysqlite_mem_free(void *ptr) +{ + PyMem_RawFree(ptr); + for (int i = 0; i < n_allocated; i++) { + if (pool[i].ptr != ptr) { + continue; + } + n_allocated--; + if (n_allocated == 0 || n_allocated == i) { + pool[i].sz = 0; + pool[i].ptr = 0; + } + else { + pool[i].sz = pool[n_allocated].sz; + pool[i].ptr = pool[n_allocated].ptr; + pool[n_allocated].sz = 0; + pool[n_allocated].ptr = 0; + } + break; + } +} + +static void * +pysqlite_mem_realloc(void *ptr, int size) +{ + void *new = PyMem_RawRealloc(ptr, size); + if (new == NULL) { + return NULL; + } + for (int i = 0; i < n_allocated; i++) { + if (pool[i].ptr == ptr) { + pool[i].ptr = new; + pool[i].sz = size; + break; + } + } + return new; +} + PyMODINIT_FUNC PyInit__sqlite3(void) { PyObject *module; @@ -370,6 +443,26 @@ PyMODINIT_FUNC PyInit__sqlite3(void) return NULL; } + /* Set SQLite up to use Python's memory allocators to give access to the + * built-in memory debugger. */ + struct sqlite3_mem_methods mem_methods; + int rc = sqlite3_config(SQLITE_CONFIG_GETMALLOC, &mem_methods); + if (rc != SQLITE_OK) { + PyErr_SetString(PyExc_ImportError, sqlite3_errstr(rc)); + return NULL; + } + + mem_methods.xMalloc = pysqlite_mem_malloc; + mem_methods.xFree = pysqlite_mem_free; + mem_methods.xSize = pysqlite_mem_size; + mem_methods.xRealloc = pysqlite_mem_realloc; + + rc = sqlite3_config(SQLITE_CONFIG_MALLOC, &mem_methods); + if (rc != SQLITE_OK) { + PyErr_SetString(PyExc_ImportError, sqlite3_errstr(rc)); + return NULL; + } + module = PyModule_Create(&_sqlite3module); if (!module ||