[bitbake-devel][PATCH 0/6] hashserv: Pipeline Upstream Queries

[bitbake-devel][PATCH 0/6] hashserv: Pipeline Upstream Queries

[Joshua Watt]

Adds a true asynchronous streaming API to the hash equivalence clients,
then uses that streaming API to pipeline requests to an upstream server.
This means that when a server wants to defer a request to an upstream,
it no longer has to wait for a complete round-trip API call. Instead, it
can use the faster streaming API to send the requests, and continue
processing other requests while a response inbound.

Joshua Watt (6):
  hashserv: client: Add asynchronous streaming API
  hashserv: server: Fix formatting
  hashserv: server: Add queued streaming API
  hashserv: server: Use streaming and queue API for upstream unihash
    queries
  hashserv: server: Use streaming and queue API for upstream exist
    queries
  hashserv: tests: Add test for upstream pipelining

 bin/bitbake-hashclient |   2 +-
 lib/hashserv/client.py | 306 ++++++++++++++++++++++++++++++++---------
 lib/hashserv/server.py | 191 +++++++++++++++++++------
 lib/hashserv/tests.py  | 105 +++++++++++++-
 4 files changed, 492 insertions(+), 112 deletions(-)

-- 
2.54.0

[bitbake-devel][PATCH 1/6] hashserv: client: Add asynchronous streaming API

[Joshua Watt]

Implments a true asynchronous streaming API for getting a unihash
(get_unihash_stream()) and checking if a unihash exists
(unihash_exists_stream()). These APIs allow a client to send a query the
to the server, then wait for the reply later. Using this API, it is
possible to interleave queries and replies.

In addition, the gc_mark_stream() API is renamed gc_mark_batch() to
match the pattern of the other APIs, and an actual gc_mark_stream() API
is implemented that allows asynchronous streaming like the others.

Note that the new stream APIs are not accessible in the "synchronous"
client API since async constructs must be used for them to make sense.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 bin/bitbake-hashclient |   2 +-
 lib/hashserv/client.py | 306 ++++++++++++++++++++++++++++++++---------
 lib/hashserv/tests.py  |   4 +-
 3 files changed, 243 insertions(+), 69 deletions(-)

diff --git a/bin/bitbake-hashclient b/bin/bitbake-hashclient
index 3a2bf5c0d..ee5c32f5f 100755
--- a/bin/bitbake-hashclient
+++ b/bin/bitbake-hashclient
@@ -240,7 +240,7 @@ def main():
         marked_hashes = 0
 
         try:
