Merge "Add feature to capture layer" into sc-dev

libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java

@@ -46,6 +46,7 @@ import android.window.TaskOrganizer;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.common.ProtoLog;
+import com.android.wm.shell.common.ScreenshotUtils;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.sizecompatui.SizeCompatUIController;
 import com.android.wm.shell.startingsurface.StartingWindowController;
@@ -55,6 +56,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * Unified task organizer for all components in the shell.
@@ -342,6 +344,19 @@ public class ShellTaskOrganizer extends TaskOrganizer {
         notifySizeCompatUI(info.getTaskInfo(), listener);
     }
 
+    /**
+     * Take a screenshot of a task.
+     */
+    public void screenshotTask(RunningTaskInfo taskInfo, Rect crop,
+            Consumer<SurfaceControl.ScreenshotHardwareBuffer> consumer) {
+        final TaskAppearedInfo info = mTasks.get(taskInfo.taskId);
+        if (info == null) {
+            return;
+        }
+        ScreenshotUtils.captureLayer(info.getLeash(), crop, consumer);
+    }
+
+
     @Override
     public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
         synchronized (mLock) {

libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java

@@ -21,11 +21,67 @@ import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.view.SurfaceControl;
 
+import java.util.function.Consumer;
+
 /**
  * Helpers for working with screenshots.
  */
 public class ScreenshotUtils {
 
+    /**
+     * Take a screenshot of the specified SurfaceControl.
+     *
+     * @param sc the SurfaceControl to take a screenshot of
+     * @param crop the crop to use when capturing the screenshot
+     * @param consumer Consumer for the captured buffer
+     */
+    public static void captureLayer(SurfaceControl sc, Rect crop,
+            Consumer<SurfaceControl.ScreenshotHardwareBuffer> consumer) {
+        consumer.accept(SurfaceControl.captureLayers(
+                new SurfaceControl.LayerCaptureArgs.Builder(sc)
+                    .setSourceCrop(crop)
+                    .setCaptureSecureLayers(true)
+                    .setAllowProtected(true)
+                    .build()));
+    }
+
+    private static class BufferConsumer implements
+            Consumer<SurfaceControl.ScreenshotHardwareBuffer> {
+        SurfaceControl mScreenshot = null;
+        SurfaceControl.Transaction mTransaction;
+        SurfaceControl mSurfaceControl;
+        int mLayer;
+
+        BufferConsumer(SurfaceControl.Transaction t, SurfaceControl sc, int layer) {
+            mTransaction = t;
+            mSurfaceControl = sc;
+            mLayer = layer;
+        }
+
+        @Override
+        public void accept(SurfaceControl.ScreenshotHardwareBuffer buffer) {
+            if (buffer == null || buffer.getHardwareBuffer() == null) {
+                return;
+            }
+            final GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(
+                buffer.getHardwareBuffer());
+            mScreenshot = new SurfaceControl.Builder()
+                .setName("ScreenshotUtils screenshot")
+                .setFormat(PixelFormat.TRANSLUCENT)
+                .setSecure(buffer.containsSecureLayers())
+                .setCallsite("ScreenshotUtils.takeScreenshot")
+                .setBLASTLayer()
+                .build();
+
+            mTransaction.setBuffer(mScreenshot, graphicBuffer);
+            mTransaction.setColorSpace(mScreenshot, buffer.getColorSpace());
+            mTransaction.reparent(mScreenshot, mSurfaceControl);
+            mTransaction.setLayer(mScreenshot, mLayer);
+            mTransaction.show(mScreenshot);
+            mTransaction.apply();
+        }
+    }
+
     /**
      * Take a screenshot of the specified SurfaceControl.
      *
@@ -38,32 +94,8 @@ public class ScreenshotUtils {
      */
     public static SurfaceControl takeScreenshot(SurfaceControl.Transaction t, SurfaceControl sc,
             Rect crop, int layer) {
-        final SurfaceControl.ScreenshotHardwareBuffer buffer = SurfaceControl.captureLayers(
+        BufferConsumer consumer = new BufferConsumer(t, sc, layer);
-                new SurfaceControl.LayerCaptureArgs.Builder(sc)
+        captureLayer(sc, crop, consumer);
-                        .setSourceCrop(crop)
+        return consumer.mScreenshot;
-                        .setCaptureSecureLayers(true)
-                        .setAllowProtected(true)
-                        .build()
-        );
-        if (buffer == null || buffer.getHardwareBuffer() == null) {
-            return null;
-        }
-        final GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(
-                buffer.getHardwareBuffer());
-        final SurfaceControl screenshot = new SurfaceControl.Builder()
-                .setName("ScreenshotUtils screenshot")
-                .setFormat(PixelFormat.TRANSLUCENT)
-                .setSecure(buffer.containsSecureLayers())
-                .setCallsite("ScreenshotUtils.takeScreenshot")
-                .setBLASTLayer()
-                .build();
-
-        t.setBuffer(screenshot, graphicBuffer);
-        t.setColorSpace(screenshot, buffer.getColorSpace());
-        t.reparent(screenshot, sc);
-        t.setLayer(screenshot, layer);
-        t.show(screenshot);
-        t.apply();
-        return screenshot;
     }
 }

libs/WindowManager/Shell/src/com/android/wm/shell/tasksurfacehelper/TaskSurfaceHelper.java

@@ -16,6 +16,13 @@
 
 package com.android.wm.shell.tasksurfacehelper;
 
+import android.app.ActivityManager.RunningTaskInfo;
+import android.graphics.Rect;
+import android.view.SurfaceControl;
+
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+
 /**
  * Interface to communicate with a Task's SurfaceControl.
  */
@@ -23,4 +30,8 @@ public interface TaskSurfaceHelper {
 
     /** Sets the METADATA_GAME_MODE for the layer corresponding to the task **/
     default void setGameModeForTask(int taskId, int gameMode) {}
+
+    /** Takes a screenshot for a task **/
+    default void screenshotTask(RunningTaskInfo taskInfo, Rect crop, Executor executor,
+            Consumer<SurfaceControl.ScreenshotHardwareBuffer> consumer) {}
 }

libs/WindowManager/Shell/src/com/android/wm/shell/tasksurfacehelper/TaskSurfaceHelperController.java

@@ -16,11 +16,16 @@
 
 package com.android.wm.shell.tasksurfacehelper;
 
+import android.app.ActivityManager.RunningTaskInfo;
+import android.graphics.Rect;
 import android.view.SurfaceControl;
 
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.ShellExecutor;
 
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+
 /**
  * Intermediary controller that communicates with {@link ShellTaskOrganizer} to send commands
  * to SurfaceControl.
@@ -49,6 +54,14 @@ public class TaskSurfaceHelperController {
         mTaskOrganizer.setSurfaceMetadata(taskId, SurfaceControl.METADATA_GAME_MODE, gameMode);
     }
 
+    /**
+     * Take screenshot of the specified task.
+     */
+    public void screenshotTask(RunningTaskInfo taskInfo, Rect crop,
+            Consumer<SurfaceControl.ScreenshotHardwareBuffer> consumer) {
+        mTaskOrganizer.screenshotTask(taskInfo, crop, consumer);
+    }
+
     private class TaskSurfaceHelperImpl implements TaskSurfaceHelper {
         @Override
         public void setGameModeForTask(int taskId, int gameMode) {
@@ -56,5 +69,14 @@ public class TaskSurfaceHelperController {
                 TaskSurfaceHelperController.this.setGameModeForTask(taskId, gameMode);
             });
         }
+
+        @Override
+        public void screenshotTask(RunningTaskInfo taskInfo, Rect crop, Executor executor,
+                Consumer<SurfaceControl.ScreenshotHardwareBuffer> consumer) {
+            mMainExecutor.execute(() -> {
+                TaskSurfaceHelperController.this.screenshotTask(taskInfo, crop,
+                        (t) -> executor.execute(() -> consumer.accept(t)));
+            });
+        }
     }
 }