Skip to content

Commit 2ca49d6

Browse files
ASPS002Ujjwal SrivastavaUjjwal Srivastava
authored
Optimise queries to save Blocked Device Info in Repository (#683)
<!-- Please provide brief information about the PR, what it contains & its purpose, new behaviors after the change. And let us know here if you need any help: https://github.com/microsoft/HydraLab/issues/new --> ## Description This pull request includes several changes to the `DeviceAgentManagementService` and its related repository to improve the handling of blocked devices. The most important changes involve replacing direct repository calls with `Optional` checks, consolidating delete operations, and simplifying imports. Improvements to blocked device handling: * [`center/src/main/java/com/microsoft/hydralab/center/service/DeviceAgentManagementService.java`](diffhunk://#diff-289317ad410c52c217b55f97dc31f0806ffad3b4cc737a56e255701ad77f3797L1163-R1166): Replaced direct repository calls with `Optional` checks to handle the presence of blocked devices more safely. * [`center/src/main/java/com/microsoft/hydralab/center/service/DeviceAgentManagementService.java`](diffhunk://#diff-289317ad410c52c217b55f97dc31f0806ffad3b4cc737a56e255701ad77f3797L1197-R1208): Simplified the `unBlockDevice` method by consolidating delete operations into a single repository method `deleteIfExists`. Repository improvements: * [`common/src/main/java/com/microsoft/hydralab/common/repository/BlockedDeviceInfoRepository.java`](diffhunk://#diff-f76c5b6f6c8de8f14263d4913754c060bd181a96e20e8aeff3728fd254a9e15aL4-R27): Added `Optional` return type for `findByBlockedDeviceSerialNumber` and consolidated delete operations into a new method `deleteIfExists`. Codebase simplification: * [`center/src/test/java/com/microsoft/hydralab/center/service/DeviceAgentManagementServiceTest.java`](diffhunk://#diff-6562c25d5fc53e481385e45d6835f0017d51ca06bcbe2a6bdbee489a0c15e51cL86-L97): Updated test cases to use the new `deleteIfExists` method. <!-- A few words to explain your changes --> ### Linked GitHub issue ID: # #685 ## Pull Request Checklist <!-- Put an x in the boxes that apply. This is simply a reminder of what we are going to look for before merging your code. --> - [x] Tests for the changes have been added (for bug fixes / features) - [x] Code compiles correctly with all tests are passed. - [x] I've read the [contributing guide](https://github.com/microsoft/HydraLab/blob/main/CONTRIBUTING.md#making-changes-to-the-code) and followed the recommended practices. - [ ] [Wikis](https://github.com/microsoft/HydraLab/wiki) or [README](https://github.com/microsoft/HydraLab/blob/main/README.md) have been reviewed and added / updated if needed (for bug fixes / features) ### Does this introduce a breaking change? *If this introduces a breaking change for Hydra Lab users, please describe the impact and migration path.* - [ ] Yes - [x] No ## How you tested it *Please make sure the change is tested, you can test it by adding UTs, do local test and share the screenshots, etc.* Yes , tested it multiple times by blocking and unblocking devices and using unit tests. **Blocking a device** <img width="1358" alt="block" src="https://github.com/user-attachments/assets/edc93fa5-f35a-4fab-86b7-eddf9c18253d" /> **Unblocking a device** <img width="1336" alt="unblock" src="https://github.com/user-attachments/assets/e0328868-ffc7-42fd-89cd-7367b8f2a988" /> Please check the type of change your PR introduces: - [x] Bugfix - [ ] Feature - [ ] Technical design - [ ] Build related changes - [ ] Refactoring (no functional changes, no api changes) - [ ] Code style update (formatting, renaming) or Documentation content changes - [ ] Other (please describe): ### Feature UI screenshots or Technical design diagrams *If this is a relatively large or complex change, kick it off by drawing the tech design with PlantUML and explaining why you chose the solution you did and what alternatives you considered, etc...* --------- Co-authored-by: Ujjwal Srivastava <[email protected]> Co-authored-by: Ujjwal Srivastava <[email protected]>
1 parent 01ae41b commit 2ca49d6

File tree

3 files changed

+25
-32
lines changed

3 files changed

+25
-32
lines changed

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

+15-14
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
import com.microsoft.hydralab.common.entity.common.AgentUser;
2323
import com.microsoft.hydralab.common.entity.common.AnalysisTask;
2424
import com.microsoft.hydralab.common.entity.common.DeviceInfo;
25-
import com.microsoft.hydralab.common.entity.common.DeviceOperation;
2625
import com.microsoft.hydralab.common.entity.common.Message;
2726
import com.microsoft.hydralab.common.entity.common.StatisticData;
2827
import com.microsoft.hydralab.common.entity.common.StorageFileInfo;
@@ -31,6 +30,7 @@
3130
import com.microsoft.hydralab.common.entity.common.TestTask;
3231
import com.microsoft.hydralab.common.entity.common.TestTaskSpec;
3332
import com.microsoft.hydralab.common.entity.common.BlockedDeviceInfo;
33+
import com.microsoft.hydralab.common.entity.common.DeviceOperation;
3434
import com.microsoft.hydralab.common.file.StorageServiceClientProxy;
3535
import com.microsoft.hydralab.common.management.device.DeviceType;
3636
import com.microsoft.hydralab.common.repository.BlockedDeviceInfoRepository;
@@ -1160,8 +1160,10 @@ public boolean isDeviceBlocked(String deviceIdentifier) {
11601160
if (blockedDevicesMap.containsKey(deviceIdentifier)){
11611161
return true;
11621162
}
1163-
if (blockedDeviceInfoRepository.existsByBlockedDeviceSerialNumber(deviceIdentifier)) {
1164-
blockedDevicesMap.put(deviceIdentifier, blockedDeviceInfoRepository.findByBlockedDeviceSerialNumber(deviceIdentifier));
1163+
Optional<BlockedDeviceInfo> blockedDeviceInfoOpt = blockedDeviceInfoRepository.findByBlockedDeviceSerialNumber(deviceIdentifier);
1164+
if (blockedDeviceInfoOpt.isPresent()){
1165+
BlockedDeviceInfo blockedDeviceInfo = blockedDeviceInfoOpt.get();
1166+
blockedDevicesMap.put(deviceIdentifier, blockedDeviceInfo);
11651167
return true;
11661168
}
11671169
return false;
@@ -1174,8 +1176,10 @@ public boolean areAllDevicesBlocked(String deviceIdentifier) {
11741176
synchronized (blockedDevicesMap) {
11751177
for (String deviceSerial : deviceSerials) {
11761178
if (!blockedDevicesMap.containsKey(deviceSerial)) {
1177-
if (blockedDeviceInfoRepository.existsByBlockedDeviceSerialNumber(deviceIdentifier)) {
1178-
blockedDevicesMap.put(deviceIdentifier, blockedDeviceInfoRepository.findByBlockedDeviceSerialNumber(deviceIdentifier));
1179+
Optional<BlockedDeviceInfo> blockedDeviceInfoOpt = blockedDeviceInfoRepository.findByBlockedDeviceSerialNumber(deviceIdentifier);
1180+
if (blockedDeviceInfoOpt.isPresent()){
1181+
BlockedDeviceInfo blockedDeviceInfo = blockedDeviceInfoOpt.get();
1182+
blockedDevicesMap.put(deviceIdentifier, blockedDeviceInfo);
11791183
continue;
11801184
}
11811185
return false;
@@ -1194,18 +1198,14 @@ public void unBlockDevice(TestTaskSpec testTaskSpec) {
11941198
BlockedDeviceInfo blockedDeviceInfo = blockedDevicesMap.get(testTaskSpec.deviceIdentifier);
11951199
if (blockedDeviceInfo.getBlockingTaskUUID().equals(testTaskSpec.unblockDeviceSecretKey)) {
11961200
blockedDevicesMap.remove(testTaskSpec.deviceIdentifier);
1197-
if (blockedDeviceInfoRepository.existsByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier)) {
1198-
blockedDeviceInfoRepository.deleteByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier);
1199-
}
1201+
blockedDeviceInfoRepository.deleteIfExists(testTaskSpec.deviceIdentifier);
12001202
testTaskSpec.unblockedDeviceSerialNumber = testTaskSpec.deviceIdentifier;
12011203

12021204
} else {
12031205
throw new IllegalArgumentException("Invalid unblock device secret key!");
12041206
}
12051207
} else {
1206-
if (blockedDeviceInfoRepository.existsByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier)) {
1207-
blockedDeviceInfoRepository.deleteByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier);
1208-
}
1208+
blockedDeviceInfoRepository.deleteIfExists(testTaskSpec.deviceIdentifier);
12091209
log.warn("Device {} is already unblocked.", testTaskSpec.deviceIdentifier);
12101210
testTaskSpec.unblockedDeviceSerialNumber = testTaskSpec.deviceIdentifier;
12111211
}
@@ -1224,7 +1224,6 @@ public void unblockFrozenBlockedDevices() {
12241224
while (iterator.hasNext()) {
12251225
Map.Entry<String, BlockedDeviceInfo> entry = iterator.next();
12261226
Instant blockedTime = entry.getValue().getBlockedTime();
1227-
String blockedDeviceSerialNumber = entry.getValue().getBlockedDeviceSerialNumber();
12281227
Duration durationBlocked = Duration.between(blockedTime, currentTime);
12291228
if (durationBlocked.compareTo(Const.DeviceGroup.BLOCKED_DEVICE_TIMEOUT) > 0) {
12301229
log.info("Unblocking device {} since it has been blocked for more than {} hours.", entry.getKey(), durationBlocked);
@@ -1246,8 +1245,9 @@ public boolean isRunOnBlockedDevice(TestTaskSpec testTaskSpec) {
12461245
BlockedDeviceInfo blockedDeviceInfo = blockedDevicesMap.get(testTaskSpec.deviceIdentifier);
12471246

12481247
if (blockedDeviceInfo == null) {
1249-
if (blockedDeviceInfoRepository.existsByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier)){
1250-
blockedDeviceInfo = blockedDeviceInfoRepository.findByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier);
1248+
Optional<BlockedDeviceInfo> blockedDeviceInfoOpt = blockedDeviceInfoRepository.findByBlockedDeviceSerialNumber(testTaskSpec.deviceIdentifier);
1249+
if (blockedDeviceInfoOpt.isPresent()) {
1250+
blockedDeviceInfo = blockedDeviceInfoOpt.get();
12511251
blockedDevicesMap.put(testTaskSpec.deviceIdentifier, blockedDeviceInfo);
12521252
return blockedDeviceInfo.getBlockingTaskUUID().equals(testTaskSpec.unblockDeviceSecretKey);
12531253
}
@@ -1292,4 +1292,5 @@ public AgentSessionInfo(Session session, AgentUser agentUser) {
12921292
this.agentUser = agentUser;
12931293
}
12941294
}
1295+
12951296
}

‎center/src/test/java/com/microsoft/hydralab/center/service/DeviceAgentManagementServiceTest.java

+2-7
Original file line numberDiff line numberDiff line change
@@ -83,17 +83,12 @@ public void blockDevice() {
8383
blockedDevicesMap.put(blockedDeviceInfo.getBlockedDeviceSerialNumber(), blockedDeviceInfo);
8484
blockedDeviceInfoRepository.save(blockedDeviceInfo);
8585

86-
if (blockedDeviceInfoRepository.existsByBlockedDeviceSerialNumber(blockedDeviceInfo.getBlockedDeviceSerialNumber())) {
87-
blockedDeviceInfoRepository.deleteByBlockedDeviceSerialNumber(blockedDeviceInfo.getBlockedDeviceSerialNumber());
88-
}
86+
blockedDeviceInfoRepository.deleteIfExists(blockedDeviceInfo.getBlockedDeviceSerialNumber());
8987

9088
List<BlockedDeviceInfo> blockedDeviceInfoList = blockedDeviceInfoRepository.findAll();
9189
for (BlockedDeviceInfo blockedDeviceInfo1 : blockedDeviceInfoList) {
92-
blockedDeviceInfoRepository.deleteByBlockedTimeBefore(Instant.now());
90+
System.out.println(blockedDeviceInfo1);
9391
}
94-
9592
}
96-
97-
9893
}
9994
}

‎common/src/main/java/com/microsoft/hydralab/common/repository/BlockedDeviceInfoRepository.java

+8-11
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,28 @@
11
package com.microsoft.hydralab.common.repository;
22

33
import com.microsoft.hydralab.common.entity.common.BlockedDeviceInfo;
4-
import org.springframework.data.jpa.repository.JpaRepository;
5-
import org.springframework.data.jpa.repository.JpaSpecificationExecutor;
6-
import org.springframework.data.jpa.repository.Modifying;
7-
import org.springframework.data.jpa.repository.Query;
4+
import org.springframework.data.jpa.repository.*;
85
import org.springframework.data.repository.query.Param;
96
import org.springframework.stereotype.Repository;
107

8+
import javax.persistence.LockModeType;
119
import javax.transaction.Transactional;
1210
import java.time.Instant;
1311
import java.util.List;
12+
import java.util.Optional;
1413

1514
@Repository
1615
public interface BlockedDeviceInfoRepository extends JpaRepository<BlockedDeviceInfo, String>, JpaSpecificationExecutor<BlockedDeviceInfo> {
1716
@Transactional
18-
BlockedDeviceInfo findByBlockedDeviceSerialNumber(String blockedDeviceSerialNumber);
19-
20-
@Transactional
21-
boolean existsByBlockedDeviceSerialNumber(String blockedDeviceSerialNumber);
17+
Optional<BlockedDeviceInfo> findByBlockedDeviceSerialNumber(@Param("serialNumber") String serialNumber);
2218

2319
@Modifying
2420
@Transactional
25-
void deleteByBlockedDeviceSerialNumber(String blockedDeviceSerialNumber);
21+
@Query("DELETE FROM BlockedDeviceInfo b WHERE b.blockedTime < :blockedTime")
22+
void deleteByBlockedTimeBefore(@Param("blockedTime") Instant blockedTime);
2623

2724
@Modifying
2825
@Transactional
29-
@Query("DELETE FROM BlockedDeviceInfo b WHERE b.blockedTime < :time")
30-
void deleteByBlockedTimeBefore(@Param("time") Instant time);
26+
@Query("DELETE FROM BlockedDeviceInfo b WHERE b.blockedDeviceSerialNumber = :blockedDeviceSerialNumber")
27+
void deleteIfExists(@Param("blockedDeviceSerialNumber") String blockedDeviceSerialNumber);
3128
}

0 commit comments

Comments
 (0)