From 6330d46d9eb2d653503dd172df70ab34839ad6b5 Mon Sep 17 00:00:00 2001
From: tonihei <tonihei@google.com>
Date: Mon, 4 Feb 2019 17:40:01 +0000
Subject: [PATCH] Improve housekeeping of ConcatenatingMediaSource callbacks.

When calling releaseSource(), all pending messages will be removed. That means
that all action-on-completion callbacks which are somewhere in flight are
just dropped without being called. This change adds code to keep track of the
current state of each callback to allow all of them being called when the
source is released.

Issue:#5464
PiperOrigin-RevId: 232312528
---
 RELEASENOTES.md                               |   2 +
 .../source/ConcatenatingMediaSource.java      | 266 +++++++++++-------
 .../source/ConcatenatingMediaSourceTest.java  |  40 ++-
 3 files changed, 192 insertions(+), 116 deletions(-)

diff --git a/RELEASENOTES.md b/RELEASENOTES.md
index 271f2d94b9..5ba3aa7c8d 100644
--- a/RELEASENOTES.md
+++ b/RELEASENOTES.md
@@ -14,6 +14,8 @@
 * OkHttp extension: Upgrade OkHttp dependency to 3.12.1.
 * MP3: Wider fix for issue where streams would play twice on some Samsung
   devices ([#4519](https://github.com/google/ExoPlayer/issues/4519)).
+* Fix issue with dropped messages when releasing a `ConcatenatingMediaSource`
+  ([#5464](https://github.com/google/ExoPlayer/issues/5464)).
 
 ### 2.9.4 ###
 
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/source/ConcatenatingMediaSource.java b/library/core/src/main/java/com/google/android/exoplayer2/source/ConcatenatingMediaSource.java
index 9874004ceb..7baea9979f 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/source/ConcatenatingMediaSource.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/source/ConcatenatingMediaSource.java
@@ -29,7 +29,6 @@ import com.google.android.exoplayer2.source.ShuffleOrder.DefaultShuffleOrder;
 import com.google.android.exoplayer2.upstream.Allocator;
 import com.google.android.exoplayer2.upstream.TransferListener;
 import com.google.android.exoplayer2.util.Assertions;
-import com.google.android.exoplayer2.util.EventDispatcher;
 import com.google.android.exoplayer2.util.Util;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -37,9 +36,11 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Concatenates multiple {@link MediaSource}s. The list of {@link MediaSource}s can be modified
@@ -52,12 +53,19 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
   private static final int MSG_REMOVE = 1;
   private static final int MSG_MOVE = 2;
   private static final int MSG_SET_SHUFFLE_ORDER = 3;
-  private static final int MSG_NOTIFY_LISTENER = 4;
+  private static final int MSG_UPDATE_TIMELINE = 4;
   private static final int MSG_ON_COMPLETION = 5;
 
   // Accessed on any thread.
+  @GuardedBy("this")
   private final List<MediaSourceHolder> mediaSourcesPublic;
-  @Nullable private Handler playbackThreadHandler;
+
+  @GuardedBy("this")
+  private final Set<HandlerAndRunnable> pendingOnCompletionActions;
+
+  @GuardedBy("this")
+  @Nullable
+  private Handler playbackThreadHandler;
 
   // Accessed on the playback thread only.
   private final List<MediaSourceHolder> mediaSourceHolders;
@@ -68,8 +76,8 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
   private final Timeline.Window window;
   private final Timeline.Period period;
 
-  private boolean listenerNotificationScheduled;
-  private EventDispatcher<Runnable> pendingOnCompletionActions;
+  private boolean timelineUpdateScheduled;
+  private Set<HandlerAndRunnable> nextTimelineUpdateOnCompletionActions;
   private ShuffleOrder shuffleOrder;
   private int windowCount;
   private int periodCount;
@@ -128,7 +136,8 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
     this.mediaSourceByUid = new HashMap<>();
     this.mediaSourcesPublic = new ArrayList<>();
     this.mediaSourceHolders = new ArrayList<>();
-    this.pendingOnCompletionActions = new EventDispatcher<>();
+    this.nextTimelineUpdateOnCompletionActions = new HashSet<>();
+    this.pendingOnCompletionActions = new HashSet<>();
     this.isAtomic = isAtomic;
     this.useLazyPreparation = useLazyPreparation;
     window = new Timeline.Window();
@@ -149,13 +158,13 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    * Appends a {@link MediaSource} to the playlist and executes a custom action on completion.
    *
    * @param mediaSource The {@link MediaSource} to be added to the list.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     source has been added to the playlist.
    */
   public final synchronized void addMediaSource(
-      MediaSource mediaSource, Handler handler, Runnable actionOnCompletion) {
-    addMediaSource(mediaSourcesPublic.size(), mediaSource, handler, actionOnCompletion);
+      MediaSource mediaSource, Handler handler, Runnable onCompletionAction) {
+    addMediaSource(mediaSourcesPublic.size(), mediaSource, handler, onCompletionAction);
   }
 
   /**
@@ -170,7 +179,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
         index,
         Collections.singletonList(mediaSource),
         /* handler= */ null,
-        /* actionOnCompletion= */ null);
+        /* onCompletionAction= */ null);
   }
 
   /**
@@ -179,14 +188,14 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    * @param index The index at which the new {@link MediaSource} will be inserted. This index must
    *     be in the range of 0 &lt;= index &lt;= {@link #getSize()}.
    * @param mediaSource The {@link MediaSource} to be added to the list.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     source has been added to the playlist.
    */
   public final synchronized void addMediaSource(
-      int index, MediaSource mediaSource, Handler handler, Runnable actionOnCompletion) {
+      int index, MediaSource mediaSource, Handler handler, Runnable onCompletionAction) {
     addPublicMediaSources(
-        index, Collections.singletonList(mediaSource), handler, actionOnCompletion);
+        index, Collections.singletonList(mediaSource), handler, onCompletionAction);
   }
 
   /**
@@ -200,7 +209,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
         mediaSourcesPublic.size(),
         mediaSources,
         /* handler= */ null,
