Skip to content

Commit

Permalink
fix: use rayon::current_num_threads() for multithreaded zstd encode
Browse files Browse the repository at this point in the history
  • Loading branch information
b-inary committed Sep 30, 2023
1 parent 6816dfd commit ff7a40a
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 10 deletions.
9 changes: 5 additions & 4 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@ license = "AGPL-3.0-or-later"

[dependencies]
bincode = { version = "2.0.0-rc.3", optional = true }
once_cell = "1.17.1"
rayon = { version = "1.7.0", optional = true }
regex = "1.8.1"
zstd = { version = "0.12.3", optional = true, features = ["zstdmt"] }
once_cell = "1.18.0"
rayon = { version = "1.8.0", optional = true }
regex = "1.9.5"
zstd = { version = "0.12.4", optional = true, default-features = false }

[features]
default = ["bincode", "rayon"]
custom-alloc = []
rayon = ["dep:rayon", "zstd?/zstdmt"]
10 changes: 4 additions & 6 deletions src/file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ use bincode::{Decode, Encode};
use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::Path;
use std::thread::available_parallelism;

const MAGIC: u32 = 0x09f15790;
const VERSION: u8 = 1;
Expand Down Expand Up @@ -88,12 +87,11 @@ pub fn save_data_into_std_write<T: FileData, W: Write>(
if let Some(compression_level) = compression_level {
let mut zstd_encoder = zstd::stream::Encoder::new(writer, compression_level)
.map_err(|e| format!("Failed to create zstd encoder: {}", e))?;
let workers = available_parallelism()
.map_err(|e| format!("Failed to get cpu cores: {}", e))?
.get() as u32;

#[cfg(feature = "rayon")]
zstd_encoder
.multithread(workers)
.map_err(|e| format!("Failed to enable multithreaded compression: {}", e))?;
.multithread(rayon::current_num_threads() as u32)
.map_err(|e| format!("Failed to enable multithreaded zstd encoder: {}", e))?;

encode_into_std_write(data, &mut zstd_encoder, "Failed to write data")?;
zstd_encoder
Expand Down

0 comments on commit ff7a40a

Please sign in to comment.