Allow instruments to be called before the target function.

README.md

@@ -32,6 +32,8 @@ Probe (foo) = 5
 18
 ```
 
+Instruments aren't limited to recording devices! Use probes to fire off any kind of event handler. Instruments are free to maintain state themselves, or read it from somewhere else, to control their own behavior or even implement feedback mechanisms. A truly evil instrument could even alter the args/kwargs passed to a function on the fly, or call arbitrary Python code to do any number of crazy things. Consequently, it's up to you to govern what instruments are added to your environment.
+
 ## Managers
 
 In a running system, we want to add, remove, start, and stop probes and instruments without having to code at an interactive prompt or restart the system; we do this with an InstrumentManager. Start by configuring the global diagnose.manager:

src/diagnose/instruments.py

@@ -23,17 +23,20 @@ class Instrument(object):
                 when constructing statsd metric names.
         * value: a Python expression to be evaluated; its result is the
                  "process variable" to be used as the instrument sees fit.
-        * internal: If True, evaluate the value in the context of the
-                    wrapped function (just before it returns). If False
-                    (the default), evaluate the value in a wrapper
-                    context, which contains the local variables:
-                        * result: the return value of the target function
-                        * start/end/elapsed: float times
-                        * now: datetime.datetime.utcnow()
-                        * args/kwargs: inputs to the target function; these are
-                          also included in locals() by their argnames.
-                        * frame: sys._getframe() of the patch wrapper
-
+        * event: a string declaring when to fire the instrument, one of:
+            * call: evaluate the value just before calling the wrapped
+                    function, in a context which contains the local variables:
+                * start: float time
+                * now: datetime.datetime.utcnow()
+                * args/kwargs: inputs to the target function; these are
+                               also included in locals() by their argnames.
+                * frame: sys._getframe() of the patch wrapper
+             * return: the default; evaluate the value just after the wrapped
+                       function returns, in a context with the additional locals:
+                * result: the return value of the target function
+                * end/elapsed: float times
+             * end: evaluate the value in the context of the wrapped function
+                    just before it returns.
         * expires: a datetime, after which point the instrument will not fire,
                    or None to mean no expiration
         * custom: a dict of any additional data for subclasses. May include
@@ -45,23 +48,23 @@ class Instrument(object):
     error_expiration = datetime.datetime(1970, 1, 1)
 
     def __init__(
-        self, name, value, internal=False, expires=None, custom=None, mgr=None, **kwargs
+        self, name, value, event="return", expires=None, custom=None, mgr=None, **kwargs
     ):
         if mgr is None:
             mgr = diagnose.manager
         self.mgr = mgr
         self.name = name
         self.value = value
-        self.internal = internal
+        self.event = event
         self.expires = expires
         self.custom = custom or {}
 
     def __str__(self):
-        return "%s(name=%r, value=%r, internal=%r, expires=%r, custom=%r)" % (
+        return "%s(name=%r, value=%r, event=%r, expires=%r, custom=%r)" % (
             self.__class__.__name__,
             self.name,
             self.value,
-            self.internal,
+            self.event,
             self.expires,
             self.custom,
         )

src/diagnose/probes.py

@@ -49,13 +49,14 @@ class FunctionProbe(object):
 
     When started, a FunctionProbe "monkey-patches" its target, replacing
     it with a wrapper function. That wrapper calls the original target
-    function, and then calls each instrument. If any instruments are "internal",
-    the probe registers a trace function (via `sys.settrace`) and then
-    calls those instruments just before the original target function returns;
-    the instrument values are evaluated in the globals and locals of that frame.
-    External instruments are evaluated after the function returns; they are
-    much less invasive since they do not call settrace, but they are
-    limited to reading the function's inputs, outputs, and timing.
+    function, calling each instrument at the event it specifies: call, end,
+    or return. If the event is "end", then the probe registers a trace function
+    (via `sys.settrace`) and then calls those instruments just before the
+    original target function returns; the instrument values are evaluated
+    in the globals and locals of that frame. "Call" and "return" instruments
+    are evaluated before or after the function returns; they are much less
+    invasive since they do not call settrace, but they are limited to reading
+    the function's inputs, outputs, and timing.
     """
 
     def __init__(self, target, instruments=None):
@@ -101,23 +102,21 @@ def probe_wrapper(*args, **kwargs):
             now = datetime.datetime.utcnow()
 
             hotspots = HotspotsFinder()
-            internals, externals = [], []
+            instruments_by_event = {"call": [], "return": [], "end": []}
             for I in self.instruments.itervalues():
                 if I.expires and now > I.expires:
                     continue
                 if I.check_call(self, *args, **kwargs):
-                    if I.internal:
-                        internals.append(I)
-                    else:
-                        externals.append(I)
+                    instruments_by_event[I.event].append(I)
+                    if I.event in ("call", "return"):
                         if "hotspots" in I.value or "hotspots" in (
                             I.custom.get("tags") or ""
                         ):
                             hotspots.enabled = True
 
             target_obj, target_func_name = self.target.rsplit(".", 1)
             is_unwrapped = base.func_code.co_name == target_func_name
-            if internals:
+            if instruments_by_event["end"]:
                 # We have instruments that require evaluation in the local
                 # context of the function. Call sys.settrace() to gain access.
                 predicate = hunter.When(
@@ -142,7 +141,7 @@ def probe_wrapper(*args, **kwargs):
                         # that without a custom Cython Query.
                         module_in=target_obj,
                     ),
-                    TraceHandler(self, internals),
+                    TraceHandler(self, instruments_by_event["end"]),
                 )
                 tracer = hunter.Tracer(
                     # There's no need to call threading.settrace() because
@@ -193,18 +192,10 @@ def probe_wrapper(*args, **kwargs):
                 tracer = None
 
             try:
-                start = time.time()
-                result = base(*args, **kwargs)
-                if hotspots.enabled:
-                    hotspots.finish()
-                end = time.time()
-                elapsed = end - start
-                if externals:
+                if instruments_by_event["call"] or instruments_by_event["return"]:
+                    start = time.time()
                     _locals = {
-                        "result": result,
                         "start": start,
-                        "end": end,
-                        "elapsed": elapsed,
                         "now": now,
                         "args": args,
                         "kwargs": kwargs,
@@ -218,14 +209,39 @@ def probe_wrapper(*args, **kwargs):
                     # Add kwargs to locals
                     _locals.update(kwargs)
 
-                    for instrument in externals:
+                for instrument in instruments_by_event["call"]:
+                    try:
+                        instrument.fire(instrument.mgr.global_namespace, _locals)
+                    except:
                         try:
-                            instrument.fire(instrument.mgr.global_namespace, _locals)
+                            instrument.handle_error(self)
                         except:
-                            try:
-                                instrument.handle_error(self)
-                            except:
-                                traceback.print_exc()
+                            traceback.print_exc()
+
+                # Execute the base function and obtain its result.
+                result = base(*args, **kwargs)
+
+                if hotspots.enabled:
+                    hotspots.finish()
+
+                if instruments_by_event["return"]:
+                    end = time.time()
+                    elapsed = end - start
+                    _locals.update({
+                        "result": result,
+                        "end": end,
+                        "elapsed": elapsed,
+                    })
+
+                for instrument in instruments_by_event["return"]:
+                    try:
+                        instrument.fire(instrument.mgr.global_namespace, _locals)
+                    except:
+                        try:
+                            instrument.handle_error(self)
+                        except:
+                            traceback.print_exc()
+
                 return result
             finally:
                 if tracer is not None:

tests/__init__.py

@@ -12,7 +12,7 @@
 
 class ProbeTestCase(unittest.TestCase):
     @contextmanager
-    def probe(self, type, name, target, value, lifespan=1, custom=None, internal=False):
+    def probe(self, type, name, target, value, lifespan=1, custom=None, event="return"):
         mgr = diagnose.manager
         instrument_id = None
         try:
@@ -23,7 +23,7 @@ def probe(self, type, name, target, value, lifespan=1, custom=None, internal=Fal
                     "type": type,
                     "name": name,
                     "value": value,
-                    "internal": internal,
+                    "event": event,
                     "custom": custom or {},
                 },
                 "lifespan": lifespan,

tests/test_instruments.py

@@ -136,26 +136,26 @@ def test_multiple_instruments(self):
             mgr = diagnose.manager
             try:
                 target = "diagnose.test_fixtures.a_func"
-                mgr.specs["a_func_internal"] = {
+                mgr.specs["a_func_end"] = {
                     "target": target,
                     "instrument": {
                         "type": "hist",
-                        "name": "a_func_internal",
+                        "name": "a_func_end",
                         "value": "output",
-                        "internal": True,
+                        "event": "end",
                         "custom": {},
                     },
                     "lifespan": 1,
                     "lastmodified": datetime.datetime.utcnow(),
                     "applied": {},
                 }
-                mgr.specs["a_func_external"] = {
+                mgr.specs["a_func_return"] = {
                     "target": target,
                     "instrument": {
                         "type": "hist",
-                        "name": "a_func_external",
+                        "name": "a_func_return",
                         "value": "result",
-                        "internal": False,
+                        "event": "return",
                         "custom": {},
                     },
                     "lifespan": 1,
@@ -165,17 +165,17 @@ def test_multiple_instruments(self):
                 mgr.apply()
                 result = a_func(78)
             finally:
-                mgr.specs.pop("a_func_internal", None)
-                mgr.specs.pop("a_func_external", None)
+                mgr.specs.pop("a_func_end", None)
+                mgr.specs.pop("a_func_return", None)
                 mgr.apply()
 
             # The call MUST succeed.
             assert result == 91
 
             # The instruments MUST each have logged an entry
             assert statsd.method_calls == [
-                call.histogram("a_func_internal", 91, tags=[]),
-                call.histogram("a_func_external", 91, tags=[]),
+                call.histogram("a_func_end", 91, tags=[]),
+                call.histogram("a_func_return", 91, tags=[]),
             ]
 
     def test_replace_instrument(self):
@@ -189,7 +189,7 @@ def test_replace_instrument(self):
                         "type": "hist",
                         "name": "a_func",
                         "value": "arg",
-                        "internal": False,
+                        "event": "return",
                         "custom": {},
                     },
                     "lifespan": 1,