Native: read each library reply from its address, one copy, freed once (closes #323)

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:

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

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