-        /* actionOnCompletion= */ null);
+        /* onCompletionAction= */ null);
   }
 
   /**
@@ -209,13 +218,13 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *
    * @param mediaSources A collection of {@link MediaSource}s to be added to the list. The media
    *     sources are added in the order in which they appear in this collection.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     sources have been added to the playlist.
    */
   public final synchronized void addMediaSources(
-      Collection<MediaSource> mediaSources, Handler handler, Runnable actionOnCompletion) {
-    addPublicMediaSources(mediaSourcesPublic.size(), mediaSources, handler, actionOnCompletion);
+      Collection<MediaSource> mediaSources, Handler handler, Runnable onCompletionAction) {
+    addPublicMediaSources(mediaSourcesPublic.size(), mediaSources, handler, onCompletionAction);
   }
 
   /**
@@ -227,7 +236,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *     sources are added in the order in which they appear in this collection.
    */
   public final synchronized void addMediaSources(int index, Collection<MediaSource> mediaSources) {
-    addPublicMediaSources(index, mediaSources, /* handler= */ null, /* actionOnCompletion= */ null);
+    addPublicMediaSources(index, mediaSources, /* handler= */ null, /* onCompletionAction= */ null);
   }
 
   /**
@@ -237,16 +246,16 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *     be in the range of 0 &lt;= index &lt;= {@link #getSize()}.
    * @param mediaSources A collection of {@link MediaSource}s to be added to the list. The media
    *     sources are added in the order in which they appear in this collection.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     sources have been added to the playlist.
    */
   public final synchronized void addMediaSources(
       int index,
       Collection<MediaSource> mediaSources,
       Handler handler,
-      Runnable actionOnCompletion) {
-    addPublicMediaSources(index, mediaSources, handler, actionOnCompletion);
+      Runnable onCompletionAction) {
+    addPublicMediaSources(index, mediaSources, handler, onCompletionAction);
   }
 
   /**
@@ -262,7 +271,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *     range of 0 &lt;= index &lt; {@link #getSize()}.
    */
   public final synchronized void removeMediaSource(int index) {
-    removePublicMediaSources(index, index + 1, /* handler= */ null, /* actionOnCompletion= */ null);
+    removePublicMediaSources(index, index + 1, /* handler= */ null, /* onCompletionAction= */ null);
   }
 
   /**
@@ -276,13 +285,13 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *
    * @param index The index at which the media source will be removed. This index must be in the
    *     range of 0 &lt;= index &lt; {@link #getSize()}.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     source has been removed from the playlist.
    */
   public final synchronized void removeMediaSource(
-      int index, Handler handler, Runnable actionOnCompletion) {
-    removePublicMediaSources(index, index + 1, handler, actionOnCompletion);
+      int index, Handler handler, Runnable onCompletionAction) {
+    removePublicMediaSources(index, index + 1, handler, onCompletionAction);
   }
 
   /**
@@ -301,7 +310,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    */
   public final synchronized void removeMediaSourceRange(int fromIndex, int toIndex) {
     removePublicMediaSources(
-        fromIndex, toIndex, /* handler= */ null, /* actionOnCompletion= */ null);
+        fromIndex, toIndex, /* handler= */ null, /* onCompletionAction= */ null);
   }
 
   /**
@@ -315,15 +324,15 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *     removed. This index must be in the range of 0 &lt;= index &lt;= {@link #getSize()}.
    * @param toIndex The final range index, pointing to the first media source that will be left
    *     untouched. This index must be in the range of 0 &lt;= index &lt;= {@link #getSize()}.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     source range has been removed from the playlist.
    * @throws IllegalArgumentException When the range is malformed, i.e. {@code fromIndex} &lt; 0,
    *     {@code toIndex} &gt; {@link #getSize()}, {@code fromIndex} &gt; {@code toIndex}
    */
   public final synchronized void removeMediaSourceRange(
-      int fromIndex, int toIndex, Handler handler, Runnable actionOnCompletion) {
-    removePublicMediaSources(fromIndex, toIndex, handler, actionOnCompletion);
+      int fromIndex, int toIndex, Handler handler, Runnable onCompletionAction) {
+    removePublicMediaSources(fromIndex, toIndex, handler, onCompletionAction);
   }
 
   /**
@@ -336,7 +345,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    */
   public final synchronized void moveMediaSource(int currentIndex, int newIndex) {
     movePublicMediaSource(
-        currentIndex, newIndex, /* handler= */ null, /* actionOnCompletion= */ null);
+        currentIndex, newIndex, /* handler= */ null, /* onCompletionAction= */ null);
   }
 
   /**
@@ -347,13 +356,13 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    *     in the range of 0 &lt;= index &lt; {@link #getSize()}.
    * @param newIndex The target index of the media source in the playlist. This index must be in the
    *     range of 0 &lt;= index &lt; {@link #getSize()}.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the media
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the media
    *     source has been moved.
    */
   public final synchronized void moveMediaSource(
-      int currentIndex, int newIndex, Handler handler, Runnable actionOnCompletion) {
-    movePublicMediaSource(currentIndex, newIndex, handler, actionOnCompletion);
+      int currentIndex, int newIndex, Handler handler, Runnable onCompletionAction) {
+    movePublicMediaSource(currentIndex, newIndex, handler, onCompletionAction);
   }
 
   /** Clears the playlist. */
