Add finished callback to retracer

Bug: b/299231709
Change-Id: I7132d4d5dc3ddc341430a8ca6413ecc19fe89cba
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierBase.java b/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierBase.java
index a15f89f..1d913e3 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierBase.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierBase.java
@@ -191,5 +191,10 @@
         MappingSupplierInternalImpl.createInternal(classNameMapper), diagnosticsHandler);
   }
 
+  @Override
+  public void finished(DiagnosticsHandler handler) {
+    finishedCallback.finished(handler);
+  }
+
   public abstract T self();
 }
diff --git a/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java b/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
index 255b693..7c74234 100644
--- a/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
@@ -21,6 +21,7 @@
 import com.android.tools.r8.retrace.PartitionMappingSupplier;
 import com.android.tools.r8.retrace.partition.testclasses.R8ZipContainerMappingFileTestClasses;
 import com.android.tools.r8.retrace.partition.testclasses.R8ZipContainerMappingFileTestClasses.Main;
+import com.android.tools.r8.utils.BooleanBox;
 import com.android.tools.r8.utils.ZipUtils.ZipBuilder;
 import com.google.common.io.ByteStreams;
 import java.io.IOException;
@@ -92,9 +93,12 @@
             .getOriginalStackTrace();
 
     assertTrue(Files.exists(pgMapFile));
+    BooleanBox calledFinished = new BooleanBox();
     assertThat(
-        originalStackTrace.retrace(createMappingSupplierFromPartitionZip(pgMapFile)),
+        originalStackTrace.retrace(
+            createMappingSupplierFromPartitionZip(pgMapFile, calledFinished::set)),
         isSame(EXPECTED));
+    assertTrue(calledFinished.get());
   }
 
   private PartitionMapConsumer createPartitionZipConsumer(Path pgMapFile) throws IOException {
@@ -130,8 +134,8 @@
     };
   }
 
-  private PartitionMappingSupplier createMappingSupplierFromPartitionZip(Path pgMapFile)
-      throws IOException {
+  private PartitionMappingSupplier createMappingSupplierFromPartitionZip(
+      Path pgMapFile, Runnable finishedCallback) throws IOException {
     ZipFile zipFile = new ZipFile(pgMapFile.toFile());
     byte[] metadata = ByteStreams.toByteArray(zipFile.getInputStream(zipFile.getEntry("METADATA")));
     return PartitionMappingSupplier.builder()
@@ -144,6 +148,15 @@
                 throw new RuntimeException(e);
               }
             })
+        .setFinishedPartitionMappingCallback(
+            handler -> {
+              try {
+                zipFile.close();
+                finishedCallback.run();
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            })
         .build();
   }
 }