Re-order CEA-6/708 samples during extraction instead of rendering

This is required before we can move CEA-6/708 parsing from the rendering
side of the sample queue to the extraction side.

This re-ordering is needed for video encodings with different decoder
and presentation orders, because the CEA-6/708 data is attached to each
frame and needs to be processed in presentation order instead of decode
order. This change re-orders frames within a group-of-pictures, but also
takes advantage of `maxNumReorderFrames/Pics` values to cap the size of
the re-ordering queue, allowing caption data to be released 'earlier'
than the end of a GoP.

Annex D of the CEA-708 spec (which also applies for CEA-608 embedded in
SEI messages), makes the need to re-order from decode to presentation
order clear.

PiperOrigin-RevId: 648648002

libraries/container/src/main/java/androidx/media3/container/ReorderingSeiMessageQueue.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package androidx.media3.container;
+
+import static androidx.annotation.RestrictTo.Scope.LIBRARY_GROUP;
+import static androidx.media3.common.util.Assertions.checkState;
+import static androidx.media3.common.util.Util.castNonNull;
+
+import androidx.annotation.RestrictTo;
+import androidx.media3.common.C;
+import androidx.media3.common.util.ParsableByteArray;
+import androidx.media3.common.util.UnstableApi;
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.PriorityQueue;
+import java.util.concurrent.atomic.AtomicLong;
+
+/** A queue of SEI messages, ordered by presentation timestamp. */
+@UnstableApi
+@RestrictTo(LIBRARY_GROUP)
+public final class ReorderingSeiMessageQueue {
+
+  /** Functional interface to handle an SEI message that is being removed from the queue. */
+  public interface SeiConsumer {
+    /** Handles an SEI message that is being removed from the queue. */
+    void consume(long presentationTimeUs, ParsableByteArray seiBuffer);
+  }
+
+  private final SeiConsumer seiConsumer;
+  private final AtomicLong tieBreakGenerator = new AtomicLong();
+
+  /**
+   * Pool of re-usable {@link SeiMessage} objects to avoid repeated allocations. Elements should be
+   * added and removed from the 'tail' of the queue (with {@link Deque#push(Object)} and {@link
+   * Deque#pop()}), to avoid unnecessary array copying.
+   */
+  private final ArrayDeque<SeiMessage> unusedSeiMessages;
+
+  private final PriorityQueue<SeiMessage> pendingSeiMessages;
+
+  private int reorderingQueueSize;
+
+  /**
+   * Creates an instance, initially with no max size.
+   *
+   * @param seiConsumer Callback to invoke when SEI messages are removed from the head of queue,
+   *     either due to exceeding the {@linkplain #setMaxSize(int) max queue size} during a call to
+   *     {@link #add(long, ParsableByteArray)}, or due to {@link #flush()}.
+   */
+  public ReorderingSeiMessageQueue(SeiConsumer seiConsumer) {
+    this.seiConsumer = seiConsumer;
+    unusedSeiMessages = new ArrayDeque<>();
+    pendingSeiMessages = new PriorityQueue<>();
+    reorderingQueueSize = C.LENGTH_UNSET;
+  }
+
+  /**
+   * Sets the max size of the re-ordering queue.
+   *
+   * <p>When the queue exceeds this size during a call to {@link #add(long, ParsableByteArray)}, the
+   * least message is passed to the {@link SeiConsumer} provided during construction.
+   *
+   * <p>If the new size is larger than the number of elements currently in the queue, items are
+   * removed from the head of the queue (least first) and passed to the {@link SeiConsumer} provided
+   * during construction.
+   */
+  public void setMaxSize(int reorderingQueueSize) {
+    checkState(reorderingQueueSize >= 0);
+    this.reorderingQueueSize = reorderingQueueSize;
+    flushQueueDownToSize(reorderingQueueSize);
+  }
+
+  /**
+   * Returns the maximum size of this queue, or {@link C#LENGTH_UNSET} if it is unbounded.
+   *
+   * <p>See {@link #setMaxSize(int)}.
+   */
+  public int getMaxSize() {
+    return reorderingQueueSize;
+  }
+
+  /**
+   * Adds a message to the queue.
+   *
+   * <p>If this causes the queue to exceed its {@linkplain #setMaxSize(int) max size}, the least
+   * message (which may be the one passed to this method) is passed to the {@link SeiConsumer}
+   * provided during construction.
+   *
+   * @param presentationTimeUs The presentation time of the SEI message.
+   * @param seiBuffer The SEI data. The data will be copied, so the provided object can be re-used.
+   */
+  public void add(long presentationTimeUs, ParsableByteArray seiBuffer) {
+    if (reorderingQueueSize == 0
+        || (reorderingQueueSize != C.LENGTH_UNSET
+            && pendingSeiMessages.size() >= reorderingQueueSize
+            && presentationTimeUs < castNonNull(pendingSeiMessages.peek()).presentationTimeUs)) {
+      seiConsumer.consume(presentationTimeUs, seiBuffer);
+      return;
+    }
+    SeiMessage seiMessage =
+        unusedSeiMessages.isEmpty() ? new SeiMessage() : unusedSeiMessages.poll();
+    seiMessage.reset(presentationTimeUs, tieBreakGenerator.getAndIncrement(), seiBuffer);
+    pendingSeiMessages.add(seiMessage);
+    if (reorderingQueueSize != C.LENGTH_UNSET) {
+      flushQueueDownToSize(reorderingQueueSize);
+    }
+  }
+
+  /**
+   * Empties the queue, passing all messages (least first) to the {@link SeiConsumer} provided
+   * during construction.
+   */
+  public void flush() {
+    flushQueueDownToSize(0);
+  }
+
+  private void flushQueueDownToSize(int targetSize) {
+    while (pendingSeiMessages.size() > targetSize) {
+      SeiMessage seiMessage = castNonNull(pendingSeiMessages.poll());
+      seiConsumer.consume(seiMessage.presentationTimeUs, seiMessage.data);
+      unusedSeiMessages.push(seiMessage);
+    }
+  }
+
+  /** Holds data from a SEI sample with its presentation timestamp. */
+  private static final class SeiMessage implements Comparable<SeiMessage> {
+
+    private final ParsableByteArray data;
+
+    private long presentationTimeUs;
+
+    /**
+     * {@link PriorityQueue} breaks ties arbitrarily. This field ensures that insertion order is
+     * preserved when messages have the same {@link #presentationTimeUs}.
+     */
+    private long tieBreak;
+
+    public SeiMessage() {
+      presentationTimeUs = C.TIME_UNSET;
+      data = new ParsableByteArray();
+    }
+
+    public void reset(long presentationTimeUs, long tieBreak, ParsableByteArray nalBuffer) {
+      checkState(presentationTimeUs >= 0);
+      this.presentationTimeUs = presentationTimeUs;
+      this.tieBreak = tieBreak;
+      this.data.reset(nalBuffer.bytesLeft());
+      System.arraycopy(
+          /* src= */ nalBuffer.getData(),
+          /* srcPos= */ nalBuffer.getPosition(),
+          /* dest= */ data.getData(),
+          /* destPos= */ 0,
+          /* length= */ nalBuffer.bytesLeft());
+    }
+
+    @Override
+    public int compareTo(SeiMessage other) {
+      int timeComparison = Long.compare(this.presentationTimeUs, other.presentationTimeUs);
+      return timeComparison != 0 ? timeComparison : Long.compare(this.tieBreak, other.tieBreak);
+    }
+  }
+}

