Capture screenshots from all displays (flag-guarded)

This changes TakeScreenshotService to gather screenshots from all
(not-virtual) connected displays when a SystemUI flag is enabled.
System server part of the screenshot flow is unchanged.

When the `multi_display_screenshot` sysui flag is enabled, one
screenshot is captured for each display, and the UI is shown accordingly
in every display. This results in different files being saved.

When the flag is disabled, the previous behaviour is left unchanged.

ScreenshotController has been been slighly refactored to be per-display.
When the flag is disabled, only the controller for the default display
is created.

There is some minimal duplication of code related to UiEventLogger that
can be removed when the flag is enabled.

Test: TakeScreenshotServiceTest, TakeScreenshotExecutorTest, ConnectedDisplayInteractorTest, ScreenshotDataTest
Bug: 290910794
Bug: 295143676
Change-Id: I24bd6436d346c2cbd5dcdc201d01c8aabb09bc17

9 files changed, 706 insertions, 76 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
index c0d807a78e90..98f2fee13de2 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
@@ -28,8 +28,18 @@ import kotlinx.coroutines.launch
 import java.util.function.Consumer
 import javax.inject.Inject
 
+/** Processes a screenshot request sent from [ScreenshotHelper]. */
+interface ScreenshotRequestProcessor {
+    /**
+     * Inspects the incoming ScreenshotData, potentially modifying it based upon policy.
+     *
+     * @param screenshot the screenshot to process
+     */
+    suspend fun process(screenshot: ScreenshotData): ScreenshotData
+}
+
 /**
- * Processes a screenshot request sent from {@link ScreenshotHelper}.
+ * Implementation of [ScreenshotRequestProcessor]
  */
 @SysUISingleton
 class RequestProcessor @Inject constructor(
@@ -38,7 +48,7 @@ class RequestProcessor @Inject constructor(
         private val flags: FeatureFlags,
         /** For the Java Async version, to invoke the callback. */
         @Application private val mainScope: CoroutineScope
-) {
+) : ScreenshotRequestProcessor {
     /**
      * Inspects the incoming request, returning a potentially modified request depending on policy.
      *
@@ -57,7 +67,6 @@ class RequestProcessor @Inject constructor(
         // regardless of the managed profile status.
 
         if (request.type != TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-
             val info = policy.findPrimaryContent(policy.getDefaultDisplayId())
             Log.d(TAG, "findPrimaryContent: $info")
 
@@ -99,12 +108,7 @@ class RequestProcessor @Inject constructor(
         }
     }
 
-    /**
-     * Inspects the incoming ScreenshotData, potentially modifying it based upon policy.
-     *
-     * @param screenshot the screenshot to process
-     */
-    suspend fun process(screenshot: ScreenshotData): ScreenshotData {
+    override suspend fun process(screenshot: ScreenshotData): ScreenshotData {
         var result = screenshot
 
         // Apply work profile screenshots policy:
@@ -116,7 +120,7 @@ class RequestProcessor @Inject constructor(
         // regardless of the managed profile status.
 
         if (screenshot.type != TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            val info = policy.findPrimaryContent(policy.getDefaultDisplayId())
+            val info = policy.findPrimaryContent(screenshot.displayId)
             Log.d(TAG, "findPrimaryContent: $info")
             result.taskId = info.taskId
             result.topComponent = info.component
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index b59106efb769..cf782b7d93c6 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -100,11 +100,14 @@ import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
 import com.android.systemui.screenshot.ScreenshotController.SavedImageData.ActionTransition;
 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback;
-import com.android.systemui.settings.DisplayTracker;
 import com.android.systemui.util.Assert;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
+import dagger.assisted.Assisted;
+import dagger.assisted.AssistedFactory;
+import dagger.assisted.AssistedInject;
+
 import java.io.File;
 import java.util.List;
 import java.util.concurrent.CancellationException;
@@ -118,7 +121,6 @@ import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
-import javax.inject.Inject;
 
 /**
  * Controls the state and flow for screenshots.
@@ -275,7 +277,7 @@ public class ScreenshotController {
     private final ScrollCaptureClient mScrollCaptureClient;
     private final PhoneWindow mWindow;
     private final DisplayManager mDisplayManager;
-    private final DisplayTracker mDisplayTracker;
+    private final int mDisplayId;
     private final ScrollCaptureController mScrollCaptureController;
     private final LongScreenshotData mLongScreenshotHolder;
     private final boolean mIsLowRamDevice;
@@ -314,7 +316,8 @@ public class ScreenshotController {
                     | ActivityInfo.CONFIG_SCREEN_LAYOUT
                     | ActivityInfo.CONFIG_ASSETS_PATHS);
 
-    @Inject
+
+    @AssistedInject
     ScreenshotController(
             Context context,
             FeatureFlags flags,
@@ -335,7 +338,7 @@ public class ScreenshotController {
             UserManager userManager,
             AssistContentRequester assistContentRequester,
             MessageContainerController messageContainerController,
-            DisplayTracker displayTracker
+            @Assisted int displayId
     ) {
         mScreenshotSmartActions = screenshotSmartActions;
         mNotificationsController = screenshotNotificationsController;
@@ -360,9 +363,9 @@ public class ScreenshotController {
             dismissScreenshot(SCREENSHOT_INTERACTION_TIMEOUT);
         });
 
+        mDisplayId = displayId;
         mDisplayManager = requireNonNull(context.getSystemService(DisplayManager.class));
-        mDisplayTracker = displayTracker;
-        final Context displayContext = context.createDisplayContext(getDefaultDisplay());
+        final Context displayContext = context.createDisplayContext(getDisplay());
         mContext = (WindowContext) displayContext.createWindowContext(TYPE_SCREENSHOT, null);
         mWindowManager = mContext.getSystemService(WindowManager.class);
         mFlags = flags;
@@ -406,7 +409,7 @@ public class ScreenshotController {
         if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN) {
             Rect bounds = getFullScreenRect();
             screenshot.setBitmap(
-                    mImageCapture.captureDisplay(mDisplayTracker.getDefaultDisplayId(), bounds));
+                    mImageCapture.captureDisplay(mDisplayId, bounds));
             screenshot.setScreenBounds(bounds);
         }
 
@@ -638,7 +641,7 @@ public class ScreenshotController {
                 setWindowFocusable(false);
             }
         }, mActionExecutor, mFlags);
-        mScreenshotView.setDefaultDisplay(mDisplayTracker.getDefaultDisplayId());
+        mScreenshotView.setDefaultDisplay(mDisplayId);
         mScreenshotView.setDefaultTimeoutMillis(mScreenshotHandler.getDefaultTimeoutMillis());
 
         mScreenshotView.setOnKeyListener((v, keyCode, event) -> {
@@ -727,8 +730,8 @@ public class ScreenshotController {
         if (mLastScrollCaptureRequest != null) {
             mLastScrollCaptureRequest.cancel(true);
         }
-        final ListenableFuture<ScrollCaptureResponse> future =
-                mScrollCaptureClient.request(mDisplayTracker.getDefaultDisplayId());
+        final ListenableFuture<ScrollCaptureResponse> future = mScrollCaptureClient.request(
+                mDisplayId);
         mLastScrollCaptureRequest = future;
         mLastScrollCaptureRequest.addListener(() ->
                 onScrollCaptureResponseReady(future, owner), mMainExecutor);
@@ -758,9 +761,8 @@ public class ScreenshotController {
             final ScrollCaptureResponse response = mLastScrollCaptureResponse;
             mScreenshotView.showScrollChip(response.getPackageName(), /* onClick */ () -> {
                 DisplayMetrics displayMetrics = new DisplayMetrics();
-                getDefaultDisplay().getRealMetrics(displayMetrics);
-                Bitmap newScreenshot = mImageCapture.captureDisplay(
-                        mDisplayTracker.getDefaultDisplayId(),
+                getDisplay().getRealMetrics(displayMetrics);
+                Bitmap newScreenshot = mImageCapture.captureDisplay(mDisplayId,
                         new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels));
 
                 mScreenshotView.prepareScrollingTransition(response, mScreenBitmap, newScreenshot,
@@ -825,7 +827,7 @@ public class ScreenshotController {
             try {
                 WindowManagerGlobal.getWindowManagerService()
                         .overridePendingAppTransitionRemote(runner,
-                                mDisplayTracker.getDefaultDisplayId());
+                                mDisplayId);
             } catch (Exception e) {
                 Log.e(TAG, "Error overriding screenshot app transition", e);
             }
@@ -1160,8 +1162,8 @@ public class ScreenshotController {
         }
     }
 
-    private Display getDefaultDisplay() {
-        return mDisplayManager.getDisplay(mDisplayTracker.getDefaultDisplayId());
+    private Display getDisplay() {
+        return mDisplayManager.getDisplay(mDisplayId);
     }
 
     private boolean allowLongScreenshots() {
@@ -1170,7 +1172,7 @@ public class ScreenshotController {
 
     private Rect getFullScreenRect() {
         DisplayMetrics displayMetrics = new DisplayMetrics();
-        getDefaultDisplay().getRealMetrics(displayMetrics);
+        getDisplay().getRealMetrics(displayMetrics);
         return new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels);
     }
 
@@ -1229,4 +1231,11 @@ public class ScreenshotController {
             };
         }
     }
+
+    /** Injectable factory to create screenshot controller instances for a specific display. */
+    @AssistedFactory
+    public interface Factory {
+        /** Creates an instance of the controller for that specific displayId. */
+        ScreenshotController create(int displayId);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
index e9be88a59990..92e933a9557b 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
@@ -6,12 +6,13 @@ import android.graphics.Insets
 import android.graphics.Rect
 import android.net.Uri
 import android.os.UserHandle
+import android.view.Display
 import android.view.WindowManager.ScreenshotSource
 import android.view.WindowManager.ScreenshotType
 import androidx.annotation.VisibleForTesting
 import com.android.internal.util.ScreenshotRequest
 
-/** ScreenshotData represents the current state of a single screenshot being acquired. */
+/** [ScreenshotData] represents the current state of a single screenshot being acquired. */
 data class ScreenshotData(
     @ScreenshotType var type: Int,
     @ScreenshotSource var source: Int,
@@ -23,6 +24,7 @@ data class ScreenshotData(
     var taskId: Int,
     var insets: Insets,
     var bitmap: Bitmap?,
+    var displayId: Int,
     /** App-provided URL representing the content the user was looking at in the screenshot. */
     var contextUrl: Uri? = null,
 ) {
@@ -31,22 +33,31 @@ data class ScreenshotData(
 
     companion object {
         @JvmStatic
-        fun fromRequest(request: ScreenshotRequest): ScreenshotData {
-            return ScreenshotData(
-                request.type,
-                request.source,
-                if (request.userId >= 0) UserHandle.of(request.userId) else null,
-                request.topComponent,
-                request.boundsInScreen,
-                request.taskId,
-                request.insets,
-                request.bitmap,
+        fun fromRequest(request: ScreenshotRequest, displayId: Int = Display.DEFAULT_DISPLAY) =
+            ScreenshotData(
+                type = request.type,
+                source = request.source,
+                userHandle = if (request.userId >= 0) UserHandle.of(request.userId) else null,
+                topComponent = request.topComponent,
+                screenBounds = request.boundsInScreen,
+                taskId = request.taskId,
+                insets = request.insets,
+                bitmap = request.bitmap,
+                displayId = displayId,
             )
-        }
 
         @VisibleForTesting
-        fun forTesting(): ScreenshotData {
-            return ScreenshotData(0, 0, null, null, null, 0, Insets.NONE, null)
-        }
+        fun forTesting() =
+            ScreenshotData(
+                type = 0,
+                source = 0,
+                userHandle = null,
+                topComponent = null,
+                screenBounds = null,
+                taskId = 0,
+                insets = Insets.NONE,
+                bitmap = null,
+                displayId = Display.DEFAULT_DISPLAY,
+            )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
new file mode 100644
index 000000000000..6c886fcd9c8f
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -0,0 +1,201 @@
+package com.android.systemui.screenshot
+
+import android.net.Uri
+import android.os.Trace
+import android.util.Log
+import android.view.Display
+import com.android.internal.logging.UiEventLogger
+import com.android.internal.util.ScreenshotRequest
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.display.data.repository.DisplayRepository
+import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
+import java.util.function.Consumer
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+/**
+ * Receives the signal to take a screenshot from [TakeScreenshotService], and calls back with the
+ * result.
+ *
+ * Captures a screenshot for each [Display] available.
+ */
+@SysUISingleton
+class TakeScreenshotExecutor
+@Inject
+constructor(
+    private val screenshotControllerFactory: ScreenshotController.Factory,
+    displayRepository: DisplayRepository,
+    @Application private val mainScope: CoroutineScope,
+    private val screenshotRequestProcessor: ScreenshotRequestProcessor,
+    private val uiEventLogger: UiEventLogger
+) {
+
+    private lateinit var displays: StateFlow<Set<Display>>
+    private val displaysCollectionJob: Job =
+        mainScope.launch {
+            displays = displayRepository.displays.stateIn(this, SharingStarted.Eagerly, emptySet())
+        }
+
+    private val screenshotControllers = mutableMapOf<Int, ScreenshotController>()
+
+    /**
+     * Executes the [ScreenshotRequest].
+     *
+     * [onSaved] is invoked only on the default display result. [RequestCallback.onFinish] is
+     * invoked only when both screenshot UIs are removed.
+     */
+    suspend fun executeScreenshots(
+        screenshotRequest: ScreenshotRequest,
+        onSaved: (Uri) -> Unit,
+        requestCallback: RequestCallback
+    ) {
+        val displayIds = getDisplaysToScreenshot()
+        val resultCallbackWrapper = MultiResultCallbackWrapper(requestCallback)
+        screenshotRequest.oneForEachDisplay(displayIds).forEach { screenshotData: ScreenshotData ->
+            dispatchToController(
+                screenshotData = screenshotData,
+                onSaved =
+                    if (screenshotData.displayId == Display.DEFAULT_DISPLAY) onSaved else { _ -> },
+                callback = resultCallbackWrapper.createCallbackForId(screenshotData.displayId)
+            )
+        }
+    }
+
+    /** Creates a [ScreenshotData] for each display. */
+    private suspend fun ScreenshotRequest.oneForEachDisplay(
+        displayIds: List<Int>
+    ): List<ScreenshotData> {
+        return displayIds
+            .map { displayId -> ScreenshotData.fromRequest(this, displayId) }
+            .map { screenshotData: ScreenshotData ->
+                screenshotRequestProcessor.process(screenshotData)
+            }
+    }
+
+    private fun dispatchToController(
+        screenshotData: ScreenshotData,
+        onSaved: (Uri) -> Unit,
+        callback: RequestCallback
+    ) {
+        uiEventLogger.log(
+            ScreenshotEvent.getScreenshotSource(screenshotData.source),
+            0,
+            screenshotData.packageNameString
+        )
+        Log.d(TAG, "Screenshot request: $screenshotData")
+        getScreenshotController(screenshotData.displayId)
+            .handleScreenshot(screenshotData, onSaved, callback)
+    }
+
+    private fun getDisplaysToScreenshot(): List<Int> {
+        return displays.value.filter { it.type in ALLOWED_DISPLAY_TYPES }.map { it.displayId }
+    }
+
+    /**
+     * Propagates the close system dialog signal to all controllers.
+     *
+     * TODO(b/295143676): Move the receiver in this class once the flag is flipped.
+     */
+    fun onCloseSystemDialogsReceived() {
+        screenshotControllers.forEach { (_, screenshotController) ->
+            if (!screenshotController.isPendingSharedTransition) {
+                screenshotController.dismissScreenshot(ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER)
+            }
+        }
+    }
+
+    /** Removes all screenshot related windows. */
+    fun removeWindows() {
+        screenshotControllers.forEach { (_, screenshotController) ->
+            screenshotController.removeWindow()
+        }
+    }
+
+    /**
+     * Destroys the executor. Afterwards, this class is not expected to work as intended anymore.
+     */
+    fun onDestroy() {
+        screenshotControllers.forEach { (_, screenshotController) ->
+            screenshotController.onDestroy()
+        }
+        screenshotControllers.clear()
+        displaysCollectionJob.cancel()
+    }
+
+    private fun getScreenshotController(id: Int): ScreenshotController {
+        return screenshotControllers.computeIfAbsent(id) { screenshotControllerFactory.create(id) }
+    }
+
+    /** For java compatibility only. see [executeScreenshots] */
+    fun executeScreenshotsAsync(
+        screenshotRequest: ScreenshotRequest,
+        onSaved: Consumer<Uri>,
+        requestCallback: RequestCallback
+    ) {
+        mainScope.launch {
+            executeScreenshots(screenshotRequest, { uri -> onSaved.accept(uri) }, requestCallback)
+        }
+    }
+
+    /**
+     * Returns a [RequestCallback] that calls [RequestCallback.onFinish] only when all callbacks for
+     * id created have finished.
+     *
+     * If any callback created calls [reportError], then following [onFinish] are not considered.
+     */
+    private class MultiResultCallbackWrapper(
+        private val originalCallback: RequestCallback,
+    ) {
+        private val idsPending = mutableSetOf<Int>()
+        private var errorReported = false
+
+        /**
+         * Creates a callback for [id].
+         *
+         * [originalCallback]'s [onFinish] will be called only when this (and the other created)
+         * callback's [onFinish] have been called.
+         */
+        fun createCallbackForId(id: Int): RequestCallback {
+            Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, TAG, "Waiting for id=$id", id)
+            idsPending += id
+            return object : RequestCallback {
+                override fun reportError() {
+                    Log.d(TAG, "ReportError id=$id")
+                    Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, TAG, id)
+                    Trace.instantForTrack(Trace.TRACE_TAG_APP, TAG, "reportError id=$id")
+                    originalCallback.reportError()
+                    errorReported = true
+                }
+
+                override fun onFinish() {
+                    Log.d(TAG, "onFinish id=$id")
+                    if (errorReported) return
+                    idsPending -= id
+                    Trace.instantForTrack(Trace.TRACE_TAG_APP, TAG, "onFinish id=$id")
+                    if (idsPending.isEmpty()) {
+                        Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, TAG, id)
+                        originalCallback.onFinish()
+                    }
+                }
+            }
+        }
+    }
+
+    private companion object {
+        val TAG = LogConfig.logTag(TakeScreenshotService::class.java)
+
+        val ALLOWED_DISPLAY_TYPES =
+            listOf(
+                Display.TYPE_EXTERNAL,
+                Display.TYPE_INTERNAL,
+                Display.TYPE_OVERLAY,
+                Display.TYPE_WIFI
+            )
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
index 1cdad83fb0aa..1e8542fe8f0c 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
@@ -21,6 +21,7 @@ import static android.content.Intent.ACTION_CLOSE_SYSTEM_DIALOGS;
 
 import static com.android.internal.util.ScreenshotHelper.SCREENSHOT_MSG_PROCESS_COMPLETE;
 import static com.android.internal.util.ScreenshotHelper.SCREENSHOT_MSG_URI;
+import static com.android.systemui.flags.Flags.MULTI_DISPLAY_SCREENSHOT;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_CALLBACK;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_DISMISS;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_SERVICE;
@@ -46,6 +47,7 @@ import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.util.Log;
+import android.view.Display;
 import android.widget.Toast;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -59,6 +61,7 @@ import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
 import javax.inject.Inject;
+import javax.inject.Provider;
 
 public class TakeScreenshotService extends Service {
     private static final String TAG = logTag(TakeScreenshotService.class);
@@ -82,12 +85,17 @@ public class TakeScreenshotService extends Service {
                 if (DEBUG_DISMISS) {
                     Log.d(TAG, "Received ACTION_CLOSE_SYSTEM_DIALOGS");
                 }
-                if (!mScreenshot.isPendingSharedTransition()) {
+                if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+                    // TODO(b/295143676): move receiver inside executor when the flag is enabled.
+                    mTakeScreenshotExecutor.get().onCloseSystemDialogsReceived();
+                } else if (!mScreenshot.isPendingSharedTransition()) {
                     mScreenshot.dismissScreenshot(SCREENSHOT_DISMISSED_OTHER);
                 }
             }
         }
     };
+    private final Provider<TakeScreenshotExecutor> mTakeScreenshotExecutor;
+
 
     /** Informs about coarse grained state of the Controller. */
     public interface RequestCallback {
@@ -99,16 +107,15 @@ public class TakeScreenshotService extends Service {
     }
 
     @Inject
-    public TakeScreenshotService(ScreenshotController screenshotController, UserManager userManager,
-            DevicePolicyManager devicePolicyManager, UiEventLogger uiEventLogger,
-            ScreenshotNotificationsController notificationsController, Context context,
-            @Background Executor bgExecutor, FeatureFlags featureFlags,
-            RequestProcessor processor) {
+    public TakeScreenshotService(ScreenshotController.Factory screenshotControllerFactory,
+            UserManager userManager, DevicePolicyManager devicePolicyManager,
+            UiEventLogger uiEventLogger, ScreenshotNotificationsController notificationsController,
+            Context context, @Background Executor bgExecutor, FeatureFlags featureFlags,
+            RequestProcessor processor, Provider<TakeScreenshotExecutor> takeScreenshotExecutor) {
         if (DEBUG_SERVICE) {
             Log.d(TAG, "new " + this);
         }
         mHandler = new Handler(Looper.getMainLooper(), this::handleMessage);
-        mScreenshot = screenshotController;
         mUserManager = userManager;
         mDevicePolicyManager = devicePolicyManager;
         mUiEventLogger = uiEventLogger;
@@ -117,6 +124,12 @@ public class TakeScreenshotService extends Service {
         mBgExecutor = bgExecutor;
         mFeatureFlags = featureFlags;
         mProcessor = processor;
+        mTakeScreenshotExecutor = takeScreenshotExecutor;
+        if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+            mScreenshot = null;
+        } else {
+            mScreenshot = screenshotControllerFactory.create(Display.DEFAULT_DISPLAY);
+        }
     }
 
     @Override
@@ -142,7 +155,11 @@ public class TakeScreenshotService extends Service {
         if (DEBUG_SERVICE) {
             Log.d(TAG, "onUnbind");
         }
-        mScreenshot.removeWindow();
+        if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+            mTakeScreenshotExecutor.get().removeWindows();
+        } else {
+            mScreenshot.removeWindow();
+        }
         unregisterReceiver(mCloseSystemDialogs);
         return false;
     }
@@ -150,7 +167,11 @@ public class TakeScreenshotService extends Service {
     @Override
     public void onDestroy() {
         super.onDestroy();
-        mScreenshot.onDestroy();
+        if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+            mTakeScreenshotExecutor.get().onDestroy();
+        } else {
+            mScreenshot.onDestroy();
+        }
         if (DEBUG_SERVICE) {
             Log.d(TAG, "onDestroy");
         }
@@ -218,10 +239,17 @@ public class TakeScreenshotService extends Service {
         }
 
         Log.d(TAG, "Processing screenshot data");
-        ScreenshotData screenshotData = ScreenshotData.fromRequest(request);
+
+
+        ScreenshotData screenshotData = ScreenshotData.fromRequest(
+                request, Display.DEFAULT_DISPLAY);
         try {
-            mProcessor.processAsync(screenshotData,
-                    (data) -> dispatchToController(data, onSaved, callback));
+            if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+                mTakeScreenshotExecutor.get().executeScreenshotsAsync(request, onSaved, callback);
+            } else {
+                mProcessor.processAsync(screenshotData, (data) ->
+                        dispatchToController(data, onSaved, callback));
+            }
         } catch (IllegalStateException e) {
             Log.e(TAG, "Failed to process screenshot request!", e);
             logFailedRequest(request);
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
index 22e238c0f2ad..7d17d4c72b76 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
@@ -20,9 +20,11 @@ import android.app.Service;
 
 import com.android.systemui.screenshot.ImageCapture;
 import com.android.systemui.screenshot.ImageCaptureImpl;
+import com.android.systemui.screenshot.RequestProcessor;
 import com.android.systemui.screenshot.ScreenshotPolicy;
 import com.android.systemui.screenshot.ScreenshotPolicyImpl;
 import com.android.systemui.screenshot.ScreenshotProxyService;
+import com.android.systemui.screenshot.ScreenshotRequestProcessor;
 import com.android.systemui.screenshot.TakeScreenshotService;
 import com.android.systemui.screenshot.appclips.AppClipsScreenshotHelperService;
 import com.android.systemui.screenshot.appclips.AppClipsService;
@@ -63,4 +65,8 @@ public abstract class ScreenshotModule {
     @IntoMap
     @ClassKey(AppClipsService.class)
     abstract Service bindAppClipsService(AppClipsService service);
+
+    @Binds
+    abstract ScreenshotRequestProcessor bindScreenshotRequestProcessor(
+            RequestProcessor requestProcessor);
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
index 43e99393b874..f8a8a6830669 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
@@ -20,6 +20,7 @@ import android.content.ComponentName
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
+import android.view.Display
 import android.view.WindowManager
 import com.android.internal.util.ScreenshotRequest
 import com.google.common.truth.Truth.assertThat
@@ -54,6 +55,16 @@ class ScreenshotDataTest {
         assertThat(data.taskId).isEqualTo(taskId)
         assertThat(data.userHandle).isEqualTo(UserHandle.of(userId))
         assertThat(data.topComponent).isEqualTo(component)
+        assertThat(data.displayId).isEqualTo(Display.DEFAULT_DISPLAY)
+    }
+
+    @Test
+    fun testConstruction_notDefaultDisplayId() {
+        val request = ScreenshotRequest.Builder(type, source).build()
+
+        val data = ScreenshotData.fromRequest(request, displayId = 42)
+
+        assertThat(data.displayId).isEqualTo(42)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
new file mode 100644
index 000000000000..97c2ed45c26d
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -0,0 +1,303 @@
+package com.android.systemui.screenshot
+
+import android.content.ComponentName
+import android.net.Uri
+import android.testing.AndroidTestingRunner
+import android.view.Display
+import android.view.Display.TYPE_EXTERNAL
+import android.view.Display.TYPE_INTERNAL
+import android.view.Display.TYPE_OVERLAY
+import android.view.Display.TYPE_VIRTUAL
+import android.view.Display.TYPE_WIFI
+import android.view.WindowManager
+import androidx.test.filters.SmallTest
+import com.android.internal.logging.testing.UiEventLoggerFake
+import com.android.internal.util.ScreenshotRequest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.display.data.repository.FakeDisplayRepository
+import com.android.systemui.display.data.repository.display
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.kotlinArgumentCaptor as ArgumentCaptor
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.nullable
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyNoMoreInteractions
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class TakeScreenshotExecutorTest : SysuiTestCase() {
+
+    private val controller0 = mock<ScreenshotController>()
+    private val controller1 = mock<ScreenshotController>()
+    private val controllerFactory = mock<ScreenshotController.Factory>()
+    private val callback = mock<TakeScreenshotService.RequestCallback>()
+
+    private val fakeDisplayRepository = FakeDisplayRepository()
+    private val requestProcessor = FakeRequestProcessor()
+    private val topComponent = ComponentName(mContext, TakeScreenshotExecutorTest::class.java)
+    private val testScope = TestScope(UnconfinedTestDispatcher())
+    private val eventLogger = UiEventLoggerFake()
+
+    private val screenshotExecutor =
+        TakeScreenshotExecutor(
+            controllerFactory,
+            fakeDisplayRepository,
+            testScope,
+            requestProcessor,
+            eventLogger,
+        )
+
+    @Before
+    fun setUp() {
+        whenever(controllerFactory.create(eq(0))).thenReturn(controller0)
+        whenever(controllerFactory.create(eq(1))).thenReturn(controller1)
+    }
+
+    @Test
+    fun executeScreenshots_severalDisplays_callsControllerForEachOne() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verify(controllerFactory).create(eq(0))
+            verify(controllerFactory).create(eq(1))
+
+            val capturer = ArgumentCaptor<ScreenshotData>()
+
+            verify(controller0).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value.displayId).isEqualTo(0)
+            // OnSaved callback should be different.
+            verify(controller1).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value.displayId).isEqualTo(1)
+
+            assertThat(eventLogger.numLogs()).isEqualTo(2)
+            assertThat(eventLogger.get(0).eventId)
+                .isEqualTo(ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER.id)
+            assertThat(eventLogger.get(0).packageName).isEqualTo(topComponent.packageName)
+            assertThat(eventLogger.get(1).eventId)
+                .isEqualTo(ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER.id)
+            assertThat(eventLogger.get(1).packageName).isEqualTo(topComponent.packageName)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_onlyVirtualDisplays_noInteractionsWithControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_VIRTUAL, id = 0), display(TYPE_VIRTUAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verifyNoMoreInteractions(controllerFactory)
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_allowedTypes_allCaptured() =
+        testScope.runTest {
+            whenever(controllerFactory.create(any())).thenReturn(controller0)
+
+            setDisplays(
+                display(TYPE_INTERNAL, id = 0),
+                display(TYPE_EXTERNAL, id = 1),
+                display(TYPE_OVERLAY, id = 2),
+                display(TYPE_WIFI, id = 3)
+            )
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verify(controller0, times(4)).handleScreenshot(any(), any(), any())
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_reportsOnFinishedOnlyWhenBothFinished() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+            val capturer1 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+
+            verify(controller0).handleScreenshot(any(), any(), capturer0.capture())
+            verify(controller1).handleScreenshot(any(), any(), capturer1.capture())
+
+            verify(callback, never()).onFinish()
+
+            capturer0.value.onFinish()
+
+            verify(callback, never()).onFinish()
+
+            capturer1.value.onFinish()
+
+            verify(callback).onFinish()
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_doesNotReportFinishedIfOneFinishesOtherFails() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+            val capturer1 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+
+            verify(controller0).handleScreenshot(any(), any(), capturer0.capture())
+            verify(controller1).handleScreenshot(any(), nullable(), capturer1.capture())
+
+            verify(callback, never()).onFinish()
+
+            capturer0.value.onFinish()
+
+            verify(callback, never()).onFinish()
+
+            capturer1.value.reportError()
+
+            verify(callback, never()).onFinish()
+            verify(callback).reportError()
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_doesNotReportFinishedAfterOneFails() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+            val capturer1 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+
+            verify(controller0).handleScreenshot(any(), any(), capturer0.capture())
+            verify(controller1).handleScreenshot(any(), any(), capturer1.capture())
+
+            verify(callback, never()).onFinish()
+
+            capturer0.value.reportError()
+
+            verify(callback, never()).onFinish()
+            verify(callback).reportError()
+
+            capturer1.value.onFinish()
+
+            verify(callback, never()).onFinish()
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun onDestroy_propagatedToControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.onDestroy()
+            verify(controller0).onDestroy()
+            verify(controller1).onDestroy()
+        }
+
+    @Test
+    fun removeWindows_propagatedToControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.removeWindows()
+            verify(controller0).removeWindow()
+            verify(controller1).removeWindow()
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun onCloseSystemDialogsReceived_propagatedToControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.onCloseSystemDialogsReceived()
+            verify(controller0).dismissScreenshot(any())
+            verify(controller1).dismissScreenshot(any())
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun onCloseSystemDialogsReceived_someControllerHavePendingTransitions() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            whenever(controller0.isPendingSharedTransition).thenReturn(true)
+            whenever(controller1.isPendingSharedTransition).thenReturn(false)
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.onCloseSystemDialogsReceived()
+            verify(controller0, never()).dismissScreenshot(any())
+            verify(controller1).dismissScreenshot(any())
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_controllerCalledWithRequestProcessorReturnValue() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0))
+            val screenshotRequest = createScreenshotRequest()
+            val toBeReturnedByProcessor = ScreenshotData.forTesting()
+            requestProcessor.toReturn = toBeReturnedByProcessor
+
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(screenshotRequest, onSaved, callback)
+
+            assertThat(requestProcessor.processed)
+                .isEqualTo(ScreenshotData.fromRequest(screenshotRequest))
+
+            val capturer = ArgumentCaptor<ScreenshotData>()
+            verify(controller0).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value).isEqualTo(toBeReturnedByProcessor)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    private suspend fun TestScope.setDisplays(vararg displays: Display) {
+        fakeDisplayRepository.emit(displays.toSet())
+        runCurrent()
+    }
+
+    private fun createScreenshotRequest() =
+        ScreenshotRequest.Builder(
+                WindowManager.TAKE_SCREENSHOT_FULLSCREEN,
+                WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
+            )
+            .setTopComponent(topComponent)
+            .build()
+
+    private class FakeRequestProcessor : ScreenshotRequestProcessor {
+        var processed: ScreenshotData? = null
+        var toReturn: ScreenshotData? = null
+
+        override suspend fun process(screenshot: ScreenshotData): ScreenshotData {
+            processed = screenshot
+            return toReturn ?: screenshot
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
index 77f742647497..a08cda6ff158 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
@@ -27,6 +27,7 @@ import android.hardware.HardwareBuffer
 import android.os.UserHandle
 import android.os.UserManager
 import android.testing.AndroidTestingRunner
+import android.view.Display
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import androidx.test.filters.SmallTest
@@ -34,6 +35,7 @@ import com.android.internal.logging.testing.UiEventLoggerFake
 import com.android.internal.util.ScreenshotRequest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags.MULTI_DISPLAY_SCREENSHOT
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_CAPTURE_FAILED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER
 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
@@ -48,6 +50,7 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.isNull
+import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.doThrow
 import org.mockito.Mockito.times
@@ -60,6 +63,8 @@ class TakeScreenshotServiceTest : SysuiTestCase() {
 
     private val application = mock<Application>()
     private val controller = mock<ScreenshotController>()
+    private val controllerFactory = mock<ScreenshotController.Factory>()
+    private val takeScreenshotExecutor = mock<TakeScreenshotExecutor>()
     private val userManager = mock<UserManager>()
     private val requestProcessor = mock<RequestProcessor>()
     private val devicePolicyManager = mock<DevicePolicyManager>()
@@ -71,21 +76,11 @@ class TakeScreenshotServiceTest : SysuiTestCase() {
     private val flags = FakeFeatureFlags()
     private val topComponent = ComponentName(mContext, TakeScreenshotServiceTest::class.java)
 
-    private val service =
-        TakeScreenshotService(
-            controller,
-            userManager,
-            devicePolicyManager,
-            eventLogger,
-            notificationsController,
-            mContext,
-            Runnable::run,
-            flags,
-            requestProcessor
-        )
+    private lateinit var service: TakeScreenshotService
 
     @Before
     fun setUp() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, false)
         whenever(devicePolicyManager.resources).thenReturn(devicePolicyResourcesManager)
         whenever(
                 devicePolicyManager.getScreenCaptureDisabled(
@@ -95,6 +90,7 @@ class TakeScreenshotServiceTest : SysuiTestCase() {
             )
             .thenReturn(false)
         whenever(userManager.isUserUnlocked).thenReturn(true)
+        whenever(controllerFactory.create(any())).thenReturn(controller)
 
         // Stub request processor as a synchronous no-op for tests with the flag enabled
         doAnswer {
@@ -113,14 +109,7 @@ class TakeScreenshotServiceTest : SysuiTestCase() {
             .whenever(requestProcessor)
             .processAsync(/* screenshot= */ any(ScreenshotData::class.java), /* callback= */ any())
 
-        service.attach(
-            mContext,
-            /* thread = */ null,
-            /* className = */ null,
-            /* token = */ null,
-            application,
-            /* activityManager = */ null
-        )
+        service = createService()
     }
 
     @Test
@@ -146,7 +135,7 @@ class TakeScreenshotServiceTest : SysuiTestCase() {
 
         verify(controller, times(1))
             .handleScreenshot(
-                eq(ScreenshotData.fromRequest(request)),
+                eq(ScreenshotData.fromRequest(request, Display.DEFAULT_DISPLAY)),
                 /* onSavedListener = */ any(),
                 /* requestCallback = */ any()
             )
@@ -295,6 +284,74 @@ class TakeScreenshotServiceTest : SysuiTestCase() {
             failureEvent.packageName
         )
     }
+
+    @Test
+    fun takeScreenshotFullScreen_multiDisplayFlagEnabled_takeScreenshotExecutor() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, true)
+        service = createService()
+
+        val request =
+            ScreenshotRequest.Builder(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_OTHER)
+                .setTopComponent(topComponent)
+                .build()
+
+        service.handleRequest(request, { /* onSaved */}, callback)
+
+        verifyZeroInteractions(controller)
+        verify(takeScreenshotExecutor, times(1)).executeScreenshotsAsync(any(), any(), any())
+
+        assertEquals("Expected one UiEvent", 0, eventLogger.numLogs())
+    }
+
+    @Test
+    fun testServiceLifecycle_multiDisplayScreenshotFlagEnabled() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, true)
+        service = createService()
+
+        service.onCreate()
+        service.onBind(null /* unused: Intent */)
+
+        service.onUnbind(null /* unused: Intent */)
+        verify(takeScreenshotExecutor, times(1)).removeWindows()
+
+        service.onDestroy()
+        verify(takeScreenshotExecutor, times(1)).onDestroy()
+    }
+
+    @Test
+    fun constructor_MultiDisplayFlagOn_screenshotControllerNotCreated() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, true)
+        clearInvocations(controllerFactory)
+
+        service = createService()
+
+        verifyZeroInteractions(controllerFactory)
+    }
+
+    private fun createService(): TakeScreenshotService {
+        val service =
+            TakeScreenshotService(
+                controllerFactory,
+                userManager,
+                devicePolicyManager,
+                eventLogger,
+                notificationsController,
+                mContext,
+                Runnable::run,
+                flags,
+                requestProcessor,
+                { takeScreenshotExecutor },
+            )
+        service.attach(
+            mContext,
+            /* thread = */ null,
+            /* className = */ null,
+            /* token = */ null,
+            application,
+            /* activityManager = */ null
+        )
+        return service
+    }
 }
 
 private fun Bitmap.equalsHardwareBitmap(other: Bitmap): Boolean {