Repository: bazelbuild/bazel
Issue: D8 dex merger fails when a synthetic class is placed on a different shard than its container class
Labels: type: bug
P1
team-Android

### Description of the bug:


 a desugared lambda is packed in a different shard than the one
from its original class D8 will fail with the following message:

`Attempt at compiling intermediate artifact without its context
`
e.g.

I'm encountering an issue when compiling my app with the following classes:
`org/chromium/base/MemoryPressureListener
``org/chromium/base/MemoryPressureListener$$ExternalSyntheticLambda0`

bazel fails with the following message 

```
Error in bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/mirroring-sample-app/dexsplits/apk/9.shard.zip:org/chromium/base/MemoryPressureListener$$ExternalSyntheticLambda0.class.dex:
Attempt at compiling intermediate artifact without its context
Merge failed: Compilation failed to complete, origin: bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/mirroring-sample-app/dexsplits/apk/9.shard.zip:org/chromium/base/MemoryPressureListener$$ExternalSyntheticLambda0.class.dex
Exception in thread "main" com.android.tools.r8.CompilationFailedException: Merge failed: Compilation failed to complete, origin: bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/mirroring-sample-app/dexsplits/apk/9.shard.zip:org/chromium/base/MemoryPressureListener$$ExternalSyntheticLambda0.class.dex
at com.google.devtools.build.android.r8.DexFileMerger.main(DexFileMerger.java:389)
Target //apps/hermosa/hermosa-services/realtime/mirroring-sample-app:apk failed to build
```
Here the container class is in shard #8 and the lambda class is in shard #9

Running with d8 dexmerger and native multidex.

### What's the simplest, easiest way to reproduce this bug? Please provide a minimal example if possible.


Its hard to provide a small reproducible example since it needs to be an app large enough to exceed the size of a dex file and the sharding has to result in a lambda class ending up in a different shard than the container class.

### Which operating system are you running Bazel on?


Linux/Mac

### What is the output of `bazel info release`?


Internal build merged with master at commit https://github.com/bazelbuild/bazel/commit/b6803a2a8367e0a6748256d2e1f162ceb5451636

### If `bazel info release` returns `development version` or `(@non-git)`, tell us how you built Bazel.


_No response_

### What's the output of `git remote get-url origin; git rev-parse master; git rev-parse HEAD` ?


_No response_

### Have you found anything relevant by searching the web?

_No response_

### Any other information, logs, or outputs that you want to share?

_No response_

