Store adPlaybackStates and shared periods by period UID

To support multi-period content we need to store AdPlaybackStates and SharedMediaPeriod by the periodUid as a key. While after this no-op CL, we still only support single-period content, storing these resources by periodUid is the ground work for multi-period support being added in an follow-up CL.

PiperOrigin-RevId: 416836445

libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/ads/ServerSideAdInsertionMediaSource.java

@@ -17,6 +17,7 @@ package androidx.media3.exoplayer.source.ads;
 
 import static androidx.media3.common.util.Assertions.checkArgument;
 import static androidx.media3.common.util.Assertions.checkNotNull;
+import static androidx.media3.common.util.Assertions.checkState;
 import static androidx.media3.common.util.Util.castNonNull;
 import static androidx.media3.exoplayer.source.ads.ServerSideAdInsertionUtil.getAdCountInGroup;
 import static androidx.media3.exoplayer.source.ads.ServerSideAdInsertionUtil.getMediaPeriodPositionUs;
@@ -36,7 +37,6 @@ import androidx.media3.common.StreamKey;
 import androidx.media3.common.Timeline;
 import androidx.media3.common.TrackGroup;
 import androidx.media3.common.TrackGroupArray;
-import androidx.media3.common.util.Assertions;
 import androidx.media3.common.util.UnstableApi;
 import androidx.media3.common.util.Util;
 import androidx.media3.datasource.TransferListener;
@@ -57,6 +57,7 @@ import androidx.media3.exoplayer.source.SampleStream;
 import androidx.media3.exoplayer.trackselection.ExoTrackSelection;
 import androidx.media3.exoplayer.upstream.Allocator;
 import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.ListMultimap;
 import java.io.IOException;
@@ -75,8 +76,8 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
  * server-side inserted ad breaks and ensures that playback continues seamlessly with the wrapped
  * media across all transitions.
  *
- * <p>The ad breaks need to be specified using {@link #setAdPlaybackState} and can be updated during
+ * <p>The ad breaks need to be specified using {@link #setAdPlaybackStates} and can be updated
- * playback.
+ * during playback.
  */
 @UnstableApi
 public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
@@ -91,7 +92,7 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
      * Called when the content source has refreshed the timeline.
      *
      * <p>If true is returned the source refresh publication is deferred, to wait for an {@link
-     * #setAdPlaybackState(AdPlaybackState) ad playback state update}. If false is returned, the
+     * #setAdPlaybackStates(ImmutableMap)} ad playback state update}. If false is returned, the
      * source refresh is immediately published.
      *
      * <p>Called on the playback thread.
@@ -104,7 +105,7 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
   }
 
   private final MediaSource mediaSource;
-  private final ListMultimap<Long, SharedMediaPeriod> mediaPeriods;
+  private final ListMultimap<Pair<Long, Object>, SharedMediaPeriod> mediaPeriods;
   private final MediaSourceEventListener.EventDispatcher mediaSourceEventDispatcherWithoutId;
   private final DrmSessionEventListener.EventDispatcher drmEventDispatcherWithoutId;
   @Nullable private final AdPlaybackStateUpdater adPlaybackStateUpdater;
@@ -115,7 +116,7 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
 
   @Nullable private SharedMediaPeriod lastUsedMediaPeriod;
   @Nullable private Timeline contentTimeline;
-  private AdPlaybackState adPlaybackState;
+  private ImmutableMap<Object, AdPlaybackState> adPlaybackStates;
 
   /**
    * Creates the media source.
@@ -131,53 +132,74 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     this.mediaSource = mediaSource;
     this.adPlaybackStateUpdater = adPlaybackStateUpdater;
     mediaPeriods = ArrayListMultimap.create();
-    adPlaybackState = AdPlaybackState.NONE;
+    adPlaybackStates = ImmutableMap.of();
     mediaSourceEventDispatcherWithoutId = createEventDispatcher(/* mediaPeriodId= */ null);
     drmEventDispatcherWithoutId = createDrmEventDispatcher(/* mediaPeriodId= */ null);
   }
 
   /**
-   * Sets the {@link AdPlaybackState} published by this source.
+   * Sets the map of {@link AdPlaybackState ad playback states} published by this source. The key is
+   * the period UID of a period in the {@link
+   * AdPlaybackStateUpdater#onAdPlaybackStateUpdateRequested(Timeline)} content timeline}.
+   *
+   * <p>Each period has an {@link AdPlaybackState} that tells where in the period the ad groups
+   * start and end. Must only contain server-side inserted ad groups. The number of ad groups and
+   * the number of ads within an ad group may only increase. The durations of ads may change and the
+   * positions of future ad groups may change. Post-roll ad groups with {@link C#TIME_END_OF_SOURCE}
+   * must be empty and can be used as a placeholder for a future ad group.
    *
    * <p>May be called from any thread.
    *
-   * <p>Must only contain server-side inserted ad groups. The number of ad groups and the number of
+   * @param adPlaybackStates The map of {@link AdPlaybackState} keyed by their period UID.
-   * ads within an ad group may only increase. The durations of ads may change and the positions of
-   * future ad groups may change. Post-roll ad groups with {@link C#TIME_END_OF_SOURCE} must be
-   * empty and can be used as a placeholder for a future ad group.
-   *
-   * @param adPlaybackState The new {@link AdPlaybackState}.
    */
