Merge "Create a Feature Flag to allow enabling early device release"
diff --git a/src/com/android/tradefed/command/CommandOptions.java b/src/com/android/tradefed/command/CommandOptions.java
index bf6654f..e397adb 100644
--- a/src/com/android/tradefed/command/CommandOptions.java
+++ b/src/com/android/tradefed/command/CommandOptions.java
@@ -196,6 +196,11 @@
     )
     private String mHostLogSuffix = null;
 
+    @Option(
+            name = "early-device-release",
+            description = "Feature flag to release the device as soon as done with it.")
+    private boolean mEnableEarlyDeviceRelease = false;
+
     /**
      * Set the help mode for the config.
      * <p/>
@@ -500,4 +505,10 @@
     public int getExtraRemotePostsubmitInstance() {
         return mExtraRemoteInstancePostsubmit;
     }
+
+    /** {@inheritDoc} */
+    @Override
+    public boolean earlyDeviceRelease() {
+        return mEnableEarlyDeviceRelease;
+    }
 }
diff --git a/src/com/android/tradefed/command/CommandScheduler.java b/src/com/android/tradefed/command/CommandScheduler.java
index b2a1657..711ca17 100644
--- a/src/com/android/tradefed/command/CommandScheduler.java
+++ b/src/com/android/tradefed/command/CommandScheduler.java
@@ -465,6 +465,7 @@
             IScheduledInvocationListener {
 
         private final IDeviceManager mDeviceManager;
+        private boolean mDeviceReleased = false;
 
         FreeDeviceHandler(IDeviceManager deviceManager,
                 IScheduledInvocationListener... listeners) {
@@ -473,12 +474,11 @@
         }
 
         @Override
-        public void invocationComplete(IInvocationContext context,
-                Map<ITestDevice, FreeDeviceState> devicesStates) {
-            for (ITestInvocationListener listener : getListeners()) {
-                ((IScheduledInvocationListener) listener).invocationComplete(context, devicesStates);
+        public void releaseDevices(
+                IInvocationContext context, Map<ITestDevice, FreeDeviceState> devicesStates) {
+            if (mDeviceReleased) {
+                return;
             }
-
             for (ITestDevice device : context.getDevices()) {
                 mDeviceManager.freeDevice(device, devicesStates.get(device));
                 remoteFreeDevice(device);
@@ -487,6 +487,17 @@
                     ((IManagedTestDevice)device).setFastbootPath(mDeviceManager.getFastbootPath());
                 }
             }
+            mDeviceReleased = true;
+        }
+
+        @Override
+        public void invocationComplete(
+                IInvocationContext context, Map<ITestDevice, FreeDeviceState> devicesStates) {
+            for (ITestInvocationListener listener : getListeners()) {
+                ((IScheduledInvocationListener) listener)
+                        .invocationComplete(context, devicesStates);
+            }
+            releaseDevices(context, devicesStates);
         }
     }
 
@@ -516,11 +527,7 @@
     }
 
     private class InvocationThread extends Thread {
-        /**
-         * time to wait for device adb shell responsive connection before declaring it unavailable
-         * for the next iteration of testing.
-         */
-        private static final long CHECK_WAIT_DEVICE_AVAIL_MS = 30 * 1000;
+
         private static final int EXPECTED_THREAD_COUNT = 1;
         private static final String INVOC_END_EVENT_ID_KEY = "id";
         private static final String INVOC_END_EVENT_ELAPSED_KEY = "elapsed-time";
@@ -554,10 +561,6 @@
 
         @Override
         public void run() {
-            Map<ITestDevice, FreeDeviceState> deviceStates = new HashMap<>();
-            for (ITestDevice device : mInvocationContext.getDevices()) {
-                deviceStates.put(device, FreeDeviceState.AVAILABLE);
-            }
             mStartTime = System.currentTimeMillis();
             ITestInvocation instance = getInvocation();
             IConfiguration config = mCmd.getConfiguration();
@@ -571,6 +574,7 @@
                 }
             }
 