@@ -364,12 +373,12 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
   /**
    * Clears the playlist and executes a custom action on completion.
    *
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the playlist
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the playlist
    *     has been cleared.
    */
-  public final synchronized void clear(Handler handler, Runnable actionOnCompletion) {
-    removeMediaSourceRange(0, getSize(), handler, actionOnCompletion);
+  public final synchronized void clear(Handler handler, Runnable onCompletionAction) {
+    removeMediaSourceRange(0, getSize(), handler, onCompletionAction);
   }
 
   /** Returns the number of media sources in the playlist. */
@@ -393,20 +402,20 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
    * @param shuffleOrder A {@link ShuffleOrder}.
    */
   public final synchronized void setShuffleOrder(ShuffleOrder shuffleOrder) {
-    setPublicShuffleOrder(shuffleOrder, /* handler= */ null, /* actionOnCompletion= */ null);
+    setPublicShuffleOrder(shuffleOrder, /* handler= */ null, /* onCompletionAction= */ null);
   }
 
   /**
    * Sets a new shuffle order to use when shuffling the child media sources.
    *
    * @param shuffleOrder A {@link ShuffleOrder}.
-   * @param handler The {@link Handler} to run {@code actionOnCompletion}.
-   * @param actionOnCompletion A {@link Runnable} which is executed immediately after the shuffle
+   * @param handler The {@link Handler} to run {@code onCompletionAction}.
+   * @param onCompletionAction A {@link Runnable} which is executed immediately after the shuffle
    *     order has been changed.
    */
   public final synchronized void setShuffleOrder(
-      ShuffleOrder shuffleOrder, Handler handler, Runnable actionOnCompletion) {
-    setPublicShuffleOrder(shuffleOrder, handler, actionOnCompletion);
+      ShuffleOrder shuffleOrder, Handler handler, Runnable onCompletionAction) {
+    setPublicShuffleOrder(shuffleOrder, handler, onCompletionAction);
   }
 
   // CompositeMediaSource implementation.
@@ -425,11 +434,11 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
     super.prepareSourceInternal(player, isTopLevelSource, mediaTransferListener);
     playbackThreadHandler = new Handler(/* callback= */ this::handleMessage);
     if (mediaSourcesPublic.isEmpty()) {
-      notifyListener();
+      updateTimelineAndScheduleOnCompletionActions();
     } else {
       shuffleOrder = shuffleOrder.cloneAndInsert(0, mediaSourcesPublic.size());
       addMediaSourcesInternal(0, mediaSourcesPublic);
-      scheduleListenerNotification();
+      scheduleTimelineUpdate();
     }
   }
 
