diff --git a/accel.c b/accel.c index db6341c8..07527628 100644 --- a/accel.c +++ b/accel.c @@ -369,18 +369,28 @@ inline int IMIN(int a, int b) { return((a) < (b) ? a : b); } static PyObject *create_numpy_array(PyObject *py_memview, char *data_format, int data_type, PyObject *py_objs); +// Returns a newly allocated UTF-8 copy of `unicode`; the caller owns it +// and must free() it. char *_PyUnicode_AsUTF8(PyObject *unicode) { PyObject *bytes = PyUnicode_AsEncodedString(unicode, "utf-8", "strict"); if (!bytes) return NULL; char *str = NULL; Py_ssize_t str_l = 0; - if (PyBytes_AsStringAndSize(bytes, &str, &str_l) < 0) { - return NULL; - } + char *out = NULL; - char *out = calloc(str_l + 1, 1); + if (PyBytes_AsStringAndSize(bytes, &str, &str_l) < 0) goto exit; + + out = calloc(str_l + 1, 1); + if (!out) { + PyErr_NoMemory(); + goto exit; + } memcpy(out, str, str_l); + +exit: + Py_DECREF(bytes); + return out; } @@ -568,6 +578,8 @@ typedef struct { PyObject *create_numpy_array_kwargs_vector[8]; PyObject *struct_unpack_args; PyObject *bson_decode_args; + PyObject *structsequence_fields_registry; + PyObject *structsequence_fields_release; } PyObjects; static PyObjects PyObj = {0}; @@ -924,14 +936,71 @@ int ensure_bson() { } +// Name of the capsule that owns a struct sequence type's field name storage. +// The capsule is held by PyObj.structsequence_fields_registry, keyed by a weak +// reference to the type, and is dropped when the type is collected. +#define STRUCTSEQUENCE_FIELDS_CAPSULE "singlestoredb.Row.fields" + + +// Frees a NULL name terminated array of struct sequence fields and the +// names in it. +static void free_structsequence_fields(PyStructSequence_Field *fields) { + if (!fields) return; + for (PyStructSequence_Field *field = fields; field->name; field++) { + free((void*)field->name); + } + free(fields); +} + + +static void structsequence_fields_capsule_destructor(PyObject *py_capsule) { + PyStructSequence_Field *fields = (PyStructSequence_Field*) + PyCapsule_GetPointer(py_capsule, STRUCTSEQUENCE_FIELDS_CAPSULE); + if (!fields) { + PyErr_Clear(); + return; + } + free_structsequence_fields(fields); +} + + +// Weak reference callback for a struct sequence type. Dropping the registry +// entry drops the last reference to the capsule, whose destructor frees the +// names. The weak reference itself is the key: by the time this runs the +// referent is already gone, so the type cannot be used to find the entry. +static PyObject *structsequence_fields_release(PyObject *self, PyObject *py_weakref) { + (void)self; + if (PyObj.structsequence_fields_registry) { + if (PyDict_DelItem(PyObj.structsequence_fields_registry, py_weakref)) { + // A callback must not raise. + PyErr_Clear(); + } + } + Py_INCREF(Py_None); + return Py_None; +} + + +static PyMethodDef structsequence_fields_release_def = { + "structsequence_fields_release", + (PyCFunction)structsequence_fields_release, + METH_O, + NULL +}; + + static void State_clear_fields(StateObject *self) { if (!self) return; DESTROY(self->offsets); DESTROY(self->scales); DESTROY(self->flags); DESTROY(self->type_codes); - DESTROY(self->encodings); - DESTROY(self->structsequence_desc.fields); + if (self->encodings) { + for (unsigned long i = 0; i < self->n_cols; i++) { + DESTROY(self->encodings[i]); + } + DESTROY(self->encodings); + } DESTROY(self->encoding_errors); if (self->py_converters) { for (unsigned long i = 0; i < self->n_cols; i++) { @@ -958,6 +1027,11 @@ static void State_clear_fields(StateObject *self) { DESTROY(self->py_invalid_values); } Py_CLEAR(self->structsequence); + // Only reached if the type was never built, or building it failed before + // the capsule took ownership. Once the capsule holds the fields, this is + // NULL and the names outlive us along with the type. + free_structsequence_fields(self->structsequence_desc.fields); + self->structsequence_desc.fields = NULL; Py_CLEAR(self->py_namedtuple); Py_CLEAR(self->py_namedtuple_args); Py_CLEAR(self->py_names_list); @@ -1126,8 +1200,14 @@ static int State_init(StateObject *self, PyObject *args, PyObject *kwds) { self->py_encodings[i] = (py_encoding == Py_None) ? NULL : py_encoding; Py_XINCREF(self->py_encodings[i]); - self->encodings[i] = (!py_encoding || py_encoding == Py_None) ? - NULL : _PyUnicode_AsUTF8(py_encoding); + // NULL is the binary-column sentinel, so an allocation failure here + // can not be left in place; it has to go to the error path. + if (py_encoding == Py_None) { + self->encodings[i] = NULL; + } else { + self->encodings[i] = _PyUnicode_AsUTF8(py_encoding); + if (!self->encodings[i]) goto error; + } self->py_invalid_values[i] = (!py_invalid_value || py_invalid_value == Py_None) ? NULL : py_converter; @@ -1201,10 +1281,42 @@ static int State_init(StateObject *self, PyObject *args, PyObject *kwds) { if (!self->structsequence_desc.fields) goto error; for (unsigned long i = 0; i < self->n_cols; i++) { self->structsequence_desc.fields[i].name = _PyUnicode_AsUTF8(self->py_names[i]); + if (!self->structsequence_desc.fields[i].name) goto error; self->structsequence_desc.fields[i].doc = NULL; } self->structsequence = PyStructSequence_NewType(&self->structsequence_desc); if (!self->structsequence) goto error; + + // The type stores the field name pointers rather than copying the + // strings, and reads them again when a row is repr'd. Rows can + // outlive this State, so the storage is handed to a capsule owned + // by a module-private registry, keyed by a weak reference to the + // type. The names are freed when the type is collected. Rows are + // instances of a heap type and so keep it alive; nothing on the + // type refers to the capsule, so Python code cannot release it + // early. + PyObject *py_fields_capsule = PyCapsule_New( + self->structsequence_desc.fields, + STRUCTSEQUENCE_FIELDS_CAPSULE, + &structsequence_fields_capsule_destructor + ); + if (!py_fields_capsule) goto error; + self->structsequence_desc.fields = NULL; + + PyObject *py_fields_weakref = PyWeakref_NewRef( + (PyObject*)self->structsequence, + PyObj.structsequence_fields_release + ); + if (!py_fields_weakref) { + Py_DECREF(py_fields_capsule); + goto error; + } + + rc = PyDict_SetItem(PyObj.structsequence_fields_registry, + py_fields_weakref, py_fields_capsule); + Py_DECREF(py_fields_weakref); + Py_DECREF(py_fields_capsule); + if (rc != 0) goto error; } // Fall through @@ -6075,6 +6187,15 @@ PyMODINIT_FUNC PyInit__singlestoredb_accel(void) { PyObj.bson_decode_args = PyTuple_New(1); if (!PyObj.bson_decode_args) goto error; + // Owns the field name storage of every live struct sequence type: + // weak reference to the type => capsule holding its names. + PyObj.structsequence_fields_registry = PyDict_New(); + if (!PyObj.structsequence_fields_registry) goto error; + + PyObj.structsequence_fields_release = PyCFunction_NewEx( + &structsequence_fields_release_def, NULL, NULL); + if (!PyObj.structsequence_fields_release) goto error; + return PyModule_Create(&_singlestoredb_accelmodule); error: diff --git a/singlestoredb/tests/test_accel_leaks.py b/singlestoredb/tests/test_accel_leaks.py new file mode 100644 index 00000000..dc48b522 --- /dev/null +++ b/singlestoredb/tests/test_accel_leaks.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python +# type: ignore +"""Test that the C accelerator does not leak memory per query.""" +import gc +import sys +import unittest + +import singlestoredb as s2 +from singlestoredb.mysql import connection as mysql_connection + +# The leak in issue #135 was one allocation per column per query, so the +# signal to look for is retention that grows with the width of the result. +# Measuring two widths and taking the difference is what makes this robust: +# anything a query costs that is flat in the column count drops out, and one +# such cost is unavoidable here. Under coverage.py's sys.monitoring backend +# every code object ever seen is retained forever, deliberately, keyed by +# id() (see `code_objects` in coverage/sysmon.py). collections.namedtuple +# compiles a fresh __new__ on each call and the accelerator builds one Row +# class per query, so a coverage run retains ~15 blocks per query on the +# namedtuples path however narrow the result is. That is the tracer's +# accounting, not our allocation, and CI runs under --cov. +NARROW_COLS = 10 +WIDE_COLS = 100 + +WARMUP = 50 +ITERATIONS = 200 + + +def query_for(n_cols): + return 'SELECT ' + ', '.join(f'{i} AS c{i}' for i in range(n_cols)) + + +# Per-column budget, in allocated blocks. A fixed accelerator measures zero; +# the leak this guards against cost one block per column, so anything above +# the noise floor of a fraction of a block is the bug coming back. +MAX_BLOCKS_PER_COLUMN = 0.05 + +# Per-query budget for the width-independent part, in allocated blocks. Room +# for caches that fill on the first few queries, plus the tracer overhead +# above, which is measured rather than assumed so the budget stays tight when +# nothing is tracing. +MAX_BLOCKS_PER_QUERY = 5.0 + +N_COLS = WIDE_COLS +WIDE_QUERY = query_for(WIDE_COLS) + + +def blocks_retained_per_namedtuple(): + """Return the blocks a tracer retains per collections.namedtuple() call. + + Zero when nothing is tracing. Non-zero under coverage, which the + accelerator then pays once per query on the namedtuples path. + """ + import collections + + fields = [f'c{i}' for i in range(WIDE_COLS)] + + def build(n): + for _ in range(n): + collections.namedtuple('Row', fields, rename=True) + + build(WARMUP) + gc.collect() + before = sys.getallocatedblocks() + + build(ITERATIONS) + gc.collect() + after = sys.getallocatedblocks() + + return max(0.0, (after - before) / ITERATIONS) + + +has_accel = mysql_connection._singlestoredb_accel is not None +# Read the parsed option rather than the environment variable: the option's +# validator already accepts true/yes/on, which int() would choke on. +pure_python = bool(s2.get_option('pure_python')) + + +@unittest.skipIf(not has_accel, 'C extension is not available') +@unittest.skipIf(pure_python, 'C extension is disabled') +class TestAccelLeaks(unittest.TestCase): + + def setUp(self): + self.conn = s2.connect() + if 'http' in self.conn.driver: + self.skipTest('HTTP interface does not use the C extension') + + def tearDown(self): + try: + self.conn.close() + except Exception: + pass + + def blocks_per_query(self, results_type, n_cols=WIDE_COLS): + """Return the allocated blocks retained per query of n_cols columns.""" + query = query_for(n_cols) + with s2.connect(results_type=results_type, pure_python=False) as conn: + with conn.cursor() as cur: + for _ in range(WARMUP): + cur.execute(query) + cur.fetchall() + + gc.collect() + before = sys.getallocatedblocks() + + for _ in range(ITERATIONS): + cur.execute(query) + cur.fetchall() + + gc.collect() + after = sys.getallocatedblocks() + + return (after - before) / ITERATIONS + + def test_no_leak_per_column(self): + """Retention must not grow with the width of the result. + + This is the shape of the issue #135 leak, and differencing two widths + cancels every per-query cost that is flat in the column count -- see + the note on the tracer overhead at the top of this module. + """ + for results_type in ('tuples', 'dicts', 'namedtuples', 'structsequences'): + with self.subTest(results_type=results_type): + narrow = self.blocks_per_query(results_type, NARROW_COLS) + wide = self.blocks_per_query(results_type, WIDE_COLS) + + per_column = (wide - narrow) / (WIDE_COLS - NARROW_COLS) + + assert per_column < MAX_BLOCKS_PER_COLUMN, \ + f'{results_type} leaks {per_column} blocks per column ' \ + f'({narrow} blocks/query at {NARROW_COLS} columns, ' \ + f'{wide} at {WIDE_COLS})' + + def test_no_leak_per_query(self): + """Retention must not grow per query either. + + The per-column check above cannot see a leak of something allocated + once per query, so budget that separately. The namedtuples path is + allowed the tracer's per-class overhead on top, measured here so the + budget stays tight when nothing is tracing. + """ + tracer_overhead = blocks_retained_per_namedtuple() + + for results_type in ('tuples', 'dicts', 'namedtuples', 'structsequences'): + with self.subTest(results_type=results_type): + budget = MAX_BLOCKS_PER_QUERY + if results_type == 'namedtuples': + budget += tracer_overhead + + leaked = self.blocks_per_query(results_type) + + assert leaked < budget, \ + f'{results_type} leaks {leaked} blocks per query ' \ + f'(budget {budget}, of which {tracer_overhead} is ' \ + f'tracer overhead)' + + def test_rows_outlive_the_result_state(self): + """Struct sequence rows must survive the state that created them. + + The type does not copy its field names, it stores the pointers and + reads them again in repr, so the names have to outlive every row + rather than the query that built them. + """ + with s2.connect( + results_type='structsequences', pure_python=False, + ) as conn: + with conn.cursor() as cur: + cur.execute(WIDE_QUERY) + rows = cur.fetchall() + + # Discard the state that built the rows, several times over. + for _ in range(10): + cur.execute('SELECT 1') + cur.fetchall() + + gc.collect() + + assert len(rows[0]) == N_COLS, len(rows[0]) + assert rows[0].c0 == 0, rows[0].c0 + assert getattr(rows[0], f'c{N_COLS - 1}') == N_COLS - 1 + assert f'c{N_COLS - 1}=' in repr(rows[0]), repr(rows[0]) + + def test_field_names_survive_stripping_the_type_dict(self): + """No class attribute may own the field name storage. + + The names are read again by repr, so a deletable attribute holding + the only reference would turn `delattr` into a use-after-free. + """ + with s2.connect( + results_type='structsequences', pure_python=False, + ) as conn: + with conn.cursor() as cur: + cur.execute(WIDE_QUERY) + rows = cur.fetchall() + + row_type = type(rows[0]) + + # Nothing in the type's dict may be the owner: every entry there is + # reachable, and most of them are deletable. + for name, value in vars(row_type).items(): + assert type(value).__name__ != 'PyCapsule', name + + # Held so the field names are still read out of the type after the + # loop below deletes the type's own __repr__ entry. + row_repr = row_type.__repr__ + + # CPython reads these three back out of the dict itself, so deleting + # them breaks a struct sequence whatever owns its names. + keep = ('n_fields', 'n_sequence_fields', 'n_unnamed_fields') + + for name in list(vars(row_type)): + if name in keep: + continue + try: + delattr(row_type, name) + except (AttributeError, TypeError): + pass + + gc.collect() + + # repr reads the names out of the C field table, not the type dict. + assert f'c{N_COLS - 1}=' in row_repr(rows[0]), row_repr(rows[0]) + + +if __name__ == '__main__': + import nose2 + nose2.main()