Update uses of ProtoLog

Flag: ACONFIG android.tracing.Flags.perfettoProtolog DEVELOPMENT
Bug: 276432490
Test: atest FrameworksServicesTests
Change-Id: I92926a44baaeea3971e46c0901898c90b6c98adb
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ProtoLogController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ProtoLogController.java
index 88525aa..93893e3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ProtoLogController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ProtoLogController.java
@@ -16,7 +16,10 @@
 
 package com.android.wm.shell;
 
-import com.android.wm.shell.protolog.ShellProtoLogImpl;
+import com.android.internal.protolog.LegacyProtoLogImpl;
+import com.android.internal.protolog.common.ILogger;
+import com.android.internal.protolog.common.IProtoLog;
+import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellInit;
 
@@ -24,19 +27,19 @@
 import java.util.Arrays;
 
 /**
- * Controls the {@link ShellProtoLogImpl} in WMShell via adb shell commands.
+ * Controls the {@link ProtoLog} in WMShell via adb shell commands.
  *
  * Use with {@code adb shell dumpsys activity service SystemUIService WMShell protolog ...}.
  */
 public class ProtoLogController implements ShellCommandHandler.ShellCommandActionHandler {
     private final ShellCommandHandler mShellCommandHandler;
-    private final ShellProtoLogImpl mShellProtoLog;
+    private final IProtoLog mShellProtoLog;
 
     public ProtoLogController(ShellInit shellInit,
             ShellCommandHandler shellCommandHandler) {
         shellInit.addInitCallback(this::onInit, this);
         mShellCommandHandler = shellCommandHandler;
-        mShellProtoLog = ShellProtoLogImpl.getSingleInstance();
+        mShellProtoLog = ProtoLog.getSingleInstance();
     }
 
     void onInit() {
@@ -45,22 +48,35 @@
 
     @Override
     public boolean onShellCommand(String[] args, PrintWriter pw) {
+        final ILogger logger = pw::println;
         switch (args[0]) {
             case "status": {
-                pw.println(mShellProtoLog.getStatus());
+                if (android.tracing.Flags.perfettoProtolog()) {
+                    pw.println("(Deprecated) legacy command. Use Perfetto commands instead.");
+                    return false;
+                }
+                ((LegacyProtoLogImpl) mShellProtoLog).getStatus();
                 return true;
             }
             case "start": {
-                mShellProtoLog.startProtoLog(pw);
+                if (android.tracing.Flags.perfettoProtolog()) {
+                    pw.println("(Deprecated) legacy command. Use Perfetto commands instead.");
+                    return false;
+                }
+                ((LegacyProtoLogImpl) mShellProtoLog).startProtoLog(pw);
                 return true;
             }
             case "stop": {
-                mShellProtoLog.stopProtoLog(pw, true /* writeToFile */);
+                if (android.tracing.Flags.perfettoProtolog()) {
+                    pw.println("(Deprecated) legacy command. Use Perfetto commands instead.");
+                    return false;
+                }
+                ((LegacyProtoLogImpl) mShellProtoLog).stopProtoLog(pw, true);
                 return true;
             }
             case "enable-text": {
                 String[] groups = Arrays.copyOfRange(args, 1, args.length);
-                int result = mShellProtoLog.startTextLogging(groups, pw);
+                int result = mShellProtoLog.startLoggingToLogcat(groups, logger);
                 if (result == 0) {
                     pw.println("Starting logging on groups: " + Arrays.toString(groups));
                     return true;
@@ -69,7 +85,7 @@
             }
             case "disable-text": {
                 String[] groups = Arrays.copyOfRange(args, 1, args.length);
-                int result = mShellProtoLog.stopTextLogging(groups, pw);
+                int result = mShellProtoLog.stopLoggingToLogcat(groups, logger);
                 if (result == 0) {
                     pw.println("Stopping logging on groups: " + Arrays.toString(groups));
                     return true;
@@ -78,19 +94,23 @@
             }
             case "enable": {
                 String[] groups = Arrays.copyOfRange(args, 1, args.length);
-                return mShellProtoLog.startTextLogging(groups, pw) == 0;
+                return mShellProtoLog.startLoggingToLogcat(groups, logger) == 0;
             }
             case "disable": {
                 String[] groups = Arrays.copyOfRange(args, 1, args.length);
-                return mShellProtoLog.stopTextLogging(groups, pw) == 0;
+                return mShellProtoLog.stopLoggingToLogcat(groups, logger) == 0;
             }
             case "save-for-bugreport": {
+                if (android.tracing.Flags.perfettoProtolog()) {
+                    pw.println("(Deprecated) legacy command");
+                    return false;
+                }
                 if (!mShellProtoLog.isProtoEnabled()) {
                     pw.println("Logging to proto is not enabled for WMShell.");
                     return false;
                 }
-                mShellProtoLog.stopProtoLog(pw, true /* writeToFile */);
-                mShellProtoLog.startProtoLog(pw);
+                ((LegacyProtoLogImpl) mShellProtoLog).stopProtoLog(pw, true /* writeToFile */);
+                ((LegacyProtoLogImpl) mShellProtoLog).startProtoLog(pw);
                 return true;
             }
             default: {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt
index 9b48a54..7a50814 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt
@@ -18,7 +18,7 @@
 
 import android.util.Log
 import com.android.internal.protolog.common.IProtoLogGroup
-import com.android.wm.shell.protolog.ShellProtoLogImpl
+import com.android.internal.protolog.common.ProtoLog
 
 /**
  * Log messages using an API similar to [com.android.internal.protolog.common.ProtoLog]. Useful for
@@ -31,42 +31,42 @@
     companion object {
         /** @see [com.android.internal.protolog.common.ProtoLog.d] */
         fun d(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (ShellProtoLogImpl.isEnabled(group)) {
+            if (ProtoLog.isEnabled(group)) {
                 Log.d(group.tag, String.format(messageString, *args))
             }
         }
 
         /** @see [com.android.internal.protolog.common.ProtoLog.v] */
         fun v(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (ShellProtoLogImpl.isEnabled(group)) {
+            if (ProtoLog.isEnabled(group)) {
                 Log.v(group.tag, String.format(messageString, *args))
             }
         }
 
         /** @see [com.android.internal.protolog.common.ProtoLog.i] */
         fun i(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (ShellProtoLogImpl.isEnabled(group)) {
+            if (ProtoLog.isEnabled(group)) {
                 Log.i(group.tag, String.format(messageString, *args))
             }
         }
 
         /** @see [com.android.internal.protolog.common.ProtoLog.w] */
         fun w(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (ShellProtoLogImpl.isEnabled(group)) {
+            if (ProtoLog.isEnabled(group)) {
                 Log.w(group.tag, String.format(messageString, *args))
             }
         }
 
         /** @see [com.android.internal.protolog.common.ProtoLog.e] */
         fun e(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (ShellProtoLogImpl.isEnabled(group)) {
+            if (ProtoLog.isEnabled(group)) {
                 Log.e(group.tag, String.format(messageString, *args))
             }
         }
 
         /** @see [com.android.internal.protolog.common.ProtoLog.wtf] */
         fun wtf(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (ShellProtoLogImpl.isEnabled(group)) {
+            if (ProtoLog.isEnabled(group)) {
                 Log.wtf(group.tag, String.format(messageString, *args))
             }
         }
diff --git a/services/core/java/com/android/server/wm/AppTransition.java b/services/core/java/com/android/server/wm/AppTransition.java
index 939cf1a..1a63f14 100644
--- a/services/core/java/com/android/server/wm/AppTransition.java
+++ b/services/core/java/com/android/server/wm/AppTransition.java
@@ -137,7 +137,6 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.policy.TransitionAnimation;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.DumpUtils.Dump;
 import com.android.internal.util.function.pooled.PooledLambda;
@@ -248,7 +247,7 @@
         mHandler = new Handler(service.mH.getLooper());
         mDisplayContent = displayContent;
         mTransitionAnimation = new TransitionAnimation(
-                context, ProtoLogImpl.isEnabled(WM_DEBUG_ANIM), TAG);
+                context, ProtoLog.isEnabled(WM_DEBUG_ANIM), TAG);
 
         mGridLayoutRecentsEnabled = SystemProperties.getBoolean("ro.recents.grid", false);
 
diff --git a/services/core/java/com/android/server/wm/RemoteAnimationController.java b/services/core/java/com/android/server/wm/RemoteAnimationController.java
index a98b9f7..3ef6eeb 100644
--- a/services/core/java/com/android/server/wm/RemoteAnimationController.java
+++ b/services/core/java/com/android/server/wm/RemoteAnimationController.java
@@ -44,7 +44,6 @@
 import android.view.WindowManager;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.FastPrintWriter;
 import com.android.server.wm.SurfaceAnimator.AnimationType;
@@ -210,7 +209,7 @@
                 Slog.e(TAG, "Failed to start remote animation", e);
                 onAnimationFinished();
             }
-            if (ProtoLogImpl.isEnabled(WM_DEBUG_REMOTE_ANIMATIONS)) {
+            if (ProtoLog.isEnabled(WM_DEBUG_REMOTE_ANIMATIONS)) {
                 ProtoLog.d(WM_DEBUG_REMOTE_ANIMATIONS, "startAnimation(): Notify animation start:");
                 writeStartDebugStatement();
             }
diff --git a/services/core/java/com/android/server/wm/SurfaceAnimator.java b/services/core/java/com/android/server/wm/SurfaceAnimator.java
index c3de4d5..d67684c 100644
--- a/services/core/java/com/android/server/wm/SurfaceAnimator.java
+++ b/services/core/java/com/android/server/wm/SurfaceAnimator.java
@@ -32,7 +32,6 @@
 import android.view.SurfaceControl.Transaction;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 
 import java.io.PrintWriter;
@@ -193,7 +192,7 @@
             return;
         }
         mAnimation.startAnimation(mLeash, t, type, mInnerAnimationFinishedCallback);
-        if (ProtoLogImpl.isEnabled(WM_DEBUG_ANIM)) {
+        if (ProtoLog.isEnabled(WM_DEBUG_ANIM)) {
             StringWriter sw = new StringWriter();
             PrintWriter pw = new PrintWriter(sw);
             mAnimation.dump(pw, "");
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index 6949a87..d3344b0 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -60,7 +60,6 @@
 
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.ToBooleanFunction;
 import com.android.server.wallpaper.WallpaperCropper.WallpaperCropUtils;
@@ -359,7 +358,7 @@
         for (int i = mWallpaperTokens.size() - 1; i >= 0; i--) {
             final WallpaperWindowToken token = mWallpaperTokens.get(i);
             token.setVisibility(false);
-            if (ProtoLogImpl.isEnabled(WM_DEBUG_WALLPAPER) && token.isVisible()) {
+            if (ProtoLog.isEnabled(WM_DEBUG_WALLPAPER) && token.isVisible()) {
                 ProtoLog.d(WM_DEBUG_WALLPAPER,
                         "Hiding wallpaper %s from %s target=%s prev=%s callers=%s",
                         token, winGoingAway, mWallpaperTarget, mPrevWallpaperTarget,
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 61fde5e..fd0289e 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -113,7 +113,6 @@
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.graphics.ColorUtils;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.ToBooleanFunction;
 import com.android.server.wm.SurfaceAnimator.Animatable;
@@ -3410,7 +3409,7 @@
                 // ActivityOption#makeCustomAnimation or WindowManager#overridePendingTransition.
                 a.restrictDuration(MAX_APP_TRANSITION_DURATION);
             }
-            if (ProtoLogImpl.isEnabled(WM_DEBUG_ANIM)) {
+            if (ProtoLog.isEnabled(WM_DEBUG_ANIM)) {
                 ProtoLog.i(WM_DEBUG_ANIM, "Loaded animation %s for %s, duration: %d, stack=%s",
                         a, this, ((a != null) ? a.getDuration() : 0), Debug.getCallers(20));
             }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 61480d2..925d67c 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -327,8 +327,8 @@
 import com.android.internal.policy.IKeyguardLockedStateListener;
 import com.android.internal.policy.IShortcutService;
 import com.android.internal.policy.KeyInterceptionInfo;
+import com.android.internal.protolog.LegacyProtoLogImpl;
 import com.android.internal.protolog.ProtoLogGroup;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.DumpUtils;
 import com.android.internal.util.FastPrintWriter;
@@ -6701,7 +6701,11 @@
 
     private void dumpLogStatus(PrintWriter pw) {
         pw.println("WINDOW MANAGER LOGGING (dumpsys window logging)");
-        pw.println(ProtoLogImpl.getSingleInstance().getStatus());
+        if (android.tracing.Flags.perfettoProtolog()) {
+            pw.println("Deprecated legacy command. Use Perfetto commands instead.");
+            return;
+        }
+        ((LegacyProtoLogImpl) ProtoLog.getSingleInstance()).getStatus();
     }
 
     private void dumpSessionsLocked(PrintWriter pw) {
diff --git a/services/core/java/com/android/server/wm/WindowManagerShellCommand.java b/services/core/java/com/android/server/wm/WindowManagerShellCommand.java
index 8fad950..0b29f96 100644
--- a/services/core/java/com/android/server/wm/WindowManagerShellCommand.java
+++ b/services/core/java/com/android/server/wm/WindowManagerShellCommand.java
@@ -48,7 +48,9 @@
 import android.view.ViewDebug;
 
 import com.android.internal.os.ByteTransferPipe;
-import com.android.internal.protolog.ProtoLogImpl;
+import com.android.internal.protolog.LegacyProtoLogImpl;
+import com.android.internal.protolog.common.IProtoLog;
+import com.android.internal.protolog.common.ProtoLog;
 import com.android.server.IoThread;
 import com.android.server.wm.LetterboxConfiguration.LetterboxBackgroundType;
 import com.android.server.wm.LetterboxConfiguration.LetterboxHorizontalReachabilityPosition;
@@ -107,11 +109,19 @@
                     // trace files can be written.
                     return mInternal.mWindowTracing.onShellCommand(this);
                 case "logging":
-                    int result = ProtoLogImpl.getSingleInstance().onShellCommand(this);
-                    if (result != 0) {
-                        pw.println("Not handled, please use "
-                                + "`adb shell dumpsys activity service SystemUIService WMShell` "
-                                + "if you are looking for ProtoLog in WMShell");
+                    IProtoLog instance = ProtoLog.getSingleInstance();
+                    int result = 0;
+                    if (instance instanceof LegacyProtoLogImpl) {
+                        result = ((LegacyProtoLogImpl) instance).onShellCommand(this);
+                        if (result != 0) {
+                            pw.println("Not handled, please use "
+                                    + "`adb shell dumpsys activity service SystemUIService "
+                                    + "WMShell` if you are looking for ProtoLog in WMShell");
+                        }
+                    } else {
+                        result = -1;
+                        pw.println("Command not supported. "
+                                + "Only supported when using legacy ProtoLog.");
                     }
                     return result;
                 case "user-rotation":
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index b8f1d15..2cde0ce 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -245,7 +245,6 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.policy.KeyInterceptionInfo;
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.ToBooleanFunction;
@@ -4679,7 +4678,7 @@
     }
 
     void onExitAnimationDone() {
-        if (ProtoLogImpl.isEnabled(WM_DEBUG_ANIM)) {
+        if (ProtoLog.isEnabled(WM_DEBUG_ANIM)) {
             final AnimationAdapter animationAdapter = mSurfaceAnimator.getAnimation();
             StringWriter sw = new StringWriter();
             if (animationAdapter != null) {
diff --git a/services/core/java/com/android/server/wm/WindowStateAnimator.java b/services/core/java/com/android/server/wm/WindowStateAnimator.java
index 6428591..7f7c249 100644
--- a/services/core/java/com/android/server/wm/WindowStateAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowStateAnimator.java
@@ -60,7 +60,6 @@
 import android.view.animation.Animation;
 import android.view.animation.AnimationUtils;
 
-import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.server.policy.WindowManagerPolicy;
 
@@ -584,7 +583,7 @@
                             mWin.mAttrs, attr, TRANSIT_OLD_NONE);
                 }
             }
-            if (ProtoLogImpl.isEnabled(WM_DEBUG_ANIM)) {
+            if (ProtoLog.isEnabled(WM_DEBUG_ANIM)) {
                 ProtoLog.v(WM_DEBUG_ANIM, "applyAnimation: win=%s"
                         + " anim=%d attr=0x%x a=%s transit=%d type=%d isEntrance=%b Callers %s",
                         this, anim, attr, a, transit, mAttrType, isEntrance, Debug.getCallers(20));
diff --git a/services/core/java/com/android/server/wm/WindowTracing.java b/services/core/java/com/android/server/wm/WindowTracing.java
index 416d042..424d504 100644
--- a/services/core/java/com/android/server/wm/WindowTracing.java
+++ b/services/core/java/com/android/server/wm/WindowTracing.java
@@ -35,7 +35,9 @@
 import android.util.proto.ProtoOutputStream;
 import android.view.Choreographer;
 
-import com.android.internal.protolog.ProtoLogImpl;
+import com.android.internal.protolog.LegacyProtoLogImpl;
+import com.android.internal.protolog.common.IProtoLog;
+import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.TraceBuffer;
 
 import java.io.File;
@@ -77,6 +79,8 @@
     private volatile boolean mEnabledLockFree;
     private boolean mScheduled;
 
+    private final IProtoLog mProtoLog;
+
     static WindowTracing createDefaultAndStartLooper(WindowManagerService service,
             Choreographer choreographer) {
         File file = new File(TRACE_FILENAME);
@@ -96,6 +100,7 @@
         mTraceFile = file;
         mBuffer = new TraceBuffer(bufferCapacity);
         setLogLevel(WindowTraceLogLevel.TRIM, null /* pw */);
+        mProtoLog = ProtoLog.getSingleInstance();
     }
 
     void startTrace(@Nullable PrintWriter pw) {
@@ -104,7 +109,6 @@
             return;
         }
         synchronized (mEnabledLock) {
-            ProtoLogImpl.getSingleInstance().startProtoLog(pw);
             logAndPrintln(pw, "Start tracing to " + mTraceFile + ".");
             mBuffer.resetBuffer();
             mEnabled = mEnabledLockFree = true;
@@ -132,7 +136,6 @@
             writeTraceToFileLocked();
             logAndPrintln(pw, "Trace written to " + mTraceFile + ".");
         }
-        ProtoLogImpl.getSingleInstance().stopProtoLog(pw, true);
     }
 
     /**
@@ -152,11 +155,15 @@
             logAndPrintln(pw, "Stop tracing to " + mTraceFile + ". Waiting for traces to flush.");
             writeTraceToFileLocked();
             logAndPrintln(pw, "Trace written to " + mTraceFile + ".");
-            ProtoLogImpl.getSingleInstance().stopProtoLog(pw, true);
+            if (!android.tracing.Flags.perfettoProtolog()) {
+                ((LegacyProtoLogImpl) mProtoLog).stopProtoLog(pw, true);
+            }
             logAndPrintln(pw, "Start tracing to " + mTraceFile + ".");
             mBuffer.resetBuffer();
             mEnabled = mEnabledLockFree = true;
-            ProtoLogImpl.getSingleInstance().startProtoLog(pw);
+            if (!android.tracing.Flags.perfettoProtolog()) {
+                ((LegacyProtoLogImpl) mProtoLog).startProtoLog(pw);
+            }
         }
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java b/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java
index af0d32c..c5bf78b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java
@@ -27,6 +27,8 @@
 
 import com.android.internal.protolog.ProtoLogGroup;
 import com.android.internal.protolog.ProtoLogImpl;
+import com.android.internal.protolog.common.IProtoLog;
+import com.android.internal.protolog.common.LogLevel;
 import com.android.internal.protolog.common.ProtoLog;
 
 import org.junit.After;
@@ -47,9 +49,9 @@
     @Ignore("b/163095037")
     @Test
     public void testProtoLogToolIntegration() {
-        ProtoLogImpl mockedProtoLog = mock(ProtoLogImpl.class);
+        IProtoLog mockedProtoLog = mock(IProtoLog.class);
         runWith(mockedProtoLog, this::testProtoLog);
-        verify(mockedProtoLog).log(eq(ProtoLogImpl.LogLevel.ERROR), eq(ProtoLogGroup.TEST_GROUP),
+        verify(mockedProtoLog).log(eq(LogLevel.ERROR), eq(ProtoLogGroup.TEST_GROUP),
                 anyInt(), eq(0b0010010111),
                 eq(com.android.internal.protolog.ProtoLogGroup.TEST_GROUP.isLogToLogcat()
                         ? "Test completed successfully: %b %d %x %f %% %s"
@@ -66,18 +68,13 @@
     /**
      * Starts protolog for the duration of {@code runnable}, with a ProtoLogImpl instance installed.
      */
-    private void runWith(ProtoLogImpl mockInstance, Runnable runnable) {
-        ProtoLogImpl original = ProtoLogImpl.getSingleInstance();
-        original.startProtoLog(null);
+    private void runWith(IProtoLog mockInstance, Runnable runnable) {
+        IProtoLog original = ProtoLog.getSingleInstance();
+        ProtoLogImpl.setSingleInstance(mockInstance);
         try {
-            ProtoLogImpl.setSingleInstance(mockInstance);
-            try {
-                runnable.run();
-            } finally {
-                ProtoLogImpl.setSingleInstance(original);
-            }
+            runnable.run();
         } finally {
-            original.stopProtoLog(null, false);
+            ProtoLogImpl.setSingleInstance(original);
         }
     }
 }