SampleQueue: Let subclasses easily invalidate format adjustment

This is a nice-regardless improvement to SampleQueue, which will
likely to used to fix the referenced issue. It makes it possible
for SampleQueue subclasses to support dynamic changes to format
adjustment in a non-hacky way.

Issue: #6903
PiperOrigin-RevId: 292314720

library/core/src/main/java/com/google/android/exoplayer2/source/SampleQueue.java

@@ -16,6 +16,7 @@
 package com.google.android.exoplayer2.source;
 
 import android.os.Looper;
+import androidx.annotation.CallSuper;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.exoplayer2.C;
@@ -81,8 +82,8 @@ public class SampleQueue implements TrackOutput {
   private Format upstreamCommittedFormat;
   private int upstreamSourceId;
 
-  private boolean pendingFormatAdjustment;
+  private boolean pendingUpstreamFormatAdjustment;
-  private Format lastUnadjustedFormat;
+  private Format unadjustedUpstreamFormat;
   private long sampleOffsetUs;
   private boolean pendingSplice;
 
@@ -146,6 +147,7 @@ public class SampleQueue implements TrackOutput {
     isLastSampleQueued = false;
     upstreamCommittedFormat = null;
     if (resetUpstreamFormat) {
+      unadjustedUpstreamFormat = null;
       upstreamFormat = null;
       upstreamFormatRequired = true;
     }
@@ -433,7 +435,7 @@ public class SampleQueue implements TrackOutput {
   public void setSampleOffsetUs(long sampleOffsetUs) {
     if (this.sampleOffsetUs != sampleOffsetUs) {
       this.sampleOffsetUs = sampleOffsetUs;
-      pendingFormatAdjustment = true;
+      invalidateUpstreamFormatAdjustment();
     }
   }
 
@@ -449,13 +451,13 @@ public class SampleQueue implements TrackOutput {
   // TrackOutput implementation. Called by the loading thread.
 
   @Override
-  public void format(Format unadjustedFormat) {
+  public final void format(Format unadjustedUpstreamFormat) {
-    Format adjustedFormat = getAdjustedSampleFormat(unadjustedFormat, sampleOffsetUs);
+    Format adjustedUpstreamFormat = getAdjustedUpstreamFormat(unadjustedUpstreamFormat);
-    boolean formatChanged = setUpstreamFormat(adjustedFormat);
+    pendingUpstreamFormatAdjustment = false;
-    lastUnadjustedFormat = unadjustedFormat;
+    this.unadjustedUpstreamFormat = unadjustedUpstreamFormat;
-    pendingFormatAdjustment = false;
+    boolean upstreamFormatChanged = setUpstreamFormat(adjustedUpstreamFormat);
-    if (upstreamFormatChangeListener != null && formatChanged) {
+    if (upstreamFormatChangeListener != null && upstreamFormatChanged) {
-      upstreamFormatChangeListener.onUpstreamFormatChanged(adjustedFormat);
+      upstreamFormatChangeListener.onUpstreamFormatChanged(adjustedUpstreamFormat);
     }
   }
 
@@ -477,8 +479,8 @@ public class SampleQueue implements TrackOutput {
       int size,
       int offset,
       @Nullable CryptoData cryptoData) {
-    if (pendingFormatAdjustment) {
+    if (pendingUpstreamFormatAdjustment) {
-      format(lastUnadjustedFormat);
+      format(unadjustedUpstreamFormat);
     }
     timeUs += sampleOffsetUs;
     if (pendingSplice) {
@@ -491,6 +493,32 @@ public class SampleQueue implements TrackOutput {
     commitSample(timeUs, flags, absoluteOffset, size, cryptoData);
   }
 
+  /**
+   * Invalidates the last upstream format adjustment. {@link #getAdjustedUpstreamFormat(Format)}
+   * will be called to adjust the upstream {@link Format} again before the next sample is queued.
+   */
+  protected final void invalidateUpstreamFormatAdjustment() {
+    pendingUpstreamFormatAdjustment = true;
+  }
+
+  /**
+   * Adjusts the upstream {@link Format} (i.e., the {@link Format} that was most recently passed to
+   * {@link #format(Format)}).
+   *
+   * <p>The default implementation incorporates the sample offset passed to {@link
+   * #setSampleOffsetUs(long)} into {@link Format#subsampleOffsetUs}.
+   *
+   * @param format The {@link Format} to adjust.
+   * @return The adjusted {@link Format}.
+   */
+  @CallSuper
+  protected Format getAdjustedUpstreamFormat(Format format) {
+    if (sampleOffsetUs != 0 && format.subsampleOffsetUs != Format.OFFSET_SAMPLE_RELATIVE) {
+      format = format.copyWithSubsampleOffsetUs(format.subsampleOffsetUs + sampleOffsetUs);
+    }
+    return format;
+  }
+
   // Internal methods.
 
   /** Rewinds the read position to the first sample in the queue. */
@@ -883,23 +911,6 @@ public class SampleQueue implements TrackOutput {
     return relativeIndex < capacity ? relativeIndex : relativeIndex - capacity;
   }
 
-  /**
-   * Adjusts a {@link Format} to incorporate a sample offset into {@link Format#subsampleOffsetUs}.
-   *
-   * @param format The {@link Format} to adjust.
-   * @param sampleOffsetUs The offset to apply.
-   * @return The adjusted {@link Format}.
-   */
-  private static Format getAdjustedSampleFormat(Format format, long sampleOffsetUs) {
-    if (format == null) {
-      return null;
-    }
-    if (sampleOffsetUs != 0 && format.subsampleOffsetUs != Format.OFFSET_SAMPLE_RELATIVE) {
-      format = format.copyWithSubsampleOffsetUs(format.subsampleOffsetUs + sampleOffsetUs);
-    }
-    return format;
-  }
-
   /** A holder for sample metadata not held by {@link DecoderInputBuffer}. */
   /* package */ static final class SampleExtrasHolder {
 

library/core/src/test/java/com/google/android/exoplayer2/source/SampleQueueTest.java

@@ -15,6 +15,7 @@
  */
 package com.google.android.exoplayer2.source;
 
+import static com.google.android.exoplayer2.C.BUFFER_FLAG_KEY_FRAME;
 import static com.google.android.exoplayer2.C.RESULT_BUFFER_READ;
 import static com.google.android.exoplayer2.C.RESULT_FORMAT_READ;
 import static com.google.android.exoplayer2.C.RESULT_NOTHING_READ;
@@ -40,6 +41,7 @@ import com.google.android.exoplayer2.upstream.DefaultAllocator;
 import com.google.android.exoplayer2.util.ParsableByteArray;
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -136,7 +138,7 @@ public final class SampleQueueTest {
 
   @Before
   @SuppressWarnings("unchecked")
-  public void setUp() throws Exception {
+  public void setUp() {
     allocator = new DefaultAllocator(false, ALLOCATION_SIZE);
     mockDrmSessionManager =
         (DrmSessionManager<ExoMediaCrypto>) Mockito.mock(DrmSessionManager.class);
@@ -149,7 +151,7 @@ public final class SampleQueueTest {
   }
 
   @After
-  public void tearDown() throws Exception {
+  public void tearDown() {
     allocator = null;
     sampleQueue = null;
     formatHolder = null;
@@ -837,12 +839,93 @@ public final class SampleQueueTest {
   }
 
   @Test
-  public void testSetSampleOffset() {
+  public void testSetSampleOffsetBeforeData() {
     long sampleOffsetUs = 1000;
     sampleQueue.setSampleOffsetUs(sampleOffsetUs);
     writeTestData();
-    assertReadTestData(null, 0, 8, sampleOffsetUs);
+    assertReadTestData(
-    assertReadEndOfStream(false);
+        /* startFormat= */ null, /* firstSampleIndex= */ 0, /* sampleCount= */ 8, sampleOffsetUs);
+    assertReadEndOfStream(/* formatRequired= */ false);
+  }
+
+  @Test
+  public void testSetSampleOffsetBetweenSamples() {
+    writeTestData();
+    long sampleOffsetUs = 1000;
+    sampleQueue.setSampleOffsetUs(sampleOffsetUs);
+
+    // Write a final sample now the offset is set.
+    long unadjustedTimestampUs = LAST_SAMPLE_TIMESTAMP + 1234;
+    writeSample(DATA, unadjustedTimestampUs, /* sampleFlags= */ 0);
+
+    assertReadTestData();
+    // We expect to read the format adjusted to account for the sample offset, followed by the final
+    // sample and then the end of stream.
+    assertReadFormat(
+        /* formatRequired= */ false, FORMAT_2.copyWithSubsampleOffsetUs(sampleOffsetUs));
+    assertReadSample(
+        unadjustedTimestampUs + sampleOffsetUs,
+        /* isKeyFrame= */ false,
+        /* isEncrypted= */ false,
+        DATA,
+        /* offset= */ 0,
+        DATA.length);
+    assertReadEndOfStream(/* formatRequired= */ false);
+  }
+
+  @Test
+  public void testAdjustUpstreamFormat() {
+    String label = "label";
+    sampleQueue =
+        new SampleQueue(allocator, mockDrmSessionManager) {
+          @Override
+          public Format getAdjustedUpstreamFormat(Format format) {
+            return super.getAdjustedUpstreamFormat(format.copyWithLabel(label));
+          }
+        };
+
+    writeFormat(FORMAT_1);
+    assertReadFormat(/* formatRequired= */ false, FORMAT_1.copyWithLabel(label));
+    assertReadEndOfStream(/* formatRequired= */ false);
+  }
+
+  @Test
+  public void testInvalidateUpstreamFormatAdjustment() {
+    AtomicReference<String> label = new AtomicReference<>("label1");
+    sampleQueue =
+        new SampleQueue(allocator, mockDrmSessionManager) {
+          @Override
+          public Format getAdjustedUpstreamFormat(Format format) {
+            return super.getAdjustedUpstreamFormat(format.copyWithLabel(label.get()));
+          }
+        };
+
+    writeFormat(FORMAT_1);
+    writeSample(DATA, /* timestampUs= */ 0, BUFFER_FLAG_KEY_FRAME);
+
+    // Make a change that'll affect the SampleQueue's format adjustment, and invalidate it.
+    label.set("label2");
+    sampleQueue.invalidateUpstreamFormatAdjustment();
+
+    writeSample(DATA, /* timestampUs= */ 1, /* sampleFlags= */ 0);
+
+    assertReadFormat(/* formatRequired= */ false, FORMAT_1.copyWithLabel("label1"));
+    assertReadSample(
+        /* timeUs= */ 0,
+        /* isKeyFrame= */ true,
+        /* isEncrypted= */ false,
+        DATA,
+        /* offset= */ 0,
+        DATA.length);
+    assertReadFormat(/* formatRequired= */ false, FORMAT_1.copyWithLabel("label2"));
+    assertReadSample(
+        /* timeUs= */ 1,
+        /* isKeyFrame= */ false,
+        /* isEncrypted= */ false,
+        DATA,
+        /* offset= */ 0,
+        DATA.length);
+    assertReadEndOfStream(/* formatRequired= */ false);
   }
 
   @Test
@@ -851,7 +934,8 @@ public final class SampleQueueTest {
     sampleQueue.splice();
     // Splice should succeed, replacing the last 4 samples with the sample being written.
     long spliceSampleTimeUs = SAMPLE_TIMESTAMPS[4];
-    writeSample(DATA, spliceSampleTimeUs, FORMAT_SPLICED, C.BUFFER_FLAG_KEY_FRAME);
+    writeFormat(FORMAT_SPLICED);
+    writeSample(DATA, spliceSampleTimeUs, C.BUFFER_FLAG_KEY_FRAME);
     assertReadTestData(null, 0, 4);
     assertReadFormat(false, FORMAT_SPLICED);
     assertReadSample(spliceSampleTimeUs, true, /* isEncrypted= */ false, DATA, 0, DATA.length);
@@ -865,7 +949,8 @@ public final class SampleQueueTest {
     sampleQueue.splice();
     // Splice should fail, leaving the last 4 samples unchanged.
     long spliceSampleTimeUs = SAMPLE_TIMESTAMPS[3];
-    writeSample(DATA, spliceSampleTimeUs, FORMAT_SPLICED, C.BUFFER_FLAG_KEY_FRAME);
+    writeFormat(FORMAT_SPLICED);
+    writeSample(DATA, spliceSampleTimeUs, C.BUFFER_FLAG_KEY_FRAME);
     assertReadTestData(SAMPLE_FORMATS[3], 4, 4);
     assertReadEndOfStream(false);
 
@@ -874,7 +959,8 @@ public final class SampleQueueTest {
     sampleQueue.splice();
     // Splice should succeed, replacing the last 4 samples with the sample being written
     spliceSampleTimeUs = SAMPLE_TIMESTAMPS[3] + 1;
-    writeSample(DATA, spliceSampleTimeUs, FORMAT_SPLICED, C.BUFFER_FLAG_KEY_FRAME);
+    writeFormat(FORMAT_SPLICED);
+    writeSample(DATA, spliceSampleTimeUs, C.BUFFER_FLAG_KEY_FRAME);
     assertReadFormat(false, FORMAT_SPLICED);
     assertReadSample(spliceSampleTimeUs, true, /* isEncrypted= */ false, DATA, 0, DATA.length);
     assertReadEndOfStream(false);
@@ -888,7 +974,8 @@ public final class SampleQueueTest {
     sampleQueue.splice();
     // Splice should succeed, replacing the last 4 samples with the sample being written.
     long spliceSampleTimeUs = SAMPLE_TIMESTAMPS[4];
-    writeSample(DATA, spliceSampleTimeUs, FORMAT_SPLICED, C.BUFFER_FLAG_KEY_FRAME);
+    writeFormat(FORMAT_SPLICED);
+    writeSample(DATA, spliceSampleTimeUs, C.BUFFER_FLAG_KEY_FRAME);
     assertReadTestData(null, 0, 4, sampleOffsetUs);
     assertReadFormat(false, FORMAT_SPLICED.copyWithSubsampleOffsetUs(sampleOffsetUs));
     assertReadSample(
@@ -938,9 +1025,13 @@ public final class SampleQueueTest {
     }
   }
 
-  /** Writes a single sample to {@code sampleQueue}. */
+  /** Writes a {@link Format} to the {@code sampleQueue}. */
-  private void writeSample(byte[] data, long timestampUs, Format format, int sampleFlags) {
+  private void writeFormat(Format format) {
     sampleQueue.format(format);
+  }
+
+  /** Writes a single sample to {@code sampleQueue}. */
+  private void writeSample(byte[] data, long timestampUs, int sampleFlags) {
     sampleQueue.sampleData(new ParsableByteArray(data), data.length);
     sampleQueue.sampleMetadata(timestampUs, sampleFlags, data.length, 0, null);
   }

library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsSampleStreamWrapper.java

@@ -1290,15 +1290,17 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
     }
 
     @Override
-    public void format(Format format) {
+    public Format getAdjustedUpstreamFormat(Format format) {
-      DrmInitData drmInitData = format.drmInitData;
+      @Nullable DrmInitData drmInitData = format.drmInitData;
       if (drmInitData != null) {
+        @Nullable
         DrmInitData overridingDrmInitData = this.overridingDrmInitData.get(drmInitData.schemeType);
         if (overridingDrmInitData != null) {
           drmInitData = overridingDrmInitData;
         }
       }
-      super.format(format.copyWithAdjustments(drmInitData, getAdjustedMetadata(format.metadata)));
+      return super.getAdjustedUpstreamFormat(
+          format.copyWithAdjustments(drmInitData, getAdjustedMetadata(format.metadata)));
     }
 
     /**