address comment

Author: mainred

src/aks-agent/HISTORY.rst

@@ -13,7 +13,7 @@ Pending
 +++++++
 
 1.0.0b14
-* Fix: chunk the data for output to fix the eror on MacOS "BlockingIOError: [Errno 35] write could not complete without blocking"
+* Fix: chunk the data for output to fix the error on macos "BlockingIOError: [Errno 35] write could not complete without blocking"
 * Fix: gracefully handle the connection reset error
 
 1.0.0b13

src/aks-agent/azext_aks_agent/agent/k8s/pod_exec.py

@@ -33,6 +33,10 @@
 
 logger = get_logger(__name__)
 
+# WebSocket buffer size - matches Kubernetes client-go implementation
+# Reference: https://github.com/kubernetes/client-go/blob/master/transport/websocket/roundtripper.go#L67
+WEBSOCKET_BUFFER_SIZE = 32 * 1024  # 32 KiB
+
 
 def _get_terminal_size() -> Tuple[int, int]:
     """
@@ -165,23 +169,33 @@ def _set_terminal_raw_mode():
         return fd, old_settings
 
 
-def _restore_terminal_mode(fd_or_handle, old_settings):
+def _restore_terminal_mode(fd_or_handle, old_settings, windows_console_state=None):
     """
     Restore terminal to original mode.
 
     Args:
         fd_or_handle: File descriptor (Unix) or handle (Windows)
         old_settings: Original terminal settings to restore
+        windows_console_state: Windows console state tuple (output_cp, input_cp, stdout_mode, stdout_handle)
     """
     try:
         if IS_WINDOWS:
             import ctypes
             kernel32 = ctypes.windll.kernel32
+
+            # Restore terminal raw mode
             kernel32.SetConsoleMode(fd_or_handle, old_settings)
+
+            # Restore Windows console settings (code pages and VT100 mode)
+            if windows_console_state is not None:
+                output_cp, input_cp, stdout_mode, stdout_handle = windows_console_state
+                kernel32.SetConsoleOutputCP(output_cp)
+                kernel32.SetConsoleCP(input_cp)
+                kernel32.SetConsoleMode(stdout_handle, stdout_mode)
         else:
             termios.tcsetattr(fd_or_handle, termios.TCSADRAIN, old_settings)
-    except (NameError, OSError, IOError):
-        pass
+    except (NameError, OSError, IOError) as e:
+        logger.debug("Failed to restore terminal mode: %s", e)
 
 
 def _is_blocking_error(error):
@@ -196,7 +210,7 @@ def _is_blocking_error(error):
     """
     import errno
     err_code = getattr(error, 'errno', None) or getattr(error, 'winerror', None)
-    return err_code in (errno.EAGAIN, errno.EWOULDBLOCK, 35) if err_code else False
+    return err_code in (errno.EAGAIN, errno.EWOULDBLOCK) if err_code else False
 
 
 def _is_connection_reset_error(error):
