Expose snapshotTask for use in WindowManager

Exposes snapshotTask which will immediately take a task snapshot
instead of checking the cache. Also adds an overload which allows
specifying scaleRatio and pixelFormat

Change-Id: I0ee3d90ae0ff508e10a8b24bd47d593f28ab1ea8
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index dee9e9f..4de36b8b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -89,6 +89,14 @@
     @VisibleForTesting
     static final int SNAPSHOT_MODE_NONE = 2;
 
+    /**
+     * Constant for <code>scaleFactor</code> when calling {@link #snapshotTask} which is
+     * interpreted as using the most appropriate scale ratio for the system.
+     * This may yield a smaller ratio on low memory devices.
+     */
+    @VisibleForTesting
+    static final float SNAPSHOT_SCALE_AUTO = -1f;
+
     private final WindowManagerService mService;
 
     private final TaskSnapshotCache mCache;
@@ -260,6 +268,84 @@
         });
     }
 
+    /**
+     * Validates the state of the Task is appropriate to capture a snapshot, collects
+     * information from the task and populates the builder.
+     *
+     * @param task the task to capture
+     * @param scaleFraction the scale fraction between 0-1.0, or {@link #SNAPSHOT_SCALE_AUTO}
+     *                      to automatically select
+     * @param pixelFormat the desired pixel format, or {@link PixelFormat#UNKNOWN} to
+     *                    automatically select
+     * @param builder the snapshot builder to populate
+     *
+     * @return true if the state of the task is ok to proceed
+     */
+    private boolean prepareTaskSnapshot(Task task, float scaleFraction, int pixelFormat,
+            TaskSnapshot.Builder builder) {
+        if (!mService.mPolicy.isScreenOn()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
+            }
+            return false;
+        }
+        final ActivityRecord activity = findAppTokenForSnapshot(task);
+        if (activity == null) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.w(TAG_WM, "Failed to take screenshot. No visible windows for " + task);
+            }
+            return false;
+        }
+        if (activity.hasCommittedReparentToAnimationLeash()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.w(TAG_WM, "Failed to take screenshot. App is animating " + activity);
+            }
+            return false;
+        }
+
+        final WindowState mainWindow = activity.findMainWindow();
+        if (mainWindow == null) {
+            Slog.w(TAG_WM, "Failed to take screenshot. No main window for " + task);
+            return false;
+        }
+
+        builder.setIsRealSnapshot(true);
+        builder.setId(System.currentTimeMillis());
+        builder.setContentInsets(getInsets(mainWindow));
+
+        final boolean isLowRamDevice = ActivityManager.isLowRamDeviceStatic();
+
+        if (scaleFraction == SNAPSHOT_SCALE_AUTO) {
+            builder.setScaleFraction(isLowRamDevice
+                    ? mPersister.getReducedScale()
+                    : mFullSnapshotScale);
+            builder.setReducedResolution(isLowRamDevice);
+        } else {
+            builder.setScaleFraction(scaleFraction);
+            builder.setReducedResolution(scaleFraction < 1.0f);
+        }
+
+        final boolean isWindowTranslucent = mainWindow.getAttrs().format != PixelFormat.OPAQUE;
+        final boolean isShowWallpaper = (mainWindow.getAttrs().flags & FLAG_SHOW_WALLPAPER) != 0;
+
+        if (pixelFormat == PixelFormat.UNKNOWN) {
+            pixelFormat = mPersister.use16BitFormat() && activity.fillsParent()
+                    && !(isWindowTranslucent && isShowWallpaper)
+                    ? PixelFormat.RGB_565
+                    : PixelFormat.RGBA_8888;
+        }
+
+        final boolean isTranslucent = PixelFormat.formatHasAlpha(pixelFormat)
+                && (!activity.fillsParent() || isWindowTranslucent);
+
+        builder.setTopActivityComponent(activity.mActivityComponent);
+        builder.setIsTranslucent(isTranslucent);
+        builder.setOrientation(activity.getTask().getConfiguration().orientation);
+        builder.setWindowingMode(task.getWindowingMode());
+        builder.setSystemUiVisibility(getSystemUiVisibility(task));
+        return true;
+    }
+
     @Nullable
     SurfaceControl.ScreenshotGraphicBuffer createTaskSnapshot(@NonNull Task task,
             float scaleFraction) {
@@ -288,63 +374,31 @@
         return screenshotBuffer;
     }
 
-    @Nullable private TaskSnapshot snapshotTask(Task task) {
-        if (!mService.mPolicy.isScreenOn()) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
-            }
+    @Nullable
+    TaskSnapshot snapshotTask(Task task) {
+        return snapshotTask(task, SNAPSHOT_SCALE_AUTO, PixelFormat.UNKNOWN);
+    }
+
+    @Nullable
+    TaskSnapshot snapshotTask(Task task, float scaleFraction, int pixelFormat) {
+        TaskSnapshot.Builder builder = new TaskSnapshot.Builder();
+
+        if (!prepareTaskSnapshot(task, scaleFraction, pixelFormat, builder)) {
+            // Failed some pre-req. Has been logged.
             return null;
         }
 
-        final ActivityRecord activity = findAppTokenForSnapshot(task);
-        if (activity == null) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.w(TAG_WM, "Failed to take screenshot. No visible windows for " + task);
-            }
-            return null;
-        }
-        if (activity.hasCommittedReparentToAnimationLeash()) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.w(TAG_WM, "Failed to take screenshot. App is animating " + activity);
-            }
-            return null;
-        }
-
-        final boolean isLowRamDevice = ActivityManager.isLowRamDeviceStatic();
-        final float scaleFraction = isLowRamDevice
-                ? mPersister.getReducedScale()
-                : mFullSnapshotScale;
-
-        final WindowState mainWindow = activity.findMainWindow();
-        if (mainWindow == null) {
-            Slog.w(TAG_WM, "Failed to take screenshot. No main window for " + task);
-            return null;
-        }
-        final boolean isWindowTranslucent = mainWindow.getAttrs().format != PixelFormat.OPAQUE;
-        final boolean isShowWallpaper = (mainWindow.getAttrs().flags & FLAG_SHOW_WALLPAPER) != 0;
-        final int pixelFormat = mPersister.use16BitFormat() && activity.fillsParent()
-                && !(isWindowTranslucent && isShowWallpaper)
-                ? PixelFormat.RGB_565
-                : PixelFormat.RGBA_8888;
         final SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer =
-                createTaskSnapshot(task, scaleFraction, pixelFormat);
+                createTaskSnapshot(task, builder.getScaleFraction(),
+                builder.getPixelFormat());
 
         if (screenshotBuffer == null) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.w(TAG_WM, "Failed to take screenshot for " + task);
-            }
+            // Failed to acquire image. Has been logged.
             return null;
         }
-        final boolean isTranslucent = PixelFormat.formatHasAlpha(pixelFormat)
-                && (!activity.fillsParent() || isWindowTranslucent);
-        return new TaskSnapshot(
-                System.currentTimeMillis() /* id */,
-                activity.mActivityComponent, screenshotBuffer.getGraphicBuffer(),
-                screenshotBuffer.getColorSpace(),
-                activity.getTask().getConfiguration().orientation,
-                getInsets(mainWindow), isLowRamDevice /* reduced */, scaleFraction /* scale */,
-                true /* isRealSnapshot */, task.getWindowingMode(), getSystemUiVisibility(task),
-                isTranslucent);
+        builder.setSnapshot(screenshotBuffer.getGraphicBuffer());
+        builder.setColorSpace(screenshotBuffer.getColorSpace());
+        return builder.build();
     }
 
     private boolean shouldDisableSnapshots() {
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index 7174e5c..e17e0d8 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -23,9 +23,20 @@
 import static com.android.server.wm.TaskSnapshotController.SNAPSHOT_MODE_REAL;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
+import android.app.ActivityManager;
+import android.app.WindowConfiguration;
+import android.content.ComponentName;
+import android.content.res.Configuration;
+import android.graphics.ColorSpace;
+import android.graphics.GraphicBuffer;
+import android.graphics.PixelFormat;
+import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
 import android.util.ArraySet;
+import android.view.View;
 
 import androidx.test.filters.SmallTest;
 
@@ -33,6 +44,7 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mockito;
 
 /**
  * Test class for {@link TaskSnapshotController}.
@@ -110,4 +122,57 @@
         assertEquals(SNAPSHOT_MODE_APP_THEME,
                 mWm.mTaskSnapshotController.getSnapshotMode(secureWindow.getTask()));
     }
+
+    @Test
+    public void testSnapshotBuilder() {
+        final GraphicBuffer buffer = Mockito.mock(GraphicBuffer.class);
+        final ColorSpace sRGB = ColorSpace.get(ColorSpace.Named.SRGB);
+        final long id = 1234L;
+        final ComponentName activityComponent = new ComponentName("package", ".Class");
+        final int windowingMode = WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+        final int systemUiVisibility = View.SYSTEM_UI_FLAG_FULLSCREEN;
+        final int pixelFormat = PixelFormat.RGBA_8888;
+        final int orientation = Configuration.ORIENTATION_PORTRAIT;
+        final float scaleFraction = 0.25f;
+        final Rect contentInsets = new Rect(1, 2, 3, 4);
+
+        try {
+            ActivityManager.TaskSnapshot.Builder builder =
+                    new ActivityManager.TaskSnapshot.Builder();
+            builder.setId(id);
+            builder.setTopActivityComponent(activityComponent);
+            builder.setSystemUiVisibility(systemUiVisibility);
+            builder.setWindowingMode(windowingMode);
+            builder.setColorSpace(sRGB);
+            builder.setReducedResolution(true);
+            builder.setOrientation(orientation);
+            builder.setContentInsets(contentInsets);
+            builder.setIsTranslucent(true);
+            builder.setScaleFraction(0.25f);
+            builder.setSnapshot(buffer);
+            builder.setIsRealSnapshot(true);
+            builder.setPixelFormat(pixelFormat);
+
+            // Not part of TaskSnapshot itself, used in screenshot process
+            assertEquals(pixelFormat, builder.getPixelFormat());
+
+            ActivityManager.TaskSnapshot snapshot = builder.build();
+            assertEquals(id, snapshot.getId());
+            assertEquals(activityComponent, snapshot.getTopActivityComponent());
+            assertEquals(systemUiVisibility, snapshot.getSystemUiVisibility());
+            assertEquals(windowingMode, snapshot.getWindowingMode());
+            assertEquals(sRGB, snapshot.getColorSpace());
+            assertTrue(snapshot.isReducedResolution());
+            assertEquals(orientation, snapshot.getOrientation());
+            assertEquals(contentInsets, snapshot.getContentInsets());
+            assertTrue(snapshot.isTranslucent());
+            assertEquals(scaleFraction, builder.getScaleFraction(), 0.001f);
+            assertSame(buffer, snapshot.getSnapshot());
+            assertTrue(snapshot.isRealSnapshot());
+        } finally {
+            if (buffer != null) {
+                buffer.destroy();
+            }
+        }
+    }
 }