-            result = client.gc_mark_stream(args.mark, stdin)
+            result = client.gc_mark_batch(args.mark, stdin)
             marked_hashes = result["count"]
         except ConnectionError:
             logger.warning(
diff --git a/lib/hashserv/client.py b/lib/hashserv/client.py
index 8cb18050a..8e9eaaca3 100644
--- a/lib/hashserv/client.py
+++ b/lib/hashserv/client.py
@@ -8,44 +8,212 @@ import socket
 import asyncio
 import bb.asyncrpc
 import json
+from abc import abstractmethod
+from collections.abc import AsyncIterable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from . import create_async_client
 
-
 logger = logging.getLogger("hashserv.client")
 
 
+class AsyncQueue(AsyncIterable):
+    class Shutdown(Exception):
+        pass
+
+    SHUTDOWN_SENTINEL = object()
+
+    def __init__(self, *args, **kwargs):
+        self.__queue = asyncio.Queue()
+        self.__shutdown = False
+        self.__is_done = False
+
+    async def done(self):
+        if self.__is_done:
+            return
+        self.__is_done = True
+        await self.__queue.put(self.SHUTDOWN_SENTINEL)
+
+    async def put(self, item):
+        if self.__is_done:
+            raise self.Shutdown
+        await self.__queue.put(item)
+
+    async def get(self):
+        if self.__shutdown:
+            raise self.Shutdown
+
+        item = await self.__queue.get()
+        if item is self.SHUTDOWN_SENTINEL:
+            self.__shutdown = True
+            raise self.Shutdown
+
+        return item
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self.get()
+        except self.Shutdown:
+            raise StopAsyncIteration
+
+
+@dataclass(eq=False, frozen=True)
+class AsyncPipe:
+    send_queue: AsyncQueue
+    recv_queue: AsyncQueue
+
+
+class Stream(AsyncIterable):
+    def __init__(self, pipe):
+        self._pipe = pipe
+
+    async def done(self):
+        await self._pipe.send_queue.done()
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self.get_result()
+        except AsyncQueue.Shutdown:
+            raise StopAsyncIteration
+
+    @abstractmethod
+    async def _send_batch_input(self, i):
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    async def get_result(self):
+        raise NotImplementedError("Not implemented")
+
+    async def batch(self, inputs):
+        """
+        Does a "batch" process of stream messages. This sends the query
+        messages as fast as possible, and simultaneously attempts to read the
+        messages back. This helps to mitigate the effects of latency to the
+        hash equivalence server be allowing multiple queries to be "in-flight"
+        at once
+
+        The input may be a generator or an async generator
+        """
+
+        async def get_inputs():
+            if isinstance(inputs, AsyncIterable):
+                async for i in inputs:
+                    yield i
+            else:
+                for i in inputs:
+                    yield i
+
+        async def send():
+            try:
+                async for i in get_inputs():
+                    await self._send_batch_input(i)
+            finally:
+                await self.done()
+
+        results = []
+
+        async def recv():
+            async for item in self:
+                results.append(item)
+
+        await asyncio.gather(send(), recv())
+        return results
+
+
+class GetUnihashStream(Stream):
+    def __init__(self, pipe):
+        super().__init__(pipe)
+
+    async def _send_batch_input(self, i):
+        method, taskhash = i
+        await self.send_query(method, taskhash)
+
+    async def send_query(self, method, taskhash):
+        await self._pipe.send_queue.put(f"{method} {taskhash}")
+
+    async def get_result(self):
+        r = await self._pipe.recv_queue.get()
+        return r if r else None
+
+
+class UnihashExistsStream(Stream):
+    def __init__(self, pipe):
+        super().__init__(pipe)
+
+    async def _send_batch_input(self, i):
+        await self.send_query(i)
+
+    async def send_query(self, unihash):
+        await self._pipe.send_queue.put(unihash)
+
+    async def get_result(self):
+        r = await self._pipe.recv_queue.get()
+        return r == "true"
+
+
+class GcMarkStream(Stream):
+    def __init__(self, pipe, mark):
+        super().__init__(pipe)
+        self.mark = mark
+
+    async def _send_batch_input(self, i):
+        def row_to_dict(row):
+            pairs = row.split()
+            return dict(zip(pairs[::2], pairs[1::2]))
+
+        await self.send_mark(row_to_dict(i))
+
+    async def send_mark(self, where):
+        await self._pipe.send_queue.put(json.dumps({"mark": self.mark, "where": where}))
+
+    async def get_result(self):
+        r = await self._pipe.recv_queue.get()
+        return json.loads(r)
+
+
 class Batch(object):
-    def __init__(self):
+    def __init__(self, send_queue, recv_queue):
+        self.send_queue = send_queue
+        self.recv_queue = recv_queue
         self.done = False
         self.cond = asyncio.Condition()
         self.pending = []
-        self.results = []
         self.sent_count = 0
+        self.recv_count = 0
 
     async def recv(self, socket):
-        while True:
-            async with self.cond:
-                await self.cond.wait_for(lambda: self.pending or self.done)
+        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
+                    if not self.pending:
+                        if self.done:
+                            return
+                        continue
 
-            r = await socket.recv()
-            self.results.append(r)
+                await self.recv_queue.put(await socket.recv())
 
-            async with self.cond:
-                self.pending.pop(0)
+                async with self.cond:
+                    self.recv_count += 1
+                    self.pending.pop(0)
+        finally:
+            await self.recv_queue.done()
 
-    async def send(self, socket, msgs):
+    async def send(self, socket):
         try:
             # In the event of a restart due to a reconnect, all in-flight
             # messages need to be resent first to keep to result count in sync
             for m in self.pending:
                 await socket.send(m)
 
-            for m in msgs:
+            async for m in self.send_queue:
                 # Add the message to the pending list before attempting to send
                 # it so that if the send fails it will be retried
                 async with self.cond:
@@ -60,19 +228,14 @@ class Batch(object):
                 self.done = True
                 self.cond.notify()
 
-    async def process(self, socket, msgs):
-        await asyncio.gather(
-            self.recv(socket),
-            self.send(socket, msgs),
-        )
+    async def stream(self, socket):
+        await asyncio.gather(self.send(socket), self.recv(socket))
 
-        if len(self.results) != self.sent_count:
-            raise ValueError(
-                f"Expected result count {len(self.results)}. Expected {self.sent_count}"
+        if self.sent_count != self.recv_count:
+            raise ConnectionError(
+                f"Sent {self.sent_count} messages but only received {self.recv_count}"
             )
 
-        return self.results
-
 
 class AsyncClient(bb.asyncrpc.AsyncClient):
     MODE_NORMAL = 0
@@ -98,29 +261,32 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
             if become:
                 await self.become_user(become)
 
-    async def send_stream_batch(self, mode, msgs):
-        """
-        Does a "batch" process of stream messages. This sends the query
-        messages as fast as possible, and simultaneously attempts to read the
-        messages back. This helps to mitigate the effects of latency to the
-        hash equivalence server be allowing multiple queries to be "in-flight"
-        at once
-
-        The implementation does more complicated tracking using a count of sent
-        messages so that `msgs` can be a generator function (i.e. its length is
-        unknown)
-
-        """
-
-        b = Batch()
+    @asynccontextmanager
+    async def send_stream(self, mode):
+        send_queue = AsyncQueue()
+        recv_queue = AsyncQueue()
+        b = Batch(send_queue, recv_queue)
 
         async def proc():
-            nonlocal b
-
             await self._set_mode(mode)
-            return await b.process(self.socket, msgs)
+            await b.stream(self.socket)
+
+        async def process():
+            try:
+                await self._send_wrapper(proc)
+            finally:
+                await recv_queue.done()
 
-        return await self._send_wrapper(proc)
+        # Create background process to process messages
+        task = asyncio.create_task(process())
+
+        try:
+            yield AsyncPipe(send_queue, recv_queue)
+        except AsyncQueue.Shutdown:
+            pass
+        finally:
+            await send_queue.done()
+            await task
 
     async def invoke(self, *args, skip_mode=False, **kwargs):
         # It's OK if connection errors cause a failure here, because the mode
@@ -173,15 +339,18 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
         self.mode = new_mode
 
     async def get_unihash(self, method, taskhash):
-        r = await self.get_unihash_batch([(method, taskhash)])
-        return r[0]
+        async with self.get_unihash_stream() as stream:
+            await stream.send_query(method, taskhash)
+            return await stream.get_result()
 
     async def get_unihash_batch(self, args):
-        result = await self.send_stream_batch(
-            self.MODE_GET_STREAM,
-            (f"{method} {taskhash}" for method, taskhash in args),
-        )
-        return [r if r else None for r in result]
+        async with self.get_unihash_stream() as stream:
+            return await stream.batch(args)
+
+    @asynccontextmanager
+    async def get_unihash_stream(self):
+        async with self.send_stream(self.MODE_GET_STREAM) as pipe:
+            yield GetUnihashStream(pipe)
 
     async def report_unihash(self, taskhash, method, outhash, unihash, extra={}):
         m = extra.copy()
@@ -204,12 +373,18 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
         )
 
     async def unihash_exists(self, unihash):
-        r = await self.unihash_exists_batch([unihash])
-        return r[0]
+        async with self.unihash_exists_stream() as stream:
+            await stream.send_query(unihash)
+            return await stream.get_result()
 
     async def unihash_exists_batch(self, unihashes):
-        result = await self.send_stream_batch(self.MODE_EXIST_STREAM, unihashes)
-        return [r == "true" for r in result]
+        async with self.unihash_exists_stream() as stream:
+            return await stream.batch(unihashes)
+
+    @asynccontextmanager
+    async def unihash_exists_stream(self):
+        async with self.send_stream(self.MODE_EXIST_STREAM) as pipe:
+            yield UnihashExistsStream(pipe)
 
     async def get_outhash(self, method, outhash, taskhash, with_unihash=True):
         return await self.invoke(
@@ -309,23 +484,22 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
         """
         return await self.invoke({"gc-mark": {"mark": mark, "where": where}})
 
-    async def gc_mark_stream(self, mark, rows):
+    async def gc_mark_batch(self, mark, rows):
         """
         Similar to `gc-mark`, but accepts a list of "where" key-value pair
         conditions. It utilizes stream mode to mark hashes, which helps reduce
         the impact of latency when communicating with the hash equivalence
         server.
         """
-        def row_to_dict(row):
-            pairs = row.split()
-            return dict(zip(pairs[::2], pairs[1::2]))
+        async with self.gc_mark_stream(mark) as stream:
+            results = await stream.batch(rows)
 
-        responses = await self.send_stream_batch(
-            self.MODE_MARK_STREAM,
-            (json.dumps({"mark": mark, "where": row_to_dict(row)}) for row in rows),
-        )
+        return {"count": sum(int(r["count"]) for r in results)}
 
-        return {"count": sum(int(json.loads(r)["count"]) for r in responses)}
+    @asynccontextmanager
+    async def gc_mark_stream(self, mark):
+        async with self.send_stream(self.MODE_MARK_STREAM) as pipe:
+            yield GcMarkStream(pipe, mark)
 
     async def gc_sweep(self, mark):
         """
@@ -372,7 +546,7 @@ class Client(bb.asyncrpc.Client):
             "get_db_query_columns",
             "gc_status",
             "gc_mark",
-            "gc_mark_stream",
+            "gc_mark_batch",
             "gc_sweep",
         )
 
diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index 7c736d6cc..3acfdcd8c 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -1055,7 +1055,7 @@ class HashEquivalenceCommonTests(object):
         # First hash is still present
         self.assertClientGetHash(self.client, taskhash, unihash)
 
-    def test_gc_stream(self):
+    def test_gc_batch(self):
         taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
         outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
         unihash = '46edb5140d2613049332d0bf3745d9fafec9c559dac8cc61813739a28007fcdf'
@@ -1078,7 +1078,7 @@ class HashEquivalenceCommonTests(object):
         self.assertClientGetHash(self.client, taskhash3, unihash3)
 
         # Mark the first unihash to be kept
-        ret = self.client.gc_mark_stream("ABC", (f"unihash {h}" for h in [unihash, unihash2]))
+        ret = self.client.gc_mark_batch("ABC", (f"unihash {h}" for h in [unihash, unihash2]))
         self.assertEqual(ret, {"count": 2})
 
         ret = self.client.gc_status()
-- 
2.54.0

[bitbake-devel][PATCH 2/6] hashserv: server: Fix formatting

[Joshua Watt]

Reformats the file with `black`

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/server.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index 3ff434785..e7e79196f 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -424,7 +424,7 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
     @permissions(READ_PERM)
     async def handle_get_stream(self, request):
         async def handler(l):
-            (method, taskhash) = l.split()
+            method, taskhash = l.split()
             # self.logger.debug('Looking up %s %s' % (method, taskhash))
             row = await self.db.get_equivalent(method, taskhash)
 
@@ -646,7 +646,7 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
 
     @permissions(DB_ADMIN_PERM)
     async def handle_gc_status(self, request):
-        (keep_rows, remove_rows, current_mark) = await self.db.gc_status()
+        keep_rows, remove_rows, current_mark = await self.db.gc_status()
         return {
             "keep": keep_rows,
             "remove": remove_rows,
@@ -903,7 +903,9 @@ class Server(bb.asyncrpc.AsyncServer):
                 d = await client.get_taskhash(method, taskhash)
                 if d is not None:
                     if is_valid_unihash(d.get("unihash")):
-                        await db.insert_unihash(d["method"], d["taskhash"], d["unihash"])
+                        await db.insert_unihash(
+                            d["method"], d["taskhash"], d["unihash"]
+                        )
                     else:
                         self.logger.warning("Upstream server returned invalid unihash")
                 self.backfill_queue.task_done()
-- 
2.54.0

[bitbake-devel][PATCH 3/6] hashserv: server: Add queued streaming API

[Joshua Watt]

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(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index e7e79196f..991b99b86 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -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)
-- 
2.54.0

[bitbake-devel][PATCH 4/6] hashserv: server: Use streaming and queue API for upstream unihash queries

[Joshua Watt]

Reworks the "get-unihash" handler to use the new server queue API and
client streaming API to efficiently stream requests to the upstream
server instead of having to wait for a roundtrip on the requests.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/server.py | 52 ++++++++++++++++++++++++++++++------------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index 991b99b86..d392ac19f 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -481,24 +481,48 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
 
     @permissions(READ_PERM)
     async def handle_get_stream(self, request):
-        async def handler(l):
-            method, taskhash = l.split()
-            # self.logger.debug('Looking up %s %s' % (method, taskhash))
-            row = await self.db.get_equivalent(method, taskhash)
-
-            if row is not None:
-                # self.logger.debug('Found equivalent task %s -> %s', (row['taskhash'], row['unihash']))
+        async def get_unihash(m):
+            method, taskhash = m.split()
+            if (row := await self.db.get_equivalent(method, taskhash)) is not None:
                 return row["unihash"]
 
-            if self.upstream_client is not None:
-                upstream = await self.upstream_client.get_unihash(method, taskhash)
-                if upstream:
-                    await self.server.backfill_queue.put((method, taskhash))
-                    return upstream
-
             return ""
 
-        return await self._stream_handler(handler)
+        if not self.upstream_client:
+            return await self._stream_handler(get_unihash)
+
+        async with self.upstream_client.get_unihash_stream() as stream:
+
+            async def get_local_result(m):
+                method, taskhash = m.split()
+                if (row := await self.db.get_equivalent(method, taskhash)) is not None:
+                    return row["unihash"]
+                return None
+
+            async def send_upstream(m):
+                method, taskhash = m.split()
+                await stream.send_query(method, taskhash)
+
+            async def get_upstream_result(m):
+                unihash = await stream.get_result()
+                if unihash:
+                    method, taskhash = m.split()
+                    await self.server.backfill_queue.put((method, taskhash))
+                return unihash
+
+            queue = asyncio.Queue()
+            upstream = UpstreamQueue(
+                queue,
+                get_local_result,
+                send_upstream,
+                get_upstream_result,
+            )
+
+            await asyncio.gather(
+                self._stream_queue_handler(upstream.handler, queue),
+                upstream.process_results(),
+            )
+        return self.NO_RESPONSE
 
     @permissions(READ_PERM)
     async def handle_exists_stream(self, request):
-- 
2.54.0

[bitbake-devel][PATCH 5/6] hashserv: server: Use streaming and queue API for upstream exist queries

[Joshua Watt]

Reworks the "unihash-exists" handler to use the new server queue API and
client streaming API to efficiently stream requests to the upstream
server instead of having to wait for a roundtrip on the requests.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/server.py | 31 ++++++++++++++++++++++++++-----
 lib/hashserv/tests.py  |  5 +++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index d392ac19f..5017320c1 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -526,17 +526,38 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
 
     @permissions(READ_PERM)
     async def handle_exists_stream(self, request):
-        async def handler(l):
+        async def exists_handler(l):
             if await self.db.unihash_exists(l):
                 return "true"
+            return "false"
+
+        if not self.upstream_client:
+            return await self._stream_handler(exists_handler)
 
-            if self.upstream_client is not None:
-                if await self.upstream_client.unihash_exists(l):
+        async with self.upstream_client.unihash_exists_stream() as stream:
+
+            async def get_local_result(m):
+                if await self.db.unihash_exists(m):
                     return "true"
+                return None
 
-            return "false"
+            async def get_upstream_result(m):
+                exists = await stream.get_result()
+                return "true" if exists else "false"
 
-        return await self._stream_handler(handler)
+            queue = asyncio.Queue()
+            upstream = UpstreamQueue(
+                queue,
+                get_local_result,
+                stream.send_query,
+                get_upstream_result,
+            )
+
+            await asyncio.gather(
+                self._stream_queue_handler(upstream.handler, queue),
+                upstream.process_results(),
+            )
+        return self.NO_RESPONSE
 
     async def report_readonly(self, data):
         method = data["method"]
diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index 3acfdcd8c..0fbc19c1d 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -372,20 +372,25 @@ class HashEquivalenceCommonTests(object):
             nonlocal side_client
 
             # check upstream server
+            self.assertTrue(self.client.unihash_exists(unihash))
             self.assertClientGetHash(self.client, taskhash, unihash)
 
             # Hash should *not* be present on the side server
+            if old_sidehash and unihash != old_sidehash:
+                self.assertFalse(side_client.unihash_exists(unihash))
             self.assertClientGetHash(side_client, taskhash, old_sidehash)
 
             # Hash should be present on the downstream server, since it
             # will defer to the upstream server. This will trigger
             # the backfill in the downstream server
+            self.assertTrue(down_client.unihash_exists(unihash))
             self.assertClientGetHash(down_client, taskhash, unihash)
 
             # After waiting for the downstream client to finish backfilling the
             # task from the upstream server, it should appear in the side server
             # since the database is populated
             down_client.backfill_wait()
+            self.assertTrue(side_client.unihash_exists(unihash))
             self.assertClientGetHash(side_client, taskhash, unihash)
 
         # Basic report
-- 
2.54.0

[bitbake-devel][PATCH 6/6] hashserv: tests: Add test for upstream pipelining

[Joshua Watt]

Adds a test that verifies that pipelining to an upstream server works as
expected.

AI-Generated: Uses Cursor, Claude Sonnet 5
Co-authored-by: Michal Sieron <michal.sieron@nokia.com>
Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/tests.py | 96 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index 0fbc19c1d..41c1f249f 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -1561,6 +1561,102 @@ class TestHashEquivalenceTCPServer(HashEquivalenceTestSetup, HashEquivalenceComm
         # case it is more reliable to resolve the IP address explicitly.
         return socket.gethostbyname("localhost") + ":0"
 
+    def test_get_stream_upstream_pipelined(self):
+        # Verify that, with an upstream configured, the get-stream handler
+        # pipelines its upstream queries instead of doing one blocking
+        # round-trip per task. A latency proxy injects RTT between this server
+        # and its upstream; a serial (one round-trip per task) implementation
+        # could not beat N * RTT, so finishing far faster proves the queries
+        # are pipelined.
+        import asyncio
+
+        LATENCY = 0.01  # 10ms each direction => ~20ms round trip
+        upstream_host, upstream_port = self.server_address.rsplit(":", 1)
+        upstream_port = int(upstream_port)
+
+        ready = threading.Event()
+        proxy = {}
+
+        def run_proxy():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            async def handle(creader, cwriter):
+                ureader, uwriter = await asyncio.open_connection(upstream_host, upstream_port)
+
+                async def pipe(r, w):
+                    try:
+                        while True:
+                            data = await r.read(65536)
+                            if not data:
+                                break
+                            await asyncio.sleep(LATENCY)
+                            w.write(data)
+                            await w.drain()
+                    except Exception:
+                        pass
+                    finally:
+                        try:
+                            w.close()
+                        except Exception:
+                            pass
+
+                await asyncio.gather(pipe(creader, uwriter), pipe(ureader, cwriter))
+
+            stop_event = asyncio.Event()
+            proxy["stop_event"] = stop_event
+            proxy["loop"] = loop
+
+            async def main():
+                server = await asyncio.start_server(handle, "127.0.0.1", 0)
+                proxy["port"] = server.sockets[0].getsockname()[1]
+                ready.set()
+                async with server:
+                    await stop_event.wait()
+
+            try:
+                loop.run_until_complete(main())
+            except Exception:
+                pass
+            finally:
+                loop.close()
+
+        proxy_thread = threading.Thread(target=run_proxy, daemon=True)
+        proxy_thread.start()
+        self.assertTrue(ready.wait(10), "latency proxy did not start")
+        self.addCleanup(proxy_thread.join, 10)
+        self.addCleanup(lambda: proxy["loop"].call_soon_threadsafe(proxy["stop_event"].set))
+        proxy_addr = "127.0.0.1:%d" % proxy["port"]
+
+        N = 400
+        expected = []
+        for i in range(N):
+            taskhash = hashlib.sha256(("task%d" % i).encode()).hexdigest()
+            outhash = hashlib.sha256(("out%d" % i).encode()).hexdigest()
+            unihash = hashlib.sha256(("uni%d" % i).encode()).hexdigest()
+            self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
+            expected.append((self.METHOD, taskhash, unihash))
+
+        # Downstream server with an EMPTY local DB whose upstream is the slow proxy.
+        down_server = self.start_server(upstream=proxy_addr)
+        down_client = self.start_client(down_server.address)
+
+        args = [(m, th) for (m, th, _uh) in expected]
+        start = time.time()
+        results = down_client.get_unihash_batch(args)
+        elapsed = time.time() - start
+
+        self.assertEqual(results, [uh for (_m, _th, uh) in expected])
+
+        # A serial (one round-trip per task) implementation cannot beat N*RTT.
+        # Allow a generous margin to avoid flakiness on loaded CI machines.
+        serial_lower_bound = N * 2 * LATENCY
+        self.assertLess(elapsed, serial_lower_bound / 4,
+                        "Upstream queries are not being pipelined "
+                        "(%.3fs for %d queries at %.0fms injected RTT)"
+                        % (elapsed, N, 2 * LATENCY * 1000))
+
+
 
 class TestHashEquivalenceWebsocketServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
     def setUp(self):
-- 
2.54.0

Re: [bitbake-devel][PATCH 0/6] hashserv: Pipeline Upstream Queries

[Richard Purdie]

Hi Joshua,

On Fri, 2026-07-24 at 15:25 -0600, Joshua Watt via lists.openembedded.org wrote:
> Adds a true asynchronous streaming API to the hash equivalence clients,
> then uses that streaming API to pipeline requests to an upstream server.
> This means that when a server wants to defer a request to an upstream,
> it no longer has to wait for a complete round-trip API call. Instead, it
> can use the faster streaming API to send the requests, and continue
> processing other requests while a response inbound.
> 
> Joshua Watt (6):
>   hashserv: client: Add asynchronous streaming API
>   hashserv: server: Fix formatting
>   hashserv: server: Add queued streaming API
>   hashserv: server: Use streaming and queue API for upstream unihash
>     queries
>   hashserv: server: Use streaming and queue API for upstream exist
>     queries
>   hashserv: tests: Add test for upstream pipelining

Thanks, this looks promising but builds didn't like it:

https://autobuilder.yoctoproject.org/valkyrie/#/builders/22/builds/4225

(and more in that a-full).

Cheers,

Richard

[bitbake-devel][PATCH v2 00/10] hashserv: Pipeline Upstream Queries

[Joshua Watt]

Adds a true asynchronous streaming API to the hash equivalence clients,
then uses that streaming API to pipeline requests to an upstream server.
This means that when a server wants to defer a request to an upstream,
it no longer has to wait for a complete round-trip API call. Instead, it
can use the faster streaming API to send the requests, and continue
processing other requests while a response inbound.

V2: Fixes several bugs related to when the connection between the client
and upstream server was disconnected, and added tests for them.


Joshua Watt (10):
  asyncrpc: Add Task Group
  asyncrpc: serv: Use Task Group
  asyncrpc: serv: Cancel all clients on server stop
  hashserv: tests: Improve test logging
  hashserv: client: Add asynchronous streaming API
  hashserv: server: Add queued streaming API
  hashserv: server: Use streaming and queue API for upstream unihash
    queries
  hashserv: server: Use streaming and queue API for upstream exist
    queries
  hashserv: tests: Add more upstream tests
  hashserv: tests: Add test for upstream pipelining

 bin/bitbake-hashclient       |   2 +-
 lib/bb/asyncrpc/__init__.py  |   1 +
 lib/bb/asyncrpc/serv.py      |  13 +-
 lib/bb/asyncrpc/taskgroup.py |  52 ++++++
 lib/hashserv/client.py       | 349 +++++++++++++++++++++++++++--------
 lib/hashserv/server.py       | 181 ++++++++++++++----
 lib/hashserv/tests.py        | 261 ++++++++++++++++++++++++--
 7 files changed, 723 insertions(+), 136 deletions(-)
 create mode 100644 lib/bb/asyncrpc/taskgroup.py

-- 
2.54.0

[bitbake-devel][PATCH v2 01/10] asyncrpc: Add Task Group

[Joshua Watt]

Adds a Task Group object which can be used to manage multiple
asynchronous tasks and correctly cancel them if an exception is raised

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/bb/asyncrpc/__init__.py  |  1 +
 lib/bb/asyncrpc/taskgroup.py | 52 ++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 lib/bb/asyncrpc/taskgroup.py

diff --git a/lib/bb/asyncrpc/__init__.py b/lib/bb/asyncrpc/__init__.py
index a4371643d..2f0956dfa 100644
--- a/lib/bb/asyncrpc/__init__.py
+++ b/lib/bb/asyncrpc/__init__.py
@@ -14,3 +14,4 @@ from .exceptions import (
     ConnectionClosedError,
     InvokeError,
 )
+from .taskgroup import TaskGroup
diff --git a/lib/bb/asyncrpc/taskgroup.py b/lib/bb/asyncrpc/taskgroup.py
new file mode 100644
index 000000000..b61f40385
--- /dev/null
+++ b/lib/bb/asyncrpc/taskgroup.py
@@ -0,0 +1,52 @@
+#
+# Copyright BitBake Contributors
+#
+# SPDX-License-Identifier: GPL-2.0-only
+#
+import asyncio
+import logging
+
+logger = logging.getLogger("asyncio.TaskGroup")
+
+
+class TaskGroup(object):
+    def __init__(self):
+        self._tasks = []
+
+    def create_task(self, coro, **kwargs):
+        self._tasks.append(asyncio.create_task(coro, **kwargs))
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        try:
+            if exc is None:
+                while self._tasks:
+                    done, pending = await asyncio.wait(
+                        self._tasks, return_when=asyncio.FIRST_COMPLETED
+                    )
+                    self._tasks = pending
+
+                    for t in done:
+                        try:
+                            await t
+                        except asyncio.CancelledError:
+                            pass
+
+        finally:
+            for t in self._tasks:
+                t.cancel()
+                try:
+                    await t
+                except:
+                    # Ignore exceptions
+                    pass
+
+        return False
+
+    @classmethod
+    async def run(cls, *coros):
+        async with cls() as group:
+            for c in coros:
+                group.create_task(c)
-- 
2.54.0

[bitbake-devel][PATCH v2 02/10] asyncrpc: serv: Use Task Group

[Joshua Watt]

Uses a task group instead of asyncio.gather(). The task group ensures
that all tasks are canceled if an exception occurs, which
asyncio.gather() does not.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/bb/asyncrpc/serv.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/bb/asyncrpc/serv.py b/lib/bb/asyncrpc/serv.py
index bd1aded8d..d3b1c6c35 100644
--- a/lib/bb/asyncrpc/serv.py
+++ b/lib/bb/asyncrpc/serv.py
@@ -334,7 +334,7 @@ class AsyncServer(object):
             self.loop.add_signal_handler(signal.SIGQUIT, self.signal_handler)
             signal.pthread_sigmask(signal.SIG_UNBLOCK, [signal.SIGTERM])
 
-            self.loop.run_until_complete(asyncio.gather(*tasks))
+            self.loop.run_until_complete(bb.asyncrpc.TaskGroup.run(*tasks))
 
             self.logger.debug("Server shutting down")
         finally:
-- 
2.54.0

[bitbake-devel][PATCH v2 03/10] asyncrpc: serv: Cancel all clients on server stop

[Joshua Watt]

When the server is stopped, all clients should be stopped, otherwise the
server will hang waiting for a disconnect.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/bb/asyncrpc/serv.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/lib/bb/asyncrpc/serv.py b/lib/bb/asyncrpc/serv.py
index d3b1c6c35..dd77ed0a4 100644
--- a/lib/bb/asyncrpc/serv.py
+++ b/lib/bb/asyncrpc/serv.py
@@ -56,7 +56,7 @@ class AsyncServerConnection(object):
             if not client_protocol:
                 return
 
-            (client_proto_name, client_proto_version) = client_protocol.split()
+            client_proto_name, client_proto_version = client_protocol.split()
             if client_proto_name != self.proto_name:
                 self.logger.debug("Rejecting invalid protocol %s" % (self.proto_name))
                 return
@@ -123,6 +123,7 @@ class StreamServer(object):
         self.handler = handler
         self.logger = logger
         self.closed = False
+        self.clients = []
 
     async def handle_stream_client(self, reader, writer):
         # writer.transport.set_write_buffer_limits(0)
@@ -131,10 +132,16 @@ class StreamServer(object):
             await socket.close()
             return
 
-        await self.handler(socket)
+        self.clients.append(socket)
+        try:
+            await self.handler(socket)
+        finally:
+            self.clients.remove(socket)
 
     async def stop(self):
         self.closed = True
+        for socket in self.clients:
+            await socket.close()
 
 
 class TCPStreamServer(StreamServer):
-- 
2.54.0

[bitbake-devel][PATCH v2 04/10] hashserv: tests: Improve test logging

[Joshua Watt]

Log client messages to a file to aid in test debugging

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/tests.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index 7c736d6cc..551b9e298 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -30,7 +30,7 @@ BIN_DIR = THIS_DIR.parent.parent / "bin"
 
 def server_prefunc(server, idx):
     logging.basicConfig(level=logging.DEBUG, filename='bbhashserv-%d.log' % idx, filemode='w',
-                        format='%(levelname)s %(filename)s:%(lineno)d %(message)s')
+                        format='%(levelname)s %(filename)s:%(lineno)d %(message)s', force=True)
     server.logger.debug("Running server %d" % idx)
     sys.stdout = open('bbhashserv-stdout-%d.log' % idx, 'w')
     sys.stderr = sys.stdout
@@ -93,6 +93,8 @@ class HashEquivalenceTestSetup(object):
         return self.start_client(self.auth_server_address, user["username"], user["token"])
 
     def setUp(self):
+        logging.basicConfig(level=logging.DEBUG, filename='bbhashtest.log', filemode='w',
+                        format='%(levelname)s %(filename)s:%(lineno)d %(message)s')
         self.temp_dir = tempfile.TemporaryDirectory(prefix='bb-hashserv')
         self.addCleanup(self.temp_dir.cleanup)
 
-- 
2.54.0

[bitbake-devel][PATCH v2 05/10] hashserv: client: Add asynchronous streaming API

[Joshua Watt]

Implments a true asynchronous streaming API for getting a unihash
(get_unihash_stream()) and checking if a unihash exists
(unihash_exists_stream()). These APIs allow a client to send a query the
to the server, then wait for the reply later. Using this API, it is
possible to interleave queries and replies.

In addition, the gc_mark_stream() API is renamed gc_mark_batch() to
match the pattern of the other APIs, and an actual gc_mark_stream() API
is implemented that allows asynchronous streaming like the others.

Note that the new stream APIs are not accessible in the "synchronous"
client API since async constructs must be used for them to make sense.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 bin/bitbake-hashclient |   2 +-
 lib/hashserv/client.py | 349 ++++++++++++++++++++++++++++++++---------
 lib/hashserv/tests.py  |   4 +-
 3 files changed, 275 insertions(+), 80 deletions(-)

diff --git a/bin/bitbake-hashclient b/bin/bitbake-hashclient
index 3a2bf5c0d..ee5c32f5f 100755
--- a/bin/bitbake-hashclient
+++ b/bin/bitbake-hashclient
@@ -240,7 +240,7 @@ def main():
         marked_hashes = 0
 
         try:
-            result = client.gc_mark_stream(args.mark, stdin)
+            result = client.gc_mark_batch(args.mark, stdin)
             marked_hashes = result["count"]
         except ConnectionError:
             logger.warning(
diff --git a/lib/hashserv/client.py b/lib/hashserv/client.py
index 8cb18050a..b447fb259 100644
--- a/lib/hashserv/client.py
+++ b/lib/hashserv/client.py
@@ -8,70 +8,252 @@ import socket
 import asyncio
 import bb.asyncrpc
 import json
+from abc import abstractmethod
+from collections.abc import AsyncIterable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from . import create_async_client
 
-
 logger = logging.getLogger("hashserv.client")
 
 
+class AsyncQueue(AsyncIterable):
+    class Shutdown(Exception):
+        pass
+
+    SHUTDOWN_SENTINEL = object()
+
+    def __init__(self, *args, **kwargs):
+        self.__queue = asyncio.Queue()
+        self.__shutdown = False
+        self.__is_done = False
+
+    async def done(self):
+        if self.__is_done:
+            return
+        self.__is_done = True
+        await self.__queue.put(self.SHUTDOWN_SENTINEL)
+
+    async def put(self, item):
+        if self.__is_done:
+            raise self.Shutdown
+        await self.__queue.put(item)
+
+    async def get(self):
+        if self.__shutdown:
+            raise self.Shutdown
+
+        item = await self.__queue.get()
+        if item is self.SHUTDOWN_SENTINEL:
+            self.__shutdown = True
+            raise self.Shutdown
+
+        return item
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self.get()
+        except self.Shutdown:
+            raise StopAsyncIteration
+
+
+@dataclass(eq=False, frozen=True)
+class AsyncPipe:
+    send_queue: AsyncQueue
+    recv_queue: AsyncQueue
+
+
+class Stream(AsyncIterable):
+    def __init__(self, pipe):
+        self._pipe = pipe
+
+    async def done(self):
+        await self._pipe.send_queue.done()
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self.get_result()
+        except AsyncQueue.Shutdown:
+            raise StopAsyncIteration
+
+    @abstractmethod
+    async def _send_batch_input(self, i):
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    async def get_result(self):
+        raise NotImplementedError("Not implemented")
+
+    async def batch(self, inputs):
+        """
+        Does a "batch" process of stream messages. This sends the query
+        messages as fast as possible, and simultaneously attempts to read the
+        messages back. This helps to mitigate the effects of latency to the
+        hash equivalence server be allowing multiple queries to be "in-flight"
+        at once
+
+        The input may be a generator or an async generator
+        """
+
+        async def get_inputs():
+            if isinstance(inputs, AsyncIterable):
+                async for i in inputs:
+                    yield i
+            else:
+                for i in inputs:
+                    yield i
+
+        async def send():
+            try:
+                async for i in get_inputs():
+                    await self._send_batch_input(i)
+            finally:
+                await self.done()
+
+        results = []
+
+        async def recv():
+            async for item in self:
+                results.append(item)
+
+        await bb.asyncrpc.TaskGroup.run(send(), recv())
+        return results
+
+
+class GetUnihashStream(Stream):
+    def __init__(self, pipe):
+        super().__init__(pipe)
+
+    async def _send_batch_input(self, i):
+        method, taskhash = i
+        await self.send_query(method, taskhash)
+
+    async def send_query(self, method, taskhash):
+        await self._pipe.send_queue.put(f"{method} {taskhash}")
+
+    async def get_result(self):
+        r = await self._pipe.recv_queue.get()
+        return r if r else None
+
+
+class UnihashExistsStream(Stream):
+    def __init__(self, pipe):
+        super().__init__(pipe)
+
+    async def _send_batch_input(self, i):
+        await self.send_query(i)
+
+    async def send_query(self, unihash):
+        await self._pipe.send_queue.put(unihash)
+
+    async def get_result(self):
+        r = await self._pipe.recv_queue.get()
+        return r == "true"
+
+
+class GcMarkStream(Stream):
+    def __init__(self, pipe, mark):
+        super().__init__(pipe)
+        self.mark = mark
+
+    async def _send_batch_input(self, i):
+        def row_to_dict(row):
+            pairs = row.split()
+            return dict(zip(pairs[::2], pairs[1::2]))
+
+        await self.send_mark(row_to_dict(i))
+
+    async def send_mark(self, where):
+        await self._pipe.send_queue.put(json.dumps({"mark": self.mark, "where": where}))
+
+    async def get_result(self):
+        r = await self._pipe.recv_queue.get()
+        return json.loads(r)
+
+
 class Batch(object):
-    def __init__(self):
-        self.done = False
+    def __init__(self, send_queue, recv_queue):
+        self.send_queue = send_queue
+        self.recv_queue = recv_queue
+        self.fill_done = False
+        self.send_done = False
         self.cond = asyncio.Condition()
         self.pending = []
-        self.results = []
         self.sent_count = 0
+        self.recv_count = 0
+        self.item = None
 
     async def recv(self, socket):
         while True:
             async with self.cond:
-                await self.cond.wait_for(lambda: self.pending or self.done)
-
+                await self.cond.wait_for(lambda: self.pending or self.send_done)
                 if not self.pending:
-                    if self.done:
+                    if self.send_done:
                         return
                     continue
 
-            r = await socket.recv()
-            self.results.append(r)
+            m = await socket.recv()
+            await self.recv_queue.put(m)
 
             async with self.cond:
+                self.recv_count += 1
                 self.pending.pop(0)
 
-    async def send(self, socket, msgs):
-        try:
-            # In the event of a restart due to a reconnect, all in-flight
-            # messages need to be resent first to keep to result count in sync
+    async def fill(self):
+        async for m in self.send_queue:
+            async with self.cond:
+                # Wait for item to be consumed
+                await self.cond.wait_for(lambda: self.item is None)
+                self.item = m
+                self.cond.notify_all()
+
+        async with self.cond:
+            self.fill_done = True
+            self.cond.notify_all()
+
+    async def send(self, socket):
+        # In the event of a restart due to a reconnect, all in-flight
+        # messages need to be resent first to keep to result count in sync
+        async with self.cond:
             for m in self.pending:
                 await socket.send(m)
 
-            for m in msgs:
-                # Add the message to the pending list before attempting to send
-                # it so that if the send fails it will be retried
-                async with self.cond:
-                    self.pending.append(m)
-                    self.cond.notify()
-                    self.sent_count += 1
-
-                await socket.send(m)
-
-        finally:
+        while True:
             async with self.cond:
-                self.done = True
-                self.cond.notify()
+                await self.cond.wait_for(
+                    lambda: self.item is not None or self.fill_done
+                )
+                if self.item is None:
+                    if self.fill_done:
+                        self.send_done = True
+                        self.cond.notify_all()
+                        return
+                    continue
 
-    async def process(self, socket, msgs):
-        await asyncio.gather(
-            self.recv(socket),
-            self.send(socket, msgs),
-        )
+                m = self.item
 
-        if len(self.results) != self.sent_count:
-            raise ValueError(
-                f"Expected result count {len(self.results)}. Expected {self.sent_count}"
-            )
+            await socket.send(m)
 
-        return self.results
+            async with self.cond:
+                self.item = None
+                self.pending.append(m)
+                self.sent_count += 1
+                self.cond.notify_all()
+
+    async def stream(self, socket):
+        await bb.asyncrpc.TaskGroup.run(self.send(socket), self.recv(socket))
+
+    def check(self):
+        if self.sent_count != self.recv_count:
+            raise ConnectionError(
+                f"Sent {self.sent_count} messages but only received {self.recv_count}"
+            )
 
 
 class AsyncClient(bb.asyncrpc.AsyncClient):
@@ -98,29 +280,34 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
             if become:
                 await self.become_user(become)
 
-    async def send_stream_batch(self, mode, msgs):
-        """
-        Does a "batch" process of stream messages. This sends the query
-        messages as fast as possible, and simultaneously attempts to read the
-        messages back. This helps to mitigate the effects of latency to the
-        hash equivalence server be allowing multiple queries to be "in-flight"
-        at once
-
-        The implementation does more complicated tracking using a count of sent
-        messages so that `msgs` can be a generator function (i.e. its length is
-        unknown)
-
-        """
-
-        b = Batch()
+    @asynccontextmanager
+    async def send_stream(self, mode):
+        send_queue = AsyncQueue()
+        recv_queue = AsyncQueue()
+        b = Batch(send_queue, recv_queue)
 
         async def proc():
-            nonlocal b
-
             await self._set_mode(mode)
-            return await b.process(self.socket, msgs)
-
-        return await self._send_wrapper(proc)
+            await b.stream(self.socket)
+
+        async def process():
+            try:
+                await self._send_wrapper(proc)
+            finally:
+                await recv_queue.done()
+
+        # Create background process to process messages
+        async with bb.asyncrpc.TaskGroup() as group:
+            group.create_task(process())
+            group.create_task(b.fill())
+            try:
+                yield AsyncPipe(send_queue, recv_queue)
+                b.check()
+            except AsyncQueue.Shutdown as e:
+                pass
+            finally:
+                await send_queue.done()
+                await recv_queue.done()
 
     async def invoke(self, *args, skip_mode=False, **kwargs):
         # It's OK if connection errors cause a failure here, because the mode
@@ -173,15 +360,18 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
         self.mode = new_mode
 
     async def get_unihash(self, method, taskhash):
-        r = await self.get_unihash_batch([(method, taskhash)])
-        return r[0]
+        async with self.get_unihash_stream() as stream:
+            await stream.send_query(method, taskhash)
+            return await stream.get_result()
 
     async def get_unihash_batch(self, args):
-        result = await self.send_stream_batch(
-            self.MODE_GET_STREAM,
-            (f"{method} {taskhash}" for method, taskhash in args),
-        )
-        return [r if r else None for r in result]
+        async with self.get_unihash_stream() as stream:
+            return await stream.batch(args)
+
+    @asynccontextmanager
+    async def get_unihash_stream(self):
+        async with self.send_stream(self.MODE_GET_STREAM) as pipe:
+            yield GetUnihashStream(pipe)
 
     async def report_unihash(self, taskhash, method, outhash, unihash, extra={}):
         m = extra.copy()
@@ -204,12 +394,18 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
         )
 
     async def unihash_exists(self, unihash):
-        r = await self.unihash_exists_batch([unihash])
-        return r[0]
+        async with self.unihash_exists_stream() as stream:
+            await stream.send_query(unihash)
+            return await stream.get_result()
 
     async def unihash_exists_batch(self, unihashes):
-        result = await self.send_stream_batch(self.MODE_EXIST_STREAM, unihashes)
-        return [r == "true" for r in result]
+        async with self.unihash_exists_stream() as stream:
+            return await stream.batch(unihashes)
+
+    @asynccontextmanager
+    async def unihash_exists_stream(self):
+        async with self.send_stream(self.MODE_EXIST_STREAM) as pipe:
+            yield UnihashExistsStream(pipe)
 
     async def get_outhash(self, method, outhash, taskhash, with_unihash=True):
         return await self.invoke(
@@ -309,23 +505,22 @@ class AsyncClient(bb.asyncrpc.AsyncClient):
         """
         return await self.invoke({"gc-mark": {"mark": mark, "where": where}})
 
-    async def gc_mark_stream(self, mark, rows):
+    async def gc_mark_batch(self, mark, rows):
         """
         Similar to `gc-mark`, but accepts a list of "where" key-value pair
         conditions. It utilizes stream mode to mark hashes, which helps reduce
         the impact of latency when communicating with the hash equivalence
         server.
         """
-        def row_to_dict(row):
-            pairs = row.split()
-            return dict(zip(pairs[::2], pairs[1::2]))
+        async with self.gc_mark_stream(mark) as stream:
+            results = await stream.batch(rows)
 
-        responses = await self.send_stream_batch(
-            self.MODE_MARK_STREAM,
-            (json.dumps({"mark": mark, "where": row_to_dict(row)}) for row in rows),
-        )
+        return {"count": sum(int(r["count"]) for r in results)}
 
-        return {"count": sum(int(json.loads(r)["count"]) for r in responses)}
+    @asynccontextmanager
+    async def gc_mark_stream(self, mark):
+        async with self.send_stream(self.MODE_MARK_STREAM) as pipe:
+            yield GcMarkStream(pipe, mark)
 
     async def gc_sweep(self, mark):
         """
@@ -372,7 +567,7 @@ class Client(bb.asyncrpc.Client):
             "get_db_query_columns",
             "gc_status",
             "gc_mark",
-            "gc_mark_stream",
+            "gc_mark_batch",
             "gc_sweep",
         )
 
diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index 551b9e298..e24bdcacb 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -1057,7 +1057,7 @@ class HashEquivalenceCommonTests(object):
         # First hash is still present
         self.assertClientGetHash(self.client, taskhash, unihash)
 
-    def test_gc_stream(self):
+    def test_gc_batch(self):
         taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
         outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
         unihash = '46edb5140d2613049332d0bf3745d9fafec9c559dac8cc61813739a28007fcdf'
@@ -1080,7 +1080,7 @@ class HashEquivalenceCommonTests(object):
         self.assertClientGetHash(self.client, taskhash3, unihash3)
 
         # Mark the first unihash to be kept
-        ret = self.client.gc_mark_stream("ABC", (f"unihash {h}" for h in [unihash, unihash2]))
+        ret = self.client.gc_mark_batch("ABC", (f"unihash {h}" for h in [unihash, unihash2]))
         self.assertEqual(ret, {"count": 2})
 
         ret = self.client.gc_status()
-- 
2.54.0

[bitbake-devel][PATCH v2 06/10] hashserv: server: Add queued streaming API

[Joshua Watt]

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 | 96 +++++++++++++++++++++++++++++++++---------
 1 file changed, 75 insertions(+), 21 deletions(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index e7e79196f..0fa81e85d 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -229,6 +229,54 @@ 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):
+        if m is None:
+            async with self.cond:
+                self.done = True
+                self.cond.notify_all()
+            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()
+
+
 class ServerClient(bb.asyncrpc.AsyncServerConnection):
     def __init__(self, socket, server):
         super().__init__(socket, "OEHASHEQUIV", server.logger)
@@ -390,35 +438,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 bb.asyncrpc.TaskGroup.run(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)
-- 
2.54.0

[bitbake-devel][PATCH v2 07/10] hashserv: server: Use streaming and queue API for upstream unihash queries

[Joshua Watt]

Reworks the "get-unihash" handler to use the new server queue API and
client streaming API to efficiently stream requests to the upstream
server instead of having to wait for a roundtrip on the requests.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/server.py | 52 ++++++++++++++++++++++++++++++------------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index 0fa81e85d..d0e6f23fc 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -477,24 +477,48 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
 
     @permissions(READ_PERM)
     async def handle_get_stream(self, request):
-        async def handler(l):
-            method, taskhash = l.split()
-            # self.logger.debug('Looking up %s %s' % (method, taskhash))
-            row = await self.db.get_equivalent(method, taskhash)
-
-            if row is not None:
-                # self.logger.debug('Found equivalent task %s -> %s', (row['taskhash'], row['unihash']))
+        async def get_unihash(m):
+            method, taskhash = m.split()
+            if (row := await self.db.get_equivalent(method, taskhash)) is not None:
                 return row["unihash"]
 
-            if self.upstream_client is not None:
-                upstream = await self.upstream_client.get_unihash(method, taskhash)
-                if upstream:
-                    await self.server.backfill_queue.put((method, taskhash))
-                    return upstream
-
             return ""
 
-        return await self._stream_handler(handler)
+        if not self.upstream_client:
+            return await self._stream_handler(get_unihash)
+
+        async with self.upstream_client.get_unihash_stream() as stream:
+
+            async def get_local_result(m):
+                method, taskhash = m.split()
+                if (row := await self.db.get_equivalent(method, taskhash)) is not None:
+                    return row["unihash"]
+                return None
+
+            async def send_upstream(m):
+                method, taskhash = m.split()
+                await stream.send_query(method, taskhash)
+
+            async def get_upstream_result(m):
+                unihash = await stream.get_result()
+                if unihash:
+                    method, taskhash = m.split()
+                    await self.server.backfill_queue.put((method, taskhash))
+                return unihash
+
+            queue = asyncio.Queue()
+            upstream = UpstreamQueue(
+                queue,
+                get_local_result,
+                send_upstream,
+                get_upstream_result,
+            )
+
+            await bb.asyncrpc.TaskGroup.run(
+                self._stream_queue_handler(upstream.handler, queue),
+                upstream.process_results(),
+            )
+        return self.NO_RESPONSE
 
     @permissions(READ_PERM)
     async def handle_exists_stream(self, request):
-- 
2.54.0

[bitbake-devel][PATCH v2 08/10] hashserv: server: Use streaming and queue API for upstream exist queries

[Joshua Watt]

Reworks the "unihash-exists" handler to use the new server queue API and
client streaming API to efficiently stream requests to the upstream
server instead of having to wait for a roundtrip on the requests.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/server.py | 31 ++++++++++++++++++++++++++-----
 lib/hashserv/tests.py  |  5 +++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py
index d0e6f23fc..0153730fa 100644
--- a/lib/hashserv/server.py
+++ b/lib/hashserv/server.py
@@ -522,17 +522,38 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection):
 
     @permissions(READ_PERM)
     async def handle_exists_stream(self, request):
-        async def handler(l):
+        async def exists_handler(l):
             if await self.db.unihash_exists(l):
                 return "true"
+            return "false"
+
+        if not self.upstream_client:
+            return await self._stream_handler(exists_handler)
 
-            if self.upstream_client is not None:
-                if await self.upstream_client.unihash_exists(l):
+        async with self.upstream_client.unihash_exists_stream() as stream:
+
+            async def get_local_result(m):
+                if await self.db.unihash_exists(m):
                     return "true"
+                return None
 
-            return "false"
+            async def get_upstream_result(m):
+                exists = await stream.get_result()
+                return "true" if exists else "false"
 
-        return await self._stream_handler(handler)
+            queue = asyncio.Queue()
+            upstream = UpstreamQueue(
+                queue,
+                get_local_result,
+                stream.send_query,
+                get_upstream_result,
+            )
+
+            await bb.asyncrpc.TaskGroup.run(
+                self._stream_queue_handler(upstream.handler, queue),
+                upstream.process_results(),
+            )
+        return self.NO_RESPONSE
 
     async def report_readonly(self, data):
         method = data["method"]
diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index e24bdcacb..a7ce7425e 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -374,20 +374,25 @@ class HashEquivalenceCommonTests(object):
             nonlocal side_client
 
             # check upstream server
+            self.assertTrue(self.client.unihash_exists(unihash))
             self.assertClientGetHash(self.client, taskhash, unihash)
 
             # Hash should *not* be present on the side server
+            if old_sidehash and unihash != old_sidehash:
+                self.assertFalse(side_client.unihash_exists(unihash))
             self.assertClientGetHash(side_client, taskhash, old_sidehash)
 
             # Hash should be present on the downstream server, since it
             # will defer to the upstream server. This will trigger
             # the backfill in the downstream server
+            self.assertTrue(down_client.unihash_exists(unihash))
             self.assertClientGetHash(down_client, taskhash, unihash)
 
             # After waiting for the downstream client to finish backfilling the
             # task from the upstream server, it should appear in the side server
             # since the database is populated
             down_client.backfill_wait()
+            self.assertTrue(side_client.unihash_exists(unihash))
             self.assertClientGetHash(side_client, taskhash, unihash)
 
         # Basic report
-- 
2.54.0

[bitbake-devel][PATCH v2 09/10] hashserv: tests: Add more upstream tests

[Joshua Watt]

Adds a test to verify that the batch API works properly with an upstream
server and also a test to verify that if the upstream server restarts or
is temporarily disconnected that the client recovers properly.

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/tests.py | 152 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 141 insertions(+), 11 deletions(-)

diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index a7ce7425e..bb227c161 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -5,12 +5,13 @@
 # SPDX-License-Identifier: GPL-2.0-only
 #
 
-from . import create_server, create_client
+from . import create_server, create_client, create_async_client
 from .server import DEFAULT_ANON_PERMS, ALL_PERMISSIONS
 from bb.asyncrpc import InvokeError
 import hashlib
 import logging
 from bb import multiprocessing
+import asyncio
 import os
 import sys
 import tempfile
@@ -41,19 +42,15 @@ class HashEquivalenceTestSetup(object):
     server_index = 0
     client_index = 0
 
-    def start_server(self, dbpath=None, upstream=None, read_only=False, prefunc=server_prefunc, anon_perms=DEFAULT_ANON_PERMS, admin_username=None, admin_password=None):
+    def start_server(self, dbpath=None, upstream=None, read_only=False, prefunc=server_prefunc, anon_perms=DEFAULT_ANON_PERMS, admin_username=None, admin_password=None, addr=None):
         self.server_index += 1
+        if addr is None:
+            addr = self.get_server_addr(self.server_index)
+
         if dbpath is None:
             dbpath = self.make_dbpath()
 
-        def cleanup_server(server):
-            if server.process.exitcode is not None:
-                return
-
-            server.process.terminate()
-            server.process.join()
-
-        server = create_server(self.get_server_addr(self.server_index),
+        server = create_server(addr,
                                dbpath,
                                upstream=upstream,
                                read_only=read_only,
@@ -63,7 +60,7 @@ class HashEquivalenceTestSetup(object):
         server.dbpath = dbpath
 
         server.serve_as_process(prefunc=prefunc, args=(self.server_index,))
-        self.addCleanup(cleanup_server, server)
+        self.addCleanup(self.stop_server, server)
 
         return server
 
@@ -83,6 +80,13 @@ class HashEquivalenceTestSetup(object):
         self.server = self.start_server()
         return self.server.address
 
+    def stop_server(self, server):
+        if not server or server.process.exitcode is not None:
+            return
+
+        server.process.terminate()
+        server.process.join()
+
     def start_auth_server(self):
         auth_server = self.start_server(self.server.dbpath, anon_perms=[], admin_username="admin", admin_password="password")
         self.auth_server_address = auth_server.address
@@ -476,6 +480,132 @@ class HashEquivalenceCommonTests(object):
         self.assertEqual(result['taskhash'], taskhash9, 'Server failed to copy unihash from upstream')
         self.assertEqual(result['method'], self.METHOD)
 
+    def test_upstream_batch(self):
+        down_server = self.start_server(upstream=self.server.address)
+        down_client = self.start_client(down_server.address)
+
+        taskhash1 = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
+        outhash1 = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
+        unihash1 = '5b521d8a12683086cc08bc2c6d94a7a2dcff17eba53b9911e145d51164689380'
+        self.client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)
+
+        taskhash2 = "e3da00593d6a7fb435c7e2114976c59c5fd6d561"
+        outhash2 = "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a"
+        unihash2 = "7aebef07d66a8c0f92d0c4f65ec8b1fbb850a3693c53827b8774b64fa9a8a9fe"
+        self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
+
+        taskhash3 = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
+        outhash3 = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
+        unihash3 = 'a69ec97f5af2e21e1a1f9cc8896965515d5559425666f734e245a3d40cee33d9'
+        self.client.report_unihash(taskhash3, self.METHOD, outhash3, unihash3)
+
+        def query_generator():
+            yield unihash1
+            yield unihash2
+            yield unihash3
+            yield "cc74784b2c0ad5b378a6b783c74c518d2c46b8b52fba29cb39a8430d742440d7"
+
+        results = down_client.unihash_exists_batch(query_generator())
+        self.assertEqual(results, [True, True, True, False])
+
+    def test_upstream_interrupted(self):
+        up_server = self.start_server()
+        down_server = self.start_server(upstream=up_server.address)
+
+        def restart_upstream():
+            nonlocal up_server
+
+            self.stop_server(up_server)
+            up_server = self.start_server(addr=up_server.address, dbpath=up_server.dbpath)
+
+        # Report some hashes
+        with self.start_client(up_server.address) as up_client:
+            taskhash1 = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
+            outhash1 = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
+            unihash1 = '5b521d8a12683086cc08bc2c6d94a7a2dcff17eba53b9911e145d51164689380'
+            up_client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)
+
+            taskhash2 = "e3da00593d6a7fb435c7e2114976c59c5fd6d561"
+            outhash2 = "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a"
+            unihash2 = "7aebef07d66a8c0f92d0c4f65ec8b1fbb850a3693c53827b8774b64fa9a8a9fe"
+            up_client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
+
+            taskhash3 = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
+            outhash3 = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
+            unihash3 = 'a69ec97f5af2e21e1a1f9cc8896965515d5559425666f734e245a3d40cee33d9'
+            up_client.report_unihash(taskhash3, self.METHOD, outhash3, unihash3)
+
+        restart_upstream()
+
+        with self.start_client(up_server.address) as up_client:
+            # Verify that reported hashes are correct after restaring server
+            self.assertTrue(up_client.unihash_exists(unihash1))
+            self.assertClientGetHash(up_client, taskhash1, unihash1)
+
+            self.assertTrue(up_client.unihash_exists(unihash2))
+            self.assertClientGetHash(up_client, taskhash2, unihash2)
+
+        async def check_unihashes():
+            async with await create_async_client(down_server.address) as down_client:
+                async with down_client.unihash_exists_stream() as stream:
+                    await stream.send_query(unihash1)
+                    r = await stream.get_result()
+                    self.assertTrue(r)
+
+                    restart_upstream()
+
+                    await stream.send_query(unihash2)
+                    r = await stream.get_result()
+                    self.assertTrue(r)
+
+                    await stream.send_query(unihash3)
+                    r = await stream.get_result()
+                    self.assertTrue(r)
+
+        asyncio.run(check_unihashes())
+
+    def test_upstream_lost(self):
+        up_server = self.start_server()
+        down_server = self.start_server(upstream=up_server.address)
+
+        def restart_upstream():
+            nonlocal up_server
+
+            self.stop_server(up_server)
+            up_server = self.start_server(addr=up_server.address, dbpath=up_server.dbpath)
+
+        # Report some hashes
+        with self.start_client(up_server.address) as up_client:
+            taskhash1 = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
+            outhash1 = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
+            unihash1 = '5b521d8a12683086cc08bc2c6d94a7a2dcff17eba53b9911e145d51164689380'
+            up_client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)
+
+            taskhash2 = "e3da00593d6a7fb435c7e2114976c59c5fd6d561"
+            outhash2 = "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a"
+            unihash2 = "7aebef07d66a8c0f92d0c4f65ec8b1fbb850a3693c53827b8774b64fa9a8a9fe"
+            up_client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
+
+            taskhash3 = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
+            outhash3 = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
+            unihash3 = 'a69ec97f5af2e21e1a1f9cc8896965515d5559425666f734e245a3d40cee33d9'
+            up_client.report_unihash(taskhash3, self.METHOD, outhash3, unihash3)
+
+        async def check_unihashes():
+            async with await create_async_client(down_server.address) as down_client:
+                with self.assertRaises(ConnectionError):
+                    async with down_client.unihash_exists_stream() as stream:
+                        await stream.send_query(unihash1)
+                        r = await stream.get_result()
+                        self.assertTrue(r)
+
+                        self.stop_server(up_server)
+
+                        await stream.send_query(unihash2)
+                        r = await stream.get_result()
+
+        asyncio.run(check_unihashes())
+
     def test_unihash_exsits(self):
         taskhash, outhash, unihash = self.create_test_hash(self.client)
         self.assertTrue(self.client.unihash_exists(unihash))