+            Exception trackDeviceException = null;
             try {
                 // Copy the command options invocation attributes to the invocation if it has not
                 // been already done.
@@ -589,17 +593,11 @@
                         new Rescheduler(mCmd.getCommandTracker()), mListeners);
             } catch (DeviceUnresponsiveException e) {
                 CLog.w("Device %s is unresponsive. Reason: %s", e.getSerial(), e.getMessage());
-                ITestDevice badDevice = mInvocationContext.getDeviceBySerial(e.getSerial());
-                if (badDevice != null) {
-                    deviceStates.put(badDevice, FreeDeviceState.UNRESPONSIVE);
-                }
+                trackDeviceException = e;
                 setLastInvocationExitCode(ExitCode.DEVICE_UNRESPONSIVE, e);
             } catch (DeviceNotAvailableException e) {
                 CLog.w("Device %s is not available. Reason: %s", e.getSerial(), e.getMessage());
-                ITestDevice badDevice = mInvocationContext.getDeviceBySerial(e.getSerial());
-                if (badDevice != null) {
-                    deviceStates.put(badDevice, FreeDeviceState.UNAVAILABLE);
-                }
+                trackDeviceException = e;
                 setLastInvocationExitCode(ExitCode.DEVICE_UNAVAILABLE, e);
             } catch (FatalHostError e) {
                 CLog.wtf(String.format("Fatal error occurred: %s, shutting down", e.getMessage()),
@@ -620,24 +618,11 @@
                 // remove invocation thread first so another invocation can be started on device
                 // when freed
                 removeInvocationThread(this);
-                for (ITestDevice device : mInvocationContext.getDevices()) {
-                    if (device.getIDevice() instanceof StubDevice) {
-                        // Never release stub and Tcp devices, otherwise they will disappear
-                        // during deallocation since they are only placeholder.
-                        deviceStates.put(device, FreeDeviceState.AVAILABLE);
-                    } else if (!TestDeviceState.ONLINE.equals(device.getDeviceState())) {
-                        // If the device is offline at the end of the test
-                        deviceStates.put(device, FreeDeviceState.UNAVAILABLE);
-                    } else if (!isDeviceResponsive(device)) {
-                        // If device cannot pass basic shell responsiveness test.
-                        deviceStates.put(device, FreeDeviceState.UNAVAILABLE);
-                    }
-                    // Reset the recovery mode at the end of the invocation.
-                    device.setRecoveryMode(RecoveryMode.AVAILABLE);
-                }
 
                 checkStrayThreads();
 
+                Map<ITestDevice, FreeDeviceState> deviceStates =
+                        createReleaseMap(mInvocationContext, trackDeviceException);
                 for (final IScheduledInvocationListener listener : mListeners) {
                     try {
                         listener.invocationComplete(mInvocationContext, deviceStates);
@@ -697,11 +682,6 @@
             logEvent(EventType.INVOCATION_END, args);
         }
 
-        /** Basic responsiveness check at the end of an invocation. */
-        private boolean isDeviceResponsive(ITestDevice device) {
-            return device.waitForDeviceShell(CHECK_WAIT_DEVICE_AVAIL_MS);
-        }
-
         ITestInvocation getInvocation() {
             return mInvocation;
         }
@@ -800,6 +780,48 @@
         }
     }
 
+    /** Create a map of the devices state so they can be released appropriately. */
+    public static Map<ITestDevice, FreeDeviceState> createReleaseMap(
+            IInvocationContext context, Throwable e) {
+        Map<ITestDevice, FreeDeviceState> deviceStates = new HashMap<>();
+        for (ITestDevice device : context.getDevices()) {
+            deviceStates.put(device, FreeDeviceState.AVAILABLE);
+        }
+
+        if (e != null) {
+            if (e instanceof DeviceUnresponsiveException) {
+                ITestDevice badDevice =
+                        context.getDeviceBySerial(((DeviceUnresponsiveException) e).getSerial());
+                if (badDevice != null) {
+                    deviceStates.put(badDevice, FreeDeviceState.UNRESPONSIVE);
+                }
+            } else if (e instanceof DeviceNotAvailableException) {
+                ITestDevice badDevice =
+                        context.getDeviceBySerial(((DeviceNotAvailableException) e).getSerial());
+                if (badDevice != null) {
+                    deviceStates.put(badDevice, FreeDeviceState.UNAVAILABLE);
+                }
+            }
+        }
+
+        for (ITestDevice device : context.getDevices()) {
+            if (device.getIDevice() instanceof StubDevice) {
+                // Never release stub and Tcp devices, otherwise they will disappear
+                // during deallocation since they are only placeholder.
+                deviceStates.put(device, FreeDeviceState.AVAILABLE);
+            } else if (!TestDeviceState.ONLINE.equals(device.getDeviceState())) {
+                // If the device is offline at the end of the test
+                deviceStates.put(device, FreeDeviceState.UNAVAILABLE);
+            } else if (!device.waitForDeviceShell(30000)) {
+                // If device cannot pass basic shell responsiveness test.
+                deviceStates.put(device, FreeDeviceState.UNAVAILABLE);
+            }
+            // Reset the recovery mode at the end of the invocation.
+            device.setRecoveryMode(RecoveryMode.AVAILABLE);
+        }
+        return deviceStates;
+    }
+
     /**
      * A {@link IDeviceMonitor} that signals scheduler to process commands when an available device
      * is added.
diff --git a/src/com/android/tradefed/command/ICommandOptions.java b/src/com/android/tradefed/command/ICommandOptions.java
index 28f12a8..4064904 100644
--- a/src/com/android/tradefed/command/ICommandOptions.java
+++ b/src/com/android/tradefed/command/ICommandOptions.java
@@ -189,4 +189,7 @@
 
     /** Whether or not to start extra instances in the remote VM in postsubmit. */
     public int getExtraRemotePostsubmitInstance();
+
+    /** Whether or not to release the device early when done with it. */
+    public boolean earlyDeviceRelease();
 }
