Add ignore_count to deferredutil

src/allmydata/util/deferredutil.py

@@ -5,6 +5,7 @@ from foolscap.api import eventually, fireEventually
 from twisted.internet import defer, reactor
 
 from allmydata.util import log
+from allmydata.util.assertutil import _assert
 from allmydata.util.pollmixin import PollMixin
 
 
@@ -77,28 +78,35 @@ class HookMixin:
     I am a helper mixin that maintains a collection of named hooks, primarily
     for use in tests. Each hook is set to an unfired Deferred using 'set_hook',
     and can then be fired exactly once at the appropriate time by '_call_hook'.
+    If 'ignore_count' is given, that number of calls to '_call_hook' will be
+    ignored before firing the hook.
 
     I assume a '_hooks' attribute that should set by the class constructor to
     a dict mapping each valid hook name to None.
     """
-    def set_hook(self, name, d=None):
+    def set_hook(self, name, d=None, ignore_count=0):
         """
         Called by the hook observer (e.g. by a test).
         If d is not given, an unfired Deferred is created and returned.
         The hook must not already be set.
         """
+        self._log("set_hook %r, ignore_count=%r" % (name, ignore_count))
         if d is None:
             d = defer.Deferred()
-        assert self._hooks[name] is None, self._hooks[name]
-        assert isinstance(d, defer.Deferred), d
-        self._hooks[name] = d
+        _assert(ignore_count >= 0, ignore_count=ignore_count)
+        _assert(name in self._hooks, name=name)
+        _assert(self._hooks[name] is None, name=name, hook=self._hooks[name])
+        _assert(isinstance(d, defer.Deferred), d=d)
+
+        self._hooks[name] = (d, ignore_count)
         return d
 
-    def _call_hook(self, res, name):
+    def _call_hook(self, res, name, async=False):
         """
-        Called to trigger the hook, with argument 'res'. This is a no-op if the
-        hook is unset. Otherwise, the hook will be unset, and then its Deferred
-        will be fired synchronously.
+        Called to trigger the hook, with argument 'res'. This is a no-op if
+        the hook is unset. If the hook's ignore_count is positive, it will be
+        decremented; if it was already zero, the hook will be unset, and then
+        its Deferred will be fired synchronously.
 
         The expected usage is "deferred.addBoth(self._call_hook, 'hookname')".
         This ensures that if 'res' is a failure, the hook will be errbacked,
@@ -106,13 +114,25 @@ class HookMixin:
         'res' is returned so that the current result or failure will be passed
         through.
         """
-        d = self._hooks[name]
-        if d is None:
-            return defer.succeed(None)
-        self._hooks[name] = None
-        _with_log(d.callback, res)
+        hook = self._hooks[name]
+        if hook is None:
+            return res  # pass on error/result
+
+        (d, ignore_count) = hook
+        self._log("call_hook %r, ignore_count=%r" % (name, ignore_count))
+        if ignore_count > 0:
+            self._hooks[name] = (d, ignore_count - 1)
+        else:
+            self._hooks[name] = None
+            if async:
+                _with_log(eventually_callback(d), res)
+            else:
+                _with_log(d.callback, res)
         return res
 
+    def _log(self, msg):
+        log.msg(msg, level=log.NOISY)
+
 
 def async_iterate(process, iterable, *extra_args, **kwargs):
     """