back pressure control

zh217 · Sep 7, 2018 · a84ae7e · a84ae7e
1 parent f839e83
commit a84ae7e
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/aiochan/channel.py b/aiochan/channel.py
@@ -589,9 +589,12 @@ def process_worker(in_q, out_q):
         elif mode == 'process':
             threading.Thread(target=process_worker, args=(in_q.sync_q, out_q.sync_q)).start()
 
+        in_flight = asyncio.Semaphore(n, loop=self.loop)
+
         async def pipe_in_worker(q):
             while True:
                 data = await self.get()
+                await in_flight.acquire()
                 if data is None:
                     await q.put(None)
                     results_chan.close()
@@ -603,6 +606,7 @@ async def pipe_in_worker(q):
         async def pipe_out_worker(q):
             while True:
                 next_item = await q.get()
+                in_flight.release()
                 if next_item is None:
                     break
                 data, async_ft = next_item
@@ -698,16 +702,20 @@ def process_worker(in_q, out_q):
         elif mode == 'process':
             threading.Thread(target=process_worker, args=(in_q.sync_q, out_q.sync_q)).start()
 
+        in_flight = asyncio.Semaphore(n, loop=self.loop)
+
         async def pipe_in_worker(q):
             while True:
                 data = await self.get()
+                await in_flight.acquire()
                 await q.put(data)
                 if data is None:
                     break
 
         async def pipe_out_worker(q):
             while True:
                 data = await q.get()
+                in_flight.release()
                 if data is None:
                     if close:
                         out.close()