Repository: dask/distributed

Issue:
`distributed.cli.tests.test_dask_worker.test_signal_handling` flaky on py3.11
![image](https://user-images.githubusercontent.com/8629629/227206328-559f7f8c-dcd5-4d4b-bc4f-303dae9da5d6.png)

cc @graingert 


Discussion:
See https://github.com/dask/distributed/issues/7756#issue-1657157057 for some additional info
pasting the additional info here:

e.g. https://github.com/dask/distributed/pull/7752/checks?check_run_id=12549109243

1. test starts a Nanny process
2. client.wait_for_workers(1)
3. os.kill(SIGINT, nanny_pid)
4. but the nanny has not heard back from the worker yet

```
    2023-04-05 18:24:32,912 - distributed.worker - info -         registered to:      tcp://127.0.0.1:49849
    2023-04-05 18:24:32,912 - distributed.worker - info - -------------------------------------------------
    2023-04-05 18:24:32,915 - distributed.core - info - starting established connection to tcp://127.0.0.1:49849
    2023-04-05 18:24:32,942 - distributed._signals - info - received signal sigint (2)
    2023-04-05 18:24:32,943 - distributed.nanny - info - closing nanny at 'tcp://127.0.0.1:49852'. reason: nanny-close
    2023-04-05 18:24:32,956 - distributed.nanny - info - nanny asking worker to close. reason: nanny-close
    2023-04-05 18:24:32,957 - distributed.worker - info - stopping worker at tcp://127.0.0.1:49854. reason: nanny-close
    2023-04-05 18:24:32,960 - distributed.core - info - connection to tcp://127.0.0.1:49849 has been closed.
    2023-04-05 18:24:32,961 - distributed.nanny - info - worker closed
    2023-04-05 18:24:33,608 - distributed.nanny - warning - restarting worker
    
    traceback (most recent call last):
      file "/users/runner/work/distributed/distributed/distributed/utils.py", line 752, in wrapper
        return await func(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
      file "/users/runner/work/distributed/distributed/distributed/nanny.py", line 549, in _on_worker_exit
        await self.instantiate()
      file "/users/runner/work/distributed/distributed/distributed/nanny.py", line 442, in instantiate
        result = await self.process.start()
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^
      file "/users/runner/work/distributed/distributed/distributed/nanny.py", line 711, in start
        await self.process.start()
    asyncio.exceptions.cancellederror
```

Resolving patch:
diff --git a/distributed/core.py b/distributed/core.py
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -543,7 +543,8 @@ async def _close_on_failure(exc: Exception) -> None:
             except Exception as exc:
                 await _close_on_failure(exc)
                 raise RuntimeError(f"{type(self).__name__} failed to start.") from exc
-            self.status = Status.running
+            if self.status == Status.init:
+                self.status = Status.running
         return self
 
     async def __aenter__(self):
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1196,7 +1196,6 @@ async def _register_with_scheduler(self) -> None:
                 middle = (_start + _end) / 2
                 self._update_latency(_end - start)
                 self.scheduler_delay = response["time"] - middle
-                self.status = Status.running
                 break
             except OSError:
                 logger.info("Waiting to connect to: %26s", self.scheduler.address)
@@ -1207,18 +1206,20 @@ async def _register_with_scheduler(self) -> None:
             msg = response["message"] if "message" in response else repr(response)
             logger.error(f"Unable to connect to scheduler: {msg}")
             raise ValueError(f"Unexpected response from register: {response!r}")
-        else:
-            await asyncio.gather(
-                *(
-                    self.plugin_add(name=name, plugin=plugin)
-                    for name, plugin in response["worker-plugins"].items()
-                )
+
+        self.batched_stream.start(comm)
+        self.status = Status.running
+
+        await asyncio.gather(
+            *(
+                self.plugin_add(name=name, plugin=plugin)
+                for name, plugin in response["worker-plugins"].items()
             )
+        )
 
-            logger.info("        Registered to: %26s", self.scheduler.address)
-            logger.info("-" * 49)
+        logger.info("        Registered to: %26s", self.scheduler.address)
+        logger.info("-" * 49)
 
-        self.batched_stream.start(comm)
         self.periodic_callbacks["keep-alive"].start()
         self.periodic_callbacks["heartbeat"].start()
         self.loop.add_callback(self.handle_scheduler, comm)


Test patch: