Add a tick event for addons and scripts, and use it for race-free reload

3 files changed, 24 insertions, 14 deletions

diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index 015adef9..34801ff7 100644
--- a/mitmproxy/builtins/script.py
+++ b/mitmproxy/builtins/script.py
@@ -4,8 +4,8 @@ import contextlib
 import os
 import shlex
 import sys
+import threading
 import traceback
-import copy
 
 from mitmproxy import exceptions
 from mitmproxy import controller
@@ -76,15 +76,14 @@ def load_script(path, args):
 
 
 class ReloadHandler(watchdog.events.FileSystemEventHandler):
-    def __init__(self, callback, master, options):
+    def __init__(self, callback):
         self.callback = callback
-        self.master, self.options = master, options
 
     def on_modified(self, event):
-        self.callback(self.master, self.options)
+        self.callback()
 
     def on_created(self, event):
-        self.callback(self.master, self.options)
+        self.callback()
 
 
 class Script:
@@ -99,7 +98,10 @@ class Script:
         self.ns = None
         self.observer = None
 
-        for i in controller.Events:
+        self.last_options = None
+        self.should_reload = threading.Event()
+
+        for i in controller.Events - set(["tick"]):
             def mkprox():
                 evt = i
 
@@ -117,21 +119,25 @@ class Script:
                 with scriptenv(self.path, self.args):
                     func(*args, **kwargs)
 
-    def reload(self, master, options):
-        with master.handlecontext():
+    def reload(self):
+        self.should_reload.set()
+
+    def tick(self):
+        if self.should_reload.is_set():
+            self.should_reload.clear()
             self.ns = None
-            self.configure(options)
+            ctx.log.info("Reloading script: %s" % self.name)
+            self.configure(self.last_options)
+        else:
+            self.run("tick")
 
     def configure(self, options):
+        self.last_options = options
         if not self.observer:
             self.observer = Observer()
             # Bind the handler to the real underlying master object
             self.observer.schedule(
-                ReloadHandler(
-                    self.reload,
-                    ctx.master,
-                    copy.copy(options),
-                ),
+                ReloadHandler(self.reload),
                 os.path.dirname(self.path) or "."
             )
             self.observer.start()
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index 72f8e001..d3ae1baa 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -33,6 +33,7 @@ Events = frozenset([
     "error",
     "log",
     "done",
+    "tick",
 
     "script_change",
 ])
@@ -108,6 +109,8 @@ class Master(object):
             self.shutdown()
 
     def tick(self, timeout):
+        with self.handlecontext():
+            self.addons("tick")
         changed = False
         try:
             mtype, obj = self.event_queue.get(timeout=timeout)
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 2447c8ea..394c3f38 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -78,6 +78,7 @@ class TestScript(mastertest.MasterTest):
             for _ in range(100):
                 with open("foo.py", "a") as f:
                     f.write(".")
+                m.addons.invoke_with_context(sc, "tick")
                 time.sleep(0.1)
                 if m.event_log:
                     return