diff --git a/peerix/__main__.py b/peerix/__main__.py index accd7b6..22352b6 100644 --- a/peerix/__main__.py +++ b/peerix/__main__.py @@ -14,6 +14,7 @@ parser = argparse.ArgumentParser(description="Peerix nix binary cache.") parser.add_argument("--verbose", action="store_const", const=logging.DEBUG, default=logging.INFO, dest="loglevel") parser.add_argument("--port", default=12304, type=int) parser.add_argument("--private-key", required=False) +parser.add_argument("--timeout", type=int, default=50) def run(): args = parser.parse_args() @@ -23,14 +24,14 @@ def run(): logging.basicConfig(level=args.loglevel) uvloop.install() - asyncio.run(main(args.port)) + asyncio.run(main(args.port, args.timeout / 1000.0)) -async def main(port: int): +async def main(port: int, timeout: float): config = Config() config.bind = [f"0.0.0.0:{port}"] - async with setup_stores(port): + async with setup_stores(port, timeout): await serve(app, config) diff --git a/peerix/app.py b/peerix/app.py index 719a2b6..e9e61dd 100644 --- a/peerix/app.py +++ b/peerix/app.py @@ -13,12 +13,12 @@ from peerix.prefix import PrefixStore @contextlib.asynccontextmanager -async def setup_stores(local_port: int): +async def setup_stores(local_port: int, timeout: float): global l_access, r_access async with local() as l: l_access = PrefixStore("local/nar", l) lp = PrefixStore("local", l) - async with remote(lp, local_port, "0.0.0.0", lp.prefix) as r: + async with remote(lp, local_port, "0.0.0.0", lp.prefix, timeout) as r: r_access = PrefixStore("v2/remote", r) yield diff --git a/peerix/remote.py b/peerix/remote.py index 3b9a246..ba01f0d 100644 --- a/peerix/remote.py +++ b/peerix/remote.py @@ -47,14 +47,16 @@ class DiscoveryProtocol(asyncio.DatagramProtocol, Store): session: aiohttp.ClientSession local_port: int prefix: str + timeout: float - def __init__(self, store: Store, session: aiohttp.ClientSession, local_port: int, prefix: str): + def __init__(self, store: Store, session: aiohttp.ClientSession, local_port: int, prefix: str, timeout: float): self.idx = 0 self.waiters = {} self.store = store self.session = session self.local_port = local_port self.prefix = prefix + self.timeout = timeout def connection_made(self, transport): self.transport = transport @@ -114,7 +116,7 @@ class DiscoveryProtocol(asyncio.DatagramProtocol, Store): try: # This must have a short timeout so it does not noticably slow down # querying of other caches. - port, url, addr = await asyncio.wait_for(fut, 0.05) + port, url, addr = await asyncio.wait_for(fut, self.timeout) except asyncio.TimeoutError: logging.debug(f"No response for {hsh}") return None @@ -164,11 +166,11 @@ class DiscoveryProtocol(asyncio.DatagramProtocol, Store): @contextlib.asynccontextmanager -async def remote(store: Store, local_port: int, local_addr: str="0.0.0.0", prefix: str="local"): +async def remote(store: Store, local_port: int, local_addr: str="0.0.0.0", prefix: str="local", timeout: float = 0.05): protocol: DiscoveryProtocol async with aiohttp.ClientSession() as session: _, protocol = await asyncio.get_running_loop().create_datagram_endpoint( - lambda: DiscoveryProtocol(store, session, local_port, prefix), + lambda: DiscoveryProtocol(store, session, local_port, prefix, timeout), local_addr=(local_addr, local_port), family=socket.AF_INET, allow_broadcast=True