[SF] Don't keep sp<DisplayDevice> when doing screenshot

If a hotplug event is processed while a screenshot is taken
the system may crash because binder thread can end up holding
the last sp<> to a DisplayDevice. In this change we store a
weak pointer and promote to a strong pointer when we are on the
main thread.

Bug: 158599281
Test: atest libsurfaceflinger_unittest
Change-Id: Ica09398a48e68ec7b6bda3b88a6dadfa27b3455d
diff --git a/services/surfaceflinger/DisplayDevice.h b/services/surfaceflinger/DisplayDevice.h
index 4dabd2b..81b3ccf 100644
--- a/services/surfaceflinger/DisplayDevice.h
+++ b/services/surfaceflinger/DisplayDevice.h
@@ -59,12 +59,14 @@
 class DisplaySurface;
 } // namespace compositionengine
 
-class DisplayDevice : public LightRefBase<DisplayDevice> {
+class DisplayDevice : public RefBase {
 public:
     constexpr static float sDefaultMinLumiance = 0.0;
     constexpr static float sDefaultMaxLumiance = 500.0;
 
     explicit DisplayDevice(DisplayDeviceCreationArgs& args);
+
+    // Must be destroyed on the main thread because it may call into HWComposer.
     virtual ~DisplayDevice();
 
     std::shared_ptr<compositionengine::Display> getCompositionDisplay() const {
@@ -246,21 +248,18 @@
 
 class DisplayRenderArea : public RenderArea {
 public:
-    DisplayRenderArea(const sp<const DisplayDevice>& display,
-                      RotationFlags rotation = ui::Transform::ROT_0)
-          : DisplayRenderArea(display, display->getBounds(),
-                              static_cast<uint32_t>(display->getWidth()),
-                              static_cast<uint32_t>(display->getHeight()),
-                              display->getCompositionDataSpace(), rotation) {}
-
-    DisplayRenderArea(sp<const DisplayDevice> display, const Rect& sourceCrop, uint32_t reqWidth,
-                      uint32_t reqHeight, ui::Dataspace reqDataSpace, RotationFlags rotation,
-                      bool allowSecureLayers = true)
-          : RenderArea(reqWidth, reqHeight, CaptureFill::OPAQUE, reqDataSpace,
-                       display->getViewport(), applyDeviceOrientation(rotation, display)),
-            mDisplay(std::move(display)),
-            mSourceCrop(sourceCrop),
-            mAllowSecureLayers(allowSecureLayers) {}
+    static std::unique_ptr<RenderArea> create(wp<const DisplayDevice> displayWeak,
+                                              const Rect& sourceCrop, ui::Size reqSize,
+                                              ui::Dataspace reqDataSpace, RotationFlags rotation,
+                                              bool allowSecureLayers = true) {
+        if (auto display = displayWeak.promote()) {
+            // Using new to access a private constructor.
+            return std::unique_ptr<DisplayRenderArea>(
+                    new DisplayRenderArea(std::move(display), sourceCrop, reqSize, reqDataSpace,
+                                          rotation, allowSecureLayers));
+        }
+        return nullptr;
+    }
 
     const ui::Transform& getTransform() const override { return mTransform; }
     Rect getBounds() const override { return mDisplay->getBounds(); }
@@ -307,6 +306,14 @@
     }
 
 private:
+    DisplayRenderArea(sp<const DisplayDevice> display, const Rect& sourceCrop, ui::Size reqSize,
+                      ui::Dataspace reqDataSpace, RotationFlags rotation, bool allowSecureLayers)
+          : RenderArea(reqSize, CaptureFill::OPAQUE, reqDataSpace, display->getViewport(),
+                       applyDeviceOrientation(rotation, display)),
+            mDisplay(std::move(display)),
+            mSourceCrop(sourceCrop),
+            mAllowSecureLayers(allowSecureLayers) {}
+
     static RotationFlags applyDeviceOrientation(RotationFlags orientationFlag,
                                                 const sp<const DisplayDevice>& device) {
         uint32_t inverseRotate90 = 0;
diff --git a/services/surfaceflinger/DisplayHardware/HWComposer.h b/services/surfaceflinger/DisplayHardware/HWComposer.h
index c355ebd..b800038 100644
--- a/services/surfaceflinger/DisplayHardware/HWComposer.h
+++ b/services/surfaceflinger/DisplayHardware/HWComposer.h
@@ -64,6 +64,8 @@
     const uint32_t id;
 };
 
+// See the comment for SurfaceFlinger::getHwComposer for the thread safety rules for accessing
+// this class.
 class HWComposer {
 public:
     struct DeviceRequestedChanges {
diff --git a/services/surfaceflinger/RegionSamplingThread.cpp b/services/surfaceflinger/RegionSamplingThread.cpp
index 27353d8..398fd40 100644
--- a/services/surfaceflinger/RegionSamplingThread.cpp
+++ b/services/surfaceflinger/RegionSamplingThread.cpp
@@ -36,6 +36,7 @@
 
 #include "DisplayDevice.h"
 #include "Layer.h"
+#include "Promise.h"
 #include "Scheduler/DispSync.h"
 #include "SurfaceFlinger.h"
 
@@ -336,8 +337,20 @@
         return;
     }
 
-    const auto device = mFlinger.getDefaultDisplayDevice();
-    const auto orientation = ui::Transform::toRotationFlags(device->getOrientation());
+    wp<const DisplayDevice> displayWeak;
+
+    ui::LayerStack layerStack;
+    ui::Transform::RotationFlags orientation;
+    ui::Size displaySize;
+
+    {
+        // TODO(b/159112860): Don't keep sp<DisplayDevice> outside of SF main thread
+        const sp<const DisplayDevice> display = mFlinger.getDefaultDisplayDevice();
+        displayWeak = display;
+        layerStack = display->getLayerStack();
+        orientation = ui::Transform::toRotationFlags(display->getOrientation());
+        displaySize = display->getSize();
+    }
 
     std::vector<RegionSamplingThread::Descriptor> descriptors;
     Region sampleRegion;
@@ -346,20 +359,18 @@
         descriptors.emplace_back(descriptor);
     }
 
-    const Rect sampledArea = sampleRegion.bounds();
-
     auto dx = 0;
     auto dy = 0;
     switch (orientation) {
         case ui::Transform::ROT_90:
-            dx = device->getWidth();
+            dx = displaySize.getWidth();
             break;
         case ui::Transform::ROT_180:
-            dx = device->getWidth();
-            dy = device->getHeight();
+            dx = displaySize.getWidth();
+            dy = displaySize.getHeight();
             break;
         case ui::Transform::ROT_270:
-            dy = device->getHeight();
+            dy = displaySize.getHeight();
             break;
         default:
             break;
@@ -368,8 +379,13 @@
     ui::Transform t(orientation);
     auto screencapRegion = t.transform(sampleRegion);
     screencapRegion = screencapRegion.translate(dx, dy);
-    DisplayRenderArea renderArea(device, screencapRegion.bounds(), sampledArea.getWidth(),
-                                 sampledArea.getHeight(), ui::Dataspace::V0_SRGB, orientation);
+
+    const Rect sampledBounds = sampleRegion.bounds();
+
+    SurfaceFlinger::RenderAreaFuture renderAreaFuture = promise::defer([=] {
+        return DisplayRenderArea::create(displayWeak, sampledBounds, sampledBounds.getSize(),
+                                         ui::Dataspace::V0_SRGB, orientation);
+    });
 
     std::unordered_set<sp<IRegionSamplingListener>, SpHash<IRegionSamplingListener>> listeners;
 
@@ -393,9 +409,9 @@
             constexpr bool roundOutwards = true;
             Rect transformed = transform.transform(bounds, roundOutwards);
 
-            // If this layer doesn't intersect with the larger sampledArea, skip capturing it
+            // If this layer doesn't intersect with the larger sampledBounds, skip capturing it
             Rect ignore;
-            if (!transformed.intersect(sampledArea, &ignore)) return;
+            if (!transformed.intersect(sampledBounds, &ignore)) return;
 
             // If the layer doesn't intersect a sampling area, skip capturing it
             bool intersectsAnyArea = false;
@@ -411,22 +427,22 @@
                   bounds.top, bounds.right, bounds.bottom);
             visitor(layer);
         };
-        mFlinger.traverseLayersInDisplay(device, filterVisitor);
+        mFlinger.traverseLayersInLayerStack(layerStack, filterVisitor);
     };
 
     sp<GraphicBuffer> buffer = nullptr;
-    if (mCachedBuffer && mCachedBuffer->getWidth() == sampledArea.getWidth() &&
-        mCachedBuffer->getHeight() == sampledArea.getHeight()) {
+    if (mCachedBuffer && mCachedBuffer->getWidth() == sampledBounds.getWidth() &&
+        mCachedBuffer->getHeight() == sampledBounds.getHeight()) {
         buffer = mCachedBuffer;
     } else {
         const uint32_t usage = GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_HW_RENDER;
-        buffer = new GraphicBuffer(sampledArea.getWidth(), sampledArea.getHeight(),
+        buffer = new GraphicBuffer(sampledBounds.getWidth(), sampledBounds.getHeight(),
                                    PIXEL_FORMAT_RGBA_8888, 1, usage, "RegionSamplingThread");
     }
 
     bool ignored;
-    mFlinger.captureScreenCommon(renderArea, traverseLayers, buffer, false /* identityTransform */,
-                                 true /* regionSampling */, ignored);
+    mFlinger.captureScreenCommon(std::move(renderAreaFuture), traverseLayers, buffer,
+                                 false /* identityTransform */, true /* regionSampling */, ignored);
 
     std::vector<Descriptor> activeDescriptors;
     for (const auto& descriptor : descriptors) {
@@ -437,7 +453,7 @@
 
     ALOGV("Sampling %zu descriptors", activeDescriptors.size());
     std::vector<float> lumas =
-            sampleBuffer(buffer, sampledArea.leftTop(), activeDescriptors, orientation);
+            sampleBuffer(buffer, sampledBounds.leftTop(), activeDescriptors, orientation);
     if (lumas.size() != activeDescriptors.size()) {
         ALOGW("collected %zu median luma values for %zu descriptors", lumas.size(),
               activeDescriptors.size());
diff --git a/services/surfaceflinger/RenderArea.h b/services/surfaceflinger/RenderArea.h
index 6b0455a..a6246d9 100644
--- a/services/surfaceflinger/RenderArea.h
+++ b/services/surfaceflinger/RenderArea.h
@@ -23,11 +23,9 @@
 
     static float getCaptureFillValue(CaptureFill captureFill);
 
-    RenderArea(uint32_t reqWidth, uint32_t reqHeight, CaptureFill captureFill,
-               ui::Dataspace reqDataSpace, const Rect& displayViewport,
-               RotationFlags rotation = ui::Transform::ROT_0)
-          : mReqWidth(reqWidth),
-            mReqHeight(reqHeight),
+    RenderArea(ui::Size reqSize, CaptureFill captureFill, ui::Dataspace reqDataSpace,
+               const Rect& displayViewport, RotationFlags rotation = ui::Transform::ROT_0)
+          : mReqSize(reqSize),
             mReqDataSpace(reqDataSpace),
             mCaptureFill(captureFill),
             mRotationFlags(rotation),
@@ -70,8 +68,8 @@
     RotationFlags getRotationFlags() const { return mRotationFlags; }
 
     // Returns the size of the physical render area.
-    int getReqWidth() const { return static_cast<int>(mReqWidth); }
-    int getReqHeight() const { return static_cast<int>(mReqHeight); }
+    int getReqWidth() const { return mReqSize.width; }
+    int getReqHeight() const { return mReqSize.height; }
 
     // Returns the composition data space of the render area.
     ui::Dataspace getReqDataSpace() const { return mReqDataSpace; }
@@ -86,8 +84,7 @@
     const Rect& getDisplayViewport() const { return mDisplayViewport; }
 
 private:
-    const uint32_t mReqWidth;
-    const uint32_t mReqHeight;
+    const ui::Size mReqSize;
     const ui::Dataspace mReqDataSpace;
     const CaptureFill mCaptureFill;
     const RotationFlags mRotationFlags;
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 36d272a..9fd0f65 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5383,27 +5383,33 @@
         renderAreaRotation = ui::Transform::ROT_0;
     }
 
-    sp<DisplayDevice> display;
+    wp<DisplayDevice> displayWeak;
+    ui::LayerStack layerStack;
+    ui::Size reqSize(reqWidth, reqHeight);
     {
         Mutex::Autolock lock(mStateLock);
-
-        display = getDisplayDeviceLocked(displayToken);
+        sp<DisplayDevice> display = getDisplayDeviceLocked(displayToken);
         if (!display) return NAME_NOT_FOUND;
+        displayWeak = display;
+        layerStack = display->getLayerStack();
 
         // set the requested width/height to the logical display viewport size
         // by default
         if (reqWidth == 0 || reqHeight == 0) {
-            reqWidth = uint32_t(display->getViewport().width());
-            reqHeight = uint32_t(display->getViewport().height());
+            reqSize = display->getViewport().getSize();
         }
     }
 
-    DisplayRenderArea renderArea(display, sourceCrop, reqWidth, reqHeight, reqDataspace,
-                                 renderAreaRotation, captureSecureLayers);
-    auto traverseLayers = std::bind(&SurfaceFlinger::traverseLayersInDisplay, this, display,
-                                    std::placeholders::_1);
-    return captureScreenCommon(renderArea, traverseLayers, outBuffer, reqPixelFormat,
-                               useIdentityTransform, outCapturedSecureLayers);
+    RenderAreaFuture renderAreaFuture = promise::defer([=] {
+        return DisplayRenderArea::create(displayWeak, sourceCrop, reqSize, reqDataspace,
+                                         renderAreaRotation, captureSecureLayers);
+    });
+
+    auto traverseLayers = [this, layerStack](const LayerVector::Visitor& visitor) {
+        traverseLayersInLayerStack(layerStack, visitor);
+    };
+    return captureScreenCommon(std::move(renderAreaFuture), traverseLayers, reqSize, outBuffer,
+                               reqPixelFormat, useIdentityTransform, outCapturedSecureLayers);
 }
 
 static Dataspace pickDataspaceFromColorMode(const ColorMode colorMode) {
@@ -5459,19 +5465,20 @@
 
 status_t SurfaceFlinger::captureScreen(uint64_t displayOrLayerStack, Dataspace* outDataspace,
                                        sp<GraphicBuffer>* outBuffer) {
-    sp<DisplayDevice> display;
-    uint32_t width;
-    uint32_t height;
+    ui::LayerStack layerStack;
+    wp<DisplayDevice> displayWeak;
+    ui::Size size;
     ui::Transform::RotationFlags captureOrientation;
     {
         Mutex::Autolock lock(mStateLock);
-        display = getDisplayByIdOrLayerStack(displayOrLayerStack);
+        sp<DisplayDevice> display = getDisplayByIdOrLayerStack(displayOrLayerStack);
         if (!display) {
             return NAME_NOT_FOUND;
         }
+        layerStack = display->getLayerStack();
+        displayWeak = display;
 
-        width = uint32_t(display->getViewport().width());
-        height = uint32_t(display->getViewport().height());
+        size = display->getViewport().getSize();
 
         const auto orientation = display->getOrientation();
         captureOrientation = ui::Transform::toRotationFlags(orientation);
@@ -5497,14 +5504,19 @@
                 pickDataspaceFromColorMode(display->getCompositionDisplay()->getState().colorMode);
     }
 
-    DisplayRenderArea renderArea(display, Rect(), width, height, *outDataspace, captureOrientation,
-                                 false /* captureSecureLayers */);
+    RenderAreaFuture renderAreaFuture = promise::defer([=] {
+        return DisplayRenderArea::create(displayWeak, Rect(), size, *outDataspace,
+                                         captureOrientation, false /* captureSecureLayers */);
+    });
 
-    auto traverseLayers = std::bind(&SurfaceFlinger::traverseLayersInDisplay, this, display,
-                                    std::placeholders::_1);
+    auto traverseLayers = [this, layerStack](const LayerVector::Visitor& visitor) {
+        traverseLayersInLayerStack(layerStack, visitor);
+    };
+
     bool ignored = false;
-    return captureScreenCommon(renderArea, traverseLayers, outBuffer, ui::PixelFormat::RGBA_8888,
-                               false /* useIdentityTransform */,
+
+    return captureScreenCommon(std::move(renderAreaFuture), traverseLayers, size, outBuffer,
+                               ui::PixelFormat::RGBA_8888, false /* useIdentityTransform */,
                                ignored /* outCapturedSecureLayers */);
 }
 
@@ -5518,9 +5530,9 @@
     class LayerRenderArea : public RenderArea {
     public:
         LayerRenderArea(SurfaceFlinger* flinger, const sp<Layer>& layer, const Rect crop,
-                        int32_t reqWidth, int32_t reqHeight, Dataspace reqDataSpace,
-                        bool childrenOnly, const Rect& displayViewport)
-              : RenderArea(reqWidth, reqHeight, CaptureFill::CLEAR, reqDataSpace, displayViewport),
+                        ui::Size reqSize, Dataspace reqDataSpace, bool childrenOnly,
+                        const Rect& displayViewport)
+              : RenderArea(reqSize, CaptureFill::CLEAR, reqDataSpace, displayViewport),
                 mLayer(layer),
                 mCrop(crop),
                 mNeedsFiltering(false),
@@ -5595,8 +5607,7 @@
         const bool mChildrenOnly;
     };
 
-    int reqWidth = 0;
-    int reqHeight = 0;
+    ui::Size reqSize;
     sp<Layer> parent;
     Rect crop(sourceCrop);
     std::unordered_set<sp<Layer>, ISurfaceComposer::SpHash<Layer>> excludeLayers;
@@ -5633,8 +5644,7 @@
             // crop was not specified, or an invalid frame scale was provided.
             return BAD_VALUE;
         }
-        reqWidth = crop.width() * frameScale;
-        reqHeight = crop.height() * frameScale;
+        reqSize = ui::Size(crop.width() * frameScale, crop.height() * frameScale);
 
         for (const auto& handle : excludeHandles) {
             sp<Layer> excludeLayer = fromHandleLocked(handle).promote();
@@ -5655,15 +5665,18 @@
     } // mStateLock
 
     // really small crop or frameScale
-    if (reqWidth <= 0) {
-        reqWidth = 1;
+    if (reqSize.width <= 0) {
+        reqSize.width = 1;
     }
-    if (reqHeight <= 0) {
-        reqHeight = 1;
+    if (reqSize.height <= 0) {
+        reqSize.height = 1;
     }
 
-    LayerRenderArea renderArea(this, parent, crop, reqWidth, reqHeight, reqDataspace, childrenOnly,
-                               displayViewport);
+    RenderAreaFuture renderAreaFuture = promise::defer([=]() -> std::unique_ptr<RenderArea> {
+        return std::make_unique<LayerRenderArea>(this, parent, crop, reqSize, reqDataspace,
+                                                 childrenOnly, displayViewport);
+    });
+
     auto traverseLayers = [parent, childrenOnly,
                            &excludeLayers](const LayerVector::Visitor& visitor) {
         parent->traverseChildrenInZOrder(LayerVector::StateSet::Drawing, [&](Layer* layer) {
@@ -5686,14 +5699,14 @@
     };
 
     bool outCapturedSecureLayers = false;
-    return captureScreenCommon(renderArea, traverseLayers, outBuffer, reqPixelFormat, false,
-                               outCapturedSecureLayers);
+    return captureScreenCommon(std::move(renderAreaFuture), traverseLayers, reqSize, outBuffer,
+                               reqPixelFormat, false, outCapturedSecureLayers);
 }
 
-status_t SurfaceFlinger::captureScreenCommon(RenderArea& renderArea,
+status_t SurfaceFlinger::captureScreenCommon(RenderAreaFuture renderAreaFuture,
                                              TraverseLayersFunction traverseLayers,
-                                             sp<GraphicBuffer>* outBuffer,
-                                             const ui::PixelFormat reqPixelFormat,
+                                             ui::Size bufferSize, sp<GraphicBuffer>* outBuffer,
+                                             ui::PixelFormat reqPixelFormat,
                                              bool useIdentityTransform,
                                              bool& outCapturedSecureLayers) {
     ATRACE_CALL();
@@ -5701,16 +5714,16 @@
     // TODO(b/116112787) Make buffer usage a parameter.
     const uint32_t usage = GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN |
             GRALLOC_USAGE_HW_RENDER | GRALLOC_USAGE_HW_TEXTURE;
-    *outBuffer =
-            getFactory().createGraphicBuffer(renderArea.getReqWidth(), renderArea.getReqHeight(),
-                                             static_cast<android_pixel_format>(reqPixelFormat), 1,
-                                             usage, "screenshot");
+    *outBuffer = getFactory().createGraphicBuffer(bufferSize.getWidth(), bufferSize.getHeight(),
+                                                  static_cast<android_pixel_format>(reqPixelFormat),
+                                                  1, usage, "screenshot");
 
-    return captureScreenCommon(renderArea, traverseLayers, *outBuffer, useIdentityTransform,
-                               false /* regionSampling */, outCapturedSecureLayers);
+    return captureScreenCommon(std::move(renderAreaFuture), traverseLayers, *outBuffer,
+                               useIdentityTransform, false /* regionSampling */,
+                               outCapturedSecureLayers);
 }
 
-status_t SurfaceFlinger::captureScreenCommon(RenderArea& renderArea,
+status_t SurfaceFlinger::captureScreenCommon(RenderAreaFuture renderAreaFuture,
                                              TraverseLayersFunction traverseLayers,
                                              const sp<GraphicBuffer>& buffer,
                                              bool useIdentityTransform, bool regionSampling,
@@ -5723,23 +5736,28 @@
 
     do {
         std::tie(result, syncFd) =
-                schedule([&] {
+                schedule([&]() -> std::pair<status_t, int> {
                     if (mRefreshPending) {
-                        ATRACE_NAME("Skipping screenshot for now");
-                        return std::make_pair(EAGAIN, -1);
+                        ALOGW("Skipping screenshot for now");
+                        return {EAGAIN, -1};
+                    }
+                    std::unique_ptr<RenderArea> renderArea = renderAreaFuture.get();
+                    if (!renderArea) {
+                        ALOGW("Skipping screen capture because of invalid render area.");
+                        return {NO_MEMORY, -1};
                     }
 
                     status_t result = NO_ERROR;
                     int fd = -1;
 
                     Mutex::Autolock lock(mStateLock);
-                    renderArea.render([&] {
-                        result = captureScreenImplLocked(renderArea, traverseLayers, buffer.get(),
+                    renderArea->render([&] {
+                        result = captureScreenImplLocked(*renderArea, traverseLayers, buffer.get(),
                                                          useIdentityTransform, forSystem, &fd,
                                                          regionSampling, outCapturedSecureLayers);
                     });
 
-                    return std::make_pair(result, fd);
+                    return {result, fd};
                 }).get();
     } while (result == EAGAIN);
 
@@ -5898,17 +5916,17 @@
     layersSortedByZ.traverseInReverseZOrder(stateSet, visitor);
 }
 
-void SurfaceFlinger::traverseLayersInDisplay(const sp<const DisplayDevice>& display,
-                                             const LayerVector::Visitor& visitor) {
+void SurfaceFlinger::traverseLayersInLayerStack(ui::LayerStack layerStack,
+                                                const LayerVector::Visitor& visitor) {
     // We loop through the first level of layers without traversing,
     // as we need to determine which layers belong to the requested display.
     for (const auto& layer : mDrawingState.layersSortedByZ) {
-        if (!layer->belongsToDisplay(display->getLayerStack(), false)) {
+        if (!layer->belongsToDisplay(layerStack, false)) {
             continue;
         }
         // relative layers are traversed in Layer::traverseInZOrder
         layer->traverseInZOrder(LayerVector::StateSet::Drawing, [&](Layer* layer) {
-            if (!layer->belongsToDisplay(display->getLayerStack(), false)) {
+            if (!layer->belongsToDisplay(layerStack, false)) {
                 return;
             }
             if (!layer->isVisible()) {
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 8d3211b..35648ce 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -711,16 +711,17 @@
     void startBootAnim();
 
     using TraverseLayersFunction = std::function<void(const LayerVector::Visitor&)>;
+    using RenderAreaFuture = std::future<std::unique_ptr<RenderArea>>;
 
     void renderScreenImplLocked(const RenderArea& renderArea, TraverseLayersFunction traverseLayers,
                                 const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
                                 bool regionSampling, int* outSyncFd);
-    status_t captureScreenCommon(RenderArea& renderArea, TraverseLayersFunction traverseLayers,
-                                 sp<GraphicBuffer>* outBuffer, const ui::PixelFormat reqPixelFormat,
+    status_t captureScreenCommon(RenderAreaFuture, TraverseLayersFunction, ui::Size bufferSize,
+                                 sp<GraphicBuffer>* outBuffer, ui::PixelFormat,
                                  bool useIdentityTransform, bool& outCapturedSecureLayers);
-    status_t captureScreenCommon(RenderArea& renderArea, TraverseLayersFunction traverseLayers,
-                                 const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
-                                 bool regionSampling, bool& outCapturedSecureLayers);
+    status_t captureScreenCommon(RenderAreaFuture, TraverseLayersFunction, const sp<GraphicBuffer>&,
+                                 bool useIdentityTransform, bool regionSampling,
+                                 bool& outCapturedSecureLayers);
     sp<DisplayDevice> getDisplayByIdOrLayerStack(uint64_t displayOrLayerStack) REQUIRES(mStateLock);
     sp<DisplayDevice> getDisplayByLayerStack(uint64_t layerStack) REQUIRES(mStateLock);
     status_t captureScreenImplLocked(const RenderArea& renderArea,
@@ -728,8 +729,7 @@
                                      const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
                                      bool forSystem, int* outSyncFd, bool regionSampling,
                                      bool& outCapturedSecureLayers);
-    void traverseLayersInDisplay(const sp<const DisplayDevice>& display,
-                                 const LayerVector::Visitor& visitor);
+    void traverseLayersInLayerStack(ui::LayerStack, const LayerVector::Visitor&);
 
     sp<StartPropertySetThread> mStartPropertySetThread;
 
diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 95e2ef5..456891e 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -233,23 +233,22 @@
     constexpr bool forSystem = true;
     constexpr bool regionSampling = false;
 
-    DisplayRenderArea renderArea(mDisplay, sourceCrop, DEFAULT_DISPLAY_WIDTH,
-                                 DEFAULT_DISPLAY_HEIGHT, ui::Dataspace::V0_SRGB,
-                                 ui::Transform::ROT_0);
+    auto renderArea = DisplayRenderArea::create(mDisplay, sourceCrop, sourceCrop.getSize(),
+                                                ui::Dataspace::V0_SRGB, ui::Transform::ROT_0);
 
     auto traverseLayers = [this](const LayerVector::Visitor& visitor) {
-        return mFlinger.traverseLayersInDisplay(mDisplay, visitor);
+        return mFlinger.traverseLayersInLayerStack(mDisplay->getLayerStack(), visitor);
     };
 
     // TODO: Eliminate expensive/real allocation if possible.
     const uint32_t usage = GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN |
             GRALLOC_USAGE_HW_RENDER | GRALLOC_USAGE_HW_TEXTURE;
-    mCaptureScreenBuffer = new GraphicBuffer(renderArea.getReqWidth(), renderArea.getReqHeight(),
+    mCaptureScreenBuffer = new GraphicBuffer(renderArea->getReqWidth(), renderArea->getReqHeight(),
                                              HAL_PIXEL_FORMAT_RGBA_8888, 1, usage, "screenshot");
 
     int fd = -1;
     status_t result =
-            mFlinger.captureScreenImplLocked(renderArea, traverseLayers, mCaptureScreenBuffer.get(),
+            mFlinger.captureScreenImplLocked(*renderArea, traverseLayers, mCaptureScreenBuffer.get(),
                                              useIdentityTransform, forSystem, &fd, regionSampling);
     if (fd >= 0) {
         close(fd);
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 4652da0..f630103 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -341,9 +341,9 @@
                                                  regionSampling, ignored);
     }
 
-    auto traverseLayersInDisplay(const sp<const DisplayDevice>& display,
-                                 const LayerVector::Visitor& visitor) {
-        return mFlinger->SurfaceFlinger::traverseLayersInDisplay(display, visitor);
+    auto traverseLayersInLayerStack(ui::LayerStack layerStack,
+                                    const LayerVector::Visitor& visitor) {
+        return mFlinger->SurfaceFlinger::traverseLayersInLayerStack(layerStack, visitor);
     }
 
     auto getDisplayNativePrimaries(const sp<IBinder>& displayToken,