Inject log group mapping into ProtoLogImpl generated classes

We then keep track of the logging state of each group that way and in
particular use this to toggle groups to logcat on and off.

Test: adb shell wm logging enable-text WM_DEBUG_WINDOW_TRANSITIONS
Bug: 327491675
Change-Id: I2a2e4cfdc51113a5aaf2adc51aa10f5df47a5051

9 files changed, 84 insertions, 23 deletions

diff --git a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
index 30de5468001a..2096ba42080f 100644
--- a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
@@ -55,37 +55,38 @@ import java.util.stream.Collectors;
  * A service for the ProtoLog logging system.
  */
 public class LegacyProtoLogImpl implements IProtoLog {
-    private final TreeMap<String, IProtoLogGroup> mLogGroups = new TreeMap<>();
-
     private static final int BUFFER_CAPACITY = 1024 * 1024;
     private static final int PER_CHUNK_SIZE = 1024;
     private static final String TAG = "ProtoLog";
     private static final long MAGIC_NUMBER_VALUE = ((long) MAGIC_NUMBER_H << 32) | MAGIC_NUMBER_L;
     static final String PROTOLOG_VERSION = "2.0.0";
-    private static final int DEFAULT_PER_CHUNK_SIZE = 0;
 
     private final File mLogFile;
     private final String mLegacyViewerConfigFilename;
     private final TraceBuffer mBuffer;
     private final LegacyProtoLogViewerConfigReader mViewerConfig;
+    private final TreeMap<String, IProtoLogGroup> mLogGroups;
     private final int mPerChunkSize;
 
     private boolean mProtoLogEnabled;
     private boolean mProtoLogEnabledLockFree;
     private final Object mProtoLogEnabledLock = new Object();
 
-    public LegacyProtoLogImpl(String outputFile, String viewerConfigFilename) {
+    public LegacyProtoLogImpl(String outputFile, String viewerConfigFilename,
+            TreeMap<String, IProtoLogGroup> logGroups) {
         this(new File(outputFile), viewerConfigFilename, BUFFER_CAPACITY,
-                new LegacyProtoLogViewerConfigReader(), PER_CHUNK_SIZE);
+                new LegacyProtoLogViewerConfigReader(), PER_CHUNK_SIZE, logGroups);
     }
 
     public LegacyProtoLogImpl(File file, String viewerConfigFilename, int bufferCapacity,
-            LegacyProtoLogViewerConfigReader viewerConfig, int perChunkSize) {
+            LegacyProtoLogViewerConfigReader viewerConfig, int perChunkSize,
+            TreeMap<String, IProtoLogGroup> logGroups) {
         mLogFile = file;
         mBuffer = new TraceBuffer(bufferCapacity);
         mLegacyViewerConfigFilename = viewerConfigFilename;
         mViewerConfig = viewerConfig;
         mPerChunkSize = perChunkSize;
+        this.mLogGroups = logGroups;
     }
 
     /**
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 53062d837cfa..bdd9a913a06c 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -22,9 +22,9 @@ import static perfetto.protos.PerfettoTrace.InternedString.IID;
 import static perfetto.protos.PerfettoTrace.InternedString.STR;
 import static perfetto.protos.PerfettoTrace.ProtoLogMessage.BOOLEAN_PARAMS;
 import static perfetto.protos.PerfettoTrace.ProtoLogMessage.DOUBLE_PARAMS;
-import static perfetto.protos.PerfettoTrace.ProtoLogMessage.STACKTRACE_IID;
 import static perfetto.protos.PerfettoTrace.ProtoLogMessage.MESSAGE_ID;
 import static perfetto.protos.PerfettoTrace.ProtoLogMessage.SINT64_PARAMS;
+import static perfetto.protos.PerfettoTrace.ProtoLogMessage.STACKTRACE_IID;
 import static perfetto.protos.PerfettoTrace.ProtoLogMessage.STR_PARAM_IIDS;
 import static perfetto.protos.PerfettoTrace.ProtoLogViewerConfig.GROUPS;
 import static perfetto.protos.PerfettoTrace.ProtoLogViewerConfig.Group.ID;
@@ -78,7 +78,6 @@ import perfetto.protos.PerfettoTrace.ProtoLogViewerConfig.MessageData;
  * A service for the ProtoLog logging system.
  */
 public class PerfettoProtoLogImpl implements IProtoLog {
-    private final TreeMap<String, IProtoLogGroup> mLogGroups = new TreeMap<>();
     private static final String LOG_TAG = "ProtoLog";
     private final AtomicInteger mTracingInstances = new AtomicInteger();
 
@@ -89,8 +88,10 @@ public class PerfettoProtoLogImpl implements IProtoLog {
     );
     private final ProtoLogViewerConfigReader mViewerConfigReader;
     private final ViewerConfigInputStreamProvider mViewerConfigInputStreamProvider;
+    private final TreeMap<String, IProtoLogGroup> mLogGroups;
 
-    public PerfettoProtoLogImpl(String viewerConfigFilePath) {
+    public PerfettoProtoLogImpl(String viewerConfigFilePath,
+            TreeMap<String, IProtoLogGroup> logGroups) {
         this(() -> {
             try {
                 return new ProtoInputStream(new FileInputStream(viewerConfigFilePath));
@@ -98,23 +99,28 @@ public class PerfettoProtoLogImpl implements IProtoLog {
                 Slog.w(LOG_TAG, "Failed to load viewer config file " + viewerConfigFilePath, e);
                 return null;
             }
-        });
+        }, logGroups);
     }
 
-    public PerfettoProtoLogImpl(ViewerConfigInputStreamProvider viewerConfigInputStreamProvider) {
+    public PerfettoProtoLogImpl(
+            ViewerConfigInputStreamProvider viewerConfigInputStreamProvider,
+            TreeMap<String, IProtoLogGroup> logGroups
+    ) {
         this(viewerConfigInputStreamProvider,
-                new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider));
+                new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider), logGroups);
     }
 
     @VisibleForTesting
     public PerfettoProtoLogImpl(
             ViewerConfigInputStreamProvider viewerConfigInputStreamProvider,
-            ProtoLogViewerConfigReader viewerConfigReader
+            ProtoLogViewerConfigReader viewerConfigReader,
+            TreeMap<String, IProtoLogGroup> logGroups
     ) {
         Producer.init(InitArguments.DEFAULTS);
         mDataSource.register(DataSourceParams.DEFAULTS);
         this.mViewerConfigInputStreamProvider = viewerConfigInputStreamProvider;
         this.mViewerConfigReader = viewerConfigReader;
+        this.mLogGroups = logGroups;
     }
 
     /**
diff --git a/core/java/com/android/internal/protolog/ProtoLogImpl.java b/core/java/com/android/internal/protolog/ProtoLogImpl.java
index 896538564c2f..487ae8145546 100644
--- a/core/java/com/android/internal/protolog/ProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/ProtoLogImpl.java
@@ -18,6 +18,7 @@ package com.android.internal.protolog;
 
 import static com.android.internal.protolog.common.ProtoLogToolInjected.Value.LEGACY_OUTPUT_FILE_PATH;
 import static com.android.internal.protolog.common.ProtoLogToolInjected.Value.LEGACY_VIEWER_CONFIG_PATH;
+import static com.android.internal.protolog.common.ProtoLogToolInjected.Value.LOG_GROUPS;
 import static com.android.internal.protolog.common.ProtoLogToolInjected.Value.VIEWER_CONFIG_PATH;
 
 import android.annotation.Nullable;
@@ -28,6 +29,8 @@ import com.android.internal.protolog.common.IProtoLogGroup;
 import com.android.internal.protolog.common.LogLevel;
 import com.android.internal.protolog.common.ProtoLogToolInjected;
 
+import java.util.TreeMap;
+
 /**
  * A service for the ProtoLog logging system.
  */
