Fix timing race condition in toggling ProtoLog to logcat

It was possible for the viewer config to be unloaded before we updated the log to proto state variable to false leading to crashes.

Change-Id: I26b2ec045df22cfbd417452794435894a9cc092b
Flag: EXEMPT small bug fix
Test: atest com.android.internal.protolog.ProcessedPerfettoProtoLogImplTest

3 files changed, 68 insertions, 3 deletions

diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index a1c987f79304..eb682dff14de 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -676,15 +676,30 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
         return internMap.get(string);
     }
 
+    protected boolean validateGroups(ILogger logger, String[] groups) {
+        for (int i = 0; i < groups.length; i++) {
+            String group = groups[i];
+            IProtoLogGroup g = mLogGroups.get(group);
+            if (g == null) {
+                logger.log("No IProtoLogGroup named " + group);
+                return false;
+            }
+        }
+        return true;
+    }
+
     private int setTextLogging(boolean value, ILogger logger, String... groups) {
+        if (!validateGroups(logger, groups)) {
+            return -1;
+        }
+
         for (int i = 0; i < groups.length; i++) {
             String group = groups[i];
             IProtoLogGroup g = mLogGroups.get(group);
             if (g != null) {
                 g.setLogToLogcat(value);
             } else {
-                logger.log("No IProtoLogGroup named " + group);
-                return -1;
+                throw new RuntimeException("No IProtoLogGroup named " + group);
             }
         }
 
diff --git a/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java
index febe1f3a72ac..d213b5de1e87 100644
--- a/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java
@@ -113,6 +113,10 @@ public class ProcessedPerfettoProtoLogImpl extends PerfettoProtoLogImpl {
      */
     @Override
     public int startLoggingToLogcat(String[] groups, @NonNull ILogger logger) {
+        if (!validateGroups(logger, groups)) {
+            return -1;
+        }
+
         mViewerConfigReader.loadViewerConfig(groups, logger);
         return super.startLoggingToLogcat(groups, logger);
     }
@@ -125,8 +129,19 @@ public class ProcessedPerfettoProtoLogImpl extends PerfettoProtoLogImpl {
      */
     @Override
     public int stopLoggingToLogcat(String[] groups, @NonNull ILogger logger) {
+        if (!validateGroups(logger, groups)) {
+            return -1;
+        }
+
+        var status = super.stopLoggingToLogcat(groups, logger);
+
+        if (status != 0) {
+            throw new RuntimeException("Failed to stop logging to logcat");
+        }
+
+        // If we successfully disabled logging, unload the viewer config.
         mViewerConfigReader.unloadViewerConfig(groups, logger);
-        return super.stopLoggingToLogcat(groups, logger);
+        return status;
     }
 
     @Deprecated
diff --git a/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java b/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java
index 2692e12c57ed..44641f7a1e12 100644
--- a/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java
+++ b/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java
@@ -24,6 +24,7 @@ import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -57,6 +58,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mockito;
+import org.mockito.stubbing.Answer;
 
 import perfetto.protos.Protolog;
 import perfetto.protos.ProtologCommon;
@@ -858,6 +860,39 @@ public class ProcessedPerfettoProtoLogImplTest {
                 .isEqualTo("This message should also be logged 567");
     }
 
+    @Test
+    public void enablesLogGroupAfterLoadingConfig() {
+        sProtoLog.stopLoggingToLogcat(
+                new String[] { TestProtoLogGroup.TEST_GROUP.name() }, (msg) -> {});
+        Truth.assertThat(TestProtoLogGroup.TEST_GROUP.isLogToLogcat()).isFalse();
+
+        doAnswer((Answer<Void>) invocation -> {
+            // logToLogcat is still false before we laod the viewer config
+            Truth.assertThat(TestProtoLogGroup.TEST_GROUP.isLogToLogcat()).isFalse();
+            return null;
+        }).when(sReader).unloadViewerConfig(any(), any());
+
+        sProtoLog.startLoggingToLogcat(
+                new String[] { TestProtoLogGroup.TEST_GROUP.name() }, (msg) -> {});
+        Truth.assertThat(TestProtoLogGroup.TEST_GROUP.isLogToLogcat()).isTrue();
+    }
+
+    @Test
+    public void disablesLogGroupBeforeUnloadingConfig() {
+        sProtoLog.startLoggingToLogcat(
+                new String[] { TestProtoLogGroup.TEST_GROUP.name() }, (msg) -> {});
+        Truth.assertThat(TestProtoLogGroup.TEST_GROUP.isLogToLogcat()).isTrue();
+
+        doAnswer((Answer<Void>) invocation -> {
+            // Already set logToLogcat to false by the time we unload the config
+            Truth.assertThat(TestProtoLogGroup.TEST_GROUP.isLogToLogcat()).isFalse();
+            return null;
+        }).when(sReader).unloadViewerConfig(any(), any());
+        sProtoLog.stopLoggingToLogcat(
+                new String[] { TestProtoLogGroup.TEST_GROUP.name() }, (msg) -> {});
+        Truth.assertThat(TestProtoLogGroup.TEST_GROUP.isLogToLogcat()).isFalse();
+    }
+
     private enum TestProtoLogGroup implements IProtoLogGroup {
         TEST_GROUP(true, true, false, "TEST_TAG");