SF: Add LayerSnapshotGuard to manage layer snapshot

Bug: 238781169
Test: presubmits

Change-Id: Ie221f9a4601512f4b256b0aa8c6b14fd20901cef
diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index c2a0e30..b517568 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -3535,14 +3535,6 @@
     return mSnapshot.get();
 }
 
-void Layer::moveSnapshotToLayerFE() {
-    mLayerFE->mSnapshot = std::move(mSnapshot);
-}
-
-void Layer::moveSnapshotToLayer() {
-    mSnapshot = std::move(mLayerFE->mSnapshot);
-}
-
 void Layer::useSurfaceDamage() {
     if (mFlinger->mForceFullDamage) {
         surfaceDamageRegion = Region::INVALID_REGION;
@@ -4007,6 +3999,28 @@
     }
 }
 
+LayerSnapshotGuard::LayerSnapshotGuard(Layer* layer) : mLayer(layer) {
+    if (mLayer) {
+        mLayer->mLayerFE->mSnapshot = std::move(mLayer->mSnapshot);
+    }
+}
+
+LayerSnapshotGuard::~LayerSnapshotGuard() {
+    if (mLayer) {
+        mLayer->mSnapshot = std::move(mLayer->mLayerFE->mSnapshot);
+    }
+}
+
+LayerSnapshotGuard::LayerSnapshotGuard(LayerSnapshotGuard&& other) : mLayer(other.mLayer) {
+    other.mLayer = nullptr;
+}
+
+LayerSnapshotGuard& LayerSnapshotGuard::operator=(LayerSnapshotGuard&& other) {
+    mLayer = other.mLayer;
+    other.mLayer = nullptr;
+    return *this;
+}
+
 // ---------------------------------------------------------------------------
 
 std::ostream& operator<<(std::ostream& stream, const Layer::FrameRate& rate) {
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index da8be6b..7e46d1c 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -379,15 +379,6 @@
 
     virtual sp<LayerFE> getCompositionEngineLayerFE() const;
 
-    // Move LayerSnapshot from this layer into its LayerFE. This must be called before passing the
-    // LayerFE to CompositionEngine. Moving the snapshot instead of sharing common state
-    // prevents use of LayerFE outside the main thread by making errors obvious (i.e. use outside
-    // the main thread results in SEGFAULTs due to nullptr dereference).
-    void moveSnapshotToLayerFE();
-    // Move LayerSnapshot into this layer from its LayerFE. This must be called after
-    // CompositionEngine has presented the layer.
-    void moveSnapshotToLayer();
-
     const LayerSnapshot* getLayerSnapshot() const;
     LayerSnapshot* editLayerSnapshot();
 
@@ -1188,6 +1179,31 @@
 
     sp<LayerFE> mLayerFE;
     std::unique_ptr<LayerSnapshot> mSnapshot = std::make_unique<LayerSnapshot>();
+
+    friend class LayerSnapshotGuard;
+};
+
+// LayerSnapshotGuard manages the movement of LayerSnapshot between a Layer and its corresponding
+// LayerFE. This class must be used whenever LayerFEs are passed to CompositionEngine. Instances of
+// LayerSnapshotGuard should only be constructed on the main thread and should not be moved outside
+// the main thread.
+//
+// Moving the snapshot instead of sharing common state prevents use of LayerFE outside the main
+// thread by making errors obvious (i.e. use outside the main thread results in SEGFAULTs due to
+// nullptr dereference).
+class LayerSnapshotGuard {
+public:
+    LayerSnapshotGuard(Layer* layer) REQUIRES(kMainThreadContext);
+    ~LayerSnapshotGuard() REQUIRES(kMainThreadContext);
+
+    LayerSnapshotGuard(const LayerSnapshotGuard&) = delete;
+    LayerSnapshotGuard& operator=(const LayerSnapshotGuard&) = delete;
+
+    LayerSnapshotGuard(LayerSnapshotGuard&& other) REQUIRES(kMainThreadContext);
+    LayerSnapshotGuard& operator=(LayerSnapshotGuard&& other) REQUIRES(kMainThreadContext);
+
+private:
+    Layer* mLayer;
 };
 
 std::ostream& operator<<(std::ostream& stream, const Layer::FrameRate& rate);
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index a04ceef..f3148cb 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -2181,7 +2181,6 @@
     mDrawingState.traverseInZOrder([&refreshArgs, &layers](Layer* layer) {
         layer->updateSnapshot(refreshArgs.updatingGeometryThisFrame);
         if (auto layerFE = layer->getCompositionEngineLayerFE()) {
-            layer->moveSnapshotToLayerFE();
             refreshArgs.layers.push_back(layerFE);
             layers.push_back(layer);
         }
@@ -2213,10 +2212,15 @@
     // the scheduler.
     const auto presentTime = systemTime();
 
-    mCompositionEngine->present(refreshArgs);
+    {
+        std::vector<LayerSnapshotGuard> layerSnapshotGuards;
+        for (Layer* layer : layers) {
+            layerSnapshotGuards.emplace_back(layer);
+        }
+        mCompositionEngine->present(refreshArgs);
+    }
 
     for (auto& layer : layers) {
-        layer->moveSnapshotToLayer();
         CompositionResult compositionResult{
                 layer->getCompositionEngineLayerFE()->stealCompositionResult()};
         layer->onPreComposition(compositionResult.refreshStartTime);
@@ -3312,21 +3316,16 @@
         }
     }
 
-    std::vector<Layer*> cursorLayers;
-    mDrawingState.traverse([&cursorLayers](Layer* layer) {
+    std::vector<LayerSnapshotGuard> layerSnapshotGuards;
+    mDrawingState.traverse([&layerSnapshotGuards](Layer* layer) {
         if (layer->getLayerSnapshot()->compositionType ==
             aidl::android::hardware::graphics::composer3::Composition::CURSOR) {
             layer->updateSnapshot(false /* updateGeometry */);
-            layer->moveSnapshotToLayerFE();
-            cursorLayers.push_back(layer);
+            layerSnapshotGuards.emplace_back(layer);
         }
     });
 
     mCompositionEngine->updateCursorAsync(refreshArgs);
-
-    for (Layer* layer : cursorLayers) {
-        layer->moveSnapshotToLayer();
-    }
 }
 
 void SurfaceFlinger::requestDisplayModes(
@@ -6322,37 +6321,38 @@
 
     bool canCaptureBlackoutContent = hasCaptureBlackoutContentPermission();
 
-    auto future = mScheduler->schedule([=, renderAreaFuture = std::move(renderAreaFuture)]() mutable
-                                       -> ftl::SharedFuture<FenceResult> {
-        ScreenCaptureResults captureResults;
-        std::unique_ptr<RenderArea> renderArea = renderAreaFuture.get();
-        if (!renderArea) {
-            ALOGW("Skipping screen capture because of invalid render area.");
-            captureResults.fenceResult = base::unexpected(NO_MEMORY);
-            captureListener->onScreenCaptureCompleted(captureResults);
-            return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
-        }
+    auto future = mScheduler->schedule(
+            [=, renderAreaFuture = std::move(renderAreaFuture)]() FTL_FAKE_GUARD(
+                    kMainThreadContext) mutable -> ftl::SharedFuture<FenceResult> {
+                ScreenCaptureResults captureResults;
+                std::unique_ptr<RenderArea> renderArea = renderAreaFuture.get();
+                if (!renderArea) {
+                    ALOGW("Skipping screen capture because of invalid render area.");
+                    captureResults.fenceResult = base::unexpected(NO_MEMORY);
+                    captureListener->onScreenCaptureCompleted(captureResults);
+                    return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
+                }
 
-        ftl::SharedFuture<FenceResult> renderFuture;
-        renderArea->render([&] {
-            renderFuture =
-                    renderScreenImpl(*renderArea, traverseLayers, buffer, canCaptureBlackoutContent,
-                                     regionSampling, grayscale, captureResults);
-        });
+                ftl::SharedFuture<FenceResult> renderFuture;
+                renderArea->render([&]() FTL_FAKE_GUARD(kMainThreadContext) {
+                    renderFuture = renderScreenImpl(*renderArea, traverseLayers, buffer,
+                                                    canCaptureBlackoutContent, regionSampling,
+                                                    grayscale, captureResults);
+                });
 
-        if (captureListener) {
-            // Defer blocking on renderFuture back to the Binder thread.
-            return ftl::Future(std::move(renderFuture))
-                    .then([captureListener, captureResults = std::move(captureResults)](
-                                  FenceResult fenceResult) mutable -> FenceResult {
-                        captureResults.fenceResult = std::move(fenceResult);
-                        captureListener->onScreenCaptureCompleted(captureResults);
-                        return base::unexpected(NO_ERROR);
-                    })
-                    .share();
-        }
-        return renderFuture;
-    });
+                if (captureListener) {
+                    // Defer blocking on renderFuture back to the Binder thread.
+                    return ftl::Future(std::move(renderFuture))
+                            .then([captureListener, captureResults = std::move(captureResults)](
+                                          FenceResult fenceResult) mutable -> FenceResult {
+                                captureResults.fenceResult = std::move(fenceResult);
+                                captureListener->onScreenCaptureCompleted(captureResults);
+                                return base::unexpected(NO_ERROR);
+                            })
+                            .share();
+                }
+                return renderFuture;
+            });
 
     // Flatten nested futures.
     auto chain = ftl::Future(std::move(future)).then([](ftl::SharedFuture<FenceResult> future) {
@@ -6446,7 +6446,7 @@
     const auto display = renderArea.getDisplayDevice();
     std::vector<Layer*> renderedLayers;
     bool disableBlurs = false;
-    traverseLayers([&](Layer* layer) {
+    traverseLayers([&](Layer* layer) FTL_FAKE_GUARD(kMainThreadContext) {
         // Layer::prepareClientComposition uses the layer's snapshot to populate the resulting
         // LayerSettings. Calling Layer::updateSnapshot ensures that LayerSettings are
         // generated with the layer's current buffer and geometry.
@@ -6477,10 +6477,11 @@
             return;
         }
 
-        layer->moveSnapshotToLayerFE();
-        std::optional<compositionengine::LayerFE::LayerSettings> settings =
-                layerFE->prepareClientComposition(targetSettings);
-        layer->moveSnapshotToLayer();
+        std::optional<compositionengine::LayerFE::LayerSettings> settings;
+        {
+            LayerSnapshotGuard layerSnapshotGuard(layer);
+            settings = layerFE->prepareClientComposition(targetSettings);
+        }
 
         if (!settings) {
             return;
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 0f8ac98..d403f47 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -798,7 +798,8 @@
     ftl::SharedFuture<FenceResult> renderScreenImpl(
             const RenderArea&, TraverseLayersFunction,
             const std::shared_ptr<renderengine::ExternalTexture>&, bool canCaptureBlackoutContent,
-            bool regionSampling, bool grayscale, ScreenCaptureResults&) EXCLUDES(mStateLock);
+            bool regionSampling, bool grayscale, ScreenCaptureResults&) EXCLUDES(mStateLock)
+            REQUIRES(kMainThreadContext);
 
     // If the uid provided is not UNSET_UID, the traverse will skip any layers that don't have a
     // matching ownerUid
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 83d852a..c72fb76 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -401,9 +401,10 @@
                                 const std::shared_ptr<renderengine::ExternalTexture>& buffer,
                                 bool forSystem, bool regionSampling) {
         ScreenCaptureResults captureResults;
-        return mFlinger->renderScreenImpl(renderArea, traverseLayers, buffer, forSystem,
-                                                regionSampling, false /* grayscale */,
-                                                captureResults);
+        return FTL_FAKE_GUARD(kMainThreadContext,
+                              mFlinger->renderScreenImpl(renderArea, traverseLayers, buffer,
+                                                         forSystem, regionSampling,
+                                                         false /* grayscale */, captureResults));
     }
 
     auto traverseLayersInLayerStack(ui::LayerStack layerStack, int32_t uid,