@@ -229,6 +229,58 @@ def permissions(*permissions, allow_anon=True, allow_self_service=False):
return wrapper
+class UpstreamQueue(object):
+ UPSTREAM_NONCE = object()
+
+ def __init__(self, queue, get_local_result, send_upstream, get_upstream_result):
+ self.queue = queue
+ self.pending = []
+ self.cond = asyncio.Condition()
+ self.done = False
+ self.get_local_result = get_local_result
+ self.send_upstream = send_upstream
+ self.get_upstream_result = get_upstream_result
+
+ async def process_results(self):
+ try:
+ while True:
+ async with self.cond:
+ await self.cond.wait_for(lambda: self.pending or self.done)
+ if not self.pending:
+ if self.done:
+ return
+ continue
+
+ value, m = self.pending.pop(0)
+
+ if value is self.UPSTREAM_NONCE:
+ value = await self.get_upstream_result(m)
+
+ await self.queue.put(value)
+ finally:
+ await self.queue.put(None)
+
+ async def handler(self, m):
+ try:
+ if m is None:
+ return
+
+ value = await self.get_local_result(m)
+ if value is None:
+ await self.send_upstream(m)
+ value = self.UPSTREAM_NONCE
+
+ async with self.cond:
+ self.pending.append((value, m))
+ self.cond.notify_all()
+
+ finally:
+ async with self.cond:
+ self.done = True
+ self.cond.notify_all()
+ # await stream.done()
+
+
class ServerClient(bb.asyncrpc.AsyncServerConnection):
def __init__(self, socket, server):
super().__init__(socket, "OEHASHEQUIV", server.logger)
@@ -390,35 +442,41 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
validate_unihash(unihash)
return await self.db.insert_unihash(method, taskhash, unihash)
- async def _stream_handler(self, handler):
+ async def _stream_queue_handler(self, handler, queue):
await self.socket.send_message("ok")
- while True:
- upstream = None
+ async def recv():
+ try:
+ while True:
+ m = await self.socket.recv()
+ if not m or m == "END":
+ break
- l = await self.socket.recv()
- if not l:
- break
+ await handler(m)
+ finally:
+ await handler(None)
- try:
- # This inner loop is very sensitive and must be as fast as
- # possible (which is why the request sample is handled manually
- # instead of using 'with', and also why logging statements are
- # commented out.
- self.request_sample = self.server.request_stats.start_sample()
- request_measure = self.request_sample.measure()
- request_measure.start()
-
- if l == "END":
+ async def process():
+ while True:
+ m = await queue.get()
+ if m is None:
break
- msg = await handler(l)
- await self.socket.send(msg)
- finally:
- request_measure.end()
- self.request_sample.end()
+ await self.socket.send(m)
+ await asyncio.gather(recv(), process())
await self.socket.send("ok")
+
+ async def _stream_handler(self, handler):
+ queue = asyncio.Queue(1000)
+
+ async def h(m):
+ if m is None:
+ await queue.put(None)
+ else:
+ await queue.put(await handler(m))
+
+ await self._stream_queue_handler(h, queue)
return self.NO_RESPONSE
@permissions(READ_PERM)
Adds an API that allows a stream handler to more precisely control when a response is sent to the client. The new API does not directly send the response from the handler to the remote client, but instead the handler is expected to put the result in a provided queue when the response is ready. In particular, this allows a stream handler to defer to an upstream server (utilizing the new client streaming API) in an efficient way that does not require waiting on a roundtrip with the upstream server. Instead, several queries to the upstream can be in-flight at once. Signed-off-by: Joshua Watt <JPEWhacker@gmail.com> --- lib/hashserv/server.py | 100 ++++++++++++++++++++++++++++++++--------- 1 file changed, 79 insertions(+), 21 deletions(-)