diff --git a/services/surfaceflinger/SurfaceTracing.cpp b/services/surfaceflinger/SurfaceTracing.cpp
index db78f1d..c4ab066 100644
--- a/services/surfaceflinger/SurfaceTracing.cpp
+++ b/services/surfaceflinger/SurfaceTracing.cpp
@@ -28,46 +28,56 @@
 
 namespace android {
 
-void SurfaceTracing::mainLoop() {
-    bool enabled = true;
-    // Upon activation, logs the first frame
-    traceLayers("tracing.enable");
-    do {
-        std::unique_lock<std::mutex> sfLock(mFlinger.mDrawingStateLock);
-        mConditionalVariable.wait(sfLock);
-        LayersTraceProto entry = traceLayersLocked(mWhere);
-        sfLock.unlock();
-        {
-            std::scoped_lock bufferLock(mTraceLock);
-            mBuffer.emplace(std::move(entry));
-            if (mWriteToFile) {
-                writeProtoFileLocked();
-                mWriteToFile = false;
-            }
+SurfaceTracing::SurfaceTracing(SurfaceFlinger& flinger)
+      : mFlinger(flinger), mSfLock(flinger.mDrawingStateLock) {}
 
-            enabled = mEnabled;
-        }
-    } while (enabled);
+void SurfaceTracing::mainLoop() {
+    addFirstEntry();
+    bool enabled = true;
+    while (enabled) {
+        LayersTraceProto entry = traceWhenNotified();
+        enabled = addTraceToBuffer(entry);
+    }
 }
 
-void SurfaceTracing::traceLayers(const char* where) {
-    std::unique_lock<std::mutex> sfLock(mFlinger.mDrawingStateLock);
-    LayersTraceProto entry = traceLayersLocked(where);
-    sfLock.unlock();
-    std::scoped_lock bufferLock(mTraceLock);
+void SurfaceTracing::addFirstEntry() {
+    LayersTraceProto entry;
+    {
+        std::scoped_lock lock(mSfLock);
+        entry = traceLayersLocked("tracing.enable");
+    }
+    addTraceToBuffer(entry);
+}
+
+LayersTraceProto SurfaceTracing::traceWhenNotified() {
+    std::unique_lock<std::mutex> lock(mSfLock);
+    mCanStartTrace.wait(lock);
+    android::base::ScopedLockAssertion assumeLock(mSfLock);
+    LayersTraceProto entry = traceLayersLocked(mWhere);
+    lock.unlock();
+    return entry;
+}
+
+bool SurfaceTracing::addTraceToBuffer(LayersTraceProto& entry) {
+    std::scoped_lock lock(mTraceLock);
     mBuffer.emplace(std::move(entry));
+    if (mWriteToFile) {
+        writeProtoFileLocked();
+        mWriteToFile = false;
+    }
+    return mEnabled;
 }
 
 void SurfaceTracing::notify(const char* where) {
-    std::lock_guard<std::mutex> sfLock(mFlinger.mDrawingStateLock);
+    std::scoped_lock lock(mSfLock);
     mWhere = where;
-    mConditionalVariable.notify_one();
+    mCanStartTrace.notify_one();
 }
 
 void SurfaceTracing::writeToFileAsync() {
-    std::lock_guard<std::mutex> bufferLock(mTraceLock);
+    std::scoped_lock lock(mTraceLock);
     mWriteToFile = true;
-    mConditionalVariable.notify_one();
+    mCanStartTrace.notify_one();
 }
 
 void SurfaceTracing::LayersTraceBuffer::reset(size_t newSize) {
@@ -102,12 +112,11 @@
 }
 
 void SurfaceTracing::enable() {
-    std::lock_guard<std::mutex> bufferLock(mTraceLock);
+    std::scoped_lock lock(mTraceLock);
 
     if (mEnabled) {
         return;
     }
-
     mBuffer.reset(mBufferSize);
     mEnabled = true;
     mThread = std::thread(&SurfaceTracing::mainLoop, this);
@@ -119,7 +128,7 @@
 }
 
 bool SurfaceTracing::disable() {
-    std::lock_guard<std::mutex> bufferLock(mTraceLock);
+    std::scoped_lock lock(mTraceLock);
 
     if (!mEnabled) {
         return false;
@@ -127,28 +136,33 @@
 
     mEnabled = false;
     mWriteToFile = true;
-    mConditionalVariable.notify_all();
+    mCanStartTrace.notify_all();
     return true;
 }
 
 bool SurfaceTracing::isEnabled() const {
-    std::lock_guard<std::mutex> bufferLock(mTraceLock);
+    std::scoped_lock lock(mTraceLock);
     return mEnabled;
 }
 
 void SurfaceTracing::setBufferSize(size_t bufferSizeInByte) {
-    std::lock_guard<std::mutex> bufferLock(mTraceLock);
+    std::scoped_lock lock(mTraceLock);
     mBufferSize = bufferSizeInByte;
     mBuffer.setSize(bufferSizeInByte);
 }
 
+void SurfaceTracing::setTraceFlags(uint32_t flags) {
+    std::scoped_lock lock(mSfLock);
+    mTraceFlags = flags;
+}
+
 LayersTraceProto SurfaceTracing::traceLayersLocked(const char* where) {
     ATRACE_CALL();
 
     LayersTraceProto entry;
     entry.set_elapsed_realtime_nanos(elapsedRealtimeNano());
     entry.set_where(where);
-    LayersProto layers(mFlinger.dumpProtoInfo(LayerVector::StateSet::Drawing));
+    LayersProto layers(mFlinger.dumpProtoInfo(LayerVector::StateSet::Drawing, mTraceFlags));
     entry.mutable_layers()->Swap(&layers);
 
     return entry;
@@ -179,8 +193,7 @@
 }
 
 void SurfaceTracing::dump(std::string& result) const {
-    std::lock_guard<std::mutex> bufferLock(mTraceLock);
-
+    std::scoped_lock lock(mTraceLock);
     base::StringAppendF(&result, "Tracing state: %s\n", mEnabled ? "enabled" : "disabled");
     base::StringAppendF(&result, "  number of entries: %zu (%.2fMB / %.2fMB)\n",
                         mBuffer.frameCount(), float(mBuffer.used()) / float(1_MB),