From 3cdbefc6f3852e65a62bd3c84fcefb46ff1b7704 Mon Sep 17 00:00:00 2001
From: Mark Renouf <mrenouf@google.com>
Date: Tue, 28 Sep 2021 10:42:13 -0400
Subject: [PATCH] Remove weak reference to prevent loss of callback

Callback was passed as a method reference and stored in
a weak reference. This bug was hidden until ag/15819081 (no longer
block for render thread) was merged. With this change, the callback
now has a higher chance of being collected early and breaking capture.

Bug: 201158402
Fix: 201158402
Test: atest ScrollCaptureConnectionTest
Change-Id: I294a7e3b4e57f60117ea14492182a690f46f02bb
---
 .../android/view/ScrollCaptureConnection.java | 25 +++++++------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/core/java/android/view/ScrollCaptureConnection.java b/core/java/android/view/ScrollCaptureConnection.java
index 5fcb011c76f1b..99141c386c0db 100644
--- a/core/java/android/view/ScrollCaptureConnection.java
+++ b/core/java/android/view/ScrollCaptureConnection.java
@@ -33,8 +33,8 @@ import android.util.Log;
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.lang.ref.Reference;
-import java.lang.ref.WeakReference;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 /**
@@ -231,31 +231,24 @@ public class ScrollCaptureConnection extends IScrollCaptureConnection.Stub {
 
     private static class SafeCallback<T> {
         private final CancellationSignal mSignal;
-        private final WeakReference<T> mTargetRef;
         private final Executor mExecutor;
-        private boolean mExecuted;
+        private final AtomicReference<T> mValue;
 
-        protected SafeCallback(CancellationSignal signal, Executor executor, T target) {
+        protected SafeCallback(CancellationSignal signal, Executor executor, T value) {
             mSignal = signal;
-            mTargetRef = new WeakReference<>(target);
+            mValue = new AtomicReference<T>(value);
             mExecutor = executor;
         }
 
-        // Provide the target to the consumer to invoke, forward on handler thread ONCE,
-        // and only if noy cancelled, and the target is still available (not collected)
-        protected final void maybeAccept(Consumer<T> targetConsumer) {
-            if (mExecuted) {
-                return;
-            }
-            mExecuted = true;
+        // Provide the value to the consumer to accept only once.
+        protected final void maybeAccept(Consumer<T> consumer) {
+            T value = mValue.getAndSet(null);
             if (mSignal.isCanceled()) {
                 return;
             }
-            T target = mTargetRef.get();
-            if (target == null) {
-                return;
+            if (value != null) {
+                mExecutor.execute(() -> consumer.accept(value));
             }
-            mExecutor.execute(() -> targetConsumer.accept(target));
         }
 
         static Runnable create(CancellationSignal signal, Executor executor, Runnable target) {