Correctly pass screenshot fences to transaction callbacks

In some instances, a screenshot may be captured before a layer has a
release callback registered. This can happen when a new buffer has
not yet been transacted before a screenshot is captured. This causes the
screenshot fence to be dropped and the possibility of tearing when
capturing a screenshot while continuously rendering.

To resolve this, buffer screenshot fences into a list of future fences
when there is no callback registered, and merge those fences when
dispatching the release callback.

Bug: 302703346
Test: SurfaceViewTests#testMovingWhiteSurfaceView 100 times
Change-Id: I91aec3cdb0973092d48cd77e59dd3999e9d9e847

11 files changed, 197 insertions, 36 deletions

diff --git a/include/ftl/details/future.h b/include/ftl/details/future.h
index df1323e8be..8d82e0ffd2 100644
--- a/include/ftl/details/future.h
+++ b/include/ftl/details/future.h
@@ -73,8 +73,18 @@ class BaseFuture<Self, T, std::future> {
     return std::get<Impl>(self()).get();
   }
 
+  template <class Rep, class Period>
+  std::future_status wait_for(const std::chrono::duration<Rep, Period>& timeout_duration) const {
+    if (std::holds_alternative<T>(self())) {
+      return std::future_status::ready;
+    }
+
+    return std::get<Impl>(self()).wait_for(timeout_duration);
+  }
+
  private:
   auto& self() { return static_cast<Self&>(*this).future_; }
+  const auto& self() const { return static_cast<const Self&>(*this).future_; }
 };
 
 template <typename Self, typename T>
@@ -90,6 +100,15 @@ class BaseFuture<Self, T, std::shared_future> {
     return std::get<Impl>(self()).get();
   }
 
+  template <class Rep, class Period>
+  std::future_status wait_for(const std::chrono::duration<Rep, Period>& timeout_duration) const {
+    if (std::holds_alternative<T>(self())) {
+      return std::future_status::ready;
+    }
+
+    return std::get<Impl>(self()).wait_for(timeout_duration);
+  }
+
  private:
   const auto& self() const { return static_cast<const Self&>(*this).future_; }
 };
diff --git a/include/ftl/future.h b/include/ftl/future.h
index c78f9b76b6..dad180ff8e 100644
--- a/include/ftl/future.h
+++ b/include/ftl/future.h
@@ -51,6 +51,7 @@ class Future final : public details::BaseFuture<Future<T, FutureImpl>, T, Future
   // Forwarding functions. Base::share is only defined when FutureImpl is std::future, whereas the
   // following are defined for either FutureImpl:
   using Base::get;
+  using Base::wait_for;
 
   // Attaches a continuation to the future. The continuation is a function that maps T to either R
   // or ftl::Future<R>. In the former case, the chain wraps the result in a future as if by
diff --git a/libs/ftl/future_test.cpp b/libs/ftl/future_test.cpp
index 5a245b681c..1140639c87 100644
--- a/libs/ftl/future_test.cpp
+++ b/libs/ftl/future_test.cpp
@@ -102,4 +102,42 @@ TEST(Future, Chain) {
   decrement_thread.join();
 }
 
+TEST(Future, WaitFor) {
+  using namespace std::chrono_literals;
+  {
+    auto future = ftl::yield(42);
+    // Check that we can wait_for multiple times without invalidating the future
+    EXPECT_EQ(future.wait_for(1s), std::future_status::ready);
+    EXPECT_EQ(future.wait_for(1s), std::future_status::ready);
+    EXPECT_EQ(future.get(), 42);
+  }
+
+  {
+    std::condition_variable cv;
+    std::mutex m;
+    bool ready = false;
+
+    std::packaged_task<int32_t()> get_int([&] {
+      std::unique_lock lk(m);
+      cv.wait(lk, [&] { return ready; });
+      return 24;
+    });
+
+    auto get_future = ftl::Future(get_int.get_future());
+    std::thread get_thread(std::move(get_int));
+
+    EXPECT_EQ(get_future.wait_for(0s), std::future_status::timeout);
+    {
+      std::unique_lock lk(m);
+      ready = true;
+    }
+    cv.notify_one();
+
+    EXPECT_EQ(get_future.wait_for(1s), std::future_status::ready);
+    EXPECT_EQ(get_future.get(), 24);
+
+    get_thread.join();
+  }
+}
+
 }  // namespace android::test
diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index c8b1059b94..219a8e20a4 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -78,11 +78,13 @@
 #include "SurfaceFlinger.h"
 #include "TimeStats/TimeStats.h"
 #include "TunnelModeEnabledReporter.h"
