SF: region sampling scheduling optimization

Schedule the region sampling thread more wisely by estimating when
the next invalidate is going to happen. This makes the region sampling
to be scheduled often at times where the main thread is idle, without
eating the budget in an invalidate message.

Bug: 181983990
Test: SF unit tests
Test: observe systrace
Change-Id: I1faca3aa7f882ed7c69e77e6a0877c10d57f0f1c
diff --git a/services/surfaceflinger/RegionSamplingThread.cpp b/services/surfaceflinger/RegionSamplingThread.cpp
index 00090d9..653aca6 100644
--- a/services/surfaceflinger/RegionSamplingThread.cpp
+++ b/services/surfaceflinger/RegionSamplingThread.cpp
@@ -56,16 +56,14 @@
     noWorkNeeded,
     idleTimerWaiting,
     waitForQuietFrame,
-    waitForZeroPhase,
     waitForSamplePhase,
     sample
 };
 
-constexpr auto timeForRegionSampling = 5000000ns;
-constexpr auto maxRegionSamplingSkips = 10;
 constexpr auto defaultRegionSamplingWorkDuration = 3ms;
 constexpr auto defaultRegionSamplingPeriod = 100ms;
 constexpr auto defaultRegionSamplingTimerTimeout = 100ms;
+constexpr auto maxRegionSamplingDelay = 100ms;
 // TODO: (b/127403193) duration to string conversion could probably be constexpr
 template <typename Rep, typename Per>
 inline std::string toNsString(std::chrono::duration<Rep, Per> t) {
@@ -99,97 +97,22 @@
     }
 }
 
-struct SamplingOffsetCallback : VSyncSource::Callback {
-    SamplingOffsetCallback(RegionSamplingThread& samplingThread, Scheduler& scheduler,
-                           std::chrono::nanoseconds targetSamplingWorkDuration)
-          : mRegionSamplingThread(samplingThread),
-            mTargetSamplingWorkDuration(targetSamplingWorkDuration),
-            mVSyncSource(scheduler.makePrimaryDispSyncSource("SamplingThreadDispSyncListener", 0ns,
-                                                             0ns,
-                                                             /*traceVsync=*/false)) {
-        mVSyncSource->setCallback(this);
-    }
-
-    ~SamplingOffsetCallback() { stopVsyncListener(); }
-
-    SamplingOffsetCallback(const SamplingOffsetCallback&) = delete;
-    SamplingOffsetCallback& operator=(const SamplingOffsetCallback&) = delete;
-
-    void startVsyncListener() {
-        std::lock_guard lock(mMutex);
-        if (mVsyncListening) return;
-
-        mPhaseIntervalSetting = Phase::ZERO;
-        mVSyncSource->setVSyncEnabled(true);
-        mVsyncListening = true;
-    }
-
-    void stopVsyncListener() {
-        std::lock_guard lock(mMutex);
-        stopVsyncListenerLocked();
-    }
-
-private:
-    void stopVsyncListenerLocked() /*REQUIRES(mMutex)*/ {
-        if (!mVsyncListening) return;
-
-        mVSyncSource->setVSyncEnabled(false);
-        mVsyncListening = false;
-    }
-
-    void onVSyncEvent(nsecs_t /*when*/, nsecs_t /*expectedVSyncTimestamp*/,
-                      nsecs_t /*deadlineTimestamp*/) final {
-        std::unique_lock<decltype(mMutex)> lock(mMutex);
-
-        if (mPhaseIntervalSetting == Phase::ZERO) {
-            ATRACE_INT(lumaSamplingStepTag, static_cast<int>(samplingStep::waitForSamplePhase));
-            mPhaseIntervalSetting = Phase::SAMPLING;
-            mVSyncSource->setDuration(mTargetSamplingWorkDuration, 0ns);
-            return;
-        }
-
-        if (mPhaseIntervalSetting == Phase::SAMPLING) {
-            mPhaseIntervalSetting = Phase::ZERO;
-            mVSyncSource->setDuration(0ns, 0ns);
-            stopVsyncListenerLocked();
-            lock.unlock();
-            mRegionSamplingThread.notifySamplingOffset();
-            return;
-        }
-    }
-
-    RegionSamplingThread& mRegionSamplingThread;
-    const std::chrono::nanoseconds mTargetSamplingWorkDuration;
-    mutable std::mutex mMutex;
-    enum class Phase {
-        ZERO,
-        SAMPLING
-    } mPhaseIntervalSetting /*GUARDED_BY(mMutex) macro doesnt work with unique_lock?*/
-            = Phase::ZERO;
-    bool mVsyncListening /*GUARDED_BY(mMutex)*/ = false;
-    std::unique_ptr<VSyncSource> mVSyncSource;
-};
-
-RegionSamplingThread::RegionSamplingThread(SurfaceFlinger& flinger, Scheduler& scheduler,
-                                           const TimingTunables& tunables)
+RegionSamplingThread::RegionSamplingThread(SurfaceFlinger& flinger, const TimingTunables& tunables)
       : mFlinger(flinger),
