Fix missing ProtoLog viewer config in some tracing instances

We were re-using the same ProtoInputStream across all trace function calls, which meant that only the first tracing instance would get the viewer config and all other tracing instances would be missing the viewer config.

This issue became very common after enabling the AOT protolog config for eng builds.

Change-Id: I05287f64e57ff8ec53e474e368a799fdfd9f62b8
Flag: EXEMPT small bug fix
Bug: 358100425
Test: atest InternalTests:com.android.internal.protolog.PerfettoProtoLogImplTest#handlesConcurrentTracingSessions

2 files changed, 44 insertions, 7 deletions

diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 78b5cfed1771..49ed55dbdb9f 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -431,15 +431,10 @@ public class PerfettoProtoLogImpl extends IProtoLogClient.Stub implements IProto
 
         Log.d(LOG_TAG, "Dumping viewer config to trace");
 
-        ProtoInputStream pis = mViewerConfigInputStreamProvider.getInputStream();
-
-        if (pis == null) {
-            Slog.w(LOG_TAG, "Failed to get viewer input stream.");
-            return;
-        }
-
         mDataSource.trace(ctx -> {
             try {
+                ProtoInputStream pis = mViewerConfigInputStreamProvider.getInputStream();
+
                 final ProtoOutputStream os = ctx.newTracePacket();
 
                 os.write(TIMESTAMP, SystemClock.elapsedRealtimeNanos());
diff --git a/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java b/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
index 4826f4241e2b..05a68e9649d4 100644
--- a/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
@@ -756,6 +756,48 @@ public class PerfettoProtoLogImplTest {
                 .isEqualTo("My null args: 0, 0, false");
     }
 
+    @Test
+    public void handlesConcurrentTracingSessions() throws IOException {
+        PerfettoTraceMonitor traceMonitor1 =
+                PerfettoTraceMonitor.newBuilder().enableProtoLog(true)
+                        .build();
+
+        PerfettoTraceMonitor traceMonitor2 =
+                PerfettoTraceMonitor.newBuilder().enableProtoLog(true)
+                        .build();
+
+        final ResultWriter writer2 = new ResultWriter()
+                .forScenario(new ScenarioBuilder()
+                        .forClass(createTempFile("temp", "").getName()).build())
+                .withOutputDir(mTracingDirectory)
+                .setRunComplete();
+
+        try {
+            traceMonitor1.start();
+            traceMonitor2.start();
+
+            mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP, 1,
+                    LogDataType.BOOLEAN, new Object[]{true});
+        } finally {
+            traceMonitor1.stop(mWriter);
+            traceMonitor2.stop(writer2);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final ProtoLogTrace protologFromMonitor1 = reader.readProtoLogTrace();
+
+        final ResultReader reader2 = new ResultReader(writer2.write(), mTraceConfig);
+        final ProtoLogTrace protologFromMonitor2 = reader2.readProtoLogTrace();
+
+        Truth.assertThat(protologFromMonitor1.messages).hasSize(1);
+        Truth.assertThat(protologFromMonitor1.messages.get(0).getMessage())
+                .isEqualTo("My Test Debug Log Message true");
+
+        Truth.assertThat(protologFromMonitor2.messages).hasSize(1);
+        Truth.assertThat(protologFromMonitor2.messages.get(0).getMessage())
+                .isEqualTo("My Test Debug Log Message true");
+    }
+
     private enum TestProtoLogGroup implements IProtoLogGroup {
         TEST_GROUP(true, true, false, "TEST_TAG");