Some refactoring ahead of smart protocol changes

breezy/bzr/smart/client.py

@@ -29,8 +29,7 @@ def __init__(self, medium, headers=None):
         """
         self._medium = medium
         if headers is None:
-            self._headers = {
-                b'Software version': breezy.__version__.encode('utf-8')}
+            self._headers = {'Software version': breezy.__version__}
         else:
             self._headers = dict(headers)
 
@@ -184,8 +183,7 @@ def _call(self, protocol_version):
             self.client._medium.reset()
             if not self._is_safe_to_send_twice():
                 raise
-            trace.warning('ConnectionReset reading response for %r, retrying'
-                          % (self.method,))
+            trace.warning('ConnectionReset reading response for %r, retrying', self.method)
             trace.log_exception_quietly()
             encoder, response_handler = self._construct_protocol(
                 protocol_version)
@@ -212,8 +210,8 @@ def _call_determining_protocol_version(self):
                 # we recognise the protocol version.
                 trace.warning(
                     'Server does not understand Bazaar network protocol %d,'
-                    ' reconnecting.  (Upgrade the server to avoid this.)'
-                    % (protocol_version,))
+                    ' reconnecting.  (Upgrade the server to avoid this.)',
+                    protocol_version)
                 self.client._medium.disconnect()
                 last_err = err
                 continue
@@ -275,7 +273,7 @@ def _send(self, protocol_version):
                 #   SmartClientRequestProtocolOne or Two, because they don't
                 #   support client-side body streams.
                 raise
-            trace.warning(f'ConnectionReset calling {self.method!r}, retrying')
+            trace.warning('ConnectionReset calling %r, retrying', self.method)
             trace.log_exception_quietly()
             encoder, response_handler = self._construct_protocol(
                 protocol_version)

breezy/bzr/smart/medium.py

@@ -37,7 +37,6 @@
 
 lazy_import(globals(), """
 import select
-import socket
 import weakref
 
 from breezy import (
@@ -922,10 +921,13 @@ def _accept_bytes(self, data):
         """See SmartClientStreamMedium.accept_bytes."""
         try:
             self._writeable_pipe.write(data)
-        except OSError as e:
-            if e.errno in (errno.EINVAL, errno.EPIPE):
+        except BrokenPipeError as ex:
+            raise ConnectionResetError(
+                "Error trying to write to subprocess") from ex
+        except OSError as ex:
+            if ex.errno == errno.EINVAL:
                 raise ConnectionResetError(
-                    "Error trying to write to subprocess", e)
+                    "Error trying to write to subprocess") from ex
             raise
         self._report_activity(len(data), 'write')
 
@@ -1110,6 +1112,7 @@ def __init__(self, host, port, base):
 
     def _ensure_connection(self):
         """Connect this medium if not already connected."""
+        import socket
         if self._connected:
             return
         if self._port is None:

breezy/bzr/smart/protocol.py

@@ -114,24 +114,12 @@ def set_headers(self, headers):
         raise NotImplementedError(self.set_headers)
 
 
-class SmartProtocolBase:
-    """Methods common to client and server."""
+def _encode_bulk_data(body):
+    """Encode body as a bulk data chunk."""
+    return b''.join((b'%d\n' % len(body), body, b'done\n'))
 
-    # TODO: this only actually accomodates a single block; possibly should
-    # support multiple chunks?
-    def _encode_bulk_data(self, body):
-        """Encode body as a bulk data chunk."""
-        return b''.join((b'%d\n' % len(body), body, b'done\n'))
 
