2 files changed, 53 insertions, 6 deletions

diff --git a/dom/html/HTMLMediaElement.cpp b/dom/html/HTMLMediaElement.cpp
index 1f1a545fa..0b9f606f1 100644
--- a/dom/html/HTMLMediaElement.cpp
+++ b/dom/html/HTMLMediaElement.cpp
@@ -817,6 +817,7 @@ NS_IMPL_CYCLE_COLLECTION_TRAVERSE_BEGIN_INHERITED(HTMLMediaElement, nsGenericHTM
   NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mErrorSink->mError)
   for (uint32_t i = 0; i < tmp->mOutputStreams.Length(); ++i) {
     NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mOutputStreams[i].mStream);
+    NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mOutputStreams[i].mPreCreatedTracks);
   }
   NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mPlayed);
   NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mTextTrackManager)
@@ -1020,6 +1021,17 @@ void HTMLMediaElement::ShutdownDecoder()
   RemoveMediaElementFromURITable();
   NS_ASSERTION(mDecoder, "Must have decoder to shut down");
   mWaitingForKeyListener.DisconnectIfExists();
+  for (OutputMediaStream& out : mOutputStreams) {
+    if (!out.mCapturingDecoder) {
+      continue;
+    }
+    if (!out.mStream) {
+      continue;
+    }
+    out.mNextAvailableTrackID = std::max<TrackID>(
+      mDecoder->NextAvailableTrackIDFor(out.mStream->GetInputStream()),
+      out.mNextAvailableTrackID);
+  }
   mDecoder->Shutdown();
   mDecoder = nullptr;
 }
@@ -2730,6 +2742,7 @@ HTMLMediaElement::CaptureStreamInternal(bool aFinishWhenEnded,
   if (mDecoder) {
     out->mCapturingDecoder = true;
     mDecoder->AddOutputStream(out->mStream->GetInputStream()->AsProcessedStream(),
+                              out->mNextAvailableTrackID,
                               aFinishWhenEnded);
   } else if (mSrcStream) {
     out->mCapturingMediaStream = true;
@@ -2743,23 +2756,26 @@ HTMLMediaElement::CaptureStreamInternal(bool aFinishWhenEnded,
 
   if (mDecoder) {
     if (HasAudio()) {
-      TrackID audioTrackId = mMediaInfo.mAudio.mTrackId;
+      TrackID audioTrackId = out->mNextAvailableTrackID++;
       RefPtr<MediaStreamTrackSource> trackSource =
         getter->GetMediaStreamTrackSource(audioTrackId);
       RefPtr<MediaStreamTrack> track =
-        out->mStream->CreateDOMTrack(audioTrackId, MediaSegment::AUDIO,
+        out->mStream->CreateDOMTrack(audioTrackId,
+                                     MediaSegment::AUDIO,
                                      trackSource);
+      out->mPreCreatedTracks.AppendElement(track);
       out->mStream->AddTrackInternal(track);
       LOG(LogLevel::Debug,
           ("Created audio track %d for captured decoder", audioTrackId));
     }
     if (IsVideo() && HasVideo() && !out->mCapturingAudioOnly) {
-      TrackID videoTrackId = mMediaInfo.mVideo.mTrackId;
+      TrackID videoTrackId = out->mNextAvailableTrackID++;
       RefPtr<MediaStreamTrackSource> trackSource =
         getter->GetMediaStreamTrackSource(videoTrackId);
       RefPtr<MediaStreamTrack> track =
         out->mStream->CreateDOMTrack(videoTrackId, MediaSegment::VIDEO,
                                      trackSource);
+      out->mPreCreatedTracks.AppendElement(track);
       out->mStream->AddTrackInternal(track);
       LOG(LogLevel::Debug,
           ("Created video track %d for captured decoder", videoTrackId));
@@ -2848,6 +2864,25 @@ NS_IMETHODIMP HTMLMediaElement::GetMozAudioCaptured(bool* aCaptured)
   return NS_OK;
 }
 
+void
+HTMLMediaElement::EndPreCreatedCapturedDecoderTracks()
+{
+  MOZ_ASSERT(NS_IsMainThread());
+  for (OutputMediaStream& ms : mOutputStreams) {
+    if (!ms.mCapturingDecoder) {
+      continue;
+    }
+    for (RefPtr<MediaStreamTrack>& t : ms.mPreCreatedTracks) {
+      if (t->Ended()) {
+        continue;
+      }
+      NS_DispatchToMainThread(NewRunnableMethod(
+        t, &MediaStreamTrack::OverrideEnded));
+    }
+    ms.mPreCreatedTracks.Clear();
+  }
+}
+
 class MediaElementSetForURI : public nsURIHashKey {
 public:
   explicit MediaElementSetForURI(const nsIURI* aKey) : nsURIHashKey(aKey) {}
@@ -3380,11 +3415,12 @@ HTMLMediaElement::WakeLockRelease()
 }
 
 HTMLMediaElement::OutputMediaStream::OutputMediaStream()
-  : mFinishWhenEnded(false)
+  : mNextAvailableTrackID(1)
+  , mFinishWhenEnded(false)
   , mCapturingAudioOnly(false)
   , mCapturingDecoder(false)
   , mCapturingMediaStream(false)
-  , mNextAvailableTrackID(1) {}
+{}
 
 HTMLMediaElement::OutputMediaStream::~OutputMediaStream()
 {
@@ -4008,6 +4044,7 @@ nsresult HTMLMediaElement::FinishDecoderSetup(MediaDecoder* aDecoder,
 
     ms.mCapturingDecoder = true;
     aDecoder->AddOutputStream(ms.mStream->GetInputStream()->AsProcessedStream(),
+                              ms.mNextAvailableTrackID,
                               ms.mFinishWhenEnded);
   }
 
diff --git a/dom/html/HTMLMediaElement.h b/dom/html/HTMLMediaElement.h
index b65049206..af944a318 100644
--- a/dom/html/HTMLMediaElement.h
+++ b/dom/html/HTMLMediaElement.h
@@ -675,6 +675,11 @@ public:
     return mAudioCaptured;
   }
 
+  /**
+   * Ensures any MediaStreamTracks captured from a MediaDecoder are ended.
+   */
+  void EndPreCreatedCapturedDecoderTracks();
+
   void MozGetMetadata(JSContext* aCx, JS::MutableHandle<JSObject*> aResult,
                       ErrorResult& aRv);
 
@@ -815,13 +820,18 @@ protected:
     ~OutputMediaStream();
 
     RefPtr<DOMMediaStream> mStream;
+    TrackID mNextAvailableTrackID;
     bool mFinishWhenEnded;
     bool mCapturingAudioOnly;
     bool mCapturingDecoder;
     bool mCapturingMediaStream;
 
+    // The following members are keeping state for a captured MediaDecoder.
+    // Tracks that were created on main thread before MediaDecoder fed them
+    // to the MediaStreamGraph.
+    nsTArray<RefPtr<MediaStreamTrack>> mPreCreatedTracks;
+
     // The following members are keeping state for a captured MediaStream.
-    TrackID mNextAvailableTrackID;
     nsTArray<Pair<nsString, RefPtr<MediaInputPort>>> mTrackPorts;
   };