ios: Add a TaskTraits extension and TaskExecutor for WebThreads.

This allows using //base/task/post_task.h for posting tasks to a
WebThread by specifying a WebThread::ID as a task trait.

Related patch for BrowserThreads: https://crrev.com/c/1181364

Bug: 867421, 878356
Cq-Include-Trybots: luci.chromium.try:ios-simulator-cronet;luci.chromium.try:ios-simulator-full-configs
Change-Id: Ic4a539eadc1fda199376c42afa44ba95e39e7041
Reviewed-on: https://chromium-review.googlesource.com/1188283
Commit-Queue: Eric Seckler <eseckler@chromium.org>
Reviewed-by: Kurt Horimoto <kkhorimoto@chromium.org>
Reviewed-by: Sylvain Defresne <sdefresne@chromium.org>
Reviewed-by: François Doray <fdoray@chromium.org>
Cr-Commit-Position: refs/heads/master@{#587089}

ios/web/BUILD.gn

@@ -212,6 +212,7 @@ source_set("ios_web_general_unittests") {
     "url_scheme_util_unittest.mm",
     "url_util_unittest.cc",
     "web_client_unittest.mm",
+    "web_thread_unittest.cc",
   ]
 }
 

ios/web/app/web_main_runner.mm

@@ -11,6 +11,7 @@
 #include "ios/web/public/global_state/ios_global_state.h"
 #include "ios/web/public/url_schemes.h"
 #import "ios/web/public/web_client.h"
+#include "ios/web/web_thread_impl.h"
 #include "mojo/core/embedder/embedder.h"
 #include "ui/base/ui_base_paths.h"
 
