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

@@ -4,9 +4,10 @@
 import logging
 import socket
 import select
+import struct
 import threading
 import subprocess
-import typing
+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
@@ -81,9 +82,14 @@ def listen_and_serve_until_canceled(self):
         Handles at most one client at a time. """
         try:
             while not self.cancellation_token.cancelled:
-                client_socket, addr = self.server_socket.accept()
-                logging.info(f"Accepted connection from {addr[0]}:{addr[1]}")
-                self._handle_client(client_socket)
+                readable, _, _ = select.select([self.server_socket, self.cancellation_token], [], [])
+
+                # ISSUE-922: socket.accept() blocks, so if cancel() did not come very fast, we'd loop over and block
+                if self.server_socket in readable:
+                    client_socket, addr = self.server_socket.accept()
+                    logging.info(f"Accepted connection from {addr[0]}:{addr[1]}")
+                    # handle client synchronously, which means that there can be at most one at a time
+                    self._handle_client(client_socket)
         except Exception as e:
             logging.exception(f"Proxying failed to listen", exc_info=e)
             raise
@@ -96,27 +102,38 @@ 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:
+        with client_socket as _, self.remote_socket_factory() as remote_socket:
+            while not self.cancellation_token.cancelled:
                 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)
+                    try:
+                        data = remote_socket.recv(self.buffer_size)
+                    except ConnectionResetError:
+                        # ISSUE-922: it seems best to propagate the error and let the client retry
+                        # alternatively it would be necessary to resend anything already received from client_socket
+                        _rst_socket(client_socket)
+                        break
                     if not data:
                         break
                     client_socket.send(data)
 
 
-if __name__ == "__main__":
+def _rst_socket(s: socket):
+    """Closing a SO_LINGER socket will RST it
+    https://stackoverflow.com/questions/46264404/how-can-i-reset-a-tcp-socket-in-python
+    """
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack('ii', 1, 0))
+    s.close()
+
+
+def main() -> None:
     """Sample application to show how this can work."""
 
 
@@ -161,13 +178,21 @@ 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()
 
-    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    client_socket.connect(("localhost", proxy.get_actual_port()))
+    for _ in range(2):
+        client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        client_socket.connect(("localhost", proxy.get_actual_port()))
 
-    print(client_socket.recv(1024))  # prints Hello World
+        print(client_socket.recv(1024))  # prints Hello World
+        print(client_socket.recv(1024))  # prints nothing
+        client_socket.close()
+    proxy.cancellation_token.cancel()
 
     thread.join()
+
+
+if __name__ == "__main__":
+    main()