ScreenCaptureListenerWrapper - hold weak ref

This change should reduce the amount of time taken by Java's garbage collection to free hardware buffers.

Without this change, cleaning up hardware buffers via GC required:
* A GC cycle to collect ScreenCaptureListener
* NativeAllocationRegistry to clean up ScreenCaptureListenerWrapper (this can be tens of seconds)
* A second GC cycle to collect the consumer ScreenCaptureListenerWrapper referenced and the hardware buffer the consumer referenced

With this change, ScreenCaptureListenerWrapper no longer holds a global reference to the consumer allowing GC to clean up unreferenced consumers immediately.

The change also includes a small refactor of the interface used to create synchronous ScreenCaptureListeners.

Bug: 283813337
Test: presumbits
Change-Id: I318d73d6cbf2210db5e4386ab57e4ce3a3629166

11 files changed, 95 insertions, 85 deletions

diff --git a/core/java/android/app/UiAutomationConnection.java b/core/java/android/app/UiAutomationConnection.java
index d96a9d104ec2..34f0964cf823 100644
--- a/core/java/android/app/UiAutomationConnection.java
+++ b/core/java/android/app/UiAutomationConnection.java
@@ -39,7 +39,6 @@ import android.os.ServiceManager;
 import android.os.UserHandle;
 import android.permission.IPermissionManager;
 import android.util.Log;
-import android.util.Pair;
 import android.view.IWindowManager;
 import android.view.InputDevice;
 import android.view.InputEvent;
@@ -52,7 +51,8 @@ import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.IAccessibilityManager;
 import android.window.ScreenCapture;
 import android.window.ScreenCapture.CaptureArgs;
-import android.window.ScreenCapture.ScreenCaptureListener;
+import android.window.ScreenCapture.ScreenshotHardwareBuffer;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 import libcore.io.IoUtils;
 
