diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..a59e571 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,64 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +env: + CARGO_TERM_COLOR: always + +jobs: + test: + strategy: + matrix: + include: + - os: ubuntu-latest + target: x86_64-unknown-linux-gnu + - os: macos-latest + target: aarch64-apple-darwin + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + targets: ${{ matrix.target }} + - name: Build + run: cargo build --target ${{ matrix.target }} + - name: Test + run: cargo test + + cross-build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + targets: aarch64-unknown-linux-gnu + - name: Install cross + run: cargo install cross --git https://github.com/cross-rs/cross + - name: Build + run: cross build --target aarch64-unknown-linux-gnu + + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: clippy, rustfmt + - name: Format check + run: cargo fmt --check + - name: Clippy + run: cargo clippy -- -D warnings + + msrv: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@master + with: + toolchain: "1.73" + - name: Check MSRV (no default features) + run: cargo check --no-default-features diff --git a/Cargo.toml b/Cargo.toml index 714081d..9452e91 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,10 +2,20 @@ name = "zk-alloc" version = "0.1.0" edition = "2021" -description = "A memory allocator designed for ZK proving workloads" +rust-version = "1.73" +description = "Bump+reset arena allocator for ZK proving workloads" license = "MIT OR Apache-2.0" +repository = "https://github.com/Barnadrot/zk-alloc" +readme = "README.md" +keywords = ["allocator", "arena", "zk", "proving", "memory"] +categories = ["memory-management", "no-std"] + +[features] +default = ["rayon-flush"] +rayon-flush = ["rayon"] [dependencies] +rayon = { version = "1.10", optional = true } [target.'cfg(not(all(target_os = "linux", target_arch = "x86_64")))'.dependencies] libc = "0.2" diff --git a/src/lib.rs b/src/lib.rs index a25d17a..db2903d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,8 +17,8 @@ use std::alloc::{GlobalAlloc, Layout}; use std::cell::Cell; -use std::sync::Once; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Once; mod syscall; @@ -103,8 +103,28 @@ pub fn begin_phase() { /// Deactivates the arena. New allocations go to the system allocator; existing arena /// pointers stay valid until the next `begin_phase()` resets the slabs. +/// +/// With the `rayon-flush` feature (default), this also drains rayon's internal +/// queues to release any crossbeam-deque blocks allocated during the phase. pub fn end_phase() { ARENA_ACTIVE.store(false, Ordering::Release); + #[cfg(feature = "rayon-flush")] + flush_rayon(); +} + +/// Drains rayon's crossbeam-deque injector to release blocks allocated during +/// the active phase. Without this, `begin_phase()` would recycle memory that +/// rayon's injector still references, causing silent corruption. +/// +/// Pushes `FLUSH_JOBS` no-op joins. Each consumes one injector slot; once a +/// block's last slot is consumed, crossbeam deallocates it. The fresh tail +/// block lands in the system allocator (arena is already inactive). +#[cfg(feature = "rayon-flush")] +fn flush_rayon() { + const FLUSH_JOBS: usize = 256; + for _ in 0..FLUSH_JOBS { + rayon::join(|| {}, || {}); + } } /// Returns (overflow_count, overflow_bytes) — allocations that fell through to System diff --git a/src/syscall.rs b/src/syscall.rs index f600fe0..f676b2a 100644 --- a/src/syscall.rs +++ b/src/syscall.rs @@ -18,7 +18,15 @@ mod imp { pub const MADV_NOHUGEPAGE: usize = 15; #[inline] - unsafe fn syscall6(nr: usize, a1: usize, a2: usize, a3: usize, a4: usize, a5: usize, a6: usize) -> isize { + unsafe fn syscall6( + nr: usize, + a1: usize, + a2: usize, + a3: usize, + a4: usize, + a5: usize, + a6: usize, + ) -> isize { let ret: isize; unsafe { std::arch::asm!( @@ -60,8 +68,22 @@ mod imp { #[inline] pub unsafe fn mmap_anonymous(size: usize) -> *mut u8 { let flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE; - let ret = unsafe { syscall6(SYS_MMAP, 0, size, PROT_READ | PROT_WRITE, flags, usize::MAX, 0) }; - if ret < 0 { ptr::null_mut() } else { ret as *mut u8 } + let ret = unsafe { + syscall6( + SYS_MMAP, + 0, + size, + PROT_READ | PROT_WRITE, + flags, + usize::MAX, + 0, + ) + }; + if ret < 0 { + ptr::null_mut() + } else { + ret as *mut u8 + } } #[inline] @@ -97,4 +119,4 @@ mod imp { } } -pub use imp::{MADV_NOHUGEPAGE, madvise, mmap_anonymous}; +pub use imp::{madvise, mmap_anonymous, MADV_NOHUGEPAGE};