Remove main thread double hop from screenshots

The SF main thread is accessed twice during screenshots,
which leads to possible inconsistent states between jumps
onto the main thread. Removing the double hop preserves
correctness and reduces the amount of computation taking place
on the main thread. The only time the main thread should be
accessed should be when getting layer snapshots. All other
work related to screenshots can take place in a binder thread.

Bug: b/294936197, b/285553970
Test: atest SurfaceFlinger_test
Flag: com.android.graphics.surfaceflinger.flags.single_hop_screenshot
Change-Id: If9fd36f82c2d550bc0821b52fef3ea88b8099116

7 files changed, 250 insertions, 95 deletions

diff --git a/services/surfaceflinger/RegionSamplingThread.cpp b/services/surfaceflinger/RegionSamplingThread.cpp
index 59eb7f5d4a..5add290e96 100644
--- a/services/surfaceflinger/RegionSamplingThread.cpp
+++ b/services/surfaceflinger/RegionSamplingThread.cpp
@@ -348,17 +348,30 @@ void RegionSamplingThread::captureSample() {
     constexpr bool kGrayscale = false;
     constexpr bool kIsProtected = false;
 
-    if (const auto fenceResult =
-                mFlinger.captureScreenshot(SurfaceFlinger::RenderAreaBuilderVariant(
-                                                   std::in_place_type<DisplayRenderAreaBuilder>,
-                                                   sampledBounds, sampledBounds.getSize(),
-                                                   ui::Dataspace::V0_SRGB,
-                                                   kHintForSeamlessTransition,
-                                                   true /* captureSecureLayers */, displayWeak),
-                                           getLayerSnapshotsFn, buffer, kRegionSampling, kGrayscale,
-                                           kIsProtected, nullptr)
+    SurfaceFlinger::RenderAreaBuilderVariant
+            renderAreaBuilder(std::in_place_type<DisplayRenderAreaBuilder>, sampledBounds,
+                              sampledBounds.getSize(), ui::Dataspace::V0_SRGB,
+                              kHintForSeamlessTransition, true /* captureSecureLayers */,
+                              displayWeak);
+
+    FenceResult fenceResult;
+    if (FlagManager::getInstance().single_hop_screenshot() &&
+        FlagManager::getInstance().ce_fence_promise()) {
+        std::vector<sp<LayerFE>> layerFEs;
+        auto displayState =
+                mFlinger.getDisplayAndLayerSnapshotsFromMainThread(renderAreaBuilder,
+                                                                   getLayerSnapshotsFn, layerFEs);
+        fenceResult =
+                mFlinger.captureScreenshot(renderAreaBuilder, buffer, kRegionSampling, kGrayscale,
+                                           kIsProtected, nullptr, displayState, layerFEs)
                         .get();
-        fenceResult.ok()) {
+    } else {
+        fenceResult =
+                mFlinger.captureScreenshotLegacy(renderAreaBuilder, getLayerSnapshotsFn, buffer,
+                                                 kRegionSampling, kGrayscale, kIsProtected, nullptr)
+                        .get();
+    }
+    if (fenceResult.ok()) {
         fenceResult.value()->waitForever(LOG_TAG);
     }
 
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index fb350107d8..a48a1b3a0b 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -8110,12 +8110,15 @@ void SurfaceFlinger::attachReleaseFenceFutureToLayer(Layer* layer, LayerFE* laye
     owningLayer->prepareReleaseCallbacks(std::move(futureFence), layerStack);
 }
 
-bool SurfaceFlinger::layersHasProtectedLayer(
-        const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers) const {
+// Loop over all visible layers to see whether there's any protected layer. A protected layer is
+// typically a layer with DRM contents, or have the GRALLOC_USAGE_PROTECTED set on the buffer.
+// A protected layer has no implication on whether it's secure, which is explicitly set by
+// application to avoid being screenshot or drawn via unsecure display.
+bool SurfaceFlinger::layersHasProtectedLayer(const std::vector<sp<LayerFE>>& layers) const {
     bool protectedLayerFound = false;
-    for (auto& [_, layerFe] : layers) {
+    for (auto& layerFE : layers) {
         protectedLayerFound |=
-                (layerFe->mSnapshot->isVisible && layerFe->mSnapshot->hasProtectedContent);
+                (layerFE->mSnapshot->isVisible && layerFE->mSnapshot->hasProtectedContent);
         if (protectedLayerFound) {
             break;
         }
@@ -8123,6 +8126,26 @@ bool SurfaceFlinger::layersHasProtectedLayer(
     return protectedLayerFound;
 }
 
+// Getting layer snapshots and display should take place on main thread.
+// Accessing display requires mStateLock, and contention for this lock
+// is reduced when grabbed from the main thread, thus also reducing
+// risk of deadlocks.
+std::optional<SurfaceFlinger::OutputCompositionState>
+SurfaceFlinger::getDisplayAndLayerSnapshotsFromMainThread(
+        RenderAreaBuilderVariant& renderAreaBuilder, GetLayerSnapshotsFunction getLayerSnapshotsFn,
+        std::vector<sp<LayerFE>>& layerFEs) {
+    return mScheduler
+            ->schedule([=, this, &renderAreaBuilder, &layerFEs]() REQUIRES(kMainThreadContext) {
+                auto layers = getLayerSnapshotsFn();
+                for (auto& [layer, layerFE] : layers) {
+                    attachReleaseFenceFutureToLayer(layer, layerFE.get(), ui::INVALID_LAYER_STACK);
+                }
+                layerFEs = extractLayerFEs(layers);
+                return getDisplayStateFromRenderAreaBuilder(renderAreaBuilder);
+            })
+            .get();
+}
+
 void SurfaceFlinger::captureScreenCommon(RenderAreaBuilderVariant renderAreaBuilder,
                                          GetLayerSnapshotsFunction getLayerSnapshotsFn,
                                          ui::Size bufferSize, ui::PixelFormat reqPixelFormat,
@@ -8138,47 +8161,85 @@ void SurfaceFlinger::captureScreenCommon(RenderAreaBuilderVariant renderAreaBuil
         return;
     }
 
-    // Loop over all visible layers to see whether there's any protected layer. A protected layer is
-    // typically a layer with DRM contents, or have the GRALLOC_USAGE_PROTECTED set on the buffer.
-    // A protected layer has no implication on whether it's secure, which is explicitly set by
-    // application to avoid being screenshot or drawn via unsecure display.
-    const bool supportsProtected = getRenderEngine().supportsProtectedContent();
-    bool hasProtectedLayer = false;
-    if (allowProtected && supportsProtected) {
-        auto layers = mScheduler->schedule([=]() { return getLayerSnapshotsFn(); }).get();
-        hasProtectedLayer = layersHasProtectedLayer(layers);
-    }
-    const bool isProtected = hasProtectedLayer && allowProtected && supportsProtected;
-    const uint32_t usage = GRALLOC_USAGE_HW_COMPOSER | GRALLOC_USAGE_HW_RENDER |
-            GRALLOC_USAGE_HW_TEXTURE |
-            (isProtected ? GRALLOC_USAGE_PROTECTED
-                         : GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN);
-    sp<GraphicBuffer> buffer =
-            getFactory().createGraphicBuffer(bufferSize.getWidth(), bufferSize.getHeight(),
-                                             static_cast<android_pixel_format>(reqPixelFormat),
-                                             1 /* layerCount */, usage, "screenshot");
-
-    const status_t bufferStatus = buffer->initCheck();
-    if (bufferStatus != OK) {
-        // Animations may end up being really janky, but don't crash here.
-        // Otherwise an irreponsible process may cause an SF crash by allocating
-        // too much.
-        ALOGE("%s: Buffer failed to allocate: %d", __func__, bufferStatus);
-        invokeScreenCaptureError(bufferStatus, captureListener);
-        return;
+    if (FlagManager::getInstance().single_hop_screenshot() &&
+        FlagManager::getInstance().ce_fence_promise()) {
+        std::vector<sp<LayerFE>> layerFEs;
+        auto displayState =
+                getDisplayAndLayerSnapshotsFromMainThread(renderAreaBuilder, getLayerSnapshotsFn,
+                                                          layerFEs);
+
+        const bool supportsProtected = getRenderEngine().supportsProtectedContent();
+        bool hasProtectedLayer = false;
+        if (allowProtected && supportsProtected) {
+            hasProtectedLayer = layersHasProtectedLayer(layerFEs);
+        }
+        const bool isProtected = hasProtectedLayer && allowProtected && supportsProtected;
+        const uint32_t usage = GRALLOC_USAGE_HW_COMPOSER | GRALLOC_USAGE_HW_RENDER |
+                GRALLOC_USAGE_HW_TEXTURE |
+                (isProtected ? GRALLOC_USAGE_PROTECTED
+                             : GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN);
+        sp<GraphicBuffer> buffer =
+                getFactory().createGraphicBuffer(bufferSize.getWidth(), bufferSize.getHeight(),
+                                                 static_cast<android_pixel_format>(reqPixelFormat),
+                                                 1 /* layerCount */, usage, "screenshot");
+
+        const status_t bufferStatus = buffer->initCheck();
+        if (bufferStatus != OK) {
+            // Animations may end up being really janky, but don't crash here.
+            // Otherwise an irreponsible process may cause an SF crash by allocating
+            // too much.
+            ALOGE("%s: Buffer failed to allocate: %d", __func__, bufferStatus);
+            invokeScreenCaptureError(bufferStatus, captureListener);
+            return;
+        }
+        const std::shared_ptr<renderengine::ExternalTexture> texture = std::make_shared<
+                renderengine::impl::ExternalTexture>(buffer, getRenderEngine(),
+                                                     renderengine::impl::ExternalTexture::Usage::
+                                                             WRITEABLE);
+        auto futureFence =
+                captureScreenshot(renderAreaBuilder, texture, false /* regionSampling */, grayscale,
+                                  isProtected, captureListener, displayState, layerFEs);
+        futureFence.get();
+
+    } else {
+        const bool supportsProtected = getRenderEngine().supportsProtectedContent();
+        bool hasProtectedLayer = false;
+        if (allowProtected && supportsProtected) {
+            auto layers = mScheduler->schedule([=]() { return getLayerSnapshotsFn(); }).get();
+            hasProtectedLayer = layersHasProtectedLayer(extractLayerFEs(layers));
+        }
+        const bool isProtected = hasProtectedLayer && allowProtected && supportsProtected;
+        const uint32_t usage = GRALLOC_USAGE_HW_COMPOSER | GRALLOC_USAGE_HW_RENDER |
+                GRALLOC_USAGE_HW_TEXTURE |
+                (isProtected ? GRALLOC_USAGE_PROTECTED
+                             : GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN);
+        sp<GraphicBuffer> buffer =
+                getFactory().createGraphicBuffer(bufferSize.getWidth(), bufferSize.getHeight(),
+                                                 static_cast<android_pixel_format>(reqPixelFormat),
+                                                 1 /* layerCount */, usage, "screenshot");
+
+        const status_t bufferStatus = buffer->initCheck();
+        if (bufferStatus != OK) {
+            // Animations may end up being really janky, but don't crash here.
+            // Otherwise an irreponsible process may cause an SF crash by allocating
+            // too much.
+            ALOGE("%s: Buffer failed to allocate: %d", __func__, bufferStatus);
+            invokeScreenCaptureError(bufferStatus, captureListener);
+            return;
+        }
+        const std::shared_ptr<renderengine::ExternalTexture> texture = std::make_shared<
+                renderengine::impl::ExternalTexture>(buffer, getRenderEngine(),
+                                                     renderengine::impl::ExternalTexture::Usage::
+                                                             WRITEABLE);
+        auto futureFence = captureScreenshotLegacy(renderAreaBuilder, getLayerSnapshotsFn, texture,
+                                                   false /* regionSampling */, grayscale,
+                                                   isProtected, captureListener);
+        futureFence.get();
     }
-    const std::shared_ptr<renderengine::ExternalTexture> texture = std::make_shared<
-            renderengine::impl::ExternalTexture>(buffer, getRenderEngine(),
-                                                 renderengine::impl::ExternalTexture::Usage::
-                                                         WRITEABLE);
-    auto futureFence =
-            captureScreenshot(renderAreaBuilder, getLayerSnapshotsFn, texture,
-                              false /* regionSampling */, grayscale, isProtected, captureListener);
-    futureFence.get();
 }
 
-const sp<const DisplayDevice> SurfaceFlinger::getRenderAreaDisplay(
-        RenderAreaBuilderVariant& renderAreaBuilder, OutputCompositionState& state) {
+std::optional<SurfaceFlinger::OutputCompositionState>
+SurfaceFlinger::getDisplayStateFromRenderAreaBuilder(RenderAreaBuilderVariant& renderAreaBuilder) {
     sp<const DisplayDevice> display = nullptr;
     {
         Mutex::Autolock lock(mStateLock);
@@ -8207,24 +8268,64 @@ const sp<const DisplayDevice> SurfaceFlinger::getRenderAreaDisplay(
         }
 
         if (display != nullptr) {
-            state = display->getCompositionDisplay()->getState();
+            return std::optional{display->getCompositionDisplay()->getState()};
         }
     }
-    return display;
+    return std::nullopt;
 }
 
-std::vector<std::pair<Layer*, sp<android::LayerFE>>>
-SurfaceFlinger::getLayerSnapshotsFromMainThread(GetLayerSnapshotsFunction getLayerSnapshotsFn) {
-    auto layers = getLayerSnapshotsFn();
-    if (FlagManager::getInstance().ce_fence_promise()) {
-        for (auto& [layer, layerFE] : layers) {
-            attachReleaseFenceFutureToLayer(layer, layerFE.get(), ui::INVALID_LAYER_STACK);
-        }
+std::vector<sp<LayerFE>> SurfaceFlinger::extractLayerFEs(
+        const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers) const {
+    std::vector<sp<LayerFE>> layerFEs;
+    layerFEs.reserve(layers.size());
+    for (const auto& [_, layerFE] : layers) {
+        layerFEs.push_back(layerFE);
     }
-    return layers;
+    return layerFEs;
 }
 
 ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
+        const RenderAreaBuilderVariant& renderAreaBuilder,
+        const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
+        bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener,
+        std::optional<OutputCompositionState>& displayState, std::vector<sp<LayerFE>>& layerFEs) {
+    ATRACE_CALL();
+
+    ScreenCaptureResults captureResults;
+    std::unique_ptr<const RenderArea> renderArea =
+            std::visit([](auto&& arg) -> std::unique_ptr<RenderArea> { return arg.build(); },
+                       renderAreaBuilder);
+
+    if (!renderArea) {
+        ALOGW("Skipping screen capture because of invalid render area.");
+        if (captureListener) {
+            captureResults.fenceResult = base::unexpected(NO_MEMORY);
+            captureListener->onScreenCaptureCompleted(captureResults);
+        }
+        return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
+    }
+
+    // Empty vector needed to pass into renderScreenImpl for legacy path
+    std::vector<std::pair<Layer*, sp<android::LayerFE>>> layers;
+    ftl::SharedFuture<FenceResult> renderFuture =
+            renderScreenImpl(std::move(renderArea), buffer, regionSampling, grayscale, isProtected,
+                             captureResults, displayState, layers, layerFEs);
+
+    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;
+}
+
+ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshotLegacy(
         RenderAreaBuilderVariant renderAreaBuilder, GetLayerSnapshotsFunction getLayerSnapshotsFn,
         const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
         bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener) {
@@ -8232,10 +8333,13 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
 
     auto takeScreenshotFn = [=, this, renderAreaBuilder = std::move(renderAreaBuilder)]() REQUIRES(
                                     kMainThreadContext) mutable -> ftl::SharedFuture<FenceResult> {
-        auto layers = getLayerSnapshotsFromMainThread(getLayerSnapshotsFn);
-
-        OutputCompositionState state;
-        const auto display = getRenderAreaDisplay(renderAreaBuilder, state);
+        auto layers = getLayerSnapshotsFn();
+        if (FlagManager::getInstance().ce_fence_promise()) {
+            for (auto& [layer, layerFE] : layers) {
+                attachReleaseFenceFutureToLayer(layer, layerFE.get(), ui::INVALID_LAYER_STACK);
+            }
+        }
+        auto displayState = getDisplayStateFromRenderAreaBuilder(renderAreaBuilder);
 
         ScreenCaptureResults captureResults;
         std::unique_ptr<const RenderArea> renderArea =
@@ -8251,9 +8355,10 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
             return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
         }
 
+        auto layerFEs = extractLayerFEs(layers);
         ftl::SharedFuture<FenceResult> renderFuture =
                 renderScreenImpl(std::move(renderArea), buffer, regionSampling, grayscale,
-                                 isProtected, captureResults, display, state, layers);
+                                 isProtected, captureResults, displayState, layers, layerFEs);
 
         if (captureListener) {
             // Defer blocking on renderFuture back to the Binder thread.
@@ -8286,11 +8391,11 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
         std::unique_ptr<const RenderArea> renderArea,
         const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
         bool grayscale, bool isProtected, ScreenCaptureResults& captureResults,
-        const sp<const DisplayDevice> display, const OutputCompositionState& state,
-        std::vector<std::pair<Layer*, sp<android::LayerFE>>>& layers) {
+        std::optional<OutputCompositionState>& displayState,
+        std::vector<std::pair<Layer*, sp<LayerFE>>>& layers, std::vector<sp<LayerFE>>& layerFEs) {
     ATRACE_CALL();
 
-    for (auto& [_, layerFE] : layers) {
+    for (auto& layerFE : layerFEs) {
         frontend::LayerSnapshot* snapshot = layerFE->mSnapshot.get();
         captureResults.capturedSecureLayers |= (snapshot->isVisible && snapshot->isSecure);
         captureResults.capturedHdrLayers |= isHdrLayer(*snapshot);
@@ -8313,7 +8418,8 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
     const bool enableLocalTonemapping = FlagManager::getInstance().local_tonemap_screenshots() &&
             !renderArea->getHintForSeamlessTransition();
 
-    if (display != nullptr) {
+    if (displayState) {
+        const auto& state = displayState.value();
         captureResults.capturedDataspace =
                 pickBestDataspace(requestedDataspace, state, captureResults.capturedHdrLayers,
                                   renderArea->getHintForSeamlessTransition());
@@ -8348,18 +8454,18 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
     captureResults.buffer = capturedBuffer->getBuffer();
 
     ui::LayerStack layerStack{ui::DEFAULT_LAYER_STACK};
-    if (!layers.empty()) {
-        const sp<LayerFE>& layerFE = layers.back().second;
+    if (!layerFEs.empty()) {
+        const sp<LayerFE>& layerFE = layerFEs.back();
         layerStack = layerFE->getCompositionState()->outputFilter.layerStack;
     }
 
-    auto copyLayerFEs = [&layers]() {
-        std::vector<sp<compositionengine::LayerFE>> layerFEs;
-        layerFEs.reserve(layers.size());
-        for (const auto& [_, layerFE] : layers) {
-            layerFEs.push_back(layerFE);
+    auto copyLayerFEs = [&layerFEs]() {
+        std::vector<sp<compositionengine::LayerFE>> ceLayerFEs;
+        ceLayerFEs.reserve(layerFEs.size());
+        for (const auto& layerFE : layerFEs) {
+            ceLayerFEs.push_back(layerFE);
         }
-        return layerFEs;
+        return ceLayerFEs;
     };
 
     auto present = [this, buffer = capturedBuffer, dataspace = captureResults.capturedDataspace,
@@ -8428,8 +8534,16 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
     //
     // TODO(b/196334700) Once we use RenderEngineThreaded everywhere we can always defer the call
     // to CompositionEngine::present.
-    auto presentFuture = mRenderEngine->isThreaded() ? ftl::defer(std::move(present)).share()
-                                                     : ftl::yield(present()).share();
+    ftl::SharedFuture<FenceResult> presentFuture;
+    if (FlagManager::getInstance().single_hop_screenshot() &&
+        FlagManager::getInstance().ce_fence_promise()) {
+        presentFuture = mRenderEngine->isThreaded()
+                ? ftl::yield(present()).share()
+                : mScheduler->schedule(std::move(present)).share();
+    } else {
+        presentFuture = mRenderEngine->isThreaded() ? ftl::defer(std::move(present)).share()
+                                                    : ftl::yield(present()).share();
+    }
 
     if (!FlagManager::getInstance().ce_fence_promise()) {
         for (auto& [layer, layerFE] : layers) {
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 209d9bcfe6..12307172f7 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -892,22 +892,35 @@ private:
     void attachReleaseFenceFutureToLayer(Layer* layer, LayerFE* layerFE, ui::LayerStack layerStack);
 
     // Checks if a protected layer exists in a list of layers.
-    bool layersHasProtectedLayer(const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers) const;
+    bool layersHasProtectedLayer(const std::vector<sp<LayerFE>>& layers) const;
+
+    using OutputCompositionState = compositionengine::impl::OutputCompositionState;
+
+    std::optional<OutputCompositionState> getDisplayAndLayerSnapshotsFromMainThread(
+            RenderAreaBuilderVariant& renderAreaBuilder,
+            GetLayerSnapshotsFunction getLayerSnapshotsFn, std::vector<sp<LayerFE>>& layerFEs);
 
     void captureScreenCommon(RenderAreaBuilderVariant, GetLayerSnapshotsFunction,
                              ui::Size bufferSize, ui::PixelFormat, bool allowProtected,
                              bool grayscale, const sp<IScreenCaptureListener>&);
 
-    using OutputCompositionState = compositionengine::impl::OutputCompositionState;
-
-    const sp<const DisplayDevice> getRenderAreaDisplay(RenderAreaBuilderVariant& renderAreaBuilder,
-                                                       OutputCompositionState& state)
-            REQUIRES(kMainThreadContext);
+    std::optional<OutputCompositionState> getDisplayStateFromRenderAreaBuilder(
+            RenderAreaBuilderVariant& renderAreaBuilder) REQUIRES(kMainThreadContext);
 
-    std::vector<std::pair<Layer*, sp<android::LayerFE>>> getLayerSnapshotsFromMainThread(
-            GetLayerSnapshotsFunction getLayerSnapshotsFn) REQUIRES(kMainThreadContext);
+    // Legacy layer raw pointer is not safe to access outside the main thread.
+    // Creates a new vector consisting only of LayerFEs, which can be safely
+    // accessed outside the main thread.
+    std::vector<sp<LayerFE>> extractLayerFEs(
+            const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers) const;
 
     ftl::SharedFuture<FenceResult> captureScreenshot(
+            const RenderAreaBuilderVariant& renderAreaBuilder,
+            const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
+            bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener,
+            std::optional<OutputCompositionState>& displayState,
+            std::vector<sp<LayerFE>>& layerFEs);
+
+    ftl::SharedFuture<FenceResult> captureScreenshotLegacy(
             RenderAreaBuilderVariant, GetLayerSnapshotsFunction,
             const std::shared_ptr<renderengine::ExternalTexture>&, bool regionSampling,
             bool grayscale, bool isProtected, const sp<IScreenCaptureListener>&);
@@ -916,9 +929,9 @@ private:
             std::unique_ptr<const RenderArea>,
             const std::shared_ptr<renderengine::ExternalTexture>&, bool regionSampling,
             bool grayscale, bool isProtected, ScreenCaptureResults&,
-            const sp<const DisplayDevice> display, const OutputCompositionState& state,
-            std::vector<std::pair<Layer*, sp<android::LayerFE>>>& layers)
-            REQUIRES(kMainThreadContext);
+            std::optional<OutputCompositionState>& displayState,
+            std::vector<std::pair<Layer*, sp<LayerFE>>>& layers,
+            std::vector<sp<LayerFE>>& layerFEs);
 
     // 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/common/FlagManager.cpp b/services/surfaceflinger/common/FlagManager.cpp
index b7ec6e0b26..4216771515 100644
--- a/services/surfaceflinger/common/FlagManager.cpp
+++ b/services/surfaceflinger/common/FlagManager.cpp
@@ -150,6 +150,7 @@ void FlagManager::dump(std::string& result) const {
     DUMP_READ_ONLY_FLAG(override_trusted_overlay);
     DUMP_READ_ONLY_FLAG(flush_buffer_slots_to_uncache);
     DUMP_READ_ONLY_FLAG(force_compile_graphite_renderengine);
+    DUMP_READ_ONLY_FLAG(single_hop_screenshot);
 
 #undef DUMP_READ_ONLY_FLAG
 #undef DUMP_SERVER_FLAG
@@ -250,6 +251,7 @@ FLAG_MANAGER_READ_ONLY_FLAG(local_tonemap_screenshots, "debug.sf.local_tonemap_s
 FLAG_MANAGER_READ_ONLY_FLAG(override_trusted_overlay, "");
 FLAG_MANAGER_READ_ONLY_FLAG(flush_buffer_slots_to_uncache, "");
 FLAG_MANAGER_READ_ONLY_FLAG(force_compile_graphite_renderengine, "");
+FLAG_MANAGER_READ_ONLY_FLAG(single_hop_screenshot, "");
 
 /// Trunk stable server flags ///
 FLAG_MANAGER_SERVER_FLAG(refresh_rate_overlay_on_external_display, "")
diff --git a/services/surfaceflinger/common/include/common/FlagManager.h b/services/surfaceflinger/common/include/common/FlagManager.h
index 8f98ed3777..22118ab585 100644
--- a/services/surfaceflinger/common/include/common/FlagManager.h
+++ b/services/surfaceflinger/common/include/common/FlagManager.h
@@ -89,6 +89,7 @@ public:
     bool override_trusted_overlay() const;
     bool flush_buffer_slots_to_uncache() const;
     bool force_compile_graphite_renderengine() const;
+    bool single_hop_screenshot() const;
 
 protected:
     // overridden for unit tests
diff --git a/services/surfaceflinger/surfaceflinger_flags_new.aconfig b/services/surfaceflinger/surfaceflinger_flags_new.aconfig
index ee12325705..f4d4ee978d 100644
--- a/services/surfaceflinger/surfaceflinger_flags_new.aconfig
+++ b/services/surfaceflinger/surfaceflinger_flags_new.aconfig
@@ -103,6 +103,17 @@ flag {
     is_fixed_read_only: true
 } # local_tonemap_screenshots
 
+flag {
+  name: "single_hop_screenshot"
+  namespace: "window_surfaces"
+  description: "Only access SF main thread once during a screenshot"
+  bug: "285553970"
+  is_fixed_read_only: true
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+ } # single_hop_screenshot
+
  flag {
   name: "override_trusted_overlay"
   namespace: "window_surfaces"
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 7889096fe2..8c72a7dba5 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -494,12 +494,13 @@ public:
         ftl::FakeGuard guard(kMainThreadContext);
 
         ScreenCaptureResults captureResults;
-        SurfaceFlinger::OutputCompositionState state = display->getCompositionDisplay()->getState();
-        auto layers = mFlinger->getLayerSnapshotsFromMainThread(getLayerSnapshotsFn);
+        auto displayState = std::optional{display->getCompositionDisplay()->getState()};
+        auto layers = getLayerSnapshotsFn();
+        auto layerFEs = mFlinger->extractLayerFEs(layers);
 
         return mFlinger->renderScreenImpl(std::move(renderArea), buffer, regionSampling,
                                           false /* grayscale */, false /* isProtected */,
-                                          captureResults, display, state, layers);
+                                          captureResults, displayState, layers, layerFEs);
     }
 
     auto traverseLayersInLayerStack(ui::LayerStack layerStack, int32_t uid,