Limit Predictions size based on count instead of time

We have an arbitrary time of 120ms for the predictions before they
expire. Our assumption was that 120ms is plenty enough for apps & sf
that if they finish beyond this, it's basically a jank. However, with
some traces, we have noticed SF not running in the main thread at all
during setPowerMode() for more than 120ms. This is causing a data loss
in jank classification.

This change addresses this by limiting the number of predictions stored
in TokenManager.

Bug: 187091879
Bug: 186874532
Test: libsurfaceflinger_unittest
Change-Id: I555bfd974585b7e0632eade776d201f1189c81e0
diff --git a/services/surfaceflinger/FrameTimeline/FrameTimeline.cpp b/services/surfaceflinger/FrameTimeline/FrameTimeline.cpp
index 0033dbe..f19e2a7 100644
--- a/services/surfaceflinger/FrameTimeline/FrameTimeline.cpp
+++ b/services/surfaceflinger/FrameTimeline/FrameTimeline.cpp
@@ -730,9 +730,11 @@
 int64_t TokenManager::generateTokenForPredictions(TimelineItem&& predictions) {
     ATRACE_CALL();
     std::scoped_lock lock(mMutex);
+    while (mPredictions.size() >= kMaxTokens) {
+        mPredictions.erase(mPredictions.begin());
+    }
     const int64_t assignedToken = mCurrentToken++;
-    mPredictions[assignedToken] = {systemTime(), predictions};
-    flushTokens(systemTime());
+    mPredictions[assignedToken] = predictions;
     return assignedToken;
 }
 
@@ -740,23 +742,11 @@
     std::scoped_lock lock(mMutex);
     auto predictionsIterator = mPredictions.find(token);
     if (predictionsIterator != mPredictions.end()) {
-        return predictionsIterator->second.predictions;
+        return predictionsIterator->second;
     }
     return {};
 }
 
-void TokenManager::flushTokens(nsecs_t flushTime) {
-    for (auto it = mPredictions.begin(); it != mPredictions.end();) {
-        if (flushTime - it->second.timestamp >= kMaxRetentionTime) {
-            it = mPredictions.erase(it);
-        } else {
-            // Tokens are ordered by time. If i'th token is within the retention time, then the
-            // i+1'th token will also be within retention time.
-            break;
-        }
-    }
-}
-
 FrameTimeline::FrameTimeline(std::shared_ptr<TimeStats> timeStats, pid_t surfaceFlingerPid,
                              JankClassificationThresholds thresholds)
       : mMaxDisplayFrames(kDefaultMaxDisplayFrames),
diff --git a/services/surfaceflinger/FrameTimeline/FrameTimeline.h b/services/surfaceflinger/FrameTimeline/FrameTimeline.h
index 0563a53..42be55a 100644
--- a/services/surfaceflinger/FrameTimeline/FrameTimeline.h
+++ b/services/surfaceflinger/FrameTimeline/FrameTimeline.h
@@ -92,11 +92,6 @@
     bool operator!=(const TimelineItem& other) const { return !(*this == other); }
 };
 
-struct TokenManagerPrediction {
-    nsecs_t timestamp = 0;
-    TimelineItem predictions;
-};
-
 struct JankClassificationThresholds {
     // The various thresholds for App and SF. If the actual timestamp falls within the threshold
     // compared to prediction, we treat it as on time.
@@ -334,11 +329,10 @@
 
     void flushTokens(nsecs_t flushTime) REQUIRES(mMutex);
 
-    std::map<int64_t, TokenManagerPrediction> mPredictions GUARDED_BY(mMutex);
+    std::map<int64_t, TimelineItem> mPredictions GUARDED_BY(mMutex);
     int64_t mCurrentToken GUARDED_BY(mMutex);
     mutable std::mutex mMutex;
-    static constexpr nsecs_t kMaxRetentionTime =
-            std::chrono::duration_cast<std::chrono::nanoseconds>(120ms).count();
+    static constexpr size_t kMaxTokens = 500;
 };
 
 class FrameTimeline : public android::frametimeline::FrameTimeline {
diff --git a/services/surfaceflinger/tests/unittests/FrameTimelineTest.cpp b/services/surfaceflinger/tests/unittests/FrameTimelineTest.cpp
index 6ed6148..c6a4115 100644
--- a/services/surfaceflinger/tests/unittests/FrameTimelineTest.cpp
+++ b/services/surfaceflinger/tests/unittests/FrameTimelineTest.cpp
@@ -73,7 +73,7 @@
         mTokenManager = &mFrameTimeline->mTokenManager;
         mTraceCookieCounter = &mFrameTimeline->mTraceCookieCounter;
         maxDisplayFrames = &mFrameTimeline->mMaxDisplayFrames;
-        maxTokenRetentionTime = mTokenManager->kMaxRetentionTime;
+        maxTokens = mTokenManager->kMaxTokens;
     }
 
     // Each tracing session can be used for a single block of Start -> Stop.
