Add a new scheduler parameter for control pending jobs size (#24)

* #19: Add a new scheduler parameter for control pending jobs size

* Clear pending queue properly.

CHANGES/19.feature

@@ -0,0 +1 @@
+Add a new scheduler parameter for control pending jobs size.

aiojobs/__init__.py

@@ -5,13 +5,13 @@
 
 
 async def create_scheduler(*, close_timeout=0.1, limit=100,
-                           exception_handler=None):
+                           pending_limit=0, exception_handler=None):
     if exception_handler is not None and not callable(exception_handler):
         raise TypeError('A callable object or None is expected, '
                         'got {!r}'.format(exception_handler))
     loop = asyncio.get_event_loop()
     return Scheduler(loop=loop, close_timeout=close_timeout,
-                     limit=limit,
+                     limit=limit, pending_limit=pending_limit,
                      exception_handler=exception_handler)
 
 __all__ = ('create_scheduler',)

aiojobs/_scheduler.py

@@ -1,5 +1,4 @@
 import asyncio
-from collections import deque
 
 from ._job import Job
 
@@ -14,16 +13,17 @@
 
 
 class Scheduler(*bases):
-    def __init__(self, *, close_timeout, limit,
+    def __init__(self, *, close_timeout, limit, pending_limit,
                  exception_handler, loop):
         self._loop = loop
         self._jobs = set()
         self._close_timeout = close_timeout
         self._limit = limit
+        self._pending_limit = pending_limit or 0  # the queue size is infinite
         self._exception_handler = exception_handler
         self._failed_tasks = asyncio.Queue(loop=loop)
         self._failed_task = loop.create_task(self._wait_failed())
-        self._pending = deque()
+        self._pending = asyncio.Queue(maxsize=pending_limit, loop=loop)
         self._closed = False
 
     def __iter__(self):
@@ -48,17 +48,21 @@ def __repr__(self):
     def limit(self):
         return self._limit
 
+    @property
+    def pending_limit(self):
+        return self._pending_limit
+
     @property
     def close_timeout(self):
         return self._close_timeout
 
     @property
     def active_count(self):
-        return len(self._jobs) - len(self._pending)
+        return len(self._jobs) - self._pending.qsize()
 
     @property
     def pending_count(self):
-        return len(self._pending)
+        return self._pending.qsize()
 
     @property
     def closed(self):
@@ -77,7 +81,8 @@ async def spawn(self, coro):
         if should_start:
             job._start()
         else:
-            self._pending.append(job)
+            # wait for free slot in queue
+            await self._pending.put(job)
         return job
 
     async def close(self):
@@ -89,7 +94,8 @@ async def close(self):
         if jobs:
             # cleanup pending queue
             # all job will be started on closing
-            self._pending.clear()
+            while not self._pending.empty():
+                self._pending.get_nowait()
             await asyncio.gather(
                 *[job._close(self._close_timeout) for job in jobs],
                 loop=self._loop, return_exceptions=True)
@@ -110,16 +116,16 @@ def exception_handler(self):
 
     def _done(self, job):
         self._jobs.discard(job)
-        if not self._pending:
+        if not self.pending_count:
             return
         # No pending jobs when limit is None
         # Safe to subtract.
         ntodo = self._limit - self.active_count
         i = 0
         while i < ntodo:
-            if not self._pending:
+            if not self.pending_count:
                 return
-            new_job = self._pending.popleft()
+            new_job = self._pending.get_nowait()
             if new_job.closed:
                 continue
             new_job._start()

tests/conftest.py

@@ -3,7 +3,12 @@
 from aiojobs import create_scheduler
 from aiojobs._scheduler import Scheduler
 
-PARAMS = dict(close_timeout=1.0, limit=100, exception_handler=None)
+PARAMS = dict(
+    close_timeout=1.0,
+    limit=100,
+    pending_limit=0,
+    exception_handler=None
+)
 
 
 @pytest.fixture

tests/test_scheduler.py

@@ -2,6 +2,7 @@
 from unittest import mock
 
 import pytest
+from async_timeout import timeout
 
 
 def test_ctor(scheduler):
@@ -228,6 +229,58 @@ async def test_limit(make_scheduler):
     assert s2.limit == 2
 
 
+async def test_pending_limit(make_scheduler):
+    s1 = await make_scheduler()
+    assert s1.pending_limit == 0
+    s2 = await make_scheduler(pending_limit=2)
+    assert s2.pending_limit == 2
+
+
+async def test_pending_queue_infinite(make_scheduler):
+    scheduler = await make_scheduler(limit=1)
+
+    async def coro(fut):
+        await fut
+
+    fut1 = asyncio.Future()
+    fut2 = asyncio.Future()
+    fut3 = asyncio.Future()
+
+    await scheduler.spawn(coro(fut1))
+    assert scheduler.pending_count == 0
+
+    await scheduler.spawn(coro(fut2))
+    assert scheduler.pending_count == 1
+
+    await scheduler.spawn(coro(fut3))
+    assert scheduler.pending_count == 2
+
+
+async def test_pending_queue_limit_wait(make_scheduler, loop):
+    scheduler = await make_scheduler(limit=1, pending_limit=1)
+
+    async def coro(fut):
+        await asyncio.sleep(0)
+        await fut
+
+    fut1 = asyncio.Future()
+    fut2 = asyncio.Future()
+    fut3 = asyncio.Future()
+
+    await scheduler.spawn(coro(fut1))
+    assert scheduler.pending_count == 0
+
+    await scheduler.spawn(coro(fut2))
+    assert scheduler.pending_count == 1
+
+    with pytest.raises(asyncio.TimeoutError):
+        # try to wait for 1 sec to add task to pending queue
+        with timeout(1, loop=loop):
+            await scheduler.spawn(coro(fut3))
+
+    assert scheduler.pending_count == 1
+
+
 async def test_scheduler_concurrency_limit(make_scheduler):
     scheduler = await make_scheduler(limit=1)