@@ -235,12 +235,12 @@ public final class UiAutomationConnection extends IUiAutomationConnection.Stub {
             final CaptureArgs captureArgs = new CaptureArgs.Builder<>()
                     .setSourceCrop(crop)
                     .build();
-            Pair<ScreenCaptureListener, ScreenCapture.ScreenshotSync> syncScreenCapture =
+            SynchronousScreenCaptureListener syncScreenCapture =
                     ScreenCapture.createSyncCaptureListener();
             mWindowManager.captureDisplay(DEFAULT_DISPLAY, captureArgs,
-                    syncScreenCapture.first);
-            final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer =
-                    syncScreenCapture.second.get();
+                    syncScreenCapture);
+            final ScreenshotHardwareBuffer screenshotBuffer =
+                    syncScreenCapture.getBuffer();
             return screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         } catch (RemoteException re) {
             re.rethrowAsRuntimeException();
diff --git a/core/java/android/window/ScreenCapture.java b/core/java/android/window/ScreenCapture.java
index fa7f577dadb8..0cc9c644f8ea 100644
--- a/core/java/android/window/ScreenCapture.java
+++ b/core/java/android/window/ScreenCapture.java
@@ -27,7 +27,6 @@ import android.os.IBinder;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.util.Log;
-import android.util.Pair;
 import android.view.SurfaceControl;
 
 import libcore.util.NativeAllocationRegistry;
@@ -73,14 +72,14 @@ public class ScreenCapture {
      */
     public static ScreenshotHardwareBuffer captureDisplay(
             DisplayCaptureArgs captureArgs) {
-        Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture = createSyncCaptureListener();
-        int status = captureDisplay(captureArgs, syncScreenCapture.first);
+        SynchronousScreenCaptureListener syncScreenCapture = createSyncCaptureListener();
+        int status = captureDisplay(captureArgs, syncScreenCapture);
         if (status != 0) {
             return null;
         }
 
         try {
-            return syncScreenCapture.second.get();
+            return syncScreenCapture.getBuffer();
         } catch (Exception e) {
             return null;
         }
@@ -133,14 +132,14 @@ public class ScreenCapture {
      * @hide
      */
     public static ScreenshotHardwareBuffer captureLayers(LayerCaptureArgs captureArgs) {
-        Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture = createSyncCaptureListener();
-        int status = captureLayers(captureArgs, syncScreenCapture.first);
+        SynchronousScreenCaptureListener syncScreenCapture = createSyncCaptureListener();
+        int status = captureLayers(captureArgs, syncScreenCapture);
         if (status != 0) {
             return null;
         }
 
         try {
-            return syncScreenCapture.second.get();
+            return syncScreenCapture.getBuffer();
         } catch (Exception e) {
             return null;
         }
@@ -743,14 +742,35 @@ public class ScreenCapture {
      * A helper method to handle the async screencapture callbacks synchronously. This should only
      * be used if the screencapture caller doesn't care that it blocks waiting for a screenshot.
      *
-     * @return a Pair that holds the {@link ScreenCaptureListener} that should be used for capture
-     * calls into SurfaceFlinger and a {@link ScreenshotSync} object to retrieve the results.
+     * @return a {@link SynchronousScreenCaptureListener} that should be used for capture
+     * calls into SurfaceFlinger.
      */
-    public static Pair<ScreenCaptureListener, ScreenshotSync> createSyncCaptureListener() {
-        final ScreenshotSync screenshotSync = new ScreenshotSync();
-        final ScreenCaptureListener screenCaptureListener = new ScreenCaptureListener(
-                screenshotSync::setScreenshotHardwareBuffer);
-        return new Pair<>(screenCaptureListener, screenshotSync);
+    public static SynchronousScreenCaptureListener createSyncCaptureListener() {
+        ScreenshotHardwareBuffer[] bufferRef = new ScreenshotHardwareBuffer[1];
+        CountDownLatch latch = new CountDownLatch(1);
+        Consumer<ScreenshotHardwareBuffer> consumer = buffer -> {
+            bufferRef[0] = buffer;
+            latch.countDown();
+        };
+
+        return new SynchronousScreenCaptureListener(consumer) {
+            // In order to avoid requiring two GC cycles to clean up the consumer and the buffer
+            // it references, the underlying JNI listener holds a weak reference to the consumer.
+            // This property exists to ensure the consumer stays alive during the listener's
+            // lifetime.
+            private Consumer<ScreenshotHardwareBuffer> mConsumer = consumer;
+
+            @Override
+            public ScreenshotHardwareBuffer getBuffer() {
+                try {
+                    latch.await(SCREENSHOT_WAIT_TIME_S, TimeUnit.SECONDS);
+                    return bufferRef[0];
+                } catch (Exception e) {
+                    Log.e(TAG, "Failed to wait for screen capture result", e);
+                    return null;
+                }
+            }
+        };
     }
 
     /**
@@ -758,28 +778,15 @@ public class ScreenCapture {
      * {@link #captureLayers(LayerCaptureArgs, ScreenCaptureListener)} or
      * {@link #captureDisplay(DisplayCaptureArgs, ScreenCaptureListener)}
      */
-    public static class ScreenshotSync {
-        private final CountDownLatch mCountDownLatch = new CountDownLatch(1);
-        private ScreenshotHardwareBuffer mScreenshotHardwareBuffer;
-
-        private void setScreenshotHardwareBuffer(
-                ScreenshotHardwareBuffer screenshotHardwareBuffer) {
-            mScreenshotHardwareBuffer = screenshotHardwareBuffer;
-            mCountDownLatch.countDown();
+    public abstract static class SynchronousScreenCaptureListener extends ScreenCaptureListener {
+        SynchronousScreenCaptureListener(Consumer<ScreenshotHardwareBuffer> consumer) {
+            super(consumer);
         }
 
         /**
          * Get the {@link ScreenshotHardwareBuffer} synchronously. This can be null if the
          * screenshot failed or if there was no callback in {@link #SCREENSHOT_WAIT_TIME_S} seconds.
          */
-        public ScreenshotHardwareBuffer get() {
-            try {
-                mCountDownLatch.await(SCREENSHOT_WAIT_TIME_S, TimeUnit.SECONDS);
-                return mScreenshotHardwareBuffer;
-            } catch (Exception e) {
-                Log.e(TAG, "Failed to wait for screen capture result", e);
-                return null;
-            }
-        }
+        public abstract ScreenshotHardwareBuffer getBuffer();
     }
 }
diff --git a/core/jni/android_window_ScreenCapture.cpp b/core/jni/android_window_ScreenCapture.cpp
index 986dbe9a204c..e729750b812e 100644
--- a/core/jni/android_window_ScreenCapture.cpp
+++ b/core/jni/android_window_ScreenCapture.cpp
@@ -81,22 +81,28 @@ class ScreenCaptureListenerWrapper : public gui::BnScreenCaptureListener {
 public:
     explicit ScreenCaptureListenerWrapper(JNIEnv* env, jobject jobject) {
         env->GetJavaVM(&mVm);
-        mConsumerObject = env->NewGlobalRef(jobject);
-        LOG_ALWAYS_FATAL_IF(!mConsumerObject, "Failed to make global ref");
+        mConsumerWeak = env->NewWeakGlobalRef(jobject);
     }
 
     ~ScreenCaptureListenerWrapper() {
-        if (mConsumerObject) {
-            getenv()->DeleteGlobalRef(mConsumerObject);
-            mConsumerObject = nullptr;
+        if (mConsumerWeak) {
+            getenv()->DeleteWeakGlobalRef(mConsumerWeak);
+            mConsumerWeak = nullptr;
         }
     }
 
     binder::Status onScreenCaptureCompleted(
             const gui::ScreenCaptureResults& captureResults) override {
         JNIEnv* env = getenv();
+
+        ScopedLocalRef<jobject> consumer{env, env->NewLocalRef(mConsumerWeak)};
+        if (consumer == nullptr) {
+            ALOGE("ScreenCaptureListenerWrapper consumer not alive.");
+            return binder::Status::ok();
+        }
+
         if (!captureResults.fenceResult.ok() || captureResults.buffer == nullptr) {
-            env->CallVoidMethod(mConsumerObject, gConsumerClassInfo.accept, nullptr);
+            env->CallVoidMethod(consumer.get(), gConsumerClassInfo.accept, nullptr);
             checkAndClearException(env, "accept");
             return binder::Status::ok();
         }
@@ -111,7 +117,7 @@ public:
                                             captureResults.capturedSecureLayers,
                                             captureResults.capturedHdrLayers);
         checkAndClearException(env, "builder");
-        env->CallVoidMethod(mConsumerObject, gConsumerClassInfo.accept, screenshotHardwareBuffer);
+        env->CallVoidMethod(consumer.get(), gConsumerClassInfo.accept, screenshotHardwareBuffer);
         checkAndClearException(env, "accept");
         env->DeleteLocalRef(jhardwareBuffer);
         env->DeleteLocalRef(screenshotHardwareBuffer);
@@ -119,7 +125,7 @@ public:
     }
 
 private:
-    jobject mConsumerObject;
+    jweak mConsumerWeak;
     JavaVM* mVm;
 
     JNIEnv* getenv() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index 3eb9fa2eef6b..0fdfbb8c0c61 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -80,8 +80,7 @@ import android.view.ViewRootImpl;
 import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.window.ScreenCapture;
-import android.window.ScreenCapture.ScreenCaptureListener;
-import android.window.ScreenCapture.ScreenshotSync;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 import androidx.annotation.MainThread;
 import androidx.annotation.Nullable;
@@ -1222,10 +1221,11 @@ public class BubbleController implements ConfigurationChangeListener,
 
     /**
      * Performs a screenshot that may exclude the bubble layer, if one is present. The screenshot
-     * can be access via the supplied {@link ScreenshotSync#get()} asynchronously.
+     * can be access via the supplied {@link SynchronousScreenCaptureListener#getBuffer()}
+     * asynchronously.
      */
     public void getScreenshotExcludingBubble(int displayId,
-            Pair<ScreenCaptureListener, ScreenshotSync> screenCaptureListener) {
+            SynchronousScreenCaptureListener screenCaptureListener) {
         try {
             ScreenCapture.CaptureArgs args = null;
             if (mStackView != null) {
@@ -1240,7 +1240,7 @@ public class BubbleController implements ConfigurationChangeListener,
                 }
             }
 
-            mWmService.captureDisplay(displayId, args, screenCaptureListener.first);
+            mWmService.captureDisplay(displayId, args, screenCaptureListener);
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to capture screenshot");
         }
@@ -2211,15 +2211,15 @@ public class BubbleController implements ConfigurationChangeListener,
 
         @Override
         @Nullable
-        public ScreenshotSync getScreenshotExcludingBubble(int displayId) {
-            Pair<ScreenCaptureListener, ScreenshotSync> screenCaptureListener =
+        public SynchronousScreenCaptureListener getScreenshotExcludingBubble(int displayId) {
+            SynchronousScreenCaptureListener screenCaptureListener =
                     ScreenCapture.createSyncCaptureListener();
 
             mMainExecutor.execute(
                     () -> BubbleController.this.getScreenshotExcludingBubble(displayId,
                             screenCaptureListener));
 
-            return screenCaptureListener.second;
+            return screenCaptureListener;
         }
 
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java
index 259f69296ac7..4d329dd5d446 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java
@@ -16,8 +16,6 @@
 
 package com.android.wm.shell.bubbles;
 
-import static android.window.ScreenCapture.ScreenshotSync;
-
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.ElementType.LOCAL_VARIABLE;
 import static java.lang.annotation.ElementType.PARAMETER;
@@ -34,6 +32,7 @@ import android.service.notification.NotificationListenerService.RankingMap;
 import android.util.Pair;
 import android.util.SparseArray;
 import android.window.ScreenCapture.ScreenshotHardwareBuffer;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
@@ -150,13 +149,14 @@ public interface Bubbles {
     boolean isAppBubbleTaskId(int taskId);
 
     /**
-     * @return a {@link ScreenshotSync} after performing a screenshot that may exclude the bubble
-     * layer, if one is present. The underlying {@link ScreenshotHardwareBuffer} can be access via
-     * {@link ScreenshotSync#get()} asynchronously and care should be taken to
-     * {@link HardwareBuffer#close()} the associated
-     * {@link ScreenshotHardwareBuffer#getHardwareBuffer()} when no longer required.
+`    * @return a {@link SynchronousScreenCaptureListener} after performing a screenshot that may
+     * exclude the bubble layer, if one is present. The underlying
+     * {@link ScreenshotHardwareBuffer} can be accessed via
+     * {@link SynchronousScreenCaptureListener#getBuffer()} asynchronously and care should be taken
+     * to {@link HardwareBuffer#close()} the associated
+     * {@link ScreenshotHardwareBuffer#getHardwareBuffer()} when no longer required.`
      */
-    ScreenshotSync getScreenshotExcludingBubble(int displayId);
+    SynchronousScreenCaptureListener getScreenshotExcludingBubble(int displayId);
 
     /**
      * @return a bubble that matches the provided shortcutId, if one exists.
diff --git a/packages/Shell/src/com/android/shell/Screenshooter.java b/packages/Shell/src/com/android/shell/Screenshooter.java
index baaddf53bcab..8a5b7daab7e0 100644
--- a/packages/Shell/src/com/android/shell/Screenshooter.java
+++ b/packages/Shell/src/com/android/shell/Screenshooter.java
@@ -21,12 +21,10 @@ import static android.view.Display.DEFAULT_DISPLAY;
 import android.graphics.Bitmap;
 import android.os.RemoteException;
 import android.util.Log;
-import android.util.Pair;
 import android.view.WindowManagerGlobal;
 import android.window.ScreenCapture;
-import android.window.ScreenCapture.ScreenCaptureListener;
 import android.window.ScreenCapture.ScreenshotHardwareBuffer;
-import android.window.ScreenCapture.ScreenshotSync;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 /**
  * Helper class used to take screenshots.
@@ -46,15 +44,15 @@ final class Screenshooter {
     static Bitmap takeScreenshot() {
         Log.d(TAG, "Taking fullscreen screenshot");
         // Take the screenshot
-        final Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture =
+        final SynchronousScreenCaptureListener syncScreenCapture =
                 ScreenCapture.createSyncCaptureListener();
         try {
             WindowManagerGlobal.getWindowManagerService().captureDisplay(DEFAULT_DISPLAY, null,
-                    syncScreenCapture.first);
+                    syncScreenCapture);
         } catch (RemoteException e) {
             e.rethrowAsRuntimeException();
         }
-        final ScreenshotHardwareBuffer screenshotBuffer = syncScreenCapture.second.get();
+        final ScreenshotHardwareBuffer screenshotBuffer = syncScreenCapture.getBuffer();
         final Bitmap screenShot = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         if (screenShot == null) {
             Log.e(TAG, "Failed to take fullscreen screenshot");
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
index 2e47ab6417ae..24fe7d593b3d 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
@@ -42,8 +42,8 @@ open class ImageCaptureImpl @Inject constructor(
             .setSourceCrop(crop)
             .build()
         val syncScreenCapture = ScreenCapture.createSyncCaptureListener()
-        windowManager.captureDisplay(displayId, captureArgs, syncScreenCapture.first)
-        val buffer = syncScreenCapture.second.get()
+        windowManager.captureDisplay(displayId, captureArgs, syncScreenCapture)
+        val buffer = syncScreenCapture.getBuffer()
         return buffer?.asBitmap()
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperService.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperService.java
index 83ff020362f1..4c013151c464 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperService.java
@@ -20,7 +20,7 @@ import android.app.Service;
 import android.content.Intent;
 import android.os.IBinder;
 import android.window.ScreenCapture.ScreenshotHardwareBuffer;
-import android.window.ScreenCapture.ScreenshotSync;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 import androidx.annotation.Nullable;
 
@@ -54,9 +54,9 @@ public class AppClipsScreenshotHelperService extends Service {
                     return null;
                 }
 
-                ScreenshotSync screenshotSync =
+                SynchronousScreenCaptureListener screenshotSync =
                         mOptionalBubbles.get().getScreenshotExcludingBubble(displayId);
-                ScreenshotHardwareBuffer screenshotHardwareBuffer = screenshotSync.get();
+                ScreenshotHardwareBuffer screenshotHardwareBuffer = screenshotSync.getBuffer();
                 if (screenshotHardwareBuffer == null) {
                     return null;
                 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java
index ab321f1b5341..ba3d39248762 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java
@@ -26,7 +26,7 @@ import android.hardware.HardwareBuffer;
 import android.os.RemoteException;
 import android.view.Display;
 import android.window.ScreenCapture.ScreenshotHardwareBuffer;
-import android.window.ScreenCapture.ScreenshotSync;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 import androidx.test.runner.AndroidJUnit4;
 
@@ -58,7 +58,7 @@ public final class AppClipsScreenshotHelperServiceTest extends SysuiTestCase {
     @Mock private Optional<Bubbles> mBubblesOptional;
     @Mock private Bubbles mBubbles;
     @Mock private ScreenshotHardwareBuffer mScreenshotHardwareBuffer;
-    @Mock private ScreenshotSync mScreenshotSync;
+    @Mock private SynchronousScreenCaptureListener mScreenshotSync;
 
     private AppClipsScreenshotHelperService mAppClipsScreenshotHelperService;
 
@@ -80,7 +80,7 @@ public final class AppClipsScreenshotHelperServiceTest extends SysuiTestCase {
         when(mBubblesOptional.isEmpty()).thenReturn(false);
         when(mBubblesOptional.get()).thenReturn(mBubbles);
         when(mBubbles.getScreenshotExcludingBubble(DEFAULT_DISPLAY)).thenReturn(mScreenshotSync);
-        when(mScreenshotSync.get()).thenReturn(null);
+        when(mScreenshotSync.getBuffer()).thenReturn(null);
 
         assertThat(getInterface().takeScreenshot(DEFAULT_DISPLAY)).isNull();
     }
@@ -90,7 +90,7 @@ public final class AppClipsScreenshotHelperServiceTest extends SysuiTestCase {
         when(mBubblesOptional.isEmpty()).thenReturn(false);
         when(mBubblesOptional.get()).thenReturn(mBubbles);
         when(mBubbles.getScreenshotExcludingBubble(DEFAULT_DISPLAY)).thenReturn(mScreenshotSync);
-        when(mScreenshotSync.get()).thenReturn(mScreenshotHardwareBuffer);
+        when(mScreenshotSync.getBuffer()).thenReturn(mScreenshotHardwareBuffer);
         when(mScreenshotHardwareBuffer.getHardwareBuffer()).thenReturn(FAKE_HARDWARE_BUFFER);
         when(mScreenshotHardwareBuffer.getColorSpace()).thenReturn(FAKE_COLOR_SPACE);
 
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 7fc86b0fdc01..ef19eef22794 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -244,6 +244,7 @@ import android.view.inputmethod.ImeTracker;
 import android.window.DisplayWindowPolicyController;
 import android.window.IDisplayAreaOrganizer;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 import android.window.TransitionRequestInfo;
 
 import com.android.internal.R;
@@ -5088,7 +5089,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             return null;
         }
 
-        Pair<ScreenCapture.ScreenCaptureListener, ScreenCapture.ScreenshotSync> syncScreenCapture =
+        SynchronousScreenCaptureListener syncScreenCapture =
                 ScreenCapture.createSyncCaptureListener();
 
         getBounds(mTmpRect);
@@ -5097,10 +5098,10 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
                 new ScreenCapture.LayerCaptureArgs.Builder(getSurfaceControl())
                         .setSourceCrop(mTmpRect).build();
 
-        ScreenCapture.captureLayers(args, syncScreenCapture.first);
+        ScreenCapture.captureLayers(args, syncScreenCapture);
 
         final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer =
-                syncScreenCapture.second.get();
+                syncScreenCapture.getBuffer();
         final Bitmap bitmap = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         if (bitmap == null) {
             Slog.w(TAG_WM, "Failed to take screenshot");
diff --git a/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java b/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java
index 1ee0959447a6..b181213df003 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java
@@ -43,16 +43,14 @@ import android.os.Handler;
 import android.os.Looper;
 import android.os.ServiceManager;
 import android.platform.test.annotations.Presubmit;
-import android.util.Pair;
 import android.view.IWindowManager;
 import android.view.PointerIcon;
 import android.view.SurfaceControl;
 import android.view.cts.surfacevalidator.BitmapPixelChecker;
 import android.view.cts.surfacevalidator.SaveBitmapHelper;
 import android.window.ScreenCapture;
-import android.window.ScreenCapture.ScreenCaptureListener;
 import android.window.ScreenCapture.ScreenshotHardwareBuffer;
-import android.window.ScreenCapture.ScreenshotSync;
+import android.window.ScreenCapture.SynchronousScreenCaptureListener;
 
 import androidx.annotation.Nullable;
 import androidx.test.filters.SmallTest;
@@ -169,10 +167,10 @@ public class ScreenshotTests {
                 .setPosition(sc, point.x, point.y)
                 .apply(true);
 
-        Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture =
+        SynchronousScreenCaptureListener syncScreenCapture =
                 ScreenCapture.createSyncCaptureListener();
-        windowManager.captureDisplay(DEFAULT_DISPLAY, null, syncScreenCapture.first);
-        ScreenshotHardwareBuffer hardwareBuffer = syncScreenCapture.second.get();
+        windowManager.captureDisplay(DEFAULT_DISPLAY, null, syncScreenCapture);
+        ScreenshotHardwareBuffer hardwareBuffer = syncScreenCapture.getBuffer();
         assertNotNull(hardwareBuffer);
 
         Bitmap screenshot = hardwareBuffer.asBitmap();