Merge "Allow ProtoLog.init to be called multiple times in the same process" into main

6 files changed, 162 insertions, 1 deletions

diff --git a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
index 2feb3d50cf64..8771cde1de46 100644
--- a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
@@ -29,6 +29,7 @@ import static com.android.internal.protolog.ProtoLogMessage.MESSAGE_HASH;
 import static com.android.internal.protolog.ProtoLogMessage.SINT64_PARAMS;
 import static com.android.internal.protolog.ProtoLogMessage.STR_PARAMS;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.ShellCommand;
 import android.os.SystemClock;
@@ -49,6 +50,7 @@ import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
@@ -419,6 +421,12 @@ public class LegacyProtoLogImpl implements IProtoLog {
         return group.isLogToLogcat() || (group.isLogToProto() && isProtoEnabled());
     }
 
+    @Override
+    @NonNull
+    public List<IProtoLogGroup> getRegisteredGroups() {
+        return mLogGroups.values().stream().toList();
+    }
+
     public void registerGroups(IProtoLogGroup... protoLogGroups) {
         for (IProtoLogGroup group : protoLogGroups) {
             mLogGroups.put(group.name(), group);
diff --git a/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java b/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java
index e8d5195d121d..b82c660fcf57 100644
--- a/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java
@@ -18,6 +18,7 @@ package com.android.internal.protolog;
 
 import static com.android.internal.protolog.ProtoLog.REQUIRE_PROTOLOGTOOL;
 
+import android.annotation.NonNull;
 import android.text.TextUtils;
 import android.util.Log;
 
@@ -26,6 +27,9 @@ import com.android.internal.protolog.common.IProtoLog;
 import com.android.internal.protolog.common.IProtoLogGroup;
 import com.android.internal.protolog.common.LogLevel;
 
+import java.util.Collections;
+import java.util.List;
+
 /**
  * Class only create and used to server temporarily for when there is source code pre-processing by
  * the ProtoLog tool, when the tracing to Perfetto flag is off, and the static REQUIRE_PROTOLOGTOOL
@@ -79,4 +83,10 @@ public class LogcatOnlyProtoLogImpl implements IProtoLog {
     public boolean isEnabled(IProtoLogGroup group, LogLevel level) {
         return true;
     }
+
+    @Override
+    @NonNull
+    public List<IProtoLogGroup> getRegisteredGroups() {
+        return Collections.emptyList();
+    }
 }
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 49ed55dbdb9f..5517967341b5 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -306,6 +306,12 @@ public class PerfettoProtoLogImpl extends IProtoLogClient.Stub implements IProto
                 || group.isLogToLogcat();
     }
 
+    @Override
+    @NonNull
+    public List<IProtoLogGroup> getRegisteredGroups() {
+        return mLogGroups.values().stream().toList();
+    }
+
     private void registerGroupsLocally(@NonNull IProtoLogGroup[] protoLogGroups) {
         final var groupsLoggingToLogcat = new ArrayList<String>();
         for (IProtoLogGroup protoLogGroup : protoLogGroups) {
diff --git a/core/java/com/android/internal/protolog/ProtoLog.java b/core/java/com/android/internal/protolog/ProtoLog.java
index f9b989477907..660d3c99f538 100644
--- a/core/java/com/android/internal/protolog/ProtoLog.java
+++ b/core/java/com/android/internal/protolog/ProtoLog.java
@@ -20,6 +20,9 @@ import com.android.internal.protolog.common.IProtoLog;
 import com.android.internal.protolog.common.IProtoLogGroup;
 import com.android.internal.protolog.common.LogLevel;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+
 /**
  * ProtoLog API - exposes static logging methods. Usage of this API is similar
  * to {@code android.utils.Log} class. Instead of plain text log messages each call consists of
@@ -49,6 +52,8 @@ public class ProtoLog {
 
     private static IProtoLog sProtoLogInstance;
 
+    private static final Object sInitLock = new Object();
+
     /**
      * Initialize ProtoLog in this process.
      * <p>
@@ -59,7 +64,17 @@ public class ProtoLog {
      */
     public static void init(IProtoLogGroup... groups) {
         if (android.tracing.Flags.perfettoProtologTracing()) {
-            sProtoLogInstance = new PerfettoProtoLogImpl(groups);
+            synchronized (sInitLock) {
+                if (sProtoLogInstance != null) {
+                    // The ProtoLog instance has already been initialized in this process
+                    final var alreadyRegisteredGroups = sProtoLogInstance.getRegisteredGroups();
+                    final var allGroups = new ArrayList<>(alreadyRegisteredGroups);
+                    allGroups.addAll(Arrays.stream(groups).toList());
+                    groups = allGroups.toArray(new IProtoLogGroup[0]);
+                }
+
+                sProtoLogInstance = new PerfettoProtoLogImpl(groups);
+            }
         } else {
             // The first call to ProtoLog is likely to flip REQUIRE_PROTOLOGTOOL, which is when this
             // static block will be executed before REQUIRE_PROTOLOGTOOL is actually set.
diff --git a/core/java/com/android/internal/protolog/common/IProtoLog.java b/core/java/com/android/internal/protolog/common/IProtoLog.java
index d5c2ac10852f..f06f08a9f02d 100644
--- a/core/java/com/android/internal/protolog/common/IProtoLog.java
+++ b/core/java/com/android/internal/protolog/common/IProtoLog.java
@@ -16,6 +16,8 @@
 
 package com.android.internal.protolog.common;
 
+import java.util.List;
+
 /**
  * Interface for ProtoLog implementations.
  */
@@ -68,4 +70,9 @@ public interface IProtoLog {
      * @return If we need to log this group and level to either ProtoLog or Logcat.
      */
     boolean isEnabled(IProtoLogGroup group, LogLevel level);
+
+    /**
+     * @return an immutable list of the registered ProtoLog groups in this ProtoLog instance.
+     */
+    List<IProtoLogGroup> getRegisteredGroups();
 }
diff --git a/tests/Internal/src/com/android/internal/protolog/ProtoLogTest.java b/tests/Internal/src/com/android/internal/protolog/ProtoLogTest.java
new file mode 100644
index 000000000000..9d56a92fad52
--- /dev/null
+++ b/tests/Internal/src/com/android/internal/protolog/ProtoLogTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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.protolog;
+
+import android.platform.test.annotations.Presubmit;
+
+import com.android.internal.protolog.common.IProtoLogGroup;
+
+import com.google.common.truth.Truth;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Test class for {@link ProtoLog}. */
+@SuppressWarnings("ConstantConditions")
+@Presubmit
+@RunWith(JUnit4.class)
+public class ProtoLogTest {
+
+    @Test
+    public void canRunProtoLogInitMultipleTimes() {
+        ProtoLog.init(TEST_GROUP_1);
+        ProtoLog.init(TEST_GROUP_1);
+        ProtoLog.init(TEST_GROUP_2);
+        ProtoLog.init(TEST_GROUP_1, TEST_GROUP_2);
+
+        final var instance = ProtoLog.getSingleInstance();
+        Truth.assertThat(instance.getRegisteredGroups())
+                .containsExactly(TEST_GROUP_1, TEST_GROUP_2);
+    }
+
+    private static final IProtoLogGroup TEST_GROUP_1 = new ProtoLogGroup("TEST_TAG_1", 1);
+    private static final IProtoLogGroup TEST_GROUP_2 = new ProtoLogGroup("TEST_TAG_2", 2);
+
+    private static class ProtoLogGroup implements IProtoLogGroup {
+        private final boolean mEnabled;
+        private volatile boolean mLogToProto;
+        private volatile boolean mLogToLogcat;
+        private final String mTag;
+        private final int mId;
+
+        ProtoLogGroup(String tag, int id) {
+            this(true, true, false, tag, id);
+        }
+
+        ProtoLogGroup(
+                boolean enabled, boolean logToProto, boolean logToLogcat, String tag, int id) {
+            this.mEnabled = enabled;
+            this.mLogToProto = logToProto;
+            this.mLogToLogcat = logToLogcat;
+            this.mTag = tag;
+            this.mId = id;
+        }
+
+        @Override
+        public String name() {
+            return mTag;
+        }
+
+        @Override
+        public boolean isEnabled() {
+            return mEnabled;
+        }
+
+        @Override
+        public boolean isLogToProto() {
+            return mLogToProto;
+        }
+
+        @Override
+        public boolean isLogToLogcat() {
+            return mLogToLogcat;
+        }
+
+        @Override
+        public boolean isLogToAny() {
+            return mLogToLogcat || mLogToProto;
+        }
+
+        @Override
+        public String getTag() {
+            return mTag;
+        }
+
+        @Override
+        public void setLogToProto(boolean logToProto) {
+            this.mLogToProto = logToProto;
+        }
+
+        @Override
+        public void setLogToLogcat(boolean logToLogcat) {
+            this.mLogToLogcat = logToLogcat;
+        }
+
+        @Override
+        public int getId() {
+            return mId;
+        }
+    }
+}