diff --git a/services/surfaceflinger/Tracing/LayerTracing.cpp b/services/surfaceflinger/Tracing/LayerTracing.cpp
index 657df16..403e105 100644
--- a/services/surfaceflinger/Tracing/LayerTracing.cpp
+++ b/services/surfaceflinger/Tracing/LayerTracing.cpp
@@ -36,6 +36,10 @@
     LayerDataSource::Initialize(*this);
 }
 
+LayerTracing::LayerTracing(std::ostream& outStream) : LayerTracing() {
+    mOutStream = std::ref(outStream);
+}
+
 LayerTracing::~LayerTracing() {
     LayerDataSource::UnregisterLayerTracing();
 }
@@ -49,10 +53,6 @@
     mTransactionTracing = &transactionTracing;
 }
 
-void LayerTracing::setOutputStream(std::ostream& outStream) {
-    mOutStream = std::ref(outStream);
-}
-
 void LayerTracing::onStart(Mode mode, uint32_t flags) {
     switch (mode) {
         case Mode::MODE_ACTIVE: {
@@ -95,7 +95,7 @@
     }
 
     auto transactionTrace = mTransactionTracing->writeToProto();
-    LayerTraceGenerator{}.generate(transactionTrace, flags);
+    LayerTraceGenerator{}.generate(transactionTrace, flags, *this);
     ALOGD("Flushed generated tracing");
 }
 
diff --git a/services/surfaceflinger/Tracing/LayerTracing.h b/services/surfaceflinger/Tracing/LayerTracing.h
index 349cc40..fe7f06d 100644
--- a/services/surfaceflinger/Tracing/LayerTracing.h
+++ b/services/surfaceflinger/Tracing/LayerTracing.h
@@ -97,11 +97,11 @@
     };
 
     LayerTracing();
+    LayerTracing(std::ostream&);
     ~LayerTracing();
     void setTakeLayersSnapshotProtoFunction(
             const std::function<perfetto::protos::LayersSnapshotProto(uint32_t)>&);
     void setTransactionTracing(TransactionTracing&);
-    void setOutputStream(std::ostream&);
 
     // Start event from perfetto data source
     void onStart(Mode mode, uint32_t flags);
diff --git a/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp b/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp
index 62c362e..23fe8fa 100644
--- a/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp
+++ b/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp
@@ -41,8 +41,7 @@
 using namespace ftl::flag_operators;
 
 bool LayerTraceGenerator::generate(const perfetto::protos::TransactionTraceFile& traceFile,
-                                   std::uint32_t traceFlags,
-                                   std::optional<std::reference_wrapper<std::ostream>> outStream,
+                                   std::uint32_t traceFlags, LayerTracing& layerTracing,
                                    bool onlyLastEntry) {
     if (traceFile.entry_size() == 0) {
         ALOGD("Trace file is empty");
@@ -51,11 +50,6 @@
 
     TransactionProtoParser parser(std::make_unique<TransactionProtoParser::FlingerDataMapper>());
 
-    LayerTracing layerTracing;
-    if (outStream) {
-        layerTracing.setOutputStream(outStream->get());
-    }
-
     // frontend
     frontend::LayerLifecycleManager lifecycleManager;
     frontend::LayerHierarchyBuilder hierarchyBuilder{{}};
diff --git a/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.h b/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.h
index 2bb6f51..e4d02ca 100644
--- a/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.h
+++ b/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.h
@@ -30,7 +30,6 @@
 class LayerTraceGenerator {
 public:
     bool generate(const perfetto::protos::TransactionTraceFile&, std::uint32_t traceFlags,
-                  std::optional<std::reference_wrapper<std::ostream>> outStream = std::nullopt,
-                  bool onlyLastEntry = false);
+                  LayerTracing& layerTracing, bool onlyLastEntry = false);
 };
 } // namespace android
diff --git a/services/surfaceflinger/Tracing/tools/main.cpp b/services/surfaceflinger/Tracing/tools/main.cpp
index a8ac36a..65e8479 100644
--- a/services/surfaceflinger/Tracing/tools/main.cpp
+++ b/services/surfaceflinger/Tracing/tools/main.cpp
@@ -52,6 +52,8 @@
             (argc == 3) ? argv[2] : "/data/misc/wmtrace/layers_trace.winscope";
     auto outStream = std::ofstream{outputLayersTracePath, std::ios::binary | std::ios::app};
 
+    auto layerTracing = LayerTracing{outStream};
+
     const bool generateLastEntryOnly =
             argc >= 4 && std::string_view(argv[3]) == "--last-entry-only";
 
@@ -60,7 +62,7 @@
     ALOGD("Generating %s...", outputLayersTracePath);
     std::cout << "Generating " << outputLayersTracePath << "\n";
 
-    if (!LayerTraceGenerator().generate(transactionTraceFile, traceFlags, outStream,
+    if (!LayerTraceGenerator().generate(transactionTraceFile, traceFlags, layerTracing,
                                         generateLastEntryOnly)) {
         std::cout << "Error: Failed to generate layers trace " << outputLayersTracePath;
         return -1;