pikers
/
piker
3
0
Fork 0
Code Issues 2 Pull Requests 5 Packages Projects Releases Wiki Activity

Fix clients map typing annot

basic_orders
Tyler Goodlet 2021-01-15 20:57:25 -05:00
parent 5acd780eb6
commit 2bf95d7ec7
1 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
piker/brokers/api.py
View File

@ -28,6 +28,7 @@ from ..log import get_logger
log = get_logger(__name__) log = get_logger(__name__)
_clients: Dict[str, 'Client'] = {}
@asynccontextmanager @asynccontextmanager
async def get_cached_client( async def get_cached_client(
@ -39,29 +40,38 @@ async def get_cached_client(
If one has not been setup do it and cache it. If one has not been setup do it and cache it.
""" """
# check if a cached client is in the local actor's statespace global _clients
ss = tractor.current_actor().statespace
clients = ss.setdefault('clients', {'_lock': trio.Lock()}) clients = ss.setdefault('clients', {'_lock': trio.Lock()})
lock = clients['_lock'] lock = clients['_lock']
client = None client = None
try: try:
log.info(f"Loading existing `{brokername}` daemon") log.info(f"Loading existing `{brokername}` client")
async with lock: async with lock:
client = clients[brokername] client = clients[brokername]
client._consumers += 1 client._consumers += 1
yield client yield client
except KeyError: except KeyError:
log.info(f"Creating new client for broker {brokername}") log.info(f"Creating new client for broker {brokername}")
async with lock: async with lock:
brokermod = get_brokermod(brokername) brokermod = get_brokermod(brokername)
exit_stack = AsyncExitStack() exit_stack = AsyncExitStack()
client = await exit_stack.enter_async_context( client = await exit_stack.enter_async_context(
brokermod.get_client() brokermod.get_client()
) )
client._consumers = 0 client._consumers = 0
client._exit_stack = exit_stack client._exit_stack = exit_stack
clients[brokername] = client clients[brokername] = client
yield client yield client
finally: finally:
client._consumers -= 1 client._consumers -= 1
if client._consumers <= 0: if client._consumers <= 0: