diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/resources/numa/NumaResourceAllocation.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/resources/numa/NumaResourceAllocation.java index f8d47396e98..e361170a51d 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/resources/numa/NumaResourceAllocation.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/resources/numa/NumaResourceAllocation.java @@ -20,6 +20,7 @@ import java.io.Serializable; import java.util.HashMap; import java.util.Map; +import java.util.Objects; import java.util.Set; /** @@ -66,4 +67,26 @@ public void addCpuNode(String cpuNodeId, int cpus) { public Map getNodeVsCpus() { return nodeVsCpus; } -} + + @Override + public String toString() { + return "NumaResourceAllocation{" + + "nodeVsMemory=" + nodeVsMemory + + ", nodeVsCpus=" + nodeVsCpus + + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + NumaResourceAllocation that = (NumaResourceAllocation) o; + return Objects.equals(nodeVsMemory, that.nodeVsMemory) && + Objects.equals(nodeVsCpus, that.nodeVsCpus); + } + + @Override + public int hashCode() { + return Objects.hash(nodeVsMemory, nodeVsCpus); + } +} \ No newline at end of file diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMLeveldbStateStoreService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMLeveldbStateStoreService.java index 1d7771a9e1f..8de94a54987 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMLeveldbStateStoreService.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMLeveldbStateStoreService.java @@ -1459,8 +1459,7 @@ public void storeAssignedResources(Container container, String keyResChng = CONTAINERS_KEY_PREFIX + container.getContainerId().toString() + CONTAINER_ASSIGNED_RESOURCES_KEY_SUFFIX + resourceType; try { - WriteBatch batch = db.createWriteBatch(); - try { + try (WriteBatch batch = db.createWriteBatch()) { ResourceMappings.AssignedResources res = new ResourceMappings.AssignedResources(); res.updateAssignedResources(assignedResources); @@ -1468,8 +1467,6 @@ public void storeAssignedResources(Container container, // New value will overwrite old values for the same key batch.put(bytes(keyResChng), res.toBytes()); db.write(batch); - } finally { - batch.close(); } } catch (DBException e) { markStoreUnHealthy(e); diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/recovery/TestNMLeveldbStateStoreService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/recovery/TestNMLeveldbStateStoreService.java index 87208f76497..c4c194c3a94 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/recovery/TestNMLeveldbStateStoreService.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/recovery/TestNMLeveldbStateStoreService.java @@ -75,6 +75,9 @@ import org.apache.hadoop.yarn.server.nodemanager.amrmproxy.AMRMProxyTokenSecretManager; import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.Container; import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.ResourceMappings; +import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.fpga.FpgaResourceAllocator.FpgaDevice; +import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.numa.NumaResourceAllocation; +import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.gpu.GpuDevice; import org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.LocalResourceTrackerState; import org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.RecoveredAMRMProxyState; import org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.RecoveredApplicationsState; @@ -1448,7 +1451,7 @@ public void testAMRMProxyStorage() throws IOException { @Test public void testStateStoreForResourceMapping() throws IOException { - // test empty when no state + // test that stateStore is initially empty List recoveredContainers = loadContainersState(stateStore.getContainerStateIterator()); assertTrue(recoveredContainers.isEmpty()); @@ -1464,38 +1467,43 @@ public void testStateStoreForResourceMapping() throws IOException { ResourceMappings resourceMappings = new ResourceMappings(); when(container.getResourceMappings()).thenReturn(resourceMappings); - // Store ResourceMapping stateStore.storeAssignedResources(container, "gpu", - Arrays.asList("1", "2", "3")); - // This will overwrite above - List gpuRes1 = Arrays.asList("1", "2", "4"); + Arrays.asList(new GpuDevice(1, 1), new GpuDevice(2, 2), + new GpuDevice(3, 3))); + + // This will overwrite the above + List gpuRes1 = Arrays.asList( + new GpuDevice(1, 1), new GpuDevice(2, 2), new GpuDevice(4, 4)); stateStore.storeAssignedResources(container, "gpu", gpuRes1); - List fpgaRes = Arrays.asList("3", "4", "5", "6"); + + List fpgaRes = Arrays.asList( + new FpgaDevice("testType", 3, 3, "testIPID"), + new FpgaDevice("testType", 4, 4, "testIPID"), + new FpgaDevice("testType", 5, 5, "testIPID"), + new FpgaDevice("testType", 6, 6, "testIPID")); stateStore.storeAssignedResources(container, "fpga", fpgaRes); - List numaRes = Arrays.asList("numa1"); + + List numaRes = Arrays.asList( + new NumaResourceAllocation("testmemNodeId", 2048, "testCpuNodeId", 10)); stateStore.storeAssignedResources(container, "numa", numaRes); - // add a invalid key restartStateStore(); recoveredContainers = loadContainersState(stateStore.getContainerStateIterator()); assertEquals(1, recoveredContainers.size()); RecoveredContainerState rcs = recoveredContainers.get(0); - List res = rcs.getResourceMappings() + List resources = rcs.getResourceMappings() .getAssignedResources("gpu"); - Assert.assertTrue(res.equals(gpuRes1)); - Assert.assertTrue( - resourceMappings.getAssignedResources("gpu").equals(gpuRes1)); - - res = rcs.getResourceMappings().getAssignedResources("fpga"); - Assert.assertTrue(res.equals(fpgaRes)); - Assert.assertTrue( - resourceMappings.getAssignedResources("fpga").equals(fpgaRes)); - - res = rcs.getResourceMappings().getAssignedResources("numa"); - Assert.assertTrue(res.equals(numaRes)); - Assert.assertTrue( - resourceMappings.getAssignedResources("numa").equals(numaRes)); + Assert.assertEquals(gpuRes1, resources); + Assert.assertEquals(gpuRes1, resourceMappings.getAssignedResources("gpu")); + + resources = rcs.getResourceMappings().getAssignedResources("fpga"); + Assert.assertEquals(fpgaRes, resources); + Assert.assertEquals(fpgaRes, resourceMappings.getAssignedResources("fpga")); + + resources = rcs.getResourceMappings().getAssignedResources("numa"); + Assert.assertEquals(numaRes, resources); + Assert.assertEquals(numaRes, resourceMappings.getAssignedResources("numa")); } @Test