diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 8f7d703..fbec1f1 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -768,7 +768,8 @@
         return -1;
     }
 
-    private synchronized void onTracingInstanceStart(ProtoLogDataSource.ProtoLogConfig config) {
+    private synchronized void onTracingInstanceStart(
+            int instanceIdx, ProtoLogDataSource.ProtoLogConfig config) {
         final LogLevel defaultLogFrom = config.getDefaultGroupConfig().logFrom;
         for (int i = defaultLogFrom.ordinal(); i < LogLevel.values().length; i++) {
             mDefaultLogLevelCounts[i]++;
@@ -801,7 +802,8 @@
         this.mTracingInstances.incrementAndGet();
     }
 
-    private synchronized void onTracingInstanceStop(ProtoLogDataSource.ProtoLogConfig config) {
+    private synchronized void onTracingInstanceStop(
+            int instanceIdx, ProtoLogDataSource.ProtoLogConfig config) {
         this.mTracingInstances.decrementAndGet();
 
         final LogLevel defaultLogFrom = config.getDefaultGroupConfig().logFrom;
diff --git a/core/java/com/android/internal/protolog/ProtoLogDataSource.java b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
index 84f3237..6dc6585 100644
--- a/core/java/com/android/internal/protolog/ProtoLogDataSource.java
+++ b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
@@ -49,12 +49,12 @@
         ProtoLogDataSource.TlsState,
         ProtoLogDataSource.IncrementalState> {
 
-    private final Consumer<ProtoLogConfig> mOnStart;
+    private final Instance.TracingInstanceStartCallback mOnStart;
     private final Runnable mOnFlush;
-    private final Consumer<ProtoLogConfig> mOnStop;
+    private final Instance.TracingInstanceStopCallback mOnStop;
 
-    public ProtoLogDataSource(Consumer<ProtoLogConfig> onStart, Runnable onFlush,
-            Consumer<ProtoLogConfig> onStop) {
+    public ProtoLogDataSource(Instance.TracingInstanceStartCallback onStart, Runnable onFlush,
+            Instance.TracingInstanceStopCallback onStop) {
         super("android.protolog");
         this.mOnStart = onStart;
         this.mOnFlush = onFlush;
@@ -267,20 +267,30 @@
 
     public static class Instance extends DataSourceInstance {
 
-        private final Consumer<ProtoLogConfig> mOnStart;
+        public interface TracingInstanceStartCallback {
+            void run(int instanceIdx, ProtoLogConfig config);
+        }
+
+        public interface TracingInstanceStopCallback {
+            void run(int instanceIdx, ProtoLogConfig config);
+        }
+
+        private final TracingInstanceStartCallback mOnStart;
         private final Runnable mOnFlush;
-        private final Consumer<ProtoLogConfig> mOnStop;
+        private final TracingInstanceStopCallback mOnStop;
         private final ProtoLogConfig mConfig;
+        private final int mInstanceIndex;
 
         public Instance(
                 DataSource<Instance, TlsState, IncrementalState> dataSource,
                 int instanceIdx,
                 ProtoLogConfig config,
-                Consumer<ProtoLogConfig> onStart,
+                TracingInstanceStartCallback onStart,
                 Runnable onFlush,
-                Consumer<ProtoLogConfig> onStop
+                TracingInstanceStopCallback onStop
         ) {
             super(dataSource, instanceIdx);
+            this.mInstanceIndex = instanceIdx;
             this.mOnStart = onStart;
             this.mOnFlush = onFlush;
             this.mOnStop = onStop;
@@ -289,7 +299,7 @@
 
         @Override
         public void onStart(StartCallbackArguments args) {
-            this.mOnStart.accept(this.mConfig);
+            this.mOnStart.run(this.mInstanceIndex, this.mConfig);
         }
 
         @Override
@@ -299,7 +309,7 @@
 
         @Override
         public void onStop(StopCallbackArguments args) {
-            this.mOnStop.accept(this.mConfig);
+            this.mOnStop.run(this.mInstanceIndex, this.mConfig);
         }
     }
 }
diff --git a/tests/Internal/src/com/android/internal/protolog/ProtologDataSourceTest.java b/tests/Internal/src/com/android/internal/protolog/ProtologDataSourceTest.java
index be9fb1b..9a062e3 100644
--- a/tests/Internal/src/com/android/internal/protolog/ProtologDataSourceTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/ProtologDataSourceTest.java
@@ -67,7 +67,7 @@
 
     @Test
     public void allEnabledTraceMode() {
-        final ProtoLogDataSource ds = new ProtoLogDataSource((c) -> {}, () -> {}, (c) -> {});
+        final ProtoLogDataSource ds = new ProtoLogDataSource((idx, c) -> {}, () -> {}, (idx, c) -> {});
 
         final ProtoLogDataSource.TlsState tlsState = createTlsState(
                 DataSourceConfigOuterClass.DataSourceConfig.newBuilder().setProtologConfig(
@@ -154,7 +154,7 @@
     private ProtoLogDataSource.TlsState createTlsState(
             DataSourceConfigOuterClass.DataSourceConfig config) {
         final ProtoLogDataSource ds =
-                Mockito.spy(new ProtoLogDataSource((c) -> {}, () -> {}, (c) -> {}));
+                Mockito.spy(new ProtoLogDataSource((idx, c) -> {}, () -> {}, (idx, c) -> {}));
 
         ProtoInputStream configStream = new ProtoInputStream(config.toByteArray());
         final ProtoLogDataSource.Instance dsInstance = Mockito.spy(