[WEKAPP-476559] adding callback on blpop in wait

weka · Jan 27, 2025 · 9dec154 · 9dec154
1 parent 6e6bd3f
commit 9dec154
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/talker/command.py b/talker/command.py
@@ -20,7 +20,7 @@
     TalkerClientSendTimeout, TalkerServerTimeout, ClientCommandTimeoutError, CommandTimeoutError,
     CommandAbortedByReboot, CommandAbortedByOverflow, CommandOrphaned,
     CommandLineTimeout, CommandExecutionError, UnknownExitCode, TalkerError, CommandPidTimeoutError, CommandAlreadyDone,
-    HostIsNotResponsive
+    HostIsNotResponsive,NoResponseForRedisCommand
 )
 
 
@@ -613,9 +613,12 @@ def wait(self, for_ack=False):
 
             if self.alive_check:
                 blpop_timeout = min(blpop_timeout, self.alive_check_interval)
-
-            result = self.talker.reactor.blpop(
-                [self._exit_code_key, self._ack_key], timeout=blpop_timeout, _cmd_id=self.job_id)
+            try:
+                result = self.talker.reactor.blpop(
+                    [self._exit_code_key, self._ack_key], timeout=blpop_timeout, _cmd_id=self.job_id, _callback=self.on_sent)
+            except NoResponseForRedisCommand as e:
+                if self.is_sent:
+                    raise e
             if result is None:
                 self.check_client_timeout()