Add shell-side transition trace perfetto tracer

Bug: 309630341

Test: adb shell -t perfetto -c - --txt -o /data/misc/perfetto-traces/trace <<EOF
   unique_session_name: transitions_tracing
   buffers: {
       size_kb: 63488
       fill_policy: RING_BUFFER
   }
   data_sources: {
       config {
           name: com.android.wm.shell.transition
       }
   }
   duration_ms: 10000
EOF

Change-Id: If510fabd57fb7a9242235ad0b42d81133b96377f
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index 45540e0..4cdc06a 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -162,6 +162,7 @@
         "com_android_wm_shell_flags_lib",
         "com.android.window.flags.window-aconfig-java",
         "WindowManager-Shell-proto",
+        "perfetto_trace_java_protos",
         "dagger2",
         "jsr330",
     ],
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
index b0d8b47..3fb0dbf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
@@ -83,6 +83,9 @@
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.transition.tracing.LegacyTransitionTracer;
+import com.android.wm.shell.transition.tracing.PerfettoTransitionTracer;
+import com.android.wm.shell.transition.tracing.TransitionTracer;
 import com.android.wm.shell.util.TransitionUtil;
 
 import java.io.PrintWriter;
@@ -184,7 +187,7 @@
     private final ShellController mShellController;
     private final ShellTransitionImpl mImpl = new ShellTransitionImpl();
     private final SleepHandler mSleepHandler = new SleepHandler();
-    private final Tracer mTracer = new Tracer();
+    private final TransitionTracer mTransitionTracer;
     private boolean mIsRegistered = false;
 
     /** List of possible handlers. Ordered by specificity (eg. tapped back to front). */
@@ -307,6 +310,12 @@
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "addHandler: Remote");
         shellInit.addInitCallback(this::onInit, this);
         mHomeTransitionObserver = observer;
+
+        if (android.tracing.Flags.perfettoTransitionTracing()) {
+            mTransitionTracer = new PerfettoTransitionTracer();
+        } else {
+            mTransitionTracer = new LegacyTransitionTracer();
+        }
     }
 
     private void onInit() {
@@ -868,7 +877,7 @@
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "Transition %s ready while"
                 + " %s is still animating. Notify the animating transition"
                 + " in case they can be merged", ready, playing);
-        mTracer.logMergeRequested(ready.mInfo.getDebugId(), playing.mInfo.getDebugId());
+        mTransitionTracer.logMergeRequested(ready.mInfo.getDebugId(), playing.mInfo.getDebugId());
         playing.mHandler.mergeAnimation(ready.mToken, ready.mInfo, ready.mStartT,
                 playing.mToken, (wct) -> onMerged(playing, ready));
     }
@@ -902,7 +911,7 @@
         for (int i = 0; i < mObservers.size(); ++i) {
             mObservers.get(i).onTransitionMerged(merged.mToken, playing.mToken);
         }
-        mTracer.logMerged(merged.mInfo.getDebugId(), playing.mInfo.getDebugId());
+        mTransitionTracer.logMerged(merged.mInfo.getDebugId(), playing.mInfo.getDebugId());
         // See if we should merge another transition.
         processReadyQueue(track);
     }
@@ -923,7 +932,7 @@
                     active.mStartT, active.mFinishT, (wct) -> onFinish(active, wct));
             if (consumed) {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " animated by firstHandler");
-                mTracer.logDispatched(active.mInfo.getDebugId(), active.mHandler);
+                mTransitionTracer.logDispatched(active.mInfo.getDebugId(), active.mHandler);
                 return;
             }
         }
@@ -948,7 +957,7 @@
             if (consumed) {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " animated by %s",
                         mHandlers.get(i));
-                mTracer.logDispatched(info.getDebugId(), mHandlers.get(i));
+                mTransitionTracer.logDispatched(info.getDebugId(), mHandlers.get(i));
                 return mHandlers.get(i);
             }
         }
@@ -978,7 +987,7 @@
         final Track track = mTracks.get(transition.getTrack());
         transition.mAborted = true;
 
-        mTracer.logAborted(transition.mInfo.getDebugId());
+        mTransitionTracer.logAborted(transition.mInfo.getDebugId());
 
         if (transition.mHandler != null) {
             // Notifies to clean-up the aborted transition.
@@ -1506,12 +1515,18 @@
         }
     }
 