+#include "Utils/FenceUtils.h"
 
 #define DEBUG_RESIZE 0
 #define EARLY_RELEASE_ENABLED false
 
 namespace android {
+using namespace std::chrono_literals;
 namespace {
 constexpr int kDumpTableRowLength = 159;
 
@@ -2911,7 +2913,8 @@ void Layer::callReleaseBufferCallback(const sp<ITransactionCompletedListener>& l
 }
 
 void Layer::onLayerDisplayed(ftl::SharedFuture<FenceResult> futureFenceResult,
-                             ui::LayerStack layerStack) {
+                             ui::LayerStack layerStack,
+                             std::function<FenceResult(FenceResult)>&& continuation) {
     // If we are displayed on multiple displays in a single composition cycle then we would
     // need to do careful tracking to enable the use of the mLastClientCompositionFence.
     //  For example we can only use it if all the displays are client comp, and we need
@@ -2946,11 +2949,41 @@ void Layer::onLayerDisplayed(ftl::SharedFuture<FenceResult> futureFenceResult,
             break;
         }
     }
+
+    if (!FlagManager::getInstance().screenshot_fence_preservation() && continuation) {
+        futureFenceResult = ftl::Future(futureFenceResult).then(std::move(continuation)).share();
+    }
+
     if (ch != nullptr) {
         ch->previousReleaseCallbackId = mPreviousReleaseCallbackId;
         ch->previousReleaseFences.emplace_back(std::move(futureFenceResult));
         ch->name = mName;
+    } else if (FlagManager::getInstance().screenshot_fence_preservation()) {
+        // If we didn't get a release callback yet, e.g. some scenarios when capturing screenshots
+        // asynchronously, then make sure we don't drop the fence.
+        mAdditionalPreviousReleaseFences.emplace_back(std::move(futureFenceResult),
+                                                      std::move(continuation));
+        std::vector<FenceAndContinuation> mergedFences;
+        sp<Fence> prevFence = nullptr;
+        // For a layer that's frequently screenshotted, try to merge fences to make sure we don't
+        // grow unbounded.
+        for (const auto& futureAndContinution : mAdditionalPreviousReleaseFences) {
+            auto result = futureAndContinution.future.wait_for(0s);
+            if (result != std::future_status::ready) {
+                mergedFences.emplace_back(futureAndContinution);
+                continue;
+            }
+
+            mergeFence(getDebugName(), futureAndContinution.chain().get().value_or(Fence::NO_FENCE),
+                       prevFence);
+        }
+        if (prevFence != nullptr) {
+            mergedFences.emplace_back(ftl::yield(FenceResult(std::move(prevFence))).share());
+        }
+
+        mAdditionalPreviousReleaseFences.swap(mergedFences);
     }
+
     if (mBufferInfo.mBuffer) {
         mPreviouslyPresentedLayerStacks.push_back(layerStack);
     }
@@ -3440,6 +3473,15 @@ bool Layer::setTransactionCompletedListeners(const std::vector<sp<CallbackHandle
             handle->acquireTimeOrFence = mCallbackHandleAcquireTimeOrFence;
             handle->frameNumber = mDrawingState.frameNumber;
             handle->previousFrameNumber = mDrawingState.previousFrameNumber;
+            if (FlagManager::getInstance().screenshot_fence_preservation() &&
+                mPreviousReleaseBufferEndpoint == handle->listener) {
+                // Add fences from previous screenshots now so that they can be dispatched to the
+                // client.
+                for (const auto& futureAndContinution : mAdditionalPreviousReleaseFences) {
+                    handle->previousReleaseFences.emplace_back(futureAndContinution.chain());
+                }
+                mAdditionalPreviousReleaseFences.clear();
+            }
 
             // Store so latched time and release fence can be set
             mDrawingState.callbackHandles.push_back(handle);
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index c772e0ea1b..dfd57c65e3 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -555,7 +555,8 @@ public:
     const compositionengine::LayerFECompositionState* getCompositionState() const;
     bool fenceHasSignaled() const;
     void onPreComposition(nsecs_t refreshStartTime);
-    void onLayerDisplayed(ftl::SharedFuture<FenceResult>, ui::LayerStack layerStack);
+    void onLayerDisplayed(ftl::SharedFuture<FenceResult>, ui::LayerStack layerStack,
+                          std::function<FenceResult(FenceResult)>&& continuation = nullptr);
 
     void setWasClientComposed(const sp<Fence>& fence) {
         mLastClientCompositionFence = fence;
@@ -932,6 +933,19 @@ public:
     // the release fences from the correct displays when we release the last buffer
     // from the layer.
     std::vector<ui::LayerStack> mPreviouslyPresentedLayerStacks;
+    struct FenceAndContinuation {
+        ftl::SharedFuture<FenceResult> future;
+        std::function<FenceResult(FenceResult)> continuation;
+
+        ftl::SharedFuture<FenceResult> chain() const {
+            if (continuation) {
+                return ftl::Future(future).then(continuation).share();
+            } else {
+                return future;
+            }
+        }
+    };
+    std::vector<FenceAndContinuation> mAdditionalPreviousReleaseFences;
     // Exposed so SurfaceFlinger can assert that it's held
     const sp<SurfaceFlinger> mFlinger;
 
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 51a11e2e49..7626fe73e9 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -8152,14 +8152,23 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
                                                      : ftl::yield(present()).share();
 
     for (auto& [layer, layerFE] : layers) {
-        layer->onLayerDisplayed(ftl::Future(presentFuture)
-                                        .then([layerFE = std::move(layerFE)](FenceResult) {
-                                            return layerFE->stealCompositionResult()
-                                                    .releaseFences.back()
-                                                    .first.get();
-                                        })
-                                        .share(),
-                                ui::INVALID_LAYER_STACK);
+        layer->onLayerDisplayed(presentFuture, ui::INVALID_LAYER_STACK,
+                                [layerFE = std::move(layerFE)](FenceResult) {
+                                    if (FlagManager::getInstance()
+                                                .screenshot_fence_preservation()) {
+                                        const auto compositionResult =
+                                                layerFE->stealCompositionResult();
+                                        const auto& fences = compositionResult.releaseFences;
+                                        // CompositionEngine may choose to cull layers that
+                                        // aren't visible, so pass a non-fence.
+                                        return fences.empty() ? Fence::NO_FENCE
+                                                              : fences.back().first.get();
+                                    } else {
+                                        return layerFE->stealCompositionResult()
+                                                .releaseFences.back()
+                                                .first.get();
+                                    }
+                                });
     }
 
     return presentFuture;
diff --git a/services/surfaceflinger/TransactionCallbackInvoker.cpp b/services/surfaceflinger/TransactionCallbackInvoker.cpp
index 6a155c17df..7b5298c82e 100644
--- a/services/surfaceflinger/TransactionCallbackInvoker.cpp
+++ b/services/surfaceflinger/TransactionCallbackInvoker.cpp
@@ -25,6 +25,7 @@
 
 #include "TransactionCallbackInvoker.h"
 #include "BackgroundExecutor.h"
+#include "Utils/FenceUtils.h"
 
 #include <cinttypes>
 
@@ -127,33 +128,8 @@ status_t TransactionCallbackInvoker::addCallbackHandle(const sp<CallbackHandle>&
     sp<IBinder> surfaceControl = handle->surfaceControl.promote();
     if (surfaceControl) {
         sp<Fence> prevFence = nullptr;
-
         for (const auto& future : handle->previousReleaseFences) {
-            sp<Fence> currentFence = future.get().value_or(Fence::NO_FENCE);
-            if (prevFence == nullptr && currentFence->getStatus() != Fence::Status::Invalid) {
-                prevFence = std::move(currentFence);
-            } else if (prevFence != nullptr) {
-                // If both fences are signaled or both are unsignaled, we need to merge
-                // them to get an accurate timestamp.
-                if (prevFence->getStatus() != Fence::Status::Invalid &&
-                    prevFence->getStatus() == currentFence->getStatus()) {
-                    char fenceName[32] = {};
-                    snprintf(fenceName, 32, "%.28s", handle->name.c_str());
-                    sp<Fence> mergedFence = Fence::merge(fenceName, prevFence, currentFence);
-                    if (mergedFence->isValid()) {
-                        prevFence = std::move(mergedFence);
-                    }
-                } else if (currentFence->getStatus() == Fence::Status::Unsignaled) {
-                    // If one fence has signaled and the other hasn't, the unsignaled
-                    // fence will approximately correspond with the correct timestamp.
-                    // There's a small race if both fences signal at about the same time
-                    // and their statuses are retrieved with unfortunate timing. However,
-                    // by this point, they will have both signaled and only the timestamp
-                    // will be slightly off; any dependencies after this point will
-                    // already have been met.
-                    prevFence = std::move(currentFence);
-                }
-            }
+            mergeFence(handle->name.c_str(), future.get().value_or(Fence::NO_FENCE), prevFence);
         }
         handle->previousReleaseFence = prevFence;
         handle->previousReleaseFences.clear();
diff --git a/services/surfaceflinger/Utils/FenceUtils.h b/services/surfaceflinger/Utils/FenceUtils.h
new file mode 100644
index 0000000000..f6f70062df
--- /dev/null
+++ b/services/surfaceflinger/Utils/FenceUtils.h
@@ -0,0 +1,51 @@
+/**
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <ui/Fence.h>
+
+namespace android {
+
+// TODO: measure if Fence::merge is cheaper
+inline void mergeFence(const char* debugName, sp<Fence>&& incomingFence, sp<Fence>& prevFence) {
+    if (prevFence == nullptr && incomingFence->getStatus() != Fence::Status::Invalid) {
+        prevFence = std::move(incomingFence);
+    } else if (prevFence != nullptr) {
+        // If both fences are signaled or both are unsignaled, we need to merge
+        // them to get an accurate timestamp.
+        if (prevFence->getStatus() != Fence::Status::Invalid &&
+            prevFence->getStatus() == incomingFence->getStatus()) {
+            char fenceName[32] = {};
+            snprintf(fenceName, 32, "%.28s", debugName);
+            sp<Fence> mergedFence = Fence::merge(fenceName, prevFence, incomingFence);
+            if (mergedFence->isValid()) {
+                prevFence = std::move(mergedFence);
+            }
+        } else if (incomingFence->getStatus() == Fence::Status::Unsignaled) {
+            // If one fence has signaled and the other hasn't, the unsignaled
+            // fence will approximately correspond with the correct timestamp.
+            // There's a small race if both fences signal at about the same time
+            // and their statuses are retrieved with unfortunate timing. However,
+            // by this point, they will have both signaled and only the timestamp
+            // will be slightly off; any dependencies after this point will
+            // already have been met.
+            prevFence = std::move(incomingFence);
+        }
+    }
+}
+
+} // namespace android
diff --git a/services/surfaceflinger/common/FlagManager.cpp b/services/surfaceflinger/common/FlagManager.cpp
index 255b517e3e..b07e7ace06 100644
--- a/services/surfaceflinger/common/FlagManager.cpp
+++ b/services/surfaceflinger/common/FlagManager.cpp
@@ -128,6 +128,7 @@ void FlagManager::dump(std::string& result) const {
     DUMP_READ_ONLY_FLAG(fp16_client_target);
     DUMP_READ_ONLY_FLAG(game_default_frame_rate);
     DUMP_READ_ONLY_FLAG(enable_layer_command_batching);
+    DUMP_READ_ONLY_FLAG(screenshot_fence_preservation);
     DUMP_READ_ONLY_FLAG(vulkan_renderengine);
 #undef DUMP_READ_ONLY_FLAG
 #undef DUMP_SERVER_FLAG
@@ -203,6 +204,7 @@ FLAG_MANAGER_READ_ONLY_FLAG(display_protected, "")
 FLAG_MANAGER_READ_ONLY_FLAG(fp16_client_target, "debug.sf.fp16_client_target")
 FLAG_MANAGER_READ_ONLY_FLAG(game_default_frame_rate, "")
 FLAG_MANAGER_READ_ONLY_FLAG(enable_layer_command_batching, "")
+FLAG_MANAGER_READ_ONLY_FLAG(screenshot_fence_preservation, "debug.sf.screenshot_fence_preservation")
 FLAG_MANAGER_READ_ONLY_FLAG(vulkan_renderengine, "debug.renderengine.vulkan")
 
 /// Trunk stable server flags ///
diff --git a/services/surfaceflinger/common/include/common/FlagManager.h b/services/surfaceflinger/common/include/common/FlagManager.h
index 15938c012c..2a30a40730 100644
--- a/services/surfaceflinger/common/include/common/FlagManager.h
+++ b/services/surfaceflinger/common/include/common/FlagManager.h
@@ -68,6 +68,7 @@ public:
     bool fp16_client_target() const;
     bool game_default_frame_rate() const;
     bool enable_layer_command_batching() const;
+    bool screenshot_fence_preservation() const;
     bool vulkan_renderengine() const;
 
 protected:
diff --git a/services/surfaceflinger/surfaceflinger_flags.aconfig b/services/surfaceflinger/surfaceflinger_flags.aconfig
index f097a13d0e..fcbef01bfa 100644
--- a/services/surfaceflinger/surfaceflinger_flags.aconfig
+++ b/services/surfaceflinger/surfaceflinger_flags.aconfig
@@ -166,3 +166,11 @@ flag {
   bug: "293371537"
   is_fixed_read_only: true
 }
+
+flag {
+  name: "screenshot_fence_preservation"
+  namespace: "core_graphics"
+  description: "Bug fix around screenshot fences"
+  bug: "302703346"
+  is_fixed_read_only: true
+}