Add back custom allocators

.github/workflows/build.yml

@@ -3,60 +3,83 @@ name: build
 on: workflow_dispatch
 
 env:
-  PY_ALL: 3.9 3.10 3.11 3.12 3.13 pypy3.9 pypy3.10
+  PY_ALL: 3.9 3.10 3.11 3.12 3.13 3.13t pypy3.9 pypy3.10
 
 jobs:
   wheels:
-    name: wheel ${{ matrix.platform || matrix.os }}(${{ matrix.target }}) - ${{ matrix.manylinux || 'auto' }}
+    name: wheel ${{ matrix.platform || matrix.os }}(${{ matrix.target }}) - ${{ matrix.manylinux || 'auto' }} - ${{ matrix.allocator || 'auto' }}
     strategy:
       fail-fast: false
       matrix:
         os: [ubuntu, macos, windows]
         target: [x86_64, aarch64]
         manylinux: [auto]
+        allocator: [auto]
         include:
           - os: ubuntu
             platform: linux
-          # FIXME: for some unknown reasons PGO builds are failing on Linux
-          # - os: ubuntu
-          #   platform: linux
-          #   target: x86_64
-          #   manylinux: auto
-          #   interpreter: pypy3.9 pypy3.10
-          - os: macos
             target: x86_64
-            interpreter: pypy3.9 pypy3.10
-          - os: macos
+            interpreter: 3.9 3.10 3.11 3.12 3.13 pypy3.9 pypy3.10
+          - os: ubuntu
+            platform: linux
+            target: x86_64
+            interpreter: 3.13t
+            allocator: mimalloc
+          - os: ubuntu
+            platform: linux
             target: aarch64
-            interpreter: pypy3.9 pypy3.10
+            container: messense/manylinux_2_24-cross:aarch64
+            interpreter: 3.9 3.10 3.11 3.12 3.13 pypy3.9 pypy3.10
           - os: ubuntu
             platform: linux
             target: aarch64
             container: messense/manylinux_2_24-cross:aarch64
+            interpreter: 3.13t
+            allocator: mimalloc
           - os: ubuntu
             platform: linux
             target: x86_64
             manylinux: musllinux_1_1
+            allocator: mimalloc
           - os: ubuntu
             platform: linux
             target: aarch64
             manylinux: musllinux_1_1
+            allocator: mimalloc
+          - os: macos
+            target: x86_64
+            interpreter: 3.9 3.10 3.11 3.12 3.13 pypy3.9 pypy3.10
+          - os: macos
+            target: x86_64
+            interpreter: 3.13t
+            allocator: mimalloc
+          - os: macos
+            target: aarch64
+            interpreter: 3.9 3.10 3.11 3.12 3.13 pypy3.9 pypy3.10
+          - os: macos
+            target: aarch64
+            interpreter: 3.13t
+            allocator: mimalloc
           - os: windows
             target: x86_64
-            interpreter: pypy3.9 pypy3.10
+            interpreter: 3.9 3.10 3.11 3.12 3.13 pypy3.9 pypy3.10
+          - os: windows
+            target: x86_64
+            interpreter: 3.13t
+            allocator: mimalloc
         exclude:
           - os: windows
             target: aarch64
 
     runs-on: ${{ matrix.os }}-latest
     steps:
       - uses: actions/checkout@v4
-      # NOTE: we don't use jemalloc anymore
-      # - name: Set jemalloc for aarch64 Linux
-      #   if: matrix.target == 'aarch64' && matrix.os == 'ubuntu'
-      #   run: |
-      #     echo "JEMALLOC_SYS_WITH_LG_PAGE=16" >> $GITHUB_ENV
+      - name: Set jemalloc for aarch64 Linux
+        if: matrix.target == 'aarch64' && matrix.os == 'ubuntu'
+        run: |
+          echo "JEMALLOC_SYS_WITH_LG_PAGE=16" >> $GITHUB_ENV
       - uses: pyo3/maturin-action@v1
+        if: matrix.allocator != 'mimalloc'
         with:
           rust-toolchain: stable
           command: build
@@ -65,77 +88,18 @@ jobs:
           manylinux: ${{ matrix.manylinux || 'auto' }}
           container: ${{ matrix.container }}
           docker-options: -e CI
-      - name: Upload wheels
-        uses: actions/upload-artifact@v4
-        with:
-          name: dist-${{ matrix.platform || matrix.os }}-${{ matrix.target }}-${{ matrix.manylinux || 'auto' }}
-          path: dist
-
-  wheels-pgo:
-    name: pgo-wheel ${{ matrix.platform || matrix.os }} (${{ matrix.interpreter }})
-    strategy:
-      fail-fast: false
-      matrix:
-        # FIXME: for some unknown reasons PGO builds are failing on Linux
-        # os: [ubuntu-latest, macos-13, macos-14, windows-latest]
-        os: [macos-13, macos-14, windows-latest]
-        manylinux: [auto]
-        interpreter: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-        include:
-          # - os: ubuntu-latest
-          #   platform: linux
-          - os: windows-latest
-            platform: windows
-
-    runs-on: ${{ matrix.os }}
-    env:
-      UV_PYTHON: ${{ matrix.interpreter }}
-    steps:
-      - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - uses: dtolnay/rust-toolchain@stable
-        with:
-          components: llvm-tools
-      - name: prepare profiling directory
-        shell: bash
-        run: mkdir -p ${{ github.workspace }}/profdata
-      - name: Build initial wheel
-        uses: PyO3/maturin-action@v1
+      - uses: pyo3/maturin-action@v1
+        if: matrix.allocator == 'mimalloc'
         with:
           rust-toolchain: stable
           command: build
