Merge changes from topic "protolog-reuse-datasource" into main

* changes:
  Update ProtoLog cache updater to take an ProtoLog instance object
  Refactor PerfettoProtoLogImpl classes to take the datasource object directly and be enabled and disabled on demand
  Add options to register lifecycle callbacks on a ProtoLogDataSource
  Create shared static ProtoLogDataSource

14 files changed, 334 insertions, 155 deletions

diff --git a/apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java b/apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java
index 7ef8c53d1d62..7168fbec07bc 100644
--- a/apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java
+++ b/apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java
@@ -19,6 +19,9 @@ import android.app.Activity;
 import android.os.Bundle;
 import android.os.ServiceManager.ServiceNotFoundException;
 import android.perftests.utils.Stats;
+import android.tracing.perfetto.DataSourceParams;
+import android.tracing.perfetto.InitArguments;
+import android.tracing.perfetto.Producer;
 
 import androidx.test.InstrumentationRegistry;
 
@@ -70,6 +73,8 @@ public class ProtoLogPerfTest {
     }
 
     private IProtoLog mProcessedProtoLogger;
+    private static ProtoLogDataSource sTestDataSource;
+    private static final String TEST_PROTOLOG_DATASOURCE_NAME = "test.android.protolog";
     private static final String MOCK_TEST_FILE_PATH = "mock/file/path";
     private static final perfetto.protos.Protolog.ProtoLogViewerConfig VIEWER_CONFIG =
             perfetto.protos.Protolog.ProtoLogViewerConfig.newBuilder()
