Send ScreenCaptureListener to native screen capture requests.

Allow for asynchronous screenshot by sending a ScreenCaptureListener to
handle screenshot callbacks. All existing calls are still synchronous
since the SurfaceControl method will block the caller until the results
from SurfaceFlinger are ready.

Test: power + volume down
Test: adb shell screencap
Bug: 162367424
Change-Id: I6cb5641dc19f32f262e3120949fc30ea104cff89

3 files changed, 163 insertions, 46 deletions

diff --git a/cmds/screencap/screencap.cpp b/cmds/screencap/screencap.cpp
index dec4a567fc81..5c08704a6623 100644
--- a/cmds/screencap/screencap.cpp
+++ b/cmds/screencap/screencap.cpp
@@ -30,8 +30,9 @@
 
 #include <binder/ProcessState.h>
 
-#include <gui/SurfaceComposerClient.h>
 #include <gui/ISurfaceComposer.h>
+#include <gui/SurfaceComposerClient.h>
+#include <gui/SyncScreenCaptureListener.h>
 
 #include <ui/DisplayInfo.h>
 #include <ui/GraphicTypes.h>
@@ -181,13 +182,18 @@ int main(int argc, char** argv)
     ProcessState::self()->setThreadPoolMaxThreadCount(0);
     ProcessState::self()->startThreadPool();
 
-    ScreenCaptureResults captureResults;
-    status_t result = ScreenshotClient::captureDisplay(displayId->value, captureResults);
+    sp<SyncScreenCaptureListener> captureListener = new SyncScreenCaptureListener();
+    status_t result = ScreenshotClient::captureDisplay(displayId->value, captureListener);
     if (result != NO_ERROR) {
         close(fd);
         return 1;
     }
 
