helix-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From hu...@apache.org
Subject [helix] branch zooscalability updated: Implement setRoutingData for MetadataStoreDirectoryService (#844)
Date Wed, 04 Mar 2020 01:37:24 GMT
This is an automated email from the ASF dual-hosted git repository.

hulee pushed a commit to branch zooscalability
in repository https://gitbox.apache.org/repos/asf/helix.git


The following commit(s) were added to refs/heads/zooscalability by this push:
     new 1ebf942  Implement setRoutingData for MetadataStoreDirectoryService (#844)
1ebf942 is described below

commit 1ebf942909f040fecdc664c90dbcc9838fae1a91
Author: Neal Sun <nealsun.0428@gmail.com>
AuthorDate: Tue Mar 3 17:37:15 2020 -0800

    Implement setRoutingData for MetadataStoreDirectoryService (#844)
    
    Implement setRoutingData endpoint. Modify TrieRoutingData construction in MetadataStoreDirectory. Fix race conditions among writing operations in MetadataStoreDirectory.
---
 .../rest/metadatastore/MetadataStoreDirectory.java |   9 ++
 .../metadatastore/ZkMetadataStoreDirectory.java    | 100 ++++++++++++++-----
 .../accessor/MetadataStoreRoutingDataWriter.java   |   2 +-
 .../accessor/ZkRoutingDataReader.java              |   2 +-
 .../accessor/ZkRoutingDataWriter.java              |  43 +++++---
 .../MetadataStoreDirectoryAccessor.java            |  25 +++++
 .../TestZkMetadataStoreDirectory.java              |  69 ++++++++++++-
 .../accessor/TestZkRoutingDataReader.java          |  64 +++++-------
 .../accessor/TestZkRoutingDataWriter.java          | 109 +++++++++++++--------
 .../MetadataStoreDirectoryAccessorTestBase.java    |  30 ++++--
 .../rest/server/TestMSDAccessorLeaderElection.java |  57 ++++++++---
 .../server/TestMetadataStoreDirectoryAccessor.java |  38 +++++++
 .../helix/msdcommon/datamodel/TrieRoutingData.java |   4 +-
 13 files changed, 409 insertions(+), 143 deletions(-)

diff --git a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/MetadataStoreDirectory.java b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/MetadataStoreDirectory.java
index 4630d50..8371e07 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/MetadataStoreDirectory.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/MetadataStoreDirectory.java
@@ -61,6 +61,15 @@ public interface MetadataStoreDirectory extends AutoCloseable {
   Map<String, List<String>> getNamespaceRoutingData(String namespace);
 
   /**
+   * Sets and overwrites routing data in the given namespace.
+   *
+   * @param namespace namespace in metadata store directory.
+   * @param routingData Routing data map: realm -> List of sharding keys
+   * @return true if successful; false otherwise.
+   */
+  boolean setNamespaceRoutingData(String namespace, Map<String, List<String>> routingData);
+
+  /**
    * Returns all path-based sharding keys in the given namespace and the realm.
    * @param namespace
    * @param realm
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/ZkMetadataStoreDirectory.java b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/ZkMetadataStoreDirectory.java
index 28a4afe..c83245f 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/ZkMetadataStoreDirectory.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/ZkMetadataStoreDirectory.java
@@ -113,10 +113,14 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
           _routingDataWriterMap.put(namespace, new ZkRoutingDataWriter(namespace, zkAddress));
 
           // Populate realmToShardingKeys with ZkRoutingDataReader
-          _realmToShardingKeysMap
-              .put(namespace, _routingDataReaderMap.get(namespace).getRoutingData());
-          _routingDataMap
-              .put(namespace, new TrieRoutingData(_realmToShardingKeysMap.get(namespace)));
+          Map<String, List<String>> rawRoutingData =
+              _routingDataReaderMap.get(namespace).getRoutingData();
+          _realmToShardingKeysMap.put(namespace, rawRoutingData);
+          try {
+            _routingDataMap.put(namespace, new TrieRoutingData(rawRoutingData));
+          } catch (InvalidRoutingDataException e) {
+            LOG.warn("TrieRoutingData is not created for namespace {}", namespace, e);
+          }
         }
       }
     }
@@ -156,6 +160,19 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
   }
 
   @Override
+  public boolean setNamespaceRoutingData(String namespace, Map<String, List<String>> routingData) {
+    if (!_routingDataWriterMap.containsKey(namespace)) {
+      throw new IllegalArgumentException(
+          "Failed to set routing data: Namespace " + namespace + " is not found!");
+    }
+    synchronized (this) {
+      boolean result = _routingDataWriterMap.get(namespace).setRoutingData(routingData);
+      refreshRoutingData(namespace);
+      return result;
+    }
+  }
+
+  @Override
   public Collection<String> getAllShardingKeysInRealm(String namespace, String realm) {
     if (!_realmToShardingKeysMap.containsKey(namespace)) {
       throw new NoSuchElementException("Namespace " + namespace + " does not exist!");
@@ -169,19 +186,31 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
 
   @Override
   public Map<String, String> getAllMappingUnderPath(String namespace, String path) {
-    if (!_routingDataMap.containsKey(namespace)) {
+    // Check _routingZkAddressMap first to see if namespace is included
+    if (!_routingZkAddressMap.containsKey(namespace)) {
       throw new NoSuchElementException(
           "Failed to get all mapping under path: Namespace " + namespace + " is not found!");
     }
+    // If namespace is included but not routing data, it means the routing data is invalid
+    if (!_routingDataMap.containsKey(namespace)) {
+      throw new IllegalStateException("Failed to get all mapping under path: Namespace " + namespace
+          + " contains either empty or invalid routing data!");
+    }
     return _routingDataMap.get(namespace).getAllMappingUnderPath(path);
   }
 
   @Override
   public String getMetadataStoreRealm(String namespace, String shardingKey) {
-    if (!_routingDataMap.containsKey(namespace)) {
+    // Check _routingZkAddressMap first to see if namespace is included
+    if (!_routingZkAddressMap.containsKey(namespace)) {
       throw new NoSuchElementException(
           "Failed to get metadata store realm: Namespace " + namespace + " is not found!");
     }
+    // If namespace is included but not routing data, it means the routing data is invalid
+    if (!_routingDataMap.containsKey(namespace)) {
+      throw new IllegalStateException("Failed to get metadata store realm: Namespace " + namespace
+          + " contains either empty or invalid routing data!");
+    }
     return _routingDataMap.get(namespace).getMetadataStoreRealm(shardingKey);
   }
 
@@ -193,7 +222,11 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
       throw new NoSuchElementException(
           "Failed to add metadata store realm: Namespace " + namespace + " is not found!");
     }
-    return _routingDataWriterMap.get(namespace).addMetadataStoreRealm(realm);
+    synchronized (this) {
+      boolean result = _routingDataWriterMap.get(namespace).addMetadataStoreRealm(realm);
+      refreshRoutingData(namespace);
+      return result;
+    }
   }
 
   @Override
@@ -204,26 +237,36 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
       throw new NoSuchElementException(
           "Failed to delete metadata store realm: Namespace " + namespace + " is not found!");
     }
-    return _routingDataWriterMap.get(namespace).deleteMetadataStoreRealm(realm);
+    synchronized (this) {
+      boolean result = _routingDataWriterMap.get(namespace).deleteMetadataStoreRealm(realm);
+      refreshRoutingData(namespace);
+      return result;
+    }
   }
 
   @Override
   public boolean addShardingKey(String namespace, String realm, String shardingKey) {
-    if (!_routingDataWriterMap.containsKey(namespace) || !_routingDataMap.containsKey(namespace)) {
+    if (!_routingDataWriterMap.containsKey(namespace)) {
       // throwing NoSuchElementException instead of IllegalArgumentException to differentiate the
       // status code in the Accessor level
       throw new NoSuchElementException(
           "Failed to add sharding key: Namespace " + namespace + " is not found!");
     }
-    if (_routingDataMap.get(namespace).containsKeyRealmPair(shardingKey, realm)) {
-      return true;
-    }
-    if (!_routingDataMap.get(namespace).isShardingKeyInsertionValid(shardingKey)) {
-      throw new IllegalArgumentException(
-          "Failed to add sharding key: Adding sharding key " + shardingKey
-              + " makes routing data invalid!");
+    synchronized (this) {
+      if (_routingDataMap.containsKey(namespace) && _routingDataMap.get(namespace)
+          .containsKeyRealmPair(shardingKey, realm)) {
+        return true;
+      }
+      if (_routingDataMap.containsKey(namespace) && !_routingDataMap.get(namespace)
+          .isShardingKeyInsertionValid(shardingKey)) {
+        throw new IllegalArgumentException(
+            "Failed to add sharding key: Adding sharding key " + shardingKey
+                + " makes routing data invalid!");
+      }
+      boolean result = _routingDataWriterMap.get(namespace).addShardingKey(realm, shardingKey);
+      refreshRoutingData(namespace);
+      return result;
     }
-    return _routingDataWriterMap.get(namespace).addShardingKey(realm, shardingKey);
   }
 
   @Override
@@ -234,7 +277,11 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
       throw new NoSuchElementException(
           "Failed to delete sharding key: Namespace " + namespace + " is not found!");
     }
-    return _routingDataWriterMap.get(namespace).deleteShardingKey(realm, shardingKey);
+    synchronized (this) {
+      boolean result = _routingDataWriterMap.get(namespace).deleteShardingKey(realm, shardingKey);
+      refreshRoutingData(namespace);
+      return result;
+    }
   }
 
   /**
@@ -251,7 +298,7 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
     // Safe to ignore the callback if any of the maps are null.
     // If routingDataMap is null, then it will be populated by the constructor anyway
     // If routingDataMap is not null, then it's safe for the callback function to update it
-    if (_routingZkAddressMap == null || _routingDataMap == null || _realmToShardingKeysMap == null
+    if (_routingZkAddressMap == null || _realmToShardingKeysMap == null
         || _routingDataReaderMap == null || _routingDataWriterMap == null) {
       LOG.warn(
           "refreshRoutingData callback called before ZKMetadataStoreDirectory was fully initialized. Skipping refresh!");
@@ -262,17 +309,22 @@ public class ZkMetadataStoreDirectory implements MetadataStoreDirectory, Routing
     if (!_routingZkAddressMap.containsKey(namespace)) {
       LOG.error(
           "Failed to refresh internally-cached routing data! Namespace not found: " + namespace);
+      return;
     }
 
+    Map<String, List<String>> rawRoutingData;
     try {
-      Map<String, List<String>> rawRoutingData =
-          _routingDataReaderMap.get(namespace).getRoutingData();
+      rawRoutingData = _routingDataReaderMap.get(namespace).getRoutingData();
       _realmToShardingKeysMap.put(namespace, rawRoutingData);
-
-      MetadataStoreRoutingData routingData = new TrieRoutingData(rawRoutingData);
-      _routingDataMap.put(namespace, routingData);
     } catch (InvalidRoutingDataException e) {
       LOG.error("Failed to refresh cached routing data for namespace {}", namespace, e);
+      return;
+    }
+
+    try {
+      _routingDataMap.put(namespace, new TrieRoutingData(rawRoutingData));
+    } catch (InvalidRoutingDataException e) {
+      LOG.warn("TrieRoutingData is not created for namespace {}", namespace, e);
     }
   }
 
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/MetadataStoreRoutingDataWriter.java b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/MetadataStoreRoutingDataWriter.java
index 349bbd0..02ce60b 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/MetadataStoreRoutingDataWriter.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/MetadataStoreRoutingDataWriter.java
@@ -63,7 +63,7 @@ public interface MetadataStoreRoutingDataWriter {
    * Sets (overwrites) the routing data with the given <realm, list of sharding keys> mapping.
    * WARNING: This overwrites all existing routing data. Use with care!
    * @param routingData
-   * @return
+   * @return true if successful; false otherwise.
    */
   boolean setRoutingData(Map<String, List<String>> routingData);
 
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataReader.java b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataReader.java
index 9251571..6c75618 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataReader.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataReader.java
@@ -99,7 +99,7 @@ public class ZkRoutingDataReader implements MetadataStoreRoutingDataReader, IZkD
     if (allRealmAddresses != null) {
       for (String realmAddress : allRealmAddresses) {
         ZNRecord record = _zkClient
-            .readData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + realmAddress);
+            .readData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + realmAddress, true);
         if (record != null) {
           List<String> shardingKeys =
               record.getListField(MetadataStoreRoutingConstants.ZNRECORD_LIST_FIELD_KEY);
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataWriter.java b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataWriter.java
index 74cc14c..fddd9ee 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataWriter.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/metadatastore/accessor/ZkRoutingDataWriter.java
@@ -41,9 +41,14 @@ import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.HttpDelete;
 import org.apache.http.client.methods.HttpPut;
 import org.apache.http.client.methods.HttpUriRequest;
+import org.apache.http.entity.ContentType;
+import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.util.EntityUtils;
+import org.codehaus.jackson.JsonGenerationException;
+import org.codehaus.jackson.map.JsonMappingException;
+import org.codehaus.jackson.map.ObjectMapper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -52,6 +57,7 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
   // Time out for http requests that are forwarded to leader instances measured in milliseconds
   private static final int HTTP_REQUEST_FORWARDING_TIMEOUT = 60 * 1000;
   private static final Logger LOG = LoggerFactory.getLogger(ZkRoutingDataWriter.class);
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
 
   private final String _namespace;
   private final HelixZkClient _zkClient;
@@ -113,7 +119,7 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
 
     String urlSuffix =
         constructUrlSuffix(MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, realm);
-    return forwardRequestToLeader(urlSuffix, HttpConstants.RestVerbs.PUT,
+    return buildAndSendRequestToLeader(urlSuffix, HttpConstants.RestVerbs.PUT,
         Response.Status.CREATED.getStatusCode());
   }
 
@@ -128,7 +134,7 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
 
     String urlSuffix =
         constructUrlSuffix(MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, realm);
-    return forwardRequestToLeader(urlSuffix, HttpConstants.RestVerbs.DELETE,
+    return buildAndSendRequestToLeader(urlSuffix, HttpConstants.RestVerbs.DELETE,
         Response.Status.OK.getStatusCode());
   }
 
@@ -144,7 +150,7 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
     String urlSuffix =
         constructUrlSuffix(MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, realm,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_SHARDING_KEYS_ENDPOINT, shardingKey);
-    return forwardRequestToLeader(urlSuffix, HttpConstants.RestVerbs.PUT,
+    return buildAndSendRequestToLeader(urlSuffix, HttpConstants.RestVerbs.PUT,
         Response.Status.CREATED.getStatusCode());
   }
 
@@ -160,7 +166,7 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
     String urlSuffix =
         constructUrlSuffix(MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, realm,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_SHARDING_KEYS_ENDPOINT, shardingKey);
-    return forwardRequestToLeader(urlSuffix, HttpConstants.RestVerbs.DELETE,
+    return buildAndSendRequestToLeader(urlSuffix, HttpConstants.RestVerbs.DELETE,
         Response.Status.OK.getStatusCode());
   }
 
@@ -209,8 +215,23 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
       return true;
     }
 
-    // TODO: Forward the request to leader
-    return true;
+    String leaderHostName = _leaderElection.getCurrentLeaderInfo().getId();
+    String url = leaderHostName + constructUrlSuffix(
+        MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT);
+    HttpPut httpPut = new HttpPut(url);
+    String routingDataJsonString;
+    try {
+      routingDataJsonString = OBJECT_MAPPER.writeValueAsString(routingData);
+    } catch (JsonGenerationException | JsonMappingException e) {
+      throw new IllegalArgumentException(e.getMessage());
+    } catch (IOException e) {
+      LOG.error(
+          "setRoutingData failed before forwarding the request to leader: an exception happened while routingData is converted to json. routingData: {}",
+          routingData, e);
+      return false;
+    }
+    httpPut.setEntity(new StringEntity(routingDataJsonString, ContentType.APPLICATION_JSON));
+    return sendRequestToLeader(httpPut, Response.Status.CREATED.getStatusCode(), leaderHostName);
   }
 
   @Override
@@ -332,12 +353,13 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
     return String.join("", allUrlParameters);
   }
 
-  private boolean forwardRequestToLeader(String urlSuffix, HttpConstants.RestVerbs request_method,
-      int expectedResponseCode) throws IllegalArgumentException {
+  private boolean buildAndSendRequestToLeader(String urlSuffix,
+      HttpConstants.RestVerbs requestMethod, int expectedResponseCode)
+      throws IllegalArgumentException {
     String leaderHostName = _leaderElection.getCurrentLeaderInfo().getId();
     String url = leaderHostName + urlSuffix;
     HttpUriRequest request;
-    switch (request_method) {
+    switch (requestMethod) {
       case PUT:
         request = new HttpPut(url);
         break;
@@ -345,8 +367,7 @@ public class ZkRoutingDataWriter implements MetadataStoreRoutingDataWriter {
         request = new HttpDelete(url);
         break;
       default:
-        LOG.error("Unsupported request_method: " + request_method.name());
-        return false;
+        throw new IllegalArgumentException("Unsupported requestMethod: " + requestMethod.name());
     }
 
     return sendRequestToLeader(request, expectedResponseCode, leaderHostName);
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/metadatastore/MetadataStoreDirectoryAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/metadatastore/MetadataStoreDirectoryAccessor.java
index 0763ec1..f20fd9a 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/metadatastore/MetadataStoreDirectoryAccessor.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/metadatastore/MetadataStoreDirectoryAccessor.java
@@ -19,18 +19,22 @@ package org.apache.helix.rest.server.resources.metadatastore;
  * under the License.
  */
 
+import java.io.IOException;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.stream.Collectors;
 import javax.annotation.PostConstruct;
+import javax.ws.rs.Consumes;
 import javax.ws.rs.DELETE;
 import javax.ws.rs.GET;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.QueryParam;
+import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 
 import com.google.common.collect.ImmutableMap;
@@ -44,6 +48,9 @@ import org.apache.helix.rest.metadatastore.ZkMetadataStoreDirectory;
 import org.apache.helix.rest.metadatastore.datamodel.MetadataStoreShardingKey;
 import org.apache.helix.rest.metadatastore.datamodel.MetadataStoreShardingKeysByRealm;
 import org.apache.helix.rest.server.resources.AbstractResource;
+import org.codehaus.jackson.JsonParseException;
+import org.codehaus.jackson.map.JsonMappingException;
+import org.codehaus.jackson.type.TypeReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -220,6 +227,24 @@ public class MetadataStoreDirectoryAccessor extends AbstractResource {
     return JSONRepresentation(responseMap);
   }
 
+  @PUT
+  @Path("/routing-data")
+  @Consumes(MediaType.APPLICATION_JSON)
+  public Response setRoutingData(String jsonContent) {
+    try {
+      Map<String, List<String>> routingData =
+          OBJECT_MAPPER.readValue(jsonContent, new TypeReference<HashMap<String, List<String>>>() {
+          });
+      _metadataStoreDirectory.setNamespaceRoutingData(_namespace, routingData);
+    } catch (JsonMappingException | JsonParseException | IllegalArgumentException e) {
+      return badRequest(e.getMessage());
+    } catch (IOException e) {
+      return serverError(e);
+    }
+
+    return created();
+  }
+
   /**
    * Gets all path-based sharding keys for a queried realm at endpoint:
    * "GET /metadata-store-realms/{realm}/sharding-keys"
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/TestZkMetadataStoreDirectory.java b/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/TestZkMetadataStoreDirectory.java
index 8eddcea..df84754 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/TestZkMetadataStoreDirectory.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/TestZkMetadataStoreDirectory.java
@@ -35,6 +35,7 @@ import org.apache.helix.msdcommon.exception.InvalidRoutingDataException;
 import org.apache.helix.rest.server.AbstractTestClass;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.zkclient.ZkClient;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -62,9 +63,11 @@ public class TestZkMetadataStoreDirectory extends AbstractTestClass {
   private MetadataStoreDirectory _metadataStoreDirectory;
 
   @BeforeClass
-  public void beforeClass() throws InvalidRoutingDataException {
+  public void beforeClass() throws Exception {
     _zkList = new ArrayList<>(ZK_SERVER_MAP.keySet());
 
+    clearRoutingData();
+
     // Populate routingZkAddrMap
     _routingZkAddrMap = new LinkedHashMap<>();
     int namespaceIndex = 0;
@@ -111,10 +114,9 @@ public class TestZkMetadataStoreDirectory extends AbstractTestClass {
   }
 
   @AfterClass
-  public void afterClass() {
+  public void afterClass() throws Exception {
     _metadataStoreDirectory.close();
-    _zkList.forEach(zk -> ZK_SERVER_MAP.get(zk).getZkClient()
-        .deleteRecursive(MetadataStoreRoutingConstants.ROUTING_DATA_PATH));
+    clearRoutingData();
     System.clearProperty(MetadataStoreRoutingConstants.MSDS_SERVER_HOSTNAME_KEY);
   }
 
@@ -147,6 +149,42 @@ public class TestZkMetadataStoreDirectory extends AbstractTestClass {
   }
 
   @Test(dependsOnMethods = "testGetAllShardingKeys")
+  public void testGetNamespaceRoutingData() {
+    Map<String, List<String>> routingDataMap = new HashMap<>();
+    routingDataMap.put(TEST_REALM_1, TEST_SHARDING_KEYS_1);
+    routingDataMap.put(TEST_REALM_2, TEST_SHARDING_KEYS_2);
+
+    for (String namespace : _routingZkAddrMap.keySet()) {
+      Assert
+          .assertEquals(_metadataStoreDirectory.getNamespaceRoutingData(namespace), routingDataMap);
+    }
+  }
+
+  @Test(dependsOnMethods = "testGetNamespaceRoutingData")
+  public void testSetNamespaceRoutingData() {
+    Map<String, List<String>> routingDataMap = new HashMap<>();
+    routingDataMap.put(TEST_REALM_1, TEST_SHARDING_KEYS_2);
+    routingDataMap.put(TEST_REALM_2, TEST_SHARDING_KEYS_1);
+
+    for (String namespace : _routingZkAddrMap.keySet()) {
+      _metadataStoreDirectory.setNamespaceRoutingData(namespace, routingDataMap);
+      Assert
+          .assertEquals(_metadataStoreDirectory.getNamespaceRoutingData(namespace), routingDataMap);
+    }
+
+    // Revert it back to the original state
+    Map<String, List<String>> originalRoutingDataMap = new HashMap<>();
+    originalRoutingDataMap.put(TEST_REALM_1, TEST_SHARDING_KEYS_1);
+    originalRoutingDataMap.put(TEST_REALM_2, TEST_SHARDING_KEYS_2);
+
+    for (String namespace : _routingZkAddrMap.keySet()) {
+      _metadataStoreDirectory.setNamespaceRoutingData(namespace, originalRoutingDataMap);
+      Assert.assertEquals(_metadataStoreDirectory.getNamespaceRoutingData(namespace),
+          originalRoutingDataMap);
+    }
+  }
+
+  @Test(dependsOnMethods = "testGetNamespaceRoutingData")
   public void testGetAllShardingKeysInRealm() {
     for (String namespace : _routingZkAddrMap.keySet()) {
       // Test two realms independently
@@ -277,4 +315,27 @@ public class TestZkMetadataStoreDirectory extends AbstractTestClass {
       return false;
     }, TestHelper.WAIT_DURATION));
   }
+
+  private void clearRoutingData() throws Exception {
+    Assert.assertTrue(TestHelper.verify(() -> {
+      for (String zk : _zkList) {
+        ZkClient zkClient = ZK_SERVER_MAP.get(zk).getZkClient();
+        if (zkClient.exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+          for (String zkRealm : zkClient
+              .getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+            zkClient.delete(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + zkRealm);
+          }
+        }
+      }
+
+      for (String zk : _zkList) {
+        ZkClient zkClient = ZK_SERVER_MAP.get(zk).getZkClient();
+        if (zkClient.exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH) && !zkClient
+            .getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH).isEmpty()) {
+          return false;
+        }
+      }
+      return true;
+    }, TestHelper.WAIT_DURATION), "Routing data path should be deleted after the tests.");
+  }
 }
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataReader.java b/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataReader.java
index 63a013b..4c95445 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataReader.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataReader.java
@@ -24,12 +24,12 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
-import org.apache.helix.AccessOption;
 import org.apache.helix.TestHelper;
 import org.apache.helix.msdcommon.constant.MetadataStoreRoutingConstants;
 import org.apache.helix.msdcommon.exception.InvalidRoutingDataException;
 import org.apache.helix.rest.server.AbstractTestClass;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.zkclient.ZkClient;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.AfterMethod;
@@ -38,24 +38,25 @@ import org.testng.annotations.Test;
 
 
 public class TestZkRoutingDataReader extends AbstractTestClass {
-  private static final String DUMMY_NAMESPACE = "NAMESPACE";
   private MetadataStoreRoutingDataReader _zkRoutingDataReader;
+  private ZkClient _zkClient;
 
   @BeforeClass
   public void beforeClass() throws Exception {
-    deleteRoutingDataPath();
-    _zkRoutingDataReader = new ZkRoutingDataReader(DUMMY_NAMESPACE, ZK_ADDR, null);
+    _zkClient = ZK_SERVER_MAP.get(_zkAddrTestNS).getZkClient();
+    _zkRoutingDataReader = new ZkRoutingDataReader(TEST_NAMESPACE, _zkAddrTestNS, null);
+    clearRoutingDataPath();
   }
 
   @AfterClass
   public void afterClass() throws Exception {
     _zkRoutingDataReader.close();
-    deleteRoutingDataPath();
+    clearRoutingDataPath();
   }
 
   @AfterMethod
-  public void afterMethod() {
-    _baseAccessor.remove(MetadataStoreRoutingConstants.ROUTING_DATA_PATH, AccessOption.PERSISTENT);
+  public void afterMethod() throws Exception {
+    clearRoutingDataPath();
   }
 
   @Test
@@ -75,10 +76,14 @@ public class TestZkRoutingDataReader extends AbstractTestClass {
         .setListField(MetadataStoreRoutingConstants.ZNRECORD_LIST_FIELD_KEY, testShardingKeys2);
 
     // Add both nodes as children nodes to ZkRoutingDataReader.ROUTING_DATA_PATH
-    _baseAccessor.create(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress1",
-        testZnRecord1, AccessOption.PERSISTENT);
-    _baseAccessor.create(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress2",
-        testZnRecord2, AccessOption.PERSISTENT);
+    _zkClient
+        .createPersistent(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress1");
+    _zkClient.writeData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress1",
+        testZnRecord1);
+    _zkClient
+        .createPersistent(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress2");
+    _zkClient.writeData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress2",
+        testZnRecord2);
 
     try {
       Map<String, List<String>> routingData = _zkRoutingDataReader.getRoutingData();
@@ -90,22 +95,8 @@ public class TestZkRoutingDataReader extends AbstractTestClass {
     }
   }
 
-  @Test
-  public void testGetRoutingDataMissingMSRD() {
-    try {
-      _zkRoutingDataReader.getRoutingData();
-      Assert.fail("Expecting InvalidRoutingDataException");
-    } catch (InvalidRoutingDataException e) {
-      Assert.assertTrue(e.getMessage().contains(
-          "Routing data directory ZNode " + MetadataStoreRoutingConstants.ROUTING_DATA_PATH
-              + " does not exist. Routing ZooKeeper address: " + ZK_ADDR));
-    }
-  }
-
-  @Test
+  @Test(dependsOnMethods = "testGetRoutingData")
   public void testGetRoutingDataMissingMSRDChildren() {
-    _baseAccessor.create(MetadataStoreRoutingConstants.ROUTING_DATA_PATH, new ZNRecord("test"),
-        AccessOption.PERSISTENT);
     try {
       Map<String, List<String>> routingData = _zkRoutingDataReader.getRoutingData();
       Assert.assertEquals(routingData.size(), 0);
@@ -114,13 +105,15 @@ public class TestZkRoutingDataReader extends AbstractTestClass {
     }
   }
 
-  @Test
+  @Test(dependsOnMethods = "testGetRoutingData")
   public void testGetRoutingDataMSRDChildEmptyValue() {
     ZNRecord testZnRecord1 = new ZNRecord("testZnRecord1");
     testZnRecord1.setListField(MetadataStoreRoutingConstants.ZNRECORD_LIST_FIELD_KEY,
         Collections.emptyList());
-    _baseAccessor.create(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress1",
-        testZnRecord1, AccessOption.PERSISTENT);
+    _zkClient
+        .createPersistent(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress1");
+    _zkClient.writeData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/testRealmAddress1",
+        testZnRecord1);
     try {
       Map<String, List<String>> routingData = _zkRoutingDataReader.getRoutingData();
       Assert.assertEquals(routingData.size(), 1);
@@ -130,17 +123,14 @@ public class TestZkRoutingDataReader extends AbstractTestClass {
     }
   }
 
-  private void deleteRoutingDataPath() throws Exception {
+  private void clearRoutingDataPath() throws Exception {
     Assert.assertTrue(TestHelper.verify(() -> {
-      ZK_SERVER_MAP.get(ZK_ADDR).getZkClient()
-          .deleteRecursively(MetadataStoreRoutingConstants.ROUTING_DATA_PATH);
-
-      if (ZK_SERVER_MAP.get(ZK_ADDR).getZkClient()
-          .exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
-        return false;
+      for (String zkRealm : _zkClient
+          .getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+        _zkClient.delete(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + zkRealm);
       }
 
-      return true;
+      return _zkClient.getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH).isEmpty();
     }, TestHelper.WAIT_DURATION), "Routing data path should be deleted after the tests.");
   }
 }
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataWriter.java b/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataWriter.java
index 069931f..31db291 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataWriter.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/metadatastore/accessor/TestZkRoutingDataWriter.java
@@ -25,22 +25,25 @@ import java.util.List;
 import java.util.Map;
 
 import com.google.common.collect.ImmutableMap;
-import org.apache.helix.AccessOption;
+import org.apache.helix.TestHelper;
 import org.apache.helix.msdcommon.constant.MetadataStoreRoutingConstants;
+import org.apache.helix.rest.common.HttpConstants;
 import org.apache.helix.rest.server.AbstractTestClass;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.zkclient.ZkClient;
 import org.apache.http.client.methods.HttpUriRequest;
-import org.junit.Assert;
+import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 
 public class TestZkRoutingDataWriter extends AbstractTestClass {
-  private static final String DUMMY_NAMESPACE = "NAMESPACE";
   private static final String DUMMY_REALM = "REALM";
-  private static final String DUMMY_SHARDING_KEY = "SHARDING_KEY";
+  private static final String DUMMY_SHARDING_KEY = "/DUMMY/SHARDING/KEY";
+
   private MetadataStoreRoutingDataWriter _zkRoutingDataWriter;
+  private ZkClient _zkClient;
 
   // MockWriter is used for testing request forwarding features in non-leader situations
   class MockWriter extends ZkRoutingDataWriter {
@@ -60,43 +63,41 @@ public class TestZkRoutingDataWriter extends AbstractTestClass {
   }
 
   @BeforeClass
-  public void beforeClass() {
-    _baseAccessor.remove(MetadataStoreRoutingConstants.ROUTING_DATA_PATH, AccessOption.PERSISTENT);
+  public void beforeClass() throws Exception {
+    _zkClient = ZK_SERVER_MAP.get(_zkAddrTestNS).getZkClient();
     System.setProperty(MetadataStoreRoutingConstants.MSDS_SERVER_HOSTNAME_KEY,
         getBaseUri().getHost() + ":" + getBaseUri().getPort());
-    _zkRoutingDataWriter = new ZkRoutingDataWriter(DUMMY_NAMESPACE, ZK_ADDR);
+    _zkRoutingDataWriter = new ZkRoutingDataWriter(TEST_NAMESPACE, _zkAddrTestNS);
+    clearRoutingDataPath();
   }
 
   @AfterClass
-  public void afterClass() {
-    _baseAccessor.remove(MetadataStoreRoutingConstants.ROUTING_DATA_PATH, AccessOption.PERSISTENT);
+  public void afterClass() throws Exception {
     System.clearProperty(MetadataStoreRoutingConstants.MSDS_SERVER_HOSTNAME_KEY);
     _zkRoutingDataWriter.close();
+    clearRoutingDataPath();
   }
 
   @Test
   public void testAddMetadataStoreRealm() {
     _zkRoutingDataWriter.addMetadataStoreRealm(DUMMY_REALM);
-    ZNRecord znRecord = _baseAccessor
-        .get(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM, null,
-            AccessOption.PERSISTENT);
+    ZNRecord znRecord =
+        _zkClient.readData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM);
     Assert.assertNotNull(znRecord);
   }
 
   @Test(dependsOnMethods = "testAddMetadataStoreRealm")
   public void testDeleteMetadataStoreRealm() {
     _zkRoutingDataWriter.deleteMetadataStoreRealm(DUMMY_REALM);
-    Assert.assertFalse(_baseAccessor
-        .exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM,
-            AccessOption.PERSISTENT));
+    Assert.assertFalse(
+        _zkClient.exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM));
   }
 
   @Test(dependsOnMethods = "testDeleteMetadataStoreRealm")
   public void testAddShardingKey() {
     _zkRoutingDataWriter.addShardingKey(DUMMY_REALM, DUMMY_SHARDING_KEY);
-    ZNRecord znRecord = _baseAccessor
-        .get(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM, null,
-            AccessOption.PERSISTENT);
+    ZNRecord znRecord =
+        _zkClient.readData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM);
     Assert.assertNotNull(znRecord);
     Assert.assertTrue(znRecord.getListField(MetadataStoreRoutingConstants.ZNRECORD_LIST_FIELD_KEY)
         .contains(DUMMY_SHARDING_KEY));
@@ -105,9 +106,8 @@ public class TestZkRoutingDataWriter extends AbstractTestClass {
   @Test(dependsOnMethods = "testAddShardingKey")
   public void testDeleteShardingKey() {
     _zkRoutingDataWriter.deleteShardingKey(DUMMY_REALM, DUMMY_SHARDING_KEY);
-    ZNRecord znRecord = _baseAccessor
-        .get(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM, null,
-            AccessOption.PERSISTENT);
+    ZNRecord znRecord =
+        _zkClient.readData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM);
     Assert.assertNotNull(znRecord);
     Assert.assertFalse(znRecord.getListField(MetadataStoreRoutingConstants.ZNRECORD_LIST_FIELD_KEY)
         .contains(DUMMY_SHARDING_KEY));
@@ -118,9 +118,8 @@ public class TestZkRoutingDataWriter extends AbstractTestClass {
     Map<String, List<String>> testRoutingDataMap =
         ImmutableMap.of(DUMMY_REALM, Collections.singletonList(DUMMY_SHARDING_KEY));
     _zkRoutingDataWriter.setRoutingData(testRoutingDataMap);
-    ZNRecord znRecord = _baseAccessor
-        .get(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM, null,
-            AccessOption.PERSISTENT);
+    ZNRecord znRecord =
+        _zkClient.readData(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + DUMMY_REALM);
     Assert.assertNotNull(znRecord);
     Assert.assertEquals(
         znRecord.getListField(MetadataStoreRoutingConstants.ZNRECORD_LIST_FIELD_KEY).size(), 1);
@@ -130,63 +129,93 @@ public class TestZkRoutingDataWriter extends AbstractTestClass {
 
   @Test(dependsOnMethods = "testSetRoutingData")
   public void testAddMetadataStoreRealmNonLeader() {
-    MockWriter mockWriter = new MockWriter(DUMMY_NAMESPACE, ZK_ADDR);
+    MockWriter mockWriter = new MockWriter(TEST_NAMESPACE, _zkAddrTestNS);
     mockWriter.addMetadataStoreRealm(DUMMY_REALM);
-    Assert.assertEquals("PUT", mockWriter.calledRequest.getMethod());
+    Assert.assertEquals(mockWriter.calledRequest.getMethod(), HttpConstants.RestVerbs.PUT.name());
     List<String> expectedUrlParams = Arrays
-        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, DUMMY_NAMESPACE,
+        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, TEST_NAMESPACE,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, DUMMY_REALM);
     String expectedUrl =
         getBaseUri().toString() + String.join("/", expectedUrlParams).replaceAll("//", "/")
             .substring(1);
-    Assert.assertEquals(expectedUrl, mockWriter.calledRequest.getURI().toString());
+    Assert.assertEquals(mockWriter.calledRequest.getURI().toString(), expectedUrl);
     mockWriter.close();
   }
 
   @Test(dependsOnMethods = "testAddMetadataStoreRealmNonLeader")
   public void testDeleteMetadataStoreRealmNonLeader() {
-    MockWriter mockWriter = new MockWriter(DUMMY_NAMESPACE, ZK_ADDR);
+    MockWriter mockWriter = new MockWriter(TEST_NAMESPACE, _zkAddrTestNS);
     mockWriter.deleteMetadataStoreRealm(DUMMY_REALM);
-    Assert.assertEquals("DELETE", mockWriter.calledRequest.getMethod());
+    Assert
+        .assertEquals(mockWriter.calledRequest.getMethod(), HttpConstants.RestVerbs.DELETE.name());
     List<String> expectedUrlParams = Arrays
-        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, DUMMY_NAMESPACE,
+        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, TEST_NAMESPACE,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, DUMMY_REALM);
     String expectedUrl =
         getBaseUri().toString() + String.join("/", expectedUrlParams).replaceAll("//", "/")
             .substring(1);
-    Assert.assertEquals(expectedUrl, mockWriter.calledRequest.getURI().toString());
+    Assert.assertEquals(mockWriter.calledRequest.getURI().toString(), expectedUrl);
     mockWriter.close();
   }
 
   @Test(dependsOnMethods = "testDeleteMetadataStoreRealmNonLeader")
   public void testAddShardingKeyNonLeader() {
-    MockWriter mockWriter = new MockWriter(DUMMY_NAMESPACE, ZK_ADDR);
+    MockWriter mockWriter = new MockWriter(TEST_NAMESPACE, _zkAddrTestNS);
     mockWriter.addShardingKey(DUMMY_REALM, DUMMY_SHARDING_KEY);
-    Assert.assertEquals("PUT", mockWriter.calledRequest.getMethod());
+    Assert.assertEquals(mockWriter.calledRequest.getMethod(), HttpConstants.RestVerbs.PUT.name());
     List<String> expectedUrlParams = Arrays
-        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, DUMMY_NAMESPACE,
+        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, TEST_NAMESPACE,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, DUMMY_REALM,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_SHARDING_KEYS_ENDPOINT, DUMMY_SHARDING_KEY);
     String expectedUrl =
         getBaseUri().toString() + String.join("/", expectedUrlParams).replaceAll("//", "/")
             .substring(1);
-    Assert.assertEquals(expectedUrl, mockWriter.calledRequest.getURI().toString());
+    Assert.assertEquals(mockWriter.calledRequest.getURI().toString(), expectedUrl);
     mockWriter.close();
   }
 
   @Test(dependsOnMethods = "testAddShardingKeyNonLeader")
   public void testDeleteShardingKeyNonLeader() {
-    MockWriter mockWriter = new MockWriter(DUMMY_NAMESPACE, ZK_ADDR);
+    MockWriter mockWriter = new MockWriter(TEST_NAMESPACE, _zkAddrTestNS);
     mockWriter.deleteShardingKey(DUMMY_REALM, DUMMY_SHARDING_KEY);
-    Assert.assertEquals("DELETE", mockWriter.calledRequest.getMethod());
+    Assert
+        .assertEquals(mockWriter.calledRequest.getMethod(), HttpConstants.RestVerbs.DELETE.name());
     List<String> expectedUrlParams = Arrays
-        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, DUMMY_NAMESPACE,
+        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, TEST_NAMESPACE,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_REALMS_ENDPOINT, DUMMY_REALM,
             MetadataStoreRoutingConstants.MSDS_GET_ALL_SHARDING_KEYS_ENDPOINT, DUMMY_SHARDING_KEY);
     String expectedUrl =
         getBaseUri().toString() + String.join("/", expectedUrlParams).replaceAll("//", "/")
             .substring(1);
-    Assert.assertEquals(expectedUrl, mockWriter.calledRequest.getURI().toString());
+    Assert.assertEquals(mockWriter.calledRequest.getURI().toString(), expectedUrl);
     mockWriter.close();
   }
+
+  @Test(dependsOnMethods = "testDeleteShardingKeyNonLeader")
+  public void testSetRoutingDataNonLeader() {
+    MockWriter mockWriter = new MockWriter(TEST_NAMESPACE, _zkAddrTestNS);
+    Map<String, List<String>> testRoutingDataMap =
+        ImmutableMap.of(DUMMY_REALM, Collections.singletonList(DUMMY_SHARDING_KEY));
+    mockWriter.setRoutingData(testRoutingDataMap);
+    Assert.assertEquals(mockWriter.calledRequest.getMethod(), HttpConstants.RestVerbs.PUT.name());
+    List<String> expectedUrlParams = Arrays
+        .asList(MetadataStoreRoutingConstants.MSDS_NAMESPACES_URL_PREFIX, TEST_NAMESPACE,
+            MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT);
+    String expectedUrl =
+        getBaseUri().toString() + String.join("/", expectedUrlParams).replaceAll("//", "/")
+            .substring(1);
+    Assert.assertEquals(mockWriter.calledRequest.getURI().toString(), expectedUrl);
+    mockWriter.close();
+  }
+
+  private void clearRoutingDataPath() throws Exception {
+    Assert.assertTrue(TestHelper.verify(() -> {
+      for (String zkRealm : _zkClient
+          .getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+        _zkClient.delete(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + zkRealm);
+      }
+
+      return _zkClient.getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH).isEmpty();
+    }, TestHelper.WAIT_DURATION), "Routing data path should be deleted after the tests.");
+  }
 }
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/MetadataStoreDirectoryAccessorTestBase.java b/helix-rest/src/test/java/org/apache/helix/rest/server/MetadataStoreDirectoryAccessorTestBase.java
index 7cebbf3..f2ed433 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/server/MetadataStoreDirectoryAccessorTestBase.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/server/MetadataStoreDirectoryAccessorTestBase.java
@@ -23,6 +23,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import org.apache.helix.TestHelper;
@@ -32,6 +33,7 @@ import org.apache.helix.rest.metadatastore.accessor.MetadataStoreRoutingDataRead
 import org.apache.helix.rest.metadatastore.accessor.ZkRoutingDataReader;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.zkclient.ZkClient;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -62,7 +64,7 @@ public class MetadataStoreDirectoryAccessorTestBase extends AbstractTestClass {
   public void beforeClass() throws Exception {
     _zkList = new ArrayList<>(ZK_SERVER_MAP.keySet());
 
-    deleteRoutingDataPath();
+    clearRoutingData();
 
     // Write dummy mappings in ZK
     // Create a node that represents a realm address and add 3 sharding keys to it
@@ -100,21 +102,29 @@ public class MetadataStoreDirectoryAccessorTestBase extends AbstractTestClass {
   @AfterClass
   public void afterClass() throws Exception {
     System.clearProperty(MetadataStoreRoutingConstants.MSDS_SERVER_HOSTNAME_KEY);
-    deleteRoutingDataPath();
+    _routingDataReader.close();
+    clearRoutingData();
   }
 
-  protected void deleteRoutingDataPath() throws Exception {
+  protected void clearRoutingData() throws Exception {
     Assert.assertTrue(TestHelper.verify(() -> {
-      _zkList.forEach(zk -> ZK_SERVER_MAP.get(zk).getZkClient()
-          .deleteRecursively(MetadataStoreRoutingConstants.ROUTING_DATA_PATH));
+      for (String zk : _zkList) {
+        ZkClient zkClient = ZK_SERVER_MAP.get(zk).getZkClient();
+        if (zkClient.exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+          for (String zkRealm : zkClient
+              .getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+            zkClient.delete(MetadataStoreRoutingConstants.ROUTING_DATA_PATH + "/" + zkRealm);
+          }
+        }
+      }
 
       for (String zk : _zkList) {
-        if (ZK_SERVER_MAP.get(zk).getZkClient()
-            .exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH)) {
+        ZkClient zkClient = ZK_SERVER_MAP.get(zk).getZkClient();
+        if (zkClient.exists(MetadataStoreRoutingConstants.ROUTING_DATA_PATH) && !zkClient
+            .getChildren(MetadataStoreRoutingConstants.ROUTING_DATA_PATH).isEmpty()) {
           return false;
         }
       }
-
       return true;
     }, TestHelper.WAIT_DURATION), "Routing data path should be deleted after the tests.");
   }
@@ -129,4 +139,8 @@ public class MetadataStoreDirectoryAccessorTestBase extends AbstractTestClass {
   protected Set<String> getAllShardingKeysInTestRealm1() throws InvalidRoutingDataException {
     return new HashSet<>(_routingDataReader.getRoutingData().get(TEST_REALM_1));
   }
+
+  protected Map<String, List<String>> getRawRoutingData() throws InvalidRoutingDataException {
+    return _routingDataReader.getRoutingData();
+  }
 }
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestMSDAccessorLeaderElection.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestMSDAccessorLeaderElection.java
index 951015b..2a3f0be 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestMSDAccessorLeaderElection.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestMSDAccessorLeaderElection.java
@@ -24,7 +24,9 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.ws.rs.core.Response;
 
@@ -45,8 +47,11 @@ import org.apache.http.HttpResponse;
 import org.apache.http.client.methods.HttpDelete;
 import org.apache.http.client.methods.HttpPut;
 import org.apache.http.client.methods.HttpUriRequest;
+import org.apache.http.entity.ContentType;
+import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClients;
+import org.codehaus.jackson.map.ObjectMapper;
 import org.glassfish.jersey.server.ResourceConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -120,8 +125,9 @@ public class TestMSDAccessorLeaderElection extends MetadataStoreDirectoryAccesso
     Set<String> expectedRealmsSet = getAllRealms();
     Assert.assertFalse(expectedRealmsSet.contains(TEST_REALM_3),
         "Metadata store directory should not have realm: " + TEST_REALM_3);
-    sendRequestAndValidate("/metadata-store-realms/" + TEST_REALM_3, HttpConstants.RestVerbs.PUT,
-        Response.Status.CREATED.getStatusCode());
+    HttpUriRequest request =
+        buildRequest("/metadata-store-realms/" + TEST_REALM_3, HttpConstants.RestVerbs.PUT, "");
+    sendRequestAndValidate(request, Response.Status.CREATED.getStatusCode());
     expectedRealmsSet.add(TEST_REALM_3);
     Assert.assertEquals(getAllRealms(), expectedRealmsSet);
     MockMetadataStoreDirectoryAccessor._mockMSDInstance.close();
@@ -131,8 +137,9 @@ public class TestMSDAccessorLeaderElection extends MetadataStoreDirectoryAccesso
   public void testDeleteMetadataStoreRealmRequestForwarding()
       throws InvalidRoutingDataException, IOException {
     Set<String> expectedRealmsSet = getAllRealms();
-    sendRequestAndValidate("/metadata-store-realms/" + TEST_REALM_3, HttpConstants.RestVerbs.DELETE,
-        Response.Status.OK.getStatusCode());
+    HttpUriRequest request =
+        buildRequest("/metadata-store-realms/" + TEST_REALM_3, HttpConstants.RestVerbs.DELETE, "");
+    sendRequestAndValidate(request, Response.Status.OK.getStatusCode());
     expectedRealmsSet.remove(TEST_REALM_3);
     Assert.assertEquals(getAllRealms(), expectedRealmsSet);
     MockMetadataStoreDirectoryAccessor._mockMSDInstance.close();
@@ -144,9 +151,10 @@ public class TestMSDAccessorLeaderElection extends MetadataStoreDirectoryAccesso
     Set<String> expectedShardingKeysSet = getAllShardingKeysInTestRealm1();
     Assert.assertFalse(expectedShardingKeysSet.contains(TEST_SHARDING_KEY),
         "Realm does not have sharding key: " + TEST_SHARDING_KEY);
-    sendRequestAndValidate(
+    HttpUriRequest request = buildRequest(
         "/metadata-store-realms/" + TEST_REALM_1 + "/sharding-keys/" + TEST_SHARDING_KEY,
-        HttpConstants.RestVerbs.PUT, Response.Status.CREATED.getStatusCode());
+        HttpConstants.RestVerbs.PUT, "");
+    sendRequestAndValidate(request, Response.Status.CREATED.getStatusCode());
     expectedShardingKeysSet.add(TEST_SHARDING_KEY);
     Assert.assertEquals(getAllShardingKeysInTestRealm1(), expectedShardingKeysSet);
     MockMetadataStoreDirectoryAccessor._mockMSDInstance.close();
@@ -156,28 +164,47 @@ public class TestMSDAccessorLeaderElection extends MetadataStoreDirectoryAccesso
   public void testDeleteShardingKeyRequestForwarding()
       throws InvalidRoutingDataException, IOException {
     Set<String> expectedShardingKeysSet = getAllShardingKeysInTestRealm1();
-    sendRequestAndValidate(
+    HttpUriRequest request = buildRequest(
         "/metadata-store-realms/" + TEST_REALM_1 + "/sharding-keys/" + TEST_SHARDING_KEY,
-        HttpConstants.RestVerbs.DELETE, Response.Status.OK.getStatusCode());
+        HttpConstants.RestVerbs.DELETE, "");
+    sendRequestAndValidate(request, Response.Status.OK.getStatusCode());
     expectedShardingKeysSet.remove(TEST_SHARDING_KEY);
     Assert.assertEquals(getAllShardingKeysInTestRealm1(), expectedShardingKeysSet);
     MockMetadataStoreDirectoryAccessor._mockMSDInstance.close();
   }
 
-  private void sendRequestAndValidate(String urlSuffix, HttpConstants.RestVerbs requestMethod,
-      int expectedResponseCode) throws IllegalArgumentException, IOException {
+  @Test(dependsOnMethods = "testDeleteShardingKeyRequestForwarding")
+  public void testSetRoutingDataRequestForwarding()
+      throws InvalidRoutingDataException, IOException {
+    Map<String, List<String>> routingData = new HashMap<>();
+    routingData.put(TEST_REALM_1, TEST_SHARDING_KEYS_2);
+    routingData.put(TEST_REALM_2, TEST_SHARDING_KEYS_1);
+    String routingDataString = new ObjectMapper().writeValueAsString(routingData);
+    HttpUriRequest request =
+        buildRequest(MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT,
+            HttpConstants.RestVerbs.PUT, routingDataString);
+    sendRequestAndValidate(request, Response.Status.CREATED.getStatusCode());
+    Assert.assertEquals(getRawRoutingData(), routingData);
+    MockMetadataStoreDirectoryAccessor._mockMSDInstance.close();
+  }
+
+  private HttpUriRequest buildRequest(String urlSuffix, HttpConstants.RestVerbs requestMethod,
+      String jsonEntity) {
     String url = _mockBaseUri + TEST_NAMESPACE_URI_PREFIX + MOCK_URL_PREFIX + urlSuffix;
-    HttpUriRequest request;
     switch (requestMethod) {
       case PUT:
-        request = new HttpPut(url);
-        break;
+        HttpPut httpPut = new HttpPut(url);
+        httpPut.setEntity(new StringEntity(jsonEntity, ContentType.APPLICATION_JSON));
+        return httpPut;
       case DELETE:
-        request = new HttpDelete(url);
-        break;
+        return new HttpDelete(url);
       default:
         throw new IllegalArgumentException("Unsupported requestMethod: " + requestMethod);
     }
+  }
+
+  private void sendRequestAndValidate(HttpUriRequest request, int expectedResponseCode)
+      throws IllegalArgumentException, IOException {
     HttpResponse response = _httpClient.execute(request);
     Assert.assertEquals(response.getStatusLine().getStatusCode(), expectedResponseCode);
 
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestMetadataStoreDirectoryAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestMetadataStoreDirectoryAccessor.java
index 94641ff..570ab90 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestMetadataStoreDirectoryAccessor.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestMetadataStoreDirectoryAccessor.java
@@ -21,6 +21,7 @@ package org.apache.helix.rest.server;
 
 import java.io.IOException;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -452,6 +453,43 @@ public class TestMetadataStoreDirectoryAccessor extends MetadataStoreDirectoryAc
     Assert.assertEquals(getAllShardingKeysInTestRealm1(), expectedShardingKeysSet);
   }
 
+  @Test(dependsOnMethods = "testDeleteShardingKey")
+  public void testSetRoutingData() throws InvalidRoutingDataException, IOException {
+    Map<String, List<String>> routingData = new HashMap<>();
+    routingData.put(TEST_REALM_1, TEST_SHARDING_KEYS_2);
+    routingData.put(TEST_REALM_2, TEST_SHARDING_KEYS_1);
+    String routingDataString = OBJECT_MAPPER.writeValueAsString(routingData);
+
+    Map<String, String> badFormatRoutingData = new HashMap<>();
+    badFormatRoutingData.put(TEST_REALM_1, TEST_REALM_2);
+    badFormatRoutingData.put(TEST_REALM_2, TEST_REALM_1);
+    String badFormatRoutingDataString = OBJECT_MAPPER.writeValueAsString(badFormatRoutingData);
+
+    // Request that gets not found response.
+    put("/namespaces/non-existing-namespace"
+            + MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT, null,
+        Entity.entity(routingDataString, MediaType.APPLICATION_JSON_TYPE),
+        Response.Status.NOT_FOUND.getStatusCode());
+
+    put(TEST_NAMESPACE_URI_PREFIX
+            + MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT, null,
+        Entity.entity("?", MediaType.APPLICATION_JSON_TYPE),
+        Response.Status.BAD_REQUEST.getStatusCode());
+
+    put(TEST_NAMESPACE_URI_PREFIX
+            + MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT, null,
+        Entity.entity(badFormatRoutingDataString, MediaType.APPLICATION_JSON_TYPE),
+        Response.Status.BAD_REQUEST.getStatusCode());
+
+    // Successful request.
+    put(TEST_NAMESPACE_URI_PREFIX
+            + MetadataStoreRoutingConstants.MSDS_GET_ALL_ROUTING_DATA_ENDPOINT, null,
+        Entity.entity(routingDataString, MediaType.APPLICATION_JSON_TYPE),
+        Response.Status.CREATED.getStatusCode());
+
+    Assert.assertEquals(getRawRoutingData(), routingData);
+  }
+
   private void verifyRealmShardingKeys(String responseBody) throws IOException {
     // It is safe to cast the object and suppress warnings.
     @SuppressWarnings("unchecked")
diff --git a/metadata-store-directory-common/src/main/java/org/apache/helix/msdcommon/datamodel/TrieRoutingData.java b/metadata-store-directory-common/src/main/java/org/apache/helix/msdcommon/datamodel/TrieRoutingData.java
index 0f53c23..94bb331 100644
--- a/metadata-store-directory-common/src/main/java/org/apache/helix/msdcommon/datamodel/TrieRoutingData.java
+++ b/metadata-store-directory-common/src/main/java/org/apache/helix/msdcommon/datamodel/TrieRoutingData.java
@@ -167,8 +167,8 @@ public class TrieRoutingData implements MetadataStoreRoutingData {
   }
 
   /*
-   * Checks for the edge case when the only sharding key in provided routing data is the delimiter
-   * or an empty string. When this is the case, the trie is valid and contains only one node, which
+   * Checks for the edge case when the only sharding key in provided routing data is the delimiter.
+   * When this is the case, the trie is valid and contains only one node, which
    * is the root node, and the root node is a leaf node with a realm address associated with it.
    * @param routingData - a mapping from "sharding keys" to "realm addresses" to be parsed into a
    *          trie


Mime
View raw message