diff --git a/src/com/android/tradefed/command/ICommandScheduler.java b/src/com/android/tradefed/command/ICommandScheduler.java
index 21a0559..3169982 100644
--- a/src/com/android/tradefed/command/ICommandScheduler.java
+++ b/src/com/android/tradefed/command/ICommandScheduler.java
@@ -49,6 +49,16 @@
         public default void invocationInitiated(IInvocationContext context) {}
 
         /**
+         * Callback associated with {@link ICommandOptions#earlyDeviceRelease()} to release the
+         * devices when done with them.
+         *
+         * @param context
+         * @param devicesStates
+         */
+        public default void releaseDevices(
+                IInvocationContext context, Map<ITestDevice, FreeDeviceState> devicesStates) {}
+
+        /**
          * Callback when entire invocation has completed, including all {@link
          * ITestInvocationListener#invocationEnded(long)} events.
          *
diff --git a/src/com/android/tradefed/invoker/TestInvocation.java b/src/com/android/tradefed/invoker/TestInvocation.java
index 61754c9..6f617f8 100644
--- a/src/com/android/tradefed/invoker/TestInvocation.java
+++ b/src/com/android/tradefed/invoker/TestInvocation.java
@@ -20,11 +20,14 @@
 import com.android.tradefed.build.BuildRetrievalError;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.command.CommandRunner.ExitCode;
+import com.android.tradefed.command.CommandScheduler;
+import com.android.tradefed.command.ICommandScheduler.IScheduledInvocationListener;
 import com.android.tradefed.config.ConfigurationException;
 import com.android.tradefed.config.GlobalConfiguration;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceUnresponsiveException;
+import com.android.tradefed.device.FreeDeviceState;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.ITestDevice.RecoveryMode;
 import com.android.tradefed.device.NativeDevice;
@@ -82,6 +85,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Map.Entry;
 
 /**
@@ -145,6 +149,7 @@
     private String mStopCause = null;
     private Long mStopRequestTime = null;
     private boolean mTestStarted = false;
+    private List<IScheduledInvocationListener> mSchedulerListeners = new ArrayList<>();
 
     /**
      * A {@link ResultForwarder} for forwarding resumed invocations.
@@ -359,6 +364,13 @@
             // Track the timestamp when we are done with devices
             InvocationMetricLogger.addInvocationMetrics(
                     InvocationMetricKey.DEVICE_DONE_TIMESTAMP, System.currentTimeMillis());
+            if (config.getCommandOptions().earlyDeviceRelease()) {
+                Map<ITestDevice, FreeDeviceState> devicesStates =
+                        CommandScheduler.createReleaseMap(context, exception);
+                for (IScheduledInvocationListener scheduleListener : mSchedulerListeners) {
+                    scheduleListener.releaseDevices(context, devicesStates);
+                }
+            }
             try {
                 // Clean up host.
                 invocationPath.doCleanUp(context, config, exception);
@@ -712,6 +724,11 @@
             IRescheduler rescheduler,
             ITestInvocationListener... extraListeners)
             throws DeviceNotAvailableException, Throwable {
+        for (ITestInvocationListener listener : extraListeners) {
+            if (listener instanceof IScheduledInvocationListener) {
+                mSchedulerListeners.add((IScheduledInvocationListener) listener);
+            }
+        }
         // Create the TestInformation for the invocation
         // TODO: Use invocation-id in the workfolder name
         Object sharedInfoObject =