Fix OOM error when pulling layer blobs

Signed-off-by: Paolo Di Tommaso <[email protected]>
seqeralabs · Apr 26, 2023 · f065379 · f065379
1 parent 818a559
commit f065379
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 19 deletions.
diff --git a/src/main/groovy/io/seqera/wave/controller/RegistryProxyController.groovy b/src/main/groovy/io/seqera/wave/controller/RegistryProxyController.groovy
@@ -235,16 +235,19 @@ class RegistryProxyController {
                 .headers(toMutableHeaders(resp.headers, override))
     }
 
-    MutableHttpResponse<?> fromDelegateResponse(final DelegateResponse delegateResponse){
+    MutableHttpResponse<?> fromDelegateResponse(final DelegateResponse response){
 
-        final Long contentLength = delegateResponse.headers
+        final Long len = response.headers
                 .find {it.key.toLowerCase()=='content-length'}?.value?.first() as long ?: null
-        final fluxInputStream = createFluxFromChunkBytes(delegateResponse.body, contentLength)
+
+        final streamedFile =  len
+                    ?  new StreamedFile(response.body, MediaType.APPLICATION_OCTET_STREAM_TYPE, Instant.now().toEpochMilli(), len)
+                    :  new StreamedFile(response.body, MediaType.APPLICATION_OCTET_STREAM_TYPE)
 
         HttpResponse
-                .status(HttpStatus.valueOf(delegateResponse.statusCode))
-                .body(fluxInputStream)
-                .headers(toMutableHeaders(delegateResponse.headers))
+                .status(HttpStatus.valueOf(response.statusCode))
+                .body(streamedFile)
+                .headers(toMutableHeaders(response.headers))
     }
 
     MutableHttpResponse<?> fromManifestResponse(DelegateResponse resp) {
@@ -270,13 +273,6 @@ class RegistryProxyController {
         }
     }
 
-    static protected StreamedFile createFluxFromChunkBytes(InputStream inputStream, Long size){
-        if( size )
-            new StreamedFile(inputStream, MediaType.APPLICATION_OCTET_STREAM_TYPE, Instant.now().toEpochMilli(), size)
-        else
-            new StreamedFile(inputStream, MediaType.APPLICATION_OCTET_STREAM_TYPE)
-    }
-
 
     static String stripUriParams(String uri) {
         final p = uri.indexOf('?')

diff --git a/src/main/groovy/io/seqera/wave/core/RegistryProxyService.groovy b/src/main/groovy/io/seqera/wave/core/RegistryProxyService.groovy
@@ -114,7 +114,7 @@ class RegistryProxyService {
 
     DelegateResponse handleRequest(RoutePath route, Map<String,List<String>> headers){
         ProxyClient proxyClient = client(route)
-        final resp1 = proxyClient.getString(route.path, headers, false)
+        final resp1 = proxyClient.getStream(route.path, headers, false)
         final redirect = resp1.headers().firstValue('Location').orElse(null)
         if( redirect && resp1.statusCode() in REDIRECT_CODES ) {
             // the redirect location can be a relative path i.e. without hostname
@@ -126,11 +126,10 @@ class RegistryProxyService {
                     headers:resp1.headers().map())
         }
 
-        final resp2 = proxyClient.getStream(route.path, headers)
         new DelegateResponse(
-                statusCode: resp2.statusCode(),
-                headers: resp2.headers().map(),
-                body: resp2.body() )
+                statusCode: resp1.statusCode(),
+                headers: resp1.headers().map(),
+                body: resp1.body() )
     }
 
     boolean isManifestPresent(String image){

diff --git a/src/main/groovy/io/seqera/wave/proxy/ProxyClient.groovy b/src/main/groovy/io/seqera/wave/proxy/ProxyClient.groovy
@@ -164,7 +164,7 @@ class ProxyClient {
                 loginService.invalidateAuthorization(image, registry.auth, credentials)
                 continue
             }
-            if( result.statusCode() in REDIRECT_CODES && followRedirect  ) {
+            if( result.statusCode() in REDIRECT_CODES && followRedirect ) {
                 final redirect = result.headers().firstValue('location').orElse(null)
                 log.trace "Redirecting (${++redirectCount}) $target ==> $redirect ${RegHelper.dumpHeaders(result.headers())}"
                 if( !redirect ) {