@@ -43,6 +46,9 @@ public class ProtoLogImpl {
     @ProtoLogToolInjected(LEGACY_OUTPUT_FILE_PATH)
     private static String sLegacyOutputFilePath;
 
+    @ProtoLogToolInjected(LOG_GROUPS)
+    private static TreeMap<String, IProtoLogGroup> sLogGroups;
+
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
     public static void d(IProtoLogGroup group, long messageHash, int paramsMask,
             @Nullable String messageString,
@@ -99,11 +105,10 @@ public class ProtoLogImpl {
     public static synchronized IProtoLog getSingleInstance() {
         if (sServiceInstance == null) {
             if (android.tracing.Flags.perfettoProtologTracing()) {
-                sServiceInstance =
-                        new PerfettoProtoLogImpl(sViewerConfigPath);
+                sServiceInstance = new PerfettoProtoLogImpl(sViewerConfigPath, sLogGroups);
             } else {
-                sServiceInstance =
-                        new LegacyProtoLogImpl(sLegacyOutputFilePath, sLegacyViewerConfigPath);
+                sServiceInstance = new LegacyProtoLogImpl(
+                        sLegacyOutputFilePath, sLegacyViewerConfigPath, sLogGroups);
             }
         }
         return sServiceInstance;
diff --git a/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java b/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java
index 3c206acf7c0f..ae3d4488dbd9 100644
--- a/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java
+++ b/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java
@@ -4,6 +4,7 @@ import static perfetto.protos.PerfettoTrace.ProtoLogViewerConfig.MESSAGES;
 import static perfetto.protos.PerfettoTrace.ProtoLogViewerConfig.MessageData.MESSAGE;
 import static perfetto.protos.PerfettoTrace.ProtoLogViewerConfig.MessageData.MESSAGE_ID;
 
+import android.util.ArrayMap;
 import android.util.proto.ProtoInputStream;
 
 import com.android.internal.protolog.common.ILogger;
@@ -57,6 +58,7 @@ public class ProtoLogViewerConfigReader {
     }
 
     private void doLoadViewerConfig() throws IOException {
+        mLogMessageMap = new ArrayMap<>();
         final ProtoInputStream pis = mViewerConfigInputStreamProvider.getInputStream();
 
         while (pis.nextField() != ProtoInputStream.NO_MORE_FIELDS) {
diff --git a/core/java/com/android/internal/protolog/common/ProtoLogToolInjected.java b/core/java/com/android/internal/protolog/common/ProtoLogToolInjected.java
index ffd0d7623852..17c82d76408e 100644
--- a/core/java/com/android/internal/protolog/common/ProtoLogToolInjected.java
+++ b/core/java/com/android/internal/protolog/common/ProtoLogToolInjected.java
@@ -22,7 +22,9 @@ import java.lang.annotation.Target;
 
 @Target({ElementType.FIELD, ElementType.PARAMETER})
 public @interface ProtoLogToolInjected {
-    enum Value { VIEWER_CONFIG_PATH, LEGACY_OUTPUT_FILE_PATH, LEGACY_VIEWER_CONFIG_PATH }
+    enum Value {
+        VIEWER_CONFIG_PATH, LEGACY_OUTPUT_FILE_PATH, LEGACY_VIEWER_CONFIG_PATH, LOG_GROUPS
+    }
 
     Value value();
 }
diff --git a/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java b/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
index a64996c6838e..d9a4c261ee15 100644
--- a/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
@@ -59,6 +59,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.util.LinkedList;
+import java.util.TreeMap;
 
 /**
  * Test class for {@link ProtoLogImpl}.
@@ -89,7 +90,7 @@ public class LegacyProtoLogImplTest {
         //noinspection ResultOfMethodCallIgnored
         mFile.delete();
         mProtoLog = new LegacyProtoLogImpl(mFile, mViewerConfigFilename,
-                1024 * 1024, mReader, 1024);
+                1024 * 1024, mReader, 1024, new TreeMap<>());
     }
 
     @After
diff --git a/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java b/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
index 270f5957a9b9..548adeff07b2 100644
--- a/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
@@ -64,6 +64,7 @@ import java.io.File;
 import java.io.IOException;
 import java.util.List;
 import java.util.Random;
+import java.util.TreeMap;
 
 import perfetto.protos.Protolog;
 import perfetto.protos.ProtologCommon;
@@ -152,7 +153,8 @@ public class PerfettoProtoLogImplTest {
                 .thenAnswer(it -> new ProtoInputStream(mViewerConfigBuilder.build().toByteArray()));
 
         mReader = Mockito.spy(new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider));
-        mProtoLog = new PerfettoProtoLogImpl(viewerConfigInputStreamProvider, mReader);
+        mProtoLog =
+                new PerfettoProtoLogImpl(viewerConfigInputStreamProvider, mReader, new TreeMap<>());
     }
 
     @After
diff --git a/tools/protologtool/src/com/android/protolog/tool/CommandOptions.kt b/tools/protologtool/src/com/android/protolog/tool/CommandOptions.kt
index a3591558bd67..2690bc5ef405 100644
--- a/tools/protologtool/src/com/android/protolog/tool/CommandOptions.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/CommandOptions.kt
@@ -25,6 +25,7 @@ class CommandOptions(args: Array<String>) {
         const val READ_LOG_CMD = "read-log"
         private val commands = setOf(TRANSFORM_CALLS_CMD, GENERATE_CONFIG_CMD, READ_LOG_CMD)
 
+        // TODO: This is always the same. I don't think it's required
         private const val PROTOLOG_CLASS_PARAM = "--protolog-class"
         private const val PROTOLOGGROUP_CLASS_PARAM = "--loggroups-class"
         private const val PROTOLOGGROUP_JAR_PARAM = "--loggroups-jar"
diff --git a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
index ea2c8dacff50..837dae92b711 100644
--- a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
@@ -24,11 +24,17 @@ import com.github.javaparser.ParseProblemException
 import com.github.javaparser.ParserConfiguration
 import com.github.javaparser.StaticJavaParser
 import com.github.javaparser.ast.CompilationUnit
+import com.github.javaparser.ast.NodeList
 import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration
+import com.github.javaparser.ast.body.InitializerDeclaration
+import com.github.javaparser.ast.expr.FieldAccessExpr
 import com.github.javaparser.ast.expr.MethodCallExpr
+import com.github.javaparser.ast.expr.NameExpr
 import com.github.javaparser.ast.expr.NullLiteralExpr
+import com.github.javaparser.ast.expr.ObjectCreationExpr
 import com.github.javaparser.ast.expr.SimpleName
 import com.github.javaparser.ast.expr.StringLiteralExpr
+import com.github.javaparser.ast.stmt.BlockStmt
 import java.io.File
 import java.io.FileInputStream
 import java.io.FileNotFoundException
@@ -96,7 +102,8 @@ object ProtoLogTool {
         outJar.putNextEntry(zipEntry(protologImplPath))
 
         outJar.write(generateProtoLogImpl(protologImplName, command.viewerConfigFilePathArg,
-            command.legacyViewerConfigFilePathArg, command.legacyOutputFilePath).toByteArray())
+            command.legacyViewerConfigFilePathArg, command.legacyOutputFilePath,
+            groups, command.protoLogGroupsClassNameArg).toByteArray())
 
         val executor = newThreadPool()
 
@@ -140,6 +147,8 @@ object ProtoLogTool {
         viewerConfigFilePath: String,
         legacyViewerConfigFilePath: String?,
         legacyOutputFilePath: String?,
+        groups: Map<String, LogGroup>,
+        protoLogGroupsClassName: String,
     ): String {
         val file = File(PROTOLOG_IMPL_SRC_PATH)
 
@@ -160,7 +169,8 @@ object ProtoLogTool {
         classNameNode.setId(protoLogImplGenName)
 
         injectConstants(classDeclaration,
-            viewerConfigFilePath, legacyViewerConfigFilePath, legacyOutputFilePath)
+            viewerConfigFilePath, legacyViewerConfigFilePath, legacyOutputFilePath, groups,
+            protoLogGroupsClassName)
 
         return code.toString()
     }
@@ -169,7 +179,9 @@ object ProtoLogTool {
         classDeclaration: ClassOrInterfaceDeclaration,
         viewerConfigFilePath: String,
         legacyViewerConfigFilePath: String?,
-        legacyOutputFilePath: String?
+        legacyOutputFilePath: String?,
+        groups: Map<String, LogGroup>,
+        protoLogGroupsClassName: String
     ) {
         classDeclaration.fields.forEach { field ->
             field.getAnnotationByClass(ProtoLogToolInjected::class.java)
@@ -197,6 +209,35 @@ object ProtoLogTool {
                                                 StringLiteralExpr(it)
                                             } ?: NullLiteralExpr())
                                 }
+                                ProtoLogToolInjected.Value.LOG_GROUPS.name -> {
+                                    val initializerBlockStmt = BlockStmt()
+                                    for (group in groups) {
+                                        initializerBlockStmt.addStatement(
+                                            MethodCallExpr()
+                                                    .setName("put")
+                                                    .setArguments(
+                                                        NodeList(StringLiteralExpr(group.key),
+                                                            FieldAccessExpr()
+                                                                    .setScope(
+                                                                        NameExpr(
+                                                                            protoLogGroupsClassName
+                                                                        ))
+                                                                    .setName(group.value.name)))
+                                        )
+                                        group.key
+                                    }
+
+                                    val treeMapCreation = ObjectCreationExpr()
+                                            .setType("TreeMap<String, IProtoLogGroup>")
+                                            .setAnonymousClassBody(NodeList(
+                                                InitializerDeclaration().setBody(
+                                                    initializerBlockStmt
+                                                )
+                                            ))
+
+                                    field.setFinal(true)
+                                    field.variables.first().setInitializer(treeMapCreation)
+                                }
                                 else -> error("Unhandled ProtoLogToolInjected value: $valueName.")
                             }
                         }