Better handling of blocked streams.

Lukasa · Lukasa · commit 9f1c3b151442 · 2015-12-03T14:46:19.000Z
diff --git a/src/priority/__init__.py b/src/priority/__init__.py
@@ -2,4 +2,4 @@
 """
 priority: HTTP/2 priority implementation for Python
 """
-from .priority import Stream, PriorityTree  # noqa
+from .priority import Stream, PriorityTree, DeadlockError  # noqa
diff --git a/src/priority/priority.py b/src/priority/priority.py
@@ -16,6 +16,14 @@
     import queue
 
 
+class DeadlockError(Exception):
+    """
+    Raised when there are no streams that can make progress: all streams are
+    blocked.
+    """
+    pass
+
+
 class Stream(object):
     """
     Priority information for a given stream.
@@ -82,18 +90,34 @@ def schedule(self):
         # Cannot be called on active streams.
         assert not self.active
 
-        level, child = self.child_queue.get()
-
-        if child.active:
-            next_stream = child.stream_id
-        else:
-            next_stream = child.schedule()
-
-        self.last_weight = level
-
-        level += (256 + child._defecit) // child.weight
-        child._defecit = (256 + child._defecit) % child.weight
-        self.child_queue.put((level, child))
+        next_stream = None
+        popped_streams = []
+
+        # Spin looking for the next active stream. Everything we pop off has
+        # to be rescheduled, even if it turns out none of them were active at
+        # this time.
+        try:
+            while next_stream is None:
+                # If the queue is empty, immediately fail.
+                val = self.child_queue.get(block=False)
+                popped_streams.append(val)
+                level, child = val
+
+                if child.active:
+                    next_stream = child.stream_id
+                else:
+                    # Guard against the possibility that the child also has no
+                    # suitable children.
+                    try:
+                        next_stream = child.schedule()
+                    except queue.Empty:
+                        continue
+        finally:
+            for level, child in popped_streams:
+                self.last_weight = level
+                level += (256 + child._defecit) // child.weight
+                child._defecit = (256 + child._defecit) % child.weight
+                self.child_queue.put((level, child))
 
         return next_stream
 
@@ -179,6 +203,7 @@ def insert_stream(self,
             assert depends_on is not None
             parent_stream = self._streams[depends_on]
             self._exclusive_insert(parent_stream, stream)
+            self._streams[stream_id] = stream
             return
 
         if not depends_on:
@@ -206,14 +231,18 @@ def unblock(self, stream_id):
         """
         Marks a given stream as unblocked, with more data to send.
         """
+        # When a stream becomes unblocked,
         self._streams[stream_id].active = True
 
     # The iterator protocol
     def __iter__(self):
         return self
 
     def __next__(self):
-        return self._root_stream.schedule()
+        try:
+            return self._root_stream.schedule()
+        except queue.Empty:
+            raise DeadlockError("No unblocked streams to schedule.")
 
     def next(self):
         return self.__next__()
diff --git a/test/test_priority.py b/test/test_priority.py
@@ -10,8 +10,12 @@
 import collections
 import itertools
 
+import pytest
+
 from hypothesis import given
-from hypothesis.strategies import integers, lists, tuples
+from hypothesis.strategies import (
+    integers, lists, tuples, composite, sampled_from
+)
 
 import priority
 
@@ -23,6 +27,83 @@
     unique_by=lambda x: x[0],
 )
 
+BLOCKED_AND_ACTIVE = lists(
+    elements=sampled_from([1, 3, 5, 7, 9, 11]),
+    unique=True,
+).map(
+    lambda blocked: (blocked, active_readme_streams_from_filter(blocked))
+)
+
+UNBLOCKED_AND_ACTIVE = lists(
+    elements=sampled_from([1, 3, 5, 7, 9, 11]),
+    unique=True,
+).map(
+    lambda unblocked: (unblocked, active_readme_streams_from_filter(
+        unblocked, blocked=False
+    ))
+)
+
+
+
+def readme_tree():
+    """
+    Provide a tree configured as the one in the readme.
+    """
+    p = priority.PriorityTree()
+    p.insert_stream(stream_id=1)
+    p.insert_stream(stream_id=3)
+    p.insert_stream(stream_id=5, depends_on=1)
+    p.insert_stream(stream_id=7, weight=32)
+    p.insert_stream(stream_id=9, depends_on=7, weight=8)
+    p.insert_stream(stream_id=11, depends_on=7, exclusive=True)
+    return p
+
+
+def active_readme_streams_from_filter(filtered, blocked=True):
+    """
+    Given a collection of filtered streams, determine which ones are active.
+    This applies only to the readme tree at this time, though in future it
+    should be possible to apply this to an arbitrary tree.
+
+    If ``blocked`` is ``True``, the filter is a set of blocked streams. If
+    ``False``, it's a collection of unblocked streams.
+    """
+    tree = {
+        1: {
+            5: {},
+        },
+        3: {},
+        7: {
+            11: {
+                9: {},
+            },
+        },
+    }
+    filtered = set(filtered)
+
+    def get_expected(tree):
+        expected = []
+
+        for stream_id in tree:
+            if stream_id not in filtered and blocked:
+                expected.append(stream_id)
+            elif stream_id in filtered and not blocked:
+                expected.append(stream_id)
+            else:
+                expected.extend(get_expected(tree[stream_id]))
+
+        return expected
+
+    return get_expected(tree)
+
+
+def active_streams_from_unblocked(unblocked):
+    """
+    Given a collection of unblocked streams, determine which ones are active.
+    This applies only to the readme tree at this time, though in future it
+    should be possible to apply this to an arbitrary tree.
+    """
+
 
 class TestStream(object):
     def test_stream_repr(self):
@@ -43,15 +124,55 @@ class TestPriorityTreeManual(object):
     Hypothesis-based ones for the same data, but getting Hypothesis to generate
     useful data in this case is going to be quite tricky.
     """
