ISSUE-922: chore(tests/containers): implement retry if port-forwarding fails

Author: jiridanek

tests/containers/kubernetes_utils.py

@@ -225,7 +225,7 @@ def deploy(self, container_name: str) -> None:
         assert len(pod_name.items) == 1
         pod: kubernetes.client.models.v1_pod.V1Pod = pod_name.items[0]
 
-        p = socket_proxy.SocketProxy(exposing_contextmanager(core_v1_api, pod), "localhost", 0)
+        p = socket_proxy.SocketProxy(lambda: exposing_contextmanager(core_v1_api, pod), "localhost", 0)
         t = threading.Thread(target=p.listen_and_serve_until_canceled)
         t.start()
         self.tf.defer(t, lambda thread: thread.join())

tests/containers/socket_proxy.py

@@ -6,7 +6,8 @@
 import select
 import threading
 import subprocess
-import typing
+import time
+from typing import Callable, ContextManager
 
 from tests.containers.cancellation_token import CancellationToken
 
@@ -50,7 +51,7 @@ def stop(self):
 class SocketProxy:
     def __init__(
             self,
-            remote_socket_factory: typing.ContextManager[socket.socket],
+            remote_socket_factory: Callable[..., ContextManager[socket.socket]],
             local_host: str = "localhost",
             local_port: int = 0,
             buffer_size: int = 4096
@@ -96,27 +97,32 @@ def get_actual_port(self) -> int:
         return self.server_socket.getsockname()[1]
 
     def _handle_client(self, client_socket):
-        with client_socket as _, self.remote_socket_factory as remote_socket:
-            while True:
-                readable, _, _ = select.select([client_socket, remote_socket, self.cancellation_token], [], [])
-
-                if self.cancellation_token.cancelled:
-                    break
-
-                if client_socket in readable:
-                    data = client_socket.recv(self.buffer_size)
-                    if not data:
-                        break
-                    remote_socket.send(data)
-
-                if remote_socket in readable:
-                    data = remote_socket.recv(self.buffer_size)
-                    if not data:
-                        break
-                    client_socket.send(data)
-
-
-if __name__ == "__main__":
+        with client_socket as _:
+            while not self.cancellation_token.cancelled:
+                try:
+                    with self.remote_socket_factory() as remote_socket:
+                        while not self.cancellation_token.cancelled:
+                            readable, _, _ = select.select([client_socket, remote_socket, self.cancellation_token], [], [])
+
+                            if client_socket in readable:
+                                data = client_socket.recv(self.buffer_size)
+                                if not data:
+                                    return
+                                remote_socket.send(data)
+
+                            if remote_socket in readable:
+                                data = remote_socket.recv(self.buffer_size)
+                                if not data:
+                                    return
+                                client_socket.send(data)
+                except ConnectionResetError as e:
+                    # data = remote_socket.recv(self.buffer_size) may fail like this
+                    # usually it happens if the pod has not been fully up, so retry is in order
+                    logging.info("failed to read, will try again", exc_info=e)
+                    time.sleep(2)
+
+
+def main() -> None:
     """Sample application to show how this can work."""
 
 
@@ -161,7 +167,7 @@ def get_actual_port(self):
         server.join()
 
 
-    proxy = SocketProxy(remote_socket_factory(), "localhost", 0)
+    proxy = SocketProxy(remote_socket_factory, "localhost", 0)
     thread = threading.Thread(target=proxy.listen_and_serve_until_canceled)
     thread.start()
 
@@ -171,3 +177,7 @@ def get_actual_port(self):
     print(client_socket.recv(1024))  # prints Hello World
 
     thread.join()
+
+
+if __name__ == "__main__":
+    main()