@@ -46,6 +47,7 @@ class WebMainRunnerImpl : public WebMainRunner {
     create_params.argc = params.argc;
     create_params.argv = params.argv;
     ios_global_state::Create(create_params);
+    web::WebThreadImpl::CreateTaskExecutor();
 
     if (delegate_) {
       delegate_->BasicStartupComplete();

ios/web/public/BUILD.gn

@@ -96,6 +96,8 @@ source_set("public") {
     "web_state/web_state_policy_decider.h",
     "web_state/web_state_policy_decider_bridge.h",
     "web_state/web_state_user_data.h",
+    "web_task_traits.cc",
+    "web_task_traits.h",
     "web_thread.h",
     "web_thread_delegate.h",
     "web_ui_ios_data_source.h",

ios/web/public/web_task_traits.cc

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "ios/web/public/web_task_traits.h"
+
+namespace web {
+
+// static
+constexpr uint8_t WebTaskTraitsExtension::kExtensionId;
+
+}  // namespace web

ios/web/public/web_task_traits.h

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef IOS_WEB_PUBLIC_WEB_TASK_TRAITS_H_
+#define IOS_WEB_PUBLIC_WEB_TASK_TRAITS_H_
+
+#include "base/task/task_traits.h"
+#include "base/task/task_traits_extension.h"
+#include "ios/web/public/web_thread.h"
+
+namespace web {
+
+// TaskTraits for running tasks on a WebThread.
+//
+// These traits enable the use of the //base/task/post_task.h APIs to post tasks
+// to a WebThread.
+//
+// To post a task to the UI thread (analogous for IO thread):
+//     base::PostTaskWithTraits(FROM_HERE, {WebThread::UI}, task);
+//
+// To obtain a TaskRunner for the UI thread (analogous for the IO thread):
+//     base::CreateSingleThreadTaskRunnerWithTraits({WebThread::UI});
+//
+// See //base/task/post_task.h for more detailed documentation.
+//
+// Posting to a WebThread must only be done after it was initialized (ref.
+// WebMainLoop::CreateThreads() phase).
+class WebTaskTraitsExtension {
+ public:
+  static constexpr uint8_t kExtensionId =
+      base::TaskTraitsExtensionStorage::kFirstEmbedderExtensionId;
+
+  struct ValidTrait {
+    ValidTrait(WebThread::ID) {}
+  };
+
+  template <
+      class... ArgTypes,
+      class CheckArgumentsAreValid = std::enable_if_t<
+          base::trait_helpers::AreValidTraits<ValidTrait, ArgTypes...>::value>>
+  constexpr WebTaskTraitsExtension(ArgTypes... args)
+      : web_thread_(base::trait_helpers::GetValueFromArgList(
+            base::trait_helpers::RequiredEnumArgGetter<WebThread::ID>(),
+            args...)) {}
+
+  constexpr base::TaskTraitsExtensionStorage Serialize() const {
+    static_assert(4 == sizeof(WebTaskTraitsExtension),
+                  "Update Serialize() and Parse() when changing "
+                  "WebTaskTraitsExtension");
+    return {kExtensionId, {static_cast<uint8_t>(web_thread_)}};
+  }
+
+  static const WebTaskTraitsExtension Parse(
+      const base::TaskTraitsExtensionStorage& extension) {
+    return WebTaskTraitsExtension(
+        static_cast<WebThread::ID>(extension.data[0]));
+  }
+
+  constexpr WebThread::ID web_thread() const { return web_thread_; }
+
+ private:
+  WebThread::ID web_thread_;
+};
+
+template <class... ArgTypes,
+          class = std::enable_if_t<base::trait_helpers::AreValidTraits<
+              WebTaskTraitsExtension::ValidTrait,
+              ArgTypes...>::value>>
+constexpr base::TaskTraitsExtensionStorage MakeTaskTraitsExtension(
+    ArgTypes&&... args) {
+  return WebTaskTraitsExtension(std::forward<ArgTypes>(args)...).Serialize();
+}
+
+}  // namespace web
+
+#endif  // IOS_WEB_PUBLIC_WEB_TASK_TRAITS_H_

ios/web/public/web_thread.h

@@ -42,20 +42,12 @@ class WebThreadDelegate;
 // one IO thread for the entire process, and various pieces of code find it
 // useful to retrieve a pointer to the IO thread's message loop.
 //
-// Invoke a task by thread ID:
+// See web_task_traits.h for posting Tasks to a WebThread.
+// TODO(crbug.com/878356): Replace uses with base::PostTaskWithTraits.
 //
-//   WebThread::PostTask(WebThread::IO, FROM_HERE, task);
-//
-// The return value is false if the task couldn't be posted because the target
-// thread doesn't exist.  If this could lead to data loss, you need to check the
-// result and restructure the code to ensure it doesn't occur.
-//
-// This class automatically handles the lifetime of different threads.
-// It's always safe to call PostTask on any thread.  If it's not yet created,
-// the task is deleted.  There are no race conditions.  If the thread that the
-// task is posted to is guaranteed to outlive the current thread, then no locks
-// are used.  You should never need to cache pointers to MessageLoops, since
-// they're not thread safe.
+// This class automatically handles the lifetime of different threads. You
+// should never need to cache pointers to MessageLoops, since they're not thread
+// safe.
 class WebThread {
  public:
   // An enumeration of the well-known threads.
@@ -69,10 +61,8 @@ class WebThread {
     // Blocking IO should happen in TaskScheduler.
     IO,
 
-    // NOTE: do not add new threads here that are only used by a small number of
-    // files. Instead you should just use a Thread class and pass its
-    // SingleThreadTaskRunner around. Named threads there are only for threads
-    // that are used in many places.
+    // NOTE: do not add new threads here. Instead you should just use
+    // base::Create*TaskRunnerWithTraits to run tasks on the TaskScheduler.
 
     // This identifier does not represent a thread.  Instead it counts the
     // number of well-known threads.  Insert new well-known threads before this
@@ -80,6 +70,9 @@ class WebThread {
     ID_COUNT
   };
 
+  // DEPRECATED: Please use the API described in web_task_traits.h instead.
+  // TODO(crbug.com/878356): Replace uses with base::PostTaskWithTraits.
+  //
   // These are the same methods as in message_loop.h, but are guaranteed to
   // either get posted to the MessageLoop if it's still alive, or be deleted
   // otherwise.
@@ -135,6 +128,10 @@ class WebThread {
   // sets identifier to its ID.
   static bool GetCurrentThreadIdentifier(ID* identifier) WARN_UNUSED_RESULT;
 
+  // DEPRECATED: Please use the API described in web_task_traits.h instead.
+  // TODO(crbug.com/878356): Replace uses with
+  // base::Create*TaskRunnerWithTraits.
+  //
   // Callers can hold on to a refcounted SingleThreadTaskRunner beyond the
   // lifetime of the thread.
   static scoped_refptr<base::SingleThreadTaskRunner> GetTaskRunnerForThread(

ios/web/test/test_web_thread_bundle.cc

@@ -46,9 +46,13 @@ TestWebThreadBundle::~TestWebThreadBundle() {
   base::RunLoop().RunUntilIdle();
 
   scoped_task_environment_.reset();
+
+  WebThreadImpl::ResetTaskExecutorForTesting();
 }
 
 void TestWebThreadBundle::Init(int options) {
+  WebThreadImpl::CreateTaskExecutor();
+
   scoped_task_environment_ =
       std::make_unique<base::test::ScopedTaskEnvironment>(
           options & TestWebThreadBundle::IO_MAINLOOP

ios/web/web_thread_impl.cc

@@ -15,7 +15,9 @@
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "base/single_thread_task_runner.h"
+#include "base/task/task_executor.h"
 #include "base/threading/thread_restrictions.h"
+#include "ios/web/public/web_task_traits.h"
 #include "ios/web/public/web_thread_delegate.h"
 #include "net/url_request/url_fetcher.h"
 
@@ -107,6 +109,97 @@ struct WebThreadGlobals {
 base::LazyInstance<WebThreadGlobals>::Leaky g_globals =
     LAZY_INSTANCE_INITIALIZER;
 
+bool PostTaskHelper(WebThread::ID identifier,
+                    const base::Location& from_here,
+                    base::OnceClosure task,
+                    base::TimeDelta delay,
+                    bool nestable) {
+  DCHECK(identifier >= 0 && identifier < WebThread::ID_COUNT);
+  // Optimization: to avoid unnecessary locks, we listed the ID enumeration in
+  // order of lifetime.  So no need to lock if we know that the target thread
+  // outlives current thread.
+  // Note: since the array is so small, ok to loop instead of creating a map,
+  // which would require a lock because std::map isn't thread safe, defeating
+  // the whole purpose of this optimization.
+  WebThread::ID current_thread = WebThread::ID_COUNT;
+  bool target_thread_outlives_current =
+      WebThread::GetCurrentThreadIdentifier(&current_thread) &&
+      current_thread >= identifier;
+
+  WebThreadGlobals& globals = g_globals.Get();
+  if (!target_thread_outlives_current)
+    globals.lock.Acquire();
+
+  base::MessageLoop* message_loop =
+      globals.threads[identifier] ? globals.threads[identifier]->message_loop()
+                                  : nullptr;
+  if (message_loop) {
+    if (nestable) {
+      message_loop->task_runner()->PostDelayedTask(from_here, std::move(task),
+                                                   delay);
+    } else {
+      message_loop->task_runner()->PostNonNestableDelayedTask(
+          from_here, std::move(task), delay);
+    }
+  }
+
+  if (!target_thread_outlives_current)
+    globals.lock.Release();
+
+  return !!message_loop;
+}
+
+class WebThreadTaskExecutor : public base::TaskExecutor {
+ public:
+  WebThreadTaskExecutor() {}
+  ~WebThreadTaskExecutor() override {}
+
+  // base::TaskExecutor implementation.
+  bool PostDelayedTaskWithTraits(const base::Location& from_here,
+                                 const base::TaskTraits& traits,
+                                 base::OnceClosure task,
+                                 base::TimeDelta delay) override {
+    return PostTaskHelper(GetWebThreadIdentifier(traits), from_here,
+                          std::move(task), delay, true);
+  }
+
+  scoped_refptr<base::TaskRunner> CreateTaskRunnerWithTraits(
+      const base::TaskTraits& traits) override {
+    return GetTaskRunnerForThread(GetWebThreadIdentifier(traits));
+  }
+
+  scoped_refptr<base::SequencedTaskRunner> CreateSequencedTaskRunnerWithTraits(
+      const base::TaskTraits& traits) override {
+    return GetTaskRunnerForThread(GetWebThreadIdentifier(traits));
+  }
+
+  scoped_refptr<base::SingleThreadTaskRunner>
+  CreateSingleThreadTaskRunnerWithTraits(
+      const base::TaskTraits& traits,
+      base::SingleThreadTaskRunnerThreadMode thread_mode) override {
+    // It's not possible to request DEDICATED access to a WebThread.
+    DCHECK_EQ(thread_mode, base::SingleThreadTaskRunnerThreadMode::SHARED);
+    return GetTaskRunnerForThread(GetWebThreadIdentifier(traits));
+  }
+
+ private:
+  WebThread::ID GetWebThreadIdentifier(const base::TaskTraits& traits) {
+    DCHECK_EQ(traits.extension_id(), WebTaskTraitsExtension::kExtensionId);
+    WebThread::ID id =
+        traits.GetExtension<WebTaskTraitsExtension>().web_thread();
+    DCHECK_LT(id, WebThread::ID_COUNT);
+    return id;
+  }
+
+  scoped_refptr<base::SingleThreadTaskRunner> GetTaskRunnerForThread(
+      WebThread::ID identifier) {
+    return g_task_runners.Get().task_runners[identifier];
+  }
+};
+
+// |g_web_thread_task_executor| is intentionally leaked on shutdown.
+WebThreadTaskExecutor* g_web_thread_task_executor = nullptr;
+
 }  // namespace
 
 WebThreadImpl::WebThreadImpl(ID identifier)
@@ -223,47 +316,6 @@ WebThreadImpl::~WebThreadImpl() {
 #endif
 }
 
-// static
-bool WebThreadImpl::PostTaskHelper(WebThread::ID identifier,
-                                   const base::Location& from_here,
-                                   base::OnceClosure task,
-                                   base::TimeDelta delay,
-                                   bool nestable) {
-  DCHECK(identifier >= 0 && identifier < ID_COUNT);
-  // Optimization: to avoid unnecessary locks, we listed the ID enumeration in
-  // order of lifetime.  So no need to lock if we know that the target thread
-  // outlives current thread.
-  // Note: since the array is so small, ok to loop instead of creating a map,
-  // which would require a lock because std::map isn't thread safe, defeating
-  // the whole purpose of this optimization.
-  WebThread::ID current_thread = ID_COUNT;
-  bool target_thread_outlives_current =
-      GetCurrentThreadIdentifier(&current_thread) &&
-      current_thread >= identifier;
-
-  WebThreadGlobals& globals = g_globals.Get();
-  if (!target_thread_outlives_current)
-    globals.lock.Acquire();
-
-  base::MessageLoop* message_loop =
-      globals.threads[identifier] ? globals.threads[identifier]->message_loop()
-                                  : nullptr;
-  if (message_loop) {
-    if (nestable) {
-      message_loop->task_runner()->PostDelayedTask(from_here, std::move(task),
-                                                   delay);
-    } else {
-      message_loop->task_runner()->PostNonNestableDelayedTask(
-          from_here, std::move(task), delay);
-    }
-  }
-
-  if (!target_thread_outlives_current)
-    globals.lock.Release();
-
-  return !!message_loop;
-}
-
 // static
 bool WebThread::IsThreadInitialized(ID identifier) {
   if (!g_globals.IsCreated())
@@ -303,8 +355,8 @@ std::string WebThread::GetDCheckCurrentlyOnErrorMessage(ID expected) {
 bool WebThread::PostTask(ID identifier,
                          const base::Location& from_here,
                          base::OnceClosure task) {
-  return WebThreadImpl::PostTaskHelper(identifier, from_here, std::move(task),
-                                       base::TimeDelta(), true);
+  return PostTaskHelper(identifier, from_here, std::move(task),
+                        base::TimeDelta(), true);
 }
 
 // static
@@ -312,16 +364,15 @@ bool WebThread::PostDelayedTask(ID identifier,
                                 const base::Location& from_here,
                                 base::OnceClosure task,
                                 base::TimeDelta delay) {
-  return WebThreadImpl::PostTaskHelper(identifier, from_here, std::move(task),
-                                       delay, true);
+  return PostTaskHelper(identifier, from_here, std::move(task), delay, true);
 }
 
 // static
 bool WebThread::PostNonNestableTask(ID identifier,
                                     const base::Location& from_here,
                                     base::OnceClosure task) {
-  return WebThreadImpl::PostTaskHelper(identifier, from_here, std::move(task),
-                                       base::TimeDelta(), false);
+  return PostTaskHelper(identifier, from_here, std::move(task),
+                        base::TimeDelta(), false);
 }
 
 // static
@@ -329,8 +380,7 @@ bool WebThread::PostNonNestableDelayedTask(ID identifier,
                                            const base::Location& from_here,
                                            base::OnceClosure task,
                                            base::TimeDelta delay) {
-  return WebThreadImpl::PostTaskHelper(identifier, from_here, std::move(task),
-                                       delay, false);
+  return PostTaskHelper(identifier, from_here, std::move(task), delay, false);
 }
 
 // static
@@ -379,4 +429,20 @@ void WebThread::SetDelegate(ID identifier, WebThreadDelegate* delegate) {
   DCHECK(!delegate || !old_pointer);
 }
 
+// static
+void WebThreadImpl::CreateTaskExecutor() {
+  DCHECK(!g_web_thread_task_executor);
+  g_web_thread_task_executor = new WebThreadTaskExecutor();
+  base::RegisterTaskExecutor(WebTaskTraitsExtension::kExtensionId,
+                             g_web_thread_task_executor);
+}
+
+// static
+void WebThreadImpl::ResetTaskExecutorForTesting() {
+  DCHECK(g_web_thread_task_executor);
+  base::UnregisterTaskExecutorForTesting(WebTaskTraitsExtension::kExtensionId);
+  delete g_web_thread_task_executor;
+  g_web_thread_task_executor = nullptr;
+}
+
 }  // namespace web

ios/web/web_thread_impl.h

@@ -25,6 +25,13 @@ class WebThreadImpl : public WebThread, public base::Thread {
   WebThreadImpl(WebThread::ID identifier, base::MessageLoop* message_loop);
   ~WebThreadImpl() override;
 
+  // Creates and registers a TaskExecutor that facilitates posting tasks to a
+  // WebThread via //base/task/post_task.h.
+  static void CreateTaskExecutor();
+
+  // Unregister and delete the TaskExecutor after a test.
+  static void ResetTaskExecutorForTesting();
+
  protected:
   void Init() override;
   void Run(base::RunLoop* run_loop) override;
@@ -41,12 +48,6 @@ class WebThreadImpl : public WebThread, public base::Thread {
   void UIThreadRun(base::RunLoop* run_loop);
   void IOThreadRun(base::RunLoop* run_loop);
 
-  static bool PostTaskHelper(WebThread::ID identifier,
-                             const base::Location& from_here,
-                             base::OnceClosure task,
-                             base::TimeDelta delay,
-                             bool nestable);
-
   // Common initialization code for the constructors.
   void Initialize();
 

ios/web/web_thread_unittest.cc

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "ios/web/public/web_thread.h"
+#include "base/task/post_task.h"
+#include "ios/web/public/test/test_web_thread_bundle.h"
+#include "ios/web/public/web_task_traits.h"
+#include "ios/web/web_thread_impl.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "testing/platform_test.h"
+
+namespace web {
+
+class WebThreadTest : public PlatformTest {
+ protected:
+  static void BasicFunction(base::OnceClosure continuation,
+                            WebThread::ID target) {
+    EXPECT_TRUE(WebThread::CurrentlyOn(target));
+    std::move(continuation).Run();
+  }
+
+  web::TestWebThreadBundle web_thread_bundle_;
+};
+
+TEST_F(WebThreadTest, PostTaskWithTraits) {
+  base::RunLoop run_loop;
+  EXPECT_TRUE(base::PostTaskWithTraits(
+      FROM_HERE, {WebThread::IO},
+      base::BindOnce(&BasicFunction, run_loop.QuitWhenIdleClosure(),
+                     WebThread::IO)));
+  run_loop.Run();
+}
+
+TEST_F(WebThreadTest, PostTaskViaTaskRunnerWithTraits) {
+  scoped_refptr<base::TaskRunner> task_runner =
+      base::CreateTaskRunnerWithTraits({WebThread::IO});
+  base::RunLoop run_loop;
+  EXPECT_TRUE(task_runner->PostTask(
+      FROM_HERE, base::BindOnce(&BasicFunction, run_loop.QuitWhenIdleClosure(),
+                                WebThread::IO)));
+  run_loop.Run();
+}
+
+TEST_F(WebThreadTest, PostTaskViaSequencedTaskRunnerWithTraits) {
+  scoped_refptr<base::SequencedTaskRunner> task_runner =
+      base::CreateSequencedTaskRunnerWithTraits({WebThread::IO});
+  base::RunLoop run_loop;
+  EXPECT_TRUE(task_runner->PostTask(
+      FROM_HERE, base::BindOnce(&BasicFunction, run_loop.QuitWhenIdleClosure(),
+                                WebThread::IO)));
+  run_loop.Run();
+}
+
+TEST_F(WebThreadTest, PostTaskViaSingleThreadTaskRunnerWithTraits) {
+  scoped_refptr<base::SingleThreadTaskRunner> task_runner =
+      base::CreateSingleThreadTaskRunnerWithTraits({WebThread::IO});
+  base::RunLoop run_loop;
+  EXPECT_TRUE(task_runner->PostTask(
+      FROM_HERE, base::BindOnce(&BasicFunction, run_loop.QuitWhenIdleClosure(),
+                                WebThread::IO)));
+  run_loop.Run();
+}
+
+}  // namespace web