(WIP) Worker: Merge ParentFrameTaskRunners into TaskRunnerHelper

third_party/WebKit/Source/core/dom/TaskRunnerHelper.cpp

Index: third_party/WebKit/Source/core/dom/TaskRunnerHelper.cpp
diff --git a/third_party/WebKit/Source/core/dom/TaskRunnerHelper.cpp b/third_party/WebKit/Source/core/dom/TaskRunnerHelper.cpp
index 8e5d0af5b9246e96d0dfae688dd9a52f04365dc5..2938de43d71e777f02bd39613e56de3004a4b3a6 100644
--- a/third_party/WebKit/Source/core/dom/TaskRunnerHelper.cpp
+++ b/third_party/WebKit/Source/core/dom/TaskRunnerHelper.cpp
@@ -5,7 +5,10 @@
 #include "core/dom/TaskRunnerHelper.h"
 
 #include "core/dom/Document.h"
+#include "core/dom/ExecutionContext.h"
 #include "core/frame/LocalFrame.h"
+#include "core/workers/WorkerOrWorkletGlobalScope.h"
+#include "core/workers/WorkerThread.h"
 #include "platform/WebFrameScheduler.h"
 #include "platform/WebTaskRunner.h"
 #include "public/platform/Platform.h"
@@ -13,7 +16,20 @@
 
 namespace blink {
 
-RefPtr<WebTaskRunner> TaskRunnerHelper::get(TaskType type, LocalFrame* frame) {
+using TaskRunnersByContext =
+    HeapHashMap<Member<ExecutionContext>,
+                CrossThreadPersistent<FrameTaskRunnersHolder>>;
+
+namespace {
+
+ThreadSpecific<TaskRunnersByContext>& getTaskRunnerMap() {
+  DEFINE_THREAD_SAFE_STATIC_LOCAL(ThreadSpecific<TaskRunnersByContext>,
+                                  taskRunners,
+                                  new ThreadSpecific<TaskRunnersByContext>);
+  return taskRunners;
+}
+
+RefPtr<WebTaskRunner> getInternal(TaskType type, LocalFrame* frame) {
   // TODO(haraken): Optimize the mapping from TaskTypes to task runners.
   switch (type) {
     case TaskType::Timer:
@@ -50,26 +66,96 @@ RefPtr<WebTaskRunner> TaskRunnerHelper::get(TaskType type, LocalFrame* frame) {
     case TaskType::Unthrottled:
       return frame ? frame->frameScheduler()->unthrottledTaskRunner()
                    : Platform::current()->currentThread()->getWebTaskRunner();
+    case TaskType::NumberOfTaskTypes:
+      break;
   }
   NOTREACHED();
   return nullptr;
 }
 
-RefPtr<WebTaskRunner> TaskRunnerHelper::get(TaskType type, Document* document) {
+}  // namespace
+
+FrameTaskRunnersHolder::FrameTaskRunnersHolder(LocalFrame* frame)
+    : ContextLifecycleObserver(frame ? frame->document() : nullptr) {
+  DCHECK(isMainThread());
+  if (frame && frame->document())
+    DCHECK(frame->document()->isContextThread());
+
+  for (unsigned i = 0; i != static_cast<unsigned>(TaskType::NumberOfTaskTypes);
+       ++i) {
+    TaskType type = static_cast<TaskType>(i);
+    m_taskRunners.insert(type, getInternal(type, frame));
+  }
+}
+
+RefPtr<WebTaskRunner> FrameTaskRunnersHolder::get(TaskType type) {
+  MutexLocker lock(m_taskRunnersMutex);
+  return m_taskRunners.at(type);
+}
+
+DEFINE_TRACE(FrameTaskRunnersHolder) {
+  ContextLifecycleObserver::trace(visitor);
+}
+
+void FrameTaskRunnersHolder::contextDestroyed(ExecutionContext*) {
+  MutexLocker lock(m_taskRunnersMutex);
+  for (auto& entry : m_taskRunners)
+    entry.value = Platform::current()->currentThread()->getWebTaskRunner();
+}
+
+// ----- FrameTaskRunnerHelper -----
+
+void FrameTaskRunnerHelper::setTaskRunners(
+    ExecutionContext* executionContext,
+    FrameTaskRunnersHolder* taskRunners) {
+  getTaskRunnerMap()->insert(executionContext, taskRunners);
+}
+
+RefPtr<WebTaskRunner> FrameTaskRunnerHelper::get(TaskType type,
+                                                 LocalFrame* frame) {
+  DCHECK(isMainThread());
+  return getInternal(type, frame);
+}
+
+RefPtr<WebTaskRunner> FrameTaskRunnerHelper::get(TaskType type,
+                                                 Document* document) {
+  DCHECK(isMainThread());
   return get(type, document ? document->frame() : nullptr);
 }
 
-RefPtr<WebTaskRunner> TaskRunnerHelper::get(
+RefPtr<WebTaskRunner> FrameTaskRunnerHelper::get(
     TaskType type,
     ExecutionContext* executionContext) {
-  return get(type, executionContext && executionContext->isDocument()
-                       ? static_cast<Document*>(executionContext)
-                       : nullptr);
+  if (!executionContext)
+    return get(type, static_cast<LocalFrame*>(nullptr));
+  if (executionContext->isDocument()) {
+    DCHECK(executionContext->isContextThread());
+    return get(type, static_cast<Document*>(executionContext));
+  }
+  if (executionContext->isWorkerOrWorkletGlobalScope()) {
+    DCHECK(executionContext->isContextThread());
+    return get(type,
+               static_cast<WorkerOrWorkletGlobalScope*>(executionContext));
+  }
+  return get(type, static_cast<LocalFrame*>(nullptr));
 }
 
-RefPtr<WebTaskRunner> TaskRunnerHelper::get(TaskType type,
-                                            ScriptState* scriptState) {
+RefPtr<WebTaskRunner> FrameTaskRunnerHelper::get(TaskType type,
+                                                 ScriptState* scriptState) {
+  DCHECK(isMainThread());
   return get(type, scriptState ? scriptState->getExecutionContext() : nullptr);
 }
 
+RefPtr<WebTaskRunner> FrameTaskRunnerHelper::get(TaskType type,
+                                                 WorkerThread* workerThread) {
+  return get(type, workerThread->globalScope());
+}
+
+RefPtr<WebTaskRunner> FrameTaskRunnerHelper::get(
+    TaskType type,
+    WorkerOrWorkletGlobalScope* globalScope) {
+  DCHECK(globalScope->isContextThread());
+  return getTaskRunnerMap()->at(globalScope)->get(type);
+}
+
 }  // namespace blink