Speed up MAS token introspection #18357
Conversation
| // TODO: Is it safe to assert unwind safety here? I think so, as we | ||
| // don't use anything that could be tainted by the panic afterwards. | ||
| // Note that `.spawn(..)` asserts unwind safety on the future too. | ||
| let res = AssertUnwindSafe(fut).catch_unwind().await; |
There was a problem hiding this comment.
alternatively, spawn the future on the runtime, await on the handle, and it will give you an Err with the panic in case it panics
There was a problem hiding this comment.
We could, though then you end up spawning two tasks per function, rather than one. Probably not a huge deal, but feels a bit bleurgh
There was a problem hiding this comment.
Spawning is cheap, let's do that instead please. Also we'll need to spawn a separate task anyway if we want to properly support cancel
There was a problem hiding this comment.
Can you expand on why you want to use new tasks please? I don't see the benefit of spawning a new task to just wait on it, semantically you end up with a bunch of tasks with different IDs all for the same work. In future, if we wanted to start tracking tasks and e.g. their resource usage then using multiple tasks makes that more complicated.
I also don't think we need a separate task for cancellation necessarily. You can change this line to do a select on both fut and the cancellation future.
rust/src/http_client.rs
Outdated
| static ref DEFERRED_CLASS: PyObject = { | ||
| Python::with_gil(|py| { | ||
| py.import("twisted.internet.defer") | ||
| .expect("module 'twisted.internet.defer' should be importable") |
There was a problem hiding this comment.
I'm not a fan of panicking like that, not sure what will happen if it's the case
There was a problem hiding this comment.
This will cause a panic the first and subsequent derefs. Given that this shouldn't ever fail I prefer having the initialisation closer to the definition for clarities sake.
I've added an explicit call to these functions in the init of HttpClient so that if it ever did fail, it'd fail at startup.
This is because openssl cannot be used in a manylinux context due to lack of stable ABI.
rust/src/http_client.rs
Outdated
| // Make sure we fail early if we can't build the lazy statics. | ||
| LazyLock::force(&RUNTIME); | ||
| LazyLock::force(&DEFERRED_CLASS); |
There was a problem hiding this comment.
Could be done in the module initialisation?
There was a problem hiding this comment.
Annoyingly, we can't import twisted reactor at this stage as it happens too early. See comment I left in HttpClient::new
| let mut stream = response.bytes_stream(); | ||
| let mut buffer = Vec::new(); | ||
| while let Some(chunk) = stream.try_next().await.context("reading body")? { | ||
| if buffer.len() + chunk.len() > response_limit { | ||
| Err(anyhow::anyhow!("Response size too large"))?; | ||
| } | ||
|
|
||
| buffer.extend_from_slice(&chunk); | ||
| } |
There was a problem hiding this comment.
I believe you can achieve the same with http_body_util::Limited; reqwest::Response implements Into<Body>
There was a problem hiding this comment.
Yeah, originally tried that but it messed up the errors (one of the exceptions stops implementing std Error). Given how straight forwards this is it felt easier than faffing with error types.
| // TODO: Is it safe to assert unwind safety here? I think so, as we | ||
| // don't use anything that could be tainted by the panic afterwards. | ||
| // Note that `.spawn(..)` asserts unwind safety on the future too. | ||
| let res = AssertUnwindSafe(fut).catch_unwind().await; |
There was a problem hiding this comment.
Spawning is cheap, let's do that instead please. Also we'll need to spawn a separate task anyway if we want to properly support cancel
| /// The tokio runtime that we're using to run async Rust libs. | ||
| static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| { | ||
| tokio::runtime::Builder::new_multi_thread() | ||
| .worker_threads(4) |
There was a problem hiding this comment.
We'll likely want to have that configurable at some point, but this is probably a sane default.
We do this by shoving it into Rust. We believe our python http client is a bit slow.
Also bumps minimum rust version to 1.81.0, released last September (over six months ago)