-  public void setAdPlaybackState(AdPlaybackState adPlaybackState) {
+  public void setAdPlaybackStates(ImmutableMap<Object, AdPlaybackState> adPlaybackStates) {
-    checkArgument(adPlaybackState.adGroupCount >= this.adPlaybackState.adGroupCount);
+    checkArgument(!adPlaybackStates.isEmpty());
+    Object adsId = checkNotNull(adPlaybackStates.values().asList().get(0).adsId);
+    for (Map.Entry<Object, AdPlaybackState> entry : adPlaybackStates.entrySet()) {
+      Object periodUid = entry.getKey();
+      AdPlaybackState adPlaybackState = entry.getValue();
+      checkArgument(Util.areEqual(adsId, adPlaybackState.adsId));
+      @Nullable AdPlaybackState oldAdPlaybackState = this.adPlaybackStates.get(periodUid);
+      if (oldAdPlaybackState != null) {
         for (int i = adPlaybackState.removedAdGroupCount; i < adPlaybackState.adGroupCount; i++) {
           AdPlaybackState.AdGroup adGroup = adPlaybackState.getAdGroup(i);
           checkArgument(adGroup.isServerSideInserted);
-      if (i < this.adPlaybackState.adGroupCount) {
+          if (i < oldAdPlaybackState.adGroupCount) {
             checkArgument(
                 getAdCountInGroup(adPlaybackState, /* adGroupIndex= */ i)
-                >= getAdCountInGroup(this.adPlaybackState, /* adGroupIndex= */ i));
+                    >= getAdCountInGroup(oldAdPlaybackState, /* adGroupIndex= */ i));
           }
           if (adGroup.timeUs == C.TIME_END_OF_SOURCE) {
             checkArgument(getAdCountInGroup(adPlaybackState, /* adGroupIndex= */ i) == 0);
           }
         }
+      }
+    }
     synchronized (this) {
       if (playbackHandler == null) {
-        this.adPlaybackState = adPlaybackState;
+        this.adPlaybackStates = adPlaybackStates;
       } else {
         playbackHandler.post(
             () -> {
               for (SharedMediaPeriod mediaPeriod : mediaPeriods.values()) {
+                @Nullable
+                AdPlaybackState adPlaybackState = adPlaybackStates.get(mediaPeriod.periodUid);
+                if (adPlaybackState != null) {
                   mediaPeriod.updateAdPlaybackState(adPlaybackState);
                 }
+              }
               if (lastUsedMediaPeriod != null) {
+                @Nullable
+                AdPlaybackState adPlaybackState =
+                    adPlaybackStates.get(lastUsedMediaPeriod.periodUid);
+                if (adPlaybackState != null) {
                   lastUsedMediaPeriod.updateAdPlaybackState(adPlaybackState);
                 }
-              this.adPlaybackState = adPlaybackState;
+              }
+              this.adPlaybackStates = adPlaybackStates;
               if (contentTimeline != null) {
                 refreshSourceInfo(
-                    new ServerSideAdInsertionTimeline(contentTimeline, adPlaybackState));
+                    new ServerSideAdInsertionTimeline(contentTimeline, adPlaybackStates));
               }
             });
       }
@@ -221,8 +243,8 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     this.contentTimeline = timeline;
     if ((adPlaybackStateUpdater == null
             || !adPlaybackStateUpdater.onAdPlaybackStateUpdateRequested(timeline))
-        && !AdPlaybackState.NONE.equals(adPlaybackState)) {
+        && !adPlaybackStates.isEmpty()) {
-      refreshSourceInfo(new ServerSideAdInsertionTimeline(timeline, adPlaybackState));
+      refreshSourceInfo(new ServerSideAdInsertionTimeline(timeline, adPlaybackStates));
     }
   }
 
@@ -240,19 +262,26 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
 
   @Override
   public MediaPeriod createPeriod(MediaPeriodId id, Allocator allocator, long startPositionUs) {
-    SharedMediaPeriod sharedPeriod;
+    @Nullable SharedMediaPeriod sharedPeriod = null;
+    Pair<Long, Object> sharedMediaPeriodKey = new Pair<>(id.windowSequenceNumber, id.periodUid);
     if (lastUsedMediaPeriod != null) {
+      if (lastUsedMediaPeriod.periodUid.equals(id.periodUid)) {
         sharedPeriod = lastUsedMediaPeriod;
-      lastUsedMediaPeriod = null;
+        mediaPeriods.put(sharedMediaPeriodKey, sharedPeriod);
-      mediaPeriods.put(id.windowSequenceNumber, sharedPeriod);
       } else {
+        lastUsedMediaPeriod.release(mediaSource);
+      }
+      lastUsedMediaPeriod = null;
+    }
+    if (sharedPeriod == null) {
       @Nullable
       SharedMediaPeriod lastExistingPeriod =
-          Iterables.getLast(mediaPeriods.get(id.windowSequenceNumber), /* defaultValue= */ null);
+          Iterables.getLast(mediaPeriods.get(sharedMediaPeriodKey), /* defaultValue= */ null);
       if (lastExistingPeriod != null
           && lastExistingPeriod.canReuseMediaPeriod(id, startPositionUs)) {
         sharedPeriod = lastExistingPeriod;
       } else {
+        AdPlaybackState adPlaybackState = checkNotNull(adPlaybackStates.get(id.periodUid));
         long streamPositionUs = getStreamPositionUs(startPositionUs, id, adPlaybackState);
         sharedPeriod =
             new SharedMediaPeriod(
@@ -260,8 +289,9 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
                     new MediaPeriodId(id.periodUid, id.windowSequenceNumber),
                     allocator,
                     streamPositionUs),
+                id.periodUid,
                 adPlaybackState);
-        mediaPeriods.put(id.windowSequenceNumber, sharedPeriod);
+        mediaPeriods.put(sharedMediaPeriodKey, sharedPeriod);
       }
     }
     MediaPeriodImpl mediaPeriod =
@@ -277,7 +307,10 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     mediaPeriodImpl.sharedPeriod.remove(mediaPeriodImpl);
     if (mediaPeriodImpl.sharedPeriod.isUnused()) {
       mediaPeriods.remove(
-          mediaPeriodImpl.mediaPeriodId.windowSequenceNumber, mediaPeriodImpl.sharedPeriod);
+          new Pair<>(
+              mediaPeriodImpl.mediaPeriodId.windowSequenceNumber,
+              mediaPeriodImpl.mediaPeriodId.periodUid),
+          mediaPeriodImpl.sharedPeriod);
       if (mediaPeriods.isEmpty()) {
         // Keep until disabled.
         lastUsedMediaPeriod = mediaPeriodImpl.sharedPeriod;
@@ -381,7 +414,11 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     } else {
       mediaPeriod.sharedPeriod.onLoadStarted(loadEventInfo, mediaLoadData);
       mediaPeriod.mediaSourceEventDispatcher.loadStarted(
-          loadEventInfo, correctMediaLoadData(mediaPeriod, mediaLoadData, adPlaybackState));
+          loadEventInfo,
+          correctMediaLoadData(
+              mediaPeriod,
+              mediaLoadData,
+              checkNotNull(adPlaybackStates.get(mediaPeriod.mediaPeriodId.periodUid))));
     }
   }
 
