Support async hooks. Fixes #4207. by nneonneo · Pull Request #5106 · mitmproxy/mitmproxy

nneonneo · 2022-02-02T07:54:20Z

Description

Add support for async hooks to AddonManager, enabling the proxy server to call hooks asynchronously. Uses inspect.isawaitable to ensure that existing sync hooks continue to work without any change. Fixes #4207.

The forking of invoke_addon and trigger is a little unfortunate, but both of these get called from sync contexts throughout mitmproxy (in particular, triggers are often called from command.command handlers, which are all sync). To ensure that we don't accidentally invoke an async handler in such cases, this patch adds a check to throw an exception if an async handler gets invoked synchronously.

This is a first stab at async hook support. There are more things that could be done here, and I'd love to hear if either of them makes sense for this PR, or left for future work:

Eliminate concurrent entirely, and possibly even the entire reply mechanism. This would be a significant, breaking API change. However, in almost every case, async hooks are better: async hooks are properly sequenced (even with respect to other sync hooks), races are essentially eliminated, and the behaviour is much more predictable. (It might be better to just deprecate concurrent now for future removal). Basically done: concurrent is kept, but made async. reply is gone.
Eliminate the sync versions of trigger and invoke_addon to allow async hooks in all contexts. This is trickier than it seems because of the number of places in which hooks are triggered in a synchronous context; a lot of code would probably need to be updated to be async. While this probably should happen eventually, I'm hoping this much smaller patch can serve 80% of use cases immediately.

There are also a few easier things that could be done:

Document which contexts will permit async hooks. Lifecycle hooks like start, done, etc. don't look like they should be async, for example, since they're often called synchronously.
Add tests, and/or modify existing tests to use async hooks instead.

Checklist

I have updated tests where applicable.
I have added an entry to the CHANGELOG.

mhils

Thanks for the very nice PR! 🍰

Eliminate concurrent entirely, and possibly even the entire reply mechanism. This would be a significant, breaking API change.

Yes, eliminating reply is something I definitely want to see happening. @concurrent is here to stay though, that's widely used by existing addons. It just needs to be changed to spawn a thread, and then await the completion of an asyncio.Event which is set by the thread on completion. This way users can continue to use requests or other blocking stuff in their hooks.

If we could tackle the reply-sectomy in this PR that would be great, if you give it a stab and notice it's too much work I won't insist. :)

Eliminate the sync versions of trigger and invoke_addon to allow async hooks in all contexts. This is trickier than it seems because of the number of places in which hooks are triggered in a synchronous context; a lot of code would probably need to be updated to be async. While this probably should happen eventually, I'm hoping this much smaller patch can serve 80% of use cases immediately.

Yes, 👍 to this. See my inline comments below.

Document which contexts will permit async hooks. Lifecycle hooks like start, done, etc. don't look like they should be async, for example, since they're often called synchronously.

My naive approach would be that we could just permit all hooks to be async. If you add asyncio.sleep(5) to the start hook that's your own problem really, on the other hand it can be useful to have I-know-what-I'm-doing awaits in there.

nneonneo · 2022-02-02T09:14:06Z

If we could tackle the reply-sectomy in this PR that would be great, if you give it a stab and notice it's too much work I won't insist. :)

I'd be happy to, but I did just realize one problem: as it currently stands, if a hook can't be made async (because it gets called synchronously somewhere), it can be made @concurrent and called synchronously. If we switch now to using asyncio internally, this particular workaround goes away. If this isn't a concern (e.g. if 8.0.0 is planned to ship with async everything), then I'll go ahead and make the change.

mhils · 2022-02-02T09:24:47Z

If we switch now to using asyncio internally, this particular workaround goes away.

Good catch. We don't use @concurrent internally, so this should only affect third-party hooks making use of @concurrent in unexpected hooks (or invoking them in non-standard ways), so I think that is tolerable. Let's not impede progress. :)

Prinzhorn · 2022-02-03T11:00:30Z

Thank you very much! I've been playing with it this morning and so far haven't run into any issues. It does consistently increase the throughput of my addon quite a bit (between 10% and 300%).

I've been using mostly https://http2.akamai.com/demo to test.

nneonneo · 2022-02-03T14:37:06Z

@mhils: I appended a patch series to this PR which eliminates the flow.reply system. There are a few semantic differences to be aware of, but these are hopefully not significant:

@concurrent produces an async hook that runs hooks in a thread pool executor, so any @concurrent hooks need to run in async context. @concurrent hooks now run in series with other hooks of the same type on a given flow (but run concurrently with other hooks on other flows).
Due to the way intercept now works, Flow.resume is no longer thread-safe, so it cannot be called from a @concurrent function. (If you do, it'll probably work but there is a potential race with wait_for_resume)
There is no more option to "take" a reply to signal to future hooks, so a few addons had those checks removed. Other than interception, there were no examples of taking a reply in the repo.

mhils

This looks exceptional - thank you so much! 🍰🎉 I'm in the somewhat dry final write-up stage of my PhD and this PR just made my day. 😃

A few minor comments below. Once those are addressed I'd be happy to get this into main right away so that it can already get some testing before 8.0. There likely are some unintended side-effects we haven't considered yet, but that's okay. Overall this PR is a massive net benefit.

Ah, also for your entertainment:

chrome_2022-02-03_16-20-46.mp4

A mismatch between the actual changes and the rendered diff tickles my spidey sense, but I'll refrain from falling into this rabbit hole now. 😅

mhils

Excellent, thanks. I'll re-review the changes tomorrow with a fresh mind so that this can go into main right away. 😃

nneonneo · 2022-02-04T16:08:54Z

Just rebased to main to fix a merge conflict.

mhils · 2022-02-04T16:10:36Z

Thanks! I'm currently on a final pass locally and will force-push again once I'm done. :)