@@ -89,6 +94,17 @@ public class ProtoLogPerfTest {
 
     @BeforeClass
     public static void init() {
+        Producer.init(InitArguments.DEFAULTS);
+
+        sTestDataSource = new ProtoLogDataSource(TEST_PROTOLOG_DATASOURCE_NAME);
+        DataSourceParams params =
+                new DataSourceParams.Builder()
+                        .setBufferExhaustedPolicy(
+                                DataSourceParams
+                                        .PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP)
+                        .build();
+        sTestDataSource.register(params);
+
         ProtoLog.init(TestProtoLogGroup.values());
     }
 
@@ -98,9 +114,10 @@ public class ProtoLogPerfTest {
         TestProtoLogGroup.TEST_GROUP.setLogToLogcat(mLogToLogcat);
 
         mProcessedProtoLogger = new ProcessedPerfettoProtoLogImpl(
+                sTestDataSource,
                 MOCK_TEST_FILE_PATH,
                 () -> new AutoClosableProtoInputStream(VIEWER_CONFIG.toByteArray()),
-                () -> {},
+                (instance) -> {},
                 TestProtoLogGroup.values()
         );
     }
diff --git a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
index 8771cde1de46..be3bbb2d1d94 100644
--- a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
@@ -70,7 +70,7 @@ public class LegacyProtoLogImpl implements IProtoLog {
     private final TraceBuffer mBuffer;
     private final LegacyProtoLogViewerConfigReader mViewerConfig;
     private final Map<String, IProtoLogGroup> mLogGroups = new TreeMap<>();
-    private final Runnable mCacheUpdater;
+    private final ProtoLogCacheUpdater mCacheUpdater;
     private final int mPerChunkSize;
 
     private boolean mProtoLogEnabled;
@@ -78,14 +78,14 @@ public class LegacyProtoLogImpl implements IProtoLog {
     private final Object mProtoLogEnabledLock = new Object();
 
     public LegacyProtoLogImpl(String outputFile, String viewerConfigFilename,
-            Runnable cacheUpdater) {
+            ProtoLogCacheUpdater cacheUpdater) {
         this(new File(outputFile), viewerConfigFilename, BUFFER_CAPACITY,
                 new LegacyProtoLogViewerConfigReader(), PER_CHUNK_SIZE, cacheUpdater);
     }
 
     public LegacyProtoLogImpl(File file, String viewerConfigFilename, int bufferCapacity,
             LegacyProtoLogViewerConfigReader viewerConfig, int perChunkSize,
-            Runnable cacheUpdater) {
+            ProtoLogCacheUpdater cacheUpdater) {
         mLogFile = file;
         mBuffer = new TraceBuffer(bufferCapacity);
         mLegacyViewerConfigFilename = viewerConfigFilename;
@@ -298,7 +298,7 @@ public class LegacyProtoLogImpl implements IProtoLog {
             }
         }
 
-        mCacheUpdater.run();
+        mCacheUpdater.update(this);
         return 0;
     }
 
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index eb682dff14de..05a33fe830e8 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -51,7 +51,6 @@ import android.os.ServiceManager;
 import android.os.ShellCommand;
 import android.os.SystemClock;
 import android.text.TextUtils;
-import android.tracing.perfetto.DataSourceParams;
 import android.tracing.perfetto.InitArguments;
 import android.tracing.perfetto.Producer;
 import android.tracing.perfetto.TracingContext;
@@ -86,7 +85,6 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.stream.Stream;
 
 /**
  * A service for the ProtoLog logging system.
@@ -98,10 +96,12 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
 
     @NonNull
     protected final ProtoLogDataSource mDataSource;
+    @Nullable
+    protected final IProtoLogConfigurationService mConfigurationService;
     @NonNull
     protected final TreeMap<String, IProtoLogGroup> mLogGroups = new TreeMap<>();
     @NonNull
-    private final Runnable mCacheUpdater;
+    private final ProtoLogCacheUpdater mCacheUpdater;
 
     @NonNull
     private final int[] mDefaultLogLevelCounts = new int[LogLevel.values().length];
@@ -117,10 +117,10 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
     private boolean mLogcatReady = false;
 
     protected PerfettoProtoLogImpl(
-            @NonNull Runnable cacheUpdater,
+            @NonNull ProtoLogDataSource dataSource,
+            @NonNull ProtoLogCacheUpdater cacheUpdater,
             @NonNull IProtoLogGroup[] groups) throws ServiceManager.ServiceNotFoundException {
-        this(cacheUpdater, groups,
-                ProtoLogDataSource::new,
+        this(dataSource, cacheUpdater, groups,
                 android.tracing.Flags.clientSideProtoLogging() ?
                     IProtoLogConfigurationService.Stub.asInterface(
                         ServiceManager.getServiceOrThrow(PROTOLOG_CONFIGURATION_SERVICE)
@@ -129,49 +129,62 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
     }
 
     protected PerfettoProtoLogImpl(
-            @NonNull Runnable cacheUpdater,
+            @NonNull ProtoLogDataSource dataSource,
+            @NonNull ProtoLogCacheUpdater cacheUpdater,
             @NonNull IProtoLogGroup[] groups,
-            @NonNull ProtoLogDataSourceBuilder dataSourceBuilder,
             @Nullable IProtoLogConfigurationService configurationService) {
-        mDataSource = dataSourceBuilder.build(
-                this::onTracingInstanceStart,
-                this::onTracingFlush,
-                this::onTracingInstanceStop);
-        Producer.init(InitArguments.DEFAULTS);
-        DataSourceParams params =
-                new DataSourceParams.Builder()
-                        .setBufferExhaustedPolicy(
-                                DataSourceParams
-                                        .PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP)
-                        .build();
-        // NOTE: Registering that datasource is an async operation, so there may be no data traced
-        // for some messages logged right after the construction of this class.
-        mDataSource.register(params);
-
-        this.mCacheUpdater = cacheUpdater;
+        mDataSource = dataSource;
+        mCacheUpdater = cacheUpdater;
+        mConfigurationService = configurationService;
 
         registerGroupsLocally(groups);
+    }
+
+    /**
+     * To be called to enable the ProtoLogImpl to start tracing to ProtoLog and register with all
+     * the expected ProtoLog components.
+     */
+    public void enable() {
+        Producer.init(InitArguments.DEFAULTS);
 
         if (android.tracing.Flags.clientSideProtoLogging()) {
-            Objects.requireNonNull(configurationService,
-                    "A null ProtoLog Configuration Service was provided!");
+            connectToConfigurationService();
+        }
 
-            try {
-                var args = createConfigurationServiceRegisterClientArgs();
+        mDataSource.registerOnStartCallback(this::onTracingInstanceStart);
+        mDataSource.registerOnFlushCallback(this::onTracingFlush);
+        mDataSource.registerOnStopCallback(this::onTracingInstanceStop);
+    }
 
-                final var groupArgs = Stream.of(groups)
-                        .map(group -> new RegisterClientArgs
-                                .GroupConfig(group.name(), group.isLogToLogcat()))
-                        .toArray(RegisterClientArgs.GroupConfig[]::new);
-                args.setGroups(groupArgs);
+    private void connectToConfigurationService() {
+        Objects.requireNonNull(mConfigurationService,
+                "A null ProtoLog Configuration Service was provided!");
 
-                configurationService.registerClient(this, args);
-            } catch (RemoteException e) {
-                throw new RuntimeException("Failed to register ProtoLog client");
-            }
+        try {
+            var args = createConfigurationServiceRegisterClientArgs();
+
+            final var groupArgs = mLogGroups.values().stream()
+                    .map(group -> new RegisterClientArgs
+                            .GroupConfig(group.name(), group.isLogToLogcat()))
+                    .toArray(RegisterClientArgs.GroupConfig[]::new);
+            args.setGroups(groupArgs);
+
+            mConfigurationService.registerClient(this, args);
+        } catch (RemoteException e) {
+            throw new RuntimeException("Failed to register ProtoLog client");
         }
     }
 
+    /**
+     * Should be called when we no longer want to use the ProtoLog logger to unlink ourselves from
+     * the datasource and the configuration service to ensure we no longer receive the callback.
+     */
+    public void disable() {
+        mDataSource.unregisterOnStartCallback(this::onTracingInstanceStart);
+        mDataSource.unregisterOnFlushCallback(this::onTracingFlush);
+        mDataSource.unregisterOnStopCallback(this::onTracingInstanceStop);
+    }
+
     @NonNull
     protected abstract RegisterClientArgs createConfigurationServiceRegisterClientArgs();
 
@@ -703,7 +716,7 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
             }
         }
 
-        mCacheUpdater.run();
+        mCacheUpdater.update(this);
         return 0;
     }
 
@@ -746,7 +759,7 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
             }
         }
 