Issue comments:
@ahumesky 
We are also seeing this in our repo using the latest rolling release (6.0.0-pre.20220922.1)
@ahumesky can this issue be added to [Bazel 6.1. release blockers](https://github.com/bazelbuild/bazel/milestone/46)? As you noted #16369 does not completely solve it. We are hesitant to upgrade to 6.0 even with that fix since we might run into the same issue with another of synthetic classes at any point.
Yes I think we can get the fix in for 6.1.

But a question for those of you running into this issue -- we're not seeing this internally, so I'm trying to figure out what differences might be causing this with bazel. Is anyone using the `min_sdk_version` attribute of android_binary (related: https://github.com/bazelbuild/bazel/issues/10556#issuecomment-1425041290), or perhaps modified the [hard coded dexing API level](https://github.com/bazelbuild/bazel/blob/master/src/tools/android/java/com/google/devtools/build/android/r8/CompatDexBuilder.java#L178)?
We are not setting either of those. We are setting min API level in the manifest using the uses-sdk tag. For an application that exhibits it, it is happening consistently. However, triggering it doesn't seem to be related to changing build rules. i.e. We have experienced with an application in which the change that triggered this behavior is purely a code change. It's a little hard for me to create a repro, given that it depends on internal code.

Also—I don't see this added to 6.1 release blockers. Is that still a possibility?
Thanks -- turns out the reason we weren't seeing this internally isn't related to API level exactly, but internally we have other code paths that avoid this problem. We're removing these codepath soon so that the code matches internally and externally, so we'll want the fix for this as well.

It's been pretty tricky to get the fix for this working in all cases, and has required reworking the PR for this a fair bit (https://github.com/bazelbuild/bazel/pull/16369#issuecomment-1442802265). I'm rerunning all our tests again to see if anything still fails with the updates to the PR, so I'm a little hesitant now to block 6.1.0 on this if there are more cases that need to be addressed. If we can't get the fix in for 6.1.0 we cherrypick this into 6.2.0.
After working through https://github.com/bazelbuild/bazel/pull/16369 and talking with the D8 team, they propose (copied from #16369):

1. Update CompatDexBuilder to use the D8 DexFilePerClassFile consumer instead of DexIndexed (this should be a no op for the build as the current build there is already doing one dex per class file, but it uses another D8 API such that it can get the synthetic info).

2. Update the D8 DexFilePerClassFile accept function to also receive the "synthetic context" information if the compiled class is a compiler synthesized class.

3. Update CompatDexBuilder to write the synthetic context information to some build meta info for later access by the DexFileSplitter. Something like this is already done for the desugar dependencies which are written to a proto in META-INF/desugar_deps.

4. Update DexFileSplitter to use the meta-inf of synthetic classes to group with the right contexts.

The D8 team would handle 1 and 2, and the bazel team and/or community could do 3 and 4. This avoids coupling to D8 implementation details and having to parse dex files and hopefully reduces the code complexity a bit.
The fix for this should be ready to submit tomorrow, and then I'll cherrypick it into the Bazel 6.3 release.
There's still this same issue with the `dex_shards` attribute and DexMapper (which is still used for mobile-install "classic" while we open source the starlark replacement). I have an in progress change that applies same fix for DexFileSplitter to DexMapper, which we'll merge into Bazel 6.3 as well.
@ahumesky I think we're still seeing this issue on Linux only when filenames contain capital letters. 
```
11:49:44 18:49:44 [Bazel] bazel-out/k8-opt-exec-ST-7a15de0b7d6a/bin/external/bazel_tools/tools/android/d8_dexmerger --input bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/dexsplits/apk/7.shard.zip --output bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/dexfiles/apk/7.shard.zip '--multidex=given_shard')
11:49:44 18:49:44 [Bazel] # Configuration: 064afe4bfa5b855024151ed67fa23ab997cd53f5f864cf3be12acef0225806a5
11:49:44 18:49:44 [Bazel] # Execution platform: @local_config_platform//:host
11:49:44 18:49:44 [Bazel] Use --sandbox_debug to see verbose messages from the sandbox and retain the sandbox build root for debugging
11:49:44 18:49:44 [Bazel] Error in bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/dexsplits/apk/7.shard.zip:com/snap/ipc/buffers/IBufferMapper$-CC.class.dex:
11:49:44 18:49:44 [Bazel] Attempt at compiling intermediate artifact without its context
11:49:44 18:49:44 [Bazel] Merge failed: Compilation failed to complete, origin: bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/dexsplits/apk/7.shard.zip:com/snap/ipc/buffers/IBufferMapper$-CC.class.dex
11:49:44 18:49:44 [Bazel] Exception in thread "main" java.lang.RuntimeException: Merge failed: Compilation failed to complete, origin: bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/dexsplits/apk/7.shard.zip:com/snap/ipc/buffers/IBufferMapper$-CC.class.dex
11:49:44 18:49:44 [Bazel] at com.google.devtools.build.android.r8.DexFileMerger.main(DexFileMerger.java:412)
11:49:44 18:49:44 [Bazel] Caused by: com.android.tools.r8.CompilationFailedException: Compilation failed to complete, origin: bazel-out/k8-fastbuild/bin/apps/hermosa/hermosa-services/realtime/dexsplits/apk/7.shard.zip:com/snap/ipc/buffers/IBufferMapper$-CC.class.dex
11:49:44 18:49:44 [Bazel] at Version.fakeStackEntry(Version_8.1.51.java:0)
11:49:44 18:49:44 [Bazel] at com.android.tools.r8.M.a(R8_8.1.51_9b737621870a00577be8badac89ea7485627e0cae5c007fd3c57c275c1a604a9:5)
11:49:44 18:49:44 [Bazel] at com.android.tools.r8.utils.R0.a(R8_8.1.51_9b737621870a00577be8badac89ea7485627e0cae5c007fd3c57c275c1a604a9:81)
11:49:44 18:49:44 [Bazel] at com.android.tools.r8.utils.R0.a(R8_8.1.51_9b737621870a00577be8badac89ea7485627e0cae5c007fd3c57c275c1a604a9:32)
11:49:44 18:49:44 [Bazel] at com.android.tools.r8.utils.R0.a(R8_8.1.51_9b737621870a00577be8badac89ea7485627e0cae5c007fd3c57c275c1a604a9:31)
11:49:44 18:49:44 [Bazel] at com.android.tools.r8.utils.R0.b(R8_8.1.51_9b737621870a00577be8badac89ea7485627e0cae5c007fd3c57c275c1a604a9:2)
11:49:44 18:49:44 [Bazel] at com.android.tools.r8.D8.run(R8_8.1.51_9b737621870a00577be8badac89ea7485627e0cae5c007fd3c57c275c1a604a9:6)
11:49:44 18:49:44 [Bazel] at com.google.devtools.build.android.r8.DexFileMerger.run(DexFileMerger.java:395)
11:49:44 18:49:44 [Bazel] at com.google.devtools.build.android.r8.DexFileMerger.main(DexFileMerger.java:409)
11:49:44 18:49:44 [Bazel] Caused by: com.android.tools.r8.utils.b: Attempt at compiling intermediate artifact without its context
```
nkoroste@ is this a custom R8 build (I can't retrace)?

Is IBufferMapper.class in any of the shards?

Pull request review comments:

Patch that resolved the issue:
diff --git a/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java b/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java
index a3b33a042c9a91..8999b9ae269c94 100644
--- a/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java
+++ b/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java
@@ -27,9 +27,13 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
@@ -39,10 +43,12 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Scanner;
 import java.util.Set;
 import java.util.TreeSet;
 
@@ -416,12 +422,22 @@ private void checkConfig() throws IOException {
     filter = filterFile == null && filterInputStream == null ? null : readPaths(filterFile);
   }
 
-  /**
-   * Parses the entries and assign each entry to an output file.
-   */
-  private void split() {
+  /** Parses the entries and assign each entry to an output file. */
+  private void split() throws IOException {
+
+    // A map of class (a "context") to its inner synthetic classes from D8.
+    SetMultimap<String, String> syntheticClassContexts = HashMultimap.create();
+
     for (ZipIn in : inputs) {
       CentralDirectory cdir = centralDirectories.get(in.getFilename());
+
+      for (DirectoryEntry entry : cdir.list()) {
+        if (entry.getFilename().equals("META-INF/synthetic-contexts.map")) {
+          parseSyntheticContextsMap(in.entryFor(entry).getContent(), syntheticClassContexts);
+          break;
+        }
+      }
+
       for (DirectoryEntry entry : cdir.list()) {
         String filename = normalizedFilename(entry.getFilename());
         if (!inputFilter.apply(filename)) {
@@ -438,17 +454,31 @@ private void split() {
         }
       }
     }
+
     Splitter splitter = new Splitter(outputs.size(), classes.size());
+
     if (filter != null) {
       // Assign files in the filter to the first output file.
-      splitter.assign(Sets.filter(filter, inputFilter));
+      splitter.assignAllToCurrentShard(Sets.filter(filter, inputFilter));
       splitter.nextShard(); // minimal initial shard
     }
+
+    Set<String> allSyntheticClasses = new HashSet<>(syntheticClassContexts.values());
+
     for (String path : classes) {
-      // Use normalized filename so the filter file doesn't have to change
-      int assignment = splitter.assign(path);
-      Preconditions.checkState(assignment >= 0 && assignment < zipOuts.length);
-      assignments.put(path, zipOuts[assignment]);
+      if (!allSyntheticClasses.contains(path)) {
+
+        // Use normalized filename so the filter file doesn't have to change
+        int assignment = splitter.assign(path);
+        Set<String> syntheticClasses = syntheticClassContexts.get(path);
+        splitter.assignAllToCurrentShard(syntheticClasses);
+        Preconditions.checkState(assignment >= 0 && assignment < zipOuts.length);
+
+        assignments.put(path, zipOuts[assignment]);
+        for (String syntheticClass : syntheticClasses) {
+          assignments.put(syntheticClass, zipOuts[assignment]);
+        }
+      }
     }
   }
 
@@ -494,6 +524,23 @@ private String fixPath(String path) {
     return path;
   }
 
+  private static void parseSyntheticContextsMap(
+      ByteBuffer byteBuffer, Multimap<String, String> syntheticClassContexts) {
+    // The ByteBuffer returned from the Splitter's zip library is not backed by an accessible array,
+    // so ByteBuffer.array() is not supported, so we must go the long way.
+    byte[] bytes = new byte[byteBuffer.remaining()];
+    byteBuffer.get(bytes);
+    Scanner scanner = new Scanner(new ByteArrayInputStream(bytes), UTF_8);
+    scanner.useDelimiter("[;\n]");
+    while (scanner.hasNext()) {
+      String syntheticClass = scanner.next();
+      String context = scanner.next();
+      // The context map uses class names, whereas SplitZip uses class file names, so add ".class"
+      // here to make it easier to work with the map in the rest of the code.
+      syntheticClassContexts.put(context + ".class", syntheticClass + ".class");
+    }
+  }
+
   private void verbose(String msg) {
     if (verbose) {
       System.out.println(msg);
diff --git a/src/tools/android/java/com/google/devtools/build/android/ziputils/Splitter.java b/src/tools/android/java/com/google/devtools/build/android/ziputils/Splitter.java
index 7f1a469f10368d..2c03b74cacac09 100644
--- a/src/tools/android/java/com/google/devtools/build/android/ziputils/Splitter.java
+++ b/src/tools/android/java/com/google/devtools/build/android/ziputils/Splitter.java
@@ -24,10 +24,10 @@ class Splitter {
 
   private final int numberOfShards;
   private final Map<String, Integer> assigned;
-  private int size = 0;
-  private int shard = 0;
+  private int size = 0; // Number of classes in the current shard
+  private int shard = 0; // Current shard number
   private String prevPath = null;
-  private int remaining;
+  private int remaining; // Number of classes remaining to be assigned shards
   private int idealSize;
   private int almostFull;
 
@@ -51,15 +51,15 @@ public Splitter(int numberOfShards, int expectedSize) {
    * remaining entries to process is adjusted, by subtracting the number of as-of-yet unassigned
    * entries from the filter.
    */
-  public void assign(Collection<String> filter) {
-    if (filter != null) {
-      for (String s : filter) {
-        if (!assigned.containsKey(s)) {
+  public void assignAllToCurrentShard(Collection<String> entries) {
+    if (entries != null) {
+      for (String e : entries) {
+        if (!assigned.containsKey(e)) {
           remaining--;
         }
-        assigned.put(s, shard);
+        assigned.put(e, shard);
       }
-      size = filter.size();
+      size += entries.size();
     }
   }
 
diff --git a/tools/android/runtime_deps/BUILD.bazel b/tools/android/runtime_deps/BUILD.bazel
index 9c50eb42dc2978..b56610231034a7 100644
--- a/tools/android/runtime_deps/BUILD.bazel
+++ b/tools/android/runtime_deps/BUILD.bazel
@@ -4,7 +4,7 @@
 # extract all Android rules and tools out of Bazel and into rules_android
 # and tools_android.
 
-load("@bazel_tools//tools/build_defs/pkg:pkg.bzl", "pkg_tar")
+load("@rules_pkg//pkg:tar.bzl", "pkg_tar")
 
 filegroup(
     name = "srcs",
@@ -84,5 +84,9 @@ pkg_tar(
         "//src/java_tools/import_deps_checker/java/com/google/devtools/build/importdeps:ImportDepsChecker_deploy.jar",
         "//src/tools/android/java/com/google/devtools/build/android:all_android_tools_deploy.jar",
     ],
+    remap_paths = {
+        "src/java_tools/import_deps_checker/java/com/google/devtools/build/importdeps:ImportDepsChecker_deploy.jar": "ImportDepsChecker_deploy.jar",
+        "src/tools/android/java/com/google/devtools/build/android:all_android_tools_deploy.jar": "all_android_tools_deploy.jar",
+    },
     visibility = ["//src/test/shell/bazel:__subpackages__"],
 )


Test patch:
diff --git a/src/test/java/com/google/devtools/build/android/ziputils/SplitterTest.java b/src/test/java/com/google/devtools/build/android/ziputils/SplitterTest.java
index 3f24eda65efeaf..568db636ccb0ca 100644
--- a/src/test/java/com/google/devtools/build/android/ziputils/SplitterTest.java
+++ b/src/test/java/com/google/devtools/build/android/ziputils/SplitterTest.java
@@ -51,7 +51,7 @@ public void testAssign() {
       filter.add("dir" + i + ARCHIVE_FILE_SEPARATOR + "file" + i + CLASS_SUFFIX);
     }
     Splitter splitter = new Splitter(size + 1, input.size());
-    splitter.assign(filter);
+    splitter.assignAllToCurrentShard(filter);
     splitter.nextShard();
     output = new LinkedHashMap<>();
     for (String path : input) {
diff --git a/src/test/shell/BUILD b/src/test/shell/BUILD
index f3054b8c6c8bf4..6ddf5da707e68f 100644
--- a/src/test/shell/BUILD
+++ b/src/test/shell/BUILD
@@ -29,6 +29,7 @@ gen_workspace_stanza(
         "rules_cc",
         "rules_java",
         "rules_license",
+        "rules_pkg",
         "rules_proto",
     ],
     template = "testenv.sh.tmpl",
diff --git a/src/test/shell/bazel/android/BUILD b/src/test/shell/bazel/android/BUILD
index 4a4d14d1f4672b..9921da218dd1f7 100644
--- a/src/test/shell/bazel/android/BUILD
+++ b/src/test/shell/bazel/android/BUILD
@@ -238,11 +238,14 @@ android_sh_test(
     name = "DexFileSplitter_synthetic_classes_test",
     size = "medium",
     srcs = ["DexFileSplitter_synthetic_classes_test.sh"],
+    # Fixes to DexMapper are not released yet.
+    create_test_with_released_tools = False,
     data = [
         ":android_helper",
         "//external:android_sdk_for_testing",
         "//src/test/shell/bazel:test-deps",
     ],
+    shard_count = 2,
     tags = [
         "no_windows",
     ],
diff --git a/src/test/shell/bazel/android/DexFileSplitter_synthetic_classes_test.sh b/src/test/shell/bazel/android/DexFileSplitter_synthetic_classes_test.sh
index bd70e9f4533437..44c4a00020f2ee 100755
--- a/src/test/shell/bazel/android/DexFileSplitter_synthetic_classes_test.sh
+++ b/src/test/shell/bazel/android/DexFileSplitter_synthetic_classes_test.sh
@@ -26,18 +26,20 @@
 # Load the test setup defined in the parent directory
 CURRENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
+source "${CURRENT_DIR}/../../integration_test_setup.sh" \
+  || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
+
 source "${CURRENT_DIR}/android_helper.sh" \
   || { echo "android_helper.sh not found!" >&2; exit 1; }
 fail_if_no_android_sdk
 
-source "${CURRENT_DIR}/../../integration_test_setup.sh" \
-  || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
-
 resolve_android_toolchains "$1"
 
-function test_DexFileSplitter_synthetic_classes_crossing_dexfiles() {
-  create_new_workspace
-  setup_android_sdk_support
+function create_test_app() {
+
+  inner_class_count="$1"
+  lambda_count="$2"
+  dex_shard_count="$3"
 
   mkdir -p java/com/testapp
 
@@ -81,7 +83,7 @@ public class MainActivity extends Activity {
 }
 EOF
 
-  generate_java_file_with_many_synthetic_classes > java/com/testapp/BigLib.java
+  generate_java_file_with_many_synthetic_classes "$1" "$2" > java/com/testapp/BigLib.java
 
   cat > java/com/testapp/BUILD <<EOF
 android_binary(
@@ -90,26 +92,22 @@ android_binary(
         "MainActivity.java",
         ":BigLib.java",
     ],
+    dex_shards = $dex_shard_count,
     manifest = "AndroidManifest.xml",
 )
 EOF
-
-  bazel build java/com/testapp || fail "Test app should have built succesfully"
-
-  dex_file_count=$(unzip -l bazel-bin/java/com/testapp/testapp.apk | grep "classes[0-9]*.dex" | wc -l)
-  if [[ ! "$dex_file_count" -ge "2" ]]; then
-    echo "Expected at least 2 dexes in app, found: $dex_file_count"
-    exit 1
-  fi
 }
 
-
 function generate_java_file_with_many_synthetic_classes() {
+
+  inner_class_count="$1"
+  lambda_count="$2"
+
   echo "package com.testapp;"
   echo "public class BigLib {"
 
   # First generate enough inner classes to fill up most of the dex
-  for (( i = 0; i < 21400; i++ )) do
+  for (( i = 0; i < $inner_class_count; i++ )) do
 
     echo "  public static class Bar$i {"
     echo "    public int bar() {"
@@ -129,7 +127,7 @@ function generate_java_file_with_many_synthetic_classes() {
   echo "    public IntSupplier[] foo() {"
   echo "      return new IntSupplier[] {"
 
-  for ((i = 0; i < 6000; i++ )) do
+  for ((i = 0; i < $lambda_count; i++ )) do
     echo "        () -> $i,"
   done
 
@@ -139,4 +137,42 @@ function generate_java_file_with_many_synthetic_classes() {
   echo "}"
 }
 
+function test_DexFileSplitter_synthetic_classes_crossing_dexfiles() {
+  create_new_workspace
+  setup_android_sdk_support
+
+  # dex_shards default is 1
+  create_test_app 21400 6000 1
+
+  bazel build java/com/testapp || fail "Test app should have built succesfully"
+
+  dex_file_count=$(unzip -l bazel-bin/java/com/testapp/testapp.apk | grep "classes[0-9]*.dex" | wc -l)
+  if [[ ! "$dex_file_count" -ge "2" ]]; then
+    echo "Expected at least 2 dexes in app, found: $dex_file_count"
+    exit 1
+  fi
+}
+
+function test_DexMapper_synthetic_classes_crossing_dexfiles() {
+  create_new_workspace
+  setup_android_sdk_support
+
+  # 3 inner classes, 6 lambdas (and thus 6 synthetics from D8) and 5 dex_shards
+  # is one magic combination to repro synthetics being separated from their
+  # context / enclosing classes.
+  create_test_app 3 6 5
+
+  echo here2
+  echo $TEST_TMPDIR/bazelrc
+  cat $TEST_TMPDIR/bazelrc
+
+  bazel build java/com/testapp || fail "Test app should have built succesfully"
+
+  dex_file_count=$(unzip -l bazel-bin/java/com/testapp/testapp.apk | grep "classes[0-9]*.dex" | wc -l)
+  if [[ ! "$dex_file_count" -eq "5" ]]; then
+    echo "Expected 5 dexes in app, found: $dex_file_count"
+    exit 1
+  fi
+}
+
 run_suite "Tests for DexFileSplitter with synthetic classes crossing dexfiles"
\ No newline at end of file
diff --git a/src/test/shell/bazel/android/android_sh_test.bzl b/src/test/shell/bazel/android/android_sh_test.bzl
index 40f21648a2f939..8c0bbc626bc40e 100644
--- a/src/test/shell/bazel/android/android_sh_test.bzl
+++ b/src/test/shell/bazel/android/android_sh_test.bzl
@@ -29,20 +29,36 @@ CHECK_FOR_ANDROID_SDK = select(
     no_match_error = "This test requires an android SDK, and one isn't present. Make sure to uncomment the android rules in the WORKSPACE.",
 )
 
-def android_sh_test(**kwargs):
+def android_sh_test(create_test_with_released_tools = True, **kwargs):
+    """Creates versions of the test with and without platforms and head android tools.
+
+    Args:
+        create_test_with_released_tools: Whether to create a version of the test with the released
+        android tools, for when the code under test relies on not-yet-released code.
+        **kwargs: Args to sh_test
+    """
     name = kwargs.pop("name")
     data = kwargs.pop("data")
     if not data:
         data = []
     data = data + CHECK_FOR_ANDROID_SDK
 
-    # Test with released android_tools version.
-    native.sh_test(
-        name = name,
-        args = ["--without_platforms"],
-        data = data,
-        **kwargs
-    )
+    if create_test_with_released_tools:
+        # Test with released android_tools version.
+        native.sh_test(
+            name = name,
+            args = ["--without_platforms"],
+            data = data,
+            **kwargs
+        )
+
+        # Test with platform-based toolchain resolution.
+        native.sh_test(
+            name = name + "_with_platforms",
+            data = data,
+            args = ["--with_platforms"],
+            **kwargs
+        )
 
     # Test with android_tools version that's built at the same revision
     # as the test itself.
@@ -54,11 +70,3 @@ def android_sh_test(**kwargs):
         ],
         **kwargs
     )
-
-    # Test with platform-based toolchain resolution.
-    native.sh_test(
-        name = name + "_with_platforms",
-        data = data,
-        args = ["--with_platforms"],
-        **kwargs
-    )
diff --git a/src/test/shell/testenv.sh.tmpl b/src/test/shell/testenv.sh.tmpl
index 7ee2e772951205..7a9aac74580a1c 100755
--- a/src/test/shell/testenv.sh.tmpl
+++ b/src/test/shell/testenv.sh.tmpl
@@ -330,6 +330,7 @@ EOF
         "rules_license"
         "rules_proto"
         "rules_python"
+        "rules_pkg"
     )
     for repo in "${repos[@]}"; do
       reponame="${repo%"_for_testing"}"
@@ -554,6 +555,14 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 EOF
 }
 
+function add_rules_pkg_to_workspace() {
+  cat >> "$1"<<EOF
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+{rules_pkg}
+EOF
+}
+
 function add_rules_proto_to_workspace() {
   cat >> "$1"<<EOF
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
@@ -575,6 +584,7 @@ EOF
   add_rules_cc_to_workspace "WORKSPACE"
   add_rules_java_to_workspace "WORKSPACE"
   add_rules_license_to_workspace "WORKSPACE"
+  add_rules_pkg_to_workspace "WORKSPACE"
   add_rules_proto_to_workspace "WORKSPACE"
 
   maybe_setup_python_windows_workspace