Skip to content

Commit 5f9ed21

Browse files
jonesbusyjonesbusy
authored
Optimize copy by passing stream with known size and digest (#552)
2 parents 0f4a187 + 362146b commit 5f9ed21

13 files changed

Lines changed: 213 additions & 14 deletions

‎src/main/java/land/oras/ContainerRef.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,15 @@ public String getBlobsUploadDigestPath(Registry registry) {
325325
return "%s/blobs/uploads/?digest=%s".formatted(getApiPrefix(registry), digest);
326326
}
327327

328+
/**
329+
* Return the blobs upload URL for POST upload to get the upload location
330+
* @param registry The registry
331+
* @return The blobs upload URL
332+
*/
333+
public String getBlobsUploadPath(Registry registry) {
334+
return "%s/blobs/uploads/".formatted(getApiPrefix(registry));
335+
}
336+
328337
/**
329338
* Return the blobs URL
330339
* @param registry The registry

‎src/main/java/land/oras/CopyUtils.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import java.io.IOException;
2424
import java.io.InputStream;
25+
import java.util.Objects;
2526
import land.oras.exception.OrasException;
2627
import org.jspecify.annotations.NonNull;
2728
import org.slf4j.Logger;
@@ -73,10 +74,14 @@ void copy(
7374

7475
// Write all layer
7576
for (Layer layer : source.collectLayers(sourceRef, contentType, true)) {
76-
try (InputStream is = source.fetchBlob(sourceRef.withDigest(layer.getDigest()))) {
77-
target.pushBlob(targetRef.withDigest(layer.getDigest()), is);
78-
LOG.debug("Copied layer {}", layer.getDigest());
79-
}
77+
Objects.requireNonNull(layer.getDigest(), "Layer digest is required for streaming copy");
78+
Objects.requireNonNull(layer.getSize(), "Layer size is required for streaming copy");
79+
target.pushBlob(
80+
targetRef.withDigest(layer.getDigest()),
81+
layer.getSize(),
82+
() -> source.fetchBlob(sourceRef.withDigest(layer.getDigest())),
83+
layer.getAnnotations());
84+
LOG.debug("Copied layer {}", layer.getDigest());
8085
}
8186

8287
// Single manifest

‎src/main/java/land/oras/OCI.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.LinkedList;
3232
import java.util.List;
3333
import java.util.Map;
34+
import java.util.function.Supplier;
3435
import land.oras.exception.OrasException;
3536
import land.oras.utils.ArchiveUtils;
3637
import land.oras.utils.Const;
@@ -375,6 +376,16 @@ public abstract Manifest pushArtifact(
375376
*/
376377
public abstract Layer pushBlob(T ref, Path blob, Map<String, String> annotations);
377378

379+
/**
380+
* Push a blob from input stream with known digest and size
381+
* @param ref The container ref with digest
382+
* @param size The size of the blob
383+
* @param stream The input stream of the blob
384+
* @param annotations The annotations
385+
* @return The layer
386+
*/
387+
public abstract Layer pushBlob(T ref, long size, Supplier<InputStream> stream, Map<String, String> annotations);
388+
378389
/**
379390
* Push the blob for the given layer
380391
* @param ref The container ref

‎src/main/java/land/oras/OCILayout.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import java.util.LinkedList;
3030
import java.util.List;
3131
import java.util.Map;
32+
import java.util.function.Supplier;
3233
import land.oras.exception.OrasException;
3334
import land.oras.utils.Const;
3435
import land.oras.utils.JsonUtils;
@@ -280,6 +281,33 @@ public Layer pushBlob(LayoutRef ref, Path blob, Map<String, String> annotations)
280281
}
281282
}
282283

284+
@Override
285+
public Layer pushBlob(LayoutRef ref, long size, Supplier<InputStream> stream, Map<String, String> annotations) {
286+
String digest = ref.getTag();
287+
if (digest == null) {
288+
throw new OrasException("Digest is required to push blob to layout");
289+
}
290+
boolean isDigest = SupportedAlgorithm.isSupported(digest);
291+
if (!isDigest) {
292+
throw new OrasException("Unsupported digest: %s".formatted(digest));
293+
}
294+
ensureAlgorithmPath(digest);
295+
try {
296+
Path blobPath = getBlobPath(ref);
297+
if (Files.exists(blobPath)) {
298+
LOG.info("Blob already exists: {}", digest);
299+
return Layer.fromFile(blobPath, ref.getAlgorithm()).withAnnotations(annotations);
300+
}
301+
try (InputStream is = stream.get()) {
302+
Files.copy(is, blobPath);
303+
}
304+
ensureDigest(ref, blobPath);
305+
return Layer.fromFile(blobPath, ref.getAlgorithm()).withAnnotations(annotations);
306+
} catch (IOException e) {
307+
throw new OrasException("Failed to push blob", e);
308+
}
309+
}
310+
283311
@Override
284312
public Layer pushBlob(LayoutRef ref, byte[] data) {
285313
try {

‎src/main/java/land/oras/Registry.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.HashMap;
3232
import java.util.List;
3333
import java.util.Map;
34+
import java.util.function.Supplier;
3435
import land.oras.auth.AuthProvider;
3536
import land.oras.auth.AuthStoreAuthenticationProvider;
3637
import land.oras.auth.HttpClient;
@@ -467,6 +468,58 @@ public Layer pushBlob(ContainerRef containerRef, Path blob, Map<String, String>
467468
return Layer.fromFile(blob, containerRef.getAlgorithm()).withAnnotations(annotations);
468469
}
469470

471+
@Override
472+
public Layer pushBlob(ContainerRef ref, long size, Supplier<InputStream> stream, Map<String, String> annotations) {
473+
String digest = ref.getDigest();
474+
if (digest == null) {
475+
throw new OrasException("Digest is required to push blob with stream");
476+
}
477+
ContainerRef containerRef = ref.forRegistry(this).checkBlocked(this);
478+
if (containerRef.isInsecure(this) && !this.isInsecure()) {
479+
return asInsecure().pushBlob(ref, size, stream, annotations);
480+
}
481+
if (hasBlob(containerRef)) {
482+
LOG.info("Blob already exists: {}", digest);
483+
return Layer.fromDigest(digest, size).withAnnotations(annotations);
484+
}
485+
// Empty post without digest
486+
URI uri = URI.create("%s://%s".formatted(getScheme(), containerRef.getBlobsUploadPath(this)));
487+
HttpClient.ResponseWrapper<String> response = client.post(
488+
uri,
489+
new byte[0],
490+
Map.of(Const.CONTENT_TYPE_HEADER, Const.APPLICATION_OCTET_STREAM_HEADER_VALUE),
491+
Scopes.of(this, containerRef),
492+
authProvider);
493+
logResponse(response);
494+
if (response.statusCode() != 202) {
495+
throw new OrasException("Failed to initiate blob upload: %s".formatted(response.response()));
496+
}
497+
String location = response.headers().get(Const.LOCATION_HEADER.toLowerCase());
498+
// Ensure location is absolute URI
499+
if (!location.startsWith("http") && !location.startsWith("https")) {
500+
location = "%s://%s/%s".formatted(getScheme(), ref.getApiRegistry(this), location.replaceFirst("^/", ""));
501+
}
502+
LOG.debug("Location header: {}", location);
503+
504+
URI uploadURI = createLocationWithDigest(location, digest);
505+
506+
response = client.upload(
507+
uploadURI,
508+
size,
509+
Map.of(Const.CONTENT_TYPE_HEADER, Const.APPLICATION_OCTET_STREAM_HEADER_VALUE),
510+
stream,
511+
Scopes.of(this, containerRef),
512+
authProvider);
513+
logResponse(response);
514+
if (response.statusCode() == 201) {
515+
LOG.debug("Successful push: {}", response.response());
516+
} else {
517+
throw new OrasException("Failed to push layer: %s".formatted(response.response()));
518+
}
519+
handleError(response);
520+
return Layer.fromDigest(digest, size).withAnnotations(annotations);
521+
}
522+
470523
@Override
471524
public Layer pushBlob(ContainerRef containerRef, byte[] data) {
472525
String digest = containerRef.getAlgorithm().digest(data);

‎src/main/java/land/oras/auth/HttpClient.java‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import java.util.List;
3636
import java.util.Map;
3737
import java.util.Objects;
38+
import java.util.function.Supplier;
3839
import java.util.regex.Matcher;
3940
import java.util.regex.Pattern;
4041
import java.util.stream.Collectors;
@@ -234,6 +235,35 @@ public ResponseWrapper<String> upload(
234235
}
235236
}
236237

238+
/**
239+
* Upload from an input stream.
240+
* @param uri The URI
241+
* @param size The size of the input stream
242+
* @param headers The headers
243+
* @param stream The input stream
244+
* @param scopes The scopes
245+
* @param authProvider The authentication provider
246+
* @return The response
247+
*/
248+
public ResponseWrapper<String> upload(
249+
URI uri,
250+
long size,
251+
Map<String, String> headers,
252+
Supplier<InputStream> stream,
253+
Scopes scopes,
254+
AuthProvider authProvider) {
255+
return executeRequest(
256+
"PUT",
257+
uri,
258+
true,
259+
headers,
260+
new byte[0],
261+
HttpResponse.BodyHandlers.ofString(),
262+
HttpRequest.BodyPublishers.fromPublisher(HttpRequest.BodyPublishers.ofInputStream(stream), size),
263+
scopes,
264+
authProvider);
265+
}
266+
237267
/**
238268
* Perform a HEAD request
239269
* @param uri The URI
@@ -296,7 +326,7 @@ public ResponseWrapper<String> post(
296326
headers,
297327
body,
298328
HttpResponse.BodyHandlers.ofString(),
299-
HttpRequest.BodyPublishers.ofByteArray(body),
329+
body.length == 0 ? HttpRequest.BodyPublishers.noBody() : HttpRequest.BodyPublishers.ofByteArray(body),
300330
scopes,
301331
authProvider);
302332
}

‎src/test/java/land/oras/ContainerRefTest.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
class ContainerRefTest {
3838

3939
@Test
40+
@Execution(ExecutionMode.SAME_THREAD)
4041
void shouldReadRegistriesConfig(@TempDir Path homeDir) throws Exception {
4142
// language=toml
4243
String config =
@@ -77,6 +78,7 @@ void shouldReadRegistriesConfig(@TempDir Path homeDir) throws Exception {
7778
}
7879

7980
@Test
81+
@Execution(ExecutionMode.SAME_THREAD)
8082
void shouldDetermineFromAlias(@TempDir Path homeDir) throws Exception {
8183

8284
// language=toml
@@ -99,6 +101,7 @@ void shouldDetermineFromAlias(@TempDir Path homeDir) throws Exception {
99101
}
100102

101103
@Test
104+
@Execution(ExecutionMode.SAME_THREAD)
102105
void shouldRewriteAllSubdomainToLocalProxy(@TempDir Path homeDir) throws Exception {
103106

104107
// language=toml
@@ -145,6 +148,7 @@ void shouldRewriteAllSubdomainToLocalProxy(@TempDir Path homeDir) throws Excepti
145148
}
146149

147150
@Test
151+
@Execution(ExecutionMode.SAME_THREAD)
148152
void shouldDetermineEffectiveRegistry(@TempDir Path homeDir) throws Exception {
149153

150154
// Use from container ref

‎src/test/java/land/oras/GitHubContainerRegistryITCase.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,6 @@ class GitHubContainerRegistryITCase {
3636
@TempDir
3737
Path tempDir;
3838

39-
@TempDir
40-
private static Path homeDir;
41-
4239
@Test
4340
void shouldPullIndex() {
4441
Registry registry = Registry.builder().build();
@@ -48,7 +45,8 @@ void shouldPullIndex() {
4845
}
4946

5047
@Test
51-
void shouldPullIndexWithAlias() throws Exception {
48+
@Execution(ExecutionMode.SAME_THREAD)
49+
void shouldPullIndexWithAlias(@TempDir Path homeDir) throws Exception {
5250
// language=toml
5351
String config =
5452
"""
@@ -57,7 +55,7 @@ void shouldPullIndexWithAlias() throws Exception {
5755
""";
5856

5957
// Setup
60-
TestUtils.createRegistriesConfFile(tempDir, config);
58+
TestUtils.createRegistriesConfFile(homeDir, config);
6159

6260
TestUtils.withHome(homeDir, () -> {
6361
Registry registry = Registry.builder().defaults().build();

‎src/test/java/land/oras/OCILayoutTest.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -618,6 +618,30 @@ void shouldPushBlob() throws IOException {
618618
assertBlobContent(path, digest, "hi");
619619
}
620620

621+
@Test
622+
void shouldFailToPushBlobViaStreamWithoutDigest() {
623+
Path path = layoutPath.resolve("shouldFailToPushBlobViaStreamWithoutDigest");
624+
byte[] content = "hi".getBytes(StandardCharsets.UTF_8);
625+
OCILayout ociLayout = OCILayout.Builder.builder().defaults(path).build();
626+
LayoutRef layoutRef = LayoutRef.parse("test");
627+
OrasException e = assertThrows(OrasException.class, () -> {
628+
ociLayout.pushBlob(layoutRef, content.length, () -> InputStream.nullInputStream(), Map.of());
629+
});
630+
assertEquals("Digest is required to push blob to layout", e.getMessage());
631+
}
632+
633+
@Test
634+
void shouldFailToPushBlobViaStreamWithInvalidDigest() {
635+
Path path = layoutPath.resolve("shouldFailToPushBlobViaStreamWithInvalidDigest");
636+
byte[] content = "hi".getBytes(StandardCharsets.UTF_8);
637+
OCILayout ociLayout = OCILayout.Builder.builder().defaults(path).build();
638+
LayoutRef layoutRef = LayoutRef.parse("test:1234");
639+
OrasException e = assertThrows(OrasException.class, () -> {
640+
ociLayout.pushBlob(layoutRef, content.length, () -> InputStream.nullInputStream(), Map.of());
641+
});
642+
assertEquals("Unsupported digest: 1234", e.getMessage());
643+
}
644+
621645
@Test
622646
void cannotPushBlobWithoutTagOrDigest() throws IOException {
623647

‎src/test/java/land/oras/PublicECRITCase.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
import org.junit.jupiter.api.parallel.ExecutionMode;
3232
import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
3333

34-
@Execution(ExecutionMode.SAME_THREAD) // Avoid 429 Too Many Requests for unauthenticated requests to public ECR
34+
@Execution(ExecutionMode.SAME_THREAD)
3535
class PublicECRITCase {
3636

3737
@TempDir

0 commit comments

Comments
 (0)