`Task` is hashable, so key on it

tokio_backup
Tyler Goodlet 2021-08-08 19:58:12 -04:00
parent dfc4082ad2
commit e1e3e6918c
1 changed files with 13 additions and 16 deletions

View File

@ -48,15 +48,12 @@ class BroadcastReceiver(ReceiveChannel):
async def receive(self):
task: Task
task = current_task()
task: Task = current_task()
# check that task does not already have a value it can receive
# immediately and/or that it has lagged.
key = id(task)
# print(key)
try:
seq = self._subs[key]
seq = self._subs[task]
except KeyError:
raise RuntimeError(
f'Task {task.name} is not registerd as subscriber')
@ -69,12 +66,12 @@ class BroadcastReceiver(ReceiveChannel):
# decrement to the last value and expect
# consumer to either handle the ``Lagged`` and come back
# or bail out on it's own (thus un-subscribing)
self._subs[key] = self._queue.maxlen - 1
self._subs[task] = self._queue.maxlen - 1
# this task was overrun by the producer side
raise Lagged(f'Task {task.name} was overrun')
self._subs[key] -= 1
self._subs[task] -= 1
return value
if self._value_received is None:
@ -97,7 +94,7 @@ class BroadcastReceiver(ReceiveChannel):
subs = self._subs.copy()
# don't decerement the sequence # for this task since we
# already retreived the last value
subs.pop(key)
subs.pop(task)
for sub_key, seq in subs.items():
self._subs[sub_key] += 1
@ -109,10 +106,10 @@ class BroadcastReceiver(ReceiveChannel):
else:
await self._value_received.wait()
seq = self._subs[key]
seq = self._subs[task]
assert seq > -1, 'Internal error?'
self._subs[key] -= 1
self._subs[task] -= 1
return self._queue[0]
# @asynccontextmanager
@ -120,14 +117,14 @@ class BroadcastReceiver(ReceiveChannel):
def subscribe(
self,
) -> BroadcastReceiver:
key = id(current_task())
self._subs[key] = -1
task: task = current_task()
self._subs[task] = -1
# XXX: we only use this clone for closure tracking
clone = self._clones[key] = self._rx.clone()
clone = self._clones[task] = self._rx.clone()
try:
yield self
finally:
self._subs.pop(key)
self._subs.pop(task)
clone.close()
# TODO: do we need anything here?
@ -135,8 +132,8 @@ class BroadcastReceiver(ReceiveChannel):
# the underlying rx channel, but couldn't we just
# use ``.clone()``s trackign then?
async def aclose(self) -> None:
key = id(current_task())
await self._clones[key].aclose()
task: Task = current_task()
await self._clones[task].aclose()
def broadcast_receiver(