@@ -111,9 +111,11 @@
         mFrameTimeline->setSfPresent(2500, presentFence1);
     }
 
-    void flushTokens(nsecs_t flushTime) {
-        std::lock_guard<std::mutex> lock(mTokenManager->mMutex);
-        mTokenManager->flushTokens(flushTime);
+    void flushTokens() {
+        for (size_t i = 0; i < maxTokens; i++) {
+            mTokenManager->generateTokenForPredictions({});
+        }
+        EXPECT_EQ(getPredictions().size(), maxTokens);
     }
 
     SurfaceFrame& getSurfaceFrame(size_t displayFrameIdx, size_t surfaceFrameIdx) {
@@ -132,7 +134,7 @@
                 a.presentTime == b.presentTime;
     }
 
-    const std::map<int64_t, TokenManagerPrediction>& getPredictions() const {
+    const std::map<int64_t, TimelineItem>& getPredictions() const {
         return mTokenManager->mPredictions;
     }
 
@@ -155,7 +157,7 @@
     TraceCookieCounter* mTraceCookieCounter;
     FenceToFenceTimeMap fenceFactory;
     uint32_t* maxDisplayFrames;
-    nsecs_t maxTokenRetentionTime;
+    size_t maxTokens;
     static constexpr pid_t kSurfaceFlingerPid = 666;
     static constexpr nsecs_t kPresentThreshold = std::chrono::nanoseconds(2ns).count();
     static constexpr nsecs_t kDeadlineThreshold = std::chrono::nanoseconds(2ns).count();
@@ -177,12 +179,11 @@
 TEST_F(FrameTimelineTest, tokenManagerRemovesStalePredictions) {
     int64_t token1 = mTokenManager->generateTokenForPredictions({0, 0, 0});
     EXPECT_EQ(getPredictions().size(), 1u);
-    flushTokens(systemTime() + maxTokenRetentionTime);
+    flushTokens();
     int64_t token2 = mTokenManager->generateTokenForPredictions({10, 20, 30});
     std::optional<TimelineItem> predictions = mTokenManager->getPredictionsForToken(token1);
 
     // token1 should have expired
-    EXPECT_EQ(getPredictions().size(), 1u);
     EXPECT_EQ(predictions.has_value(), false);
 
     predictions = mTokenManager->getPredictionsForToken(token2);
@@ -212,7 +213,7 @@
 
 TEST_F(FrameTimelineTest, createSurfaceFrameForToken_expiredToken) {
     int64_t token1 = mTokenManager->generateTokenForPredictions({0, 0, 0});
-    flushTokens(systemTime() + maxTokenRetentionTime);
+    flushTokens();
     auto surfaceFrame =
             mFrameTimeline->createSurfaceFrameForToken({token1, sInputEventId}, sPidOne, sUidOne,
                                                        sLayerIdOne, sLayerNameOne, sLayerNameOne,
@@ -707,7 +708,7 @@
                                                        sLayerNameOne, /*isBuffer*/ true);
     surfaceFrame1->setAcquireFenceTime(45);
     // Trigger a prediction expiry
-    flushTokens(systemTime() + maxTokenRetentionTime);
+    flushTokens();
     mFrameTimeline->setSfWakeUp(sfToken1, 52, refreshRate);
 
     surfaceFrame1->setPresentState(SurfaceFrame::PresentState::Presented);
@@ -1065,7 +1066,7 @@
     tracingSession->StartBlocking();
     int64_t displayFrameToken1 = mTokenManager->generateTokenForPredictions({10, 25, 30});
     // Flush the token so that it would expire
-    flushTokens(systemTime() + maxTokenRetentionTime);
+    flushTokens();
 
     // Set up the display frame
     mFrameTimeline->setSfWakeUp(displayFrameToken1, 20, Fps::fromPeriodNsecs(11));
@@ -1283,7 +1284,7 @@
             mTokenManager->generateTokenForPredictions({appStartTime, appEndTime, appPresentTime});
 
     // Flush the token so that it would expire
-    flushTokens(systemTime() + maxTokenRetentionTime);
+    flushTokens();
     auto surfaceFrame1 =
             mFrameTimeline->createSurfaceFrameForToken({surfaceFrameToken, /*inputEventId*/ 0},
                                                        sPidOne, sUidOne, sLayerIdOne, sLayerNameOne,
@@ -1359,7 +1360,7 @@
             mTokenManager->generateTokenForPredictions({appStartTime, appEndTime, appPresentTime});
 
     // Flush the token so that it would expire
-    flushTokens(systemTime() + maxTokenRetentionTime);
+    flushTokens();
     auto surfaceFrame1 =
             mFrameTimeline->createSurfaceFrameForToken({surfaceFrameToken, /*inputEventId*/ 0},
                                                        sPidOne, sUidOne, sLayerIdOne, sLayerNameOne,