Make ui::Compositor BeginFrame subscription robust

ui/compositor/compositor_unittest.cc

Index: ui/compositor/compositor_unittest.cc
diff --git a/ui/compositor/compositor_unittest.cc b/ui/compositor/compositor_unittest.cc
index a53bdc2ab0f2b5d1095bde808917bb7947ca3c72..4d355aed7b280ebaa1db0706a73d476bbf76cfd0 100644
--- a/ui/compositor/compositor_unittest.cc
+++ b/ui/compositor/compositor_unittest.cc
@@ -19,6 +19,10 @@ using testing::_;
 namespace ui {
 namespace {
 
+ACTION_P2(RemoveObserver, compositor, observer) {
+  compositor->RemoveBeginFrameObserver(observer);
+}
+
 class MockCompositorBeginFrameObserver : public CompositorBeginFrameObserver {
  public:
   MOCK_METHOD1(OnSendBeginFrame, void(const cc::BeginFrameArgs&));
@@ -127,6 +131,57 @@ TEST_F(CompositorTest, AddAndRemoveBeginFrameObserver) {
   compositor()->RemoveBeginFrameObserver(&test_observer2);
 }
 
+TEST_F(CompositorTest, RemoveBeginFrameObserverWhileSendingBeginFrame) {
+  cc::BeginFrameArgs args = cc::CreateBeginFrameArgsForTesting(
+      BEGINFRAME_FROM_HERE, base::TimeTicks::FromInternalValue(33));
+
+  // Simulate to trigger new BeginFrame by using |args|.
+  compositor()->SendBeginFramesToChildren(args);
+
+  // When |missed_begin_frame_args_| is sent, its type is set to MISSED.
+  cc::BeginFrameArgs expected_args(args);
+  expected_args.type = cc::BeginFrameArgs::MISSED;
+
+  // Add both observers, and simulate removal of |test_observer2| during
+  // BeginFrame dispatch (implicitly triggered when the observer is added).
+  MockCompositorBeginFrameObserver test_observer;
+  MockCompositorBeginFrameObserver test_observer2;
+  EXPECT_CALL(test_observer, OnSendBeginFrame(expected_args));
+  EXPECT_CALL(test_observer2, OnSendBeginFrame(expected_args))
+      .WillOnce(RemoveObserver(compositor(), &test_observer2));
+
+  // When a new observer is added, Compositor immediately calls OnSendBeginFrame
+  // with |missed_begin_frame_args_|.
+  compositor()->AddBeginFrameObserver(&test_observer);
+  compositor()->AddBeginFrameObserver(&test_observer2);
+  Mock::VerifyAndClearExpectations(&test_observer);
+  Mock::VerifyAndClearExpectations(&test_observer2);
+
+  // |test_observer2| was removed during the previous implicit BeginFrame
+  // dispatch, and should not get the new frame.
+  expected_args.type = cc::BeginFrameArgs::NORMAL;
+  EXPECT_CALL(test_observer, OnSendBeginFrame(expected_args));
+  EXPECT_CALL(test_observer2, OnSendBeginFrame(expected_args)).Times(0);
+  compositor()->SendBeginFramesToChildren(args);
+  Mock::VerifyAndClearExpectations(&test_observer);
+  Mock::VerifyAndClearExpectations(&test_observer2);
+
+  // Now remove |test_observer| during explicit BeginFrame dispatch.
+  EXPECT_CALL(test_observer, OnSendBeginFrame(expected_args))
+      .WillOnce(RemoveObserver(compositor(), &test_observer));
+  EXPECT_CALL(test_observer2, OnSendBeginFrame(expected_args)).Times(0);
+  compositor()->SendBeginFramesToChildren(args);
+  Mock::VerifyAndClearExpectations(&test_observer);
+  Mock::VerifyAndClearExpectations(&test_observer2);
+
+  // No observers should get the new frame.
+  EXPECT_CALL(test_observer, OnSendBeginFrame(expected_args)).Times(0);
+  EXPECT_CALL(test_observer2, OnSendBeginFrame(expected_args)).Times(0);
+  compositor()->SendBeginFramesToChildren(args);
+  Mock::VerifyAndClearExpectations(&test_observer);
+  Mock::VerifyAndClearExpectations(&test_observer2);
+}
+
 TEST_F(CompositorTest, ReleaseWidgetWithOutputSurfaceNeverCreated) {
   compositor()->SetVisible(false);
   EXPECT_EQ(gfx::kNullAcceleratedWidget,