YARN-11272. Federation StateStore: Support storage/retrieval of Reservations With Zk. (#4781)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/impl/ZKFederationStateStoreOpDurations.java

@@ -74,6 +74,21 @@ public final class ZKFederationStateStoreOpDurations implements MetricsSource {
   @Metric("Duration for a get PolicyConfigurations call")
   private MutableRate getPoliciesConfigurations;
 
+  @Metric("Duration for a add reservation homeSubCluster call")
+  private MutableRate addReservationHomeSubCluster;
+
+  @Metric("Duration for a get reservation homeSubCluster call")
+  private MutableRate getReservationHomeSubCluster;
+
+  @Metric("Duration for a get reservations homeSubCluster call")
+  private MutableRate getReservationsHomeSubCluster;
+
+  @Metric("Duration for a delete reservation homeSubCluster call")
+  private MutableRate deleteReservationHomeSubCluster;
+
+  @Metric("Duration for a update reservation homeSubCluster call")
+  private MutableRate updateReservationHomeSubCluster;
+
   protected static final MetricsInfo RECORD_INFO =
       info("ZKFederationStateStoreOpDurations", "Durations of ZKFederationStateStore calls");
 
@@ -152,4 +167,24 @@ public final class ZKFederationStateStoreOpDurations implements MetricsSource {
   public void addGetPoliciesConfigurationsDuration(long startTime, long endTime) {
     getPoliciesConfigurations.add(endTime - startTime);
   }
+
+  public void addReservationHomeSubClusterDuration(long startTime, long endTime) {
+    addReservationHomeSubCluster.add(endTime - startTime);
+  }
+
+  public void addGetReservationHomeSubClusterDuration(long startTime, long endTime) {
+    getReservationHomeSubCluster.add(endTime - startTime);
+  }
+
+  public void addGetReservationsHomeSubClusterDuration(long startTime, long endTime) {
+    getReservationsHomeSubCluster.add(endTime - startTime);
+  }
+
+  public void addDeleteReservationHomeSubClusterDuration(long startTime, long endTime) {
+    deleteReservationHomeSubCluster.add(endTime - startTime);
+  }
+
+  public void addUpdateReservationHomeSubClusterDuration(long startTime, long endTime) {
+    updateReservationHomeSubCluster.add(endTime - startTime);
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/impl/ZookeeperFederationStateStore.java

@@ -26,7 +26,6 @@ import java.util.List;
 import java.util.TimeZone;
 
 import org.apache.hadoop.classification.VisibleForTesting;
-import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.util.curator.ZKCuratorManager;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
@@ -65,6 +64,7 @@ import org.apache.hadoop.yarn.server.federation.store.records.SubClusterPolicyCo
 import org.apache.hadoop.yarn.server.federation.store.records.SubClusterRegisterRequest;
 import org.apache.hadoop.yarn.server.federation.store.records.SubClusterRegisterResponse;
 import org.apache.hadoop.yarn.server.federation.store.records.SubClusterState;
+import org.apache.hadoop.yarn.server.federation.store.records.ReservationHomeSubCluster;
 import org.apache.hadoop.yarn.server.federation.store.records.UpdateApplicationHomeSubClusterRequest;
 import org.apache.hadoop.yarn.server.federation.store.records.UpdateApplicationHomeSubClusterResponse;
 import org.apache.hadoop.yarn.server.federation.store.records.AddReservationHomeSubClusterRequest;
@@ -84,7 +84,9 @@ import org.apache.hadoop.yarn.server.federation.store.utils.FederationApplicatio
 import org.apache.hadoop.yarn.server.federation.store.utils.FederationMembershipStateStoreInputValidator;
 import org.apache.hadoop.yarn.server.federation.store.utils.FederationPolicyStoreInputValidator;
 import org.apache.hadoop.yarn.server.federation.store.utils.FederationStateStoreUtils;
+import org.apache.hadoop.yarn.server.federation.store.utils.FederationReservationHomeSubClusterStoreInputValidator;
 import org.apache.hadoop.yarn.server.records.Version;
+import org.apache.hadoop.yarn.api.records.ReservationId;
 import org.apache.hadoop.yarn.util.Clock;
 import org.apache.hadoop.yarn.util.SystemClock;
 import org.apache.zookeeper.data.ACL;
@@ -105,8 +107,11 @@ import org.apache.hadoop.thirdparty.protobuf.InvalidProtocolBufferException;
  * |     |----- APP1
  * |     |----- APP2
  * |--- POLICY
- *       |----- QUEUE1
- *       |----- QUEUE1
+ * |     |----- QUEUE1
+ * |     |----- QUEUE1
+ * |--- RESERVATION
+ * |     |----- RESERVATION1
+ * |     |----- RESERVATION2
  */
 public class ZookeeperFederationStateStore implements FederationStateStore {
 
@@ -116,6 +121,7 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
   private final static String ROOT_ZNODE_NAME_MEMBERSHIP = "memberships";
   private final static String ROOT_ZNODE_NAME_APPLICATION = "applications";
   private final static String ROOT_ZNODE_NAME_POLICY = "policies";
+  private final static String ROOT_ZNODE_NAME_RESERVATION = "reservation";
 
   /** Interface to Zookeeper. */
   private ZKCuratorManager zkManager;
@@ -126,6 +132,7 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
   private String appsZNode;
   private String membershipZNode;
   private String policiesZNode;
+  private String reservationsZNode;
 
   private volatile Clock clock = SystemClock.getInstance();
 
@@ -151,6 +158,7 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
     membershipZNode = getNodePath(baseZNode, ROOT_ZNODE_NAME_MEMBERSHIP);
     appsZNode = getNodePath(baseZNode, ROOT_ZNODE_NAME_APPLICATION);
     policiesZNode = getNodePath(baseZNode, ROOT_ZNODE_NAME_POLICY);
+    reservationsZNode = getNodePath(baseZNode, ROOT_ZNODE_NAME_RESERVATION);
 
     // Create base znode for each entity
     try {
@@ -158,6 +166,7 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
       zkManager.createRootDirRecursively(membershipZNode, zkAcl);
       zkManager.createRootDirRecursively(appsZNode, zkAcl);
       zkManager.createRootDirRecursively(policiesZNode, zkAcl);
+      zkManager.createRootDirRecursively(reservationsZNode, zkAcl);
     } catch (Exception e) {
       String errMsg = "Cannot create base directories: " + e.getMessage();
       FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
@@ -686,6 +695,30 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
     return cal.getTimeInMillis();
   }
 
+  private void putReservation(final ReservationId reservationId,
+      final SubClusterId subClusterId, boolean update) throws YarnException {
+    String reservationZNode = getNodePath(reservationsZNode, reservationId.toString());
+    SubClusterIdProto proto = ((SubClusterIdPBImpl)subClusterId).getProto();
+    byte[] data = proto.toByteArray();
+    put(reservationZNode, data, update);
+  }
+
+  private SubClusterId getReservation(final ReservationId reservationId)
+      throws YarnException {
+    String reservationIdZNode = getNodePath(reservationsZNode, reservationId.toString());
+    SubClusterId subClusterId = null;
+    byte[] data = get(reservationIdZNode);
+    if (data != null) {
+      try {
+        subClusterId = new SubClusterIdPBImpl(SubClusterIdProto.parseFrom(data));
+      } catch (InvalidProtocolBufferException e) {
+        String errMsg = "Cannot parse reservation at " + reservationId;
+        FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+      }
+    }
+    return subClusterId;
+  }
+
   @VisibleForTesting
   public ZKFederationStateStoreOpDurations getOpDurations() {
     return opDurations;
@@ -694,30 +727,128 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
   @Override
   public AddReservationHomeSubClusterResponse addReservationHomeSubCluster(
       AddReservationHomeSubClusterRequest request) throws YarnException {
-    throw new NotImplementedException("Code is not implemented");
+
+    long start = clock.getTime();
+    FederationReservationHomeSubClusterStoreInputValidator.validate(request);
+    ReservationHomeSubCluster reservationHomeSubCluster = request.getReservationHomeSubCluster();
+    ReservationId reservationId = reservationHomeSubCluster.getReservationId();
+
+    // Try to write the subcluster
+    SubClusterId homeSubCluster = reservationHomeSubCluster.getHomeSubCluster();
+    try {
+      putReservation(reservationId, homeSubCluster, false);
+    } catch (Exception e) {
+      String errMsg = "Cannot add reservation home subcluster for " + reservationId;
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+
+    // Check for the actual subcluster
+    try {
+      homeSubCluster = getReservation(reservationId);
+    } catch (Exception e) {
+      String errMsg = "Cannot check app home subcluster for " + reservationId;
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+    long end = clock.getTime();
+    opDurations.addReservationHomeSubClusterDuration(start, end);
+    return AddReservationHomeSubClusterResponse.newInstance(homeSubCluster);
   }
 
   @Override
   public GetReservationHomeSubClusterResponse getReservationHomeSubCluster(
       GetReservationHomeSubClusterRequest request) throws YarnException {
-    throw new NotImplementedException("Code is not implemented");
+
+    long start = clock.getTime();
+    FederationReservationHomeSubClusterStoreInputValidator.validate(request);
+    ReservationId reservationId = request.getReservationId();
+    SubClusterId homeSubCluster = getReservation(reservationId);
+
+    if (homeSubCluster == null) {
+      String errMsg = "Reservation " + reservationId + " does not exist";
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+
+    ReservationHomeSubCluster reservationHomeSubCluster =
+        ReservationHomeSubCluster.newInstance(reservationId, homeSubCluster);
+    long end = clock.getTime();
+    opDurations.addGetReservationHomeSubClusterDuration(start, end);
+    return GetReservationHomeSubClusterResponse.newInstance(reservationHomeSubCluster);
   }
 
   @Override
   public GetReservationsHomeSubClusterResponse getReservationsHomeSubCluster(
       GetReservationsHomeSubClusterRequest request) throws YarnException {
-    throw new NotImplementedException("Code is not implemented");
+    long start = clock.getTime();
+    List<ReservationHomeSubCluster> result = new ArrayList<>();
+
+    try {
+      for (String child : zkManager.getChildren(reservationsZNode)) {
+        ReservationId reservationId = ReservationId.parseReservationId(child);
+        SubClusterId homeSubCluster = getReservation(reservationId);
+        ReservationHomeSubCluster app =
+            ReservationHomeSubCluster.newInstance(reservationId, homeSubCluster);
+        result.add(app);
+      }
+    } catch (Exception e) {
+      String errMsg = "Cannot get apps: " + e.getMessage();
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+    long end = clock.getTime();
+    opDurations.addGetReservationsHomeSubClusterDuration(start, end);
+    return GetReservationsHomeSubClusterResponse.newInstance(result);
   }
 
   @Override
   public DeleteReservationHomeSubClusterResponse deleteReservationHomeSubCluster(
       DeleteReservationHomeSubClusterRequest request) throws YarnException {
-    throw new NotImplementedException("Code is not implemented");
+    long start = clock.getTime();
+    FederationReservationHomeSubClusterStoreInputValidator.validate(request);
+    ReservationId reservationId = request.getReservationId();
+    String reservationZNode = getNodePath(reservationsZNode, reservationId.toString());
+
+    boolean exists = false;
+    try {
+      exists = zkManager.exists(reservationZNode);
+    } catch (Exception e) {
+      String errMsg = "Cannot check reservation: " + e.getMessage();
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+
+    if (!exists) {
+      String errMsg = "Reservation " + reservationId + " does not exist";
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+
+    try {
+      zkManager.delete(reservationZNode);
+    } catch (Exception e) {
+      String errMsg = "Cannot delete reservation: " + e.getMessage();
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+    long end = clock.getTime();
+    opDurations.addDeleteReservationHomeSubClusterDuration(start, end);
+    return DeleteReservationHomeSubClusterResponse.newInstance();
   }
 
   @Override
   public UpdateReservationHomeSubClusterResponse updateReservationHomeSubCluster(
       UpdateReservationHomeSubClusterRequest request) throws YarnException {
-    throw new NotImplementedException("Code is not implemented");
+
+    long start = clock.getTime();
+    FederationReservationHomeSubClusterStoreInputValidator.validate(request);
+    ReservationHomeSubCluster reservationHomeSubCluster = request.getReservationHomeSubCluster();
+    ReservationId reservationId = reservationHomeSubCluster.getReservationId();
+    SubClusterId homeSubCluster = getReservation(reservationId);
+
+    if (homeSubCluster == null) {
+      String errMsg = "Reservation " + reservationId + " does not exist";
+      FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
+    }
+
+    SubClusterId newSubClusterId = reservationHomeSubCluster.getHomeSubCluster();
+    putReservation(reservationId, newSubClusterId, true);
+    long end = clock.getTime();
+    opDurations.addUpdateReservationHomeSubClusterDuration(start, end);
+    return UpdateReservationHomeSubClusterResponse.newInstance();
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/store/impl/FederationStateStoreBaseTest.java

@@ -24,13 +24,17 @@ import java.util.List;
 import java.util.TimeZone;
 
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.test.LambdaTestUtils;
+import org.apache.hadoop.util.Time;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.hadoop.yarn.api.records.ReservationId;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.federation.store.FederationStateStore;
 import org.apache.hadoop.yarn.server.federation.store.exception.FederationStateStoreException;
 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.ReservationHomeSubCluster;
 import org.apache.hadoop.yarn.server.federation.store.records.DeleteApplicationHomeSubClusterRequest;
 import org.apache.hadoop.yarn.server.federation.store.records.DeleteApplicationHomeSubClusterResponse;
 import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationHomeSubClusterRequest;
@@ -56,6 +60,14 @@ import org.apache.hadoop.yarn.server.federation.store.records.SubClusterRegister
 import org.apache.hadoop.yarn.server.federation.store.records.SubClusterState;
 import org.apache.hadoop.yarn.server.federation.store.records.UpdateApplicationHomeSubClusterRequest;
 import org.apache.hadoop.yarn.server.federation.store.records.UpdateApplicationHomeSubClusterResponse;
+import org.apache.hadoop.yarn.server.federation.store.records.AddReservationHomeSubClusterRequest;
+import org.apache.hadoop.yarn.server.federation.store.records.AddReservationHomeSubClusterResponse;
+import org.apache.hadoop.yarn.server.federation.store.records.GetReservationHomeSubClusterResponse;
+import org.apache.hadoop.yarn.server.federation.store.records.GetReservationHomeSubClusterRequest;
+import org.apache.hadoop.yarn.server.federation.store.records.DeleteReservationHomeSubClusterRequest;
+import org.apache.hadoop.yarn.server.federation.store.records.DeleteReservationHomeSubClusterResponse;
+import org.apache.hadoop.yarn.server.federation.store.records.UpdateReservationHomeSubClusterRequest;
+import org.apache.hadoop.yarn.server.federation.store.records.UpdateReservationHomeSubClusterResponse;
 import org.apache.hadoop.yarn.util.MonotonicClock;
 import org.junit.After;
 import org.junit.Assert;
@@ -599,4 +611,153 @@ public abstract class FederationStateStoreBaseTest {
     return stateStore;
   }
 
+  SubClusterId queryReservationHomeSC(ReservationId reservationId)
+      throws YarnException {
+
+    GetReservationHomeSubClusterRequest request =
+        GetReservationHomeSubClusterRequest.newInstance(reservationId);
+
+    GetReservationHomeSubClusterResponse response =
+        stateStore.getReservationHomeSubCluster(request);
+
+    return response.getReservationHomeSubCluster().getHomeSubCluster();
+  }
+
+  @Test
+  public void testAddReservationHomeSubCluster() throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+    SubClusterId subClusterId = SubClusterId.newInstance("SC");
+
+    ReservationHomeSubCluster reservationHomeSubCluster =
+        ReservationHomeSubCluster.newInstance(reservationId, subClusterId);
+
+    AddReservationHomeSubClusterRequest request =
+        AddReservationHomeSubClusterRequest.newInstance(reservationHomeSubCluster);
+    AddReservationHomeSubClusterResponse response =
+        stateStore.addReservationHomeSubCluster(request);
+
+    Assert.assertEquals(subClusterId, response.getHomeSubCluster());
+    Assert.assertEquals(subClusterId, queryReservationHomeSC(reservationId));
+  }
+
+  private void addReservationHomeSC(ReservationId reservationId, SubClusterId subClusterId)
+      throws YarnException {
+
+    ReservationHomeSubCluster reservationHomeSubCluster =
+        ReservationHomeSubCluster.newInstance(reservationId, subClusterId);
+    AddReservationHomeSubClusterRequest request =
+        AddReservationHomeSubClusterRequest.newInstance(reservationHomeSubCluster);
+    stateStore.addReservationHomeSubCluster(request);
+  }
+
+  @Test
+  public void testAddReservationHomeSubClusterReservationAlreadyExists() throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+    SubClusterId subClusterId1 = SubClusterId.newInstance("SC1");
+    addReservationHomeSC(reservationId, subClusterId1);
+
+    SubClusterId subClusterId2 = SubClusterId.newInstance("SC2");
+    ReservationHomeSubCluster reservationHomeSubCluster2 =
+        ReservationHomeSubCluster.newInstance(reservationId, subClusterId2);
+    AddReservationHomeSubClusterRequest request2 =
+        AddReservationHomeSubClusterRequest.newInstance(reservationHomeSubCluster2);
+    AddReservationHomeSubClusterResponse response =
+        stateStore.addReservationHomeSubCluster(request2);
+
+    Assert.assertNotNull(response);
+    Assert.assertEquals(subClusterId1, response.getHomeSubCluster());
+    Assert.assertEquals(subClusterId1, queryReservationHomeSC(reservationId));
+  }
+
+  @Test
+  public void testAddReservationHomeSubClusterAppAlreadyExistsInTheSameSC()
+      throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+    SubClusterId subClusterId1 = SubClusterId.newInstance("SC1");
+    addReservationHomeSC(reservationId, subClusterId1);
+
+    ReservationHomeSubCluster reservationHomeSubCluster2 =
+        ReservationHomeSubCluster.newInstance(reservationId, subClusterId1);
+    AddReservationHomeSubClusterRequest request2 =
+        AddReservationHomeSubClusterRequest.newInstance(reservationHomeSubCluster2);
+    AddReservationHomeSubClusterResponse response =
+        stateStore.addReservationHomeSubCluster(request2);
+
+    Assert.assertNotNull(response);
+    Assert.assertEquals(subClusterId1, response.getHomeSubCluster());
+    Assert.assertEquals(subClusterId1, queryReservationHomeSC(reservationId));
+  }
+
+  @Test
+  public void testDeleteReservationHomeSubCluster() throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+    SubClusterId subClusterId1 = SubClusterId.newInstance("SC");
+    addReservationHomeSC(reservationId, subClusterId1);
+
+    DeleteReservationHomeSubClusterRequest delReservationRequest =
+        DeleteReservationHomeSubClusterRequest.newInstance(reservationId);
+    DeleteReservationHomeSubClusterResponse delReservationResponse =
+        stateStore.deleteReservationHomeSubCluster(delReservationRequest);
+
+    Assert.assertNotNull(delReservationResponse);
+
+    LambdaTestUtils.intercept(YarnException.class,
+        "Reservation " + reservationId + " does not exist",
+        () -> queryReservationHomeSC(reservationId));
+  }
+
+  @Test
+  public void testDeleteReservationHomeSubClusterUnknownApp() throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+
+    DeleteReservationHomeSubClusterRequest delReservationRequest =
+        DeleteReservationHomeSubClusterRequest.newInstance(reservationId);
+
+    LambdaTestUtils.intercept(YarnException.class,
+        "Reservation " + reservationId + " does not exist",
+        () -> stateStore.deleteReservationHomeSubCluster(delReservationRequest));
+  }
+
+  @Test
+  public void testUpdateReservationHomeSubCluster() throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+    SubClusterId subClusterId1 = SubClusterId.newInstance("SC");
+    addReservationHomeSC(reservationId, subClusterId1);
+
+    SubClusterId subClusterId2 = SubClusterId.newInstance("SC2");
+    ReservationHomeSubCluster reservationHomeSubCluster =
+        ReservationHomeSubCluster.newInstance(reservationId, subClusterId2);
+
+    UpdateReservationHomeSubClusterRequest updateReservationRequest =
+        UpdateReservationHomeSubClusterRequest.newInstance(reservationHomeSubCluster);
+
+    UpdateReservationHomeSubClusterResponse updateReservationResponse =
+        stateStore.updateReservationHomeSubCluster(updateReservationRequest);
+
+    Assert.assertNotNull(updateReservationResponse);
+    Assert.assertEquals(subClusterId2, queryReservationHomeSC(reservationId));
+  }
+
+  @Test
+  public void testUpdateReservationHomeSubClusterUnknownApp() throws Exception {
+
+    ReservationId reservationId = ReservationId.newInstance(Time.now(), 1);
+    SubClusterId subClusterId1 = SubClusterId.newInstance("SC1");
+
+    ReservationHomeSubCluster reservationHomeSubCluster =
+        ReservationHomeSubCluster.newInstance(reservationId, subClusterId1);
+
+    UpdateReservationHomeSubClusterRequest updateReservationRequest =
+        UpdateReservationHomeSubClusterRequest.newInstance(reservationHomeSubCluster);
+
+    LambdaTestUtils.intercept(YarnException.class,
+        "Reservation " + reservationId + " does not exist",
+        () -> stateStore.updateReservationHomeSubCluster(updateReservationRequest));
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/store/impl/TestSQLFederationStateStore.java

@@ -17,6 +17,7 @@
 
 package org.apache.hadoop.yarn.server.federation.store.impl;
 
+import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.exceptions.YarnException;
@@ -74,4 +75,39 @@ public class TestSQLFederationStateStore extends FederationStateStoreBaseTest {
     Assert.assertEquals(1,
         FederationStateStoreClientMetrics.getNumConnections());
   }
+
+  @Test(expected = NotImplementedException.class)
+  public void testAddReservationHomeSubCluster() throws Exception {
+    super.testAddReservationHomeSubCluster();
+  }
+
+  @Test(expected = NotImplementedException.class)
+  public void testAddReservationHomeSubClusterReservationAlreadyExists() throws Exception {
+    super.testAddReservationHomeSubClusterReservationAlreadyExists();
+  }
+
+  @Test(expected = NotImplementedException.class)
+  public void testAddReservationHomeSubClusterAppAlreadyExistsInTheSameSC() throws Exception {
+    super.testAddReservationHomeSubClusterAppAlreadyExistsInTheSameSC();
+  }
+
+  @Test(expected = NotImplementedException.class)
+  public void testDeleteReservationHomeSubCluster() throws Exception {
+    super.testDeleteReservationHomeSubCluster();
+  }
+
+  @Test(expected = NotImplementedException.class)
+  public void testDeleteReservationHomeSubClusterUnknownApp() throws Exception {
+    super.testDeleteReservationHomeSubClusterUnknownApp();
+  }
+
+  @Test(expected = NotImplementedException.class)
+  public void testUpdateReservationHomeSubCluster() throws Exception {
+    super.testUpdateReservationHomeSubCluster();
+  }
+
+  @Test(expected = NotImplementedException.class)
+  public void testUpdateReservationHomeSubClusterUnknownApp() throws Exception {
+    super.testUpdateReservationHomeSubClusterUnknownApp();
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/store/impl/TestZookeeperFederationStateStore.java

@@ -115,6 +115,11 @@ public class TestZookeeperFederationStateStore
     zkStateStoreOpDurations.addGetPolicyConfigurationDuration(start, end);
     zkStateStoreOpDurations.addSetPolicyConfigurationDuration(start, end);
     zkStateStoreOpDurations.addGetPoliciesConfigurationsDuration(start, end);
+    zkStateStoreOpDurations.addReservationHomeSubClusterDuration(start, end);
+    zkStateStoreOpDurations.addGetReservationHomeSubClusterDuration(start, end);
+    zkStateStoreOpDurations.addGetReservationsHomeSubClusterDuration(start, end);
+    zkStateStoreOpDurations.addDeleteReservationHomeSubClusterDuration(start, end);
+    zkStateStoreOpDurations.addUpdateReservationHomeSubClusterDuration(start, end);
 
     zkStateStoreOpDurations.getMetrics(collector, true);
     assertEquals("Incorrect number of perf metrics", 1, collector.getRecords().size());
@@ -137,6 +142,11 @@ public class TestZookeeperFederationStateStore
     MetricsRecords.assertMetric(record, "GetPolicyConfigurationAvgTime",  expectAvgTime);
     MetricsRecords.assertMetric(record, "SetPolicyConfigurationAvgTime",  expectAvgTime);
     MetricsRecords.assertMetric(record, "GetPoliciesConfigurationsAvgTime",  expectAvgTime);
+    MetricsRecords.assertMetric(record, "AddReservationHomeSubClusterAvgTime",  expectAvgTime);
+    MetricsRecords.assertMetric(record, "GetReservationHomeSubClusterAvgTime",  expectAvgTime);
+    MetricsRecords.assertMetric(record, "GetReservationsHomeSubClusterAvgTime",  expectAvgTime);
+    MetricsRecords.assertMetric(record, "DeleteReservationHomeSubClusterAvgTime",  expectAvgTime);
+    MetricsRecords.assertMetric(record, "UpdateReservationHomeSubClusterAvgTime",  expectAvgTime);
 
     long expectOps = 1;
     MetricsRecords.assertMetric(record, "AddAppHomeSubClusterNumOps",  expectOps);
@@ -152,5 +162,10 @@ public class TestZookeeperFederationStateStore
     MetricsRecords.assertMetric(record, "GetPolicyConfigurationNumOps",  expectOps);
     MetricsRecords.assertMetric(record, "SetPolicyConfigurationNumOps",  expectOps);
     MetricsRecords.assertMetric(record, "GetPoliciesConfigurationsNumOps",  expectOps);
+    MetricsRecords.assertMetric(record, "AddReservationHomeSubClusterNumOps",  expectOps);
+    MetricsRecords.assertMetric(record, "GetReservationHomeSubClusterNumOps",  expectOps);
+    MetricsRecords.assertMetric(record, "GetReservationsHomeSubClusterNumOps",  expectOps);
+    MetricsRecords.assertMetric(record, "DeleteReservationHomeSubClusterNumOps",  expectOps);
+    MetricsRecords.assertMetric(record, "UpdateReservationHomeSubClusterNumOps",  expectOps);
   }
 }