-    def _serialise_offsets(self, offsets):
-        """Serialise a readv offset list."""
-        txt = []
-        for start, length in offsets:
-            txt.append(b'%d,%d' % (start, length))
-        return b'\n'.join(txt)
-
-
-class SmartServerRequestProtocolOne(SmartProtocolBase):
+class SmartServerRequestProtocolOne:
     """Server-side encoding and decoding logic for smart version 1."""
 
     def __init__(self, backing_transport, write_func, root_client_path='/',
@@ -229,7 +217,7 @@ def _send_response(self, response):
         if body is not None:
             if not isinstance(body, bytes):
                 raise ValueError(body)
-            data = self._encode_bulk_data(body)
+            data = _encode_bulk_data(body)
             self._write_func(data)
 
     def _write_protocol_version(self):
@@ -292,7 +280,7 @@ def _send_response(self, response):
             if response.body_stream is not None:
                 raise AssertionError(
                     'body_stream and body cannot both be set')
-            data = self._encode_bulk_data(response.body)
+            data = _encode_bulk_data(response.body)
             self._write_func(data)
         elif response.body_stream is not None:
             _send_stream(response.body_stream, self._write_func)
@@ -318,6 +306,14 @@ def _send_chunks(stream, write_func):
                 f'Chunks must be str or FailedSmartServerResponse, got {chunk!r}')
 
 
+def _serialise_offsets(offsets):
+    """Serialise a readv offset list."""
+    txt = []
+    for start, length in offsets:
+        txt.append(b'%d,%d' % (start, length))
+    return b'\n'.join(txt)
+
+
 class _NeedMoreBytes(Exception):
     """Raise this inside a _StatefulDecoder to stop decoding until more bytes
     have been received.
@@ -614,7 +610,7 @@ def _state_read_body_buffer(self):
         return result
 
 
-class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
+class SmartClientRequestProtocolOne(Requester,
                                     message.ResponseHandler):
     """The client-side protocol for smart version 1."""
 
@@ -658,7 +654,7 @@ def call_with_body_bytes(self, args, body):
             if 'hpssdetail' in debug.debug_flags:
                 mutter('hpss body content: %s', body)
         self._write_args(args)
-        bytes = self._encode_bulk_data(body)
+        bytes = _encode_bulk_data(body)
         self._request.accept_bytes(bytes)
         self._request.finished_writing()
         self._last_verb = args[0]
@@ -676,8 +672,8 @@ def call_with_body_readv_array(self, args, body):
                        self._request._medium._path)
             self._request_start_time = osutils.perf_counter()
         self._write_args(args)
-        readv_bytes = self._serialise_offsets(body)
-        bytes = self._encode_bulk_data(readv_bytes)
+        readv_bytes = _serialise_offsets(body)
+        bytes = _encode_bulk_data(readv_bytes)
         self._request.accept_bytes(bytes)
         self._request.finished_writing()
         if 'hpss' in debug.debug_flags:
@@ -1127,7 +1123,7 @@ def _write_prefixed_bencode(self, structure):
         self._write_func(bytes)
 
     def _write_headers(self, headers):
-        self._write_prefixed_bencode(headers)
+        self._write_prefixed_bencode({k.encode('utf-8'): v.encode('utf-8') for (k, v) in headers.items()})
 
     def _write_structure(self, args):
         self._write_func(b's')
@@ -1164,7 +1160,7 @@ def __init__(self, write_func):
         _ProtocolThreeEncoder.__init__(self, write_func)
         self.response_sent = False
         self._headers = {
-            b'Software version': breezy.__version__.encode('utf-8')}
+            'Software version': breezy.__version__}
         if 'hpss' in debug.debug_flags:
             self._thread_id = _thread.get_ident()
             self._response_start_time = None

breezy/bzr/tests/test_smart_transport.py

@@ -1917,8 +1917,8 @@ class TestSmartProtocol(tests.TestCase):
 
     request_encoder: object
     response_decoder: Type[protocol._StatefulDecoder]
-    server_protocol_class: Type[protocol.SmartProtocolBase]
-    client_protocol_class: Optional[Type[protocol.SmartProtocolBase]] = None
+    server_protocol_class = Type[object]
+    client_protocol_class = Type[object]
 
     def make_client_protocol_and_output(self, input_bytes=None):
         """:returns: a Request"""
@@ -1980,7 +1980,7 @@ def assertOffsetSerialisation(self, expected_offsets, expected_serialised,
         readv_cmd = vfs.ReadvRequest(None, '/')
         offsets = readv_cmd._deserialise_offsets(expected_serialised)
         self.assertEqual(expected_offsets, offsets)
-        serialised = requester._serialise_offsets(offsets)
+        serialised = protocol._serialise_offsets(offsets)
         self.assertEqual(expected_serialised, serialised)
 
     def build_protocol_waiting_for_body(self):
@@ -3196,7 +3196,7 @@ def test_call_smoke_test(self):
         correct bytes for that invocation.
         """
         requester, output = self.make_client_encoder_and_output()
