Merge changes from topics "jank-tags-faas", "transition-finish-faas", "transition-type-faas" into udc-dev

* changes:
  Write start FINISHED_TRANSITIONS token when dumping aborted transition
  Dump jank tag to event log
  Dump the transition type in always on transition trace
  Dump finish time of transitions in transitions trace

6 files changed, 212 insertions, 154 deletions

diff --git a/core/java/com/android/internal/jank/EventLogTags.logtags b/core/java/com/android/internal/jank/EventLogTags.logtags
index ad47b810aa50..66ee131badac 100644
--- a/core/java/com/android/internal/jank/EventLogTags.logtags
+++ b/core/java/com/android/internal/jank/EventLogTags.logtags
@@ -3,7 +3,7 @@
 option java_package com.android.internal.jank;
 
 # Marks a request to start tracing a CUJ. Doesn't mean the request was executed.
-37001 jank_cuj_events_begin_request (CUJ Type|1|5),(Unix Time Ns|2|3),(Elapsed Time Ns|2|3),(Uptime Ns|2|3)
+37001 jank_cuj_events_begin_request (CUJ Type|1|5),(Unix Time Ns|2|3),(Elapsed Time Ns|2|3),(Uptime Ns|2|3),(Tag|3)
 # Marks a request to end tracing a CUJ. Doesn't mean the request was executed.
 37002 jank_cuj_events_end_request (CUJ Type|1|5),(Unix Time Ns|2|3),(Elapsed Time Ns|2|3),(Uptime Time Ns|2|3)
 # Marks a request to cancel tracing a CUJ. Doesn't mean the request was executed.
diff --git a/core/java/com/android/internal/jank/InteractionJankMonitor.java b/core/java/com/android/internal/jank/InteractionJankMonitor.java
index 928a09700e2e..6344568480b7 100644
--- a/core/java/com/android/internal/jank/InteractionJankMonitor.java
+++ b/core/java/com/android/internal/jank/InteractionJankMonitor.java
@@ -590,7 +590,7 @@ public class InteractionJankMonitor {
             final Configuration config = builder.build();
             postEventLogToWorkerThread((unixNanos, elapsedNanos, realtimeNanos) -> {
                 EventLogTags.writeJankCujEventsBeginRequest(
-                        config.mCujType, unixNanos, elapsedNanos, realtimeNanos);
+                        config.mCujType, unixNanos, elapsedNanos, realtimeNanos, config.mTag);
             });
             final TrackerResult result = new TrackerResult();
             final boolean success = config.getHandler().runWithScissors(
diff --git a/core/proto/android/server/windowmanagertransitiontrace.proto b/core/proto/android/server/windowmanagertransitiontrace.proto
index ab8738444aaa..9e53a9162429 100644
--- a/core/proto/android/server/windowmanagertransitiontrace.proto
+++ b/core/proto/android/server/windowmanagertransitiontrace.proto
@@ -36,8 +36,10 @@ message TransitionTraceProto {
     MAGIC_NUMBER_H = 0x45434152;  /* RACE (little-endian ASCII) */
   }
 
-  required fixed64 magic_number = 1;  /* Must be the first field, set to value in MagicNumber */
-  repeated Transition sent_transitions = 2;
+  // Must be the first field, set to value in MagicNumber
+  required fixed64 magic_number = 1;
+  // Transitions that don't have a finish time are considered aborted
+  repeated Transition finished_transitions = 2;
 
   // Additional debugging info only collected and dumped when explicitly requested to trace
   repeated TransitionState transition_states = 3;
@@ -50,7 +52,9 @@ message Transition {
   required uint64 finish_transaction_id = 3;
   required int64 create_time_ns = 4;
   required int64 send_time_ns = 5;
-  repeated Target targets = 6;
+  optional int64 finish_time_ns = 6; // consider aborted if not provided
+  required int32 type = 7;
+  repeated Target targets = 8;
 }
 
 message Target {
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 34bf8c2839a0..873a83d5527e 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -829,6 +829,7 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
         }
         mLogger.mFinishTimeNs = SystemClock.elapsedRealtimeNanos();
         mController.mLoggerHandler.post(mLogger::logOnFinish);
+        mController.mTransitionTracer.logFinishedTransition(this);
         // Close the transactions now. They were originally copied to Shell in case we needed to
         // apply them due to a remote failure. Since we don't need to apply them anymore, free them
         // immediately.
@@ -1045,6 +1046,7 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
         }
         ProtoLog.v(ProtoLogGroup.WM_DEBUG_WINDOW_TRANSITIONS, "Aborting Transition: %d", mSyncId);
         mState = STATE_ABORT;