-          args: --release --out pgo_wheel --interpreter ${{ matrix.interpreter }}
-          manylinux: ${{ matrix.manylinux || 'auto' }}
-          docker-options: -e CI
-        env:
-          RUSTFLAGS: "-Cprofile-generate=${{ github.workspace }}/profdata"
-      - run: |
-          RUST_HOST=$(rustc -Vv | grep host | cut -d ' ' -f 2) rustup run stable bash -c 'echo LLVM_PROFDATA=$RUSTUP_HOME/toolchains/$RUSTUP_TOOLCHAIN/lib/rustlib/$RUST_HOST/bin/llvm-profdata >> "$GITHUB_ENV"'
-        shell: bash
-      - name: Generate PGO data
-        shell: bash
-        run: |
-          uv python install ${{ env.UV_PYTHON }}
-          uv venv .venv
-          uv pip install -r .github/workflows/_pgo_deps.txt
-          uv pip install granian --no-index --no-deps --find-links pgo_wheel --force-reinstall
-          PGO_RUN=y LLVM_PROFILE_FILE=${{ github.workspace }}/profdata/granian_%m_%p.profraw uv run --no-sync pytest tests
-      - name: merge PGO data
-        run: ${{ env.LLVM_PROFDATA }} merge --failure-mode=all -o ${{ github.workspace }}/merged.profdata ${{ github.workspace }}/profdata
-      - name: Build PGO wheel
-        uses: PyO3/maturin-action@v1
-        with:
-          command: build
-          args: --release --out dist --interpreter ${{ matrix.interpreter }}
+          args: --release --out dist --features mimalloc --interpreter ${{ matrix.interpreter || env.PY_ALL }}
+          target: ${{ matrix.target }}
           manylinux: ${{ matrix.manylinux || 'auto' }}
-          rust-toolchain: stable
+          container: ${{ matrix.container }}
           docker-options: -e CI
-        env:
-          RUSTFLAGS: "-Cprofile-use=${{ github.workspace }}/merged.profdata"
       - name: Upload wheels
         uses: actions/upload-artifact@v4
         with:
-          name: dist-pgo-${{ matrix.platform || matrix.os }}-${{ matrix.interpreter }}
+          name: dist-${{ matrix.platform || matrix.os }}-${{ matrix.target }}-${{ matrix.manylinux || 'auto' }}-${{ matrix.allocator || 'auto' }}
           path: dist

.github/workflows/lint.yml

@@ -7,14 +7,12 @@ on:
       - master
 
 env:
-  PYTHON_VERSION: 3.12
+  UV_PYTHON: 3.12
 
 jobs:
   lint:
     runs-on: ubuntu-latest
 
-    env:
-      UV_PYTHON: ${{ matrix.python-version }}
     steps:
     - uses: actions/checkout@v4
     - uses: astral-sh/setup-uv@v5

Cargo.toml

@@ -40,6 +40,7 @@ hyper = { version = "=1.6", features = ["http1", "http2", "server"] }
 hyper-util = { version = "=0.1", features = ["server-auto", "tokio"] }
 itertools = "0.14"
 log = "0.4"
+mimalloc = { version = "0.1.43", default-features = false, features = ["local_dynamic_tls"], optional = true }
 pem = "=3.0"
 percent-encoding = "=2.3"
 pin-project = "1.1"
@@ -54,9 +55,15 @@ tokio-stream = "0.1"
 tokio-tungstenite = "=0.26"
 tokio-util = { version = "0.7", features = ["codec"] }
 
+[target.'cfg(not(any(target_env = "musl", target_os = "freebsd", target_os = "openbsd", target_os = "windows")))'.dependencies]
+tikv-jemallocator = { version = "0.6.0", default-features = false, features = ["disable_initial_exec_tls"] }
+
 [build-dependencies]
 pyo3-build-config = "=0.23"
 
+[features]
+mimalloc = ["dep:mimalloc"]
+
 [profile.release]
 codegen-units = 1
 debug = false

src/callbacks.rs

@@ -123,7 +123,7 @@ impl CallbackScheduler {
     }
 
     #[inline]
-    fn throw(&self, state: Arc<CallbackSchedulerState>, err: PyObject) {
+    fn throw(&self, _py: Python, state: Arc<CallbackSchedulerState>, err: PyObject) {
         let aiotask = self.aio_task.as_ptr();
 
         unsafe {
@@ -187,7 +187,7 @@ impl CallbackScheduler {
     }
 
     #[inline]
-    fn throw(&self, state: Arc<CallbackSchedulerState>, err: PyObject) {
+    fn throw(&self, py: Python, state: Arc<CallbackSchedulerState>, err: PyObject) {
         let aiotask = self.aio_task.as_ptr();
 
         unsafe {
@@ -311,7 +311,7 @@ impl CallbackSchedulerWaker {
                 .state
                 .sched
                 .get()
-                .throw(self.state.clone(), err.into_py_any(py).unwrap()),
+                .throw(py, self.state.clone(), err.into_py_any(py).unwrap()),
         }
     }
 }

src/lib.rs

@@ -1,3 +1,17 @@
+#[cfg(not(any(
+    target_env = "musl",
+    target_os = "freebsd",
+    target_os = "openbsd",
+    target_os = "windows",
+    feature = "mimalloc"
+)))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+#[cfg(feature = "mimalloc")]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 use pyo3::prelude::*;
 use std::sync::OnceLock;