fix WM instance leaks

Somehow weak references to WindowTracingPerfetto's methods become null,
even though WindowTracingPerfetto is not garbage collected.

Directly referencing (weakly) WindowTracingPerfetto,
instead of its methods, solves the issue.

Flag: android.tracing.perfetto_wm_tracing
Bug: 323165543
Test: presubmit
Change-Id: I73eb10b76d7e6ab3666ac977159d04ad0b38fd17

2 files changed, 12 insertions, 18 deletions

diff --git a/services/core/java/com/android/server/wm/WindowTracingDataSource.java b/services/core/java/com/android/server/wm/WindowTracingDataSource.java
index 2c5a4538863e..73ecbb47f0d2 100644
--- a/services/core/java/com/android/server/wm/WindowTracingDataSource.java
+++ b/services/core/java/com/android/server/wm/WindowTracingDataSource.java
@@ -35,7 +35,6 @@ import android.util.proto.ProtoInputStream;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Consumer;
 
 public final class WindowTracingDataSource extends DataSource<WindowTracingDataSource.Instance,
         WindowTracingDataSource.TlsState, Void> {
@@ -77,15 +76,11 @@ public final class WindowTracingDataSource extends DataSource<WindowTracingDataS
     private static final String TAG = "WindowTracingDataSource";
 
     @NonNull
-    private final WeakReference<Consumer<Config>> mOnStartCallback;
-    @NonNull
-    private final WeakReference<Consumer<Config>> mOnStopCallback;
+    private final WeakReference<WindowTracingPerfetto> mWindowTracing;
 
-    public WindowTracingDataSource(@NonNull Consumer<Config> onStart,
-            @NonNull Consumer<Config> onStop) {
+    public WindowTracingDataSource(WindowTracingPerfetto windowTracing) {
         super(DATA_SOURCE_NAME);
-        mOnStartCallback = new WeakReference(onStart);
-        mOnStopCallback = new WeakReference(onStop);
+        mWindowTracing = new WeakReference<>(windowTracing);
 
         Producer.init(InitArguments.DEFAULTS);
         DataSourceParams params =
@@ -103,17 +98,17 @@ public final class WindowTracingDataSource extends DataSource<WindowTracingDataS
         return new Instance(this, instanceIndex, config != null ? config : CONFIG_DEFAULT) {
             @Override
             protected void onStart(StartCallbackArguments args) {
-                Consumer<Config> callback = mOnStartCallback.get();
-                if (callback != null) {
-                    callback.accept(mConfig);
+                WindowTracingPerfetto windowTracing = mWindowTracing.get();
+                if (windowTracing != null) {
+                    windowTracing.onStart(mConfig);
                 }
             }
 
             @Override
             protected void onStop(StopCallbackArguments args) {
-                Consumer<Config> callback = mOnStopCallback.get();
-                if (callback != null) {
-                    callback.accept(mConfig);
+                WindowTracingPerfetto windowTracing = mWindowTracing.get();
+                if (windowTracing != null) {
+                    windowTracing.onStop(mConfig);
                 }
             }
         };
diff --git a/services/core/java/com/android/server/wm/WindowTracingPerfetto.java b/services/core/java/com/android/server/wm/WindowTracingPerfetto.java
index cf948ca76f99..22d6c863fd4f 100644
--- a/services/core/java/com/android/server/wm/WindowTracingPerfetto.java
+++ b/services/core/java/com/android/server/wm/WindowTracingPerfetto.java
@@ -35,8 +35,7 @@ class WindowTracingPerfetto extends WindowTracing {
 
     private final AtomicInteger mCountSessionsOnFrame = new AtomicInteger();
     private final AtomicInteger mCountSessionsOnTransaction = new AtomicInteger();
-    private final WindowTracingDataSource mDataSource = new WindowTracingDataSource(
-            this::onStart, this::onStop);
+    private final WindowTracingDataSource mDataSource = new WindowTracingDataSource(this);
 
     WindowTracingPerfetto(WindowManagerService service, Choreographer choreographer) {
         this(service, choreographer, service.mGlobalLock);
@@ -156,7 +155,7 @@ class WindowTracingPerfetto extends WindowTracing {
         return mCountSessionsOnTransaction.get() > 0;
     }
 
-    private void onStart(WindowTracingDataSource.Config config) {
+    void onStart(WindowTracingDataSource.Config config) {
         if (config.mLogFrequency == WindowTracingLogFrequency.FRAME) {
             mCountSessionsOnFrame.incrementAndGet();
         } else if (config.mLogFrequency == WindowTracingLogFrequency.TRANSACTION) {
@@ -168,7 +167,7 @@ class WindowTracingPerfetto extends WindowTracing {
         log(WHERE_START_TRACING);
     }
 
-    private void onStop(WindowTracingDataSource.Config config) {
+    void onStop(WindowTracingDataSource.Config config) {
         if (config.mLogFrequency == WindowTracingLogFrequency.FRAME) {
             mCountSessionsOnFrame.decrementAndGet();
         } else if (config.mLogFrequency == WindowTracingLogFrequency.TRANSACTION) {