@@ -212,9 +226,9 @@ def _is_connection_reset_error(error):
     """
     import errno
     err_code = getattr(error, 'errno', None) or getattr(error, 'winerror', None)
-    # Unix: errno 104 (ECONNRESET), errno 32 (EPIPE)
+    # Unix: ECONNRESET, EPIPE
     # Windows: WinError 10054 (WSAECONNRESET)
-    return err_code in (errno.ECONNRESET, errno.EPIPE, 104, 32, 10054) if err_code else False
+    return err_code in (errno.ECONNRESET, errno.EPIPE, 10054) if err_code else False
 
 
 def _heartbeat_worker(exec_stream, stop_event):
@@ -278,10 +292,17 @@ def exec_command_in_pod(pod_name: str, command: List[str],  # pylint: disable=to
     original_sigwinch = None
     fd = None
     fl = None
+    cleanup_done = False
+    windows_console_state = None
 
     def cleanup():
         """Cleanup function to ensure proper resource cleanup."""
-        nonlocal terminal_state, original_sigwinch, fd, fl, resp
+        nonlocal terminal_state, original_sigwinch, fd, fl, resp, cleanup_done, windows_console_state
+
+        # Prevent duplicate cleanup
+        if cleanup_done:
+            return
+        cleanup_done = True
 
         # Restore signal handler
         if original_sigwinch and not IS_WINDOWS:
@@ -290,9 +311,9 @@ def cleanup():
             except (ValueError, OSError):
                 pass
 
-        # Restore terminal mode
+        # Restore terminal mode and Windows console settings
         if terminal_state is not None:
-            _restore_terminal_mode(terminal_state[0], terminal_state[1])
+            _restore_terminal_mode(terminal_state[0], terminal_state[1], windows_console_state)
 
         # Restore stdin to blocking mode
         if not IS_WINDOWS and fd is not None and fl is not None:
@@ -309,10 +330,10 @@ def cleanup():
                 logger.debug("Error closing WebSocket connection: %s", e)
 
     # Register cleanup for SIGTERM
-    def signal_handler(signum, frame):
+    def signal_handler(signum, _frame):
         logger.info("Received signal %d, cleaning up...", signum)
-        cleanup()
-        sys.exit(0)
+        # Raise SystemExit to trigger finally block and normal cleanup
+        raise SystemExit(0)
 
     original_sigterm = None
     if hasattr(signal, 'SIGTERM'):
@@ -359,19 +380,24 @@ def signal_handler(signum, frame):
             if IS_WINDOWS:
                 import ctypes
                 kernel32 = ctypes.windll.kernel32
+
+                # Save original console settings
+                original_output_cp = kernel32.GetConsoleOutputCP()
+                original_input_cp = kernel32.GetConsoleCP()
+                STD_OUTPUT_HANDLE = -11
+                stdout_handle = kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
+                original_mode = ctypes.c_uint32()
+                kernel32.GetConsoleMode(stdout_handle, ctypes.byref(original_mode))
+                windows_console_state = (original_output_cp, original_input_cp,
+                                         original_mode.value, stdout_handle)
+
                 # Set console output code page to UTF-8 (65001)
                 kernel32.SetConsoleOutputCP(65001)
                 # Set console input code page to UTF-8
                 kernel32.SetConsoleCP(65001)
                 # Enable VT100 processing for ANSI escape sequences
-                # Get stdout handle
-                STD_OUTPUT_HANDLE = -11
-                stdout_handle = kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
-                mode = ctypes.c_uint32()
-                kernel32.GetConsoleMode(stdout_handle, ctypes.byref(mode))
-                # Enable ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x0004)
                 ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
-                kernel32.SetConsoleMode(stdout_handle, mode.value | ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+                kernel32.SetConsoleMode(stdout_handle, original_mode.value | ENABLE_VIRTUAL_TERMINAL_PROCESSING)
 
             # Put terminal in raw mode to disable local echo
             if tty:
@@ -426,16 +452,13 @@ def resize_handler(signum, frame):
                 # Handle stdout
                 if resp.peek_stdout():
                     stdout_data = resp.read_stdout()
-                    # Use the same buffer size as Kubernetes client-go (32 Kib)
-                    # Ref: https://github.com/kubernetes/client-go/blob/master/transport/websocket/roundtripper.go#L67
-                    # "BlockingIOError: [Errno 35] write could not complete without blocking", which is caused by wrting
-                    # to full pipe buffer and is easily reproducible on macOS, whose default output buffer size is 64 Kib
-                    # whose default output buffer size is 64 Kib. but using 32Kib here for consistency
-                    _BUF = 32 * 1024
                     data = stdout_data.encode()
 
-                    for start in range(0, len(data), _BUF):
-                        chunk = data[start: start + _BUF]
+                    # Write in chunks to avoid blocking on full pipe buffer
+                    # This prevents "BlockingIOError: [Errno 35] write could not complete without blocking"
+                    # which is easily reproducible on macOS (default pipe buffer: 64 KiB)
+                    for start in range(0, len(data), WEBSOCKET_BUFFER_SIZE):
+                        chunk = data[start: start + WEBSOCKET_BUFFER_SIZE]
 
                         while True:
                             try:
@@ -444,7 +467,7 @@ def resize_handler(signum, frame):
                             except BlockingIOError as exc:
                                 if exc.errno not in (errno.EAGAIN, errno.EWOULDBLOCK):
                                     raise                              # unexpected error
-                                time.sleep(0)                          # yield → let LM Studio drain pipe
+                                time.sleep(0)                          # yield to let the system drain the pipe buffer
 
                 # Handle stderr
                 if resp.peek_stderr():
@@ -462,8 +485,6 @@ def resize_handler(signum, frame):
                             if data:
                                 try:
                                     resp.write_stdin(data)
-                                except BlockingIOError as e:
-                                    logger.debug("stdin write blocked: %s", e)
                                 except OSError as e:
                                     if _is_blocking_error(e):
                                         logger.debug("stdin write blocked: %s", e)
@@ -475,12 +496,15 @@ def resize_handler(signum, frame):
                     else:
                         # Unix/Linux/macOS: Use select for non-blocking input
                         if select.select([sys.stdin], [], [], 0)[0]:
-                            data = sys.stdin.read()
+                            # Read up to 1024 bytes to avoid blocking
+                            # Even with O_NONBLOCK set, stdin.read() without args can block
+                            try:
+                                data = os.read(sys.stdin.fileno(), 1024).decode('utf-8', errors='replace')
+                            except BlockingIOError:
+                                data = None
                             if data:
                                 try:
                                     resp.write_stdin(data)
-                                except BlockingIOError as e:
-                                    logger.debug("stdin write blocked: %s", e)
                                 except OSError as e:
                                     if _is_blocking_error(e):
                                         logger.debug("stdin write blocked: %s", e)