From 4a6c5c7598c6d6faa4098203869dde8481b0472a Mon Sep 17 00:00:00 2001
From: Colin Kho <colinkho@meta.com>
Date: Mon, 19 Aug 2024 10:28:14 -0700
Subject: [PATCH] Cache ChunkIndex once load is done

---
 .../exoplayer/source/chunk/InitializationChunk.java       | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/chunk/InitializationChunk.java b/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/chunk/InitializationChunk.java
index aad438ef77..00b6b95df2 100644
--- a/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/chunk/InitializationChunk.java
+++ b/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/chunk/InitializationChunk.java
@@ -39,6 +39,7 @@ public final class InitializationChunk extends Chunk {
   private final ChunkExtractor chunkExtractor;
 
   private @MonotonicNonNull TrackOutputProvider trackOutputProvider;
+  private @Nullable ChunkIndex chunkIndex;
   private long nextLoadPosition;
   private volatile boolean loadCanceled;
 
@@ -105,14 +106,19 @@ public final class InitializationChunk extends Chunk {
         while (!loadCanceled && chunkExtractor.read(input)) {}
       } finally {
         nextLoadPosition = input.getPosition() - dataSpec.position;
+        chunkIndex = chunkExtractor.getChunkIndex();
       }
     } finally {
       DataSourceUtil.closeQuietly(dataSource);
     }
   }
 
+  /**
+   * Returns the {@link ChunkIndex} most recently obtained from the latest {@link #load()}, or null
+   * if a {@link ChunkIndex} has not been obtained.
+   */
   @Nullable
   public ChunkIndex getChunkIndex() {
-    return chunkExtractor.getChunkIndex();
+    return chunkIndex;
   }
 }