+        mController.mTransitionTracer.logAbortedTransition(this);
         // Syncengine abort will call through to onTransactionReady()
         mSyncEngine.abort(mSyncId);
         mController.dispatchLegacyAppTransitionCancelled();
@@ -1249,8 +1251,7 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
                         "Calling onTransitionReady: %s", info);
                 mLogger.mSendTimeNs = SystemClock.elapsedRealtimeNanos();
                 mLogger.mInfo = info;
-                mController.mTransitionTracer.logSentTransition(
-                        this, mTargets, mLogger.mCreateTimeNs, mLogger.mSendTimeNs, info);
+                mController.mTransitionTracer.logSentTransition(this, mTargets, info);
                 mController.getTransitionPlayer().onTransitionReady(
                         mToken, info, transaction, mFinishTransaction);
                 if (Trace.isTagEnabled(TRACE_TAG_WINDOW_MANAGER)) {
diff --git a/services/core/java/com/android/server/wm/TransitionController.java b/services/core/java/com/android/server/wm/TransitionController.java
index 86bb6b58d14c..c74f167c2222 100644
--- a/services/core/java/com/android/server/wm/TransitionController.java
+++ b/services/core/java/com/android/server/wm/TransitionController.java
@@ -976,6 +976,8 @@ class TransitionController {
         WindowContainerTransaction mStartWCT;
         int mSyncId;
         TransitionInfo mInfo;
+        ProtoOutputStream mProtoOutputStream = new ProtoOutputStream();
+        long mProtoToken;
 
         private String buildOnSendLog() {
             StringBuilder sb = new StringBuilder("Sent Transition #").append(mSyncId)
diff --git a/services/core/java/com/android/server/wm/TransitionTracer.java b/services/core/java/com/android/server/wm/TransitionTracer.java
index 7b1975d234c4..57c0d65b28ea 100644
--- a/services/core/java/com/android/server/wm/TransitionTracer.java
+++ b/services/core/java/com/android/server/wm/TransitionTracer.java
@@ -18,14 +18,6 @@ package com.android.server.wm;
 
 import static android.os.Build.IS_USER;
 
-import static com.android.server.wm.shell.ChangeInfo.CHANGE_FLAGS;
-import static com.android.server.wm.shell.ChangeInfo.HAS_CHANGED;
-import static com.android.server.wm.shell.ChangeInfo.TRANSIT_MODE;
-import static com.android.server.wm.shell.ChangeInfo.WINDOWING_MODE;
-import static com.android.server.wm.shell.ChangeInfo.WINDOW_IDENTIFIER;
-import static com.android.server.wm.shell.TransitionInfoChange.LAYER_ID;
-import static com.android.server.wm.shell.TransitionInfoChange.MODE;
-import static com.android.server.wm.shell.TransitionState.CHANGE;
 import static com.android.server.wm.shell.TransitionTraceProto.MAGIC_NUMBER;
 import static com.android.server.wm.shell.TransitionTraceProto.MAGIC_NUMBER_H;
 import static com.android.server.wm.shell.TransitionTraceProto.MAGIC_NUMBER_L;
@@ -66,20 +58,67 @@ public class TransitionTracer {
 
     /**
      * Records key information about a transition that has been sent to Shell to be played.
+     * More information will be appended to the same proto object once the transition is finished or
+     * aborted.
+     * Transition information won't be added to the trace buffer until
+     * {@link #logFinishedTransition} or {@link #logAbortedTransition} is called for this
+     * transition.
+     *
      * @param transition The transition that has been sent to Shell.
      * @param targets Information about the target windows of the transition.
-     * @param createTimeNs System elapsed time (nanoseconds since boot including sleep time) at
- *                     which the transition to be recorded was created.
-     * @param sendTimeNs System elapsed time (nanoseconds since boot including sleep time) at which
-     * @param info
+     * @param info The TransitionInfo send over to Shell to execute the transition.
      */
     public void logSentTransition(Transition transition, ArrayList<ChangeInfo> targets,
-            long createTimeNs, long sendTimeNs, TransitionInfo info) {
-        mTraceBuffer.pushSentTransition(transition, targets, createTimeNs, sendTimeNs);
+            TransitionInfo info) {
+        // Dump the info to proto that will not be available when the transition finishes or
+        // is canceled
+        final ProtoOutputStream outputStream = transition.mLogger.mProtoOutputStream;
+        transition.mLogger.mProtoToken = outputStream
+                .start(com.android.server.wm.shell.TransitionTraceProto.FINISHED_TRANSITIONS);
+        outputStream.write(com.android.server.wm.shell.Transition.START_TRANSACTION_ID,
+                transition.getStartTransaction().getId());
+        outputStream.write(com.android.server.wm.shell.Transition.FINISH_TRANSACTION_ID,
+                transition.getFinishTransaction().getId());
+        dumpTransitionTargetsToProto(outputStream, transition, targets);
+
         logTransitionInfo(transition, info);
     }
 
     /**
+     * Completes the information dumped in {@link #logSentTransition} for a transition
+     * that has finished or aborted, and add the proto object to the trace buffer.
+     *
+     * @param transition The transition that has finished.
+     */
+    public void logFinishedTransition(Transition transition) {
+        if (transition.mLogger.mProtoToken == 0) {
+            // Transition finished but never sent, so open token never added
+            final ProtoOutputStream outputStream = transition.mLogger.mProtoOutputStream;
+            transition.mLogger.mProtoToken = outputStream
+                    .start(com.android.server.wm.shell.TransitionTraceProto.FINISHED_TRANSITIONS);
+        }
+
+        // Dump the rest of the transition's info that wasn't dumped during logSentTransition
+        dumpFinishedTransitionToProto(transition.mLogger.mProtoOutputStream, transition);
+        transition.mLogger.mProtoOutputStream.end(transition.mLogger.mProtoToken);
+        mTraceBuffer.pushTransitionProto(transition.mLogger.mProtoOutputStream);
+    }
+
+    /**
+     * Same as {@link #logFinishedTransition} but don't add the transition to the trace buffer
+     * unless actively tracing.
+     *
+     * @param transition The transition that has been aborted
+     */
+    public void logAbortedTransition(Transition transition) {
+        // We don't care about aborted transitions unless actively tracing
+        if (!mActiveTracingEnabled) {
+            return;
+        }
+        logFinishedTransition(transition);
+    }
+
+    /**
      * Records the current state of a transition in the transition trace (if it is running).
      * @param transition the transition that we want to record the state of.
      */
@@ -87,7 +126,9 @@ public class TransitionTracer {
         if (!mActiveTracingEnabled) {
             return;
         }
-        mTraceBuffer.pushTransitionState(transition);
+        final ProtoOutputStream outputStream = new ProtoOutputStream();
+        dumpTransitionStateToProto(outputStream, transition);
+        mTraceBuffer.pushTransitionState(outputStream);
     }
 
     /**
@@ -99,171 +140,180 @@ public class TransitionTracer {
         if (!mActiveTracingEnabled) {
             return;
         }
-        mTraceBuffer.pushTransitionInfo(transition, info);
+        final ProtoOutputStream outputStream = new ProtoOutputStream();
+        dumpTransitionInfoToProto(outputStream, transition, info);
+        mTraceBuffer.pushTransitionInfo(outputStream);
     }
 
-    private class TransitionTraceBuffer {
-        private final TraceBuffer mBuffer = new TraceBuffer(ALWAYS_ON_TRACING_CAPACITY);
-        private final TraceBuffer mStateBuffer = new TraceBuffer(ACTIVE_TRACING_BUFFER_CAPACITY);
-        private final TraceBuffer mTransitionInfoBuffer =
-                new TraceBuffer(ACTIVE_TRACING_BUFFER_CAPACITY);
-
-        public void pushSentTransition(
-                Transition transition,
-                ArrayList<ChangeInfo> targets,
-                long createTimeNs,
-                long sendTimeNs
-        ) {
-            Trace.beginSection("TransitionTraceBuffer#pushSentTransition");
-            final ProtoOutputStream outputStream = new ProtoOutputStream();
-            final long transitionToken = outputStream
-                    .start(com.android.server.wm.shell.TransitionTraceProto.SENT_TRANSITIONS);
-
-            if (mActiveTracingEnabled) {
-                outputStream.write(com.android.server.wm.shell.Transition.ID,
-                        transition.getSyncId());
-            }
-
-            outputStream.write(com.android.server.wm.shell.Transition.START_TRANSACTION_ID,
-                    transition.getStartTransaction().getId());
-            outputStream.write(com.android.server.wm.shell.Transition.FINISH_TRANSACTION_ID,
-                    transition.getFinishTransaction().getId());
-
-            outputStream.write(com.android.server.wm.shell.Transition.CREATE_TIME_NS, createTimeNs);
-            outputStream.write(com.android.server.wm.shell.Transition.SEND_TIME_NS, sendTimeNs);
+    private void dumpTransitionTargetsToProto(ProtoOutputStream outputStream,
+            Transition transition, ArrayList<ChangeInfo> targets) {
+        Trace.beginSection("TransitionTracer#dumpTransitionTargetsToProto");
+        if (mActiveTracingEnabled) {
+            outputStream.write(com.android.server.wm.shell.Transition.ID,
+                    transition.getSyncId());
+        }
 
-            for (int i = 0; i < targets.size(); ++i) {
-                final long changeToken = outputStream
-                        .start(com.android.server.wm.shell.Transition.TARGETS);
+        outputStream.write(com.android.server.wm.shell.Transition.TYPE, transition.mType);
 
-                final Transition.ChangeInfo target = targets.get(i);
+        for (int i = 0; i < targets.size(); ++i) {
+            final long changeToken = outputStream
+                    .start(com.android.server.wm.shell.Transition.TARGETS);
 
-                final int mode = target.getTransitMode(target.mContainer);
-                final int layerId;
-                if (target.mContainer.mSurfaceControl.isValid()) {
-                    layerId = target.mContainer.mSurfaceControl.getLayerId();
-                } else {
-                    layerId = -1;
-                }
+            final Transition.ChangeInfo target = targets.get(i);
 
-                outputStream.write(com.android.server.wm.shell.Target.MODE, mode);
-                outputStream.write(com.android.server.wm.shell.Target.LAYER_ID, layerId);
+            final int mode = target.getTransitMode(target.mContainer);
+            final int layerId;
+            if (target.mContainer.mSurfaceControl.isValid()) {
+                layerId = target.mContainer.mSurfaceControl.getLayerId();
+            } else {
+                layerId = -1;
+            }
 
-                if (mActiveTracingEnabled) {
-                    // What we use in the WM trace
-                    final int windowId = System.identityHashCode(target.mContainer);
-                    outputStream.write(com.android.server.wm.shell.Target.WINDOW_ID, windowId);
-                }
+            outputStream.write(com.android.server.wm.shell.Target.MODE, mode);
+            outputStream.write(com.android.server.wm.shell.Target.LAYER_ID, layerId);
 
-                outputStream.end(changeToken);
+            if (mActiveTracingEnabled) {
+                // What we use in the WM trace
+                final int windowId = System.identityHashCode(target.mContainer);
+                outputStream.write(com.android.server.wm.shell.Target.WINDOW_ID, windowId);
             }
 
-            outputStream.end(transitionToken);
-            mBuffer.add(outputStream);
-
-            Trace.endSection();
+            outputStream.end(changeToken);
         }
 
-        private void pushTransitionState(Transition transition) {
-            Trace.beginSection("TransitionTraceBuffer#pushTransitionState");
-            final ProtoOutputStream outputStream = new ProtoOutputStream();
-            final long stateToken = outputStream
-                    .start(com.android.server.wm.shell.TransitionTraceProto.TRANSITION_STATES);
+        Trace.endSection();
+    }
 
-            outputStream.write(com.android.server.wm.shell.TransitionState.TIME_NS,
-                    SystemClock.elapsedRealtimeNanos());
-            outputStream.write(com.android.server.wm.shell.TransitionState.TRANSITION_ID,
-                    transition.getSyncId());
-            outputStream.write(com.android.server.wm.shell.TransitionState.TRANSITION_TYPE,
-                    transition.mType);
-            outputStream.write(com.android.server.wm.shell.TransitionState.STATE,
-                    transition.getState());
-            outputStream.write(com.android.server.wm.shell.TransitionState.FLAGS,
-                    transition.getFlags());
-
-            for (int i = 0; i < transition.mChanges.size(); ++i) {
-                final WindowContainer window = transition.mChanges.keyAt(i);
-                final ChangeInfo changeInfo = transition.mChanges.valueAt(i);
-                writeChange(outputStream, window, changeInfo);
-            }
+    private void dumpFinishedTransitionToProto(
+            ProtoOutputStream outputStream,
+            Transition transition
+    ) {
+        Trace.beginSection("TransitionTracer#dumpFinishedTransitionToProto");
 
-            for (int i = 0; i < transition.mChanges.size(); ++i) {
-                final WindowContainer window = transition.mChanges.keyAt(i);
-                final ChangeInfo changeInfo = transition.mChanges.valueAt(i);
-                writeChange(outputStream, window, changeInfo);
-            }
+        outputStream.write(com.android.server.wm.shell.Transition.CREATE_TIME_NS,
+                transition.mLogger.mCreateTimeNs);
+        outputStream.write(com.android.server.wm.shell.Transition.SEND_TIME_NS,
+                transition.mLogger.mSendTimeNs);
+        outputStream.write(com.android.server.wm.shell.Transition.FINISH_TIME_NS,
+                transition.mLogger.mFinishTimeNs);
 
-            for (int i = 0; i < transition.mParticipants.size(); ++i) {
-                final WindowContainer window = transition.mParticipants.valueAt(i);
-                window.writeIdentifierToProto(outputStream,
-                        com.android.server.wm.shell.TransitionState.PARTICIPANTS);
-            }
+        Trace.endSection();
+    }
 
-            outputStream.end(stateToken);
+    private void dumpTransitionStateToProto(ProtoOutputStream outputStream, Transition transition) {
+        Trace.beginSection("TransitionTracer#dumpTransitionStateToProto");
+
+        final long stateToken = outputStream
+                .start(com.android.server.wm.shell.TransitionTraceProto.TRANSITION_STATES);
+
+        outputStream.write(com.android.server.wm.shell.TransitionState.TIME_NS,
+                SystemClock.elapsedRealtimeNanos());
+        outputStream.write(com.android.server.wm.shell.TransitionState.TRANSITION_ID,
+                transition.getSyncId());
+        outputStream.write(com.android.server.wm.shell.TransitionState.TRANSITION_TYPE,
+                transition.mType);
+        outputStream.write(com.android.server.wm.shell.TransitionState.STATE,
+                transition.getState());
+        outputStream.write(com.android.server.wm.shell.TransitionState.FLAGS,
+                transition.getFlags());
+
+        for (int i = 0; i < transition.mChanges.size(); ++i) {
+            final WindowContainer window = transition.mChanges.keyAt(i);
+            final ChangeInfo changeInfo = transition.mChanges.valueAt(i);
+            dumpChangeInfoToProto(outputStream, window, changeInfo);
+        }
 
-            mStateBuffer.add(outputStream);
-            Trace.endSection();
+        for (int i = 0; i < transition.mParticipants.size(); ++i) {
+            final WindowContainer window = transition.mParticipants.valueAt(i);
+            window.writeIdentifierToProto(outputStream,
+                    com.android.server.wm.shell.TransitionState.PARTICIPANTS);
         }
 
-        private void pushTransitionInfo(Transition transition, TransitionInfo info) {
-            Trace.beginSection("TransitionTraceBuffer#pushTransitionInfo");
-            final ProtoOutputStream outputStream = new ProtoOutputStream();
-            final long transitionInfoToken = outputStream
-                    .start(com.android.server.wm.shell.TransitionTraceProto.TRANSITION_INFO);
+        outputStream.end(stateToken);
+        Trace.endSection();
+    }
 
-            outputStream.write(com.android.server.wm.shell.TransitionInfo.TRANSITION_ID,
-                    transition.getSyncId());
-            for (int i = 0; i < info.getChanges().size(); ++i) {
-                TransitionInfo.Change change = info.getChanges().get(i);
-                writeTransitionInfoChange(outputStream, change);
-            }
+    private void dumpChangeInfoToProto(ProtoOutputStream outputStream, WindowContainer window,
+            ChangeInfo changeInfo) {
+        Trace.beginSection("TransitionTraceBuffer#writeChange");
+        final long changeEntryToken =
+                outputStream.start(com.android.server.wm.shell.TransitionState.CHANGE);
+
+        final int transitMode = changeInfo.getTransitMode(window);
+        final boolean hasChanged = changeInfo.hasChanged();
+        final int changeFlags = changeInfo.getChangeFlags(window);
+        final int windowingMode = changeInfo.mWindowingMode;
+
+        outputStream.write(com.android.server.wm.shell.ChangeInfo.TRANSIT_MODE, transitMode);
+        outputStream.write(com.android.server.wm.shell.ChangeInfo.HAS_CHANGED, hasChanged);
+        outputStream.write(com.android.server.wm.shell.ChangeInfo.CHANGE_FLAGS, changeFlags);
+        outputStream.write(com.android.server.wm.shell.ChangeInfo.WINDOWING_MODE, windowingMode);
+        window.writeIdentifierToProto(
+                outputStream, com.android.server.wm.shell.ChangeInfo.WINDOW_IDENTIFIER);
+
+        outputStream.end(changeEntryToken);
+        Trace.endSection();
+    }
 
-            outputStream.end(transitionInfoToken);
-            mTransitionInfoBuffer.add(outputStream);
-            Trace.endSection();
+    private void dumpTransitionInfoToProto(ProtoOutputStream outputStream,
+            Transition transition, TransitionInfo info) {
+        Trace.beginSection("TransitionTracer#dumpTransitionInfoToProto");
+        final long transitionInfoToken = outputStream
+                .start(com.android.server.wm.shell.TransitionTraceProto.TRANSITION_INFO);
+
+        outputStream.write(com.android.server.wm.shell.TransitionInfo.TRANSITION_ID,
+                transition.getSyncId());
+        for (int i = 0; i < info.getChanges().size(); ++i) {
+            TransitionInfo.Change change = info.getChanges().get(i);
+            dumpTransitionInfoChangeToProto(outputStream, change);
         }
 
-        private void writeChange(ProtoOutputStream outputStream, WindowContainer window,
-                ChangeInfo changeInfo) {
-            Trace.beginSection("TransitionTraceBuffer#writeChange");
-            final long changeEntryToken = outputStream.start(CHANGE);
+        outputStream.end(transitionInfoToken);
+        Trace.endSection();
+    }
 
-            final int transitMode = changeInfo.getTransitMode(window);
-            final boolean hasChanged = changeInfo.hasChanged();
-            final int changeFlags = changeInfo.getChangeFlags(window);
-            final int windowingMode = changeInfo.mWindowingMode;
+    private void dumpTransitionInfoChangeToProto(
+            ProtoOutputStream outputStream,
+            TransitionInfo.Change change
+    ) {
+        Trace.beginSection("TransitionTracer#dumpTransitionInfoChangeToProto");
+        final long changeEntryToken = outputStream
+                .start(com.android.server.wm.shell.TransitionInfo.CHANGE);
 
-            outputStream.write(TRANSIT_MODE, transitMode);
-            outputStream.write(HAS_CHANGED, hasChanged);
-            outputStream.write(CHANGE_FLAGS, changeFlags);
-            outputStream.write(WINDOWING_MODE, windowingMode);
-            window.writeIdentifierToProto(outputStream, WINDOW_IDENTIFIER);
+        outputStream.write(com.android.server.wm.shell.TransitionInfoChange.LAYER_ID,
+                change.getLeash().getLayerId());
+        outputStream.write(com.android.server.wm.shell.TransitionInfoChange.MODE, change.getMode());
 
-            outputStream.end(changeEntryToken);
-            Trace.endSection();
-        }
+        outputStream.end(changeEntryToken);
+        Trace.endSection();
+    }
+
+    private class TransitionTraceBuffer {
+        private final TraceBuffer mTransitionBuffer = new TraceBuffer(ALWAYS_ON_TRACING_CAPACITY);
+        private final TraceBuffer mStateBuffer = new TraceBuffer(ACTIVE_TRACING_BUFFER_CAPACITY);
+        private final TraceBuffer mTransitionInfoBuffer =
+                new TraceBuffer(ACTIVE_TRACING_BUFFER_CAPACITY);
 
-        private void writeTransitionInfoChange(
-                ProtoOutputStream outputStream,
-                TransitionInfo.Change change
-        ) {
-            Trace.beginSection("TransitionTraceBuffer#writeTransitionInfoChange");
-            final long changeEntryToken = outputStream
-                    .start(com.android.server.wm.shell.TransitionInfo.CHANGE);
+        private void pushTransitionProto(ProtoOutputStream outputStream) {
+            mTransitionBuffer.add(outputStream);
+        }
 
-            outputStream.write(LAYER_ID, change.getLeash().getLayerId());
-            outputStream.write(MODE, change.getMode());
+        private void pushTransitionState(ProtoOutputStream outputStream) {
+            mStateBuffer.add(outputStream);
+        }
 
-            outputStream.end(changeEntryToken);
-            Trace.endSection();
+        private void pushTransitionInfo(ProtoOutputStream outputStream) {
+            mTransitionInfoBuffer.add(outputStream);
         }
 
         public void writeToFile(File file, ProtoOutputStream proto) throws IOException {
-            mBuffer.writeTraceToFile(file, proto);
+            mTransitionBuffer.writeTraceToFile(file, proto);
         }
 
         public void reset() {
-            mBuffer.resetBuffer();
+            mTransitionBuffer.resetBuffer();
+            mStateBuffer.resetBuffer();
+            mTransitionInfoBuffer.resetBuffer();
         }
     }
 
@@ -280,7 +330,7 @@ public class TransitionTracer {
         LogAndPrintln.i(pw, "Starting shell transition trace.");
         synchronized (mEnabledLock) {
             mActiveTracingEnabled = true;
-            mTraceBuffer.mBuffer.setCapacity(ACTIVE_TRACING_BUFFER_CAPACITY);
+            mTraceBuffer.mTransitionBuffer.setCapacity(ACTIVE_TRACING_BUFFER_CAPACITY);
             mTraceBuffer.reset();
         }
         Trace.endSection();
@@ -309,7 +359,8 @@ public class TransitionTracer {
         synchronized (mEnabledLock) {
             mActiveTracingEnabled = false;
             writeTraceToFileLocked(pw, outputFile);
-            mTraceBuffer.mBuffer.setCapacity(ALWAYS_ON_TRACING_CAPACITY);
+            mTraceBuffer.reset();
+            mTraceBuffer.mTransitionBuffer.setCapacity(ALWAYS_ON_TRACING_CAPACITY);
         }
         Trace.endSection();
     }