@@ -485,6 +494,9 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
       playbackThreadHandler.removeCallbacksAndMessages(null);
       playbackThreadHandler = null;
     }
+    timelineUpdateScheduled = false;
+    nextTimelineUpdateOnCompletionActions.clear();
+    dispatchOnCompletionActions(pendingOnCompletionActions);
   }
 
   @Override
@@ -524,8 +536,9 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
       int index,
       Collection<MediaSource> mediaSources,
       @Nullable Handler handler,
-      @Nullable Runnable actionOnCompletion) {
-    Assertions.checkArgument((handler == null) == (actionOnCompletion == null));
+      @Nullable Runnable onCompletionAction) {
+    Assertions.checkArgument((handler == null) == (onCompletionAction == null));
+    Handler playbackThreadHandler = this.playbackThreadHandler;
     for (MediaSource mediaSource : mediaSources) {
       Assertions.checkNotNull(mediaSource);
     }
@@ -535,12 +548,12 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
     }
     mediaSourcesPublic.addAll(index, mediaSourceHolders);
     if (playbackThreadHandler != null && !mediaSources.isEmpty()) {
+      HandlerAndRunnable callbackAction = createOnCompletionAction(handler, onCompletionAction);
       playbackThreadHandler
-          .obtainMessage(
-              MSG_ADD, new MessageData<>(index, mediaSourceHolders, handler, actionOnCompletion))
+          .obtainMessage(MSG_ADD, new MessageData<>(index, mediaSourceHolders, callbackAction))
           .sendToTarget();
-    } else if (actionOnCompletion != null && handler != null) {
-      handler.post(actionOnCompletion);
+    } else if (onCompletionAction != null && handler != null) {
+      handler.post(onCompletionAction);
     }
   }
 
@@ -549,16 +562,17 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
       int fromIndex,
       int toIndex,
       @Nullable Handler handler,
-      @Nullable Runnable actionOnCompletion) {
-    Assertions.checkArgument((handler == null) == (actionOnCompletion == null));
+      @Nullable Runnable onCompletionAction) {
+    Assertions.checkArgument((handler == null) == (onCompletionAction == null));
+    Handler playbackThreadHandler = this.playbackThreadHandler;
     Util.removeRange(mediaSourcesPublic, fromIndex, toIndex);
     if (playbackThreadHandler != null) {
+      HandlerAndRunnable callbackAction = createOnCompletionAction(handler, onCompletionAction);
       playbackThreadHandler
-          .obtainMessage(
-              MSG_REMOVE, new MessageData<>(fromIndex, toIndex, handler, actionOnCompletion))
+          .obtainMessage(MSG_REMOVE, new MessageData<>(fromIndex, toIndex, callbackAction))
           .sendToTarget();
-    } else if (actionOnCompletion != null && handler != null) {
-      handler.post(actionOnCompletion);
+    } else if (onCompletionAction != null && handler != null) {
+      handler.post(onCompletionAction);
     }
   }
 
@@ -567,23 +581,24 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
       int currentIndex,
       int newIndex,
       @Nullable Handler handler,
-      @Nullable Runnable actionOnCompletion) {
-    Assertions.checkArgument((handler == null) == (actionOnCompletion == null));
+      @Nullable Runnable onCompletionAction) {
+    Assertions.checkArgument((handler == null) == (onCompletionAction == null));
+    Handler playbackThreadHandler = this.playbackThreadHandler;
     mediaSourcesPublic.add(newIndex, mediaSourcesPublic.remove(currentIndex));
     if (playbackThreadHandler != null) {
+      HandlerAndRunnable callbackAction = createOnCompletionAction(handler, onCompletionAction);
       playbackThreadHandler
-          .obtainMessage(
-              MSG_MOVE, new MessageData<>(currentIndex, newIndex, handler, actionOnCompletion))
+          .obtainMessage(MSG_MOVE, new MessageData<>(currentIndex, newIndex, callbackAction))
           .sendToTarget();
-    } else if (actionOnCompletion != null && handler != null) {
-      handler.post(actionOnCompletion);
+    } else if (onCompletionAction != null && handler != null) {
+      handler.post(onCompletionAction);
     }
   }
 
   @GuardedBy("this")
   private void setPublicShuffleOrder(
-      ShuffleOrder shuffleOrder, @Nullable Handler handler, @Nullable Runnable actionOnCompletion) {
-    Assertions.checkArgument((handler == null) == (actionOnCompletion == null));
+      ShuffleOrder shuffleOrder, @Nullable Handler handler, @Nullable Runnable onCompletionAction) {
+    Assertions.checkArgument((handler == null) == (onCompletionAction == null));
     Handler playbackThreadHandler = this.playbackThreadHandler;
     if (playbackThreadHandler != null) {
       int size = getSize();
@@ -593,20 +608,33 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
                 .cloneAndClear()
                 .cloneAndInsert(/* insertionIndex= */ 0, /* insertionCount= */ size);
       }
+      HandlerAndRunnable callbackAction = createOnCompletionAction(handler, onCompletionAction);
       playbackThreadHandler
           .obtainMessage(
               MSG_SET_SHUFFLE_ORDER,
-              new MessageData<>(/* index= */ 0, shuffleOrder, handler, actionOnCompletion))
+              new MessageData<>(/* index= */ 0, shuffleOrder, callbackAction))
           .sendToTarget();
     } else {
       this.shuffleOrder =
           shuffleOrder.getLength() > 0 ? shuffleOrder.cloneAndClear() : shuffleOrder;
-      if (actionOnCompletion != null && handler != null) {
-        handler.post(actionOnCompletion);
+      if (onCompletionAction != null && handler != null) {
+        handler.post(onCompletionAction);
       }
     }
   }
 