@@ -399,7 +436,11 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     } else {
       mediaPeriod.sharedPeriod.onLoadFinished(loadEventInfo);
       mediaPeriod.mediaSourceEventDispatcher.loadCompleted(
-          loadEventInfo, correctMediaLoadData(mediaPeriod, mediaLoadData, adPlaybackState));
+          loadEventInfo,
+          correctMediaLoadData(
+              mediaPeriod,
+              mediaLoadData,
+              checkNotNull(adPlaybackStates.get(mediaPeriod.mediaPeriodId.periodUid))));
     }
   }
 
@@ -417,7 +458,11 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     } else {
       mediaPeriod.sharedPeriod.onLoadFinished(loadEventInfo);
       mediaPeriod.mediaSourceEventDispatcher.loadCanceled(
-          loadEventInfo, correctMediaLoadData(mediaPeriod, mediaLoadData, adPlaybackState));
+          loadEventInfo,
+          correctMediaLoadData(
+              mediaPeriod,
+              mediaLoadData,
+              checkNotNull(adPlaybackStates.get(mediaPeriod.mediaPeriodId.periodUid))));
     }
   }
 
@@ -441,7 +486,10 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
       }
       mediaPeriod.mediaSourceEventDispatcher.loadError(
           loadEventInfo,
-          correctMediaLoadData(mediaPeriod, mediaLoadData, adPlaybackState),
+          correctMediaLoadData(
+              mediaPeriod,
+              mediaLoadData,
+              checkNotNull(adPlaybackStates.get(mediaPeriod.mediaPeriodId.periodUid))),
           error,
           wasCanceled);
     }