-- 
2.54.0

[bitbake-devel][PATCH v2 10/10] hashserv: tests: Add test for upstream pipelining

[Joshua Watt]

Adds a test that verifies that pipelining to an upstream server works as
expected.

AI-Generated: Uses Cursor, Claude Sonnet 5
Co-authored-by: Michal Sieron <michal.sieron@nokia.com>
Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
---
 lib/hashserv/tests.py | 96 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py
index bb227c161..6201ce3bd 100644
--- a/lib/hashserv/tests.py
+++ b/lib/hashserv/tests.py
@@ -1693,6 +1693,102 @@ class TestHashEquivalenceTCPServer(HashEquivalenceTestSetup, HashEquivalenceComm
         # case it is more reliable to resolve the IP address explicitly.
         return socket.gethostbyname("localhost") + ":0"
 
+    def test_get_stream_upstream_pipelined(self):
+        # Verify that, with an upstream configured, the get-stream handler
+        # pipelines its upstream queries instead of doing one blocking
+        # round-trip per task. A latency proxy injects RTT between this server
+        # and its upstream; a serial (one round-trip per task) implementation
+        # could not beat N * RTT, so finishing far faster proves the queries
+        # are pipelined.
+        import asyncio
+
+        LATENCY = 0.01  # 10ms each direction => ~20ms round trip
+        upstream_host, upstream_port = self.server_address.rsplit(":", 1)
+        upstream_port = int(upstream_port)
+
+        ready = threading.Event()
+        proxy = {}
+
+        def run_proxy():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            async def handle(creader, cwriter):
+                ureader, uwriter = await asyncio.open_connection(upstream_host, upstream_port)
+
+                async def pipe(r, w):
+                    try:
+                        while True:
+                            data = await r.read(65536)
+                            if not data:
+                                break
+                            await asyncio.sleep(LATENCY)
+                            w.write(data)
+                            await w.drain()
+                    except Exception:
+                        pass
+                    finally:
+                        try:
+                            w.close()
+                        except Exception:
+                            pass
+
+                await bb.asyncrpc.TaskGroup.run(pipe(creader, uwriter), pipe(ureader, cwriter))
+
+            stop_event = asyncio.Event()
+            proxy["stop_event"] = stop_event
+            proxy["loop"] = loop
+
+            async def main():
+                server = await asyncio.start_server(handle, "127.0.0.1", 0)
+                proxy["port"] = server.sockets[0].getsockname()[1]
+                ready.set()
+                async with server:
+                    await stop_event.wait()
+
+            try:
+                loop.run_until_complete(main())
+            except Exception:
+                pass
+            finally:
+                loop.close()
+
+        proxy_thread = threading.Thread(target=run_proxy, daemon=True)
+        proxy_thread.start()
+        self.assertTrue(ready.wait(10), "latency proxy did not start")
+        self.addCleanup(proxy_thread.join, 10)
+        self.addCleanup(lambda: proxy["loop"].call_soon_threadsafe(proxy["stop_event"].set))
+        proxy_addr = "127.0.0.1:%d" % proxy["port"]
+
+        N = 400
+        expected = []
+        for i in range(N):
+            taskhash = hashlib.sha256(("task%d" % i).encode()).hexdigest()
+            outhash = hashlib.sha256(("out%d" % i).encode()).hexdigest()
+            unihash = hashlib.sha256(("uni%d" % i).encode()).hexdigest()
+            self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
+            expected.append((self.METHOD, taskhash, unihash))
+
+        # Downstream server with an EMPTY local DB whose upstream is the slow proxy.
+        down_server = self.start_server(upstream=proxy_addr)
+        down_client = self.start_client(down_server.address)
+
+        args = [(m, th) for (m, th, _uh) in expected]
+        start = time.time()
+        results = down_client.get_unihash_batch(args)
+        elapsed = time.time() - start
+
+        self.assertEqual(results, [uh for (_m, _th, uh) in expected])
+
+        # A serial (one round-trip per task) implementation cannot beat N*RTT.
+        # Allow a generous margin to avoid flakiness on loaded CI machines.
+        serial_lower_bound = N * 2 * LATENCY
+        self.assertLess(elapsed, serial_lower_bound / 4,
+                        "Upstream queries are not being pipelined "
+                        "(%.3fs for %d queries at %.0fms injected RTT)"
+                        % (elapsed, N, 2 * LATENCY * 1000))
+
+
 
 class TestHashEquivalenceWebsocketServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
     def setUp(self):
-- 
2.54.0