diff mercurial/sshpeer.py @ 52990:5ec596c91086

branching: merge stable into default
author Rapha?l Gom?s <rgomes@octobus.net>
date Wed, 19 Feb 2025 23:18:17 +0100
parents 24ee91ba9aa8 f6b30f4b5e07
children
line wrap: on
line diff
--- a/mercurial/sshpeer.py	Fri Feb 03 11:01:23 2023 +0100
+++ b/mercurial/sshpeer.py	Wed Feb 19 23:18:17 2025 +0100
@@ -11,6 +11,8 @@
 import typing
 import uuid
 
+from typing import Callable, Optional
+
 from .i18n import _
 from . import (
     error,
@@ -53,6 +55,31 @@
                 display(_(b"remote: "), l, b'\n')
 
 
+def _write_all(
+    write_once: Callable[[bytes], Optional[int]],
+    data: bytes,
+) -> Optional[int]:
+    """write data with a non blocking function
+
+    In case not all data were written, keep writing until everything is
+    written.
+    """
+    to_write = len(data)
+    written = write_once(data)
+    if written is None:
+        written = 0
+    if written < to_write:
+        data = memoryview(data)
+        while written < to_write:
+            wrote = write_once(data[written:])
+            # XXX if number of written bytes is "None", the destination is
+            # full. Some `select` call would be better than the current active
+            # polling.
+            if wrote is not None:
+                written += wrote
+    return written
+
+
 class doublepipe:
     """Operate a side-channel pipe in addition of a main one
 
@@ -97,9 +124,14 @@
             act = fds
         return (self._main.fileno() in act, self._side.fileno() in act)
 
-    def write(self, data):
+    def _write_once(self, data: bytes) -> Optional[int]:
+        """Write as much data as possible in a non blocking way"""
         return self._call(b'write', data)
 
+    def write(self, data: bytes) -> Optional[int]:
+        """write all data in a blocking way"""
+        return _write_all(self._write_once, data)
+
     def read(self, size):
         r = self._call(b'read', size)
         if size != 0 and not r:
@@ -130,6 +162,8 @@
         # data can be '' or 0
         if (data is not None and not data) or self._main.closed:
             _forwardoutput(self._ui, self._side)
+            if methname == b'write':
+                return 0
             return b''
         while True:
             mainready, sideready = self._wait()
@@ -314,7 +348,7 @@
         ui.debug(b'sending hello command\n')
         ui.debug(b'sending between command\n')
 
-        stdin.write(b''.join(handshake))
+        _write_all(stdin.write, b''.join(handshake))
         stdin.flush()
     except OSError:
         badresponse()