Repository: canonical/operator
Pull request #666 (2021-11-08 16:27:23)

Memory-efficient push via pebble
Uses a generator to efficiently send file data in chunks to pebble.

To view the diff of just the new parts: https://github.com/Vultaire/operator/compare/harness-pebble-fs...issues/576_push

Vultaire: @petevg For transparency, keep in mind that most the tests here are simply from #645.  I only added one extra test here, for push/pull of a larger file which would trigger chunking.  However, in the push case I feel like that may be sufficient as the other tests do provide a significant harness for these changes.

Thanks for the review!
jnsgruk: Despite failing tests, I think this is ready for broader review while the minor test issues are worked through.
jameinel: I found it a bit hard to do a review of just this PR since it depends on #645, if possible can we land that one first and then review this one?
jnsgruk: > (AFAICT the only actual change here above #645 is just to turn the core inner loop from returning a list to using a generator so we don't have to pull in all the bytes ahead of time, which sounds great, but hard to actually tease apart what is changing.)

Yep, I think you can see just the changes for this PR [here](https://github.com/Vultaire/operator/compare/harness-pebble-fs...issues/576_push)
petevg: I'm grabbing this to look into the failing tests.

The root cause seems to be that the mocked out object we're using in our unit tests doesn't have a .decode property, though it might go deeper than that. (IIRC, there was more convo around this earlier in the PR, or in another PR ...)
jnsgruk: Thanks @petevg!

Patch:
diff --git a/ops/pebble.py b/ops/pebble.py
index c9a71b4..883de86 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -1069,6 +1069,8 @@ class _WebsocketReader(io.BufferedIOBase):
 class Client:
     """Pebble API client."""
 
+    _chunk_size = 8192
+
     def __init__(self, socket_path=None, opener=None, base_url='http://localhost', timeout=5.0):
         """Initialize a client instance.
 
@@ -1472,7 +1474,7 @@ class Client:
 
         # Then read the rest of the response and feed it to the parser.
         while True:
-            chunk = response.read(8192)
+            chunk = response.read(self._chunk_size)
             if not chunk:
                 break
             parser.feed(chunk)
@@ -1546,14 +1548,7 @@ class Client:
             'files': [info],
         }
 
-        if hasattr(source, 'read'):
-            content = source.read()
-        else:
-            content = source
-        if isinstance(content, str):
-            content = content.encode(encoding)
-
-        data, content_type = self._encode_multipart(metadata, path, content)
+        data, content_type = self._encode_multipart(metadata, path, source, encoding)
 
         headers = {
             'Accept': 'application/json',
@@ -1579,28 +1574,46 @@ class Client:
             d['group'] = group
         return d
 
-    @staticmethod
-    def _encode_multipart(metadata, path, content):
+    def _encode_multipart(self, metadata, path, source, encoding):
         # Python's stdlib mime/multipart handling is screwy and doesn't handle
         # binary properly, so roll our own.
+
+        if isinstance(source, str):
+            source = io.StringIO(source)
+        elif isinstance(source, bytes):
+            source = io.BytesIO(source)
+
         boundary = binascii.hexlify(os.urandom(16))
         path_escaped = path.replace('"', '\\"').encode('utf-8')  # NOQA: test_quote_backslashes
-        parts = []
-        parts.extend([
-            b'--', boundary, b'\r\n',
-            b'Content-Type: application/json\r\n',
-            b'Content-Disposition: form-data; name="request"\r\n',
-            b'\r\n',
-            json.dumps(metadata).encode('utf-8'), b'\r\n',
-            b'--', boundary, b'\r\n',
-            b'Content-Type: application/octet-stream\r\n',
-            b'Content-Disposition: form-data; name="files"; filename="', path_escaped, b'"\r\n',
-            b'\r\n',
-            content, b'\r\n',
-            b'--', boundary, b'--\r\n',
-        ])
         content_type = 'multipart/form-data; boundary="' + boundary.decode('utf-8') + '"'
-        return b''.join(parts), content_type
+
+        def generator():
+            yield b''.join([
+                b'--', boundary, b'\r\n',
+                b'Content-Type: application/json\r\n',
+                b'Content-Disposition: form-data; name="request"\r\n',
+                b'\r\n',
+                json.dumps(metadata).encode('utf-8'), b'\r\n',
+                b'--', boundary, b'\r\n',
+                b'Content-Type: application/octet-stream\r\n',
+                b'Content-Disposition: form-data; name="files"; filename="',
+                path_escaped, b'"\r\n',
+                b'\r\n',
+            ])
+
+            content = source.read(self._chunk_size)
+            while content:
+                if isinstance(content, str):
+                    content = content.encode(encoding)
+                yield content
+                content = source.read(self._chunk_size)
+
+            yield b''.join([
+                b'\r\n',
+                b'--', boundary, b'--\r\n',
+            ])
+
+        return generator(), content_type
 
     def list_files(self, path: str, *, pattern: str = None,
                    itself: bool = False) -> typing.List[FileInfo]:


Test patch:
diff --git a/test/test_pebble.py b/test/test_pebble.py
index a54fabd..db21281 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -1635,6 +1635,7 @@ bad path
         self.assertEqual(request[:3], ('POST', '/v1/files', None))
 
         headers, body = request[3:]
+
         content_type = headers['Content-Type']
         req, filename, content = self._parse_write_multipart(content_type, body)
         self.assertEqual(filename, '/foo/bar')
@@ -1800,7 +1801,9 @@ bad path
         parser = email.parser.BytesFeedParser()
         parser.feed(b'Content-Type: multipart/form-data; boundary='
                     + boundary.encode('utf-8') + b'\r\n\r\n')
-        parser.feed(body)
+        for b in body:
+            # With the "memory efficient push" changes, body is an iterable.
+            parser.feed(b)
         message = parser.close()
 
         req = None
diff --git a/test/test_testing.py b/test/test_testing.py
index 9b0407b..d56b5b9 100644
--- a/test/test_testing.py
+++ b/test/test_testing.py
@@ -3051,6 +3051,19 @@ class _PebbleStorageAPIsTestMixin:
             received_data = infile.read()
         self.assertEqual(original_data, received_data)
 
+    def test_push_and_pull_larger_file(self):
+        # Intent: to ensure things work appropriately with larger files.
+        # Larger files may be sent/received in multiple chunks; this should help for
+        # checking that such logic is correct.
+        data_size = 1024 * 1024
+        original_data = os.urandom(data_size)
+
+        client = self.client
+        client.push(self.prefix + '/test', original_data, encoding=None)
+        with client.pull(self.prefix + '/test', encoding=None) as infile:
+            received_data = infile.read()
+        self.assertEqual(original_data, received_data)
+
     def test_push_to_non_existent_subdir(self):
         data = 'data'
         client = self.client