-        mScheduler(scheduler),
         mTunables(tunables),
         mIdleTimer(
                 "RegSampIdle",
                 std::chrono::duration_cast<std::chrono::milliseconds>(
                         mTunables.mSamplingTimerTimeout),
                 [] {}, [this] { checkForStaleLuma(); }),
-        mPhaseCallback(std::make_unique<SamplingOffsetCallback>(*this, mScheduler,
-                                                                tunables.mSamplingDuration)),
-        lastSampleTime(0ns) {
+        mLastSampleTime(0ns) {
     mThread = std::thread([this]() { threadMain(); });
     pthread_setname_np(mThread.native_handle(), "RegionSampling");
     mIdleTimer.start();
 }
 
-RegionSamplingThread::RegionSamplingThread(SurfaceFlinger& flinger, Scheduler& scheduler)
-      : RegionSamplingThread(flinger, scheduler,
+RegionSamplingThread::RegionSamplingThread(SurfaceFlinger& flinger)
+      : RegionSamplingThread(flinger,
                              TimingTunables{defaultRegionSamplingWorkDuration,
                                             defaultRegionSamplingPeriod,
                                             defaultRegionSamplingTimerTimeout}) {}
@@ -224,48 +147,46 @@
 void RegionSamplingThread::checkForStaleLuma() {
     std::lock_guard lock(mThreadControlMutex);
 
-    if (mDiscardedFrames > 0) {
-        ATRACE_INT(lumaSamplingStepTag, static_cast<int>(samplingStep::waitForZeroPhase));
-        mDiscardedFrames = 0;
-        mPhaseCallback->startVsyncListener();
+    if (mSampleRequestTime.has_value()) {
+        ATRACE_INT(lumaSamplingStepTag, static_cast<int>(samplingStep::waitForSamplePhase));
+        mSampleRequestTime.reset();
+        mFlinger.scheduleRegionSamplingThread();
     }
 }
 
-void RegionSamplingThread::notifyNewContent() {
-    doSample();
+void RegionSamplingThread::onCompositionComplete(
+        std::optional<std::chrono::steady_clock::time_point> samplingDeadline) {
+    doSample(samplingDeadline);
 }
 
-void RegionSamplingThread::notifySamplingOffset() {
-    doSample();
-}
-
-void RegionSamplingThread::doSample() {
+void RegionSamplingThread::doSample(
+        std::optional<std::chrono::steady_clock::time_point> samplingDeadline) {
     std::lock_guard lock(mThreadControlMutex);
-    auto now = std::chrono::nanoseconds(systemTime(SYSTEM_TIME_MONOTONIC));
-    if (lastSampleTime + mTunables.mSamplingPeriod > now) {
+    const auto now = std::chrono::steady_clock::now();
+    if (mLastSampleTime + mTunables.mSamplingPeriod > now) {
+        // content changed, but we sampled not too long ago, so we need to sample some time in the
+        // future.
         ATRACE_INT(lumaSamplingStepTag, static_cast<int>(samplingStep::idleTimerWaiting));
-        if (mDiscardedFrames == 0) mDiscardedFrames++;
+        mSampleRequestTime = now;
         return;
     }
-    if (mDiscardedFrames < maxRegionSamplingSkips) {
+    if (!mSampleRequestTime.has_value() || now - *mSampleRequestTime < maxRegionSamplingDelay) {
         // If there is relatively little time left for surfaceflinger
         // until the next vsync deadline, defer this sampling work
         // to a later frame, when hopefully there will be more time.
-        const DisplayStatInfo stats = mScheduler.getDisplayStatInfo(systemTime());
-        if (std::chrono::nanoseconds(stats.vsyncTime) - now < timeForRegionSampling) {
+        if (samplingDeadline.has_value() && now + mTunables.mSamplingDuration > *samplingDeadline) {
             ATRACE_INT(lumaSamplingStepTag, static_cast<int>(samplingStep::waitForQuietFrame));
-            mDiscardedFrames++;
+            mSampleRequestTime = mSampleRequestTime.value_or(now);
             return;
         }
     }
 
     ATRACE_INT(lumaSamplingStepTag, static_cast<int>(samplingStep::sample));
 
-    mDiscardedFrames = 0;
-    lastSampleTime = now;
+    mSampleRequestTime.reset();
+    mLastSampleTime = now;
 
     mIdleTimer.reset();
-    mPhaseCallback->stopVsyncListener();
 
     mSampleRequested = true;
     mCondition.notify_one();
diff --git a/services/surfaceflinger/RegionSamplingThread.h b/services/surfaceflinger/RegionSamplingThread.h
index 86632db..2231853 100644
--- a/services/surfaceflinger/RegionSamplingThread.h
+++ b/services/surfaceflinger/RegionSamplingThread.h
@@ -63,9 +63,8 @@
     struct EnvironmentTimingTunables : TimingTunables {
         EnvironmentTimingTunables();
     };
-    explicit RegionSamplingThread(SurfaceFlinger& flinger, Scheduler& scheduler,
-                                  const TimingTunables& tunables);
-    explicit RegionSamplingThread(SurfaceFlinger& flinger, Scheduler& scheduler);
+    explicit RegionSamplingThread(SurfaceFlinger& flinger, const TimingTunables& tunables);
+    explicit RegionSamplingThread(SurfaceFlinger& flinger);
 
     ~RegionSamplingThread();
 
@@ -76,12 +75,11 @@
     // Remove the listener to stop receiving median luma notifications.
     void removeListener(const sp<IRegionSamplingListener>& listener);
 
-    // Notifies sampling engine that new content is available. This will trigger a sampling
-    // pass at some point in the future.
-    void notifyNewContent();
-
-    // Notifies the sampling engine that it has a good timing window in which to sample.
-    void notifySamplingOffset();
+    // Notifies sampling engine that composition is done and new content is
+    // available, and the deadline for the sampling work on the main thread to
+    // be completed without eating the budget of another frame.
+    void onCompositionComplete(
+            std::optional<std::chrono::steady_clock::time_point> samplingDeadline);
 
 private:
     struct Descriptor {
@@ -99,7 +97,7 @@
             const sp<GraphicBuffer>& buffer, const Point& leftTop,
             const std::vector<RegionSamplingThread::Descriptor>& descriptors, uint32_t orientation);
 
-    void doSample();
+    void doSample(std::optional<std::chrono::steady_clock::time_point> samplingDeadline);
     void binderDied(const wp<IBinder>& who) override;
     void checkForStaleLuma();
 
@@ -107,20 +105,18 @@
     void threadMain();
 
     SurfaceFlinger& mFlinger;
-    Scheduler& mScheduler;
     const TimingTunables mTunables;
     scheduler::OneShotTimer mIdleTimer;
 
-    std::unique_ptr<SamplingOffsetCallback> const mPhaseCallback;
-
     std::thread mThread;
 
     std::mutex mThreadControlMutex;
     std::condition_variable_any mCondition;
     bool mRunning GUARDED_BY(mThreadControlMutex) = true;
     bool mSampleRequested GUARDED_BY(mThreadControlMutex) = false;
-    uint32_t mDiscardedFrames GUARDED_BY(mThreadControlMutex) = 0;
-    std::chrono::nanoseconds lastSampleTime GUARDED_BY(mThreadControlMutex);
+    std::optional<std::chrono::steady_clock::time_point> mSampleRequestTime
+            GUARDED_BY(mThreadControlMutex);
+    std::chrono::steady_clock::time_point mLastSampleTime GUARDED_BY(mThreadControlMutex);
 
     std::mutex mSamplingMutex;
     std::unordered_map<wp<IBinder>, Descriptor, WpHash> mDescriptors GUARDED_BY(mSamplingMutex);
diff --git a/services/surfaceflinger/Scheduler/MessageQueue.cpp b/services/surfaceflinger/Scheduler/MessageQueue.cpp
index 7ff0ddf..4d51125 100644
--- a/services/surfaceflinger/Scheduler/MessageQueue.cpp
+++ b/services/surfaceflinger/Scheduler/MessageQueue.cpp
@@ -23,7 +23,6 @@
 #include <utils/threads.h>
 
 #include <gui/DisplayEventReceiver.h>
-#include <gui/IDisplayEventConnection.h>
 
 #include "EventThread.h"
 #include "FrameTimeline.h"
@@ -33,27 +32,32 @@
 namespace android::impl {
 
 void MessageQueue::Handler::dispatchRefresh() {
-    if ((android_atomic_or(eventMaskRefresh, &mEventMask) & eventMaskRefresh) == 0) {
+    if ((mEventMask.fetch_or(eventMaskRefresh) & eventMaskRefresh) == 0) {
         mQueue.mLooper->sendMessage(this, Message(MessageQueue::REFRESH));
     }
 }
 
 void MessageQueue::Handler::dispatchInvalidate(int64_t vsyncId, nsecs_t expectedVSyncTimestamp) {
-    if ((android_atomic_or(eventMaskInvalidate, &mEventMask) & eventMaskInvalidate) == 0) {
+    if ((mEventMask.fetch_or(eventMaskInvalidate) & eventMaskInvalidate) == 0) {
         mVsyncId = vsyncId;
         mExpectedVSyncTime = expectedVSyncTimestamp;
         mQueue.mLooper->sendMessage(this, Message(MessageQueue::INVALIDATE));
     }
 }
 
+bool MessageQueue::Handler::invalidatePending() {
+    constexpr auto pendingMask = eventMaskInvalidate | eventMaskRefresh;
+    return (mEventMask.load() & pendingMask) != 0;
+}
+
 void MessageQueue::Handler::handleMessage(const Message& message) {
     switch (message.what) {
         case INVALIDATE:
-            android_atomic_and(~eventMaskInvalidate, &mEventMask);
+            mEventMask.fetch_and(~eventMaskInvalidate);
             mQueue.mFlinger->onMessageReceived(message.what, mVsyncId, mExpectedVSyncTime);
             break;
         case REFRESH:
-            android_atomic_and(~eventMaskRefresh, &mEventMask);
+            mEventMask.fetch_and(~eventMaskRefresh);
             mQueue.mFlinger->onMessageReceived(message.what, mVsyncId, mExpectedVSyncTime);
             break;
     }
@@ -106,7 +110,7 @@
     {
         std::lock_guard lock(mVsync.mutex);
         mVsync.lastCallbackTime = std::chrono::nanoseconds(vsyncTime);
-        mVsync.mScheduled = false;
+        mVsync.scheduled = false;
     }
     mHandler->dispatchInvalidate(mVsync.tokenManager->generateTokenForPredictions(
                                          {targetWakeupTime, readyTime, vsyncTime}),
@@ -131,9 +135,10 @@
     ATRACE_CALL();
     std::lock_guard lock(mVsync.mutex);
     mVsync.workDuration = workDuration;
-    if (mVsync.mScheduled) {
-        mVsync.registration->schedule({mVsync.workDuration.get().count(), /*readyDuration=*/0,
-                                       mVsync.lastCallbackTime.count()});
+    if (mVsync.scheduled) {
+        mVsync.expectedWakeupTime = mVsync.registration->schedule(
+                {mVsync.workDuration.get().count(),
+                 /*readyDuration=*/0, mVsync.lastCallbackTime.count()});
     }
 }
 
@@ -176,10 +181,11 @@
     }
 
     std::lock_guard lock(mVsync.mutex);
-    mVsync.mScheduled = true;
-    mVsync.registration->schedule({.workDuration = mVsync.workDuration.get().count(),
-                                   .readyDuration = 0,
-                                   .earliestVsync = mVsync.lastCallbackTime.count()});
+    mVsync.scheduled = true;
+    mVsync.expectedWakeupTime =
+            mVsync.registration->schedule({.workDuration = mVsync.workDuration.get().count(),
+                                           .readyDuration = 0,
+                                           .earliestVsync = mVsync.lastCallbackTime.count()});
 }
 
 void MessageQueue::refresh() {
@@ -200,4 +206,19 @@
     }
 }
 
+std::optional<std::chrono::steady_clock::time_point> MessageQueue::nextExpectedInvalidate() {
+    if (mHandler->invalidatePending()) {
+        return std::chrono::steady_clock::now();
+    }
+
+    std::lock_guard lock(mVsync.mutex);
+    if (mVsync.scheduled) {
+        LOG_ALWAYS_FATAL_IF(!mVsync.expectedWakeupTime.has_value(), "callback was never scheduled");
+        const auto expectedWakeupTime = std::chrono::nanoseconds(*mVsync.expectedWakeupTime);
+        return std::optional<std::chrono::steady_clock::time_point>(expectedWakeupTime);
+    }
+
+    return std::nullopt;
+}
+
 } // namespace android::impl
diff --git a/services/surfaceflinger/Scheduler/MessageQueue.h b/services/surfaceflinger/Scheduler/MessageQueue.h
index 2934af0..58ce9b9 100644
--- a/services/surfaceflinger/Scheduler/MessageQueue.h
+++ b/services/surfaceflinger/Scheduler/MessageQueue.h
@@ -72,6 +72,7 @@
     virtual void postMessage(sp<MessageHandler>&&) = 0;
     virtual void invalidate() = 0;
     virtual void refresh() = 0;
+    virtual std::optional<std::chrono::steady_clock::time_point> nextExpectedInvalidate() = 0;
 };
 
 // ---------------------------------------------------------------------------
@@ -81,9 +82,13 @@
 class MessageQueue : public android::MessageQueue {
 protected:
     class Handler : public MessageHandler {
-        enum { eventMaskInvalidate = 0x1, eventMaskRefresh = 0x2, eventMaskTransaction = 0x4 };
+        enum : uint32_t {
+            eventMaskInvalidate = 0x1,
+            eventMaskRefresh = 0x2,
+            eventMaskTransaction = 0x4
+        };
         MessageQueue& mQueue;
-        int32_t mEventMask;
+        std::atomic<uint32_t> mEventMask;
         std::atomic<int64_t> mVsyncId;
         std::atomic<nsecs_t> mExpectedVSyncTime;
 
@@ -92,6 +97,7 @@
         void handleMessage(const Message& message) override;
         virtual void dispatchRefresh();
         virtual void dispatchInvalidate(int64_t vsyncId, nsecs_t expectedVSyncTimestamp);
+        virtual bool invalidatePending();
     };
 
     friend class Handler;
@@ -107,7 +113,8 @@
         TracedOrdinal<std::chrono::nanoseconds> workDuration
                 GUARDED_BY(mutex) = {"VsyncWorkDuration-sf", std::chrono::nanoseconds(0)};
         std::chrono::nanoseconds lastCallbackTime GUARDED_BY(mutex) = std::chrono::nanoseconds{0};
-        bool mScheduled GUARDED_BY(mutex) = false;
+        bool scheduled GUARDED_BY(mutex) = false;
+        std::optional<nsecs_t> expectedWakeupTime GUARDED_BY(mutex);
         TracedOrdinal<int> value = {"VSYNC-sf", 0};
     };
 
@@ -141,6 +148,8 @@
 
     // sends REFRESH message at next VSYNC
     void refresh() override;
+
+    std::optional<std::chrono::steady_clock::time_point> nextExpectedInvalidate() override;
 };
 
 } // namespace impl
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 267dd6d..41aacc0 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -1909,6 +1909,7 @@
         mRefreshPending = true;
         onMessageRefresh();
     }
+    notifyRegionSamplingThread();
 }
 
 bool SurfaceFlinger::handleMessageTransaction() {
@@ -2311,10 +2312,6 @@
         }
     }
 
-    if (mLumaSampling && mRegionSamplingThread) {
-        mRegionSamplingThread->notifyNewContent();
-    }
-
     // Even though ATRACE_INT64 already checks if tracing is enabled, it doesn't prevent the
     // side-effect of getTotalSize(), so we check that again here
     if (ATRACE_ENABLED()) {
@@ -3067,8 +3064,7 @@
                            configs.late.sfWorkDuration);
 
     mRegionSamplingThread =
-            new RegionSamplingThread(*this, *mScheduler,
-                                     RegionSamplingThread::EnvironmentTimingTunables());
+            new RegionSamplingThread(*this, RegionSamplingThread::EnvironmentTimingTunables());
     mFpsReporter = new FpsReporter(*mFrameTimeline, *this);
     // Dispatch a mode change request for the primary display on scheduler
     // initialization, so that the EventThreads always contain a reference to a
@@ -6752,6 +6748,19 @@
 
     return layer;
 }
+
+void SurfaceFlinger::scheduleRegionSamplingThread() {
+    static_cast<void>(schedule([&] { notifyRegionSamplingThread(); }));
+}
+
+void SurfaceFlinger::notifyRegionSamplingThread() {
+    if (!mLumaSampling || !mRegionSamplingThread) {
+        return;
+    }
+
+    mRegionSamplingThread->onCompositionComplete(mEventQueue->nextExpectedInvalidate());
+}
+
 } // namespace android
 
 #if defined(__gl_h_)
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index cd20bd1..ea40cb5 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -1429,6 +1429,9 @@
             REQUIRES(mStateLock);
 
     std::atomic<ui::Transform::RotationFlags> mDefaultDisplayTransformHint;
+
+    void scheduleRegionSamplingThread();
+    void notifyRegionSamplingThread();
 };
 
 } // namespace android