-
     @Override
     public boolean onShellCommand(String[] args, PrintWriter pw) {
         switch (args[0]) {
             case "tracing": {
-                mTracer.onShellCommand(Arrays.copyOfRange(args, 1, args.length), pw);
+                if (!android.tracing.Flags.perfettoTransitionTracing()) {
+                    ((LegacyTransitionTracer) mTransitionTracer)
+                            .onShellCommand(Arrays.copyOfRange(args, 1, args.length), pw);
+                } else {
+                    pw.println("Command not supported. Use the Perfetto command instead to start "
+                            + "and stop this trace instead.");
+                    return false;
+                }
                 return true;
             }
             default: {
@@ -1524,8 +1539,10 @@
 
     @Override
     public void printShellCommandHelp(PrintWriter pw, String prefix) {
-        pw.println(prefix + "tracing");
-        mTracer.printShellCommandHelp(pw, prefix + "  ");
+        if (!android.tracing.Flags.perfettoTransitionTracing()) {
+            pw.println(prefix + "tracing");
+            ((LegacyTransitionTracer) mTransitionTracer).printShellCommandHelp(pw, prefix + "  ");
+        }
     }
 
     private void dump(@NonNull PrintWriter pw, String prefix) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Tracer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/LegacyTransitionTracer.java
similarity index 88%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/transition/Tracer.java
rename to libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/LegacyTransitionTracer.java
index 5919aad..9c84886 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Tracer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/LegacyTransitionTracer.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.transition;
+package com.android.wm.shell.transition.tracing;
 
 import static android.os.Build.IS_USER;
 
@@ -29,6 +29,7 @@
 
 import com.android.internal.util.TraceBuffer;
 import com.android.wm.shell.sysui.ShellCommandHandler;
+import com.android.wm.shell.transition.Transitions;
 
 import com.google.protobuf.nano.MessageNano;
 
@@ -45,7 +46,8 @@
 /**
  * Helper class to collect and dump transition traces.
  */
-public class Tracer implements ShellCommandHandler.ShellCommandActionHandler {
+public class LegacyTransitionTracer
+        implements ShellCommandHandler.ShellCommandActionHandler, TransitionTracer {
     private static final int ALWAYS_ON_TRACING_CAPACITY = 15 * 1024; // 15 KB
     private static final int ACTIVE_TRACING_BUFFER_CAPACITY = 5000 * 1024; // 5 MB
 
@@ -60,33 +62,33 @@
 
     private final TraceBuffer.ProtoProvider mProtoProvider =
             new TraceBuffer.ProtoProvider<MessageNano,
-                com.android.wm.shell.nano.WmShellTransitionTraceProto,
-                com.android.wm.shell.nano.Transition>() {
-        @Override
-        public int getItemSize(MessageNano proto) {
-            return proto.getCachedSize();
-        }
+                    com.android.wm.shell.nano.WmShellTransitionTraceProto,
+                    com.android.wm.shell.nano.Transition>() {
+                @Override
+                public int getItemSize(MessageNano proto) {
+                    return proto.getCachedSize();
+                }
 
-        @Override
-        public byte[] getBytes(MessageNano proto) {
-            return MessageNano.toByteArray(proto);
-        }
+                @Override
+                public byte[] getBytes(MessageNano proto) {
+                    return MessageNano.toByteArray(proto);
+                }
 
-        @Override
-        public void write(
-                com.android.wm.shell.nano.WmShellTransitionTraceProto encapsulatingProto,
-                Queue<com.android.wm.shell.nano.Transition> buffer, OutputStream os)
+                @Override
+                public void write(
+                        com.android.wm.shell.nano.WmShellTransitionTraceProto encapsulatingProto,
+                        Queue<com.android.wm.shell.nano.Transition> buffer, OutputStream os)
                         throws IOException {
-            encapsulatingProto.transitions = buffer.toArray(
-                    new com.android.wm.shell.nano.Transition[0]);
-            os.write(getBytes(encapsulatingProto));
-        }
-    };
+                    encapsulatingProto.transitions = buffer.toArray(
+                            new com.android.wm.shell.nano.Transition[0]);
+                    os.write(getBytes(encapsulatingProto));
+                }
+            };
     private final TraceBuffer<MessageNano,
             com.android.wm.shell.nano.WmShellTransitionTraceProto,
-            com.android.wm.shell.nano.Transition> mTraceBuffer
-                    = new TraceBuffer(ALWAYS_ON_TRACING_CAPACITY, mProtoProvider,
-                            (proto) -> handleOnEntryRemovedFromTrace(proto));
+            com.android.wm.shell.nano.Transition> mTraceBuffer =
+                new TraceBuffer(ALWAYS_ON_TRACING_CAPACITY, mProtoProvider,
+                        this::handleOnEntryRemovedFromTrace);
     private final Map<Object, Runnable> mRemovedFromTraceCallbacks = new HashMap<>();
 
     private final Map<Transitions.TransitionHandler, Integer> mHandlerIds = new HashMap<>();
@@ -99,6 +101,7 @@
      * @param transitionId The id of the transition being dispatched.
      * @param handler The handler the transition is being dispatched to.
      */
+    @Override
     public void logDispatched(int transitionId, Transitions.TransitionHandler handler) {
         final int handlerId;
         if (mHandlerIds.containsKey(handler)) {
@@ -130,6 +133,7 @@
      *
      * @param mergeRequestedTransitionId The id of the transition we are requesting to be merged.
      */
+    @Override
     public void logMergeRequested(int mergeRequestedTransitionId, int playingTransitionId) {
         com.android.wm.shell.nano.Transition proto = new com.android.wm.shell.nano.Transition();
         proto.id = mergeRequestedTransitionId;
@@ -145,6 +149,7 @@
      * @param mergedTransitionId The id of the transition that was merged.
      * @param playingTransitionId The id of the transition the transition was merged into.
      */
+    @Override
     public void logMerged(int mergedTransitionId, int playingTransitionId) {
         com.android.wm.shell.nano.Transition proto = new com.android.wm.shell.nano.Transition();
         proto.id = mergedTransitionId;
@@ -159,6 +164,7 @@
      *
      * @param transitionId The id of the transition that was aborted.
      */
+    @Override
     public void logAborted(int transitionId) {
         com.android.wm.shell.nano.Transition proto = new com.android.wm.shell.nano.Transition();
         proto.id = transitionId;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/PerfettoTransitionTracer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/PerfettoTransitionTracer.java
new file mode 100644
index 0000000..99df6a3
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/PerfettoTransitionTracer.java
@@ -0,0 +1,174 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.transition.tracing;
+
+import android.internal.perfetto.protos.PerfettoTrace;
+import android.os.SystemClock;
+import android.tracing.perfetto.DataSourceInstance;
+import android.tracing.perfetto.DataSourceParams;
+import android.tracing.perfetto.InitArguments;
+import android.tracing.perfetto.Producer;
+import android.tracing.perfetto.TracingContext;
+import android.tracing.transition.TransitionDataSource;
+import android.util.proto.ProtoOutputStream;
+
+import com.android.wm.shell.transition.Transitions;
+
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Helper class to collect and dump transition traces.
+ */
+public class PerfettoTransitionTracer implements TransitionTracer {
+    private final AtomicInteger mActiveTraces = new AtomicInteger(0);
+    private final TransitionDataSource mDataSource = new TransitionDataSource(
+            mActiveTraces::incrementAndGet,
+            this::onFlush,
+            mActiveTraces::decrementAndGet);
+
+    public PerfettoTransitionTracer() {
+        Producer.init(InitArguments.DEFAULTS);
+        mDataSource.register(DataSourceParams.DEFAULTS);
+    }
+
+    /**
+     * Adds an entry in the trace to log that a transition has been dispatched to a handler.
+     *
+     * @param transitionId The id of the transition being dispatched.
+     * @param handler The handler the transition is being dispatched to.
+     */
+    @Override
+    public void logDispatched(int transitionId, Transitions.TransitionHandler handler) {
+        if (!isTracing()) {
+            return;
+        }
+
+        mDataSource.trace(ctx -> {
+            final int handlerId = getHandlerId(handler, ctx);
+
+            final ProtoOutputStream os = ctx.newTracePacket();
+            final long token = os.start(PerfettoTrace.TracePacket.SHELL_TRANSITION);
+            os.write(PerfettoTrace.ShellTransition.ID, transitionId);
+            os.write(PerfettoTrace.ShellTransition.DISPATCH_TIME_NS,
+                    SystemClock.elapsedRealtimeNanos());
+            os.write(PerfettoTrace.ShellTransition.HANDLER, handlerId);
+            os.end(token);
+        });
+    }
+
+    private static int getHandlerId(Transitions.TransitionHandler handler,
+            TracingContext<DataSourceInstance, TransitionDataSource.TlsState, Void> ctx) {
+        final Map<String, Integer> handlerMapping =
+                ctx.getCustomTlsState().handlerMapping;
+        final int handlerId;
+        if (handlerMapping.containsKey(handler.getClass().getName())) {
+            handlerId = handlerMapping.get(handler.getClass().getName());
+        } else {
+            // + 1 to avoid 0 ids which can be confused with missing value when dumped to proto
+            handlerId = handlerMapping.size() + 1;
+            handlerMapping.put(handler.getClass().getName(), handlerId);
+        }
+        return handlerId;
+    }
+
+    /**
+     * Adds an entry in the trace to log that a request to merge a transition was made.
+     *
+     * @param mergeRequestedTransitionId The id of the transition we are requesting to be merged.
+     */
+    @Override
+    public void logMergeRequested(int mergeRequestedTransitionId, int playingTransitionId) {
+        if (!isTracing()) {
+            return;
+        }
+
+        mDataSource.trace(ctx -> {
+            final ProtoOutputStream os = ctx.newTracePacket();
+            final long token = os.start(PerfettoTrace.TracePacket.SHELL_TRANSITION);
+            os.write(PerfettoTrace.ShellTransition.ID, mergeRequestedTransitionId);
+            os.write(PerfettoTrace.ShellTransition.MERGE_REQUEST_TIME_NS,
+                    SystemClock.elapsedRealtimeNanos());
+            os.write(PerfettoTrace.ShellTransition.MERGE_TARGET, playingTransitionId);
+            os.end(token);
+        });
+    }
+
+    /**
+     * Adds an entry in the trace to log that a transition was merged by the handler.
+     *
+     * @param mergedTransitionId The id of the transition that was merged.
+     * @param playingTransitionId The id of the transition the transition was merged into.
+     */
+    @Override
+    public void logMerged(int mergedTransitionId, int playingTransitionId) {
+        if (!isTracing()) {
+            return;
+        }
+
+        mDataSource.trace(ctx -> {
+            final ProtoOutputStream os = ctx.newTracePacket();
+            final long token = os.start(PerfettoTrace.TracePacket.SHELL_TRANSITION);
+            os.write(PerfettoTrace.ShellTransition.ID, mergedTransitionId);
+            os.write(PerfettoTrace.ShellTransition.MERGE_TIME_NS,
+                    SystemClock.elapsedRealtimeNanos());
+            os.write(PerfettoTrace.ShellTransition.MERGE_TARGET, playingTransitionId);
+            os.end(token);
+        });
+    }
+
+    /**
+     * Adds an entry in the trace to log that a transition was aborted.
+     *
+     * @param transitionId The id of the transition that was aborted.
+     */
+    @Override
+    public void logAborted(int transitionId) {
+        if (!isTracing()) {
+            return;
+        }
+
+        mDataSource.trace(ctx -> {
+            final ProtoOutputStream os = ctx.newTracePacket();
+            final long token = os.start(PerfettoTrace.TracePacket.SHELL_TRANSITION);
+            os.write(PerfettoTrace.ShellTransition.ID, transitionId);
+            os.write(PerfettoTrace.ShellTransition.SHELL_ABORT_TIME_NS,
+                    SystemClock.elapsedRealtimeNanos());
+            os.end(token);
+        });
+    }
+
+    private boolean isTracing() {
+        return mActiveTraces.get() > 0;
+    }
+
+    private void onFlush() {
+        mDataSource.trace(ctx -> {
+            final ProtoOutputStream os = ctx.newTracePacket();
+
+            final Map<String, Integer> handlerMapping = ctx.getCustomTlsState().handlerMapping;
+            for (String handler : handlerMapping.keySet()) {
+                final long token = os.start(PerfettoTrace.TracePacket.SHELL_HANDLER_MAPPINGS);
+                os.write(PerfettoTrace.ShellHandlerMapping.ID, handlerMapping.get(handler));
+                os.write(PerfettoTrace.ShellHandlerMapping.NAME, handler);
+                os.end(token);
+            }
+
+            ctx.flush();
+        });
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/TransitionTracer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/TransitionTracer.java
new file mode 100644
index 0000000..5857ad8
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/tracing/TransitionTracer.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.transition.tracing;
+
+import com.android.wm.shell.transition.Transitions;
+
+public interface TransitionTracer {
+    /**
+     * Adds an entry in the trace to log that a transition has been dispatched to a handler.
+     *
+     * @param transitionId The id of the transition being dispatched.
+     * @param handler The handler the transition is being dispatched to.
+     */
+    void logDispatched(int transitionId, Transitions.TransitionHandler handler);
+
+    /**
+     * Adds an entry in the trace to log that a request to merge a transition was made.
+     *
+     * @param mergeRequestedTransitionId The id of the transition we are requesting to be merged.
+     */
+    void logMergeRequested(int mergeRequestedTransitionId, int playingTransitionId);
+
+    /**
+     * Adds an entry in the trace to log that a transition was merged by the handler.
+     *
+     * @param mergedTransitionId The id of the transition that was merged.
+     * @param playingTransitionId The id of the transition the transition was merged into.
+     */
+    void logMerged(int mergedTransitionId, int playingTransitionId);
+
+    /**
+     * Adds an entry in the trace to log that a transition was aborted.
+     *
+     * @param transitionId The id of the transition that was aborted.
+     */
+    void logAborted(int transitionId);
+}