+  @GuardedBy("this")
+  @Nullable
+  private HandlerAndRunnable createOnCompletionAction(
+      @Nullable Handler handler, @Nullable Runnable runnable) {
+    if (handler == null || runnable == null) {
+      return null;
+    }
+    HandlerAndRunnable handlerAndRunnable = new HandlerAndRunnable(handler, runnable);
+    pendingOnCompletionActions.add(handlerAndRunnable);
+    return handlerAndRunnable;
+  }
+
   // Internal methods. Called on the playback thread.
 
   @SuppressWarnings("unchecked")
@@ -617,7 +645,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
             (MessageData<Collection<MediaSourceHolder>>) Util.castNonNull(msg.obj);
         shuffleOrder = shuffleOrder.cloneAndInsert(addMessage.index, addMessage.customData.size());
         addMediaSourcesInternal(addMessage.index, addMessage.customData);
-        scheduleListenerNotification(addMessage.handler, addMessage.actionOnCompletion);
+        scheduleTimelineUpdate(addMessage.onCompletionAction);
         break;
       case MSG_REMOVE:
         MessageData<Integer> removeMessage = (MessageData<Integer>) Util.castNonNull(msg.obj);
@@ -631,29 +659,27 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
         for (int index = toIndex - 1; index >= fromIndex; index--) {
           removeMediaSourceInternal(index);
         }
-        scheduleListenerNotification(removeMessage.handler, removeMessage.actionOnCompletion);
+        scheduleTimelineUpdate(removeMessage.onCompletionAction);
         break;
       case MSG_MOVE:
         MessageData<Integer> moveMessage = (MessageData<Integer>) Util.castNonNull(msg.obj);
         shuffleOrder = shuffleOrder.cloneAndRemove(moveMessage.index, moveMessage.index + 1);
         shuffleOrder = shuffleOrder.cloneAndInsert(moveMessage.customData, 1);
         moveMediaSourceInternal(moveMessage.index, moveMessage.customData);
-        scheduleListenerNotification(moveMessage.handler, moveMessage.actionOnCompletion);
+        scheduleTimelineUpdate(moveMessage.onCompletionAction);
         break;
       case MSG_SET_SHUFFLE_ORDER:
         MessageData<ShuffleOrder> shuffleOrderMessage =
             (MessageData<ShuffleOrder>) Util.castNonNull(msg.obj);
         shuffleOrder = shuffleOrderMessage.customData;
-        scheduleListenerNotification(
-            shuffleOrderMessage.handler, shuffleOrderMessage.actionOnCompletion);
+        scheduleTimelineUpdate(shuffleOrderMessage.onCompletionAction);
         break;
-      case MSG_NOTIFY_LISTENER:
-        notifyListener();
+      case MSG_UPDATE_TIMELINE:
+        updateTimelineAndScheduleOnCompletionActions();
         break;
       case MSG_ON_COMPLETION:
-        EventDispatcher<Runnable> actionsOnCompletion =
-            (EventDispatcher<Runnable>) Util.castNonNull(msg.obj);
-        actionsOnCompletion.dispatch(Runnable::run);
+        Set<HandlerAndRunnable> actions = (Set<HandlerAndRunnable>) Util.castNonNull(msg.obj);
+        dispatchOnCompletionActions(actions);
         break;
       default:
         throw new IllegalStateException();
@@ -661,36 +687,48 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
     return true;
   }
 