+    ScreenCaptureResults captureResults = captureListener->waitForResults();
+    if (captureResults.result != NO_ERROR) {
+        close(fd);
+        return 1;
+    }
     ui::Dataspace dataspace = captureResults.capturedDataspace;
     sp<GraphicBuffer> buffer = captureResults.buffer;
 
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index 2bd3c06c318d..d55c25fc1a4f 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -65,6 +65,8 @@ import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.ArrayList;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Handle to an on-screen Surface managed by the system compositor. The SurfaceControl is
@@ -87,10 +89,10 @@ public final class SurfaceControl implements Parcelable {
     private static native void nativeWriteToParcel(long nativeObject, Parcel out);
     private static native void nativeRelease(long nativeObject);
     private static native void nativeDisconnect(long nativeObject);
-    private static native ScreenshotHardwareBuffer nativeCaptureDisplay(
-            DisplayCaptureArgs captureArgs);
-    private static native ScreenshotHardwareBuffer nativeCaptureLayers(
-            LayerCaptureArgs captureArgs);
+    private static native int nativeCaptureDisplay(DisplayCaptureArgs captureArgs,
+            ScreenCaptureListener captureListener);
+    private static native int nativeCaptureLayers(LayerCaptureArgs captureArgs,
+            ScreenCaptureListener captureListener);
     private static native long nativeMirrorSurface(long mirrorOfObject);
     private static native long nativeCreateTransaction();
     private static native long nativeGetNativeTransactionFinalizer();
@@ -613,6 +615,39 @@ public final class SurfaceControl implements Parcelable {
     }
 
     /**
+     * @hide
+     */
+    public interface ScreenCaptureListener {
+        /**
+         * The callback invoked when the screen capture is complete.
+         * @param hardwareBuffer Data containing info about the screen capture.
+         */
+        void onScreenCaptureComplete(ScreenshotHardwareBuffer hardwareBuffer);
+    }
+
+    private static class SyncScreenCaptureListener implements ScreenCaptureListener {
+        private static final int SCREENSHOT_WAIT_TIME_S = 1;
+        private ScreenshotHardwareBuffer mScreenshotHardwareBuffer;
+        private final CountDownLatch mCountDownLatch = new CountDownLatch(1);
+
+        @Override
+        public void onScreenCaptureComplete(ScreenshotHardwareBuffer hardwareBuffer) {
+            mScreenshotHardwareBuffer = hardwareBuffer;
+            mCountDownLatch.countDown();
+        }
+
+        private ScreenshotHardwareBuffer waitForScreenshot() {
+            try {
+                mCountDownLatch.await(SCREENSHOT_WAIT_TIME_S, TimeUnit.SECONDS);
+            } catch (Exception e) {
+                Log.e(TAG, "Failed to wait for screen capture result", e);
+            }
+
+            return mScreenshotHardwareBuffer;
+        }
+    }
+
+    /**
      * A common arguments class used for various screenshot requests. This contains arguments that
      * are shared between {@link DisplayCaptureArgs} and {@link LayerCaptureArgs}
      * @hide
@@ -687,7 +722,7 @@ public final class SurfaceControl implements Parcelable {
     /**
      * The arguments class used to make display capture requests.
      *
-     * @see #nativeCaptureDisplay(DisplayCaptureArgs)
+     * @see #nativeCaptureDisplay(DisplayCaptureArgs, ScreenCaptureListener)
      * @hide
      */
     public static class DisplayCaptureArgs extends CaptureArgs {
@@ -2228,13 +2263,30 @@ public final class SurfaceControl implements Parcelable {
     }
 
     /**
+     * @param captureArgs Arguments about how to take the screenshot
+     * @param captureListener A listener to receive the screenshot callback
+     * @hide
+     */
+    public static int captureDisplay(@NonNull DisplayCaptureArgs captureArgs,
+            @NonNull ScreenCaptureListener captureListener) {
+        return nativeCaptureDisplay(captureArgs, captureListener);
+    }
+
+    /**
      * Captures all the surfaces in a display and returns a {@link ScreenshotHardwareBuffer} with
      * the content.
      *
      * @hide
      */
     public static ScreenshotHardwareBuffer captureDisplay(DisplayCaptureArgs captureArgs) {
-        return nativeCaptureDisplay(captureArgs);
+        SyncScreenCaptureListener screenCaptureListener = new SyncScreenCaptureListener();
+
+        int status = captureDisplay(captureArgs, screenCaptureListener);
+        if (status != 0) {
+            return null;
+        }
+
+        return screenCaptureListener.waitForScreenshot();
     }
 
     /**
@@ -2279,14 +2331,21 @@ public final class SurfaceControl implements Parcelable {
                 .setPixelFormat(format)
                 .build();
 
-        return nativeCaptureLayers(captureArgs);
+        return captureLayers(captureArgs);
     }
 
     /**
      * @hide
      */
     public static ScreenshotHardwareBuffer captureLayers(LayerCaptureArgs captureArgs) {
-        return nativeCaptureLayers(captureArgs);
+        SyncScreenCaptureListener screenCaptureListener = new SyncScreenCaptureListener();
+
+        int status = captureLayers(captureArgs, screenCaptureListener);
+        if (status != 0) {
+            return null;
+        }
+
+        return screenCaptureListener.waitForScreenshot();
     }
 
     /**
@@ -2303,7 +2362,17 @@ public final class SurfaceControl implements Parcelable {
                 .setExcludeLayers(exclude)
                 .build();
 
-        return nativeCaptureLayers(captureArgs);
+        return captureLayers(captureArgs);
+    }
+
+    /**
+     * @param captureArgs Arguments about how to take the screenshot
+     * @param captureListener A listener to receive the screenshot callback
+     * @hide
+     */
+    public static int captureLayers(@NonNull LayerCaptureArgs captureArgs,
+            @NonNull ScreenCaptureListener captureListener) {
+        return nativeCaptureLayers(captureArgs, captureListener);
     }
 
     /**
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 4ef15d7098ff..9efe4b15ce7a 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -30,6 +30,7 @@
 #include <android_runtime/android_hardware_HardwareBuffer.h>
 #include <android_runtime/android_view_Surface.h>
 #include <android_runtime/android_view_SurfaceSession.h>
+#include <gui/IScreenCaptureListener.h>
 #include <gui/ISurfaceComposer.h>
 #include <gui/Surface.h>
 #include <gui/SurfaceComposerClient.h>
@@ -188,6 +189,11 @@ static struct {
 
 static struct {
     jclass clazz;
+    jmethodID onScreenCaptureComplete;
+} gScreenCaptureListenerClassInfo;
+
+static struct {
+    jclass clazz;
     jmethodID ctor;
     jfieldID defaultConfig;
     jfieldID primaryRefreshRateMin;
@@ -226,6 +232,56 @@ constexpr ui::Dataspace pickDataspaceFromColorMode(const ui::ColorMode colorMode
     }
 }
 
+class ScreenCaptureListenerWrapper : public BnScreenCaptureListener {
+public:
+    explicit ScreenCaptureListenerWrapper(JNIEnv* env, jobject jobject) {
+        env->GetJavaVM(&mVm);
+        screenCaptureListenerObject = env->NewGlobalRef(jobject);
+        LOG_ALWAYS_FATAL_IF(!screenCaptureListenerObject, "Failed to make global ref");
+    }
+
+    ~ScreenCaptureListenerWrapper() {
+        if (screenCaptureListenerObject) {
+            getenv()->DeleteGlobalRef(screenCaptureListenerObject);
+            screenCaptureListenerObject = nullptr;
+        }
+    }
+
+    status_t onScreenCaptureComplete(const ScreenCaptureResults& captureResults) override {
+        JNIEnv* env = getenv();
+        if (captureResults.result != NO_ERROR || captureResults.buffer == nullptr) {
+            env->CallVoidMethod(screenCaptureListenerObject,
+                                gScreenCaptureListenerClassInfo.onScreenCaptureComplete, nullptr);
+            return NO_ERROR;
+        }
+        jobject jhardwareBuffer = android_hardware_HardwareBuffer_createFromAHardwareBuffer(
+                env, captureResults.buffer->toAHardwareBuffer());
+        const jint namedColorSpace =
+                fromDataspaceToNamedColorSpaceValue(captureResults.capturedDataspace);
+        jobject screenshotHardwareBuffer =
+                env->CallStaticObjectMethod(gScreenshotHardwareBufferClassInfo.clazz,
+                                            gScreenshotHardwareBufferClassInfo.builder,
+                                            jhardwareBuffer, namedColorSpace,
+                                            captureResults.capturedSecureLayers);
+        env->CallVoidMethod(screenCaptureListenerObject,
+                            gScreenCaptureListenerClassInfo.onScreenCaptureComplete,
+                            screenshotHardwareBuffer);
+        env->DeleteLocalRef(jhardwareBuffer);
+        env->DeleteLocalRef(screenshotHardwareBuffer);
+        return NO_ERROR;
+    }
+
+private:
+    jobject screenCaptureListenerObject;
+    JavaVM* mVm;
+
+    JNIEnv* getenv() {
+        JNIEnv* env;
+        mVm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6);
+        return env;
+    }
+};
+
 // ----------------------------------------------------------------------------
 
 static jlong nativeCreateTransaction(JNIEnv* env, jclass clazz) {
@@ -327,36 +383,28 @@ static DisplayCaptureArgs displayCaptureArgsFromObject(JNIEnv* env,
     return captureArgs;
 }
 
-static jobject nativeCaptureDisplay(JNIEnv* env, jclass clazz, jobject displayCaptureArgsObject) {
+static jint nativeCaptureDisplay(JNIEnv* env, jclass clazz, jobject displayCaptureArgsObject,
+                                 jobject screenCaptureListenerObject) {
     const DisplayCaptureArgs captureArgs =
             displayCaptureArgsFromObject(env, displayCaptureArgsObject);
 
     if (captureArgs.displayToken == NULL) {
-        return NULL;
-    }
-
-    ScreenCaptureResults captureResults;
-    status_t res = ScreenshotClient::captureDisplay(captureArgs, captureResults);
-    if (res != NO_ERROR) {
-        return NULL;
+        return BAD_VALUE;
     }
 
-    jobject jhardwareBuffer = android_hardware_HardwareBuffer_createFromAHardwareBuffer(
-            env, captureResults.buffer->toAHardwareBuffer());
-    const jint namedColorSpace =
-            fromDataspaceToNamedColorSpaceValue(captureResults.capturedDataspace);
-    return env->CallStaticObjectMethod(gScreenshotHardwareBufferClassInfo.clazz,
-                                       gScreenshotHardwareBufferClassInfo.builder, jhardwareBuffer,
-                                       namedColorSpace, captureResults.capturedSecureLayers);
+    sp<IScreenCaptureListener> captureListener =
+            new ScreenCaptureListenerWrapper(env, screenCaptureListenerObject);
+    return ScreenshotClient::captureDisplay(captureArgs, captureListener);
 }
 
-static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureArgsObject) {
+static jint nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureArgsObject,
+                                jobject screenCaptureListenerObject) {
     LayerCaptureArgs captureArgs;
     getCaptureArgs(env, layerCaptureArgsObject, captureArgs);
     SurfaceControl* layer = reinterpret_cast<SurfaceControl*>(
             env->GetLongField(layerCaptureArgsObject, gLayerCaptureArgsClassInfo.layer));
     if (layer == nullptr) {
-        return nullptr;
+        return BAD_VALUE;
     }
 
     captureArgs.layerHandle = layer->getHandle();
@@ -380,19 +428,9 @@ static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptu
         env->ReleaseLongArrayElements(excludeObjectArray, const_cast<jlong*>(objects), JNI_ABORT);
     }
 
-    ScreenCaptureResults captureResults;
-    status_t res = ScreenshotClient::captureLayers(captureArgs, captureResults);
-    if (res != NO_ERROR) {
-        return NULL;
-    }
-
-    jobject jhardwareBuffer = android_hardware_HardwareBuffer_createFromAHardwareBuffer(
-            env, captureResults.buffer->toAHardwareBuffer());
-    const jint namedColorSpace =
-            fromDataspaceToNamedColorSpaceValue(captureResults.capturedDataspace);
-    return env->CallStaticObjectMethod(gScreenshotHardwareBufferClassInfo.clazz,
-                                       gScreenshotHardwareBufferClassInfo.builder, jhardwareBuffer,
-                                       namedColorSpace, captureResults.capturedSecureLayers);
+    sp<IScreenCaptureListener> captureListener =
+            new ScreenCaptureListenerWrapper(env, screenCaptureListenerObject);
+    return ScreenshotClient::captureLayers(captureArgs, captureListener);
 }
 
 static void nativeApplyTransaction(JNIEnv* env, jclass clazz, jlong transactionObj, jboolean sync) {
@@ -1664,12 +1702,10 @@ static const JNINativeMethod sSurfaceControlMethods[] = {
     {"nativeSetOverrideScalingMode", "(JJI)V",
             (void*)nativeSetOverrideScalingMode },
     {"nativeCaptureDisplay",
-            "(Landroid/view/SurfaceControl$DisplayCaptureArgs;)"
-            "Landroid/view/SurfaceControl$ScreenshotHardwareBuffer;",
+            "(Landroid/view/SurfaceControl$DisplayCaptureArgs;Landroid/view/SurfaceControl$ScreenCaptureListener;)I",
             (void*)nativeCaptureDisplay },
     {"nativeCaptureLayers",
-            "(Landroid/view/SurfaceControl$LayerCaptureArgs;)"
-            "Landroid/view/SurfaceControl$ScreenshotHardwareBuffer;",
+            "(Landroid/view/SurfaceControl$LayerCaptureArgs;Landroid/view/SurfaceControl$ScreenCaptureListener;)I",
             (void*)nativeCaptureLayers },
     {"nativeSetInputWindowInfo", "(JJLandroid/view/InputWindowHandle;)V",
             (void*)nativeSetInputWindowInfo },
@@ -1875,6 +1911,12 @@ int register_android_view_SurfaceControl(JNIEnv* env)
     gLayerCaptureArgsClassInfo.childrenOnly =
             GetFieldIDOrDie(env, layerCaptureArgsClazz, "mChildrenOnly", "Z");
 
+    jclass screenCaptureListenerClazz =
+            FindClassOrDie(env, "android/view/SurfaceControl$ScreenCaptureListener");
+    gScreenCaptureListenerClassInfo.clazz = MakeGlobalRefOrDie(env, screenCaptureListenerClazz);
+    gScreenCaptureListenerClassInfo.onScreenCaptureComplete =
+            GetMethodIDOrDie(env, screenCaptureListenerClazz, "onScreenCaptureComplete",
+                             "(Landroid/view/SurfaceControl$ScreenshotHardwareBuffer;)V");
     return err;
 }