Build option to add architecture suffix to tags when publishing multi-platform images

jib-core/src/main/java/com/google/cloud/tools/jib/builder/steps/PushImageStep.java

@@ -36,6 +36,7 @@
 import java.util.Collections;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.stream.Collectors;
 
 /**
  * Pushes a manifest or a manifest list for a tag. If not a manifest list, returns the manifest
@@ -55,31 +56,26 @@ static ImmutableList<PushImageStep> makeList(
       Image builtImage,
       boolean manifestAlreadyExists)
       throws IOException {
-    boolean singlePlatform = buildContext.getContainerConfiguration().getPlatforms().size() == 1;
-    Set<String> tags = buildContext.getAllTargetImageTags();
-    int numPushers = singlePlatform ? tags.size() : 1;
+    // Gets the image manifest to push.
+    BuildableManifestTemplate manifestTemplate =
+        new ImageToJsonTranslator(builtImage)
+            .getManifestTemplate(
+                buildContext.getTargetFormat(), containerConfigurationDigestAndSize);
+    DescriptorDigest manifestDigest = Digests.computeJsonDigest(manifestTemplate);
+    Set<String> imageQualifiers = getImageQualifiers(buildContext, builtImage, manifestDigest);
 
     EventHandlers eventHandlers = buildContext.getEventHandlers();
     try (TimerEventDispatcher ignored =
             new TimerEventDispatcher(eventHandlers, "Preparing manifest pushers");
         ProgressEventDispatcher progressDispatcher =
-            progressEventDispatcherFactory.create("launching manifest pushers", numPushers)) {
+            progressEventDispatcherFactory.create(
+                "launching manifest pushers", imageQualifiers.size())) {
 
       if (JibSystemProperties.skipExistingImages() && manifestAlreadyExists) {
         eventHandlers.dispatch(LogEvent.info("Skipping pushing manifest; already exists."));
         return ImmutableList.of();
       }
 
-      // Gets the image manifest to push.
-      BuildableManifestTemplate manifestTemplate =
-          new ImageToJsonTranslator(builtImage)
-              .getManifestTemplate(
-                  buildContext.getTargetFormat(), containerConfigurationDigestAndSize);
-
-      DescriptorDigest manifestDigest = Digests.computeJsonDigest(manifestTemplate);
-
-      Set<String> imageQualifiers =
-          singlePlatform ? tags : Collections.singleton(manifestDigest.toString());
       return imageQualifiers.stream()
           .map(
               qualifier ->
@@ -95,6 +91,20 @@ static ImmutableList<PushImageStep> makeList(
     }
   }
 
+  private static Set<String> getImageQualifiers(
+      BuildContext buildContext, Image builtImage, DescriptorDigest manifestDigest) {
+    boolean singlePlatform = buildContext.getContainerConfiguration().getPlatforms().size() == 1;
+    Set<String> tags = buildContext.getAllTargetImageTags();
+    if (singlePlatform) {
+      return tags;
+    }
+    if (buildContext.getEnablePlatformTags()) {
+      String architecture = builtImage.getArchitecture();
+      return tags.stream().map(tag -> tag + "-" + architecture).collect(Collectors.toSet());
+    }
+    return Collections.singleton(manifestDigest.toString());
+  }
+
   static ImmutableList<PushImageStep> makeListForManifestList(
       BuildContext buildContext,
       ProgressEventDispatcher.Factory progressEventDispatcherFactory,

jib-core/src/main/java/com/google/cloud/tools/jib/configuration/BuildContext.java

@@ -83,6 +83,7 @@ public static class Builder {
     @Nullable private ExecutorService executorService;
     private boolean alwaysCacheBaseImage = false;
     private ImmutableListMultimap<String, String> registryMirrors = ImmutableListMultimap.of();
+    private boolean enablePlatformTags = false;
 
     private Builder() {}
 
@@ -120,6 +121,19 @@ public Builder setAdditionalTargetImageTags(Set<String> tags) {
       return this;
     }
 
+    /**
+     * Sets whether to automatically add architecture suffix to tags for platform-specific images
+     * when building multi-platform images. For example, when building amd64 and arm64 images for a
+     * given tag, the final tags will be {@code <tag>-amd64} and {@code <tag>-arm64}.
+     *
+     * @param enablePlatformTags whether to append architecture suffix to tags
+     * @return this
+     */
+    public Builder setEnablePlatformTags(boolean enablePlatformTags) {
+      this.enablePlatformTags = enablePlatformTags;
+      return this;
+    }
+
     /**
      * Sets configuration parameters for the container.
      *
@@ -328,7 +342,8 @@ public BuildContext build() throws CacheDirectoryCreationException {
               executorService == null ? Executors.newCachedThreadPool() : executorService,
               executorService == null, // shutDownExecutorService
               alwaysCacheBaseImage,
-              registryMirrors);
+              registryMirrors,
+              enablePlatformTags);
 
         case 1:
           throw new IllegalStateException(missingFields.get(0) + " is required but not set");
@@ -386,6 +401,7 @@ public static Builder builder() {
   private final boolean shutDownExecutorService;
   private final boolean alwaysCacheBaseImage;
   private final ImmutableListMultimap<String, String> registryMirrors;
+  private final boolean enablePlatformTags;
 
   /** Instantiate with {@link #builder}. */
   private BuildContext(
@@ -405,7 +421,8 @@ private BuildContext(
       ExecutorService executorService,
       boolean shutDownExecutorService,
       boolean alwaysCacheBaseImage,
-      ImmutableListMultimap<String, String> registryMirrors) {
+      ImmutableListMultimap<String, String> registryMirrors,
+      boolean enablePlatformTags) {
     this.baseImageConfiguration = baseImageConfiguration;
     this.targetImageConfiguration = targetImageConfiguration;
     this.additionalTargetImageTags = additionalTargetImageTags;
@@ -423,12 +440,17 @@ private BuildContext(
     this.shutDownExecutorService = shutDownExecutorService;
     this.alwaysCacheBaseImage = alwaysCacheBaseImage;
     this.registryMirrors = registryMirrors;
+    this.enablePlatformTags = enablePlatformTags;
   }
 
   public ImageConfiguration getBaseImageConfiguration() {
     return baseImageConfiguration;
   }
 
+  public boolean getEnablePlatformTags() {
+    return enablePlatformTags;
+  }
+
   public ImageConfiguration getTargetImageConfiguration() {
     return targetImageConfiguration;
   }

jib-core/src/test/java/com/google/cloud/tools/jib/builder/steps/PushImageStepTest.java

@@ -16,27 +16,37 @@
 
 package com.google.cloud.tools.jib.builder.steps;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.tools.jib.api.DescriptorDigest;
 import com.google.cloud.tools.jib.api.RegistryException;
 import com.google.cloud.tools.jib.api.buildplan.Platform;
+import com.google.cloud.tools.jib.blob.BlobDescriptor;
 import com.google.cloud.tools.jib.builder.ProgressEventDispatcher;
 import com.google.cloud.tools.jib.configuration.BuildContext;
 import com.google.cloud.tools.jib.configuration.ContainerConfiguration;
 import com.google.cloud.tools.jib.event.EventHandlers;
 import com.google.cloud.tools.jib.global.JibSystemProperties;
+import com.google.cloud.tools.jib.image.Image;
 import com.google.cloud.tools.jib.image.json.V22ManifestListTemplate;
 import com.google.cloud.tools.jib.image.json.V22ManifestListTemplate.ManifestDescriptorTemplate;
+import com.google.cloud.tools.jib.image.json.V22ManifestTemplate;
 import com.google.cloud.tools.jib.registry.RegistryClient;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import java.io.IOException;
-import org.junit.Assert;
+import java.util.List;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.contrib.java.lang.system.RestoreSystemProperties;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
 
 /** Tests for {@link PushImageStep}. */
@@ -50,20 +60,21 @@ public class PushImageStepTest {
   @Mock private BuildContext buildContext;
   @Mock private RegistryClient registryClient;
   @Mock private ContainerConfiguration containerConfig;
+  @Mock private DescriptorDigest mockDescriptorDigest;
 
   private final V22ManifestListTemplate manifestList = new V22ManifestListTemplate();
 
   @Before
   public void setUp() {
-    Mockito.when(buildContext.getAllTargetImageTags()).thenReturn(ImmutableSet.of("tag1", "tag2"));
-    Mockito.when(buildContext.getEventHandlers()).thenReturn(EventHandlers.NONE);
-    Mockito.when(buildContext.getContainerConfiguration()).thenReturn(containerConfig);
-    Mockito.when(containerConfig.getPlatforms())
+    when(buildContext.getAllTargetImageTags()).thenReturn(ImmutableSet.of("tag1", "tag2"));
+    when(buildContext.getEventHandlers()).thenReturn(EventHandlers.NONE);
+    when(buildContext.getContainerConfiguration()).thenReturn(containerConfig);
+    doReturn(V22ManifestTemplate.class).when(buildContext).getTargetFormat();
+    when(containerConfig.getPlatforms())
         .thenReturn(
             ImmutableSet.of(new Platform("amd64", "linux"), new Platform("arm64", "windows")));
-    Mockito.when(progressDispatcherFactory.create(Mockito.anyString(), Mockito.anyLong()))
-        .thenReturn(progressDispatcher);
-    Mockito.when(progressDispatcher.newChildProducer()).thenReturn(progressDispatcherFactory);
+    when(progressDispatcherFactory.create(anyString(), anyLong())).thenReturn(progressDispatcher);
+    when(progressDispatcher.newChildProducer()).thenReturn(progressDispatcherFactory);
 
     ManifestDescriptorTemplate manifest = new ManifestDescriptorTemplate();
     manifest.setSize(100);
@@ -73,40 +84,86 @@ public void setUp() {
 
   @Test
   public void testMakeListForManifestList() throws IOException, RegistryException {
-    ImmutableList<PushImageStep> pushImageStepList =
+    List<PushImageStep> pushImageStepList =
         PushImageStep.makeListForManifestList(
             buildContext, progressDispatcherFactory, registryClient, manifestList, false);
 
-    Assert.assertEquals(2, pushImageStepList.size());
+    assertThat(pushImageStepList).hasSize(2);
     for (PushImageStep pushImageStep : pushImageStepList) {
       BuildResult buildResult = pushImageStep.call();
-      Assert.assertEquals(
-          "sha256:64303e82b8a80ef20475dc7f807b81f172cacce1a59191927f3a7ea5222f38ae",
-          buildResult.getImageDigest().toString());
-      Assert.assertEquals(
-          "sha256:64303e82b8a80ef20475dc7f807b81f172cacce1a59191927f3a7ea5222f38ae",
-          buildResult.getImageId().toString());
+      assertThat(buildResult.getImageDigest().toString())
+          .isEqualTo("sha256:64303e82b8a80ef20475dc7f807b81f172cacce1a59191927f3a7ea5222f38ae");
+      assertThat(buildResult.getImageId().toString())
+          .isEqualTo("sha256:64303e82b8a80ef20475dc7f807b81f172cacce1a59191927f3a7ea5222f38ae");
     }
   }
 
+  @Test
+  public void testMakeList_multiPlatform_platformTags() throws IOException, RegistryException {
+    Image image = Image.builder(V22ManifestTemplate.class).setArchitecture("wasm").build();
+
+    when(buildContext.getEnablePlatformTags()).thenReturn(true);
+
+    List<PushImageStep> pushImageStepList =
+        PushImageStep.makeList(
+            buildContext,
+            progressDispatcherFactory,
+            registryClient,
+            new BlobDescriptor(mockDescriptorDigest),
+            image,
+            false);
+
+    ArgumentCaptor<String> tagCatcher = ArgumentCaptor.forClass(String.class);
+    when(registryClient.pushManifest(any(), tagCatcher.capture())).thenReturn(null);
+
+    assertThat(pushImageStepList).hasSize(2);
+    pushImageStepList.get(0).call();
+    pushImageStepList.get(1).call();
+
+    assertThat(tagCatcher.getAllValues()).containsExactly("tag1-wasm", "tag2-wasm");
+  }
+
+  @Test
+  public void testMakeList_multiPlatform_nonPlatformTags() throws IOException, RegistryException {
+    Image image = Image.builder(V22ManifestTemplate.class).setArchitecture("wasm").build();
+    when(buildContext.getEnablePlatformTags()).thenReturn(false);
+
+    List<PushImageStep> pushImageStepList =
+        PushImageStep.makeList(
+            buildContext,
+            progressDispatcherFactory,
+            registryClient,
+            new BlobDescriptor(mockDescriptorDigest),
+            image,
+            false);
+
+    ArgumentCaptor<String> tagCatcher = ArgumentCaptor.forClass(String.class);
+    when(registryClient.pushManifest(any(), tagCatcher.capture())).thenReturn(null);
+
+    assertThat(pushImageStepList).hasSize(1);
+    pushImageStepList.get(0).call();
+    assertThat(tagCatcher.getAllValues())
+        .containsExactly("sha256:0dd75658cf52608fbd72eb95ff5fc5946966258c3676b35d336bfcc7ac5006f1");
+  }
+
   @Test
   public void testMakeListForManifestList_singlePlatform() throws IOException {
-    Mockito.when(containerConfig.getPlatforms())
+    when(containerConfig.getPlatforms())
         .thenReturn(ImmutableSet.of(new Platform("amd64", "linux")));
 
-    ImmutableList<PushImageStep> pushImageStepList =
+    List<PushImageStep> pushImageStepList =
         PushImageStep.makeListForManifestList(
             buildContext, progressDispatcherFactory, registryClient, manifestList, false);
-    Assert.assertEquals(0, pushImageStepList.size());
+    assertThat(pushImageStepList).isEmpty();
   }
 
   @Test
   public void testMakeListForManifestList_manifestListAlreadyExists() throws IOException {
     System.setProperty(JibSystemProperties.SKIP_EXISTING_IMAGES, "true");
 
-    ImmutableList<PushImageStep> pushImageStepList =
+    List<PushImageStep> pushImageStepList =
         PushImageStep.makeListForManifestList(
             buildContext, progressDispatcherFactory, registryClient, manifestList, true);
-    Assert.assertEquals(0, pushImageStepList.size());
+    assertThat(pushImageStepList).isEmpty();
   }
 }

jib-core/src/test/java/com/google/cloud/tools/jib/configuration/BuildContextTest.java

@@ -128,6 +128,7 @@ public void testBuilder() throws Exception {
             .setApplicationLayersCacheDirectory(expectedApplicationLayersCacheDirectory)
             .setBaseImageLayersCacheDirectory(expectedBaseImageLayersCacheDirectory)
             .setTargetFormat(ImageFormat.OCI)
+            .setEnablePlatformTags(true)
             .setAllowInsecureRegistries(true)
             .setLayerConfigurations(expectedLayerConfigurations)
             .setToolName(expectedCreatedBy)
@@ -177,6 +178,7 @@ public void testBuilder() throws Exception {
     Assert.assertEquals(expectedCreatedBy, buildContext.getToolName());
     Assert.assertEquals(expectedRegistryMirrors, buildContext.getRegistryMirrors());
     Assert.assertNotNull(buildContext.getExecutorService());
+    Assert.assertTrue(buildContext.getEnablePlatformTags());
   }
 
   @Test
@@ -220,6 +222,7 @@ public void testBuilder_default() throws CacheDirectoryCreationException {
     Assert.assertEquals(Collections.emptyList(), buildContext.getLayerConfigurations());
     Assert.assertEquals("jib", buildContext.getToolName());
     Assert.assertEquals(0, buildContext.getRegistryMirrors().size());
+    Assert.assertFalse(buildContext.getEnablePlatformTags());
   }
 
   @Test