diff --git a/services/surfaceflinger/tests/unittests/MessageQueueTest.cpp b/services/surfaceflinger/tests/unittests/MessageQueueTest.cpp
index 18ebbab..dbd51fe 100644
--- a/services/surfaceflinger/tests/unittests/MessageQueueTest.cpp
+++ b/services/surfaceflinger/tests/unittests/MessageQueueTest.cpp
@@ -104,8 +104,12 @@
     const auto timing = scheduler::VSyncDispatch::ScheduleTiming{.workDuration = mDuration.count(),
                                                                  .readyDuration = 0,
                                                                  .earliestVsync = 0};
-    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(0));
+    EXPECT_FALSE(mEventQueue.nextExpectedInvalidate().has_value());
+
+    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(1234));
     EXPECT_NO_FATAL_FAILURE(mEventQueue.invalidate());
+    EXPECT_TRUE(mEventQueue.nextExpectedInvalidate().has_value());
+    EXPECT_EQ(1234, mEventQueue.nextExpectedInvalidate().value().time_since_epoch().count());
 }
 
 TEST_F(MessageQueueTest, invalidateTwice) {
@@ -114,11 +118,15 @@
                                                                  .readyDuration = 0,
                                                                  .earliestVsync = 0};
 
-    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(0));
+    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(1234));
     EXPECT_NO_FATAL_FAILURE(mEventQueue.invalidate());
