PubSub functionality kills stream-based connections due to race conditions
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
The PubSub functionality is prone to race conditions and may bring down stream-based connections when it raises errors.
**Reproducer**
```python
from distributed import Event, Pub, Sub
from distributed.utils_test import freeze_batched_send, gen_cluster
@gen_cluster(client=True, nthreads=[("", 1)])
async def test_race(c, s, a):
sub = Sub("a")
in_event = Event()
in_event_2 = Event()
block_event = Event()
block_event_2 = Event()
def f(x, in_event, block_event, in_event_2, block_event_2):
pub = Pub("a")
in_event.set()
block_event.wait()
pub.put({"status": "OK"})
in_event_2.set()
block_event_2.wait()
del pub
return x
fut = c.submit(f, 1, in_event, block_event, in_event_2, block_event_2)
await in_event.wait()
with freeze_batched_send(s.client_comms[c.id]):
await block_event.set()
await in_event_2.wait()
del sub
await block_event_2.set()
await fut
await c.submit(lambda x: x, 2)
```
This reproducer kills the client connection due to call in a non-existent stream handler (`remove-pubpus-subscribers`) and also kills the client connection with
```
Traceback (most recent call last):
File "/Users/hendrikmakait/projects/dask/distributed/distributed/core.py", line 970, in _handle_comm
result = await result
^^^^^^^^^^^^
File "/Users/hendrikmakait/projects/dask/distributed/distributed/scheduler.py", line 5716, in add_client
await self.handle_stream(comm=comm, extra={"client": client})
File "/Users/hendrikmakait/projects/dask/distributed/distributed/core.py", line 1053, in handle_stream
handler(**merge(extra, msg))
File "/Users/hendrikmakait/projects/dask/distributed/distributed/pubsub.py", line 86, in remove_subscriber
self.client_subscribers[name].remove(client)
KeyError: 'Client-89725a46-27eb-11ef-9134-be79fecea867'
```
if the appropriate handler is used.
I suspect that there are more possible race conditions hidden in this code.
Contributor guide
Assessment
This issue has not been assessed yet.