async_trigger -> trigger_event invoke_addon -> invoke_addon_sync (API breakage) async_invoke_addon -> invoke_addon

Previously, it was throwing the following error: File mitmproxy/addons/dumper.py, line 133, in _echo_request_line elif flow.client_conn.peername: AttributeError: 'bool' object has no attribute 'peername'

This should improve behaviour, since calls to @Concurrent will now be serialized relative to other hooks on the same flow (but will still run in parallel with hooks on different flows). Unlike a plain async hook, @Concurrent allows blocking sync APIs to run concurrently (e.g. requests). This is patch 1/4 of the reply-ectomy.

In principle, a flow is killable as long as the connection handler is still checking the error status of the flow. This is patch 2/4 of the reply-ectomy.

This is patch 3/4 of the reply-ectomy.

The major, breaking change is that it is no longer possible to "take" a reply in order to block the effect of a later addon hook. This is patch 4/4 of the reply-ectomy.

We previously relied on the state of `Flow.reply` to check if a flow can be killed, but this doesn't work anymore with `Flow.reply` being removed. Instead, we now reintroduce the `Flow.live` attribute, which signals if we are on a live connection. Killing still is not ideal (see comment in `Flow.killable`), but this paves the way.

mhils · 2022-02-04T16:43:33Z

@nneonneo: Done! I think we are in a good state here. The only major issue I discovered was that Flow.killable turned true for flows loaded from disk, which is now mitigated by reintroducing Flow.live. I think it's also not a bad idea that flow-modifying addons (map local/map remote/modify headers/ASGI addons) are not applied when flows are read from disk, so I've added that to 372a632. Let me know if you have any comments/thoughts, otherwise this is good to go from my end! 😃

nneonneo · 2022-02-04T20:14:01Z

LGTM! Let's ship it! :)

nneonneo force-pushed the async_addons branch from 806cce1 to 06a8953 Compare February 2, 2022 07:57

mhils reviewed Feb 2, 2022

View reviewed changes

Comment thread mitmproxy/addonmanager.py Outdated

Comment thread mitmproxy/addonmanager.py Outdated

Comment thread examples/addons/nonblocking.py Outdated

nneonneo force-pushed the async_addons branch from e258301 to de48612 Compare February 2, 2022 09:14

nneonneo force-pushed the async_addons branch from de48612 to b02f2d8 Compare February 3, 2022 14:27

mhils reviewed Feb 3, 2022

View reviewed changes

Comment thread mitmproxy/addons/modifybody.py Outdated

Comment thread mitmproxy/addons/onboarding.py Outdated

Comment thread mitmproxy/flow.py

nneonneo force-pushed the async_addons branch 2 times, most recently from 582f8c1 to a3dc25b Compare February 3, 2022 19:27

mhils reviewed Feb 3, 2022

View reviewed changes

nneonneo force-pushed the async_addons branch from a3dc25b to 044b307 Compare February 4, 2022 16:08

nneonneo and others added 13 commits February 4, 2022 17:30

Support async hooks. Fixes mitmproxy#4207.

cee4b72

Tests for async hook support

687ab4b

Rename new async helper functions.

ee4999e

async_trigger -> trigger_event invoke_addon -> invoke_addon_sync (API breakage) async_invoke_addon -> invoke_addon

fix test.helper_tools.dumperview script.

e186ccb

Previously, it was throwing the following error: File mitmproxy/addons/dumper.py, line 133, in _echo_request_line elif flow.client_conn.peername: AttributeError: 'bool' object has no attribute 'peername'

Use async for tctx.cycle/tctx.invoke.

caf4930

Flow.kill: don't depend on reply status.

fd43ca1

In principle, a flow is killable as long as the connection handler is still checking the error status of the flow. This is patch 2/4 of the reply-ectomy.

Flow.intercept: use an Event instead of the reply system

ede269f

This is patch 3/4 of the reply-ectomy.

Entirely remove the reply system.

794c18b

The major, breaking change is that it is no longer possible to "take" a reply in order to block the effect of a later addon hook. This is patch 4/4 of the reply-ectomy.

Remove no-longer-used tctx.invoke function

6dc6269

perf: fast path for sync hooks

068204f

simplify concurrency tests

4f5930e

mhils force-pushed the async_addons branch from 044b307 to 372a632 Compare February 4, 2022 16:31

mhils merged commit 4468046 into mitmproxy:main Feb 4, 2022

nneonneo deleted the async_addons branch May 4, 2022 01:08

Uh oh!

Conversation

nneonneo commented Feb 2, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist

Uh oh!

mhils left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nneonneo commented Feb 2, 2022

Uh oh!

mhils commented Feb 2, 2022

Uh oh!

Prinzhorn commented Feb 3, 2022

Uh oh!

nneonneo commented Feb 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mhils left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mhils left a comment

Choose a reason for hiding this comment

Uh oh!

nneonneo commented Feb 4, 2022

Uh oh!

mhils commented Feb 4, 2022

Uh oh!

mhils commented Feb 4, 2022

Uh oh!

nneonneo commented Feb 4, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nneonneo commented Feb 2, 2022 •

edited

Loading

nneonneo commented Feb 3, 2022 •

edited

Loading