From e22573d1d86a8c92f357cd33801b4eb31e72e2b4 Mon Sep 17 00:00:00 2001
From: Haoran Zhang <haoranzhang@google.com>
Date: Wed, 24 May 2023 19:19:39 +0000
Subject: [PATCH] [DO NOT MERGE] Fix Memory Leak: Part II of rewrite render
 info callback

Bug:270049350
Test: atest CtsAutoFillServiceTestCases

Change-Id: I72d6e59e68e1c5e2f4cc733a34e78e7ff3f1daf0
---
 .../InlineSuggestionRequestConsumer.java      | 55 +++++++++++++++++++
 .../com/android/server/autofill/Session.java  | 35 ++++++++----
 2 files changed, 78 insertions(+), 12 deletions(-)
 create mode 100644 services/autofill/java/com/android/server/autofill/InlineSuggestionRequestConsumer.java

diff --git a/services/autofill/java/com/android/server/autofill/InlineSuggestionRequestConsumer.java b/services/autofill/java/com/android/server/autofill/InlineSuggestionRequestConsumer.java
new file mode 100644
index 0000000000000..a3efb25a57351
--- /dev/null
+++ b/services/autofill/java/com/android/server/autofill/InlineSuggestionRequestConsumer.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2023 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.server.autofill;
+
+import android.util.Slog;
+import android.view.inputmethod.InlineSuggestionsRequest;
+
+import java.lang.ref.WeakReference;
+import java.util.function.Consumer;
+
+class InlineSuggestionRequestConsumer implements Consumer<InlineSuggestionsRequest> {
+
+    static final String TAG = "InlineSuggestionRequestConsumer";
+
+    private final WeakReference<Session.AssistDataReceiverImpl> mAssistDataReceiverWeakReference;
+    private final WeakReference<ViewState>  mViewStateWeakReference;
+
+    InlineSuggestionRequestConsumer(WeakReference<Session.AssistDataReceiverImpl>
+            assistDataReceiverWeakReference,
+            WeakReference<ViewState>  viewStateWeakReference) {
+        mAssistDataReceiverWeakReference = assistDataReceiverWeakReference;
+        mViewStateWeakReference = viewStateWeakReference;
+    }
+
+    @Override
+    public void accept(InlineSuggestionsRequest inlineSuggestionsRequest) {
+        Session.AssistDataReceiverImpl assistDataReceiver = mAssistDataReceiverWeakReference.get();
+        ViewState viewState = mViewStateWeakReference.get();
+        if (assistDataReceiver == null) {
+            Slog.wtf(TAG, "assistDataReceiver is null when accepting new inline suggestion"
+                    + "requests");
+            return;
+        }
+
+        if (viewState == null) {
+            Slog.wtf(TAG, "view state is null when accepting new inline suggestion requests");
+            return;
+        }
+        assistDataReceiver.handleInlineSuggestionRequest(inlineSuggestionsRequest, viewState);
+    }
+}
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 2d2b8f896b520..db17cce5f5a30 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -307,7 +307,7 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
      * Id of the View currently being displayed.
      */
     @GuardedBy("mLock")
-    @Nullable AutofillId mCurrentViewId;
+    private @Nullable AutofillId mCurrentViewId;
 
     @GuardedBy("mLock")
     private IAutoFillManagerClient mClient;
@@ -595,7 +595,7 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
      * TODO(b/151867668): improve how asynchronous data dependencies are handled, without using
      * CountDownLatch.
      */
-    private final class AssistDataReceiverImpl extends IAssistDataReceiver.Stub {
+    final class AssistDataReceiverImpl extends IAssistDataReceiver.Stub {
         @GuardedBy("mLock")
         private boolean mWaitForInlineRequest;
         @GuardedBy("mLock")
@@ -610,17 +610,28 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
             mPendingFillRequest = null;
             mWaitForInlineRequest = isInlineRequest;
             mPendingInlineSuggestionsRequest = null;
-            return isInlineRequest ? (inlineSuggestionsRequest) -> {
-                synchronized (mLock) {
-                    if (!mWaitForInlineRequest || mPendingInlineSuggestionsRequest != null) {
-                        return;
-                    }
-                    mWaitForInlineRequest = inlineSuggestionsRequest != null;
-                    mPendingInlineSuggestionsRequest = inlineSuggestionsRequest;
-                    maybeRequestFillLocked();
-                    viewState.resetState(ViewState.STATE_PENDING_CREATE_INLINE_REQUEST);
+            if (isInlineRequest) {
+                WeakReference<AssistDataReceiverImpl> assistDataReceiverWeakReference =
+                        new WeakReference<AssistDataReceiverImpl>(this);
+                WeakReference<ViewState> viewStateWeakReference =
+                        new WeakReference<ViewState>(viewState);
+                return new InlineSuggestionRequestConsumer(assistDataReceiverWeakReference,
+                    viewStateWeakReference);
+            }
+            return null;
+        }
+
+        void handleInlineSuggestionRequest(InlineSuggestionsRequest inlineSuggestionsRequest,
+                ViewState viewState) {
+            synchronized (mLock) {
+                if (!mWaitForInlineRequest || mPendingInlineSuggestionsRequest != null) {
+                    return;
                 }
-            } : null;
+                mWaitForInlineRequest = inlineSuggestionsRequest != null;
+                mPendingInlineSuggestionsRequest = inlineSuggestionsRequest;
+                maybeRequestFillLocked();
+                viewState.resetState(ViewState.STATE_PENDING_CREATE_INLINE_REQUEST);
+            }
         }
 
         @GuardedBy("mLock")