From 61af2796b2c5eaf12cf97333e1f9ce4526fe98c6 Mon Sep 17 00:00:00 2001
From: "James E. Blair" <jeblair@redhat.com>
Date: Thu, 3 Oct 2019 10:27:08 -0700
Subject: [PATCH] Add streaming download support

So we don't run out of memory when fetching blobs, use cherrypy's
streaming support.

Change-Id: I6c05240a81f0b6b2dd0992e42d097883ca4fdc2e
---
 zuul_registry/filesystem.py   | 11 +++++++++++
 zuul_registry/main.py         | 25 +++++++++++++++++++------
 zuul_registry/storage.py      |  5 +++++
 zuul_registry/storageutils.py | 13 +++++++++++++
 zuul_registry/swift.py        |  8 ++++++++
 5 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/zuul_registry/filesystem.py b/zuul_registry/filesystem.py
index 22a123c..641fbc3 100644
--- a/zuul_registry/filesystem.py
+++ b/zuul_registry/filesystem.py
@@ -57,6 +57,17 @@ class FilesystemDriver(storageutils.StorageDriver):
         with open(path, 'rb') as f:
             return f.read()
 
+    def stream_object(self, path):
+        path = os.path.join(self.root, path)
+        if not os.path.exists(path):
+            return None
+        with open(path, 'rb') as f:
+            while True:
+                chunk = f.read(4096)
+                if not chunk:
+                    return
+                yield chunk
+
     def delete_object(self, path):
         path = os.path.join(self.root, path)
         if os.path.exists(path):
diff --git a/zuul_registry/main.py b/zuul_registry/main.py
index e888275..8a8a9c1 100644
--- a/zuul_registry/main.py
+++ b/zuul_registry/main.py
@@ -106,20 +106,29 @@ class RegistryAPI:
 
     @cherrypy.expose
     @cherrypy.config(**{'tools.auth_basic.checkpassword': require_read})
-    def get_blob(self, repository, digest):
+    def head_blob(self, repository, digest):
         namespace = self.get_namespace()
-        method = cherrypy.request.method
-        self.log.info('%s blob %s %s', method, repository, digest)
+        self.log.info('Head blob %s %s', repository, digest)
         size = self.storage.blob_size(namespace, digest)
         if size is None:
             return self.not_found()
         res = cherrypy.response
         res.headers['Docker-Content-Digest'] = digest
-        if method != 'HEAD':
-            data = self.storage.get_blob(namespace, digest)
-            return data
         return {}
 
+    @cherrypy.expose
+    @cherrypy.config(**{'tools.auth_basic.checkpassword': require_read,
+                        'response.stream': True})
+    def get_blob(self, repository, digest):
+        namespace = self.get_namespace()
+        self.log.info('Get blob %s %s', repository, digest)
+        size = self.storage.blob_size(namespace, digest)
+        if size is None:
+            return self.not_found()
+        res = cherrypy.response
+        res.headers['Docker-Content-Digest'] = digest
+        return self.storage.stream_blob(namespace, digest)
+
     @cherrypy.expose
     @cherrypy.config(**{'tools.auth_basic.checkpassword': require_write})
     def start_upload(self, repository, digest=None):
@@ -260,6 +269,10 @@ class RegistryServer:
                           conditions=dict(method=['GET']),
                           controller=api, action='get_manifest')
         route_map.connect('api', '/v2/{repository:.*}/blobs/{digest}',
+                          conditions=dict(method=['HEAD']),
+                          controller=api, action='head_blob')
+        route_map.connect('api', '/v2/{repository:.*}/blobs/{digest}',
+                          conditions=dict(method=['GET']),
                           controller=api, action='get_blob')
 
         conf = {
diff --git a/zuul_registry/storage.py b/zuul_registry/storage.py
index 2a0999d..bb75fe5 100644
--- a/zuul_registry/storage.py
+++ b/zuul_registry/storage.py
@@ -139,6 +139,11 @@ class Storage:
         path = os.path.join(namespace, 'blobs', digest, 'data')
         return self.backend.get_object(path)
 
+    def stream_blob(self, namespace, digest):
+        path = os.path.join(namespace, 'blobs',
+                            self._path_from_digest(digest), 'data')
+        return self.backend.stream_object(path)
+
     def start_upload(self, namespace):
         """Start an upload.
 
diff --git a/zuul_registry/storageutils.py b/zuul_registry/storageutils.py
index cf8979e..42aed71 100644
--- a/zuul_registry/storageutils.py
+++ b/zuul_registry/storageutils.py
@@ -91,6 +91,19 @@ class StorageDriver(metaclass=ABCMeta):
         """
         pass
 
+    @abstractmethod
+    def stream_object(self, path):
+        """Retrieve an object, streaming.
+
+        Return a generator with the content of the object at `path`.
+
+        :arg str path: The object path.
+
+        :returns: The contents of the object.
+        :rtype: generator of bytearray
+        """
+        pass
+
     @abstractmethod
     def delete_object(self, path):
         """Delete an object.
diff --git a/zuul_registry/swift.py b/zuul_registry/swift.py
index 87bb666..709a572 100644
--- a/zuul_registry/swift.py
+++ b/zuul_registry/swift.py
@@ -120,6 +120,14 @@ class SwiftDriver(storageutils.StorageDriver):
             return None
         return ret.content
 
+    def stream_object(self, path):
+        try:
+            ret = retry_function(
+                lambda: self.conn.session.get(self.get_url(path), stream=True))
+        except keystoneauth1.exceptions.http.NotFound:
+            return None
+        return ret.iter_content(chunk_size=4096)
+
     def delete_object(self, path):
         retry_function(
             lambda: self.conn.session.delete(