diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index 755e585..cf4d215 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -2122,7 +2122,9 @@
     writeToProtoCommonState(layerProto, LayerVector::StateSet::Drawing, traceFlags);
 
     if (traceFlags & LayerTracing::TRACE_COMPOSITION) {
-        writeCompositionStateToProto(layerProto);
+        ui::LayerStack layerStack =
+                (mSnapshot) ? mSnapshot->outputFilter.layerStack : ui::INVALID_LAYER_STACK;
+        writeCompositionStateToProto(layerProto, layerStack);
     }
 
     for (const sp<Layer>& layer : mDrawingChildren) {
@@ -2132,14 +2134,15 @@
     return layerProto;
 }
 
-void Layer::writeCompositionStateToProto(LayerProto* layerProto) {
+void Layer::writeCompositionStateToProto(LayerProto* layerProto, ui::LayerStack layerStack) {
     ftl::FakeGuard guard(mFlinger->mStateLock); // Called from the main thread.
+    ftl::FakeGuard mainThreadGuard(kMainThreadContext);
 
     // Only populate for the primary display.
-    if (const auto display = mFlinger->getDefaultDisplayDeviceLocked()) {
+    if (const auto display = mFlinger->getDisplayFromLayerStack(layerStack)) {
         const auto compositionType = getCompositionType(*display);
         layerProto->set_hwc_composition_type(static_cast<HwcCompositionType>(compositionType));
-        LayerProtoHelper::writeToProto(getVisibleRegion(display.get()),
+        LayerProtoHelper::writeToProto(getVisibleRegion(display),
                                        [&]() { return layerProto->mutable_visible_region(); });
     }
 }
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index 1af648a..72c44cc 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -610,7 +610,7 @@
     bool isRemovedFromCurrentState() const;
 
     LayerProto* writeToProto(LayersProto& layersProto, uint32_t traceFlags);
-    void writeCompositionStateToProto(LayerProto* layerProto);
+    void writeCompositionStateToProto(LayerProto* layerProto, ui::LayerStack layerStack);
 
     // Write states that are modified by the main thread. This includes drawing
     // state as well as buffer data. This should be called in the main or tracing
diff --git a/services/surfaceflinger/LayerProtoHelper.cpp b/services/surfaceflinger/LayerProtoHelper.cpp
index b5ae1a7..5d92485 100644
--- a/services/surfaceflinger/LayerProtoHelper.cpp
+++ b/services/surfaceflinger/LayerProtoHelper.cpp
@@ -332,7 +332,7 @@
     if (mTraceFlags & LayerTracing::TRACE_COMPOSITION) {
         auto it = mLegacyLayers.find(layer.id);
         if (it != mLegacyLayers.end()) {
-            it->second->writeCompositionStateToProto(layerProto);
+            it->second->writeCompositionStateToProto(layerProto, snapshot->outputFilter.layerStack);
         }
     }
 
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index a271083..968115c 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -8903,6 +8903,15 @@
     static_cast<void>(mScheduler->scheduleDelayed([&]() { scheduleRepaint(); }, ms2ns(delayInMs)));
 }
 
+const DisplayDevice* SurfaceFlinger::getDisplayFromLayerStack(ui::LayerStack layerStack) {
+    for (const auto& [_, display] : mDisplays) {
+        if (display->getLayerStack() == layerStack) {
+            return display.get();
+        }
+    }
+    return nullptr;
+}
+
 } // namespace android
 
 #if defined(__gl_h_)
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 74d00dd..4199910 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -323,6 +323,8 @@
     bool mIgnoreHwcPhysicalDisplayOrientation = false;
 
     void forceFutureUpdate(int delayInMs);
+    const DisplayDevice* getDisplayFromLayerStack(ui::LayerStack)
+            REQUIRES(mStateLock, kMainThreadContext);
 
 protected:
     // We're reference counted, never destroy SurfaceFlinger directly