diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 95d7b96..3688e02 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -15,6 +15,10 @@ jobs: test-linux: runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.13', '3.14'] steps: - name: Checkout repository uses: actions/checkout@v3 @@ -28,22 +32,28 @@ jobs: context: . push: false load: true - tags: sauerkraut:test + tags: sauerkraut:test-${{ matrix.python-version }} + build-args: | + PYTHON_VERSION=${{ matrix.python-version }} - name: Run tests run: | - docker run --rm sauerkraut:test python3 /sauerkraut/test/test.py + docker run --rm sauerkraut:test-${{ matrix.python-version }} python3 /sauerkraut/test/test.py test-macos: runs-on: macos-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.13', '3.14'] steps: - name: Checkout repository uses: actions/checkout@v3 - - name: Set up Python 3.13 + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: - python-version: '3.13' + python-version: ${{ matrix.python-version }} - name: Install dependencies run: | diff --git a/CMakeLists.txt b/CMakeLists.txt index 82bc4ad..97821aa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -12,7 +12,7 @@ set(CMAKE_BUILD_TYPE Debug) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -Wall") # Find Python (ensure the version matches your Python version) -find_package(Python 3.13 COMPONENTS Interpreter Development REQUIRED) +find_package(Python 3.13...<3.15 COMPONENTS Interpreter Development REQUIRED) # Set consistent output directories set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) diff --git a/Dockerfile b/Dockerfile index e8fbc13..5a3e273 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,6 @@ FROM ubuntu:22.04 AS base LABEL maintainer="Zane Fink " +ARG PYTHON_VERSION=3.13 WORKDIR / ENV DEBIAN_FRONTEND=noninteractive RUN ln -fs /usr/share/zoneinfo/UTC /etc/localtime && \ @@ -18,7 +19,7 @@ ENV PATH=/opt/conda/bin:$PATH RUN conda init bash && \ conda config --add channels conda-forge && \ conda config --set channel_priority strict && \ - conda create -n sauerkraut python=3.13 -y -c conda-forge --override-channels && \ + conda create -n sauerkraut python=${PYTHON_VERSION} -y -c conda-forge --override-channels && \ echo "conda activate sauerkraut" >> ~/.bashrc # Make the conda environment available in PATH ENV PATH=/opt/conda/envs/sauerkraut/bin:$PATH diff --git a/README.md b/README.md index 7f79bd5..ad5cbc8 100644 --- a/README.md +++ b/README.md @@ -127,4 +127,4 @@ python3 copy_then_serialize.py ## Compatibility Sauerkraut leverages intimate knowledge of CPython internals, and as such is vulnerable to changes in the CPython API and VM. -Currently, Sauerkraut supports Python 3.13 and the development version of Python 3.14. \ No newline at end of file +Currently, Sauerkraut supports Python 3.13 and 3.14. diff --git a/pyproject.toml b/pyproject.toml index ed87549..ef72168 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ classifiers = [ "Topic :: System :: Distributed Computing", ] dependencies = [ - "bytecode", + "bytecode @ git+https://github.com/ZwFink/bytecode.git", "numpy", ] diff --git a/sauerkraut/CMakeLists.txt b/sauerkraut/CMakeLists.txt index 7dad7cd..9587147 100644 --- a/sauerkraut/CMakeLists.txt +++ b/sauerkraut/CMakeLists.txt @@ -48,6 +48,10 @@ if(APPLE) greenlet_compat ) target_link_options(sauerkraut PRIVATE -undefined dynamic_lookup) + set_target_properties(sauerkraut PROPERTIES + BUILD_WITH_INSTALL_RPATH ON + INSTALL_RPATH "@loader_path" + ) else() target_link_libraries(sauerkraut PRIVATE diff --git a/sauerkraut/include/py_structs.h b/sauerkraut/include/py_structs.h index 7d5e401..de7716c 100644 --- a/sauerkraut/include/py_structs.h +++ b/sauerkraut/include/py_structs.h @@ -13,6 +13,8 @@ typedef struct _CodeUnit { uint8_t oparg; } _CodeUnit; +typedef _CodeUnit _Py_CODEUNIT; + struct _frame { PyObject_HEAD PyFrameObject *f_back; /* previous frame, or NULL */ @@ -23,6 +25,9 @@ struct _frame { char f_trace_opcodes; /* Emit per-opcode trace events? */ PyObject *f_extra_locals; /* Dict for locals set by users using f_locals, could be NULL */ PyObject *f_locals_cache; /* Backwards compatibility for PyEval_GetLocals */ +#if SAUERKRAUT_PY314 + PyObject *f_overwritten_fast_locals; +#endif PyObject *_f_frame_data[1]; /* Frame data if this frame object owns the frame */ }; @@ -32,19 +37,29 @@ _PyThreadState_PushFrame(PyThreadState *tstate, size_t size); typedef struct _PyInterpreterFrame { _PyStackRef f_executable; /* Deferred or strong reference (code object or None) */ struct _PyInterpreterFrame *previous; +#if SAUERKRAUT_PY314 + _PyStackRef f_funcobj; /* Deferred or strong reference. Only valid if not on C stack */ +#else PyObject *f_funcobj; /* Strong reference. Only valid if not on C stack */ +#endif PyObject *f_globals; /* Borrowed reference. Only valid if not on C stack */ PyObject *f_builtins; /* Borrowed reference. Only valid if not on C stack */ PyObject *f_locals; /* Strong reference, may be NULL. Only valid if not on C stack */ PyFrameObject *frame_obj; /* Strong reference, may be NULL. Only valid if not on C stack */ - _CodeUnit *instr_ptr; /* Instruction currently executing (or about to begin) */ - #if SAUERKRAUT_PY314 + _Py_CODEUNIT *instr_ptr; /* Instruction currently executing (or about to begin) */ +#if SAUERKRAUT_PY314 _PyStackRef *stackpointer; - #elif SAUERKRAUT_PY313 +#ifdef Py_GIL_DISABLED + int32_t tlbc_index; /* Index of thread-local bytecode containing instr_ptr */ +#endif +#elif SAUERKRAUT_PY313 int stacktop; - #endif +#endif uint16_t return_offset; /* Only relevant during a function call */ char owner; +#if SAUERKRAUT_PY314 + uint8_t visited; +#endif /* Locals and stack */ _PyStackRef localsplus[1]; } _PyInterpreterFrame; diff --git a/sauerkraut/include/pyref.h b/sauerkraut/include/pyref.h index 49576ab..f8726b7 100644 --- a/sauerkraut/include/pyref.h +++ b/sauerkraut/include/pyref.h @@ -73,6 +73,11 @@ class py_strongref { T *operator->() { return this->borrow(); } + + void reset() { + Py_XDECREF((PyObject*) obj); + obj = NULL; + } }; template diff --git a/sauerkraut/include/utils.h b/sauerkraut/include/utils.h index 989ac6e..c7becec 100644 --- a/sauerkraut/include/utils.h +++ b/sauerkraut/include/utils.h @@ -193,6 +193,143 @@ namespace utils { return PyDict_Check(obj); } + struct StackRefObject { + PyObject *obj; + bool owned; + }; + + class ScopedStackRefObject { + StackRefObject ref; + public: + ScopedStackRefObject(_PyStackRef stackref); + ~ScopedStackRefObject() { + if (ref.owned && ref.obj) { + Py_DECREF(ref.obj); + } + } + ScopedStackRefObject(const ScopedStackRefObject&) = delete; + ScopedStackRefObject& operator=(const ScopedStackRefObject&) = delete; + PyObject* get() const { return ref.obj; } + explicit operator bool() const { return ref.obj != nullptr; } + }; + +#if SAUERKRAUT_PY314 + constexpr uintptr_t STACKREF_TAG_BITS = 0x7; + constexpr uintptr_t STACKREF_TAG_INT = 0x3; + constexpr uintptr_t STACKREF_TAG_REFCNT = 0x1; + constexpr int STACKREF_TAGGED_SHIFT = 2; + + inline bool stackref_is_null(_PyStackRef ref) { + if (ref.bits == 0) return true; + // Check for deferred NULL: when only the deferred tag is set (bits=0x1) + // After clearing tags, the pointer would be NULL + if ((ref.bits & STACKREF_TAG_BITS) == STACKREF_TAG_REFCNT && + (ref.bits & ~STACKREF_TAG_BITS) == 0) { + return true; + } + return false; + } + + inline bool stackref_is_tagged_int(_PyStackRef ref) { + return (ref.bits & STACKREF_TAG_BITS) == STACKREF_TAG_INT; + } + + inline bool stackref_is_deferred(_PyStackRef ref) { + return (ref.bits & STACKREF_TAG_BITS) == STACKREF_TAG_REFCNT; + } + + inline intptr_t stackref_untag_int(_PyStackRef ref) { + return ((intptr_t)ref.bits) >> STACKREF_TAGGED_SHIFT; + } + + inline PyObject *stackref_as_pyobject(_PyStackRef ref) { + if (stackref_is_null(ref) || stackref_is_tagged_int(ref)) { + return NULL; + } + uintptr_t bits = ref.bits; + if (stackref_is_deferred(ref)) { + bits &= ~STACKREF_TAG_BITS; + } + return (PyObject *)bits; + } + + inline StackRefObject stackref_to_object_for_serialization(_PyStackRef ref) { + if (stackref_is_null(ref)) { + return {NULL, false}; + } + if (stackref_is_tagged_int(ref)) { + auto value = (Py_ssize_t) stackref_untag_int(ref); + return {PyLong_FromSsize_t(value), true}; + } + return {stackref_as_pyobject(ref), false}; + } + + inline void stackref_decref(_PyStackRef ref) { + if (stackref_is_null(ref) || stackref_is_tagged_int(ref)) { + return; + } + Py_XDECREF(stackref_as_pyobject(ref)); + } + + inline PyObject *get_funcobj(sauerkraut::PyInterpreterFrame *frame) { + return stackref_as_pyobject(frame->f_funcobj); + } + + inline void set_funcobj(sauerkraut::PyInterpreterFrame *frame, PyObject *obj) { + frame->f_funcobj.bits = (uintptr_t)obj; + } + + inline void set_stack_position(sauerkraut::PyInterpreterFrame *frame, int nlocalsplus, int stack_depth) { + frame->stackpointer = frame->localsplus + nlocalsplus + stack_depth; + } + + inline void init_frame_visited(sauerkraut::PyInterpreterFrame *frame) { + frame->visited = 0; + } +#else + inline bool stackref_is_null(_PyStackRef ref) { + return ref.bits == 0; + } + + inline bool stackref_is_tagged_int(_PyStackRef) { + return false; + } + + inline PyObject *stackref_as_pyobject(_PyStackRef ref) { + return (PyObject *) ref.bits; + } + + inline StackRefObject stackref_to_object_for_serialization(_PyStackRef ref) { + if (ref.bits == 0) { + return {NULL, false}; + } + return {(PyObject *) ref.bits, false}; + } + + inline void stackref_decref(_PyStackRef ref) { + Py_XDECREF((PyObject *) ref.bits); + } + + inline PyObject *get_funcobj(sauerkraut::PyInterpreterFrame *frame) { + return frame->f_funcobj; + } + + inline void set_funcobj(sauerkraut::PyInterpreterFrame *frame, PyObject *obj) { + frame->f_funcobj = obj; + } + + inline void set_stack_position(sauerkraut::PyInterpreterFrame *frame, int nlocalsplus, int stack_depth) { + frame->stacktop = nlocalsplus + stack_depth; + } + + inline void init_frame_visited(sauerkraut::PyInterpreterFrame *) { + // No-op for Python 3.13 + } +#endif + + inline ScopedStackRefObject::ScopedStackRefObject(_PyStackRef stackref) + : ref(stackref_to_object_for_serialization(stackref)) {} + int get_code_stacksize(PyCodeObject *code) { return code->co_stacksize; } @@ -468,7 +605,11 @@ namespace utils { _PyStackRef *stack_pointer = iframe->localsplus + code->co_nlocalsplus; for(int i = 0; i < stack_depth; i++) { - PyObject *stack_obj = (PyObject*) stack_pointer[i].bits; + _PyStackRef stack_ref = stack_pointer[i]; + PyObject *stack_obj = stackref_as_pyobject(stack_ref); + if (stack_obj == NULL && stackref_is_tagged_int(stack_ref)) { + stack_obj = Py_None; + } #ifdef DEBUG assert(NULL != stack_obj); if(locals.find((intptr_t) stack_obj) != locals.end()) { @@ -502,7 +643,7 @@ namespace utils { LocalNameMap local_idx_map; for(int i = 0; i < code->co_nlocalsplus; i++) { - PyObject *local = ((PyObject*) iframe->localsplus[i].bits); + PyObject *local = stackref_as_pyobject(iframe->localsplus[i]); std::string name = PyUnicode_AsUTF8(PyTuple_GetItem(locals_plus_names, i)); local_idx_map[name] = i; } @@ -575,14 +716,14 @@ namespace utils { } if (local_index >= 0) { - PyObject *old_local = (PyObject*) iframe->localsplus[local_index].bits; + _PyStackRef old_ref = iframe->localsplus[local_index]; // Increment reference count of new value before assigning Py_INCREF(value); iframe->localsplus[local_index].bits = (intptr_t) value; // Decrement reference count of old value - Py_XDECREF(old_local); + stackref_decref(old_ref); } } } @@ -590,4 +731,4 @@ namespace utils { } -#endif // UTILS_HH_INCLUDED \ No newline at end of file +#endif // UTILS_HH_INCLUDED diff --git a/sauerkraut/liveness.py b/sauerkraut/liveness.py index fab6c03..991472c 100644 --- a/sauerkraut/liveness.py +++ b/sauerkraut/liveness.py @@ -3,8 +3,14 @@ import bytecode as bc from bytecode import Instr, BasicBlock, ControlFlowGraph, Bytecode -_USE_INSTRS = ("LOAD_NAME", "LOAD_FAST", "LOAD_FAST_CHECK", "LOAD_FAST_AND_CLEAR") -_SUPER_USE_INSTRS = "LOAD_FAST_LOAD_FAST" +_USE_INSTRS = ( + "LOAD_NAME", + "LOAD_FAST", + "LOAD_FAST_CHECK", + "LOAD_FAST_AND_CLEAR", + "LOAD_FAST_BORROW", +) +_SUPER_USE_INSTRS = ("LOAD_FAST_LOAD_FAST", "LOAD_FAST_BORROW_LOAD_FAST_BORROW") _SUPER_DEF_INSTRS = "STORE_FAST_STORE_FAST" _DEF_INSTRS = ("STORE_NAME", "STORE_FAST") @@ -96,7 +102,11 @@ def _get_uses_and_defs(self, block: BasicBlock) -> Tuple[Set[str], Set[str]]: if isinstance(arg1, str): defined_vars.add(arg1) elif instr.name == "STORE_FAST_LOAD_FAST": - print(f"STORE_FAST_LOAD_FAST: {instr.arg}") + arg0, arg1 = instr.arg + if isinstance(arg0, str): + defined_vars.add(arg0) # First operand is stored + if isinstance(arg1, str): + used_vars.add(arg1) # Second operand is loaded return used_vars, defined_vars diff --git a/sauerkraut/sauerkraut.C b/sauerkraut/sauerkraut.C index 89aa2c0..b0e33a7 100644 --- a/sauerkraut/sauerkraut.C +++ b/sauerkraut/sauerkraut.C @@ -33,18 +33,42 @@ class sauerkraut_modulestate { pyobject_strongref liveness_module; pyobject_strongref get_dead_variables_at_offset; PyCodeImmutableCache code_immutable_cache; - sauerkraut_modulestate() { - deepcopy_module = PyImport_ImportModule("copy"); - deepcopy = PyObject_GetAttrString(*deepcopy_module, "deepcopy"); - pickle_module = PyImport_ImportModule("pickle"); - pickle_dumps = PyObject_GetAttrString(*pickle_module, "dumps"); - pickle_loads = PyObject_GetAttrString(*pickle_module, "loads"); - - dill_module = PyImport_ImportModule("dill"); - dill_dumps = PyObject_GetAttrString(*dill_module, "dumps"); - dill_loads = PyObject_GetAttrString(*dill_module, "loads"); - liveness_module = PyImport_ImportModule("sauerkraut.liveness"); - get_dead_variables_at_offset = PyObject_GetAttrString(*liveness_module, "get_dead_variables_at_offset"); + sauerkraut_modulestate() = default; + + bool init() { + auto import_module = [](const char* name, pyobject_strongref& dest) -> bool { + dest = PyImport_ImportModule(name); + return static_cast(dest); + }; + + auto get_attr = [](pyobject_strongref& module, const char* attr, pyobject_strongref& dest) -> bool { + dest = PyObject_GetAttrString(*module, attr); + return static_cast(dest); + }; + + if (!import_module("copy", deepcopy_module) || + !get_attr(deepcopy_module, "deepcopy", deepcopy)) { + return false; + } + + if (!import_module("pickle", pickle_module) || + !get_attr(pickle_module, "dumps", pickle_dumps) || + !get_attr(pickle_module, "loads", pickle_loads)) { + return false; + } + + if (!import_module("dill", dill_module) || + !get_attr(dill_module, "dumps", dill_dumps) || + !get_attr(dill_module, "loads", dill_loads)) { + return false; + } + + if (!import_module("sauerkraut.liveness", liveness_module) || + !get_attr(liveness_module, "get_dead_variables_at_offset", get_dead_variables_at_offset)) { + return false; + } + + return true; } pyobject_strongref get_dead_variables(py_weakref code, int offset) { @@ -65,7 +89,7 @@ class sauerkraut_modulestate { } // it's not in the cache, so we need to compute the invariants - auto funcobj = make_strongref(frame->f_frame->f_funcobj); + auto funcobj = make_strongref(utils::py::get_funcobj(frame->f_frame)); code_immutable_cache[name_str] = std::make_tuple(funcobj, code, frame->f_frame->f_globals); } @@ -93,6 +117,22 @@ class sauerkraut_modulestate { return get_code_immutables(frame.f_frame); } + void clear() { + // Clear the cache first - this decrefs Python objects while interpreter is still valid + code_immutable_cache.clear(); + // Clear all module references + deepcopy.reset(); + deepcopy_module.reset(); + pickle_module.reset(); + pickle_dumps.reset(); + pickle_loads.reset(); + dill_module.reset(); + dill_dumps.reset(); + dill_loads.reset(); + liveness_module.reset(); + get_dead_variables_at_offset.reset(); + } + }; class dumps_functor { @@ -217,20 +257,23 @@ PyObject *deepcopy_object(py_weakref obj) { return copy_obj; } -static void cleanup_interpreter_frame(_PyInterpreterFrame *interp, int nlocalsplus, int stack_depth) { - Py_XDECREF((PyObject*)interp->f_executable.bits); - Py_XDECREF(interp->f_funcobj); +static void decref_interpreter_frame_refs(_PyInterpreterFrame *interp, int nlocalsplus, int stack_depth) { + utils::py::stackref_decref(interp->f_executable); + Py_XDECREF(utils::py::get_funcobj(interp)); Py_XDECREF(interp->f_locals); for (int i = 0; i < nlocalsplus; i++) { - Py_XDECREF((PyObject*)interp->localsplus[i].bits); + utils::py::stackref_decref(interp->localsplus[i]); } _PyStackRef *stack_base = interp->localsplus + nlocalsplus; for (int i = 0; i < stack_depth; i++) { - Py_XDECREF((PyObject*)stack_base[i].bits); + utils::py::stackref_decref(stack_base[i]); } +} +static void cleanup_interpreter_frame(_PyInterpreterFrame *interp, int nlocalsplus, int stack_depth) { + decref_interpreter_frame_refs(interp, nlocalsplus, stack_depth); free(interp); } @@ -245,24 +288,9 @@ typedef struct frame_copy_capsule { ~frame_copy_capsule() { if (frame) { if (owns_interpreter_frame && frame->f_frame) { - auto *interp = frame->f_frame; - - Py_XDECREF((PyObject*)interp->f_executable.bits); - Py_XDECREF(interp->f_funcobj); - Py_XDECREF(interp->f_locals); - - for (int i = 0; i < nlocalsplus; i++) { - Py_XDECREF((PyObject*)interp->localsplus[i].bits); - } - - _PyStackRef *stack_base = interp->localsplus + nlocalsplus; - for (int i = 0; i < stack_depth; i++) { - Py_XDECREF((PyObject*)stack_base[i].bits); - } - // f_globals, f_builtins are borrowed refs; frame_obj is weak (no Py_NewRef) - - free(interp); + decref_interpreter_frame_refs(frame->f_frame, nlocalsplus, stack_depth); + free(frame->f_frame); frame->f_frame = NULL; } Py_XDECREF(frame); @@ -292,13 +320,17 @@ PyObject *frame_copy_capsule_create(py_weakref frame, utils::py:: return PyCapsule_New(copy_capsule, copy_frame_capsule_name, frame_copy_capsule_destroy); } -void copy_localsplus(py_weakref to_copy, - py_weakref new_frame, +void copy_localsplus(py_weakref to_copy, + py_weakref new_frame, int nlocals, int deepcopy) { if (deepcopy) { for (int i = 0; i < nlocals; i++) { - py_weakref local{(PyObject*)to_copy->localsplus[i].bits}; - PyObject *local_copy = deepcopy_object(local); + utils::py::ScopedStackRefObject local_obj(to_copy->localsplus[i]); + if (!local_obj) { + new_frame->localsplus[i].bits = 0; + continue; + } + PyObject *local_copy = deepcopy_object(make_weakref(local_obj.get())); new_frame->localsplus[i].bits = (uintptr_t)local_copy; } } else { @@ -306,17 +338,21 @@ void copy_localsplus(py_weakref to_copy, } } -void copy_stack(py_weakref to_copy, - py_weakref new_frame, +void copy_stack(py_weakref to_copy, + py_weakref new_frame, int stack_size, int deepcopy) { _PyStackRef *src_stack_base = utils::py::get_stack_base(*to_copy); _PyStackRef *dest_stack_base = utils::py::get_stack_base(*new_frame); if(deepcopy) { for(int i = 0; i < stack_size; i++) { - auto stack_obj = make_weakref((PyObject*)src_stack_base[i].bits); - PyObject *stack_obj_copy = deepcopy_object(stack_obj); - dest_stack_base[i].bits = (uintptr_t) stack_obj_copy; + utils::py::ScopedStackRefObject stack_obj(src_stack_base[i]); + if (!stack_obj) { + dest_stack_base[i].bits = 0; + continue; + } + PyObject *stack_obj_copy = deepcopy_object(make_weakref(stack_obj.get())); + dest_stack_base[i].bits = (uintptr_t)stack_obj_copy; } } else { memcpy(dest_stack_base, src_stack_base, stack_size * sizeof(_PyStackRef)); @@ -352,29 +388,35 @@ PyFrameObject *create_copied_frame(py_weakref tstate, return NULL; } + // PyFrame_New incref'd locals and stored them in the embedded frame. + // Clear them before replacing f_frame to avoid leaking that reference. + if (new_frame->f_frame && new_frame->f_frame->f_locals) { + Py_DECREF(new_frame->f_frame->f_locals); + new_frame->f_frame->f_locals = NULL; + } + new_frame->f_frame = stack_frame; py_weakref new_frame_ref{new_frame->f_frame}; new_frame_ref->owner = to_copy->owner; new_frame_ref->previous = set_previous ? *to_copy : NULL; - new_frame_ref->f_funcobj = deepcopy_object(make_weakref(to_copy->f_funcobj)); + utils::py::set_funcobj(*new_frame_ref, deepcopy_object(make_weakref(utils::py::get_funcobj(*to_copy)))); new_frame_ref->f_executable.bits = (uintptr_t)deepcopy_object(make_weakref((PyObject*)to_copy->f_executable.bits)); new_frame_ref->f_globals = to_copy->f_globals; new_frame_ref->f_builtins = to_copy->f_builtins; new_frame_ref->f_locals = to_copy->f_locals; new_frame_ref->return_offset = to_copy->return_offset; new_frame_ref->frame_obj = new_frame; - #if SAUERKRAUT_PY314 - new_frame->f_frame->stackpointer = NULL; - #elif SAUERKRAUT_PY313 - new_frame->f_frame->stacktop = 0; - #endif auto offset = utils::py::get_instr_offset(to_copy); new_frame->f_frame->instr_ptr = (_CodeUnit*) (code_obj->co_code_adaptive + offset); copy_localsplus(to_copy, new_frame_ref, nlocals, deepcopy_localsplus); copy_stack(to_copy, new_frame_ref, stack_size, 1); + // Set stack position after copying stack + utils::py::set_stack_position(new_frame->f_frame, nlocals, stack_size); + utils::py::init_frame_visited(new_frame->f_frame); + if(push_frame) { return *prepare_frame_for_execution(new_frame); } else { @@ -397,13 +439,9 @@ PyFrameObject *push_frame_for_running(PyThreadState *tstate, _PyInterpreterFrame auto offset = utils::py::get_instr_offset(to_push->frame_obj); stack_frame->owner = to_push->owner; - // needs to be the currently executing frame - py_weakref current_frame{(PyFrameObject*) PyEval_GetFrame()}; - if(!current_frame) { - stack_frame->previous = NULL; - } else { - stack_frame->previous = current_frame->f_frame; - } + // Set previous to NULL so that when the frame returns, it exits the interpreter loop + // rather than trying to continue in some other frame + stack_frame->previous = NULL; stack_frame->f_funcobj = to_push->f_funcobj; stack_frame->f_executable.bits = to_push->f_executable.bits; stack_frame->f_globals = to_push->f_globals; @@ -413,11 +451,9 @@ PyFrameObject *push_frame_for_running(PyThreadState *tstate, _PyInterpreterFrame stack_frame->instr_ptr = (_CodeUnit*) (code->co_code_adaptive + (offset)); auto stack_depth = utils::py::get_current_stack_depth(to_push); copy_stack(to_push, stack_frame, stack_depth, 0); - #if SAUERKRAUT_PY314 - stack_frame->stackpointer = stack_frame->localsplus + code->co_nlocalsplus + stack_depth; - #elif SAUERKRAUT_PY313 - stack_frame->stacktop = code->co_nlocalsplus + stack_depth; - #endif + utils::py::set_stack_position(stack_frame, code->co_nlocalsplus, stack_depth); + utils::py::init_frame_visited(stack_frame); + stack_frame->return_offset = to_push->return_offset; pyframe_object->f_frame = stack_frame; return *prepare_frame_for_execution(pyframe_object); @@ -429,6 +465,7 @@ struct SerializationOptions { Py_ssize_t sizehint = 0; bool exclude_dead_locals = true; bool exclude_immutables = false; + serdes::SerializationArgs to_ser_args() const { serdes::SerializationArgs args; if (sizehint > 0) { @@ -437,6 +474,14 @@ struct SerializationOptions { args.set_exclude_immutables(exclude_immutables); return args; } + + void populate(int serialize_int, PyObject* exclude_locals_obj, + int exclude_dead_locals_int, int exclude_immutables_int) { + serialize = (serialize_int != 0); + exclude_dead_locals = (exclude_dead_locals_int != 0); + exclude_immutables = (exclude_immutables_int != 0); + exclude_locals = pyobject_strongref(exclude_locals_obj); + } }; static pyobject_strongref combine_exclusions(py_weakref frame, PyObject* exclude_locals, bool exclude_dead_locals) { @@ -498,7 +543,9 @@ static PyObject *_copy_frame_object(py_weakref frame, const Seria auto stack_state = utils::py::get_stack_state((PyObject*)*frame); PyFrameObject *new_frame = create_copied_frame(tstate, to_copy, copy_code_obj, LocalCopy, 0, 1, 0, stack_state.size(), 1); - PyObject *capsule = frame_copy_capsule_create(new_frame, stack_state, true); + int nlocalsplus = copy_code_obj->co_nlocalsplus; + int stack_depth = stack_state.size(); + PyObject *capsule = frame_copy_capsule_create(new_frame, stack_state, true, nlocalsplus, stack_depth); Py_DECREF(new_frame); // Drop our ref; capsule holds its own Py_DECREF(copy_code_obj); Py_DECREF(LocalCopy); @@ -509,20 +556,24 @@ static PyObject *_copy_frame_object(py_weakref frame, const Seria static PyObject *_copy_serialize_frame_object(py_weakref frame, const SerializationOptions& options) { - using namespace utils; + if(options.exclude_immutables) { + sauerkraut_state->cache_code_immutables(frame); + } + + // First copy the frame, then serialize from the copy + // This ensures we have a consistent snapshot of the frame state + PyObject *capsule = _copy_frame_object(frame, options); + if (capsule == NULL) { + return NULL; + } + serdes::SerializationArgs args = options.to_ser_args(); - if (!apply_exclusions(frame, options, args)) { + Py_DECREF(capsule); return NULL; } - - if(options.exclude_immutables) { - sauerkraut_state->cache_code_immutables(frame); - } - auto stack_state = utils::py::get_stack_state((PyObject*)*frame); - std::unique_ptr capsule(frame_copy_capsule_create_direct(frame, stack_state)); - - PyObject *ret = _serialize_frame_direct_from_capsule(capsule.get(), args); + PyObject *ret = _serialize_frame_from_capsule(capsule, args); + Py_DECREF(capsule); // Done with the capsule return ret; } @@ -553,38 +604,33 @@ static bool parse_sizehint(PyObject* sizehint_obj, Py_ssize_t& sizehint) { } static bool parse_serialization_options(PyObject* args, PyObject* kwargs, SerializationOptions& options) { - static char* kwlist[] = {"serialize", "exclude_locals", - "exclude_immutables", "sizehint", + static char* kwlist[] = {"serialize", "exclude_locals", + "exclude_immutables", "sizehint", "exclude_dead_locals", NULL}; int serialize = 0; PyObject* sizehint_obj = NULL; PyObject* exclude_locals = NULL; int exclude_dead_locals = 1; int exclude_immutables = 0; - - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|pOpOp", kwlist, - &serialize, &exclude_locals, - &exclude_immutables, &sizehint_obj, + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|pOpOp", kwlist, + &serialize, &exclude_locals, + &exclude_immutables, &sizehint_obj, &exclude_dead_locals)) { return false; } - - options.serialize = (serialize != 0); - options.exclude_dead_locals = (exclude_dead_locals != 0); - options.exclude_locals = pyobject_strongref(exclude_locals); - options.exclude_immutables = (exclude_immutables != 0); + + options.populate(serialize, exclude_locals, exclude_dead_locals, exclude_immutables); return parse_sizehint(sizehint_obj, options.sizehint); } static PyObject *run_and_cleanup_frame(PyFrameObject *frame) { PyObject *res = PyEval_EvalFrame(frame); - PyCodeObject *code = PyFrame_GetCode(frame); - // Clear f_frame before it becomes a dangling pointer + // The stack frame is automatically cleaned up by Python after PyEval_EvalFrame. + // We just need to clear f_frame to avoid dangling pointer when the frame is GC'd. frame->f_frame = NULL; - Py_SET_REFCNT(code, 0); - Py_SET_REFCNT(frame, 0); return res; } @@ -604,22 +650,21 @@ static PyObject *copy_current_frame(PyObject *self, PyObject *args, PyObject *kw static PyObject *copy_frame(PyObject *self, PyObject *args, PyObject *kwargs) { PyObject *frame = NULL; SerializationOptions options; - - static char *kwlist[] = {"frame", "exclude_locals", "sizehint", + + static char *kwlist[] = {"frame", "exclude_locals", "sizehint", "serialize", "exclude_dead_locals", "exclude_immutables", NULL}; int serialize = 0; PyObject* sizehint_obj = NULL; + PyObject* exclude_locals = NULL; int exclude_dead_locals = 1; int exclude_immutables = 0; - - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|OOppp", kwlist, - &frame, &options.exclude_locals, &sizehint_obj, &serialize, &exclude_dead_locals, &exclude_immutables)) { + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|OOppp", kwlist, + &frame, &exclude_locals, &sizehint_obj, &serialize, &exclude_dead_locals, &exclude_immutables)) { return NULL; } - - options.serialize = (serialize != 0); - options.exclude_dead_locals = (exclude_dead_locals != 0); - options.exclude_immutables = (exclude_immutables != 0); + + options.populate(serialize, exclude_locals, exclude_dead_locals, exclude_immutables); if (!parse_sizehint(sizehint_obj, options.sizehint)) { return NULL; } @@ -810,7 +855,7 @@ static void init_pyinterpreterframe(sauerkraut::PyInterpreterFrame *interp_frame interp_frame->f_executable.bits = (uintptr_t)Py_NewRef(code.borrow()); if(frame_obj.f_executable.immutables_included()) { - interp_frame->f_funcobj = Py_NewRef(frame_obj.f_funcobj.value().borrow()); + utils::py::set_funcobj(interp_frame, Py_NewRef(frame_obj.f_funcobj.value().borrow())); if(NULL != frame_obj.f_globals) { interp_frame->f_globals = frame_obj.f_globals.borrow(); } else { @@ -819,10 +864,10 @@ static void init_pyinterpreterframe(sauerkraut::PyInterpreterFrame *interp_frame } else { auto invariants = sauerkraut_state->get_code_immutables(frame_obj); if(invariants) { - interp_frame->f_funcobj = Py_NewRef(std::get<0>(invariants.value()).borrow()); + utils::py::set_funcobj(interp_frame, Py_NewRef(std::get<0>(invariants.value()).borrow())); interp_frame->f_globals = Py_NewRef(std::get<2>(invariants.value()).borrow()); } else { - interp_frame->f_funcobj = NULL; + utils::py::set_funcobj(interp_frame, NULL); interp_frame->f_globals = NULL; } } @@ -851,17 +896,14 @@ static void init_pyinterpreterframe(sauerkraut::PyInterpreterFrame *interp_frame for(size_t i = localsplus.size(); i < (size_t)code->co_nlocalsplus; i++) { interp_frame->localsplus[i].bits = 0; } - interp_frame->instr_ptr = (sauerkraut::PyBitcodeInstruction*) + interp_frame->instr_ptr = (sauerkraut::PyBitcodeInstruction*) (utils::py::get_code_adaptive(code) + frame_obj.instr_offset/2);//utils::py::get_offset_for_skipping_call(); interp_frame->return_offset = frame_obj.return_offset; - #if SAUERKRAUT_PY314 - interp_frame->stackpointer = frame_stack_base + stack.size(); - #elif SAUERKRAUT_PY313 - interp_frame->stacktop = code->co_nlocalsplus + stack.size(); - #endif + utils::py::set_stack_position(interp_frame, code->co_nlocalsplus, stack.size()); // TODO: Check what happens when we make the owner the frame object instead of the thread. // Might allow us to skip a copy when calling this frame interp_frame->owner = frame_obj.owner; + utils::py::init_frame_visited(interp_frame); // Weak ref to avoid circular reference with capsule interp_frame->frame_obj = *frame; frame->f_frame = interp_frame; @@ -901,7 +943,6 @@ static PyObject *_deserialize_frame(PyObject *bytes, bool inplace=false) { auto serframe = pyframe_buffer::GetPyFrame(data); auto deserframe = frame_serdes.deserialize(serframe); - // FRAME_OWNED_BY_THREAD assert(deserframe.f_frame.owner == 0); pycode_strongref code; if(deserframe.f_frame.f_executable.immutables_included()) { @@ -910,6 +951,10 @@ static PyObject *_deserialize_frame(PyObject *bytes, bool inplace=false) { auto cached_invariants = sauerkraut_state->get_code_immutables(deserframe); if(cached_invariants) { code = make_strongref((PyCodeObject*)std::get<1>(cached_invariants.value()).borrow()); + } else { + PyErr_SetString(PyExc_RuntimeError, + "Cannot deserialize frame: immutables were excluded but cache lookup failed."); + return NULL; } } @@ -932,14 +977,49 @@ static PyObject *_deserialize_frame(PyObject *bytes, bool inplace=false) { static PyObject *run_frame_direct(py_weakref frame) { PyThreadState *tstate = PyThreadState_Get(); pycode_strongref code = pycode_strongref::steal(PyFrame_GetCode(*frame)); - PyFrameObject *to_run = push_frame_for_running(tstate, frame->f_frame, code.borrow()); - if (to_run == NULL) { + _PyInterpreterFrame *heap_frame = frame->f_frame; + + // Allocate a new frame on the eval stack + _PyInterpreterFrame *stack_frame = utils::py::ThreadState_PushFrame(tstate, code->co_framesize); + if (stack_frame == NULL) { PySys_WriteStderr(": failed to create frame on the framestack\n"); return NULL; } - PyObject *res = run_and_cleanup_frame(to_run); - return res; + // Copy all fields from the heap frame to the stack frame + stack_frame->f_executable.bits = heap_frame->f_executable.bits; + stack_frame->previous = NULL; // No previous frame - we're the root + utils::py::set_funcobj(stack_frame, utils::py::get_funcobj(heap_frame)); + stack_frame->f_globals = heap_frame->f_globals; + stack_frame->f_builtins = heap_frame->f_builtins; + stack_frame->f_locals = heap_frame->f_locals; + stack_frame->frame_obj = *frame; + stack_frame->instr_ptr = heap_frame->instr_ptr; + stack_frame->return_offset = heap_frame->return_offset; + stack_frame->owner = heap_frame->owner; + utils::py::init_frame_visited(stack_frame); + + // Copy localsplus (shallow copy - no refcount changes needed as we're moving refs) + int nlocalsplus = code->co_nlocalsplus; + memcpy(stack_frame->localsplus, heap_frame->localsplus, nlocalsplus * sizeof(_PyStackRef)); + + // Copy stack + int stack_depth = utils::py::get_current_stack_depth(heap_frame); + _PyStackRef *heap_stack = utils::py::get_stack_base(heap_frame); + _PyStackRef *stack_stack = utils::py::get_stack_base(stack_frame); + memcpy(stack_stack, heap_stack, stack_depth * sizeof(_PyStackRef)); + + // Set stack pointer + utils::py::set_stack_position(stack_frame, nlocalsplus, stack_depth); + + // Update the frame object to point to the new stack frame + frame->f_frame = stack_frame; + + // Skip past the CALL instruction + prepare_frame_for_execution(frame); + + PyObject *res = run_and_cleanup_frame(*frame); + return res; } @@ -961,9 +1041,12 @@ static PyObject *deserialize_frame(PyObject *self, PyObject *args, PyObject *kwa if (run) { PyFrameObject *frame = (PyFrameObject*)deser_result; if (!handle_replace_locals(replace_locals, frame)) { + Py_DECREF(frame); return NULL; } - return run_and_cleanup_frame(frame); + PyObject *result = run_and_cleanup_frame(frame); + Py_DECREF(frame); + return result; } else { // replace_locals should be applied via run_frame return deser_result; @@ -1038,21 +1121,20 @@ static PyObject *serialize_frame(PyObject *self, PyObject *args, PyObject *kwarg static PyObject *copy_frame_from_greenlet(PyObject *self, PyObject *args, PyObject *kwargs) { PyObject *greenlet = NULL; SerializationOptions options; - int exclude_immutables_int = 0; - + static char *kwlist[] = {"greenlet", "exclude_locals", "sizehint", "serialize", "exclude_dead_locals", "exclude_immutables", NULL}; int serialize = 0; PyObject* sizehint_obj = NULL; + PyObject* exclude_locals = NULL; int exclude_dead_locals = 1; - - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|OOppp", kwlist, - &greenlet, &options.exclude_locals, - &sizehint_obj, &serialize, &exclude_dead_locals, &exclude_immutables_int)) { + int exclude_immutables = 0; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|OOppp", kwlist, + &greenlet, &exclude_locals, + &sizehint_obj, &serialize, &exclude_dead_locals, &exclude_immutables)) { return NULL; } - options.serialize = (serialize != 0); - options.exclude_dead_locals = (exclude_dead_locals != 0); - options.exclude_immutables = (exclude_immutables_int != 0); + options.populate(serialize, exclude_locals, exclude_dead_locals, exclude_immutables); if (!parse_sizehint(sizehint_obj, options.sizehint)) { return NULL; } @@ -1096,7 +1178,10 @@ static PyMethodDef MyMethods[] = { }; static void sauerkraut_free(void *m) { - delete sauerkraut_state; + if (sauerkraut_state) { + delete sauerkraut_state; + sauerkraut_state = nullptr; + } } static struct PyModuleDef sauerkraut_mod = { @@ -1113,7 +1198,13 @@ static struct PyModuleDef sauerkraut_mod = { PyMODINIT_FUNC PyInit__sauerkraut(void) { sauerkraut_state = new sauerkraut_modulestate(); + if (!sauerkraut_state->init()) { + delete sauerkraut_state; + sauerkraut_state = NULL; + return NULL; + } greenlet::init_greenlet(); + return PyModule_Create(&sauerkraut_mod); } diff --git a/sauerkraut/serdes/include/serdes.h b/sauerkraut/serdes/include/serdes.h index 3f73337..e06b688 100644 --- a/sauerkraut/serdes/include/serdes.h +++ b/sauerkraut/serdes/include/serdes.h @@ -410,9 +410,15 @@ namespace serdes { _PyStackRef *stack_base = utils::py::get_stack_base(&obj); for(size_t i = 0; i < (size_t) stack_depth; i++) { - PyObject *stack_obj = (PyObject*) (stack_base[i].bits); - auto stack_obj_ser = po_serializer.serialize(builder, stack_obj); + auto stack_obj = utils::py::stackref_to_object_for_serialization(stack_base[i]); + if (stack_obj.obj == NULL) { + continue; + } + auto stack_obj_ser = po_serializer.serialize(builder, stack_obj.obj); stack.push_back(stack_obj_ser); + if (stack_obj.owned) { + Py_DECREF(stack_obj.obj); + } } auto stack_offset = builder.CreateVector(stack); @@ -425,12 +431,13 @@ namespace serdes { auto n_locals = utils::py::get_code_nlocals((PyCodeObject*)obj.f_executable.bits); auto exclude_local_bitmask = ser_args.exclude_locals.value_or(std::vector(n_locals, false)); std::vector localsplus; - + std::vector uint8_bitmask; uint8_bitmask.reserve(n_locals); int non_excluded_count = 0; for (int i = 0; i < n_locals; i++) { - if((PyObject*)obj.localsplus[i].bits == NULL || exclude_local_bitmask[i]) { + _PyStackRef local_ref = obj.localsplus[i]; + if(utils::py::stackref_is_null(local_ref) || exclude_local_bitmask[i]) { // a local can be NULL if it has not been initialized for the first time uint8_bitmask.push_back(1); } else { @@ -438,25 +445,31 @@ namespace serdes { non_excluded_count++; } } - + localsplus.reserve(non_excluded_count); // Only serialize non-excluded locals for(int i = 0; i < n_locals; i++) { auto local = obj.localsplus[i]; - PyObject *local_pyobj = (PyObject*)local.bits; + auto local_pyobj = utils::py::stackref_to_object_for_serialization(local); - if(NULL == local_pyobj || exclude_local_bitmask[i]) { + if(NULL == local_pyobj.obj || exclude_local_bitmask[i]) { + if (local_pyobj.owned) { + Py_DECREF(local_pyobj.obj); + } continue; } - auto local_ser = po_serializer.serialize(builder, local_pyobj); + auto local_ser = po_serializer.serialize(builder, local_pyobj.obj); localsplus.push_back(local_ser); + if (local_pyobj.owned) { + Py_DECREF(local_pyobj.obj); + } } - + auto localsplus_offset = builder.CreateVector(localsplus); auto bitmask_offset = builder.CreateVector(uint8_bitmask); - + return std::make_pair(localsplus_offset, bitmask_offset); } @@ -468,12 +481,17 @@ namespace serdes { template offsets::PyInterpreterFrameOffset serialize(Builder &builder, sauerkraut::PyInterpreterFrame &obj, int stack_depth, serdes::SerializationArgs& ser_args) { offsets::PyCodeObjectOffset f_executable_ser; - offsets::PyObjectOffset f_func_obj_ser; - offsets::PyObjectOffset f_globals_ser; + offsets::PyObjectOffset f_func_obj_ser = 0; + offsets::PyObjectOffset f_globals_ser = 0; + bool has_f_funcobj = false; f_executable_ser = code_serializer.serialize(builder, (PyCodeObject*)obj.f_executable.bits, ser_args); if(!ser_args.exclude_immutables) { - f_func_obj_ser = po_serializer.serialize(builder, obj.f_funcobj); + PyObject *func_obj = utils::py::get_funcobj(&obj); + if (func_obj != NULL) { + f_func_obj_ser = po_serializer.serialize(builder, func_obj); + has_f_funcobj = true; + } f_globals_ser = po_serializer.serialize_dill(builder, obj.f_globals); } @@ -489,7 +507,9 @@ namespace serdes { frame_builder.add_f_locals(f_locals_ser.value()); } if(!ser_args.exclude_immutables) { - frame_builder.add_f_funcobj(f_func_obj_ser); + if (has_f_funcobj) { + frame_builder.add_f_funcobj(f_func_obj_ser); + } frame_builder.add_f_globals(f_globals_ser); } @@ -525,10 +545,10 @@ namespace serdes { auto localsplus = obj->locals_plus(); auto exclusion_bitmask = obj->locals_exclusion_bitmask(); - + int total_locals = exclusion_bitmask->size(); deser.localsplus.reserve(total_locals); - + int localsplus_idx = 0; for(int i = 0; i < total_locals; i++) { if(exclusion_bitmask->Get(i) != 0) { @@ -635,4 +655,4 @@ namespace serdes { } -#endif // SERDES_HH_INCLUDED \ No newline at end of file +#endif // SERDES_HH_INCLUDED diff --git a/setup.py b/setup.py index 7628960..2c87d3e 100644 --- a/setup.py +++ b/setup.py @@ -143,15 +143,18 @@ def build_extension(self, ext): # MODULE (.so bundle), helper libs as SHARED (.dylib) main_ext = ".so" helper_ext = ".dylib" - target_ext = ".so" + main_target_ext = ".so" + helper_target_ext = ".dylib" elif sys.platform == "win32": main_ext = ".pyd" helper_ext = ".pyd" - target_ext = ".pyd" + main_target_ext = ".pyd" + helper_target_ext = ".pyd" else: main_ext = ".so" helper_ext = ".so" - target_ext = ".so" + main_target_ext = ".so" + helper_target_ext = ".so" package_dir = os.path.join(install_dir, "sauerkraut") os.makedirs(package_dir, exist_ok=True) @@ -159,14 +162,14 @@ def build_extension(self, ext): # Copy main module main_src = os.path.join(build_temp, f"_sauerkraut{main_ext}") if os.path.exists(main_src): - dst = os.path.join(package_dir, f"_sauerkraut{target_ext}") + dst = os.path.join(package_dir, f"_sauerkraut{main_target_ext}") self.copy_file(main_src, dst) # Copy helper libraries for lib_name in ["greenlet_compat", "serdes"]: src = os.path.join(build_temp, f"{lib_name}{helper_ext}") if os.path.exists(src): - target_name = f"{lib_name}{target_ext}" + target_name = f"{lib_name}{helper_target_ext}" dst = os.path.join(package_dir, target_name) self.copy_file(src, dst) @@ -239,5 +242,5 @@ def get_outputs(self): }, zip_safe=False, python_requires=">=3.13", - install_requires=["bytecode", "numpy"], + install_requires=["bytecode @ git+https://github.com/ZwFink/bytecode.git", "numpy"], )