Update trace buffer to support SysUI & Launcher

- Move TraceBuffer to sysui-accessible package, and make it generic
  to support both lite and nano proto callers
- Expose shell command to start/stop sysui tracing

Bug: 144854916
Test: atest TraceBufferTest

Change-Id: Id117024d943f148a91631fd9fcae1fd70fca8ab5

9 files changed, 169 insertions, 36 deletions

diff --git a/core/java/com/android/internal/statusbar/IStatusBar.aidl b/core/java/com/android/internal/statusbar/IStatusBar.aidl
index a2736333383e..21067127a8e7 100644
--- a/core/java/com/android/internal/statusbar/IStatusBar.aidl
+++ b/core/java/com/android/internal/statusbar/IStatusBar.aidl
@@ -210,4 +210,14 @@ oneway interface IStatusBar
      * Cancels toast with token {@code token} in {@code packageName}.
      */
     void hideToast(String packageName, IBinder token);
+
+    /**
+     * Notifies SystemUI to start tracing.
+     */
+    void startTracing();
+
+    /**
+     * Notifies SystemUI to stop tracing.
+     */
+    void stopTracing();
 }
diff --git a/core/java/com/android/internal/statusbar/IStatusBarService.aidl b/core/java/com/android/internal/statusbar/IStatusBarService.aidl
index a9f7b8455807..a147a417bdda 100644
--- a/core/java/com/android/internal/statusbar/IStatusBarService.aidl
+++ b/core/java/com/android/internal/statusbar/IStatusBarService.aidl
@@ -123,4 +123,19 @@ interface IStatusBarService
      * Dismiss the warning that the device is about to go to sleep due to user inactivity.
      */
     void dismissInattentiveSleepWarning(boolean animated);
+
+    /**
+     * Notifies SystemUI to start tracing.
+     */
+    void startTracing();
+
+    /**
+     * Notifies SystemUI to stop tracing.
+     */
+    void stopTracing();
+
+    /**
+     * Returns whether SystemUI tracing is enabled.
+     */
+    boolean isTracing();
 }
diff --git a/services/core/java/com/android/server/utils/TraceBuffer.java b/core/java/com/android/internal/util/TraceBuffer.java
index 0567960e05e4..fe8a59e160bd 100644
--- a/services/core/java/com/android/server/utils/TraceBuffer.java
+++ b/core/java/com/android/internal/util/TraceBuffer.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.utils;
+package com.android.internal.util;
 
 import android.util.proto.ProtoOutputStream;
 
@@ -27,19 +27,87 @@ import java.io.OutputStream;
 import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.Queue;
+import java.util.function.Consumer;
 
 /**
  * Buffer used for tracing and logging.
+ *
+ * @param <P> The class type of the proto provider
+ * @param <S> The proto class type of the encapsulating proto
+ * @param <T> The proto class type of the individual entry protos in the buffer
+ *
+ * {@hide}
  */