-    def test_priority_tree_initially_outputs_all_stream_ids(self, readme_tree):
+    @given(BLOCKED_AND_ACTIVE)
+    def test_priority_tree_initially_outputs_all_stream_ids(self,
+                                                            blocked_expected):
         """
         The first iterations of the priority tree initially output the active
         streams, in order of stream ID, regardless of weight.
         """
-        expected = [1, 3, 7]
-        result = [next(readme_tree) for _ in range(len(expected))]
+        tree = readme_tree()
+        blocked = blocked_expected[0]
+        expected = blocked_expected[1]
+
+        for stream_id in blocked:
+            tree.block(stream_id)
+
+        result = [next(tree) for _ in range(len(expected))]
         assert expected == result
 
+    @given(UNBLOCKED_AND_ACTIVE)
+    def test_priority_tree_blocking_is_isomorphic(self,
+                                                  allowed_expected):
+        """
+        Blocking everything and then unblocking certain ones has the same
+        effect as blocking specific streams.
+        """
+        tree = readme_tree()
+        allowed = allowed_expected[0]
+        expected = allowed_expected[1]
+
+        for stream_id in range(1, 12, 2):
+            tree.block(stream_id)
+
+        for stream_id in allowed:
+            tree.unblock(stream_id)
+
+        result = [next(tree) for _ in range(len(expected))]
+        assert expected == result
+
+    def test_priority_tree_raises_deadlock_error_if_all_blocked(self):
+        """
+        Assuming all streams are blocked and none can progress, asking for the
+        one with the next highest priority fires a DeadlockError.
+        """
+        tree = readme_tree()
+        for stream_id in range(1, 12, 2):
+            tree.block(stream_id)
+
+        with pytest.raises(priority.DeadlockError):
+            next(tree)
+
 
 class TestPriorityTreeOutput(object):
     """
