worker: Extract wait_shutdown_started()

master/buildbot/worker/base.py

@@ -28,7 +28,6 @@ from buildbot.process.properties import Properties
 from buildbot.util import Notifier
 from buildbot.util import bytes2unicode
 from buildbot.util import service
-from buildbot.util.eventual import eventually
 
 
 @implementer(IWorker)
@@ -534,12 +533,7 @@ class AbstractWorker(service.BuildbotService):
     @defer.inlineCallbacks
     def _waitForCompleteShutdownImpl(self, conn):
         if conn:
-            d = defer.Deferred()
-
-            def _disconnected():
-                eventually(d.callback, None)
-            conn.notifyOnDisconnect(_disconnected)
-            yield d
+            yield conn.wait_shutdown_started()
             yield conn.waitShutdown()
         elif self._pending_conn_shutdown_notifier is not None:
             yield self._pending_conn_shutdown_notifier.wait()

master/buildbot/worker/protocols/base.py

@@ -13,8 +13,11 @@
 #
 # Copyright Buildbot Team Members
 
+from twisted.internet import defer
+
 from buildbot.util import service
 from buildbot.util import subscription
+from buildbot.util.eventual import eventually
 
 
 class Listener(service.ReconfigurableServiceMixin, service.AsyncMultiService):
@@ -36,8 +39,14 @@ class Connection:
                     v = proxyclass(v)
             newargs[k] = v
         return newargs
+
     # disconnection handling
 
+    def wait_shutdown_started(self):
+        d = defer.Deferred()
+        self.notifyOnDisconnect(lambda: eventually(d.callback, None))
+        return d
+
     def waitShutdown(self):
         return self._disconnectSubs.waitForDeliveriesToFinish()