diff --git a/src/arcadedb/_native/lib.py b/src/arcadedb/_native/lib.py index 58b261ba..4659df1a 100644 --- a/src/arcadedb/_native/lib.py +++ b/src/arcadedb/_native/lib.py @@ -327,16 +327,60 @@ def _thread_handle(): return t.handle +# One reply is read straight from its address: the text becomes a ``str`` in one copy (no ``bytes`` in between), and the library's +# memory is freed in the ``finally`` of the same call, once, whatever the decode does (a reply never outlives the call that took it, +# so no view of it can). ``PyUnicode_FromStringAndSize`` decodes UTF-8 as ``bytes.decode("utf-8")`` does, and raises the same +# UnicodeDecodeError; the function pointers are the ones of this module, ``ctypes.pythonapi`` itself is not changed. +_str_from_utf8 = ctypes.PYFUNCTYPE(ctypes.py_object, ctypes.c_void_p, ctypes.c_ssize_t)( + ("PyUnicode_FromStringAndSize", ctypes.pythonapi) +) +_strlen = ctypes.CFUNCTYPE(ctypes.c_size_t, ctypes.c_void_p)( + ( + "strlen", + ctypes.cdll.msvcrt if sys.platform == "win32" else ctypes.CDLL(None), + ) +) +_TAIL_KEY = b',"bin":' +_TAIL_WINDOW = 40 + + +def _take(lib, th, ptr, tail: bool): + """Read the reply at ``ptr`` as ``(text, floats)`` and free it: ``floats`` is None unless ``tail`` and the reply ends with + ``,"bin":N}`` (then N float32 follow the text's NUL, 4-byte aligned, and the text is returned without the ``,"bin":N`` part). + """ + if not ptr: + raise ArcadeDBError("The native library returned no reply") + try: + end = _strlen(ptr) + floats = None + if tail: + # only the last bytes can hold the marker: peek at them (at most 40 bytes), never at the whole reply + start = end - _TAIL_WINDOW if end > _TAIL_WINDOW else 0 + peek = ctypes.string_at(ptr + start, end - start) + if peek.endswith(b"}"): + i = peek.rfind(_TAIL_KEY) + if i >= 0: + n = int(peek[i + 7 : -1]) + import numpy as np + + off = (end + 1 + 3) & ~3 + floats = np.frombuffer( + ctypes.string_at(ptr + off, 4 * n), dtype=np.float32 + ) + # the text ends at the marker, closed by the "}" that the JSON object needs: written over the marker's comma + # in the reply's own memory (it is freed below), so the decode reads the JSON as it stands + ctypes.c_char.from_address(ptr + start + i).value = b"}" + end = start + i + 1 + return _str_from_utf8(ptr, end), floats + finally: + lib.arcadedb_free(th, ptr) + + def call(name: str, *args) -> dict: """Run one entry point; returns the decoded reply, raises ArcadeDBError-free (callers map errors).""" lib = _lib or _load() th = _thread_handle() - ptr = getattr(lib, name)(th, *args) - try: - text = ctypes.string_at(ptr).decode("utf-8") - finally: - lib.arcadedb_free(th, ptr) - return text + return _take(lib, th, getattr(lib, name)(th, *args), False)[0] def call_bin(name: str, *args): @@ -345,24 +389,7 @@ def call_bin(name: str, *args): """ lib = _lib or _load() th = _thread_handle() - ptr = getattr(lib, name)(th, *args) - try: - raw = ctypes.string_at(ptr) - floats = None - if raw.endswith(b"}"): - i = raw.rfind(b',"bin":', max(0, len(raw) - 40)) - if i >= 0: - n = int(raw[i + 7 : -1]) - import numpy as np - - off = (len(raw) + 1 + 3) & ~3 - floats = np.frombuffer( - ctypes.string_at(ptr + off, 4 * n), dtype=np.float32 - ) - raw = raw[:i] + b"}" - finally: - lib.arcadedb_free(th, ptr) - return raw.decode("utf-8"), floats + return _take(lib, th, getattr(lib, name)(th, *args), True) def alloc(size: int): @@ -383,22 +410,12 @@ def exec_query(handle, flags, language, command, params, count): lib = _lib or _load() t = _local.__dict__.get("thread") th = t.handle if t is not None else _thread_handle() - ptr = lib.arcadedb_n_exec(th, handle, flags, language, command, params, count) - try: - raw = ctypes.string_at(ptr) - floats = None - # a negative start counts from the end and stops at 0 for a short reply - i = raw.rfind(b',"bin":', -40) - if i >= 0 and raw.endswith(b"}"): - n = int(raw[i + 7 : -1]) - import numpy as np - - off = (len(raw) + 1 + 3) & ~3 - floats = np.frombuffer(ctypes.string_at(ptr + off, 4 * n), dtype=np.float32) - raw = raw[:i] + b"}" - finally: - lib.arcadedb_free(th, ptr) - return raw.decode("utf-8"), floats + return _take( + lib, + th, + lib.arcadedb_n_exec(th, handle, flags, language, command, params, count), + True, + ) def call_void(name: str, *args) -> None: diff --git a/tests/conftest.py b/tests/conftest.py index ec3147c7..2f0bfcce 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -65,6 +65,8 @@ collect_ignore.append("test_native_async_callbacks.py") # the short call path of a point read (codec.encode_params, lib.exec_query, first()), observed on the native library collect_ignore.append("test_native_point_read.py") + # how a reply is read out of the library (lib._take): the copies and the frees, observed with spies on ctypes.string_at and arcadedb_free + collect_ignore.append("test_native_reply_reader.py") # an RID built from an RID keeps the text, so the fast named-parameter path of codec.encode_params can encode it (#265) collect_ignore.append("test_native_rid_params.py") # to_json_list walks the rows with codec.json_native only when a reply restored a tagged value (#219), observed on the native library diff --git a/tests/test_native_reply_reader.py b/tests/test_native_reply_reader.py new file mode 100644 index 00000000..0b789097 --- /dev/null +++ b/tests/test_native_reply_reader.py @@ -0,0 +1,195 @@ +"""How the native build reads a reply out of the library (``lib._take``; ``call``, ``call_bin`` and ``exec_query`` all end there). + +A reply is a malloc'd UTF-8 text, optionally followed by a float32 tail, and it must be freed exactly once. It used to be copied as +``bytes`` (``ctypes.string_at``), cut for the tail marker, decoded to ``str``, and only then parsed. It is now decoded straight from its +address in one copy, and only the last 40 bytes (the marker) and the float tail are ever copied as ``bytes``. + +The oracle is a private copy of the former reader (``_former``): every case is read by both, from two identical replies, and the text, +the floats and the number of frees must agree. The copy counts are read from a spy on ``ctypes.string_at``: the former reader called it +once with the whole reply, the new one never does for a text and asks only for the marker window and the tail. The free counts are read +from a spy on ``arcadedb_free``: once, after the use, also when the decode or the tail fails. conftest.py collects this file on the +native backend only: it needs the native library. +""" + +import ctypes +import json + +import numpy as np +import pytest + +from arcadedb._native import lib +from arcadedb.exceptions import ArcadeDBError + +_string_at = ctypes.string_at + + +def _reply(raw: bytes, floats=None): + """A reply in the library's own memory, as the library hands one out: the text, its NUL, then (4-byte aligned) the floats.""" + tail = b"" + if floats is not None: + raw = raw[:-1] + b',"bin":%d}' % len(floats) + pad = -(len(raw) + 1) % 4 + tail = b"\0" * pad + np.asarray(floats, dtype=np.float32).tobytes() + data = raw + b"\0" + tail + ptr = lib.alloc(len(data)) + ctypes.memmove(ptr, data, len(data)) + return ptr, len(raw) + + +def _former(th, ptr): + """The reader this file replaces (``call_bin`` before the change), for the oracle.""" + try: + raw = _string_at(ptr) + floats = None + if raw.endswith(b"}"): + i = raw.rfind(b',"bin":', max(0, len(raw) - 40)) + if i >= 0: + n = int(raw[i + 7 : -1]) + off = (len(raw) + 1 + 3) & ~3 + floats = np.frombuffer(_string_at(ptr + off, 4 * n), dtype=np.float32) + raw = raw[:i] + b"}" + finally: + lib._lib.arcadedb_free(th, ptr) + return raw.decode("utf-8"), floats + + +@pytest.fixture +def frees(temp_db, monkeypatch): + """Every pointer given to ``arcadedb_free``, in order.""" + seen = [] + real = lib._lib.arcadedb_free + + def spy(th, ptr): + seen.append(ptr) + return real(th, ptr) + + monkeypatch.setattr(lib._lib, "arcadedb_free", spy) + return seen + + +@pytest.fixture +def string_at_sizes(monkeypatch): + """The size asked of each ``ctypes.string_at`` call (None: the whole NUL-terminated string).""" + seen = [] + + def spy(ptr, size=-1): + seen.append(None if size == -1 else size) + return _string_at(ptr, size) + + monkeypatch.setattr(ctypes, "string_at", spy) + return seen + + +LONG = json.dumps({"ok": True, "rows": [{"t": "x" * 5000}], "meta": []}).encode() + +CASES = { + "empty object": (b"{}", None), + "ascii": (b'{"ok":true,"rows":[{"a":1}],"drained":true}', None), + "non-ascii": ('{"ok":true,"t":"café 日本語 \U0001f600"}'.encode(), None), + "long text": (LONG, None), + "tail": ( + b'{"ok":true,"rows":[{"v":["B",0,4]}],"meta":[],"drained":true}', + [1, 2, 3, 4], + ), + "tail, short text": (b'{"v":1}', [0.5, 1.5]), + "tail, non-ascii text": ( + '{"ok":true,"rows":[{"t":"日本語 café \U0001f600","v":["B",0,3]}],"drained":true}'.encode(), + [1.5, -2.5, 3.25], + ), + "tail, long text": (LONG, np.arange(16, dtype=np.float32)), + "tail of no floats": (b'{"ok":true}', []), + **{ + f"tail, text of {k} digits": (b'{"a":' + b"1" * k + b"}", [7.0, 8.0]) + for k in range(1, 5) # the four paddings that align the floats to 4 bytes + }, + "marker deep in the text": ( + b'{"a":"' + b"x" * 100 + b'","b":1,"bin":2,"c":' + b"y" * 60 + b"}", + None, + ), +} + + +@pytest.mark.parametrize("name", list(CASES)) +@pytest.mark.parametrize("tail", [True, False]) +def test_every_case_reads_as_the_former_reader_does(temp_db, frees, name, tail): + raw, floats = CASES[name] + th = lib._thread_handle() + a, _ = _reply(raw, floats) + b, _ = _reply(raw, floats) + want_text, want_floats = _former(th, a) + if not tail: # call(): the text as it stands, tail marker included + want_text = _string_at(b).decode("utf-8") + del frees[:] + got_text, got_floats = lib._take(lib._lib, th, b, tail) + assert got_text == want_text + assert frees == [b] # once, and the reply that was read + if tail and floats is not None: + assert got_floats.dtype == np.float32 + assert np.array_equal(got_floats, want_floats) + assert ( + json.loads(got_text) is not None + ) # the text is a whole JSON value without the marker + else: + assert got_floats is None + + +def test_a_text_reply_is_not_copied_as_bytes(frees, string_at_sizes): + text = lib.call("arcadedb_n_diag") + assert json.loads(text) + assert ( + string_at_sizes == [] + ) # the former reader called string_at once, for the whole reply + assert len(frees) == 1 + + +def test_a_reply_with_a_tail_copies_only_the_marker_window_and_the_floats( + temp_db, frees, string_at_sizes +): + ptr, _ = _reply(LONG, np.arange(16, dtype=np.float32)) + th = lib._thread_handle() + text, floats = lib._take(lib._lib, th, ptr, True) + assert len(text) > 5000 and floats.size == 16 + assert sorted(string_at_sizes) == [ + 40, + 64, + ] # the last 40 bytes, 16 floats: never the 5 kB of text + assert frees == [ptr] + + +def test_a_reply_that_fails_to_decode_is_still_freed_once(temp_db, frees): + th = lib._thread_handle() + for raw in (b'{"a":"\xff"}', b'{"a":"\xc3"}', b'{"a":"\xed\xa0\x80"}'): + ptr, _ = _reply(raw) + del frees[:] + with pytest.raises(UnicodeDecodeError): + lib._take(lib._lib, th, ptr, False) + assert frees == [ptr] + ptr, _ = _reply(raw) + del frees[:] + with pytest.raises(UnicodeDecodeError): + lib._take(lib._lib, th, ptr, True) + assert frees == [ptr] + + +def test_a_reply_whose_tail_fails_is_still_freed_once(temp_db, frees): + th = lib._thread_handle() + ptr, _ = _reply(b'{"ok":true,"bin":x}') + with pytest.raises(ValueError): + lib._take(lib._lib, th, ptr, True) + assert frees == [ptr] + + +def test_no_reply_raises_and_frees_nothing(temp_db, frees): + with pytest.raises(ArcadeDBError): + lib._take(lib._lib, lib._thread_handle(), None, False) + assert frees == [] + + +def test_the_public_readers_free_each_reply_once(temp_db, frees): + temp_db.command("sql", "CREATE DOCUMENT TYPE R") + del frees[:] + text, floats = lib.call_bin("arcadedb_n_diag") + assert json.loads(text) and floats is None + assert len(frees) == 1 + lib.call("arcadedb_n_diag") + assert len(frees) == 2