libraries/container/src/test/java/androidx/media3/container/ReorderingSeiMessageQueueTest.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package androidx.media3.container;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+import androidx.annotation.Nullable;
+import androidx.media3.common.util.ParsableByteArray;
+import androidx.media3.common.util.Util;
+import androidx.media3.test.utils.TestUtil;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Objects;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Tests for {@link ReorderingSeiMessageQueue}. */
+@RunWith(AndroidJUnit4.class)
+public final class ReorderingSeiMessageQueueTest {
+
+  @Test
+  public void noMaxSize_queueOnlyEmitsOnExplicitFlushCall() {
+    ArrayList<SeiMessage> emittedMessages = new ArrayList<>();
+    ReorderingSeiMessageQueue reorderingQueue =
+        new ReorderingSeiMessageQueue(
+            (presentationTimeUs, seiBuffer) ->
+                emittedMessages.add(new SeiMessage(presentationTimeUs, seiBuffer)));
+
+    // Deliberately re-use a single ParsableByteArray instance to ensure the implementation is
+    // making copies as required.
+    ParsableByteArray scratchData = new ParsableByteArray();
+    byte[] data1 = TestUtil.buildTestData(5);
+    scratchData.reset(data1);
+    reorderingQueue.add(345, scratchData);
+    byte[] data2 = TestUtil.buildTestData(10);
+    scratchData.reset(data2);
+    reorderingQueue.add(123, scratchData);
+
+    assertThat(emittedMessages).isEmpty();
+
+    reorderingQueue.flush();
+
+    assertThat(emittedMessages)
+        .containsExactly(new SeiMessage(123, data2), new SeiMessage(345, data1))
+        .inOrder();
+  }
+
+  @Test
+  public void setMaxSize_emitsImmediatelyIfQueueIsOversized() {
+    ArrayList<SeiMessage> emittedMessages = new ArrayList<>();
+    ReorderingSeiMessageQueue reorderingQueue =
+        new ReorderingSeiMessageQueue(
+            (presentationTimeUs, seiBuffer) ->
+                emittedMessages.add(new SeiMessage(presentationTimeUs, seiBuffer)));
+    ParsableByteArray scratchData = new ParsableByteArray();
+    byte[] data1 = TestUtil.buildTestData(5);
+    scratchData.reset(data1);
+    reorderingQueue.add(345, scratchData);
+    byte[] data2 = TestUtil.buildTestData(10);
+    scratchData.reset(data2);
+    reorderingQueue.add(123, scratchData);
+
+    assertThat(emittedMessages).isEmpty();
+
+    reorderingQueue.setMaxSize(1);
+
+    assertThat(emittedMessages).containsExactly(new SeiMessage(123, data2));
+  }
+
+  @Test
+  public void withMaxSize_addEmitsWhenQueueIsFull() {
+    ArrayList<SeiMessage> emittedMessages = new ArrayList<>();
+    ReorderingSeiMessageQueue reorderingQueue =
+        new ReorderingSeiMessageQueue(
+            (presentationTimeUs, seiBuffer) ->
+                emittedMessages.add(new SeiMessage(presentationTimeUs, seiBuffer)));
+    reorderingQueue.setMaxSize(1);
+
+    // Deliberately re-use a single ParsableByteArray instance to ensure the implementation is
+    // copying as required.
+    ParsableByteArray scratchData = new ParsableByteArray();
+    byte[] data1 = TestUtil.buildTestData(5);
+    scratchData.reset(data1);
+    reorderingQueue.add(345, scratchData);
+
+    assertThat(emittedMessages).isEmpty();
+
+    byte[] data2 = TestUtil.buildTestData(10);
+    scratchData.reset(data2);
+    reorderingQueue.add(123, scratchData);
+
+    assertThat(emittedMessages).containsExactly(new SeiMessage(123, data2));
+  }
+
+  /**
+   * Tests that if a message smaller than all current queue items is added when the queue is full,
+   * the same {@link ParsableByteArray} instance is passed straight to the output to avoid
+   * unnecessary array copies or allocations.
+   */
+  @Test
+  public void withMaxSize_addEmitsWhenQueueIsFull_skippingQueueReusesPbaInstance() {
+    ReorderingSeiMessageQueue.SeiConsumer mockSeiConsumer =
+        mock(ReorderingSeiMessageQueue.SeiConsumer.class);
+    ReorderingSeiMessageQueue reorderingQueue = new ReorderingSeiMessageQueue(mockSeiConsumer);
+    reorderingQueue.setMaxSize(1);
+
+    ParsableByteArray scratchData = new ParsableByteArray();
+    byte[] data1 = TestUtil.buildTestData(5);
+    scratchData.reset(data1);
+    reorderingQueue.add(345, scratchData);
+
+    verifyNoInteractions(mockSeiConsumer);
+
+    byte[] data2 = TestUtil.buildTestData(10);
+    scratchData.reset(data2);
+    reorderingQueue.add(123, scratchData);
+
+    verify(mockSeiConsumer).consume(eq(123L), same(scratchData));
+  }
+
+  private static final class SeiMessage {
+    public final long presentationTimeUs;
+    public final byte[] data;
+
+    public SeiMessage(long presentationTimeUs, ParsableByteArray seiBuffer) {
+      this(
+          presentationTimeUs,
+          Arrays.copyOfRange(seiBuffer.getData(), seiBuffer.getPosition(), seiBuffer.limit()));
+    }
+
+    public SeiMessage(long presentationTimeUs, byte[] seiBuffer) {
+      this.presentationTimeUs = presentationTimeUs;
+      this.data = seiBuffer;
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(presentationTimeUs, Arrays.hashCode(data));
+    }
+
+    @Override
+    public boolean equals(@Nullable Object obj) {
+      if (!(obj instanceof SeiMessage)) {
+        return false;
+      }
+      SeiMessage that = (SeiMessage) obj;
+      return this.presentationTimeUs == that.presentationTimeUs
+          && Arrays.equals(this.data, that.data);
+    }
+
+    @Override
+    public String toString() {
+      return "SeiMessage { ts=" + presentationTimeUs + ",data=0x" + Util.toHexString(data) + " }";
+    }
+  }
+}