Preacquire DRM sessions from the loading side of SampleQueue

Issue: #4133
PiperOrigin-RevId: 362478801

RELEASENOTES.md

@@ -47,6 +47,8 @@
         video tracks (previously separate acquire and release events were
         dispatched for each track in each period).
     *   Include the session state in DRM session-acquired listener methods.
+    *   Prepare DRM sessions (and fetch keys) ahead of the playback position
+        ([#4133](https://github.com/google/ExoPlayer/issues/4133)).
 *   Text
     *   Parse SSA/ASS bold & italic info in `Style:` lines
         ([#8435](https://github.com/google/ExoPlayer/issues/8435)).
@@ -65,8 +67,8 @@
         media item and so that it is not triggered after a timeline change.
     *   Trigger `onMediaItemTransition` event for all reasons except
         `MEDIA_ITEM_TRANSITION_REASON_REPEAT`.
-*   Allow the use of platform extractors through [MediaParser]
-    (https://developer.android.com/reference/android/media/MediaParser).
+*   Allow the use of platform extractors through
+    [MediaParser](https://developer.android.com/reference/android/media/MediaParser).
     Only supported on API 30+.
     *   You can use it for progressive media by passing a
         `MediaParserExtractorAdapter.FACTORY` when creating the

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

@@ -34,6 +34,7 @@ import com.google.android.exoplayer2.drm.DrmInitData;
 import com.google.android.exoplayer2.drm.DrmSession;
 import com.google.android.exoplayer2.drm.DrmSessionEventListener;
 import com.google.android.exoplayer2.drm.DrmSessionManager;
+import com.google.android.exoplayer2.drm.DrmSessionManager.DrmSessionReference;
 import com.google.android.exoplayer2.extractor.TrackOutput;
 import com.google.android.exoplayer2.upstream.Allocator;
 import com.google.android.exoplayer2.upstream.DataReader;
@@ -63,7 +64,7 @@ public class SampleQueue implements TrackOutput {
 
   private final SampleDataQueue sampleDataQueue;
   private final SampleExtrasHolder extrasHolder;
-  private final SpannedData<Format> formatSpans;
+  private final SpannedData<SharedSampleMetadata> sharedSampleMetadata;
   @Nullable private final DrmSessionManager drmSessionManager;
   @Nullable private final DrmSessionEventListener.EventDispatcher drmEventDispatcher;
   @Nullable private final Looper playbackLooper;
@@ -156,7 +157,8 @@ public class SampleQueue implements TrackOutput {
     flags = new int[capacity];
     sizes = new int[capacity];
     cryptoDatas = new CryptoData[capacity];
-    formatSpans = new SpannedData<>();
+    sharedSampleMetadata =
+        new SpannedData<>(/* removeCallback= */ metadata -> metadata.drmSessionReference.release());
     startTimeUs = Long.MIN_VALUE;
     largestDiscardedTimestampUs = Long.MIN_VALUE;
     largestQueuedTimestampUs = Long.MIN_VALUE;
@@ -198,7 +200,7 @@ public class SampleQueue implements TrackOutput {
     largestDiscardedTimestampUs = Long.MIN_VALUE;
     largestQueuedTimestampUs = Long.MIN_VALUE;
     isLastSampleQueued = false;
-    formatSpans.clear();
+    sharedSampleMetadata.clear();
     if (resetUpstreamFormat) {
       unadjustedUpstreamFormat = null;
       upstreamFormat = null;
@@ -371,7 +373,7 @@ public class SampleQueue implements TrackOutput {
           || isLastSampleQueued
           || (upstreamFormat != null && upstreamFormat != downstreamFormat);
     }
-    if (formatSpans.get(getReadIndex()) != downstreamFormat) {
+    if (sharedSampleMetadata.get(getReadIndex()).format != downstreamFormat) {
       // A format can be read.
       return true;
     }
@@ -690,7 +692,7 @@ public class SampleQueue implements TrackOutput {
       }
     }
 
-    Format format = formatSpans.get(getReadIndex());
+    Format format = sharedSampleMetadata.get(getReadIndex()).format;
     if (formatRequired || format != downstreamFormat) {
       onFormatResult(format, formatHolder);
       return C.RESULT_FORMAT_READ;
@@ -723,7 +725,10 @@ public class SampleQueue implements TrackOutput {
       return false;
     }
 
-    @Nullable Format upstreamCommittedFormat = formatSpans.getEndValue();
+    @Nullable SharedSampleMetadata upstreamCommittedMetadata = sharedSampleMetadata.getEndValue();
+    @Nullable
+    Format upstreamCommittedFormat =
+        upstreamCommittedMetadata != null ? upstreamCommittedMetadata.format : null;
     if (Util.areEqual(format, upstreamCommittedFormat)) {
       // The format has changed back to the format of the last committed sample. If they are
       // different objects, we revert back to using upstreamCommittedFormat as the upstreamFormat
@@ -799,8 +804,18 @@ public class SampleQueue implements TrackOutput {
     cryptoDatas[relativeEndIndex] = cryptoData;
     sourceIds[relativeEndIndex] = upstreamSourceId;
 
-    if (!Util.areEqual(upstreamFormat, formatSpans.getEndValue())) {
-      formatSpans.appendSpan(getWriteIndex(), checkNotNull(upstreamFormat));
+    @Nullable SharedSampleMetadata upstreamCommittedMetadata = sharedSampleMetadata.getEndValue();
+    if (upstreamCommittedMetadata == null
+        || !upstreamCommittedMetadata.format.equals(upstreamFormat)) {
+      DrmSessionReference drmSessionReference =
+          drmSessionManager != null
+              ? drmSessionManager.preacquireSession(
+                  checkNotNull(playbackLooper), drmEventDispatcher, upstreamFormat)
+              : DrmSessionReference.EMPTY;
+
+      sharedSampleMetadata.appendSpan(
+          getWriteIndex(),
+          new SharedSampleMetadata(checkNotNull(upstreamFormat), drmSessionReference));
     }
 
     length++;
@@ -863,7 +878,7 @@ public class SampleQueue implements TrackOutput {
     length -= discardCount;
     largestQueuedTimestampUs = max(largestDiscardedTimestampUs, getLargestTimestamp(length));
     isLastSampleQueued = discardCount == 0 && isLastSampleQueued;
-    formatSpans.discardFrom(discardFromIndex);
+    sharedSampleMetadata.discardFrom(discardFromIndex);
     if (length != 0) {
       int relativeLastWriteIndex = getRelativeIndex(length - 1);
       return offsets[relativeLastWriteIndex] + sizes[relativeLastWriteIndex];
@@ -1003,7 +1018,7 @@ public class SampleQueue implements TrackOutput {
     if (readPosition < 0) {
       readPosition = 0;
     }
-    formatSpans.discardTo(absoluteFirstIndex);
+    sharedSampleMetadata.discardTo(absoluteFirstIndex);
 
     if (length == 0) {
       int relativeLastDiscardIndex = (relativeFirstIndex == 0 ? capacity : relativeFirstIndex) - 1;
@@ -1057,4 +1072,15 @@ public class SampleQueue implements TrackOutput {
     public long offset;
     @Nullable public CryptoData cryptoData;
   }
+
+  /** A holder for metadata that applies to a span of contiguous samples. */
+  private static final class SharedSampleMetadata {
+    public final Format format;
+    public final DrmSessionReference drmSessionReference;
+
+    private SharedSampleMetadata(Format format, DrmSessionReference drmSessionReference) {
+      this.format = format;
+      this.drmSessionReference = drmSessionReference;
+    }
+  }
 }

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

@@ -22,6 +22,7 @@ import static java.lang.Math.min;
 import android.util.SparseArray;
 import androidx.annotation.Nullable;
 import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.util.Consumer;
 
 /**
  * Stores value objects associated with spans of integer keys.
@@ -39,10 +40,21 @@ import com.google.android.exoplayer2.C;
   private int memoizedReadIndex;
 
   private final SparseArray<V> spans;
+  private final Consumer<V> removeCallback;
 
   /** Constructs an empty instance. */
   public SpannedData() {
+    this(/* removeCallback= */ value -> {});
+  }
+
+  /**
+   * Constructs an empty instance that invokes {@code removeCallback} on each value that is removed
+   * from the collection.
+   */
+  public SpannedData(Consumer<V> removeCallback) {
     spans = new SparseArray<>();
+    this.removeCallback = removeCallback;
+    memoizedReadIndex = C.INDEX_UNSET;
   }
 
   /**
@@ -71,7 +83,8 @@ import com.google.android.exoplayer2.C;
    * Adds a new span to the end starting at {@code startKey} and containing {@code value}.
    *
    * <p>{@code startKey} must be greater than or equal to the start key of the previous span. If
-   * they're equal, the previous span is overwritten.
+   * they're equal, the previous span is overwritten and it's passed to {@code removeCallback} (if
+   * set).
    */
   public void appendSpan(int startKey, V value) {
     if (memoizedReadIndex == C.INDEX_UNSET) {
@@ -79,7 +92,13 @@ import com.google.android.exoplayer2.C;
       memoizedReadIndex = 0;
     }
 
-    checkArgument(spans.size() == 0 || startKey >= spans.keyAt(spans.size() - 1));
+    if (spans.size() > 0) {
+      int lastStartKey = spans.keyAt(spans.size() - 1);
+      checkArgument(startKey >= lastStartKey);
+      if (lastStartKey == startKey) {
+        removeCallback.accept(spans.valueAt(spans.size() - 1));
+      }
+    }
     spans.append(startKey, value);
   }
 
@@ -102,6 +121,7 @@ import com.google.android.exoplayer2.C;
    */
   public void discardTo(int discardToKey) {
     for (int i = 0; i < spans.size() - 1 && discardToKey >= spans.keyAt(i + 1); i++) {
+      removeCallback.accept(spans.valueAt(i));
       spans.removeAt(i);
       if (memoizedReadIndex > 0) {
         memoizedReadIndex--;
@@ -116,6 +136,7 @@ import com.google.android.exoplayer2.C;
    */
   public void discardFrom(int discardFromKey) {
     for (int i = spans.size() - 1; i >= 0 && discardFromKey < spans.keyAt(i); i--) {
+      removeCallback.accept(spans.valueAt(i));
       spans.removeAt(i);
     }
     memoizedReadIndex = spans.size() > 0 ? min(memoizedReadIndex, spans.size() - 1) : C.INDEX_UNSET;
@@ -123,6 +144,9 @@ import com.google.android.exoplayer2.C;
 
   /** Remove all spans. */
   public void clear() {
+    for (int i = 0; i < spans.size(); i++) {
+      removeCallback.accept(spans.valueAt(i));
+    }
     memoizedReadIndex = C.INDEX_UNSET;
     spans.clear();
   }

library/core/src/test/java/com/google/android/exoplayer2/analytics/AnalyticsCollectorTest.java

@@ -86,6 +86,7 @@ import com.google.android.exoplayer2.drm.ExoMediaDrm;
 import com.google.android.exoplayer2.drm.MediaDrmCallback;
 import com.google.android.exoplayer2.drm.MediaDrmCallbackException;
 import com.google.android.exoplayer2.metadata.Metadata;
+import com.google.android.exoplayer2.robolectric.RobolectricUtil;
 import com.google.android.exoplayer2.robolectric.TestPlayerRunHelper;
 import com.google.android.exoplayer2.source.ConcatenatingMediaSource;
 import com.google.android.exoplayer2.source.LoadEventInfo;
@@ -108,6 +109,7 @@ import com.google.android.exoplayer2.testutil.TestExoPlayerBuilder;
 import com.google.android.exoplayer2.testutil.TestUtil;
 import com.google.android.exoplayer2.trackselection.TrackSelectionArray;
 import com.google.android.exoplayer2.util.Clock;
+import com.google.android.exoplayer2.util.ConditionVariable;
 import com.google.android.exoplayer2.util.MimeTypes;
 import com.google.android.exoplayer2.util.Util;
 import com.google.common.collect.ImmutableList;
@@ -1436,19 +1438,43 @@ public final class AnalyticsCollectorTest {
   }
 
   @Test
-  public void drmEvents_periodWithSameDrmData_keysReused() throws Exception {
+  public void drmEvents_periodsWithSameDrmData_keysReusedButLoadEventReportedTwice()
+      throws Exception {
+    BlockingDrmCallback mediaDrmCallback = BlockingDrmCallback.returnsEmpty();
+    DrmSessionManager blockingDrmSessionManager =
+        new DefaultDrmSessionManager.Builder()
+            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
+            .setMultiSession(true)
+            .build(mediaDrmCallback);
     MediaSource mediaSource =
         new ConcatenatingMediaSource(
-            new FakeMediaSource(SINGLE_PERIOD_TIMELINE, drmSessionManager, VIDEO_FORMAT_DRM_1),
-            new FakeMediaSource(SINGLE_PERIOD_TIMELINE, drmSessionManager, VIDEO_FORMAT_DRM_1));
-    TestAnalyticsListener listener = runAnalyticsTest(mediaSource);
+            new FakeMediaSource(
+                SINGLE_PERIOD_TIMELINE, blockingDrmSessionManager, VIDEO_FORMAT_DRM_1),
+            new FakeMediaSource(
+                SINGLE_PERIOD_TIMELINE, blockingDrmSessionManager, VIDEO_FORMAT_DRM_1));
+    TestAnalyticsListener listener =
+        runAnalyticsTest(
+            mediaSource,
+            // Wait for the media to be fully buffered before unblocking the DRM key request. This
+            // ensures both periods report the same load event (because period1's DRM session is
+            // already preacquired by the time the key load completes).
+            new ActionSchedule.Builder(TAG)
+                .waitForIsLoading(false)
+                .waitForIsLoading(true)
+                .waitForIsLoading(false)
+                .executeRunnable(mediaDrmCallback.keyCondition::open)
+                .build());
 
     populateEventIds(listener.lastReportedTimeline);
     assertThat(listener.getEvents(EVENT_DRM_SESSION_MANAGER_ERROR)).isEmpty();
     assertThat(listener.getEvents(EVENT_DRM_SESSION_ACQUIRED))
         .containsExactly(period0, period1)
         .inOrder();
-    assertThat(listener.getEvents(EVENT_DRM_KEYS_LOADED)).containsExactly(period0);
+    // This includes both period0 and period1 because period1's DrmSession was preacquired before
+    // the key load completed. There's only one key load (a second would block forever). We can't
+    // assume the order these events will arrive in because it depends on the iteration order of a
+    // HashSet of EventDispatchers inside DefaultDrmSession.
+    assertThat(listener.getEvents(EVENT_DRM_KEYS_LOADED)).containsExactly(period0, period1);
     // The period1 release event is lost because it's posted to "ExoPlayerTest thread" after that
     // thread has been quit during clean-up.
     assertThat(listener.getEvents(EVENT_DRM_SESSION_RELEASED)).containsExactly(period0);
@@ -1480,11 +1506,21 @@ public final class AnalyticsCollectorTest {
 
   @Test
   public void drmEvents_errorHandling() throws Exception {
+    BlockingDrmCallback mediaDrmCallback = BlockingDrmCallback.alwaysFailing();
     DrmSessionManager failingDrmSessionManager =
-        new DefaultDrmSessionManager.Builder().build(new FailingDrmCallback());
+        new DefaultDrmSessionManager.Builder()
+            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
+            .setMultiSession(true)
+            .build(mediaDrmCallback);
     MediaSource mediaSource =
         new FakeMediaSource(SINGLE_PERIOD_TIMELINE, failingDrmSessionManager, VIDEO_FORMAT_DRM_1);
-    TestAnalyticsListener listener = runAnalyticsTest(mediaSource);
+    TestAnalyticsListener listener =
+        runAnalyticsTest(
+            mediaSource,
+            new ActionSchedule.Builder(TAG)
+                .waitForIsLoading(false)
+                .executeRunnable(mediaDrmCallback.keyCondition::open)
+                .build());
 
     populateEventIds(listener.lastReportedTimeline);
     assertThat(listener.getEvents(EVENT_DRM_SESSION_MANAGER_ERROR)).containsExactly(period0);
@@ -2341,21 +2377,56 @@ public final class AnalyticsCollectorTest {
   }
 
   /**
-   * A {@link MediaDrmCallback} that throws exceptions for both {@link
-   * #executeProvisionRequest(UUID, ExoMediaDrm.ProvisionRequest)} and {@link
-   * #executeKeyRequest(UUID, ExoMediaDrm.KeyRequest)}.
+   * A {@link MediaDrmCallback} that blocks each provision and key request until the associated
+   * {@link ConditionVariable} field is opened, and then returns an empty byte array. The {@link
+   * ConditionVariable} must be explicitly opened for each request.
    */
-  private static final class FailingDrmCallback implements MediaDrmCallback {
+  private static final class BlockingDrmCallback implements MediaDrmCallback {
+
+    public final ConditionVariable provisionCondition;
+    public final ConditionVariable keyCondition;
+
+    private final boolean alwaysFail;
+
+    private BlockingDrmCallback(boolean alwaysFail) {
+      this.provisionCondition = RobolectricUtil.createRobolectricConditionVariable();
+      this.keyCondition = RobolectricUtil.createRobolectricConditionVariable();
+
+      this.alwaysFail = alwaysFail;
+    }
+
+    /** Returns a callback that always returns an empty byte array from its execute methods. */
+    public static BlockingDrmCallback returnsEmpty() {
+      return new BlockingDrmCallback(/* alwaysFail= */ false);
+    }
+
+    /** Returns a callback that always throws an exception from its execute methods. */
+    public static BlockingDrmCallback alwaysFailing() {
+      return new BlockingDrmCallback(/* alwaysFail= */ true);
+    }
+
     @Override
     public byte[] executeProvisionRequest(UUID uuid, ExoMediaDrm.ProvisionRequest request)
         throws MediaDrmCallbackException {
-      throw new RuntimeException("executeProvision failed");
+      provisionCondition.blockUninterruptible();
+      provisionCondition.close();
+      if (alwaysFail) {
+        throw new RuntimeException("executeProvisionRequest failed");
+      } else {
+        return new byte[0];
+      }
     }
 
     @Override
     public byte[] executeKeyRequest(UUID uuid, ExoMediaDrm.KeyRequest request)
         throws MediaDrmCallbackException {
-      throw new RuntimeException("executeKey failed");
+      keyCondition.blockUninterruptible();
+      keyCondition.close();
+      if (alwaysFail) {
+        throw new RuntimeException("executeKeyRequest failed");
+      } else {
+        return new byte[0];
+      }
     }
   }
 }

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

@@ -16,129 +16,161 @@
 package com.google.android.exoplayer2.source;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.android.exoplayer2.drm.DrmSessionManager.DrmSessionReference;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 
 /** Tests for {@link SpannedData}. */
 @RunWith(AndroidJUnit4.class)
 public final class SpannedDataTest {
 
-  private static final String VALUE_1 = "value 1";
-  private static final String VALUE_2 = "value 2";
-  private static final String VALUE_3 = "value 3";
+  @Rule public final MockitoRule mockito = MockitoJUnit.rule();
+
+  @Mock private DrmSessionReference value1;
+  @Mock private DrmSessionReference value2;
+  @Mock private DrmSessionReference value3;
 
   @Test
   public void appendMultipleSpansThenRead() {
-    SpannedData<String> spannedData = new SpannedData<>();
-
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
-    spannedData.appendSpan(/* startKey= */ 4, VALUE_3);
-
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(1)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(2)).isEqualTo(VALUE_2);
-    assertThat(spannedData.get(3)).isEqualTo(VALUE_2);
-    assertThat(spannedData.get(4)).isEqualTo(VALUE_3);
-    assertThat(spannedData.get(5)).isEqualTo(VALUE_3);
+    SpannedData<DrmSessionReference> spannedData =
+        new SpannedData<>(/* removeCallback= */ DrmSessionReference::release);
+
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
+    spannedData.appendSpan(/* startKey= */ 4, value3);
+
+    assertThat(spannedData.get(0)).isEqualTo(value1);
+    assertThat(spannedData.get(1)).isEqualTo(value1);
+    assertThat(spannedData.get(2)).isEqualTo(value2);
+    assertThat(spannedData.get(3)).isEqualTo(value2);
+    assertThat(spannedData.get(4)).isEqualTo(value3);
+    assertThat(spannedData.get(5)).isEqualTo(value3);
+
+    verify(value1, never()).release();
+    verify(value2, never()).release();
+    verify(value3, never()).release();
   }
 
   @Test
   public void append_emptySpansDiscarded() {
-    SpannedData<String> spannedData = new SpannedData<>();
+    SpannedData<DrmSessionReference> spannedData = new SpannedData<>();
 
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_3);
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
+    spannedData.appendSpan(/* startKey= */ 2, value3);
 
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(1)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(2)).isEqualTo(VALUE_3);
-    assertThat(spannedData.get(3)).isEqualTo(VALUE_3);
+    assertThat(spannedData.get(0)).isEqualTo(value1);
+    assertThat(spannedData.get(1)).isEqualTo(value1);
+    assertThat(spannedData.get(2)).isEqualTo(value3);
+    assertThat(spannedData.get(3)).isEqualTo(value3);
   }
 
   @Test
   public void discardTo() {
-    SpannedData<String> spannedData = new SpannedData<>();
+    SpannedData<DrmSessionReference> spannedData =
+        new SpannedData<>(/* removeCallback= */ DrmSessionReference::release);
 
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
-    spannedData.appendSpan(/* startKey= */ 4, VALUE_3);
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
+    spannedData.appendSpan(/* startKey= */ 4, value3);
 
     spannedData.discardTo(2);
 
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_2);
-    assertThat(spannedData.get(2)).isEqualTo(VALUE_2);
+    verify(value1).release();
+    verify(value2, never()).release();
+    assertThat(spannedData.get(0)).isEqualTo(value2);
+    assertThat(spannedData.get(2)).isEqualTo(value2);
 
     spannedData.discardTo(4);
 
-    assertThat(spannedData.get(3)).isEqualTo(VALUE_3);
-    assertThat(spannedData.get(4)).isEqualTo(VALUE_3);
+    verify(value2).release();
+    verify(value3, never()).release();
+    assertThat(spannedData.get(3)).isEqualTo(value3);
+    assertThat(spannedData.get(4)).isEqualTo(value3);
   }
 
   @Test
   public void discardTo_prunesEmptySpans() {
-    SpannedData<String> spannedData = new SpannedData<>();
+    SpannedData<DrmSessionReference> spannedData = new SpannedData<>();
 
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_3);
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
+    spannedData.appendSpan(/* startKey= */ 2, value3);
 
     spannedData.discardTo(2);
 
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_3);
-    assertThat(spannedData.get(2)).isEqualTo(VALUE_3);
+    assertThat(spannedData.get(0)).isEqualTo(value3);
+    assertThat(spannedData.get(2)).isEqualTo(value3);
   }
 
   @Test
   public void discardFromThenAppend_keepsValueIfSpanEndsUpNonEmpty() {
-    SpannedData<String> spannedData = new SpannedData<>();
+    SpannedData<DrmSessionReference> spannedData =
+        new SpannedData<>(/* removeCallback= */ DrmSessionReference::release);
 
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
-    spannedData.appendSpan(/* startKey= */ 4, VALUE_3);
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
+    spannedData.appendSpan(/* startKey= */ 4, value3);
 
     spannedData.discardFrom(2);
-    assertThat(spannedData.getEndValue()).isEqualTo(VALUE_2);
 
-    spannedData.appendSpan(/* startKey= */ 3, VALUE_3);
+    verify(value3).release();
+    assertThat(spannedData.getEndValue()).isEqualTo(value2);
+
+    spannedData.appendSpan(/* startKey= */ 3, value3);
 
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(1)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(2)).isEqualTo(VALUE_2);
-    assertThat(spannedData.get(3)).isEqualTo(VALUE_3);
+    verify(value1, never()).release();
+    verify(value2, never()).release();
+    assertThat(spannedData.get(0)).isEqualTo(value1);
+    assertThat(spannedData.get(1)).isEqualTo(value1);
+    assertThat(spannedData.get(2)).isEqualTo(value2);
+    assertThat(spannedData.get(3)).isEqualTo(value3);
   }
 
   @Test
   public void discardFromThenAppend_prunesEmptySpan() {
-    SpannedData<String> spannedData = new SpannedData<>();
+    SpannedData<DrmSessionReference> spannedData =
+        new SpannedData<>(/* removeCallback= */ DrmSessionReference::release);
 
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
 
     spannedData.discardFrom(2);
 
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_3);
+    verify(value2, never()).release();
 
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(1)).isEqualTo(VALUE_1);
-    assertThat(spannedData.get(2)).isEqualTo(VALUE_3);
+    spannedData.appendSpan(/* startKey= */ 2, value3);
+
+    verify(value2).release();
+    assertThat(spannedData.get(0)).isEqualTo(value1);
+    assertThat(spannedData.get(1)).isEqualTo(value1);
+    assertThat(spannedData.get(2)).isEqualTo(value3);
   }
 
   @Test
   public void clear() {
-    SpannedData<String> spannedData = new SpannedData<>();
+    SpannedData<DrmSessionReference> spannedData =
+        new SpannedData<>(/* removeCallback= */ DrmSessionReference::release);
 
-    spannedData.appendSpan(/* startKey= */ 0, VALUE_1);
-    spannedData.appendSpan(/* startKey= */ 2, VALUE_2);
+    spannedData.appendSpan(/* startKey= */ 0, value1);
+    spannedData.appendSpan(/* startKey= */ 2, value2);
 
     spannedData.clear();
 
-    spannedData.appendSpan(/* startKey= */ 1, VALUE_3);
+    verify(value1).release();
+    verify(value2).release();
+
+    spannedData.appendSpan(/* startKey= */ 1, value3);
 
-    assertThat(spannedData.get(0)).isEqualTo(VALUE_3);
-    assertThat(spannedData.get(1)).isEqualTo(VALUE_3);
+    assertThat(spannedData.get(0)).isEqualTo(value3);
+    assertThat(spannedData.get(1)).isEqualTo(value3);
   }
 }