Skip to content

Commit

Permalink
release the GIL in run_block_generator() and run_block_generator2(), …
Browse files Browse the repository at this point in the history
…allowing these to be run in a ThreadPoolExecutor (#712)
  • Loading branch information
arvidn authored Sep 17, 2024
1 parent b1d463a commit adfc667
Showing 1 changed file with 48 additions and 38 deletions.
86 changes: 48 additions & 38 deletions wheel/src/run_generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ pub fn py_to_slice<'a>(buf: PyBuffer<u8>) -> &'a [u8] {

#[pyfunction]
pub fn run_block_generator<'a>(
_py: Python<'a>,
py: Python<'a>,
program: PyBuffer<u8>,
block_refs: &Bound<'_, PyList>,
max_cost: Cost,
Expand All @@ -29,37 +29,42 @@ pub fn run_block_generator<'a>(
) -> (Option<u32>, Option<OwnedSpendBundleConditions>) {
let mut allocator = make_allocator(flags);

let refs = block_refs.into_iter().map(|b| {
let buf = b
.extract::<PyBuffer<u8>>()
.expect("block_refs should be a list of buffers");
py_to_slice::<'a>(buf)
});
let refs = block_refs
.into_iter()
.map(|b| {
let buf = b
.extract::<PyBuffer<u8>>()
.expect("block_refs should be a list of buffers");
py_to_slice::<'a>(buf)
})
.collect::<Vec<&'a [u8]>>();
let program = py_to_slice::<'a>(program);
let run_block = if (flags & ANALYZE_SPENDS) == 0 {
native_run_block_generator::<_, EmptyVisitor, _>
} else {
native_run_block_generator::<_, MempoolVisitor, _>
};

match run_block(&mut allocator, program, refs, max_cost, flags, constants) {
Ok(spend_bundle_conds) => (
None,
Some(OwnedSpendBundleConditions::from(
&allocator,
spend_bundle_conds,
)),
),
Err(ValidationErr(_, error_code)) => {
// a validation error occurred
(Some(error_code.into()), None)
py.allow_threads(|| {
match run_block(&mut allocator, program, refs, max_cost, flags, constants) {
Ok(spend_bundle_conds) => (
None,
Some(OwnedSpendBundleConditions::from(
&allocator,
spend_bundle_conds,
)),
),
Err(ValidationErr(_, error_code)) => {
// a validation error occurred
(Some(error_code.into()), None)
}
}
}
})
}

#[pyfunction]
pub fn run_block_generator2<'a>(
_py: Python<'a>,
py: Python<'a>,
program: PyBuffer<u8>,
block_refs: &Bound<'_, PyList>,
max_cost: Cost,
Expand All @@ -68,12 +73,15 @@ pub fn run_block_generator2<'a>(
) -> (Option<u32>, Option<OwnedSpendBundleConditions>) {
let mut allocator = make_allocator(flags);

let refs = block_refs.into_iter().map(|b| {
let buf = b
.extract::<PyBuffer<u8>>()
.expect("block_refs must be list of buffers");
py_to_slice::<'a>(buf)
});
let refs = block_refs
.into_iter()
.map(|b| {
let buf = b
.extract::<PyBuffer<u8>>()
.expect("block_refs must be list of buffers");
py_to_slice::<'a>(buf)
})
.collect::<Vec<&'a [u8]>>();

let program = py_to_slice::<'a>(program);
let run_block = if (flags & ANALYZE_SPENDS) == 0 {
Expand All @@ -82,17 +90,19 @@ pub fn run_block_generator2<'a>(
native_run_block_generator2::<_, MempoolVisitor, _>
};

match run_block(&mut allocator, program, refs, max_cost, flags, constants) {
Ok(spend_bundle_conds) => (
None,
Some(OwnedSpendBundleConditions::from(
&allocator,
spend_bundle_conds,
)),
),
Err(ValidationErr(_, error_code)) => {
// a validation error occurred
(Some(error_code.into()), None)
py.allow_threads(|| {
match run_block(&mut allocator, program, refs, max_cost, flags, constants) {
Ok(spend_bundle_conds) => (
None,
Some(OwnedSpendBundleConditions::from(
&allocator,
spend_bundle_conds,
)),
),
Err(ValidationErr(_, error_code)) => {
// a validation error occurred
(Some(error_code.into()), None)
}
}
}
})
}

0 comments on commit adfc667

Please sign in to comment.