-public class TraceBuffer {
+public class TraceBuffer<P, S extends P, T extends P> {
     private final Object mBufferLock = new Object();
 
-    private final Queue<ProtoOutputStream> mBuffer = new ArrayDeque<>();
+    private final ProtoProvider<P, S, T> mProtoProvider;
+    private final Queue<T> mBuffer = new ArrayDeque<>();
+    private final Consumer mProtoDequeuedCallback;
     private int mBufferUsedSize;
     private int mBufferCapacity;
 
+    /**
+     * An interface to get protos from different sources (ie. fw-proto/proto-lite/nano-proto) for
+     * the trace buffer.
+     *
+     * @param <P> The class type of the proto provider
+     * @param <S> The proto class type of the encapsulating proto
+     * @param <T> The proto class type of the individual protos in the buffer
+     */
+    public interface ProtoProvider<P, S extends P, T extends P> {
+        /**
+         * @return The size of the given proto.
+         */
+        int getItemSize(P proto);
+
+        /**
+         * @return The bytes of the given proto.
+         */
+        byte[] getBytes(P proto);
+
+        /**
+         * Writes the given encapsulating proto and buffer of protos to the given output
+         * stream.
+         */
+        void write(S encapsulatingProto, Queue<T> buffer, OutputStream os) throws IOException;
+    }
+
+    /**
+     * An implementation of the ProtoProvider that uses only the framework ProtoOutputStream.
+     */
+    private static class ProtoOutputStreamProvider implements
+            ProtoProvider<ProtoOutputStream, ProtoOutputStream, ProtoOutputStream> {
+        @Override
+        public int getItemSize(ProtoOutputStream proto) {
+            return proto.getRawSize();
+        }
+
+        @Override
+        public byte[] getBytes(ProtoOutputStream proto) {
+            return proto.getBytes();
+        }
+
+        @Override
+        public void write(ProtoOutputStream encapsulatingProto, Queue<ProtoOutputStream> buffer,
+                OutputStream os) throws IOException {
+            os.write(encapsulatingProto.getBytes());
+            for (ProtoOutputStream protoOutputStream : buffer) {
+                byte[] protoBytes = protoOutputStream.getBytes();
+                os.write(protoBytes);
+            }
+        }
+    }
+
     public TraceBuffer(int bufferCapacity) {
+        this(bufferCapacity, new ProtoOutputStreamProvider(), null);
+    }
+
+    public TraceBuffer(int bufferCapacity, ProtoProvider protoProvider,
+            Consumer<T> protoDequeuedCallback) {
         mBufferCapacity = bufferCapacity;
+        mProtoProvider = protoProvider;
+        mProtoDequeuedCallback = protoDequeuedCallback;
         resetBuffer();
     }
 
@@ -65,8 +133,8 @@ public class TraceBuffer {
      * @throws IllegalStateException if the element cannot be added because it is larger
      *                               than the buffer size.
      */
-    public void add(ProtoOutputStream proto) {
-        int protoLength = proto.getRawSize();
+    public void add(T proto) {
+        int protoLength = mProtoProvider.getItemSize(proto);
         if (protoLength > mBufferCapacity) {
             throw new IllegalStateException("Trace object too large for the buffer. Buffer size:"
                     + mBufferCapacity + " Object size: " + protoLength);
@@ -79,26 +147,22 @@ public class TraceBuffer {
         }
     }
 
-    boolean contains(byte[] other) {
+    @VisibleForTesting
+    public boolean contains(byte[] other) {
         return mBuffer.stream()
-                .anyMatch(p -> Arrays.equals(p.getBytes(), other));
+                .anyMatch(p -> Arrays.equals(mProtoProvider.getBytes(p), other));
     }
 
     /**
-     * Writes the trace buffer to disk inside the encapsulatingProto..
+     * Writes the trace buffer to disk inside the encapsulatingProto.
      */
-    public void writeTraceToFile(File traceFile, ProtoOutputStream encapsulatingProto)
+    public void writeTraceToFile(File traceFile, S encapsulatingProto)
             throws IOException {
         synchronized (mBufferLock) {
             traceFile.delete();
             try (OutputStream os = new FileOutputStream(traceFile)) {
                 traceFile.setReadable(true /* readable */, false /* ownerOnly */);
-                os.write(encapsulatingProto.getBytes());
-                for (ProtoOutputStream protoOutputStream : mBuffer) {
-                    encapsulatingProto = protoOutputStream;
-                    byte[] protoBytes = encapsulatingProto.getBytes();
-                    os.write(protoBytes);
-                }
+                mProtoProvider.write(encapsulatingProto, mBuffer, os);
                 os.flush();
             }
         }
@@ -115,12 +179,16 @@ public class TraceBuffer {
 
         while (availableSpace < protoLength) {
 
-            ProtoOutputStream item = mBuffer.poll();
+            P item = mBuffer.poll();
             if (item == null) {
                 throw new IllegalStateException("No element to discard from buffer");
             }
-            mBufferUsedSize -= item.getRawSize();
+            mBufferUsedSize -= mProtoProvider.getItemSize(item);
             availableSpace = getAvailableSpace();
+
+            if (mProtoDequeuedCallback != null) {
+                mProtoDequeuedCallback.accept(item);
+            }
         }
     }
 
@@ -129,13 +197,18 @@ public class TraceBuffer {
      */
     public void resetBuffer() {
         synchronized (mBufferLock) {
+            if (mProtoDequeuedCallback != null) {
+                for (T item : mBuffer) {
+                    mProtoDequeuedCallback.accept(item);
+                }
+            }
             mBuffer.clear();
             mBufferUsedSize = 0;
         }
     }
 
     @VisibleForTesting
-    int getBufferSize() {
+    public int getBufferSize() {
         return mBufferUsedSize;
     }
 
@@ -144,16 +217,9 @@ public class TraceBuffer {
      */
     public String getStatus() {
         synchronized (mBufferLock) {
-            return "Buffer size: "
-                    + mBufferCapacity
-                    + " bytes"
-                    + "\n"
-                    + "Buffer usage: "
-                    + mBufferUsedSize
-                    + " bytes"
-                    + "\n"
-                    + "Elements in the buffer: "
-                    + mBuffer.size();
+            return "Buffer size: " + mBufferCapacity + " bytes" + "\n"
+                    + "Buffer usage: " + mBufferUsedSize + " bytes" + "\n"
+                    + "Elements in the buffer: " + mBuffer.size();
         }
     }
 }
diff --git a/services/core/java/com/android/server/protolog/ProtoLogImpl.java b/services/core/java/com/android/server/protolog/ProtoLogImpl.java
index 1653b3d2ae28..c9d42c854b54 100644
--- a/services/core/java/com/android/server/protolog/ProtoLogImpl.java
+++ b/services/core/java/com/android/server/protolog/ProtoLogImpl.java
@@ -38,7 +38,7 @@ import android.util.proto.ProtoOutputStream;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.protolog.common.IProtoLogGroup;
 import com.android.server.protolog.common.LogDataType;
-import com.android.server.utils.TraceBuffer;
+import com.android.internal.util.TraceBuffer;
 import com.android.server.wm.ProtoLogGroup;
 
 import java.io.File;
diff --git a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
index 3f7d373c1848..54b7d3a496c2 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
@@ -79,7 +79,6 @@ public class StatusBarManagerService extends IStatusBarService.Stub implements D
 
     private final Context mContext;
 
-    private final WindowManagerService mWindowManager;
     private Handler mHandler = new Handler();
     private NotificationDelegate mNotificationDelegate;
     private volatile IStatusBar mBar;
@@ -93,6 +92,7 @@ public class StatusBarManagerService extends IStatusBarService.Stub implements D
     private final Object mLock = new Object();
     private final DeathRecipient mDeathRecipient = new DeathRecipient();
     private int mCurrentUserId;
+    private boolean mTracingEnabled;
 
     private SparseArray<UiState> mDisplayUiState = new SparseArray<>();
 
@@ -176,11 +176,10 @@ public class StatusBarManagerService extends IStatusBarService.Stub implements D
     }
 
     /**
-     * Construct the service, add the status bar view to the window manager
+     * Construct the service
      */
-    public StatusBarManagerService(Context context, WindowManagerService windowManager) {
+    public StatusBarManagerService(Context context) {
         mContext = context;
-        mWindowManager = windowManager;
 
         LocalServices.addService(StatusBarManagerInternal.class, mInternalService);
         LocalServices.addService(GlobalActionsProvider.class, mGlobalActionsProvider);
@@ -720,6 +719,31 @@ public class StatusBarManagerService extends IStatusBarService.Stub implements D
         }
     }
 
+    @Override
+    public void startTracing() {
+        if (mBar != null) {
+            try {
+                mBar.startTracing();
+                mTracingEnabled = true;
+            } catch (RemoteException ex) {}
+        }
+    }
+
+    @Override
+    public void stopTracing() {
+        if (mBar != null) {
+            try {
+                mTracingEnabled = false;
+                mBar.stopTracing();
+            } catch (RemoteException ex) {}
+        }
+    }
+
+    @Override
+    public boolean isTracing() {
+        return mTracingEnabled;
+    }
+
     // TODO(b/117478341): make it aware of multi-display if needed.
     @Override
     public void disable(int what, IBinder token, String pkg) {
diff --git a/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java b/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java
index d7f86cd65ae9..a79c19f7bc17 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java
@@ -72,6 +72,8 @@ public class StatusBarShellCommand extends ShellCommand {
                     return runDisableForSetup();
                 case "send-disable-flag":
                     return runSendDisableFlag();
+                case "tracing":
+                    return runTracing();
                 default:
                     return handleDefaultCommands(cmd);
             }
@@ -185,6 +187,18 @@ public class StatusBarShellCommand extends ShellCommand {
         return 0;
     }
 
+    private int runTracing() {
+        switch (getNextArg()) {
+            case "start":
+                mInterface.startTracing();
+                break;
+            case "stop":
+                mInterface.stopTracing();
+                break;
+        }
+        return 0;
+    }
+
     @Override
     public void onHelp() {
         final PrintWriter pw = getOutPrintWriter();
@@ -233,6 +247,9 @@ public class StatusBarShellCommand extends ShellCommand {
         pw.println("        clock               - disable clock appearing in status bar");
         pw.println("        notification-icons  - disable notification icons from status bar");
         pw.println("");
+        pw.println("  tracing (start | stop)");
+        pw.println("    Start or stop SystemUI tracing");
+        pw.println("");
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/WindowTracing.java b/services/core/java/com/android/server/wm/WindowTracing.java
index 0be90fcb2b84..ba3dc607f6cc 100644
--- a/services/core/java/com/android/server/wm/WindowTracing.java
+++ b/services/core/java/com/android/server/wm/WindowTracing.java
@@ -35,7 +35,7 @@ import android.util.proto.ProtoOutputStream;
 import android.view.Choreographer;
 
 import com.android.server.protolog.ProtoLogImpl;
-import com.android.server.utils.TraceBuffer;
+import com.android.internal.util.TraceBuffer;
 
 import java.io.File;
 import java.io.IOException;
@@ -64,7 +64,7 @@ class WindowTracing {
 
     private final Object mEnabledLock = new Object();
     private final File mTraceFile;
-    private final com.android.server.utils.TraceBuffer mBuffer;
+    private final TraceBuffer mBuffer;
     private final Choreographer.FrameCallback mFrameCallback = (frameTimeNanos) ->
             log("onFrame" /* where */);
 
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index 258d7628e502..8b61125b2612 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -1327,7 +1327,7 @@ public final class SystemServer {
             if (!isWatch) {
                 t.traceBegin("StartStatusBarManagerService");
                 try {
-                    statusBar = new StatusBarManagerService(context, wm);
+                    statusBar = new StatusBarManagerService(context);
                     ServiceManager.addService(Context.STATUS_BAR_SERVICE, statusBar);
                 } catch (Throwable e) {
                     reportWtf("starting StatusBarManagerService", e);
diff --git a/services/tests/servicestests/src/com/android/server/utils/TraceBufferTest.java b/services/tests/servicestests/src/com/android/server/utils/TraceBufferTest.java
index 09b75e71d946..0e50f2a1b6ad 100644
--- a/services/tests/servicestests/src/com/android/server/utils/TraceBufferTest.java
+++ b/services/tests/servicestests/src/com/android/server/utils/TraceBufferTest.java
@@ -30,6 +30,7 @@ import android.util.proto.ProtoOutputStream;
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.util.Preconditions;
+import com.android.internal.util.TraceBuffer;
 
 import org.junit.After;
 import org.junit.Before;