-        requester.set_headers({b'header name': b'header value'})
+        requester.set_headers({'header name': 'header value'})
         requester.call(b'one arg')
         self.assertEqual(
             b'bzr message 3 (bzr 1.6)\n'  # protocol version
@@ -3212,7 +3212,7 @@ def test_call_with_body_bytes_smoke_test(self):
         call_with_body_bytes emits the correct bytes for that invocation.
         """
         requester, output = self.make_client_encoder_and_output()
-        requester.set_headers({b'header name': b'header value'})
+        requester.set_headers({'header name': 'header value'})
         requester.call_with_body_bytes((b'one arg',), b'body bytes')
         self.assertEqual(
             b'bzr message 3 (bzr 1.6)\n'  # protocol version
@@ -3246,7 +3246,7 @@ def test_call_with_body_stream_smoke_test(self):
         call_with_body_stream emits the correct bytes for that invocation.
         """
         requester, output = self.make_client_encoder_and_output()
-        requester.set_headers({b'header name': b'header value'})
+        requester.set_headers({'header name': 'header value'})
         stream = [b'chunk 1', b'chunk two']
         requester.call_with_body_stream((b'one arg',), stream)
         self.assertEqual(
@@ -3307,25 +3307,14 @@ def stream_checker():
             self.assertTrue(requester.body_stream_started)
             in_stream[0] = True
             yield b'content'
-        flush_called = []
-        orig_flush = requester.flush
 
-        def tracked_flush():
-            flush_called.append(in_stream[0])
-            if in_stream[0]:
-                self.assertTrue(requester.body_stream_started)
-            else:
-                self.assertFalse(requester.body_stream_started)
-            return orig_flush()
-        requester.flush = tracked_flush
         requester.call_with_body_stream((b'one arg',), stream_checker())
         self.assertEqual(
             b'bzr message 3 (bzr 1.6)\n'  # protocol version
             b'\x00\x00\x00\x02de'  # headers
             b's\x00\x00\x00\x0bl7:one arge'  # args
             b'b\x00\x00\x00\x07content'  # body
             b'e', output.getvalue())
-        self.assertEqual([False, True, True], flush_called)
 
 
 class StubMediumRequest:
@@ -3387,7 +3376,10 @@ def stream_that_fails():
             raise Exception('Boom!')
         response = _mod_request.SuccessfulSmartServerResponse(
             (b'args',), body_stream=stream_that_fails())
-        encoder.send_response(response)
+        try:
+            encoder.send_response(response)
+        except Exception as e:
+            self.assertEqual('Boom!', e.args[0])
         expected_response = (
             b'bzr message 3 (bzr 1.6)\n'  # protocol marker
             b'\x00\x00\x00\x02de'  # headers dict (empty)
@@ -3749,8 +3741,8 @@ def test_call_default_headers(self):
         """
         smart_client = client._SmartClient('dummy medium')
         self.assertEqual(
-            breezy.__version__.encode('utf-8'),
-            smart_client._headers[b'Software version'])
+            breezy.__version__,
+            smart_client._headers['Software version'])
         # XXX: need a test that smart_client._headers is passed to the request
         # encoder.