diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 66267bd..73bf3dc 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -94,6 +94,7 @@
 import android.app.ActivityManager.StackId;
 import android.content.res.Configuration;
 import android.graphics.Bitmap;
+import android.graphics.GraphicBuffer;
 import android.graphics.Matrix;
 import android.graphics.Rect;
 import android.graphics.RectF;
@@ -2105,12 +2106,55 @@
      * @param wallpaperOnly true if only the wallpaper layer should be included in the screenshot
      * @param includeDecor whether to include window decors, like the status or navigation bar
      *                     background of the window
-     * @param toAshmem whether to convert the resulting bitmap to ashmem; this should be set to
-     *                 true if the Bitmap is sent over binder, and false otherwise
      */
     Bitmap screenshotApplications(IBinder appToken, int width, int height,
             boolean includeFullDisplay, float frameScale, Bitmap.Config config,
-            boolean wallpaperOnly, boolean includeDecor, boolean toAshmem) {
+            boolean wallpaperOnly, boolean includeDecor) {
+        Bitmap bitmap = screenshotApplications(appToken, width, height, includeFullDisplay,
+                frameScale, wallpaperOnly, includeDecor, SurfaceControl::screenshot);
+
+        if (DEBUG_SCREENSHOT) {
+            // TEST IF IT's ALL BLACK
+            int[] buffer = new int[bitmap.getWidth() * bitmap.getHeight()];
+            bitmap.getPixels(buffer, 0, bitmap.getWidth(), 0, 0, bitmap.getWidth(),
+                    bitmap.getHeight());
+            boolean allBlack = true;
+            final int firstColor = buffer[0];
+            for (int i = 0; i < buffer.length; i++) {
+                if (buffer[i] != firstColor) {
+                    allBlack = false;
+                    break;
+                }
+            }
+            if (allBlack) {
+                final WindowState appWin = mScreenshotApplicationState.appWin;
+                final int maxLayer = mScreenshotApplicationState.maxLayer;
+                final int minLayer = mScreenshotApplicationState.minLayer;
+                Slog.i(TAG_WM, "Screenshot " + appWin + " was monochrome(" +
+                        Integer.toHexString(firstColor) + ")! mSurfaceLayer=" +
+                        (appWin != null ?
+                                appWin.mWinAnimator.mSurfaceController.getLayer() : "null") +
+                        " minLayer=" + minLayer + " maxLayer=" + maxLayer);
+            }
+        }
+
+        // Create a copy of the screenshot that is immutable and backed in ashmem.
+        // This greatly reduces the overhead of passing the bitmap between processes.
+        Bitmap ret = bitmap.createAshmemBitmap(config);
+        bitmap.recycle();
+        return ret;
+    }
+
+    GraphicBuffer screenshotApplicationsToBuffer(IBinder appToken, int width, int height,
+            boolean includeFullDisplay, float frameScale, boolean wallpaperOnly,
+            boolean includeDecor) {
+        return screenshotApplications(appToken, width, height, includeFullDisplay, frameScale,
+                wallpaperOnly, includeDecor, SurfaceControl::screenshotToBuffer);
+    }
+
+    private <E> E screenshotApplications(IBinder appToken, int width, int height,
+            boolean includeFullDisplay, float frameScale, boolean wallpaperOnly,
+            boolean includeDecor, Screenshoter<E> screenshoter) {
         int dw = mDisplayInfo.logicalWidth;
         int dh = mDisplayInfo.logicalHeight;
         if (dw == 0 || dh == 0) {
@@ -2119,7 +2163,7 @@
             return null;
         }
 
-        Bitmap bm = null;
+        E bitmap;
 
         mScreenshotApplicationState.reset(appToken == null && !wallpaperOnly);
         final Rect frame = new Rect();
@@ -2327,48 +2371,15 @@
             SurfaceControl.openTransaction();
             SurfaceControl.closeTransactionSync();
 
-            bm = SurfaceControl.screenshot(crop, width, height, minLayer, maxLayer,
+            bitmap = screenshoter.screenshot(crop, width, height, minLayer, maxLayer,
                     inRotation, rot);
-            if (bm == null) {
+            if (bitmap == null) {
                 Slog.w(TAG_WM, "Screenshot failure taking screenshot for (" + dw + "x" + dh
                         + ") to layer " + maxLayer);
                 return null;
             }
         }
-
-        if (DEBUG_SCREENSHOT) {
-            // TEST IF IT's ALL BLACK
-            int[] buffer = new int[bm.getWidth() * bm.getHeight()];
-            bm.getPixels(buffer, 0, bm.getWidth(), 0, 0, bm.getWidth(), bm.getHeight());
-            boolean allBlack = true;
-            final int firstColor = buffer[0];
-            for (int i = 0; i < buffer.length; i++) {
-                if (buffer[i] != firstColor) {
-                    allBlack = false;
-                    break;
-                }
-            }
-            if (allBlack) {
-                final WindowState appWin = mScreenshotApplicationState.appWin;
-                final int maxLayer = mScreenshotApplicationState.maxLayer;
-                final int minLayer = mScreenshotApplicationState.minLayer;
-                Slog.i(TAG_WM, "Screenshot " + appWin + " was monochrome(" +
-                        Integer.toHexString(firstColor) + ")! mSurfaceLayer=" +
-                        (appWin != null ?
-                                appWin.mWinAnimator.mSurfaceController.getLayer() : "null") +
-                        " minLayer=" + minLayer + " maxLayer=" + maxLayer);
-            }
-        }
-
-        // Create a copy of the screenshot that is immutable and backed in ashmem.
-        // This greatly reduces the overhead of passing the bitmap between processes.
-        if (toAshmem) {
-            Bitmap ret = bm.createAshmemBitmap(config);
-            bm.recycle();
-            return ret;
-        } else {
-            return bm;
-        }
+        return bitmap;
     }
 
     // TODO: Can this use createRotationMatrix()?
@@ -2721,4 +2732,13 @@
             return mName;
         }
     }
+
+    /**
+     * Interface to screenshot into various types, i.e. {@link Bitmap} and {@link GraphicBuffer}.
+     */
+    @FunctionalInterface
+    private interface Screenshoter<E> {
+        E screenshot(Rect sourceCrop, int width, int height, int minLayer, int maxLayer,
+                boolean useIdentityTransform, int rotation);
+    }
 }