From 932481ad47c650cb00070394829a7cc268fdd00e Mon Sep 17 00:00:00 2001
From: Jean-Paul Calderone <exarkun@twistedmatrix.com>
Date: Wed, 20 Jan 2021 12:58:03 -0500
Subject: [PATCH] A helper for doing something repeatedly for a while

---
 src/allmydata/test/test_deferredutil.py | 55 +++++++++++++++++++++++++
 src/allmydata/util/deferredutil.py      | 30 ++++++++++++++
 2 files changed, 85 insertions(+)

diff --git a/src/allmydata/test/test_deferredutil.py b/src/allmydata/test/test_deferredutil.py
index 6ebc93556..2a155089f 100644
--- a/src/allmydata/test/test_deferredutil.py
+++ b/src/allmydata/test/test_deferredutil.py
@@ -74,3 +74,58 @@ class DeferredUtilTests(unittest.TestCase, deferredutil.WaitForDelayedCallsMixin
         d = defer.succeed(None)
         d.addBoth(self.wait_for_delayed_calls)
         return d
+
+
+class UntilTests(unittest.TestCase):
+    """
+    Tests for ``deferredutil.until``.
+    """
+    def test_exception(self):
+        """
+        If the action raises an exception, the ``Deferred`` returned by ``until``
+        fires with a ``Failure``.
+        """
+        self.assertFailure(
+            deferredutil.until(lambda: 1/0, lambda: True),
+            ZeroDivisionError,
+        )
+
+    def test_stops_on_condition(self):
+        """
+        The action is called repeatedly until ``condition`` returns ``True``.
+        """
+        calls = []
+        def action():
+            calls.append(None)
+
+        def condition():
+            return len(calls) == 3
+
+        self.assertIs(
+            self.successResultOf(
+                deferredutil.until(action, condition),
+            ),
+            None,
+        )
+        self.assertEqual(3, len(calls))
+
+    def test_waits_for_deferred(self):
+        """
+        If the action returns a ``Deferred`` then it is called again when the
+        ``Deferred`` fires.
+        """
+        counter = [0]
+        r1 = defer.Deferred()
+        r2 = defer.Deferred()
+        results = [r1, r2]
+        def action():
+            counter[0] += 1
+            return results.pop(0)
+
+        def condition():
+            return False
+
+        deferredutil.until(action, condition)
+        self.assertEqual([1], counter)
+        r1.callback(None)
+        self.assertEqual([2], counter)
diff --git a/src/allmydata/util/deferredutil.py b/src/allmydata/util/deferredutil.py
index 1d13f61e6..ed2a11ee4 100644
--- a/src/allmydata/util/deferredutil.py
+++ b/src/allmydata/util/deferredutil.py
@@ -15,7 +15,18 @@ if PY2:
 
 import time
 
+try:
+    from typing import (
+        Callable,
+        Any,
+    )
+except ImportError:
+    pass
+
 from foolscap.api import eventually
+from eliot.twisted import (
+    inline_callbacks,
+)
 from twisted.internet import defer, reactor, error
 from twisted.python.failure import Failure
 
@@ -201,3 +212,22 @@ class WaitForDelayedCallsMixin(PollMixin):
         d.addErrback(log.err, "error while waiting for delayed calls")
         d.addBoth(lambda ign: res)
         return d
+
+@inline_callbacks
+def until(
+        action,     # type: Callable[[], defer.Deferred[Any]]
+        condition,  # type: Callable[[], bool]
+):
+    # type: (...) -> defer.Deferred[None]
+    """
+    Run a Deferred-returning function until a condition is true.
+
+    :param action: The action to run.
+    :param condition: The predicate signaling stop.
+
+    :return: A Deferred that fires after the condition signals stop.
+    """
+    while True:
+        yield action()
+        if condition():
+            break