ime tracing: migrate to perfetto

Add ImeTracingPerfettoImpl to be used when the flag
"android.tracing.perfetto_ime" is enabled.

Bug: 276433199
Test: enable flag android.tracing.perfetto_ime \
   + capture/visualize IME traces with winscope
Change-Id: I93dd127165ce16b2cac1e6e1cda4f6f7e029f961
diff --git a/core/java/android/tracing/inputmethod/InputMethodDataSource.java b/core/java/android/tracing/inputmethod/InputMethodDataSource.java
new file mode 100644
index 0000000..5c5ad69
--- /dev/null
+++ b/core/java/android/tracing/inputmethod/InputMethodDataSource.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2024 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 android.tracing.inputmethod;
+
+import android.annotation.NonNull;
+import android.tracing.perfetto.DataSource;
+import android.tracing.perfetto.DataSourceInstance;
+import android.tracing.perfetto.StartCallbackArguments;
+import android.tracing.perfetto.StopCallbackArguments;
+import android.util.proto.ProtoInputStream;
+
+/**
+ * @hide
+ */
+public final class InputMethodDataSource
+        extends DataSource<DataSourceInstance, Void, Void> {
+    public static final String DATA_SOURCE_NAME = "android.inputmethod";
+
+    @NonNull
+    private final Runnable mOnStartCallback;
+    @NonNull
+    private final Runnable mOnStopCallback;
+
+    public InputMethodDataSource(@NonNull Runnable onStart, @NonNull Runnable onStop) {
+        super(DATA_SOURCE_NAME);
+        mOnStartCallback = onStart;
+        mOnStopCallback = onStop;
+    }
+
+    @Override
+    public DataSourceInstance createInstance(ProtoInputStream configStream, int instanceIndex) {
+        return new DataSourceInstance(this, instanceIndex) {
+            @Override
+            protected void onStart(StartCallbackArguments args) {
+                mOnStartCallback.run();
+            }
+
+            @Override
+            protected void onStop(StopCallbackArguments args) {
+                mOnStopCallback.run();
+            }
+        };
+    }
+}
diff --git a/core/java/com/android/internal/inputmethod/ImeTracing.java b/core/java/com/android/internal/inputmethod/ImeTracing.java
index ee9c3aa..cd4ccda 100644
--- a/core/java/com/android/internal/inputmethod/ImeTracing.java
+++ b/core/java/com/android/internal/inputmethod/ImeTracing.java
@@ -60,7 +60,9 @@
      */
     public static ImeTracing getInstance() {
         if (sInstance == null) {
-            if (isSystemProcess()) {
+            if (android.tracing.Flags.perfettoIme()) {
+                sInstance = new ImeTracingPerfettoImpl();
+            } else if (isSystemProcess()) {
                 sInstance = new ImeTracingServerImpl();
             } else {
                 sInstance = new ImeTracingClientImpl();
@@ -78,7 +80,7 @@
      * and {@see #IME_TRACING_FROM_IMS}
      * @param where
      */
-    public void sendToService(byte[] protoDump, int source, String where) {
+    protected void sendToService(byte[] protoDump, int source, String where) {
         InputMethodManagerGlobal.startProtoDump(protoDump, source, where,
                 e -> Log.e(TAG, "Exception while sending ime-related dump to server", e));
     }
diff --git a/core/java/com/android/internal/inputmethod/ImeTracingPerfettoImpl.java b/core/java/com/android/internal/inputmethod/ImeTracingPerfettoImpl.java
new file mode 100644
index 0000000..91b80dd
--- /dev/null
+++ b/core/java/com/android/internal/inputmethod/ImeTracingPerfettoImpl.java
@@ -0,0 +1,178 @@
+/*
+ * Copyright (C) 2024 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.internal.inputmethod;
+
+import static android.tracing.perfetto.DataSourceParams.PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_STALL_AND_ABORT;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.internal.perfetto.protos.Inputmethodeditor.InputMethodClientsTraceProto;
+import android.internal.perfetto.protos.Inputmethodeditor.InputMethodManagerServiceTraceProto;
+import android.internal.perfetto.protos.Inputmethodeditor.InputMethodServiceTraceProto;
+import android.internal.perfetto.protos.TracePacketOuterClass.TracePacket;
+import android.internal.perfetto.protos.WinscopeExtensionsImplOuterClass.WinscopeExtensionsImpl;
+import android.os.SystemClock;
+import android.os.Trace;
+import android.tracing.inputmethod.InputMethodDataSource;
+import android.tracing.perfetto.DataSourceParams;
+import android.tracing.perfetto.InitArguments;
+import android.tracing.perfetto.Producer;
+import android.util.proto.ProtoOutputStream;
+import android.view.inputmethod.InputMethodManager;
+
+import java.io.PrintWriter;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * An implementation of {@link ImeTracing} for perfetto tracing.
+ */
+final class ImeTracingPerfettoImpl extends ImeTracing {
+    private final AtomicInteger mTracingSessionsCount = new AtomicInteger(0);
+    private final AtomicBoolean mIsClientDumpInProgress = new AtomicBoolean(false);
+    private final AtomicBoolean mIsServiceDumpInProgress = new AtomicBoolean(false);
+    private final AtomicBoolean mIsManagerServiceDumpInProgress = new AtomicBoolean(false);
+    private final InputMethodDataSource mDataSource = new InputMethodDataSource(
+            mTracingSessionsCount::incrementAndGet,
+            mTracingSessionsCount::decrementAndGet);
+
+    ImeTracingPerfettoImpl() {
+        Producer.init(InitArguments.DEFAULTS);
+        mDataSource.register(
+                new DataSourceParams(PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_STALL_AND_ABORT));
+    }
+
+
+    @Override
+    public void triggerClientDump(String where, InputMethodManager immInstance,
+            @Nullable byte[] icProto) {
+        if (!isEnabled() || !isAvailable()) {
+            return;
+        }
+
+        if (!mIsClientDumpInProgress.compareAndSet(false, true)) {
+            return;
+        }
+
+        if (immInstance == null) {
+            return;
+        }
+
+        try {
+            Trace.beginSection("inputmethod_client_dump");
+            mDataSource.trace((ctx) -> {
+                final ProtoOutputStream os = ctx.newTracePacket();
+                os.write(TracePacket.TIMESTAMP, SystemClock.elapsedRealtimeNanos());
+                final long tokenWinscopeExtensions =
+                        os.start(TracePacket.WINSCOPE_EXTENSIONS);
+                final long tokenExtensionsField =
+                        os.start(WinscopeExtensionsImpl.INPUTMETHOD_CLIENTS);
+                os.write(InputMethodClientsTraceProto.WHERE, where);
+                final long tokenClient =
+                        os.start(InputMethodClientsTraceProto.CLIENT);
+                immInstance.dumpDebug(os, icProto);
+                os.end(tokenClient);
+                os.end(tokenExtensionsField);
+                os.end(tokenWinscopeExtensions);
+            });
+        } finally {
+            mIsClientDumpInProgress.set(false);
+            Trace.endSection();
+        }
+    }
+
+    @Override
+    public void triggerServiceDump(String where,
+            @NonNull ServiceDumper dumper, @Nullable byte[] icProto) {
+        if (!isEnabled() || !isAvailable()) {
+            return;
+        }
+
+        if (!mIsServiceDumpInProgress.compareAndSet(false, true)) {
+            return;
+        }
+
+        try {
+            Trace.beginSection("inputmethod_service_dump");
+            mDataSource.trace((ctx) -> {
+                final ProtoOutputStream os = ctx.newTracePacket();
+                os.write(TracePacket.TIMESTAMP, SystemClock.elapsedRealtimeNanos());
+                final long tokenWinscopeExtensions =
+                        os.start(TracePacket.WINSCOPE_EXTENSIONS);
+                final long tokenExtensionsField =
+                        os.start(WinscopeExtensionsImpl.INPUTMETHOD_SERVICE);
+                os.write(InputMethodServiceTraceProto.WHERE, where);
+                dumper.dumpToProto(os, icProto);
+                os.end(tokenExtensionsField);
+                os.end(tokenWinscopeExtensions);
+            });
+        } finally {
+            mIsServiceDumpInProgress.set(false);
+            Trace.endSection();
+        }
+    }
+
+    @Override
+    public void triggerManagerServiceDump(@NonNull String where, @NonNull ServiceDumper dumper) {
+        if (!isEnabled() || !isAvailable()) {
+            return;
+        }
+
+        if (!mIsManagerServiceDumpInProgress.compareAndSet(false, true)) {
+            return;
+        }
+
+        try {
+            Trace.beginSection("inputmethod_manager_service_dump");
+            mDataSource.trace((ctx) -> {
+                final ProtoOutputStream os = ctx.newTracePacket();
+                os.write(TracePacket.TIMESTAMP, SystemClock.elapsedRealtimeNanos());
+                final long tokenWinscopeExtensions =
+                        os.start(TracePacket.WINSCOPE_EXTENSIONS);
+                final long tokenExtensionsField =
+                        os.start(WinscopeExtensionsImpl.INPUTMETHOD_MANAGER_SERVICE);
+                os.write(InputMethodManagerServiceTraceProto.WHERE, where);
+                dumper.dumpToProto(os, null);
+                os.end(tokenExtensionsField);
+                os.end(tokenWinscopeExtensions);
+            });
+        } finally {
+            mIsManagerServiceDumpInProgress.set(false);
+            Trace.endSection();
+        }
+    }
+
+    @Override
+    public boolean isEnabled() {
+        return mTracingSessionsCount.get() > 0;
+    }
+
+    @Override
+    public void startTrace(@Nullable PrintWriter pw) {
+        // Intentionally left empty. Tracing start/stop is managed through Perfetto.
+    }
+
+    @Override
+    public void stopTrace(@Nullable PrintWriter pw) {
+        // Intentionally left empty. Tracing start/stop is managed through Perfetto.
+    }
+
+    @Override
+    public void addToBuffer(ProtoOutputStream proto, int source) {
+        // Intentionally left empty. Only used for legacy tracing.
+    }
+}