Use separate task to get initial PSRD from hub

Author: brunorijsman

TODO

@@ -1,6 +1,14 @@
-TODO: Add ETSI QKD get-status to system test
+TODO: (Re)start get PSRD task when pool runs low
+
+TODO: Remove fully consumed blocks
+
+TODO: Add pool thresholds to configuration (and document this)
 
-TODO: Replenish PSRD pool when it runs low
+TODO: Make dske.yaml the default configuration file, and add a --config option to change it
+
+TODO: Add test cases for replenishing pools
+
+TODO: Add ETSI QKD get-status to system test
 
 TODO: Cleanup Shamir's Secret Sharing (SSS) code (no back-and-forth conversion to RawShare)
 
@@ -68,4 +76,4 @@ TODO: Avoid usage: __main__.py in --help output for client / hub
 
 TODO: Add site map to Google Search Console
 
-TODO: Consistently use httpx for client side (not requests for synchronous calls)
+TODO: Consistently use httpx for client side (not requests for synchronous calls)

client/http_client.py

@@ -77,7 +77,7 @@ async def get(
         try:
             response = await self._httpx_client.get(url, params=params, auth=auth)
         except httpx.HTTPError as exc:
-            LOGGER.error(f"Call GET {url} exception {str(exc)}")
+            LOGGER.error(f"Call GET {exc.request.url} exception {str(exc)}")
             raise exceptions.HTTPError(
                 method="GET",
                 url=url,
@@ -86,7 +86,7 @@ async def get(
                 exception=str(exc),
             ) from exc
         if response.status_code != 200:
-            LOGGER.error(f"Call GET {url} {response.status_code}")
+            LOGGER.error(f"Call GET {response.request.url} {response.status_code}")
             raise exceptions.HTTPError(
                 method="GET",
                 url=url,
@@ -95,7 +95,7 @@ async def get(
                 status_code=response.status_code,
                 response=response.content,
             )
-        LOGGER.info(f"Call GET {url} {response.status_code}")
+        LOGGER.info(f"Call GET {response.request.url} {response.status_code}")
         if api_response_class is None:
             return None
         try:

client/peer_hub.py

@@ -19,10 +19,29 @@
 from common.utils import bytes_to_str, str_to_bytes
 from .http_client import HttpClient
 
-# TODO: Decide on logic on how the PSRD block size is decided. Does the client decide? Does
-#       the hub decide?
-_PSRD_BLOCK_SIZE_IN_BYTES = 1000
-_FAILED_REQUEST_RETRY_DELAY_IN_SECONDS = 1.0
+# TODO: Make the following configurable.
+
+_START_REQUEST_PSRD_THRESHOLD = 500
+"""
+Start requesting more PSRD blocks from the hub when the amount of PSRD in the pool falls below this
+threshold.
+"""
+
+_STOP_REQUEST_PSRD_THRESHOLD = 2000
+"""
+Stop requesting more PSRD blocks from the hub when the amount of PSRD in the pool rises above or
+equal to this threshold.
+"""
+
+_GET_PSRD_BLOCK_SIZE = 2000
+"""
+When requesting more PSRD blocks from the hub, request blocks of this size.
+"""
+
+_GET_PSRD_RETRY_DELAY = 1.0
+"""
+If a get PSRD request to the hub fails, wait this many seconds before retrying.
+"""
 
 
 class PeerHub:
@@ -34,6 +53,7 @@ class PeerHub:
     _http_client: HttpClient
     _base_url: str
     _startup_task: asyncio.Task | None = None
+    _request_psrd_task: asyncio.Task | None = None
     _registered: bool
     _local_pool: Pool
     _peer_pool: Pool
@@ -45,6 +65,8 @@ def __init__(self, client, base_url):
         self._base_url = base_url
         if self._base_url.endswith("/"):
             self._base_url = self._base_url[:-1]
+        self._startup_task = None
+        self._request_psrd_task = None
         self._registered = False
         self._local_pool = Pool(Pool.Owner.LOCAL)
         self._peer_pool = Pool(Pool.Owner.PEER)
@@ -82,7 +104,6 @@ def start(self) -> None:
         """
         assert self._startup_task is None
         self._startup_task = asyncio.create_task(self.start_task())
-        # TODO: Do we need a done_callback? Set start_task to None?
 
     async def start_task(self) -> None:
         """
@@ -91,16 +112,17 @@ async def start_task(self) -> None:
         - Request the initial block of Pre-Shared Random Data (PSRD) from the hub (also periodically
           retrying if it fails).
         """
+        LOGGER.info(f"Begin start task for peer hub at {self._base_url}")
         try:
             while not await self.attempt_registration():
-                await asyncio.sleep(_FAILED_REQUEST_RETRY_DELAY_IN_SECONDS)
-            for owner in (Pool.Owner.LOCAL, Pool.Owner.PEER):
-                while not await self.attempt_request_psrd(owner):
-                    await asyncio.sleep(_FAILED_REQUEST_RETRY_DELAY_IN_SECONDS)
+                await asyncio.sleep(_GET_PSRD_RETRY_DELAY)
         except asyncio.CancelledError:
-            # The task is cancelled when the client is shut down before startup is complete.
-            # TODO: Do we need to do anything here? The un-registration is done elsewhere
-            pass
+            self._startup_task = None
+            LOGGER.info(f"Cancel start task for peer hub at {self._base_url}")
+            return
+        self.start_request_psrd_task_if_needed()
+        self._startup_task = None
+        LOGGER.info(f"Finish start task for peer hub at {self._base_url}")
 
     async def attempt_registration(self) -> bool:
         """
@@ -125,7 +147,59 @@ async def unregister(self) -> None:
         """
         Unregister this client from the peer hub.
         """
-        # TODO: Implement this
+        # TODO: Implement this and call it from somewhere
+
+    def start_request_psrd_task_if_needed(self) -> None:
+        """
+        Start the request PSRD task if needed.
+        """
+        if self._request_psrd_task is None:
+            if self.at_least_one_pool_below_start_threshold():
+                self._request_psrd_task = asyncio.create_task(self.request_psrd_task())
+
+    def at_least_one_pool_below_start_threshold(self) -> bool:
+        """
+        Check if the number of bytes available at least one of the pools is below the threshold
+        for starting to get more PSRD from the hub.
+        """
+        return (
+            self._local_pool.bytes_available < _START_REQUEST_PSRD_THRESHOLD
+            or self._peer_pool.bytes_available < _START_REQUEST_PSRD_THRESHOLD
+        )
+
+    def all_pools_above_stop_threshold(self) -> bool:
+        """
+        Check if the number of bytes available in all of the pools is above the threshold for
+        stopping to get more PSRD from the hub.
+        """
+        return (
+            self._local_pool.bytes_available >= _STOP_REQUEST_PSRD_THRESHOLD
+            and self._peer_pool.bytes_available >= _STOP_REQUEST_PSRD_THRESHOLD
+        )
+
+    async def request_psrd_task(self) -> None:
+        """
+        Task for requesting Pre-Shared Random Data (PSRD) from the peer hub. This tasks runs as long
+        as the local pool or the peer pool need more PSRD.
+        """
+        LOGGER.info(f"Begin request PSRD task for peer hub {self._hub_name}")
+        try:
+            while not self.all_pools_above_stop_threshold():
+                need_delay = False
+                if self._local_pool.bytes_available < _STOP_REQUEST_PSRD_THRESHOLD:
+                    if not await self.attempt_request_psrd(Pool.Owner.LOCAL):
+                        need_delay = True
+                if self._peer_pool.bytes_available < _STOP_REQUEST_PSRD_THRESHOLD:
+                    if not await self.attempt_request_psrd(Pool.Owner.PEER):
+                        need_delay = True
+                if need_delay:
+                    await asyncio.sleep(_GET_PSRD_RETRY_DELAY)
+        except asyncio.CancelledError:
+            self._request_psrd_task = None
+            LOGGER.info(f"Cancel request PSRD task for peer hub {self._hub_name}")
+            return
+        self._request_psrd_task = None
+        LOGGER.info(f"Finish request PSRD task for peer hub {self._hub_name}")
 
     async def attempt_request_psrd(self, owner: Pool.Owner) -> bool:
         """
@@ -144,7 +218,7 @@ async def attempt_request_psrd(self, owner: Pool.Owner) -> bool:
         params = {
             "client_name": self._client.name,
             "pool_owner": pool_owner_str,
-            "size": _PSRD_BLOCK_SIZE_IN_BYTES,
+            "size": _GET_PSRD_BLOCK_SIZE,
         }
         try:
             api_block = await self._http_client.get(url, params, APIBlock)

common/block.py

@@ -45,6 +45,13 @@ def uuid(self):
         """
         return self._block_uuid
 
+    @property
+    def bytes_available_for_allocation(self):
+        """
+        Return the number of bytes available for allocation in the block.
+        """
+        return self._size - self._allocated.count()
+
     def to_mgmt(self):
         """
         Get the management status.

common/pool.py

@@ -48,6 +48,13 @@ def owner(self) -> Owner:
         """
         return self._owner
 
+    @property
+    def bytes_available(self):
+        """
+        Return the number of bytes available for allocation in the pool.
+        """
+        return sum(block.bytes_available_for_allocation for block in self._blocks)
+
     def to_mgmt(self) -> dict:
         """
         Get the management status.