-        mCacheUpdater.run();
+        mCacheUpdater.update(this);
 
         this.mTracingInstances.incrementAndGet();
 
@@ -786,7 +799,7 @@ public abstract class PerfettoProtoLogImpl extends IProtoLogClient.Stub implemen
             }
         }
 
-        mCacheUpdater.run();
+        mCacheUpdater.update(this);
         Log.d(LOG_TAG, "Finished onTracingInstanceStop");
     }
 
diff --git a/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java
index 967a5ed1744d..e0a77d2be724 100644
--- a/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/ProcessedPerfettoProtoLogImpl.java
@@ -42,10 +42,11 @@ public class ProcessedPerfettoProtoLogImpl extends PerfettoProtoLogImpl {
     private final String mViewerConfigFilePath;
 
     public ProcessedPerfettoProtoLogImpl(
+            @NonNull ProtoLogDataSource datasource,
             @NonNull String viewerConfigFilePath,
-            @NonNull Runnable cacheUpdater,
+            @NonNull ProtoLogCacheUpdater cacheUpdater,
             @NonNull IProtoLogGroup[] groups) throws ServiceManager.ServiceNotFoundException {
-        this(viewerConfigFilePath, new ViewerConfigInputStreamProvider() {
+        this(datasource, viewerConfigFilePath, new ViewerConfigInputStreamProvider() {
                     @NonNull
                     @Override
                     public AutoClosableProtoInputStream getInputStream() {
@@ -64,11 +65,12 @@ public class ProcessedPerfettoProtoLogImpl extends PerfettoProtoLogImpl {
 
     @VisibleForTesting
     public ProcessedPerfettoProtoLogImpl(
+            @NonNull ProtoLogDataSource datasource,
             @NonNull String viewerConfigFilePath,
             @NonNull ViewerConfigInputStreamProvider viewerConfigInputStreamProvider,
-            @NonNull Runnable cacheUpdater,
+            @NonNull ProtoLogCacheUpdater cacheUpdater,
             @NonNull IProtoLogGroup[] groups) throws ServiceManager.ServiceNotFoundException {
-        super(cacheUpdater, groups);
+        super(datasource, cacheUpdater, groups);
 
         this.mViewerConfigFilePath = viewerConfigFilePath;
 
@@ -80,15 +82,15 @@ public class ProcessedPerfettoProtoLogImpl extends PerfettoProtoLogImpl {
 
     @VisibleForTesting
     public ProcessedPerfettoProtoLogImpl(
+            @NonNull ProtoLogDataSource datasource,
             @NonNull String viewerConfigFilePath,
             @NonNull ViewerConfigInputStreamProvider viewerConfigInputStreamProvider,
             @NonNull ProtoLogViewerConfigReader viewerConfigReader,
-            @NonNull Runnable cacheUpdater,
+            @NonNull ProtoLogCacheUpdater cacheUpdater,
             @NonNull IProtoLogGroup[] groups,
-            @NonNull ProtoLogDataSourceBuilder dataSourceBuilder,
             @Nullable IProtoLogConfigurationService configurationService)
             throws ServiceManager.ServiceNotFoundException {
-        super(cacheUpdater, groups, dataSourceBuilder, configurationService);
+        super(datasource, cacheUpdater, groups, configurationService);
 
         this.mViewerConfigFilePath = viewerConfigFilePath;
 
diff --git a/core/java/com/android/internal/protolog/ProtoLog.java b/core/java/com/android/internal/protolog/ProtoLog.java
index d117e93d7de7..c81af959f36c 100644
--- a/core/java/com/android/internal/protolog/ProtoLog.java
+++ b/core/java/com/android/internal/protolog/ProtoLog.java
@@ -17,6 +17,9 @@
 package com.android.internal.protolog;
 
 import android.os.ServiceManager;
+import android.tracing.perfetto.DataSourceParams;
+import android.tracing.perfetto.InitArguments;
+import android.tracing.perfetto.Producer;
 
 import com.android.internal.protolog.common.IProtoLog;
 import com.android.internal.protolog.common.IProtoLogGroup;
@@ -54,6 +57,8 @@ public class ProtoLog {
 
     private static IProtoLog sProtoLogInstance;
 
+    private static ProtoLogDataSource sDataSource;
+
     private static final Object sInitLock = new Object();
 
     /**
@@ -69,26 +74,45 @@ public class ProtoLog {
         // files to extract out the log strings. Otherwise, the trace calls are replaced with calls
         // directly to the generated tracing implementations.
         if (android.tracing.Flags.perfettoProtologTracing()) {
-            synchronized (sInitLock) {
-                final var allGroups = new HashSet<>(Arrays.stream(groups).toList());
-                if (sProtoLogInstance != null) {
-                    // The ProtoLog instance has already been initialized in this process
-                    final var alreadyRegisteredGroups = sProtoLogInstance.getRegisteredGroups();
-                    allGroups.addAll(alreadyRegisteredGroups);
-                }
-
-                try {
-                    sProtoLogInstance = new UnprocessedPerfettoProtoLogImpl(
-                            allGroups.toArray(new IProtoLogGroup[0]));
-                } catch (ServiceManager.ServiceNotFoundException e) {
-                    throw new RuntimeException(e);
-                }
-            }
+            initializePerfettoProtoLog(groups);
         } else {
             sProtoLogInstance = new LogcatOnlyProtoLogImpl();
         }
     }
 
+    private static void initializePerfettoProtoLog(IProtoLogGroup... groups) {
+        var datasource = getSharedSingleInstanceDataSource();
+
+        synchronized (sInitLock) {
+            final var allGroups = new HashSet<>(Arrays.stream(groups).toList());
+            final var previousProtoLogImpl = sProtoLogInstance;
+            if (previousProtoLogImpl != null) {
+                // The ProtoLog instance has already been initialized in this process
+                final var alreadyRegisteredGroups = previousProtoLogImpl.getRegisteredGroups();
+                allGroups.addAll(alreadyRegisteredGroups);
+            }
+
+            sProtoLogInstance = createAndEnableNewPerfettoProtoLogImpl(
+                    datasource, allGroups.toArray(new IProtoLogGroup[0]));
+            if (previousProtoLogImpl instanceof PerfettoProtoLogImpl) {
+                ((PerfettoProtoLogImpl) previousProtoLogImpl).disable();
+            }
+        }
+    }
+
+    private static PerfettoProtoLogImpl createAndEnableNewPerfettoProtoLogImpl(
+            ProtoLogDataSource datasource, IProtoLogGroup[] groups) {
+        try {
+            var unprocessedPerfettoProtoLogImpl =
+                    new UnprocessedPerfettoProtoLogImpl(datasource, groups);
+            unprocessedPerfettoProtoLogImpl.enable();
+
+            return unprocessedPerfettoProtoLogImpl;
+        } catch (ServiceManager.ServiceNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     /**
      * DEBUG level log.
      *
@@ -190,6 +214,32 @@ public class ProtoLog {
         return sProtoLogInstance;
     }
 
+    /**
+     * Gets or creates if it doesn't exist yet the protolog datasource to use in this process.
+     * We should re-use the same datasource to avoid registering the datasource multiple times in
+     * the same process, since there is no way to unregister the datasource after registration.
+     *
+     * @return The single ProtoLog datasource instance to be shared across all ProtoLog tracing
+     *         objects.
+     */
+    public static synchronized ProtoLogDataSource getSharedSingleInstanceDataSource() {
+        if (sDataSource == null) {
+            Producer.init(InitArguments.DEFAULTS);
+            sDataSource = new ProtoLogDataSource();
+            DataSourceParams params =
+                    new DataSourceParams.Builder()
+                            .setBufferExhaustedPolicy(
+                                    DataSourceParams
+                                            .PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP)
+                            .build();
+            // NOTE: Registering that datasource is an async operation, so there may be no data
+            // traced for some messages logged right after the construction of this class.
+            sDataSource.register(params);
+        }
+
+        return sDataSource;
+    }
+
     private static void logStringMessage(LogLevel logLevel, IProtoLogGroup group,
             String stringMessage, Object... args) {
         if (sProtoLogInstance == null) {
diff --git a/core/java/com/android/internal/protolog/ProtoLogCacheUpdater.java b/core/java/com/android/internal/protolog/ProtoLogCacheUpdater.java
new file mode 100644
index 000000000000..4f8655c6377b
--- /dev/null
+++ b/core/java/com/android/internal/protolog/ProtoLogCacheUpdater.java
@@ -0,0 +1,28 @@
+/*
+ * 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 com.android.internal.protolog.common.IProtoLog;
+
+public interface ProtoLogCacheUpdater {
+    /**
+     * Update the cache based on the latest state of the active tracing configurations.
+     *
+     * @param protoLogInstance the instance to use to query the latest state of tracing.
+     */
+    void update(IProtoLog protoLogInstance);
+}
diff --git a/core/java/com/android/internal/protolog/ProtoLogConfigurationServiceImpl.java b/core/java/com/android/internal/protolog/ProtoLogConfigurationServiceImpl.java
index e9a8770deb73..23f7c2a3d987 100644
--- a/core/java/com/android/internal/protolog/ProtoLogConfigurationServiceImpl.java
+++ b/core/java/com/android/internal/protolog/ProtoLogConfigurationServiceImpl.java
@@ -34,9 +34,6 @@ import android.content.Context;
 import android.os.RemoteException;
 import android.os.ResultReceiver;
 import android.os.ShellCallback;
-import android.tracing.perfetto.DataSourceParams;
-import android.tracing.perfetto.InitArguments;
-import android.tracing.perfetto.Producer;
 import android.util.Log;
 import android.util.proto.ProtoInputStream;
 import android.util.proto.ProtoOutputStream;
@@ -110,39 +107,31 @@ public class ProtoLogConfigurationServiceImpl extends IProtoLogConfigurationServ
     private final ViewerConfigFileTracer mViewerConfigFileTracer;
 
     public ProtoLogConfigurationServiceImpl() {
-        this(ProtoLogDataSource::new, ProtoLogConfigurationServiceImpl::dumpViewerConfig);
+        this(ProtoLog.getSharedSingleInstanceDataSource(),
+                ProtoLogConfigurationServiceImpl::dumpViewerConfig);
     }
 
     @VisibleForTesting
-    public ProtoLogConfigurationServiceImpl(@NonNull ProtoLogDataSourceBuilder dataSourceBuilder) {
-        this(dataSourceBuilder, ProtoLogConfigurationServiceImpl::dumpViewerConfig);
+    public ProtoLogConfigurationServiceImpl(@NonNull ProtoLogDataSource datasource) {
+        this(datasource, ProtoLogConfigurationServiceImpl::dumpViewerConfig);
     }
 
     @VisibleForTesting
     public ProtoLogConfigurationServiceImpl(@NonNull ViewerConfigFileTracer tracer) {
-        this(ProtoLogDataSource::new, tracer);
+        this(ProtoLog.getSharedSingleInstanceDataSource(), tracer);
     }
 
     @VisibleForTesting
     public ProtoLogConfigurationServiceImpl(
-            @NonNull ProtoLogDataSourceBuilder dataSourceBuilder,
+            @NonNull ProtoLogDataSource datasource,
             @NonNull ViewerConfigFileTracer tracer) {
-        mDataSource = dataSourceBuilder.build(
-            this::onTracingInstanceStart,
-            this::onTracingInstanceFlush,
-            this::onTracingInstanceStop
-        );
-
-        // Initialize the Perfetto producer and register the Perfetto ProtoLog datasource to be
-        // receive the lifecycle callbacks of the datasource and write the viewer configs if and
-        // when required to the datasource.
-        Producer.init(InitArguments.DEFAULTS);
-        final var params = new DataSourceParams.Builder()
-                .setBufferExhaustedPolicy(DataSourceParams.PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP)
-                .build();
-        mDataSource.register(params);
-
         mViewerConfigFileTracer = tracer;
+
+        datasource.registerOnStartCallback(this::onTracingInstanceStart);
+        datasource.registerOnFlushCallback(this::onTracingInstanceFlush);
+        datasource.registerOnStopCallback(this::onTracingInstanceStop);
+
+        mDataSource = datasource;
     }
 
     public static class RegisterClientArgs extends IRegisterClientArgs.Stub {
diff --git a/core/java/com/android/internal/protolog/ProtoLogDataSource.java b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
index 0afb135ac6d9..ea452494b88d 100644
--- a/core/java/com/android/internal/protolog/ProtoLogDataSource.java
+++ b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
@@ -46,36 +46,30 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 
 public class ProtoLogDataSource extends DataSource<ProtoLogDataSource.Instance,
         ProtoLogDataSource.TlsState,
         ProtoLogDataSource.IncrementalState> {
     private static final String DATASOURCE_NAME = "android.protolog";
 
+    private final Map<Integer, ProtoLogConfig> mRunningInstances = new TreeMap<>();
+
     @NonNull
-    private final Instance.TracingInstanceStartCallback mOnStart;
+    private final Set<Instance.TracingInstanceStartCallback> mOnStartCallbacks = new HashSet<>();
     @NonNull
-    private final Runnable mOnFlush;
+    private final Set<Runnable> mOnFlushCallbacks = new HashSet<>();
     @NonNull
-    private final Instance.TracingInstanceStopCallback mOnStop;
+    private final Set<Instance.TracingInstanceStopCallback> mOnStopCallbacks = new HashSet<>();
 
-    public ProtoLogDataSource(
-            @NonNull Instance.TracingInstanceStartCallback onStart,
-            @NonNull Runnable onFlush,
-            @NonNull Instance.TracingInstanceStopCallback onStop) {
-        this(onStart, onFlush, onStop, DATASOURCE_NAME);
+    public ProtoLogDataSource() {
+        this(DATASOURCE_NAME);
     }
 
     @VisibleForTesting
     public ProtoLogDataSource(
-            @NonNull Instance.TracingInstanceStartCallback onStart,
-            @NonNull Runnable onFlush,
-            @NonNull Instance.TracingInstanceStopCallback onStop,
             @NonNull String dataSourceName) {
         super(dataSourceName);
-        this.mOnStart = onStart;
-        this.mOnFlush = onFlush;
-        this.mOnStop = onStop;
     }
 
     @Override
@@ -106,7 +100,8 @@ public class ProtoLogDataSource extends DataSource<ProtoLogDataSource.Instance,
         }
 
         return new Instance(
-                this, instanceIndex, config, mOnStart, mOnFlush, mOnStop);
+                this, instanceIndex, config, this::executeOnStartCallbacks,
+                this::executeOnFlushCallbacks, this::executeOnStopCallbacks);
     }
 
     @Override
@@ -128,6 +123,84 @@ public class ProtoLogDataSource extends DataSource<ProtoLogDataSource.Instance,
         return new IncrementalState();
     }
 
+    /**
+     * Register an onStart callback that will be called when a tracing instance is started.
+     * The callback will be called immediately for all currently running tracing instances on
+     * registration.
+     * @param onStartCallback The callback to call on starting a tracing instance.
+     */
+    public synchronized void registerOnStartCallback(
+            Instance.TracingInstanceStartCallback onStartCallback) {
+        mOnStartCallbacks.add(onStartCallback);
+
+        for (var instanceIndex : mRunningInstances.keySet()) {
+            var config = mRunningInstances.get(instanceIndex);
+            onStartCallback.run(instanceIndex, config);
+        }
+    }
+
+    /**
+     * Register an onFlush callback that will be called when a tracing instance is about to flush.
+     * @param onFlushCallback The callback to call on flushing a tracing instance
+     */
+    public void registerOnFlushCallback(Runnable onFlushCallback) {
+        mOnFlushCallbacks.add(onFlushCallback);
+    }
+
+    /**
+     * Register an onStop callback that will be called when a tracing instance is being stopped.
+     * @param onStopCallback The callback to call on stopping a tracing instance.
+     */
+    public void registerOnStopCallback(Instance.TracingInstanceStopCallback onStopCallback) {
+        mOnStopCallbacks.add(onStopCallback);
+    }
+
+    /**
+     * Unregister an onStart callback.
+     * @param onStartCallback The callback object to unregister.
+     */
+    public void unregisterOnStartCallback(Instance.TracingInstanceStartCallback onStartCallback) {
+        mOnStartCallbacks.add(onStartCallback);
+    }
+
+    /**
+     * Unregister an onFlush callback.
+     * @param onFlushCallback The callback object to unregister.
+     */
+    public void unregisterOnFlushCallback(Runnable onFlushCallback) {
+        mOnFlushCallbacks.add(onFlushCallback);
+    }
+
+    /**
+     * Unregister an onStop callback.
+     * @param onStopCallback The callback object to unregister.
+     */
+    public void unregisterOnStopCallback(Instance.TracingInstanceStopCallback onStopCallback) {
+        mOnStopCallbacks.add(onStopCallback);
+    }
+
+    private synchronized void executeOnStartCallbacks(int instanceIdx, ProtoLogConfig config) {
+        mRunningInstances.put(instanceIdx, config);
+
+        for (var onStart : mOnStartCallbacks) {
+            onStart.run(instanceIdx, config);
+        }
+    }
+
+    private void executeOnFlushCallbacks() {
+        for (var onFlush : mOnFlushCallbacks) {
+            onFlush.run();
+        }
+    }
+
+    private synchronized void executeOnStopCallbacks(int instanceIdx, ProtoLogConfig config) {
+        mRunningInstances.remove(instanceIdx, config);
+
+        for (var onStop : mOnStopCallbacks) {
+            onStop.run(instanceIdx, config);
+        }
+    }
+
     public static class TlsState {
         @NonNull
         private final ProtoLogConfig mConfig;
diff --git a/core/java/com/android/internal/protolog/ProtoLogImpl.java b/core/java/com/android/internal/protolog/ProtoLogImpl.java
index 3378d08e7761..c2d4b21181b7 100644
--- a/core/java/com/android/internal/protolog/ProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/ProtoLogImpl.java
@@ -56,7 +56,7 @@ public class ProtoLogImpl {
     private static TreeMap<String, IProtoLogGroup> sLogGroups;
 
     @ProtoLogToolInjected(CACHE_UPDATER)
-    private static Runnable sCacheUpdater;
+    private static ProtoLogCacheUpdater sCacheUpdater;
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
     public static void d(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
@@ -93,7 +93,12 @@ public class ProtoLogImpl {
      * and log level.
      */
     public static boolean isEnabled(IProtoLogGroup group, LogLevel level) {
-        return getSingleInstance().isEnabled(group, level);
+        return isEnabled(getSingleInstance(), group, level);
+    }
+
+    private static boolean isEnabled(
+            IProtoLog protoLogInstance, IProtoLogGroup group, LogLevel level) {
+        return protoLogInstance.isEnabled(group, level);
     }
 
     /**
@@ -106,36 +111,37 @@ public class ProtoLogImpl {
 
             final var groups = sLogGroups.values().toArray(new IProtoLogGroup[0]);
             if (android.tracing.Flags.perfettoProtologTracing()) {
-                sServiceInstance = createProtoLogImpl(groups);
+                var viewerConfigFile = new File(sViewerConfigPath);
+                if (!viewerConfigFile.exists()) {
+                    // TODO(b/353530422): Remove - temporary fix to unblock b/352290057
+                    // In robolectric tests the viewer config file isn't current available, so we
+                    // cannot use the ProcessedPerfettoProtoLogImpl.
+                    Log.e(LOG_TAG, "Failed to find viewer config file " + sViewerConfigPath
+                            + " when setting up " + ProtoLogImpl.class.getSimpleName() + ". "
+                            + "ProtoLog will not work here!");
+
+                    sServiceInstance = new NoViewerConfigProtoLogImpl();
+                } else {
+                    var datasource = ProtoLog.getSharedSingleInstanceDataSource();
+                    try {
+                        var processedProtoLogImpl =
+                                new ProcessedPerfettoProtoLogImpl(datasource, sViewerConfigPath,
+                                        sCacheUpdater, groups);
+                        sServiceInstance = processedProtoLogImpl;
+                        processedProtoLogImpl.enable();
+                    } catch (ServiceManager.ServiceNotFoundException e) {
+                        throw new RuntimeException(e);
+                    }
+                }
             } else {
                 sServiceInstance = createLegacyProtoLogImpl(groups);
             }
 
-            sCacheUpdater.run();
+            sCacheUpdater.update(sServiceInstance);
         }
         return sServiceInstance;
     }
 
-    private static IProtoLog createProtoLogImpl(IProtoLogGroup[] groups) {
-        try {
-            File f = new File(sViewerConfigPath);
-            if (!f.exists()) {
-                // TODO(b/353530422): Remove - temporary fix to unblock b/352290057
-                // In robolectric tests the viewer config file isn't current available, so we cannot
-                // use the ProcessedPerfettoProtoLogImpl.
-                Log.e(LOG_TAG, "Failed to find viewer config file " + sViewerConfigPath
-                        + " when setting up " + ProtoLogImpl.class.getSimpleName() + ". "
-                        + "ProtoLog will not work here!");
-
-                return new NoViewerConfigProtoLogImpl();
-            } else {
-                return new ProcessedPerfettoProtoLogImpl(sViewerConfigPath, sCacheUpdater, groups);
-            }
-        } catch (ServiceManager.ServiceNotFoundException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     private static LegacyProtoLogImpl createLegacyProtoLogImpl(IProtoLogGroup[] groups) {
         var protologImpl = new LegacyProtoLogImpl(
                 sLegacyOutputFilePath, sLegacyViewerConfigPath, sCacheUpdater);
diff --git a/core/java/com/android/internal/protolog/UnprocessedPerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/UnprocessedPerfettoProtoLogImpl.java
index f3fe58070fa9..ebb07a04270d 100644
--- a/core/java/com/android/internal/protolog/UnprocessedPerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/UnprocessedPerfettoProtoLogImpl.java
@@ -23,14 +23,10 @@ import com.android.internal.protolog.ProtoLogConfigurationServiceImpl.RegisterCl
 import com.android.internal.protolog.common.IProtoLogGroup;
 
 public class UnprocessedPerfettoProtoLogImpl extends PerfettoProtoLogImpl {
-    public UnprocessedPerfettoProtoLogImpl(@NonNull IProtoLogGroup[] groups)
+    public UnprocessedPerfettoProtoLogImpl(
+            @NonNull ProtoLogDataSource dataSource, @NonNull IProtoLogGroup[] groups)
             throws ServiceManager.ServiceNotFoundException {
-        this(() -> {}, groups);
-    }
-
-    public UnprocessedPerfettoProtoLogImpl(@NonNull Runnable cacheUpdater,
-            @NonNull IProtoLogGroup[] groups) throws ServiceManager.ServiceNotFoundException {
-        super(cacheUpdater, groups);
+        super(dataSource, (instance) -> {}, groups);
         readyToLogToLogcat();
     }
 
diff --git a/tests/Tracing/src/com/android/internal/protolog/LegacyProtoLogImplTest.java b/tests/Tracing/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
index 8913e8c1996e..05308464cb9b 100644
--- a/tests/Tracing/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
+++ b/tests/Tracing/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
@@ -89,7 +89,7 @@ public class LegacyProtoLogImplTest {
         //noinspection ResultOfMethodCallIgnored
         mFile.delete();
         mProtoLog = new LegacyProtoLogImpl(mFile, mViewerConfigFilename,
-                1024 * 1024, mReader, 1024, () -> {});
+                1024 * 1024, mReader, 1024, (instance) -> {});
     }
 
     @After
diff --git a/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java b/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java
index 44641f7a1e12..ed256e72b415 100644
--- a/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java
+++ b/tests/Tracing/src/com/android/internal/protolog/ProcessedPerfettoProtoLogImplTest.java
@@ -41,6 +41,7 @@ import android.tools.traces.io.ResultWriter;
 import android.tools.traces.monitors.PerfettoTraceMonitor;
 import android.tools.traces.protolog.ProtoLogTrace;
 import android.tracing.perfetto.DataSource;
+import android.tracing.perfetto.DataSourceParams;
 
 import androidx.test.platform.app.InstrumentationRegistry;
 
@@ -95,14 +96,14 @@ public class ProcessedPerfettoProtoLogImplTest {
     );
 
     private static ProtoLogConfigurationService sProtoLogConfigurationService;
+    private static ProtoLogDataSource sTestDataSource;
     private static PerfettoProtoLogImpl sProtoLog;
     private static Protolog.ProtoLogViewerConfig.Builder sViewerConfigBuilder;
-    private static Runnable sCacheUpdater;
+    private static ProtoLogCacheUpdater sCacheUpdater;
 
     private static ProtoLogViewerConfigReader sReader;
 
-    public ProcessedPerfettoProtoLogImplTest() throws IOException {
-    }
+    public ProcessedPerfettoProtoLogImplTest() throws IOException { }
 
     @BeforeClass
     public static void setUp() throws Exception {
@@ -155,12 +156,18 @@ public class ProcessedPerfettoProtoLogImplTest {
                 .thenAnswer(it -> new AutoClosableProtoInputStream(
                         sViewerConfigBuilder.build().toByteArray()));
 
-        sCacheUpdater = () -> {};
+        sCacheUpdater = (instance) -> {};
         sReader = Mockito.spy(new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider));
+        sTestDataSource = new ProtoLogDataSource(TEST_PROTOLOG_DATASOURCE_NAME);
+        DataSourceParams params =
+                new DataSourceParams.Builder()
+                        .setBufferExhaustedPolicy(
+                                DataSourceParams
+                                        .PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP)
+                        .build();
+        sTestDataSource.register(params);
+        busyWaitForDataSourceRegistration(TEST_PROTOLOG_DATASOURCE_NAME);
 
-        final ProtoLogDataSourceBuilder dataSourceBuilder =
-                (onStart, onFlush, onStop) -> new ProtoLogDataSource(
-                        onStart, onFlush, onStop, TEST_PROTOLOG_DATASOURCE_NAME);
         final ViewerConfigFileTracer tracer = (dataSource, viewerConfigFilePath) -> {
             Utils.dumpViewerConfig(dataSource, () -> {
                 if (!viewerConfigFilePath.equals(MOCK_VIEWER_CONFIG_FILE)) {
@@ -171,14 +178,13 @@ public class ProcessedPerfettoProtoLogImplTest {
             });
         };
         sProtoLogConfigurationService =
-                new ProtoLogConfigurationServiceImpl(dataSourceBuilder, tracer);
+                new ProtoLogConfigurationServiceImpl(sTestDataSource, tracer);
 
-        sProtoLog = new ProcessedPerfettoProtoLogImpl(
+        sProtoLog = new ProcessedPerfettoProtoLogImpl(sTestDataSource,
                 MOCK_VIEWER_CONFIG_FILE, viewerConfigInputStreamProvider, sReader,
-                () -> sCacheUpdater.run(), TestProtoLogGroup.values(), dataSourceBuilder,
+                (instance) -> sCacheUpdater.update(instance), TestProtoLogGroup.values(),
                 sProtoLogConfigurationService);
-
-        busyWaitForDataSourceRegistration(TEST_PROTOLOG_DATASOURCE_NAME);
+        sProtoLog.enable();
     }
 
     @Before
@@ -606,7 +612,7 @@ public class ProcessedPerfettoProtoLogImplTest {
     @Test
     public void cacheIsUpdatedWhenTracesStartAndStop() {
         final AtomicInteger cacheUpdateCallCount = new AtomicInteger(0);
-        sCacheUpdater = cacheUpdateCallCount::incrementAndGet;
+        sCacheUpdater = (instance) -> cacheUpdateCallCount.incrementAndGet();
 
         PerfettoTraceMonitor traceMonitor1 = PerfettoTraceMonitor.newBuilder()
                 .enableProtoLog(true,
diff --git a/tests/Tracing/src/com/android/internal/protolog/ProtologDataSourceTest.java b/tests/Tracing/src/com/android/internal/protolog/ProtologDataSourceTest.java
index ce519b7a1576..49249333b72b 100644
--- a/tests/Tracing/src/com/android/internal/protolog/ProtologDataSourceTest.java
+++ b/tests/Tracing/src/com/android/internal/protolog/ProtologDataSourceTest.java
@@ -67,9 +67,6 @@ public class ProtologDataSourceTest {
 
     @Test
     public void allEnabledTraceMode() {
-        final ProtoLogDataSource ds =
-                new ProtoLogDataSource((idx, c) -> {}, () -> {}, (idx, c) -> {});
-
         final ProtoLogDataSource.TlsState tlsState = createTlsState(
                 DataSourceConfigOuterClass.DataSourceConfig.newBuilder().setProtologConfig(
                         ProtologConfig.ProtoLogConfig.newBuilder()
@@ -154,8 +151,7 @@ public class ProtologDataSourceTest {
 
     private ProtoLogDataSource.TlsState createTlsState(
             DataSourceConfigOuterClass.DataSourceConfig config) {
-        final ProtoLogDataSource ds =
-                Mockito.spy(new ProtoLogDataSource((idx, c) -> {}, () -> {}, (idx, c) -> {}));
+        final ProtoLogDataSource ds = Mockito.spy(new ProtoLogDataSource());
 
         ProtoInputStream configStream = new ProtoInputStream(config.toByteArray());
         final ProtoLogDataSource.Instance dsInstance = Mockito.spy(
diff --git a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
index 108942ee754a..9222ff4bf52c 100644
--- a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
@@ -16,6 +16,7 @@
 
 package com.android.protolog.tool
 
+import com.android.internal.protolog.common.IProtoLog
 import com.android.internal.protolog.common.LogLevel
 import com.android.internal.protolog.common.ProtoLogToolInjected
 import com.android.protolog.tool.CommandOptions.Companion.USAGE
@@ -319,6 +320,7 @@ object ProtoLogTool {
                             MethodCallExpr()
                                 .setName("isEnabled")
                                 .setArguments(NodeList(
+                                    NameExpr("protoLogInstance"),
                                     FieldAccessExpr()
                                         .setScope(NameExpr(protoLogGroupsClassName))
                                         .setName(group.value.name),
@@ -332,6 +334,7 @@ object ProtoLogTool {
         }
 
         cacheClass.addMethod("update").setPrivate(true).setStatic(true)
+            .addParameter(IProtoLog::class.java, "protoLogInstance")
             .setBody(updateBlockStmt)
 
         classDeclaration.addMember(cacheClass)