aboutsummaryrefslogtreecommitdiffstats
path: root/mitmproxy/script/concurrent.py
blob: 1d935585be44d93d4f67ca0d63d02e501ad96bd8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
"""
This module provides a @concurrent decorator primitive to
offload computations from mitmproxy's main master thread.
"""

from mitmproxy import eventsequence
from mitmproxy.coretypes import basethread


class ScriptThread(basethread.BaseThread):
    name = "ScriptThread"


def concurrent(fn):
    if fn.__name__ not in eventsequence.Events - {"load", "configure", "tick"}:
        raise NotImplementedError(
            "Concurrent decorator not supported for '%s' method." % fn.__name__
        )

    def _concurrent(obj):
        def run():
            fn(obj)
            if obj.reply.state == "taken":
                if not obj.reply.has_message:
                    obj.reply.ack()
                obj.reply.commit()
        obj.reply.take()
        ScriptThread(
            "script.concurrent (%s)" % fn.__name__,
            target=run
        ).start()
    # Support @concurrent for class-based addons
    if "." in fn.__qualname__:
        return staticmethod(_concurrent)
    else:
        return _concurrent