Fix UnicodeEncodeError for non-ASCII file paths (#1086)

scalene/scalene_mapfile.py

@@ -9,6 +9,22 @@
 Filename = NewType("Filename", str)
 
 
+def decode_sample(raw: bytes) -> str:
+    """Decode one sample record written by the native side.
+
+    Records embed source filenames, which the native side encodes as UTF-8
+    with surrogate escapes (see ``encodeFilename`` in
+    ``src/source/pywhere.cpp``). Decoding the same way is the exact inverse,
+    so paths with non-ASCII characters -- and even paths with undecodable
+    bytes, which CPython represents as lone surrogates -- round-trip.
+
+    This used to decode as ASCII, which raised ``UnicodeDecodeError`` and
+    took down the sample-processing loop for any allocation attributed to a
+    file whose path contained e.g. an umlaut (issue #1086).
+    """
+    return raw.decode("utf-8", errors="surrogateescape")
+
+
 class ScaleneMapFile:
 
     # Things that need to be in sync with the C++ side
@@ -249,7 +265,7 @@ def _read_windows(self) -> bool:
             self._buf[length:] = b"\x00" * (self.MAX_BUFSIZE - length)
 
             # Validate the sample has expected format before accepting
-            sample_preview = self._buf[:100].decode("ascii", errors="replace")
+            sample_preview = self._buf[:100].decode("utf-8", errors="replace")
             if "," not in sample_preview:
                 # Malformed sample - skip it but advance position
                 self._lastpos = struct.pack("<Q", end_pos)  # type: ignore[assignment]
@@ -264,4 +280,4 @@ def _read_windows(self) -> bool:
 
     def get_str(self) -> str:
         """Get the string from the buffer."""
-        return self._buf.rstrip(b"\x00").split(b"\n")[0].decode("ascii")
+        return decode_sample(bytes(self._buf.rstrip(b"\x00").split(b"\n")[0]))

src/include/traceconfig.hpp

@@ -74,11 +74,17 @@ class TraceConfig {
     if (scalene_pkg_path_obj && scalene_pkg_path_obj != Py_None) {
       scalene_pkg_path_owner = scalene_pkg_path_obj;
       Py_IncRef(scalene_pkg_path_owner);
-      auto enc =
-          PyUnicode_AsEncodedString(scalene_pkg_path_obj, "utf-8", "strict");
+      // surrogateescape (not strict): a path holding undecodable bytes shows
+      // up in Python as lone surrogates, which a strict encoder rejects —
+      // leaving an exception pending that would later surface at an
+      // unrelated line (issue #1086).
+      auto enc = PyUnicode_AsEncodedString(scalene_pkg_path_obj, "utf-8",
+                                           "surrogateescape");
       if (enc) {
         scalene_pkg_path = std::string(PyBytes_AsString(enc));
         Py_DecRef(enc);
+      } else {
+        PyErr_Clear();
       }
     } else {
       scalene_pkg_path_owner = nullptr;

src/source/pywhere.cpp

@@ -172,6 +172,37 @@ PyObject* set_scalene_done_false(PyObject* self, PyObject* args) {
   Py_RETURN_NONE;
 }
 
+// Encode a Python filename (a ``str``, e.g. ``code->co_filename``) as
+// bytes we can hand to C string functions.
+//
+// This used to be ``PyUnicode_AsASCIIString``, which fails for any path
+// containing a non-ASCII character — a very ordinary situation (issue
+// #1086: a module at ``.../überschüsse.py``). Two things went wrong there:
+// the caller lost the filename, and — worse — the failed conversion left a
+// ``UnicodeEncodeError`` set on the thread. Since this code runs from the
+// allocator interposer and from trace callbacks, that stray exception
+// surfaced later at an arbitrary, unrelated Python line (or as a
+// ``SystemError``/crash), which is why the bug looked flaky.
+//
+// UTF-8 with ``surrogateescape`` is the right encoding here: it is the
+// inverse of how CPython decodes filesystem paths, so every path that
+// Python can represent — including undecodable bytes, which appear as lone
+// surrogates — round-trips. The Python side decodes these records the same
+// way (see ``decode_sample`` in scalene/scalene_mapfile.py). On failure we
+// return nullptr with no exception left pending; callers must NULL-check.
+static PyObject* encodeFilename(PyObject* unicode) {
+  if (unicode == nullptr) {
+    return nullptr;
+  }
+  PyObject* bytes =
+      PyUnicode_AsEncodedString(unicode, "utf-8", "surrogateescape");
+  if (bytes == nullptr) {
+    // Non-str object, or an encoder failure we can't do anything about.
+    PyErr_Clear();
+  }
+  return bytes;
+}
+
 // Core stack-walk used by both whereInPython and whereInPythonWithStack.
 // If ``stack_buf`` is non-null and ``stack_buf_size`` > 0, every traced
 // frame encountered (leaf-first) is appended to the buffer as the
@@ -224,7 +255,7 @@ static int whereInPythonImpl(std::string& filename, int& lineno, int& bytei,
     PyPtr<PyCodeObject> code =
         PyFrame_GetCode(static_cast<PyFrameObject*>(frame));
     PyPtr<> co_filename =
-        PyUnicode_AsASCIIString(static_cast<PyCodeObject*>(code)->co_filename);
+        encodeFilename(static_cast<PyCodeObject*>(code)->co_filename);
 
     if (!(static_cast<PyObject*>(co_filename))) {
       if (stack_bytes_written != nullptr) *stack_bytes_written = written;
@@ -454,16 +485,27 @@ static unchanging_modules module_pointers;
 static std::atomic<bool> module_pointers_ready{false};
 
 static bool on_stack(char* outer_filename, int lineno, PyFrameObject* frame) {
+  if (outer_filename == nullptr) {
+    // Nothing to compare against. Release the reference the caller handed
+    // us, matching what the loop below does on every exit path.
+    Py_XDECREF(frame);
+    return false;
+  }
   while (frame != NULL) {
     int iter_lineno = PyFrame_GetLineNumber(frame);
 
     PyPtr<PyCodeObject> code =
         PyFrame_GetCode(static_cast<PyFrameObject*>(frame));
 
     PyPtr<> co_filename(
-        PyUnicode_AsASCIIString(static_cast<PyCodeObject*>(code)->co_filename));
-    auto fname = PyBytes_AsString(static_cast<PyObject*>(co_filename));
-    if (iter_lineno == lineno && strstr(fname, outer_filename)) {
+        encodeFilename(static_cast<PyCodeObject*>(code)->co_filename));
+    // NULL-check before strstr: encoding can fail, and dereferencing the
+    // result unconditionally is how #1086 turned into a segfault.
+    auto fname = static_cast<PyObject*>(co_filename)
+                     ? PyBytes_AsString(static_cast<PyObject*>(co_filename))
+                     : nullptr;
+    if (fname != nullptr && iter_lineno == lineno &&
+        strstr(fname, outer_filename)) {
       Py_XDECREF(frame);
       return true;
     }
@@ -939,11 +981,11 @@ static int trace_func(PyObject* obj, PyFrameObject* frame, int what,
                         static_cast<PyCodeObject*>(code)->co_filename) == 0) {
     return 0;
   }
-  PyPtr<> last_fname_unicode(PyUnicode_AsASCIIString(last_fname));
+  PyPtr<> last_fname_unicode(encodeFilename(last_fname));
   auto last_fname_s =
-      PyBytes_AsString(static_cast<PyObject*>(last_fname_unicode));
-  PyPtr<> co_filename(
-      PyUnicode_AsASCIIString(static_cast<PyCodeObject*>(code)->co_filename));
+      static_cast<PyObject*>(last_fname_unicode)
+          ? PyBytes_AsString(static_cast<PyObject*>(last_fname_unicode))
+          : nullptr;
 
   // Needed because decref will be called in on_stack
   Py_INCREF(frame);

tests/test_issue1086_non_ascii_paths.py

@@ -0,0 +1,194 @@
+"""Regression test for issue #1086: non-ASCII characters in a profiled file's path.
+
+Issue: https://github.com/plasma-umass/scalene/issues/1086
+
+``pywhere.cpp`` converted ``code->co_filename`` with
+``PyUnicode_AsASCIIString`` on three hot paths (the stack walker used by
+the allocator interposer, ``on_stack``, and the settrace line callback).
+For a perfectly ordinary path like ``.../überschüsse.py`` that conversion
+fails, returning NULL *and leaving a ``UnicodeEncodeError`` set on the
+thread*. The whole path is encoded, so a single non-ASCII character
+anywhere in it is enough — the reporter's own case was a module in the
+import tree *at a path* with non-ASCII characters, i.e. an ordinary
+``.py`` file under a directory like ``optionale_verlängerung/``.
+Because those paths run from the native allocator hook and from
+trace callbacks, the stray exception surfaced later at an arbitrary,
+unrelated Python line — which is why the reporter saw it land in
+``sysconfig.get_path`` one run and inside pydantic the next. The Python
+side then compounded it by decoding the native sample records as ASCII.
+
+The fix encodes/decodes filenames as UTF-8 with surrogate escapes and
+NULL-checks every conversion. See ``encodeFilename`` in
+``src/source/pywhere.cpp`` and ``decode_sample`` in
+``scalene/scalene_mapfile.py``.
+
+Only ``--memory`` runs were affected (the reporter noted ``--cpu-only``
+worked), since that is what loads the native interposer, so the
+end-to-end test here profiles with memory enabled.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import unicodedata
+from pathlib import Path
+
+import pytest
+
+from scalene.scalene_mapfile import decode_sample
+
+# A module name with two umlauts, as in the issue's reproducer.
+NON_ASCII_MODULE = "überschüsse"
+
+WORKLOAD = """\
+def f():
+    total = 0
+    data = []
+    for _ in range(300):
+        data.append([j for j in range(2000)])
+        total += sum(data[-1])
+    return total
+"""
+
+# What WORKLOAD's f() returns: sum(range(2000)) * 300.
+EXPECTED_OUTPUT = str(sum(range(2000)) * 300)
+
+
+def _main_source(module: str) -> str:
+    return f"import {module}\nprint({module}.f())\n"
+
+
+# Where the non-ASCII character sits in the imported module's path. The
+# native code encoded the *whole* path, so any component poisons it --
+# the reporter hit both spellings (``überschüsse.py`` directly, and
+# ``optionale_verlängerung.py`` living under a non-ASCII project tree).
+#
+# Each entry is (id, directory components, module basename).
+PATH_LAYOUTS = [
+    # Non-ASCII in the leaf filename only; every directory is ASCII.
+    ("leaf-only", ("program",), NON_ASCII_MODULE),
+    # Non-ASCII in a *directory* only -- the module and its immediate
+    # parent are plain ASCII, so nothing but an interior path component
+    # carries the umlaut. This is the reporter's "a module in its import
+    # tree which is at a path with non-ascii characters" case.
+    ("directory-only", ("optionale_verlängerung", "pakete"), "workload"),
+    # Both, matching the original repro most closely.
+    ("both", ("prögramm",), NON_ASCII_MODULE),
+]
+
+
+def test_decode_sample_round_trips_non_ascii_filenames() -> None:
+    """Sample records naming a non-ASCII path must decode, not raise.
+
+    This is the Python half of the fix, and it is deterministic: the old
+    ``bytes.decode("ascii")`` raised ``UnicodeDecodeError`` here and killed
+    the loop that drains malloc/free samples.
+    """
+    for path in (
+        # Non-ASCII in the leaf.
+        f"/tmp/project/{NON_ASCII_MODULE}.py",
+        # Non-ASCII only in an interior directory, ASCII leaf.
+        "/tmp/optionale_verlängerung/pakete/workload.py",
+    ):
+        record = f"M,1234,4096,0.5,999,0x7f00,{path},7,0"
+        assert decode_sample(record.encode("utf-8")) == record
+    # Undecodable filesystem bytes reach Python as lone surrogates; those
+    # must survive the round trip too rather than raising.
+    surrogate_path = "/tmp/\udcff/x.py"
+    surrogate_record = f"M,1,1,0.0,1,0x0,{surrogate_path},1,0"
+    assert (
+        decode_sample(surrogate_record.encode("utf-8", errors="surrogateescape"))
+        == surrogate_record
+    )
+
+
+@pytest.mark.parametrize(
+    ("dirs", "module"),
+    [pytest.param(dirs, module, id=name) for name, dirs, module in PATH_LAYOUTS],
+)
+def test_memory_profile_of_non_ascii_path(
+    tmp_path: Path, dirs: tuple, module: str
+) -> None:
+    """Profiling a program that imports a module at a non-ASCII path works.
+
+    See ``PATH_LAYOUTS`` for the placements covered — notably the
+    ``directory-only`` case, where the non-ASCII character appears solely
+    in an interior directory and both the module file and its immediate
+    parent are ASCII.
+
+    The crash assertions run on every attempt — a regression fails them
+    immediately and deterministically, since the profiled program dies
+    outright. Only the "file shows up in the profile" check is retried,
+    against the usual sampling flake (see ``_scalene_subprocess.py``).
+    """
+    workdir = tmp_path.joinpath(*dirs)
+    workdir.mkdir(parents=True)
+    module_path = workdir / f"{module}.py"
+    module_path.write_text(WORKLOAD, encoding="utf-8")
+    main_path = workdir / "main.py"
+    main_path.write_text(_main_source(module), encoding="utf-8")
+
+    attempts = 3
+    last_files: list = []
+    last_output = ""
+    for attempt in range(1, attempts + 1):
+        out = tmp_path / f"profile_{attempt}.json"
+        try:
+            proc = subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "scalene",
+                    "run",
+                    "--memory",
+                    "--no-browser",
+                    "-o",
+                    str(out),
+                    str(main_path),
+                ],
+                cwd=str(workdir),
+                check=False,
+                capture_output=True,
+                text=True,
+                timeout=300,
+            )
+        except subprocess.TimeoutExpired:
+            # Scalene startup occasionally wedges under CI contention.
+            continue
+        last_output = combined = proc.stdout + proc.stderr
+
+        assert "UnicodeEncodeError" not in combined, (
+            "Non-ASCII path raised UnicodeEncodeError in the profiled program "
+            f"(issue #1086):\n{combined[-3000:]}"
+        )
+        assert "UnicodeDecodeError" not in combined, (
+            f"Non-ASCII path raised UnicodeDecodeError:\n{combined[-3000:]}"
+        )
+        # The workload prints its result; a crash of the profiled program
+        # (the #1086 failure mode) means this never appears.
+        assert EXPECTED_OUTPUT in proc.stdout, (
+            f"Profiled program did not run to completion:\n{combined[-3000:]}"
+        )
+        assert proc.returncode == 0, (
+            f"scalene exited {proc.returncode}:\n{combined[-3000:]}"
+        )
+
+        if not (out.exists() and out.stat().st_size > 0):
+            continue
+        files = json.loads(out.read_text(encoding="utf-8")).get("files", {})
+        last_files = list(files)
+        # The recorded path must come back intact, not mangled or escaped —
+        # this exercises the native encode / Python decode pair end to end.
+        # Normalize first: a filesystem may hand back the decomposed (NFD)
+        # spelling of "ü" regardless of how we wrote it.
+        want = unicodedata.normalize("NFC", str(module_path))
+        if any(unicodedata.normalize("NFC", name) == want for name in files):
+            return
+    pytest.skip(
+        "Scalene recorded no samples for the non-ASCII module after "
+        f"{attempts} attempts (suspected sampling flake). The program itself "
+        f"ran cleanly, so the #1086 crash is not present. "
+        f"Profiled files: {last_files}. Last output:\n{last_output[-1000:]}"
+    )