diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/utils/FederationStateStoreFacade.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/utils/FederationStateStoreFacade.java index f1c8218263..66a0b605e1 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/utils/FederationStateStoreFacade.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/utils/FederationStateStoreFacade.java @@ -48,6 +48,7 @@ import org.apache.hadoop.yarn.exceptions.YarnException; import org.apache.hadoop.yarn.exceptions.YarnRuntimeException; import org.apache.hadoop.yarn.server.federation.store.FederationStateStore; import org.apache.hadoop.yarn.server.federation.store.records.AddApplicationHomeSubClusterRequest; +import org.apache.hadoop.yarn.server.federation.store.records.AddApplicationHomeSubClusterResponse; import org.apache.hadoop.yarn.server.federation.store.records.ApplicationHomeSubCluster; import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationHomeSubClusterRequest; import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationHomeSubClusterResponse; @@ -298,13 +299,15 @@ public final class FederationStateStoreFacade { * * @param appHomeSubCluster the mapping of the application to it's home * sub-cluster + * @return the stored Subcluster from StateStore * @throws YarnException if the call to the state store is unsuccessful */ - public void addApplicationHomeSubCluster( + public SubClusterId addApplicationHomeSubCluster( ApplicationHomeSubCluster appHomeSubCluster) throws YarnException { - stateStore.addApplicationHomeSubCluster( - AddApplicationHomeSubClusterRequest.newInstance(appHomeSubCluster)); - return; + AddApplicationHomeSubClusterResponse response = + stateStore.addApplicationHomeSubCluster( + AddApplicationHomeSubClusterRequest.newInstance(appHomeSubCluster)); + return response.getHomeSubCluster(); } /** diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/TestFederationStateStoreFacade.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/TestFederationStateStoreFacade.java index 53f4f8442e..d46bef09b1 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/TestFederationStateStoreFacade.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/TestFederationStateStoreFacade.java @@ -28,6 +28,7 @@ import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; import org.apache.hadoop.yarn.server.federation.store.FederationStateStore; import org.apache.hadoop.yarn.server.federation.store.impl.MemoryFederationStateStore; +import org.apache.hadoop.yarn.server.federation.store.records.ApplicationHomeSubCluster; import org.apache.hadoop.yarn.server.federation.store.records.SubClusterId; import org.apache.hadoop.yarn.server.federation.store.records.SubClusterInfo; import org.apache.hadoop.yarn.server.federation.store.records.SubClusterPolicyConfiguration; @@ -145,4 +146,33 @@ public class TestFederationStateStoreFacade { } } + @Test + public void testAddApplicationHomeSubCluster() throws YarnException { + + // Inserting into FederationStateStore + ApplicationId appId = ApplicationId.newInstance(clusterTs, numApps + 1); + SubClusterId subClusterId1 = SubClusterId.newInstance("Home1"); + + ApplicationHomeSubCluster appHomeSubCluster = + ApplicationHomeSubCluster.newInstance(appId, subClusterId1); + + SubClusterId result = + facade.addApplicationHomeSubCluster(appHomeSubCluster); + + Assert.assertEquals(facade.getApplicationHomeSubCluster(appId), result); + Assert.assertEquals(subClusterId1, result); + + // Inserting into FederationStateStore. + // The application is already present. + // FederationFacade will return Home1 as SubClusterId. + SubClusterId subClusterId2 = SubClusterId.newInstance("Home2"); + appHomeSubCluster = + ApplicationHomeSubCluster.newInstance(appId, subClusterId2); + + result = facade.addApplicationHomeSubCluster(appHomeSubCluster); + + Assert.assertEquals(facade.getApplicationHomeSubCluster(appId), result); + Assert.assertEquals(subClusterId1, result); + } + }