Merge/replace with_${python_obj} context managers with use of allow_threads

[stuhood]

The with_$x context managers add a bunch of noise, and are much less central than calling allow_threads, which we should try to do in nearly all cases.

Not calling allow_threads before calling various methods of the Scheduler/Session (particularly anything to do with the Graph) can cause deadlocks... not to mention meaning that we're holding the GIL over native code that might not need it. For example: #11723

It's possible that these methods could be merged, replaced, or enhanced to also release the GIL when called. But a challenge is that each of the with_$x methods cannot call allow_threads, because each of them need the GIL to dereference the pointer.

---

One strawdesign that might work would be something like:

fn with_unwrapped<U, F, T>(py: Python, unwrappable: U, f: F) -> T
where
  U: Unwrappable,
  F: FnOnce(U::Output) -> T,
{
  let unwrapped = unwrappable.unwrapped();
  py.allow_threads(|| {
    f(unwrapped)
  })
}

We'd then implement the hyptothetical Unwrappable trait for tuples of 1, 2, 3-ish python objects (PyScheduler, PySession, etc) allowing you to make a single call to unwrap all of them at once. The single call to the method would allow it to call allow_threads, since you would not need to nest them.

[Eric-Arellano]

Will this need to be done all-or-nothing, or can be incremental?

I'd love to work on this one, I'm enjoying learning more about FFI and the GIL.

[stuhood]

Will this need to be done all-or-nothing, or can be incremental?

It could be incremental: we could add function(s) next to the with_$x methods, and move code to them.

I've updated the description to call out one of the challenges and a potential design.

[stuhood]

OH. 😮 .. Alternatively, if we're able to implement FromPyObject for &Scheduler and etc, I think that we could change the signature of e.g. garbage_collect_store to:

fn garbage_collect_store(py: Python, scheduler: &Scheduler, target_size_bytes: usize) -> PyUnitResult {
    py.allow_threads(|| {
      scheduler
        .core
        .store()
        .garbage_collect(target_size_bytes, store::ShrinkBehavior::Fast)
    })
    .map_err(|e| PyErr::new::<exc::Exception, _>(py, (e,)))
    .map(|()| None)
}

...and get rid of with_$x! Doesn't help with allow_threads, but.

EDIT: Filed dgrunwald/rust-cpython#260 about a potential way to remove allow_threads in some cases.

[Eric-Arellano]

@stuhood I'm confused on this one. The pyclasses already implement FromPyObject. I'm wanting to close this issue out as prework for PyO3, starting with this because it's entirely self-contained fwict:

pants/src/rust/engine/src/externs/interface.rs

Lines 830 to 841 in 4301cb4

	fn nailgun_server_await_shutdown(py: Python, nailgun_server_ptr: PyNailgunServer) -> PyUnitResult {
	with_nailgun_server(py, nailgun_server_ptr, |nailgun_server, executor| {
	let executor = executor.clone();
	let shutdown_result = if let Some(server) = nailgun_server.borrow_mut().take() {
	py.allow_threads(|| executor.block_on(server.shutdown()))
	} else {
	Ok(())
	};
	shutdown_result.map_err(|e| PyErr::new::<exc::Exception, _>(py, (e,)))?;
	Ok(None)
	})
	}

[stuhood]

The pyclasses already implement FromPyObject.

They implement FromPyObject for themselves, but not for their inner type: the change to the method that you linked would be to directly take a nailgun::Server or a reference to a nailgun::Server. That one is a bit more challenging because it has multiple fields.

fn nailgun_server_await_shutdown(py: Python, nailgun_server: &nailgun::Server) -> PyUnitResult {

[Eric-Arellano]

I think I'm going to try to stick to the current implementation while porting to PyO3. Once that lands, it should be easier to make this change.