@@ -457,7 +505,10 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
       mediaSourceEventDispatcherWithoutId.upstreamDiscarded(mediaLoadData);
     } else {
       mediaPeriod.mediaSourceEventDispatcher.upstreamDiscarded(
-          correctMediaLoadData(mediaPeriod, mediaLoadData, adPlaybackState));
+          correctMediaLoadData(
+              mediaPeriod,
+              mediaLoadData,
+              checkNotNull(adPlaybackStates.get(mediaPeriod.mediaPeriodId.periodUid))));
     }
   }
 
@@ -472,7 +523,10 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     } else {
       mediaPeriod.sharedPeriod.onDownstreamFormatChanged(mediaPeriod, mediaLoadData);
       mediaPeriod.mediaSourceEventDispatcher.downstreamFormatChanged(
-          correctMediaLoadData(mediaPeriod, mediaLoadData, adPlaybackState));
+          correctMediaLoadData(
+              mediaPeriod,
+              mediaLoadData,
+              checkNotNull(adPlaybackStates.get(mediaPeriod.mediaPeriodId.periodUid))));
     }
   }
 
@@ -491,7 +545,8 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     if (mediaPeriodId == null) {
       return null;
     }
-    List<SharedMediaPeriod> periods = mediaPeriods.get(mediaPeriodId.windowSequenceNumber);
+    List<SharedMediaPeriod> periods =
+        mediaPeriods.get(new Pair<>(mediaPeriodId.windowSequenceNumber, mediaPeriodId.periodUid));
     if (periods.isEmpty()) {
       return null;
     }
@@ -560,6 +615,7 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     private final MediaPeriod actualMediaPeriod;
     private final List<MediaPeriodImpl> mediaPeriods;
     private final Map<Long, Pair<LoadEventInfo, MediaLoadData>> activeLoads;
+    private final Object periodUid;
 
     private AdPlaybackState adPlaybackState;
     @Nullable private MediaPeriodImpl loadingPeriod;
@@ -569,8 +625,10 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
     public @NullableType SampleStream[] sampleStreams;
     public @NullableType MediaLoadData[] lastDownstreamFormatChangeData;
 
