From 7b7c64c5ff864fc9a1dd28fc5c168e5d42f41b41 Mon Sep 17 00:00:00 2001
From: DeviceInfra <no-reply@google.com>
Date: Thu, 19 Sep 2024 05:25:08 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 676373446
---
 .../scheduler/simple/SimpleScheduler.java     | 61 +++++++++++--------
 1 file changed, 35 insertions(+), 26 deletions(-)

diff --git a/src/java/com/google/devtools/mobileharness/infra/controller/scheduler/simple/SimpleScheduler.java b/src/java/com/google/devtools/mobileharness/infra/controller/scheduler/simple/SimpleScheduler.java
index 6eb7ce43a..951fe76ee 100644
--- a/src/java/com/google/devtools/mobileharness/infra/controller/scheduler/simple/SimpleScheduler.java
+++ b/src/java/com/google/devtools/mobileharness/infra/controller/scheduler/simple/SimpleScheduler.java
@@ -48,6 +48,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
 import javax.inject.Inject;
 
 /**
@@ -74,9 +75,11 @@ public class SimpleScheduler extends AbstractScheduler implements Runnable {
   private final Object allocationLock = new Object();
 
   /** {{@link DeviceLocator#universalId()}, {@link Allocation}} mapping. */
+  @GuardedBy("allocationLock")
   private final Map<String, Allocation> deviceAllocations = new HashMap<>();
 
   /** {TestID, Allocation} mapping. */
+  @GuardedBy("allocationLock")
   private final Map<String, Allocation> testAllocations = new HashMap<>();
 
   private final Sleeper sleeper;
@@ -105,14 +108,16 @@ public void run() {
         for (SimpleJobInfo job : jobs.values()) {
           sleeper.sleep(SCHEDULING_SMALL_INTERVAL);
           for (TestLocator testLocator : job.getTests().values()) {
-            if (!testAllocations.containsKey(testLocator.getId())) {
-              // Found a new test.
-              if (allocate(job.getScheduleUnit(), testLocator)) {
-                hasNewAllocation = true;
+            synchronized (allocationLock) {
+              if (!testAllocations.containsKey(testLocator.getId())) {
+                // Found a new test.
+                if (allocate(job.getScheduleUnit(), testLocator)) {
+                  hasNewAllocation = true;
+                }
+                // No matter successfully allocate devices to the new test or not, skips the
+                // remaining tests and allocates devices for the next job.
+                break;
               }
-              // No matter successfully allocate devices to the new test or not, skips the remaining
-              // tests and allocates devices for the next job.
-              break;
             }
           }
         }
@@ -449,10 +454,12 @@ private boolean allocateSingleDeviceJob(JobScheduleUnit job, TestLocator test) {
   /** Checks whether the device can meet the job requirement. If so, allocates it to the test. */
   private boolean checkAndAllocateSingleDevice(
       JobScheduleUnit job, TestLocator test, DeviceScheduleUnit device, boolean fireEvent) {
-    if (!deviceAllocations.containsKey(device.locator().universalId())
-        && ifDeviceSupports(device, job)) {
-      // Found a suitable and idle device for the new test.
-      return allocate(test, device, fireEvent);
+    synchronized (allocationLock) {
+      if (!deviceAllocations.containsKey(device.locator().universalId())
+          && ifDeviceSupports(device, job)) {
+        // Found a suitable and idle device for the new test.
+        return allocate(test, device, fireEvent);
+      }
     }
     return false;
   }
@@ -467,21 +474,23 @@ private boolean allocateAdhocTestbedJob(JobScheduleUnit job, TestLocator test) {
       labInfos = labList;
     }
     for (SimpleLabInfo labInfo : labInfos) {
-      ImmutableList<DeviceScheduleUnit> filteredDevices =
-          labInfo.getDevices().stream()
-              // Filter out already allocated devices
-              .filter(device -> !deviceAllocations.containsKey(device.locator().universalId()))
-              // Filter out devices that don't support any desired types
-              .filter(device -> !Collections.disjoint(device.types().getAll(), types))
-              // Filter out devices that the user does not own
-              .filter(device -> device.owners().support(job.jobUser().getRunAs()))
-              .collect(toImmutableList());
-      if (!filteredDevices.isEmpty()) {
-        ImmutableList<DeviceScheduleUnit> deviceList =
-            adhocTestbedSchedulingUtil.findSubDevicesSupportingJob(filteredDevices, job);
-        // The order matters in the allocated list as it needs to match the spec.
-        if (!deviceList.isEmpty()) {
-          return allocate(test, deviceList, true);
+      synchronized (allocationLock) {
+        ImmutableList<DeviceScheduleUnit> filteredDevices =
+            labInfo.getDevices().stream()
+                // Filter out already allocated devices
+                .filter(device -> !deviceAllocations.containsKey(device.locator().universalId()))
+                // Filter out devices that don't support any desired types
+                .filter(device -> !Collections.disjoint(device.types().getAll(), types))
+                // Filter out devices that the user does not own
+                .filter(device -> device.owners().support(job.jobUser().getRunAs()))
+                .collect(toImmutableList());
+        if (!filteredDevices.isEmpty()) {
+          ImmutableList<DeviceScheduleUnit> deviceList =
+              adhocTestbedSchedulingUtil.findSubDevicesSupportingJob(filteredDevices, job);
+          // The order matters in the allocated list as it needs to match the spec.
+          if (!deviceList.isEmpty()) {
+            return allocate(test, deviceList, true);
+          }
         }
       }
     }