Set allocation

cuebot/src/main/java/com/imageworks/spcue/service/HostManagerService.java

@@ -23,6 +23,7 @@
 import java.util.List;
 
 import org.apache.log4j.Logger;
+import org.springframework.dao.EmptyResultDataAccessException;
 import org.springframework.transaction.annotation.Propagation;
 import org.springframework.transaction.annotation.Transactional;
 
@@ -152,7 +153,28 @@ public DispatchHost createHost(HostReport report) {
 
     @Transactional(propagation = Propagation.REQUIRED)
     public DispatchHost createHost(RenderHost rhost) {
-        return createHost(rhost, getDefaultAllocationDetail());
+        // Find suitable allocation with facility and tags.
+        AllocationEntity alloc = null;
+        if (rhost.getTagsCount() > 0) {
+            String facility = rhost.getFacility();
+            for (String tag : rhost.getTagsList()) {
+                try {
+                    alloc = allocationDao.findAllocationEntity(facility, tag);
+                    logger.info("set " + rhost.getName() +
+                            " to the given allocation " + alloc.getName());
+                    break;
+                }
+                catch (EmptyResultDataAccessException e) {
+                    // Allocation doesn't exist. ignore.
+                }
+            }
+        }
+        if (alloc == null) {
+            alloc = getDefaultAllocationDetail();
+            logger.info("set " + rhost.getName() +
+                    " to the default allocation " + alloc.getName());
+        }
+        return createHost(rhost, alloc);
     }
 
     @Transactional(propagation = Propagation.REQUIRED)

cuebot/src/test/java/com/imageworks/spcue/test/dispatcher/HostReportHandlerTests.java

@@ -27,9 +27,11 @@
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.transaction.annotation.Transactional;
 
+import com.imageworks.spcue.AllocationEntity;
 import com.imageworks.spcue.DispatchHost;
 import com.imageworks.spcue.dispatcher.Dispatcher;
 import com.imageworks.spcue.dispatcher.HostReportHandler;
+import com.imageworks.spcue.FacilityInterface;
 import com.imageworks.spcue.grpc.host.HardwareState;
 import com.imageworks.spcue.grpc.host.LockState;
 import com.imageworks.spcue.grpc.report.CoreDetail;
@@ -58,6 +60,7 @@ public class HostReportHandlerTests extends TransactionalTest {
     Dispatcher dispatcher;
 
     private static final String HOSTNAME = "beta";
+    private static final String NEW_HOSTNAME = "gamma";
 
     @Before
     public void setTestMode() {
@@ -106,6 +109,29 @@ private static RenderHost getRenderHost() {
                 .build();
     }
 
+    private static RenderHost getNewRenderHost(String tags) {
+        return RenderHost.newBuilder()
+                .setName(NEW_HOSTNAME)
+                .setBootTime(1192369572)
+                .setFreeMcp(76020)
+                .setFreeMem(53500)
+                .setFreeSwap(20760)
+                .setLoad(0)
+                .setTotalMcp(195430)
+                .setTotalMem(8173264)
+                .setTotalSwap(20960)
+                .setNimbyEnabled(false)
+                .setNumProcs(2)
+                .setCoresPerProc(100)
+                .addTags(tags)
+                .setState(HardwareState.UP)
+                .setFacility("spi")
+                .putAttributes("SP_OS", "Linux")
+                .putAttributes("freeGpu", String.format("%d", CueUtil.MB512))
+                .putAttributes("totalGpu", String.format("%d", CueUtil.MB512))
+                .build();
+    }
+
     @Test
     @Transactional
     @Rollback(true)
@@ -121,5 +147,71 @@ public void testHandleHostReport() {
         DispatchHost host = getHost();
         assertEquals(host.lockState, LockState.OPEN);
     }
+
+    @Test
+    @Transactional
+    @Rollback(true)
+    public void testHandleHostReportWithNewAllocation() {
+        FacilityInterface facility = adminManager.getFacility(
+                "AAAAAAAA-AAAA-AAAA-AAAA-AAAAAAAAAAA0");
+        assertEquals(facility.getName(), "spi");
+
+        AllocationEntity detail = new AllocationEntity();
+        detail.name = "test";
+        detail.tag = "test";
+        adminManager.createAllocation(facility, detail);
+        detail = adminManager.findAllocationDetail("spi", "test");
+
+        boolean isBoot = true;
+        CoreDetail cores = getCoreDetail(200, 200, 0, 0);
+        HostReport report = HostReport.newBuilder()
+                .setHost(getNewRenderHost("test"))
+                .setCoreInfo(cores)
+                .build();
+
+        hostReportHandler.handleHostReport(report, isBoot);
+        DispatchHost host = hostManager.findDispatchHost(NEW_HOSTNAME);
+        assertEquals(host.getAllocationId(), detail.id);
+    }
+
+    @Test
+    @Transactional
+    @Rollback(true)
+    public void testHandleHostReportWithExistentAllocation() {
+        AllocationEntity alloc = adminManager.getAllocationDetail(
+                "00000000-0000-0000-0000-000000000006");
+        assertEquals(alloc.getName(), "spi.general");
+
+        boolean isBoot = true;
+        CoreDetail cores = getCoreDetail(200, 200, 0, 0);
+        HostReport report = HostReport.newBuilder()
+                .setHost(getNewRenderHost("general"))
+                .setCoreInfo(cores)
+                .build();
+
+        hostReportHandler.handleHostReport(report, isBoot);
+        DispatchHost host = hostManager.findDispatchHost(NEW_HOSTNAME);
+        assertEquals(host.getAllocationId(), alloc.id);
+    }
+
+    @Test
+    @Transactional
+    @Rollback(true)
+    public void testHandleHostReportWithNonExistentTags() {
+        AllocationEntity alloc = adminManager.getAllocationDetail(
+                "00000000-0000-0000-0000-000000000002");
+        assertEquals(alloc.getName(), "lax.unassigned");
+
+        boolean isBoot = true;
+        CoreDetail cores = getCoreDetail(200, 200, 0, 0);
+        HostReport report = HostReport.newBuilder()
+                .setHost(getNewRenderHost("nonexistent"))
+                .setCoreInfo(cores)
+                .build();
+
+        hostReportHandler.handleHostReport(report, isBoot);
+        DispatchHost host = hostManager.findDispatchHost(NEW_HOSTNAME);
+        assertEquals(host.getAllocationId(), alloc.id);
+    }
 }