From 533987a47d753c8a3bb9f74df8cf70c5eda74e7e Mon Sep 17 00:00:00 2001
From: trav90 <travawine@protonmail.ch>
Date: Sun, 12 Mar 2017 04:06:02 -0500
Subject: [PATCH] Implement the ability to get the currently running task queue

This will be useful for asserting things.
---
 dom/media/AbstractThread.cpp |  5 ++++-
 dom/media/MediaDecoder.cpp   |  2 ++
 dom/media/MediaTaskQueue.cpp | 19 +++++++++++++++++--
 dom/media/MediaTaskQueue.h   | 26 ++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/dom/media/AbstractThread.cpp b/dom/media/AbstractThread.cpp
index 7adcf7e1b0..6d249ba5c5 100644
--- a/dom/media/AbstractThread.cpp
+++ b/dom/media/AbstractThread.cpp
@@ -6,6 +6,7 @@
 
 #include "AbstractThread.h"
 
+#include "MediaTaskQueue.h"
 #include "nsThreadUtils.h"
 
 #include "mozilla/ClearOnShutdown.h"
@@ -27,7 +28,9 @@ template<>
 bool
 AbstractThreadImpl<nsIThread>::IsCurrentThreadIn()
 {
-  return NS_GetCurrentThread() == mTarget;
+  bool in = NS_GetCurrentThread() == mTarget;
+  MOZ_ASSERT_IF(in, MediaTaskQueue::GetCurrentQueue() == nullptr);
+  return in;
 }
 
 AbstractThread*
diff --git a/dom/media/MediaDecoder.cpp b/dom/media/MediaDecoder.cpp
index 1db731611e..cf811b47ba 100644
--- a/dom/media/MediaDecoder.cpp
+++ b/dom/media/MediaDecoder.cpp
@@ -120,6 +120,8 @@ void
 MediaDecoder::InitStatics()
 {
   AbstractThread::InitStatics();
+
+  MediaTaskQueue::InitStatics();
 }
 
 NS_IMPL_ISUPPORTS(MediaMemoryTracker, nsIMemoryReporter)
diff --git a/dom/media/MediaTaskQueue.cpp b/dom/media/MediaTaskQueue.cpp
index 25374b9921..24baba1503 100644
--- a/dom/media/MediaTaskQueue.cpp
+++ b/dom/media/MediaTaskQueue.cpp
@@ -10,6 +10,16 @@
 
 namespace mozilla {
 
+ThreadLocal<MediaTaskQueue*> MediaTaskQueue::sCurrentQueueTLS;
+
+/* static */ void
+MediaTaskQueue::InitStatics()
+{
+  if (!sCurrentQueueTLS.init()) {
+    MOZ_CRASH();
+  }
+}
+
 MediaTaskQueue::MediaTaskQueue(TemporaryRef<SharedThreadPool> aPool)
   : mPool(aPool)
   , mQueueMonitor("MediaTaskQueue::Queue")
@@ -196,7 +206,9 @@ bool
 MediaTaskQueue::IsCurrentThreadIn()
 {
   MonitorAutoLock mon(mQueueMonitor);
-  return NS_GetCurrentThread() == mRunningThread;
+  bool in = NS_GetCurrentThread() == mRunningThread;
+  MOZ_ASSERT_IF(in, GetCurrentQueue() == this);
+  return in;
 }
 
 nsresult
@@ -223,7 +235,10 @@ MediaTaskQueue::Runner::Run()
   // fences enforced. This means that if the object we're calling wasn't
   // designed to be threadsafe, it will be, provided we're only calling it
   // in this task queue.
-  event->Run();
+  {
+    AutoTaskGuard g(mQueue);
+    event->Run();
+  }
 
   // Drop the reference to event. The event will hold a reference to the
   // object it's calling, and we don't want to keep it alive, it may be
diff --git a/dom/media/MediaTaskQueue.h b/dom/media/MediaTaskQueue.h
index efa34a621d..f82a928ec2 100644
--- a/dom/media/MediaTaskQueue.h
+++ b/dom/media/MediaTaskQueue.h
@@ -10,6 +10,7 @@
 #include <queue>
 #include "mozilla/RefPtr.h"
 #include "mozilla/Monitor.h"
+#include "mozilla/ThreadLocal.h"
 #include "SharedThreadPool.h"
 #include "nsThreadUtils.h"
 #include "MediaPromise.h"
@@ -28,7 +29,14 @@ typedef MediaPromise<bool, bool, false> ShutdownPromise;
 // They may be executed on different threads, and a memory barrier is used
 // to make this threadsafe for objects that aren't already threadsafe.
 class MediaTaskQueue : public AbstractThread {
+  static ThreadLocal<MediaTaskQueue*> sCurrentQueueTLS;
 public:
+  static void InitStatics();
+
+  // Returns the task queue that the caller is currently running in, or null
+  // if the caller is not running in a MediaTaskQueue.
+  static MediaTaskQueue* GetCurrentQueue() { return sCurrentQueueTLS.get(); }
+
   explicit MediaTaskQueue(TemporaryRef<SharedThreadPool> aPool);
 
   nsresult Dispatch(TemporaryRef<nsIRunnable> aRunnable);
@@ -103,6 +111,24 @@ protected:
   // is the thread currently running in the task queue.
   RefPtr<nsIThread> mRunningThread;
 
+  // RAII class that gets instantiated for each dispatched task.
+  class AutoTaskGuard
+  {
+  public:
+    explicit AutoTaskGuard(MediaTaskQueue* aQueue)
+    {
+      // NB: We don't hold the lock to aQueue here. Don't do anything that
+      // might require it.
+      MOZ_ASSERT(sCurrentQueueTLS.get() == nullptr);
+      sCurrentQueueTLS.set(aQueue);
+    }
+
+    ~AutoTaskGuard()
+    {
+      sCurrentQueueTLS.set(nullptr);
+    }
+  };
+
   // True if we've dispatched an event to the pool to execute events from
   // the queue.
   bool mIsRunning;