+    EXPECT_TRUE(mEventQueue.nextExpectedInvalidate().has_value());
+    EXPECT_EQ(1234, mEventQueue.nextExpectedInvalidate().value().time_since_epoch().count());
 
-    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(0));
+    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(4567));
     EXPECT_NO_FATAL_FAILURE(mEventQueue.invalidate());
+    EXPECT_TRUE(mEventQueue.nextExpectedInvalidate().has_value());
+    EXPECT_EQ(4567, mEventQueue.nextExpectedInvalidate().value().time_since_epoch().count());
 }
 
 TEST_F(MessageQueueTest, invalidateTwiceWithCallback) {
@@ -127,8 +135,10 @@
                                                                  .readyDuration = 0,
                                                                  .earliestVsync = 0};
 
-    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(0));
+    EXPECT_CALL(mVSyncDispatch, schedule(mCallbackToken, timing)).WillOnce(Return(1234));
     EXPECT_NO_FATAL_FAILURE(mEventQueue.invalidate());
+    EXPECT_TRUE(mEventQueue.nextExpectedInvalidate().has_value());
+    EXPECT_EQ(1234, mEventQueue.nextExpectedInvalidate().value().time_since_epoch().count());
 
     const auto startTime = 100;
     const auto endTime = startTime + mDuration.count();
