[SSM] Increment Work Id after idle work complete

- Increment work item id after OnIdle() callback finishes.
- DoIdleWork should only appear immediately after a DoWork call
  so no need to increment work id before doing idle work.
- Add test cases to ThreadControllerWithMessagePumpImpl covering
  all cases where work id is being counted.

Bug: 40256995
Change-Id: I89af6dd43fb4e4e75405c40625b2203bfa29fbb4
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5455554
Reviewed-by: Deep Roy <dproy@chromium.org>
Reviewed-by: Francois Pierre Doray <fdoray@chromium.org>
Commit-Queue: Rodney Ding <rodneyding@google.com>
Cr-Commit-Position: refs/heads/main@{#1289550}

base/task/sequence_manager/thread_controller_with_message_pump_impl.cc

@@ -573,13 +573,14 @@ bool ThreadControllerWithMessagePumpImpl::DoIdleWork() {
 #endif  // BUILDFLAG(IS_WIN)
 
   if (main_thread_only().task_source->OnIdle()) {
+    work_id_provider_->IncrementWorkId();
     // The OnIdle() callback resulted in more immediate work, so schedule a
     // DoWork callback. For some message pumps returning true from here is
     // sufficient to do that but not on mac.
     pump_->ScheduleWork();
     return false;
   }
-
+  work_id_provider_->IncrementWorkId();
   // This is mostly redundant with the identical call in BeforeWait (upcoming)
   // but some uninstrumented MessagePump impls don't call BeforeWait so it must
   // also be done here.

base/task/sequence_manager/thread_controller_with_message_pump_impl_unittest.cc

@@ -13,6 +13,7 @@
 #include "base/functional/bind.h"
 #include "base/memory/raw_ptr.h"
 #include "base/memory/scoped_refptr.h"
+#include "base/message_loop/message_pump.h"
 #include "base/task/sequence_manager/task_queue.h"
 #include "base/task/sequence_manager/thread_controller_power_monitor.h"
 #include "base/task/single_thread_task_runner.h"
@@ -2094,4 +2095,140 @@ TEST_F(ThreadControllerWithMessagePumpTest, MessagePumpPhasesWithQueuingTime) {
   RunLoop().Run();
 }
 
+TEST_F(ThreadControllerWithMessagePumpNoBatchesTest,
+       WorkIdIncrementedForEveryWorkItem) {
+  SingleThreadTaskRunner::CurrentDefaultHandle handle(
+      MakeRefCounted<FakeTaskRunner>());
+  WorkIdProvider* work_id_provider = WorkIdProvider::GetForCurrentThread();
+
+  work_id_provider->SetCurrentWorkIdForTesting(0u);
+
+  EXPECT_CALL(*message_pump_, Run(_))
+      .WillOnce([&](MessagePump::Delegate* delegate) {
+        // Each task will increment work id by 2, once on begin work and another
+        // on end work.
+        delegate->DoWork();
+        EXPECT_EQ(work_id_provider->GetWorkId(), 2u);
+        delegate->DoWork();
+        EXPECT_EQ(work_id_provider->GetWorkId(), 4u);
+      });
+
+  for (int task_count = 0; task_count < 2; task_count++) {
+    task_source_.AddTask(FROM_HERE, DoNothing(), TimeTicks());
+  }
+
+  RunLoop run_loop;
+  run_loop.Run();
+  testing::Mock::VerifyAndClearExpectations(message_pump_);
+}
+
+TEST_F(ThreadControllerWithMessagePumpTest,
+       WorkIdIncrementedForEveryWorkItemInBatches) {
+  SingleThreadTaskRunner::CurrentDefaultHandle handle(
+      MakeRefCounted<FakeTaskRunner>());
+  thread_controller_.SetWorkBatchSize(2);
+
+  WorkIdProvider* work_id_provider = WorkIdProvider::GetForCurrentThread();
+  work_id_provider->SetCurrentWorkIdForTesting(0u);
+
+  EXPECT_CALL(*message_pump_, Run(_))
+      .WillOnce([&](MessagePump::Delegate* delegate) {
+        delegate->DoWork();
+        // Each task will increment work id by 2, once on begin work and another
+        // on end work.
+        EXPECT_EQ(work_id_provider->GetWorkId(), 4u);
+      });
+
+  for (int task_count = 0; task_count < 2; task_count++) {
+    task_source_.AddTask(FROM_HERE, DoNothing(), TimeTicks());
+  }
+
+  RunLoop run_loop;
+  run_loop.Run();
+  testing::Mock::VerifyAndClearExpectations(message_pump_);
+}
+
+TEST_F(ThreadControllerWithMessagePumpTest, WorkIdIncrementedForIdleWork) {
+  SingleThreadTaskRunner::CurrentDefaultHandle handle(
+      MakeRefCounted<FakeTaskRunner>());
+  WorkIdProvider* work_id_provider = WorkIdProvider::GetForCurrentThread();
+
+  work_id_provider->SetCurrentWorkIdForTesting(0u);
+
+  EXPECT_CALL(*message_pump_, Run(_))
+      .WillOnce([&](MessagePump::Delegate* delegate) {
+        delegate->DoIdleWork();
+        EXPECT_EQ(work_id_provider->GetWorkId(), 1u);
+      });
+
+  task_source_.AddTask(FROM_HERE, DoNothing());
+
+  RunLoop run_loop;
+  run_loop.Run();
+  testing::Mock::VerifyAndClearExpectations(message_pump_);
+}
+
+TEST_F(ThreadControllerWithMessagePumpTest, WorkIdIncrementedScopedDoWorkItem) {
+  SingleThreadTaskRunner::CurrentDefaultHandle handle(
+      MakeRefCounted<FakeTaskRunner>());
+  WorkIdProvider* work_id_provider = WorkIdProvider::GetForCurrentThread();
+
+  work_id_provider->SetCurrentWorkIdForTesting(0u);
+
+  EXPECT_CALL(*message_pump_, Run(_))
+      .WillOnce([&](MessagePump::Delegate* delegate) {
+        MessagePump::Delegate::ScopedDoWorkItem scoped_do_work_item =
+            delegate->BeginWorkItem();
+        // ScopedDoWorkItem will increment work id by 1 on construction and
+        // another on destruction.
+        EXPECT_EQ(work_id_provider->GetWorkId(), 1u);
+      });
+
+  RunLoop run_loop;
+  run_loop.Run();
+  testing::Mock::VerifyAndClearExpectations(message_pump_);
+  // Delegate::Run() itself will increment work id to account for pump overhead.
+  EXPECT_EQ(work_id_provider->GetWorkId(), 3u);
+}
+
+TEST_F(ThreadControllerWithMessagePumpTest,
+       WorkIdIncrementedDelegateBeforeWait) {
+  SingleThreadTaskRunner::CurrentDefaultHandle handle(
+      MakeRefCounted<FakeTaskRunner>());
+  WorkIdProvider* work_id_provider = WorkIdProvider::GetForCurrentThread();
+
+  work_id_provider->SetCurrentWorkIdForTesting(0u);
+
+  EXPECT_CALL(*message_pump_, Run(_))
+      .WillOnce([&](MessagePump::Delegate* delegate) {
+        // Delegate::BeforeWait will increment work id by 1 before waiting for
+        // work.
+        delegate->BeforeWait();
+        EXPECT_EQ(work_id_provider->GetWorkId(), 1u);
+      });
+
+  RunLoop run_loop;
+  run_loop.Run();
+  testing::Mock::VerifyAndClearExpectations(message_pump_);
+}
+
+TEST_F(ThreadControllerWithMessagePumpTest, WorkIdIncrementedDelegateRun) {
+  SingleThreadTaskRunner::CurrentDefaultHandle handle(
+      MakeRefCounted<FakeTaskRunner>());
+  WorkIdProvider* work_id_provider = WorkIdProvider::GetForCurrentThread();
+
+  work_id_provider->SetCurrentWorkIdForTesting(0u);
+
+  EXPECT_CALL(*message_pump_, Run(_))
+      .WillOnce([&](MessagePump::Delegate* delegate) {
+        EXPECT_EQ(work_id_provider->GetWorkId(), 0u);
+      });
+
+  RunLoop run_loop;
+  run_loop.Run();
+  testing::Mock::VerifyAndClearExpectations(message_pump_);
+  // Delegate::Run() itself will increment work id to account for pump overhead.
+  EXPECT_EQ(work_id_provider->GetWorkId(), 1u);
+}
+
 }  // namespace base::sequence_manager::internal