-    public SharedMediaPeriod(MediaPeriod actualMediaPeriod, AdPlaybackState adPlaybackState) {
+    public SharedMediaPeriod(
+        MediaPeriod actualMediaPeriod, Object periodUid, AdPlaybackState adPlaybackState) {
       this.actualMediaPeriod = actualMediaPeriod;
+      this.periodUid = periodUid;
       this.adPlaybackState = adPlaybackState;
       mediaPeriods = new ArrayList<>();
       activeLoads = new HashMap<>();
@@ -931,19 +989,27 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
 
   private static final class ServerSideAdInsertionTimeline extends ForwardingTimeline {
 
-    private final AdPlaybackState adPlaybackState;
+    private final ImmutableMap<Object, AdPlaybackState> adPlaybackStates;
 
     public ServerSideAdInsertionTimeline(
-        Timeline contentTimeline, AdPlaybackState adPlaybackState) {
+        Timeline contentTimeline, ImmutableMap<Object, AdPlaybackState> adPlaybackStates) {
       super(contentTimeline);
-      Assertions.checkState(contentTimeline.getPeriodCount() == 1);
+      checkState(contentTimeline.getPeriodCount() == 1);
-      Assertions.checkState(contentTimeline.getWindowCount() == 1);
+      checkState(contentTimeline.getWindowCount() == 1);
-      this.adPlaybackState = adPlaybackState;
+      Period period = new Period();
+      for (int i = 0; i < contentTimeline.getPeriodCount(); i++) {
+        contentTimeline.getPeriod(/* periodIndex= */ i, period, /* setIds= */ true);
+        checkState(adPlaybackStates.containsKey(checkNotNull(period.uid)));
+      }
+      this.adPlaybackStates = adPlaybackStates;
     }
 
     @Override
     public Window getWindow(int windowIndex, Window window, long defaultPositionProjectionUs) {
       super.getWindow(windowIndex, window, defaultPositionProjectionUs);
+      Object firstPeriodUid =
+          checkNotNull(getPeriod(/* periodIndex= */ 0, new Period(), /* setIds= */ true).uid);
+      AdPlaybackState adPlaybackState = checkNotNull(adPlaybackStates.get(firstPeriodUid));
       long positionInPeriodUs =
           getMediaPeriodPositionUsForContent(
               window.positionInFirstPeriodUs,
@@ -968,7 +1034,8 @@ public final class ServerSideAdInsertionMediaSource extends BaseMediaSource
 
     @Override
     public Period getPeriod(int periodIndex, Period period, boolean setIds) {
-      super.getPeriod(periodIndex, period, setIds);
+      super.getPeriod(periodIndex, period, /* setIds= */ true);
+      AdPlaybackState adPlaybackState = checkNotNull(adPlaybackStates.get(period.uid));
       long durationUs = period.durationUs;
       if (durationUs == C.TIME_UNSET) {
         durationUs = adPlaybackState.contentDurationUs;

libraries/exoplayer/src/test/java/androidx/media3/exoplayer/source/ads/ServerSideAdInsertionMediaSourceTest.java

@@ -15,6 +15,7 @@
  */
 package androidx.media3.exoplayer.source.ads;
 
+import static androidx.media3.common.util.Assertions.checkNotNull;
 import static androidx.media3.exoplayer.source.ads.ServerSideAdInsertionUtil.addAdGroupToAdPlaybackState;
 import static androidx.media3.test.utils.robolectric.RobolectricUtil.runMainLooperUntil;
 import static androidx.media3.test.utils.robolectric.TestPlayerRunHelper.playUntilPosition;
@@ -32,6 +33,7 @@ import static org.mockito.Mockito.verify;
 
 import android.content.Context;
 import android.graphics.SurfaceTexture;
+import android.util.Pair;
 import android.view.Surface;
 import androidx.media3.common.AdPlaybackState;
 import androidx.media3.common.MediaItem;
@@ -50,7 +52,9 @@ import androidx.media3.test.utils.robolectric.PlaybackOutput;
 import androidx.media3.test.utils.robolectric.ShadowMediaCodecConfig;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.common.collect.ImmutableMap;
 import java.util.concurrent.atomic.AtomicReference;
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -104,8 +108,8 @@ public final class ServerSideAdInsertionMediaSourceTest {
             .withContentResumeOffsetUs(/* adGroupIndex= */ 1, /* contentResumeOffsetUs= */ 400_000)
             .withContentResumeOffsetUs(/* adGroupIndex= */ 2, /* contentResumeOffsetUs= */ 200_000);
     AtomicReference<Timeline> timelineReference = new AtomicReference<>();
+    mediaSource.setAdPlaybackStates(ImmutableMap.of(new Pair<>(0, 0), adPlaybackState));
 
-    mediaSource.setAdPlaybackState(adPlaybackState);
     mediaSource.prepareSource(
         (source, timeline) -> timelineReference.set(timeline),
         /* mediaTransferListener= */ null,
@@ -143,6 +147,26 @@ public final class ServerSideAdInsertionMediaSourceTest {
     assertThat(window.durationUs).isEqualTo(9_800_000);
   }
 
+  @Test
+  public void timeline_missingAdPlaybackStateByPeriodUid_isAssertedAndThrows() {
+    ServerSideAdInsertionMediaSource mediaSource =
+        new ServerSideAdInsertionMediaSource(
+            new FakeMediaSource(), /* adPlaybackStateUpdater= */ null);
+    // The map of adPlaybackStates does not contain a valid period UID as key.
+    mediaSource.setAdPlaybackStates(
+        ImmutableMap.of(new Object(), new AdPlaybackState(/* adsId= */ new Object())));
+
+    Assert.assertThrows(
+        IllegalStateException.class,
+        () ->
+            mediaSource.prepareSource(
+                (source, timeline) -> {
+                  /* Do nothing. */
+                },
+                /* mediaTransferListener= */ null,
+                PlayerId.UNSET));
+  }
+
   @Test
   public void playbackWithPredefinedAds_playsSuccessfulWithoutRendererResets() throws Exception {
     Context context = ApplicationProvider.getApplicationContext();
@@ -154,10 +178,6 @@ public final class ServerSideAdInsertionMediaSourceTest {
     player.setVideoSurface(new Surface(new SurfaceTexture(/* texName= */ 1)));
     PlaybackOutput playbackOutput = PlaybackOutput.register(player, renderersFactory);
 
-    ServerSideAdInsertionMediaSource mediaSource =
-        new ServerSideAdInsertionMediaSource(
-            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
-            /* adPlaybackStateUpdater= */ null);
     AdPlaybackState adPlaybackState = new AdPlaybackState(/* adsId= */ new Object());
     adPlaybackState =
         addAdGroupToAdPlaybackState(
@@ -171,17 +191,32 @@ public final class ServerSideAdInsertionMediaSourceTest {
             /* fromPositionUs= */ 400_000,
             /* toPositionUs= */ 700_000,
             /* contentResumeOffsetUs= */ 1_000_000);
-    adPlaybackState =
+    AdPlaybackState firstAdPlaybackState =
         addAdGroupToAdPlaybackState(
             adPlaybackState,
             /* fromPositionUs= */ 900_000,
             /* toPositionUs= */ 1_000_000,
             /* contentResumeOffsetUs= */ 0);
-    mediaSource.setAdPlaybackState(adPlaybackState);
+
+    AtomicReference<ServerSideAdInsertionMediaSource> mediaSourceRef = new AtomicReference<>();
+    mediaSourceRef.set(
+        new ServerSideAdInsertionMediaSource(
+            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
+            contentTimeline -> {
+              Object periodUid =
+                  checkNotNull(
+                      contentTimeline.getPeriod(
+                              /* periodIndex= */ 0, new Timeline.Period(), /* setIds= */ true)
+                          .uid);
+              mediaSourceRef
+                  .get()
+                  .setAdPlaybackStates(ImmutableMap.of(periodUid, firstAdPlaybackState));
+              return true;
+            }));
 
     AnalyticsListener listener = mock(AnalyticsListener.class);
     player.addAnalyticsListener(listener);
-    player.setMediaSource(mediaSource);
+    player.setMediaSource(mediaSourceRef.get());
     player.prepare();
     player.play();
     runUntilPlaybackState(player, Player.STATE_ENDED);
@@ -205,6 +240,7 @@ public final class ServerSideAdInsertionMediaSourceTest {
   @Test
   public void playbackWithNewlyInsertedAds_playsSuccessfulWithoutRendererResets() throws Exception {
     Context context = ApplicationProvider.getApplicationContext();
+    AtomicReference<Object> periodUid = new AtomicReference<>();
     CapturingRenderersFactory renderersFactory = new CapturingRenderersFactory(context);
     ExoPlayer player =
         new ExoPlayer.Builder(context, renderersFactory)
@@ -213,33 +249,43 @@ public final class ServerSideAdInsertionMediaSourceTest {
     player.setVideoSurface(new Surface(new SurfaceTexture(/* texName= */ 1)));
     PlaybackOutput playbackOutput = PlaybackOutput.register(player, renderersFactory);
 
-    ServerSideAdInsertionMediaSource mediaSource =
+    AdPlaybackState firstAdPlaybackState =
-        new ServerSideAdInsertionMediaSource(
-            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
-            /* adPlaybackStateUpdater= */ null);
-    AdPlaybackState adPlaybackState = new AdPlaybackState(/* adsId= */ new Object());
-    adPlaybackState =
         addAdGroupToAdPlaybackState(
-            adPlaybackState,
+            new AdPlaybackState(/* adsId= */ new Object()),
             /* fromPositionUs= */ 900_000,
             /* toPositionUs= */ 1_000_000,
             /* contentResumeOffsetUs= */ 0);
-    mediaSource.setAdPlaybackState(adPlaybackState);
+    AtomicReference<ServerSideAdInsertionMediaSource> mediaSourceRef = new AtomicReference<>();
-
+    mediaSourceRef.set(
+        new ServerSideAdInsertionMediaSource(
+            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
+            /* adPlaybackStateUpdater= */ contentTimeline -> {
+              periodUid.set(
+                  checkNotNull(
+                      contentTimeline.getPeriod(
+                              /* periodIndex= */ 0, new Timeline.Period(), /* setIds= */ true)
+                          .uid));
+              mediaSourceRef
+                  .get()
+                  .setAdPlaybackStates(ImmutableMap.of(periodUid.get(), firstAdPlaybackState));
+              return true;
+            }));
     AnalyticsListener listener = mock(AnalyticsListener.class);
     player.addAnalyticsListener(listener);
-    player.setMediaSource(mediaSource);
+    player.setMediaSource(mediaSourceRef.get());
     player.prepare();
 
     // Add ad at the current playback position during playback.
     runUntilPlaybackState(player, Player.STATE_READY);
-    adPlaybackState =
+    AdPlaybackState secondAdPlaybackState =
         addAdGroupToAdPlaybackState(
-            adPlaybackState,
+            firstAdPlaybackState,
             /* fromPositionUs= */ 0,
             /* toPositionUs= */ 500_000,
             /* contentResumeOffsetUs= */ 0);
-    mediaSource.setAdPlaybackState(adPlaybackState);
+    mediaSourceRef
+        .get()
+        .setAdPlaybackStates(ImmutableMap.of(periodUid.get(), secondAdPlaybackState));
     runUntilPendingCommandsAreFullyHandled(player);
 
     player.play();
@@ -265,6 +311,7 @@ public final class ServerSideAdInsertionMediaSourceTest {
   public void playbackWithAdditionalAdsInAdGroup_playsSuccessfulWithoutRendererResets()
       throws Exception {
     Context context = ApplicationProvider.getApplicationContext();
+    AtomicReference<Object> periodUid = new AtomicReference<>();
     CapturingRenderersFactory renderersFactory = new CapturingRenderersFactory(context);
     ExoPlayer player =
         new ExoPlayer.Builder(context, renderersFactory)
@@ -273,32 +320,45 @@ public final class ServerSideAdInsertionMediaSourceTest {
     player.setVideoSurface(new Surface(new SurfaceTexture(/* texName= */ 1)));
     PlaybackOutput playbackOutput = PlaybackOutput.register(player, renderersFactory);
 
-    ServerSideAdInsertionMediaSource mediaSource =
+    AdPlaybackState firstAdPlaybackState =
-        new ServerSideAdInsertionMediaSource(
-            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
-            /* adPlaybackStateUpdater= */ null);
-    AdPlaybackState adPlaybackState = new AdPlaybackState(/* adsId= */ new Object());
-    adPlaybackState =
         addAdGroupToAdPlaybackState(
-            adPlaybackState,
+            new AdPlaybackState(/* adsId= */ new Object()),
             /* fromPositionUs= */ 0,
             /* toPositionUs= */ 500_000,
             /* contentResumeOffsetUs= */ 0);
-    mediaSource.setAdPlaybackState(adPlaybackState);
+    AtomicReference<ServerSideAdInsertionMediaSource> mediaSourceRef = new AtomicReference<>();
+    mediaSourceRef.set(
+        new ServerSideAdInsertionMediaSource(
+            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
+            /* adPlaybackStateUpdater= */ contentTimeline -> {
+              if (periodUid.get() == null) {
+                periodUid.set(
+                    checkNotNull(
+                        contentTimeline.getPeriod(
+                                /* periodIndex= */ 0, new Timeline.Period(), /* setIds= */ true)
+                            .uid));
+                mediaSourceRef
+                    .get()
+                    .setAdPlaybackStates(ImmutableMap.of(periodUid.get(), firstAdPlaybackState));
+              }
+              return true;
+            }));
 
     AnalyticsListener listener = mock(AnalyticsListener.class);
     player.addAnalyticsListener(listener);
-    player.setMediaSource(mediaSource);
+    player.setMediaSource(mediaSourceRef.get());
     player.prepare();
 
     // Wait until playback is ready with first ad and then replace by 3 ads.
     runUntilPlaybackState(player, Player.STATE_READY);
-    adPlaybackState =
+    AdPlaybackState secondAdPlaybackState =
-        adPlaybackState
+        firstAdPlaybackState
             .withAdCount(/* adGroupIndex= */ 0, /* adCount= */ 3)
             .withAdDurationsUs(
                 /* adGroupIndex= */ 0, /* adDurationsUs...= */ 50_000, 250_000, 200_000);
-    mediaSource.setAdPlaybackState(adPlaybackState);
+    mediaSourceRef
+        .get()
+        .setAdPlaybackStates(ImmutableMap.of(periodUid.get(), secondAdPlaybackState));
     runUntilPendingCommandsAreFullyHandled(player);
 
     player.play();
@@ -327,10 +387,6 @@ public final class ServerSideAdInsertionMediaSourceTest {
         new ExoPlayer.Builder(context).setClock(new FakeClock(/* isAutoAdvancing= */ true)).build();
     player.setVideoSurface(new Surface(new SurfaceTexture(/* texName= */ 1)));
 
-    ServerSideAdInsertionMediaSource mediaSource =
-        new ServerSideAdInsertionMediaSource(
-            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
-            /* adPlaybackStateUpdater= */ null);
     AdPlaybackState adPlaybackState = new AdPlaybackState(/* adsId= */ new Object());
     adPlaybackState =
         addAdGroupToAdPlaybackState(
@@ -344,17 +400,32 @@ public final class ServerSideAdInsertionMediaSourceTest {
             /* fromPositionUs= */ 600_000,
             /* toPositionUs= */ 700_000,
             /* contentResumeOffsetUs= */ 1_000_000);
-    adPlaybackState =
+    AdPlaybackState firstAdPlaybackState =
         addAdGroupToAdPlaybackState(
             adPlaybackState,
             /* fromPositionUs= */ 900_000,
             /* toPositionUs= */ 1_000_000,
             /* contentResumeOffsetUs= */ 0);
-    mediaSource.setAdPlaybackState(adPlaybackState);
+
+    AtomicReference<ServerSideAdInsertionMediaSource> mediaSourceRef = new AtomicReference<>();
+    mediaSourceRef.set(
+        new ServerSideAdInsertionMediaSource(
+            new DefaultMediaSourceFactory(context).createMediaSource(MediaItem.fromUri(TEST_ASSET)),
+            /* adPlaybackStateUpdater= */ contentTimeline -> {
+              Object periodUid =
+                  checkNotNull(
+                      contentTimeline.getPeriod(
+                              /* periodIndex= */ 0, new Timeline.Period(), /* setIds= */ true)
+                          .uid);
+              mediaSourceRef
+                  .get()
+                  .setAdPlaybackStates(ImmutableMap.of(periodUid, firstAdPlaybackState));
+              return true;
+            }));
 
     AnalyticsListener listener = mock(AnalyticsListener.class);
     player.addAnalyticsListener(listener);
-    player.setMediaSource(mediaSource);
+    player.setMediaSource(mediaSourceRef.get());
     player.prepare();
     // Play to the first content part, then seek past the midroll.
     playUntilPosition(player, /* windowIndex= */ 0, /* positionMs= */ 100);