-  private void scheduleListenerNotification() {
-    scheduleListenerNotification(/* handler= */ null, /* actionOnCompletion= */ null);
+  private void scheduleTimelineUpdate() {
+    scheduleTimelineUpdate(/* onCompletionAction= */ null);
   }
 
-  private void scheduleListenerNotification(
-      @Nullable Handler handler, @Nullable Runnable actionOnCompletion) {
-    if (!listenerNotificationScheduled) {
-      Assertions.checkNotNull(playbackThreadHandler)
-          .obtainMessage(MSG_NOTIFY_LISTENER)
-          .sendToTarget();
-      listenerNotificationScheduled = true;
+  private void scheduleTimelineUpdate(@Nullable HandlerAndRunnable onCompletionAction) {
+    if (!timelineUpdateScheduled) {
+      getPlaybackThreadHandlerOnPlaybackThread().obtainMessage(MSG_UPDATE_TIMELINE).sendToTarget();
+      timelineUpdateScheduled = true;
     }
-    if (actionOnCompletion != null && handler != null) {
-      pendingOnCompletionActions.addListener(handler, actionOnCompletion);
+    if (onCompletionAction != null) {
+      nextTimelineUpdateOnCompletionActions.add(onCompletionAction);
     }
   }
 
-  private void notifyListener() {
-    listenerNotificationScheduled = false;
-    EventDispatcher<Runnable> actionsOnCompletion = pendingOnCompletionActions;
-    pendingOnCompletionActions = new EventDispatcher<>();
+  private void updateTimelineAndScheduleOnCompletionActions() {
+    timelineUpdateScheduled = false;
+    Set<HandlerAndRunnable> onCompletionActions = nextTimelineUpdateOnCompletionActions;
+    nextTimelineUpdateOnCompletionActions = new HashSet<>();
     refreshSourceInfo(
         new ConcatenatedTimeline(
             mediaSourceHolders, windowCount, periodCount, shuffleOrder, isAtomic),
         /* manifest= */ null);
-    Assertions.checkNotNull(playbackThreadHandler)
-        .obtainMessage(MSG_ON_COMPLETION, actionsOnCompletion)
+    getPlaybackThreadHandlerOnPlaybackThread()
+        .obtainMessage(MSG_ON_COMPLETION, onCompletionActions)
         .sendToTarget();
   }
 
+  @SuppressWarnings("GuardedBy")
+  private Handler getPlaybackThreadHandlerOnPlaybackThread() {
+    // Write access to this value happens on the playback thread only, so playback thread reads
+    // don't need to be synchronized.
+    return Assertions.checkNotNull(playbackThreadHandler);
+  }
+
+  private synchronized void dispatchOnCompletionActions(
+      Set<HandlerAndRunnable> onCompletionActions) {
+    for (HandlerAndRunnable pendingAction : onCompletionActions) {
+      pendingAction.dispatch();
+    }
+    pendingOnCompletionActions.removeAll(onCompletionActions);
+  }
+
   private void addMediaSourcesInternal(
       int index, Collection<MediaSourceHolder> mediaSourceHolders) {
     for (MediaSourceHolder mediaSourceHolder : mediaSourceHolders) {
@@ -787,7 +825,7 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
       }
     }
     mediaSourceHolder.isPrepared = true;
-    scheduleListenerNotification();
+    scheduleTimelineUpdate();
   }
 
   private void removeMediaSourceInternal(int index) {
@@ -900,15 +938,12 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
 
     public final int index;
     public final T customData;
-    @Nullable public final Handler handler;
-    @Nullable public final Runnable actionOnCompletion;
+    @Nullable public final HandlerAndRunnable onCompletionAction;
 
-    public MessageData(
-        int index, T customData, @Nullable Handler handler, @Nullable Runnable actionOnCompletion) {
+    public MessageData(int index, T customData, @Nullable HandlerAndRunnable onCompletionAction) {
       this.index = index;
       this.customData = customData;
-      this.handler = handler;
-      this.actionOnCompletion = actionOnCompletion;
+      this.onCompletionAction = onCompletionAction;
     }
   }
 
@@ -1161,5 +1196,20 @@ public class ConcatenatingMediaSource extends CompositeMediaSource<MediaSourceHo
       // Do nothing.
     }
   }
+
+  private static final class HandlerAndRunnable {
+
+    private final Handler handler;
+    private final Runnable runnable;
+
+    public HandlerAndRunnable(Handler handler, Runnable runnable) {
+      this.handler = handler;
+      this.runnable = runnable;
+    }
+
+    public void dispatch() {
+      handler.post(runnable);
+    }
+  }
 }
 
diff --git a/library/core/src/test/java/com/google/android/exoplayer2/source/ConcatenatingMediaSourceTest.java b/library/core/src/test/java/com/google/android/exoplayer2/source/ConcatenatingMediaSourceTest.java
index bb86c6186d..2665b61f8a 100644
--- a/library/core/src/test/java/com/google/android/exoplayer2/source/ConcatenatingMediaSourceTest.java
+++ b/library/core/src/test/java/com/google/android/exoplayer2/source/ConcatenatingMediaSourceTest.java
@@ -17,6 +17,7 @@ package com.google.android.exoplayer2.source;
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
 import android.os.ConditionVariable;
@@ -25,6 +26,7 @@ import com.google.android.exoplayer2.C;
 import com.google.android.exoplayer2.Player;
 import com.google.android.exoplayer2.Timeline;
 import com.google.android.exoplayer2.source.MediaSource.MediaPeriodId;
+import com.google.android.exoplayer2.source.MediaSource.SourceInfoRefreshListener;
 import com.google.android.exoplayer2.source.ShuffleOrder.DefaultShuffleOrder;
 import com.google.android.exoplayer2.testutil.DummyMainThread;
 import com.google.android.exoplayer2.testutil.FakeMediaSource;
@@ -42,7 +44,6 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mockito;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 
@@ -416,7 +417,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationAddSingle() {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
 
     mediaSource.addMediaSource(createFakeMediaSource(), new Handler(), runnable);
     verify(runnable).run();
@@ -424,7 +425,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationAddMultiple() {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
 
     mediaSource.addMediaSources(
         Arrays.asList(new MediaSource[] {createFakeMediaSource(), createFakeMediaSource()}),
@@ -435,7 +436,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationAddSingleWithIndex() {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
 
     mediaSource.addMediaSource(/* index */ 0, createFakeMediaSource(), new Handler(), runnable);
     verify(runnable).run();
@@ -443,7 +444,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationAddMultipleWithIndex() {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
 
     mediaSource.addMediaSources(
         /* index */ 0,
@@ -455,7 +456,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationRemove() {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
 
     mediaSource.addMediaSource(createFakeMediaSource());
     mediaSource.removeMediaSource(/* index */ 0, new Handler(), runnable);
@@ -464,7 +465,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationMove() {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
 
     mediaSource.addMediaSources(
         Arrays.asList(new MediaSource[] {createFakeMediaSource(), createFakeMediaSource()}));
@@ -588,6 +589,29 @@ public final class ConcatenatingMediaSourceTest {
     }
   }
 
+  @Test
+  public void testCustomCallbackIsCalledAfterRelease() throws IOException {
+    DummyMainThread dummyMainThread = new DummyMainThread();
+    ConditionVariable callbackCalledCondition = new ConditionVariable();
+    try {
+      dummyMainThread.runOnMainThread(
+          () -> {
+            SourceInfoRefreshListener listener = mock(SourceInfoRefreshListener.class);
+            mediaSource.addMediaSources(Arrays.asList(createMediaSources(2)));
+            mediaSource.prepareSource(listener, /* mediaTransferListener= */ null);
+            mediaSource.moveMediaSource(
+                /* currentIndex= */ 0,
+                /* newIndex= */ 1,
+                new Handler(),
+                callbackCalledCondition::open);
+            mediaSource.releaseSource(listener);
+          });
+      assertThat(callbackCalledCondition.block(MediaSourceTestRunner.TIMEOUT_MS)).isTrue();
+    } finally {
+      dummyMainThread.release();
+    }
+  }
+
   @Test
   public void testPeriodCreationWithAds() throws IOException, InterruptedException {
     // Create concatenated media source with ad child source.
@@ -973,7 +997,7 @@ public final class ConcatenatingMediaSourceTest {
 
   @Test
   public void testCustomCallbackBeforePreparationSetShuffleOrder() throws Exception {
-    Runnable runnable = Mockito.mock(Runnable.class);
+    Runnable runnable = mock(Runnable.class);
     mediaSource.setShuffleOrder(
         new ShuffleOrder.UnshuffledShuffleOrder(/* length= */ 0), new Handler(), runnable);