Retry Py4J on empty response (#551)

HyukjinKwon · web-flow · commit c0ef8fef0ed4 · 2025-01-15T10:27:04.000+09:00
This PR retries Py4J on empty response. The other end might not be dead but the JVM itself might keep running. In this case, new connection has to be made.
diff --git a/py4j-python/src/py4j/clientserver.py b/py4j-python/src/py4j/clientserver.py
@@ -537,7 +537,8 @@ def send_command(self, command):
                 # Happens when a the other end is dead. There might be an empty
                 # answer before the socket raises an error.
                 if answer.strip() == "":
-                    raise Py4JNetworkError("Answer from Java side is empty")
+                    raise Py4JNetworkError(
+                        "Answer from Java side is empty", when=proto.EMPTY_RESPONSE)
                 if answer.startswith(proto.RETURN_MESSAGE):
                     return answer[1:]
                 else:
@@ -560,6 +561,8 @@ def send_command(self, command):
                             proto.ERROR_RETURN_MESSAGE.encode("utf-8"))
         except Exception as e:
             logger.info("Error while receiving.", exc_info=True)
+            if isinstance(e, Py4JNetworkError) and e.when == proto.EMPTY_RESPONSE:
+                raise
             raise Py4JNetworkError(
                 "Error while sending or receiving", e, proto.ERROR_ON_RECEIVE)
 
diff --git a/py4j-python/src/py4j/java_gateway.py b/py4j-python/src/py4j/java_gateway.py
@@ -1049,8 +1049,12 @@ def send_command(self, command, retry=True, binary=False):
                     reset = True
                 connection.close(reset)
             if self._should_retry(retry, connection, pne):
+                if pne.when == proto.ERROR_ON_SEND or pne.when == proto.EMPTY_RESPONSE:
+                    # For empty response, just try once more because now we reset
+                    # the connection, and should work if the other end is alive.
+                    retry = False
                 logging.info("Exception while sending command.", exc_info=True)
-                response = self.send_command(command, binary=binary)
+                response = self.send_command(command, retry, binary=binary)
             else:
                 logging.exception(
                     "Exception while sending command.")
@@ -1107,7 +1111,8 @@ def _create_connection_guard(self, connection):
         return GatewayConnectionGuard(self, connection)
 
     def _should_retry(self, retry, connection, pne=None):
-        return pne and pne.when == proto.ERROR_ON_SEND
+        return retry and pne and (
+            pne.when == proto.ERROR_ON_SEND or pne.when == proto.EMPTY_RESPONSE)
 
     def close(self):
         """Closes all currently opened connections.
@@ -1254,10 +1259,12 @@ def send_command(self, command):
             # Happens when a the other end is dead. There might be an empty
             # answer before the socket raises an error.
             if answer.strip() == "":
-                raise Py4JNetworkError("Answer from Java side is empty")
+                raise Py4JNetworkError("Answer from Java side is empty", when=proto.EMPTY_RESPONSE)
             return answer
         except Exception as e:
             logger.info("Error while receiving.", exc_info=True)
+            if isinstance(e, Py4JNetworkError) and e.when == proto.EMPTY_RESPONSE:
+                raise
             raise Py4JNetworkError(
                 "Error while receiving", e, proto.ERROR_ON_RECEIVE)
 
diff --git a/py4j-python/src/py4j/protocol.py b/py4j-python/src/py4j/protocol.py
@@ -172,6 +172,7 @@
 # ERRORS
 ERROR_ON_SEND = "on_send"
 ERROR_ON_RECEIVE = "on_receive"
+EMPTY_RESPONSE = "empty_response"
 
 
 def escape_new_line(original):