@@ -141,6 +151,8 @@
     EXPECT_CALL(*mHandler, dispatchInvalidate(vsyncId, presentTime)).Times(1);
     EXPECT_NO_FATAL_FAILURE(mEventQueue.triggerVsyncCallback(presentTime, startTime, endTime));
 
+    EXPECT_FALSE(mEventQueue.nextExpectedInvalidate().has_value());
+
     const auto timingAfterCallback =
             scheduler::VSyncDispatch::ScheduleTiming{.workDuration = mDuration.count(),
                                                      .readyDuration = 0,
diff --git a/services/surfaceflinger/tests/unittests/mock/MockMessageQueue.h b/services/surfaceflinger/tests/unittests/mock/MockMessageQueue.h
index 453c93a..0e7b320 100644
--- a/services/surfaceflinger/tests/unittests/mock/MockMessageQueue.h
+++ b/services/surfaceflinger/tests/unittests/mock/MockMessageQueue.h
@@ -39,6 +39,7 @@
                  void(scheduler::VSyncDispatch&, frametimeline::TokenManager&,
                       std::chrono::nanoseconds));
     MOCK_METHOD1(setDuration, void(std::chrono::nanoseconds workDuration));
+    MOCK_METHOD0(nextExpectedInvalidate, std::optional<std::chrono::steady_clock::time_point>());
 };
 
 } // namespace android::mock