User/ujsrivastava/block unblock device test spec

center/src/main/java/com/microsoft/hydralab/center/controller/TestTaskController.java

@@ -5,6 +5,7 @@
 
 import com.alibaba.fastjson.JSONArray;
 import com.alibaba.fastjson.JSONObject;
+import com.google.common.base.Strings;
 import com.microsoft.hydralab.center.service.DeviceAgentManagementService;
 import com.microsoft.hydralab.center.service.SysUserService;
 import com.microsoft.hydralab.center.service.TestDataService;
@@ -80,6 +81,18 @@ public Result<Object> runTestTask(@CurrentSecurityContext SysUser requestor,
             testTaskSpec.teamId = testFileSet.getTeamId();
             testTaskSpec.teamName = testFileSet.getTeamName();
             testTaskSpec.testTaskId = UUID.randomUUID().toString();
+
+            if (testTaskSpec.blockDevice && testTaskSpec.unblockDevice) {
+                throw new IllegalArgumentException("Cannot block and unblock device in the same test task.");
+            }
+            if (testTaskSpec.unblockDevice && Strings.isNullOrEmpty(testTaskSpec.unblockDeviceSecretKey)) {
+                throw new IllegalArgumentException("Unblock secret key is required when unblocking a device.");
+            }
+
+            if (testTaskSpec.unblockDevice && testTaskSpec.deviceIdentifier.startsWith("G.")) {
+                throw new IllegalArgumentException("deviceIdentifier should not be a group when unblocking a device.");
+            }
+
             if (!sysUserService.checkUserAdmin(requestor)) {
                 if (!userTeamManagementService.checkRequestorTeamRelation(requestor, testTaskSpec.teamId)) {
                     return Result.error(HttpStatus.UNAUTHORIZED.value(), "Unauthorized, the TestFileSet doesn't belong to user's Teams");
@@ -89,7 +102,7 @@ public Result<Object> runTestTask(@CurrentSecurityContext SysUser requestor,
             //if the queue is not empty, the task will be added to the queue directly
             if (testTaskService.isQueueEmpty()
                     || Task.RunnerType.APK_SCANNER.name().equals(testTaskSpec.runningType)
-                    || testTaskService.isDeviceFree(testTaskSpec.deviceIdentifier)) {
+                    || deviceAgentManagementService.isRunOnBlockedDevice(testTaskSpec) || testTaskService.isDeviceFree(testTaskSpec.deviceIdentifier)) {
                 result = deviceAgentManagementService.runTestTaskBySpec(testTaskSpec);
                 if (result.get(Const.Param.TEST_DEVICE_SN) == null) {
                     //if there is no alive device, the task will be added to the queue directly

center/src/main/java/com/microsoft/hydralab/center/service/DeviceAgentManagementService.java

@@ -29,6 +29,7 @@
 import com.microsoft.hydralab.common.entity.common.TestRun;
 import com.microsoft.hydralab.common.entity.common.TestTask;
 import com.microsoft.hydralab.common.entity.common.TestTaskSpec;
+import com.microsoft.hydralab.common.entity.common.BlockedDeviceInfo;
 import com.microsoft.hydralab.common.file.StorageServiceClientProxy;
 import com.microsoft.hydralab.common.management.device.DeviceType;
 import com.microsoft.hydralab.common.repository.StatisticDataRepository;
@@ -58,6 +59,8 @@
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
@@ -67,8 +70,10 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.Iterator;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
@@ -81,7 +86,7 @@ public class DeviceAgentManagementService {
      * Connected session count
      */
     private final AtomicInteger onlineCount = new AtomicInteger(0);
-
+    private static volatile AtomicBoolean isUnblockingDevices = new AtomicBoolean(false);
     //save agent session <sessionId,session&agentUser>
     private final ConcurrentHashMap<String, AgentSessionInfo> agentSessionMap = new ConcurrentHashMap<>();
     //save agent info <agentId,agentInfo>
@@ -94,6 +99,9 @@ public class DeviceAgentManagementService {
     private final ConcurrentHashMap<String, AccessInfo> accessInfoMap = new ConcurrentHashMap<>();
     //save agent update info <agentId,updateTask>
     private final ConcurrentHashMap<String, AgentUpdateTask> agentUpdateMap = new ConcurrentHashMap<>();
+    // save blocked devices <deviceSerial, blockedDeviceInfo>
+    private final ConcurrentHashMap<String, BlockedDeviceInfo> blockedDevicesMap = new ConcurrentHashMap<>();
+
     @Resource
     MetricUtil metricUtil;
     @Resource
@@ -695,7 +703,7 @@ public List<AgentDeviceGroup> getAllAppiumAgents() {
 
     public JSONObject runTestTaskBySpec(TestTaskSpec testTaskSpec) {
         JSONObject result;
-
+        unblockFrozenBlockedDevices();
         if (Task.RunnerType.APPIUM_CROSS.name().equals(testTaskSpec.runningType)) {
             result = runAppiumTestTask(testTaskSpec);
         } else if (Task.RunnerType.T2C_JSON.name().equals(testTaskSpec.runningType)) {
@@ -921,12 +929,22 @@ private JSONObject runTestTaskByGroup(TestTaskSpec testTaskSpec) {
                 Assert.isTrue(!isAll, "Device/Agent Offline!");
                 continue;
             }
+
+            if (isDeviceBlocked(deviceSerial)) {
+                Assert.isTrue(!isAll, "Some of the devices in the device group are blocked!");
+                continue;
+            }
+
             isAllOffline = false;
             if (device.isOnline()) {
                 List<String> devices = testAgentDevicesMap.getOrDefault(device.getAgentId(), new ArrayList<>());
                 devices.add(device.getSerialNum());
                 testAgentDevicesMap.put(device.getAgentId(), devices);
                 testTaskSpec.agentIds.add(device.getAgentId());
+                if (testTaskSpec.blockDevice) {
+                    testTaskSpec.deviceIdentifier = deviceSerial;
+                    blockDevice(testTaskSpec);
+                }
                 if (isSingle) {
                     break;
                 }
@@ -964,16 +982,27 @@ private JSONObject runTestTaskByDevice(TestTaskSpec testTaskSpec) {
         Message message = new Message();
         message.setBody(testTaskSpec);
         message.setPath(Const.Path.TEST_TASK_RUN);
+
         Assert.isTrue(device.isAlive(), "Device/Agent Offline!");
-        if (device.isTesting()) {
+        if (device.isTesting() || (!isRunOnBlockedDevice(testTaskSpec) && isDeviceBlocked(testTaskSpec.deviceIdentifier))) {
             return result;
         }
         AgentSessionInfo agentSessionInfoByAgentId = getAgentSessionInfoByAgentId(device.getAgentId());
+
         Assert.notNull(agentSessionInfoByAgentId, "Device/Agent Offline!");
+
         if (isAgentUpdating(agentSessionInfoByAgentId.agentUser.getId())) {
             return result;
         }
         updateDeviceStatus(device.getSerialNum(), DeviceInfo.TESTING, testTaskSpec.testTaskId);
+
+        if (testTaskSpec.unblockDevice) {
+            unBlockDevice(testTaskSpec);
+        }
+
+        if (testTaskSpec.blockDevice) {
+            blockDevice(testTaskSpec);
+        }
         testTaskSpec.agentIds.add(device.getAgentId());
         sendMessageToSession(agentSessionInfoByAgentId.session, message);
         result.put(Const.Param.TEST_DEVICE_SN, testTaskSpec.deviceIdentifier);
@@ -1108,6 +1137,98 @@ public int getAliveDeviceNum() {
         return (int) deviceListMap.values().stream().filter(DeviceInfo::isAlive).count();
     }
 
+    public void blockDevice(TestTaskSpec testTaskSpec) {
+        synchronized (blockedDevicesMap) {
+            if (blockedDevicesMap.containsKey(testTaskSpec.deviceIdentifier)) {
+                log.warn("Device {} is already blocked!", testTaskSpec.deviceIdentifier);
+                return;
+            }
+
+            BlockedDeviceInfo blockedDeviceInfo = new BlockedDeviceInfo();
+            blockedDeviceInfo.setBlockedTime(Instant.now());
+            blockedDeviceInfo.setBlockingTaskUUID(testTaskSpec.testTaskId);
+            blockedDeviceInfo.setBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier);
+            blockedDevicesMap.put(testTaskSpec.deviceIdentifier,blockedDeviceInfo);
+            testTaskSpec.blockedDeviceSerialNumber = testTaskSpec.deviceIdentifier;
+            testTaskSpec.unblockDeviceSecretKey = testTaskSpec.testTaskId;
+        }
+    }
+
+    public boolean isDeviceBlocked(String deviceIdentifier) {
+        synchronized (blockedDevicesMap) {
+            return blockedDevicesMap.containsKey(deviceIdentifier);
+        }
+    }
+
+    public boolean areAllDevicesBlocked(String deviceIdentifier) {
+        if (deviceIdentifier.startsWith(Const.DeviceGroup.GROUP_NAME_PREFIX)) {
+            Set<String> deviceSerials = queryDeviceByGroup(deviceIdentifier);
+            synchronized (blockedDevicesMap) {
+                for (String deviceSerial : deviceSerials) {
+                    if (!blockedDevicesMap.containsKey(deviceSerial)) {
+                        return false;
+                    }
+                }
+            }
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    public void unBlockDevice(TestTaskSpec testTaskSpec) {
+        synchronized (blockedDevicesMap) {
+            if (blockedDevicesMap.containsKey(testTaskSpec.deviceIdentifier)) {
+                BlockedDeviceInfo blockedDeviceInfo = blockedDevicesMap.get(testTaskSpec.deviceIdentifier);
+                if (blockedDeviceInfo.getBlockingTaskUUID().equals(testTaskSpec.unblockDeviceSecretKey)) {
+                    blockedDevicesMap.remove(testTaskSpec.deviceIdentifier);
+                    testTaskSpec.unblockedDeviceSerialNumber = testTaskSpec.deviceIdentifier;
+                } else {
+                    throw new IllegalArgumentException("Invalid unblock device secret key!");
+                }
+            } else {
+                log.warn("Device {} is already unblocked.", testTaskSpec.deviceIdentifier);
+                testTaskSpec.unblockedDeviceSerialNumber = testTaskSpec.deviceIdentifier;
+            }
+        }
+    }
+
+    public void unblockFrozenBlockedDevices() {
+        if (isUnblockingDevices.get()){
+            return;
+        }
+        isUnblockingDevices.set(true);
+        synchronized (blockedDevicesMap) {
+            Instant currentTime = Instant.now();
+            Iterator<Map.Entry<String, BlockedDeviceInfo>> iterator = blockedDevicesMap.entrySet().iterator();
+            while (iterator.hasNext()) {
+                Map.Entry<String, BlockedDeviceInfo> entry = iterator.next();
+                Instant blockedTime = entry.getValue().getBlockedTime();
+                Duration durationBlocked = Duration.between(blockedTime, currentTime);
+                if (durationBlocked.compareTo(Const.DeviceGroup.BLOCKED_DEVICE_TIMEOUT) > 0) {
+                    log.info("Unblocking device {} since it has been blocked for more than {} hours.", entry.getKey(), durationBlocked);
+                    iterator.remove();
+                }
+            }
+        }
+        isUnblockingDevices.set(false);
+    }
+
+    public boolean isRunOnBlockedDevice(TestTaskSpec testTaskSpec) {
+        if (testTaskSpec.unblockDeviceSecretKey == null || testTaskSpec.unblockDeviceSecretKey.isEmpty()) {
+            return false;
+        }
+
+        synchronized (blockedDevicesMap) {
+            BlockedDeviceInfo blockedDeviceInfo = blockedDevicesMap.get(testTaskSpec.deviceIdentifier);
+
+            if (blockedDeviceInfo == null) {
+                return false;
+            }
+            return blockedDeviceInfo.getBlockingTaskUUID().equals(testTaskSpec.unblockDeviceSecretKey);
+        }
+    }
+
     @Scheduled(cron = "0 * * * * *")
     public void recordStatisticData() {
         int currentAgentNum = getAliveAgentNum();

center/src/main/java/com/microsoft/hydralab/center/service/TestTaskService.java

@@ -66,8 +66,16 @@ public Boolean isDeviceFree(String deviceIdentifier) {
                 relatedIdentifiers.addAll(deviceAgentManagementService.queryGroupByDevice(deviceIdentifier));
             }
         } else if (deviceIdentifier.startsWith(Const.DeviceGroup.GROUP_NAME_PREFIX)) {
+            if (deviceAgentManagementService.areAllDevicesBlocked(deviceIdentifier)) {
+                logger.warn("All Devices in the DeviceGroup " + deviceIdentifier + " are blocked currently.");
+                return false;
+            }
             relatedIdentifiers.addAll(deviceAgentManagementService.queryDeviceByGroup(deviceIdentifier));
         } else {
+            if (deviceAgentManagementService.isDeviceBlocked(deviceIdentifier)) {
+                logger.warn("Device " + deviceIdentifier + " is blocked currently.");
+                return false;
+            }
             relatedIdentifiers.addAll(deviceAgentManagementService.queryGroupByDevice(deviceIdentifier));
         }
         synchronized (taskQueue) {
@@ -92,6 +100,7 @@ public void runTask() {
             return;
         }
         isRunning.set(true);
+
         synchronized (taskQueue) {
             Iterator<TestTaskSpec> queueIterator = taskQueue.iterator();
             while (queueIterator.hasNext()) {

common/src/main/java/com/microsoft/hydralab/common/entity/common/BlockedDeviceInfo.java

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+package com.microsoft.hydralab.common.entity.common;
+
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+
+import java.time.Instant;
+
+@Getter
+@Setter
+@ToString
+public class BlockedDeviceInfo {
+    public Instant blockedTime;
+    public String blockingTaskUUID;
+    public String blockedDeviceSerialNumber;
+}

common/src/main/java/com/microsoft/hydralab/common/entity/common/Task.java

@@ -86,6 +86,13 @@ public class Task implements Serializable {
     private boolean disableRecording = false;
     @Column(columnDefinition = "boolean default false")
     private boolean isSucceed = false;
+    public String blockedDeviceSerialNumber;
+    public String unblockedDeviceSerialNumber;
+    public String unblockDeviceSecretKey;
+    @Transient
+    public boolean blockDevice = false;
+    @Transient
+    public boolean unblockDevice = false;
 
     @Transient
     private List<TestRun> deviceTestResults;
@@ -164,6 +171,16 @@ public TestTaskSpec convertToTaskSpec() {
         testTaskSpec.disableRecording = isDisableRecording();
         testTaskSpec.retryTime = getRetryTime();
 
+        if (isBlockDevice()) {
+            testTaskSpec.blockDevice = true;
+            testTaskSpec.blockedDeviceSerialNumber = getBlockedDeviceSerialNumber();
+            testTaskSpec.unblockDeviceSecretKey = getUnblockDeviceSecretKey();
+        }
+        if (isUnblockDevice()) {
+            testTaskSpec.unblockDevice = true;
+            testTaskSpec.unblockedDeviceSerialNumber = getUnblockedDeviceSerialNumber();
+            testTaskSpec.unblockDeviceSecretKey = getUnblockDeviceSecretKey();
+        }
         return testTaskSpec;
     }
 
@@ -198,6 +215,18 @@ public Task(TestTaskSpec testTaskSpec) {
         setTeamName(testTaskSpec.teamName);
         setNotifyUrl(testTaskSpec.notifyUrl);
         setDisableRecording(testTaskSpec.disableRecording);
+
+        if (testTaskSpec.blockDevice) {
+            setBlockDevice(true);
+            setBlockedDeviceSerialNumber(testTaskSpec.blockedDeviceSerialNumber);
+            setUnblockDeviceSecretKey(testTaskSpec.unblockDeviceSecretKey);
+        }
+
+        if (testTaskSpec.unblockDevice) {
+            setUnblockDevice(true);
+            setUnblockedDeviceSerialNumber(testTaskSpec.unblockedDeviceSerialNumber);
+            setUnblockDeviceSecretKey(testTaskSpec.unblockDeviceSecretKey);
+        }
     }
 
     public Task() {

common/src/main/java/com/microsoft/hydralab/common/entity/common/TestTaskSpec.java

@@ -57,6 +57,11 @@ public class TestTaskSpec {
     public boolean enableNetworkMonitor;
     public String networkMonitorRule;
     public boolean enableTestOrchestrator = false;
+    public boolean blockDevice = false;
+    public boolean unblockDevice = false;
+    public String blockedDeviceSerialNumber;
+    public String unblockedDeviceSerialNumber;
+    public String unblockDeviceSecretKey;
 
     public void updateWithDefaultValues() {
         determineScopeOfTestCase();

common/src/main/java/com/microsoft/hydralab/common/util/Const.java

@@ -3,6 +3,7 @@
 
 package com.microsoft.hydralab.common.util;
 
+import java.time.Duration;
 import java.util.List;
 
 public interface Const {
@@ -32,6 +33,7 @@ interface DeviceGroup {
         String SINGLE_TYPE = "SINGLE";
         String REST_TYPE = "REST";
         String ALL_TYPE = "ALL";
+        Duration BLOCKED_DEVICE_TIMEOUT = Duration.ofHours(4);
     }
 
     interface AgentConfig {