Examples with HelixExternalViewBasedRouting com.linkedin.pinot.routing.HelixExternalViewBasedRouting used on opensource projects

Example 1 with HelixExternalViewBasedRouting

use of com.linkedin.pinot.routing.HelixExternalViewBasedRouting in project pinot by linkedin.

the class HelixBrokerStarterTest method testResourceAndTagAssignment.

@Test
public void testResourceAndTagAssignment() throws Exception {
    IdealState idealState;
    Assert.assertEquals(_helixAdmin.getInstancesInClusterWithTag(HELIX_CLUSTER_NAME, "DefaultTenant_BROKER").size(), 6);
    idealState = _helixAdmin.getResourceIdealState(HELIX_CLUSTER_NAME, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE);
    Assert.assertEquals(idealState.getInstanceSet(DINING_TABLE_NAME).size(), SEGMENT_COUNT);
    ExternalView externalView = _helixAdmin.getResourceExternalView(HELIX_CLUSTER_NAME, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE);
    Assert.assertEquals(externalView.getStateMap(DINING_TABLE_NAME).size(), SEGMENT_COUNT);
    HelixExternalViewBasedRouting helixExternalViewBasedRouting = _helixBrokerStarter.getHelixExternalViewBasedRouting();
    Field brokerRoutingTableField;
    brokerRoutingTableField = HelixExternalViewBasedRouting.class.getDeclaredField("_brokerRoutingTable");
    brokerRoutingTableField.setAccessible(true);
    final Map<String, List<ServerToSegmentSetMap>> brokerRoutingTable = (Map<String, List<ServerToSegmentSetMap>>) brokerRoutingTableField.get(helixExternalViewBasedRouting);
    // Wait up to 30s for routing table to reach the expected size
    waitForPredicate(new Callable<Boolean>() {

        @Override
        public Boolean call() throws Exception {
            return brokerRoutingTable.size() == 1;
        }
    }, 30000L);
    Assert.assertEquals(Arrays.toString(brokerRoutingTable.keySet().toArray()), "[dining_OFFLINE]");
    final String tableName = "coffee";
    JSONObject buildCreateOfflineTableV2JSON = ControllerRequestBuilderUtil.buildCreateOfflineTableJSON(tableName, "testServer", "testBroker", 1);
    AbstractTableConfig config = AbstractTableConfig.init(buildCreateOfflineTableV2JSON.toString());
    _pinotResourceManager.addTable(config);
    Assert.assertEquals(_helixAdmin.getInstancesInClusterWithTag(HELIX_CLUSTER_NAME, "DefaultTenant_BROKER").size(), 6);
    idealState = _helixAdmin.getResourceIdealState(HELIX_CLUSTER_NAME, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE);
    Assert.assertEquals(idealState.getInstanceSet(COFFEE_TABLE_NAME).size(), SEGMENT_COUNT);
    Assert.assertEquals(idealState.getInstanceSet(DINING_TABLE_NAME).size(), SEGMENT_COUNT);
    // Wait up to 30s for broker external view to reach the expected size
    waitForPredicate(new Callable<Boolean>() {

        @Override
        public Boolean call() throws Exception {
            return _helixAdmin.getResourceExternalView(HELIX_CLUSTER_NAME, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE).getStateMap(COFFEE_TABLE_NAME).size() == SEGMENT_COUNT;
        }
    }, 30000L);
    externalView = _helixAdmin.getResourceExternalView(HELIX_CLUSTER_NAME, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE);
    Assert.assertEquals(externalView.getStateMap(COFFEE_TABLE_NAME).size(), SEGMENT_COUNT);
    // Wait up to 30s for routing table to reach the expected size
    waitForPredicate(new Callable<Boolean>() {

        @Override
        public Boolean call() throws Exception {
            return brokerRoutingTable.size() == 2;
        }
    }, 30000L);
    Object[] tableArray = brokerRoutingTable.keySet().toArray();
    Arrays.sort(tableArray);
    Assert.assertEquals(Arrays.toString(tableArray), "[coffee_OFFLINE, dining_OFFLINE]");
    Set<String> serverSet = brokerRoutingTable.get(DINING_TABLE_NAME).get(0).getServerSet();
    Assert.assertEquals(brokerRoutingTable.get(DINING_TABLE_NAME).get(0).getSegmentSet(serverSet.iterator().next()).size(), 5);
    final String dataResource = DINING_TABLE_NAME;
    addOneSegment(dataResource);
    // Wait up to 30s for external view to reach the expected size
    waitForPredicate(new Callable<Boolean>() {

        @Override
        public Boolean call() throws Exception {
            return _helixAdmin.getResourceExternalView(HELIX_CLUSTER_NAME, DINING_TABLE_NAME).getPartitionSet().size() == SEGMENT_COUNT;
        }
    }, 30000L);
    externalView = _helixAdmin.getResourceExternalView(HELIX_CLUSTER_NAME, DINING_TABLE_NAME);
    Assert.assertEquals(externalView.getPartitionSet().size(), SEGMENT_COUNT);
    tableArray = brokerRoutingTable.keySet().toArray();
    Arrays.sort(tableArray);
    Assert.assertEquals(Arrays.toString(tableArray), "[coffee_OFFLINE, dining_OFFLINE]");
    // Wait up to 30s for routing table to reach the expected size
    waitForPredicate(new Callable<Boolean>() {

        @Override
        public Boolean call() throws Exception {
            ServerToSegmentSetMap routingTable = brokerRoutingTable.get(DINING_TABLE_NAME).get(0);
            String firstServer = routingTable.getServerSet().iterator().next();
            return routingTable.getSegmentSet(firstServer).size() == SEGMENT_COUNT;
        }
    }, 30000L);
    serverSet = brokerRoutingTable.get(DINING_TABLE_NAME).get(0).getServerSet();
    Assert.assertEquals(brokerRoutingTable.get(DINING_TABLE_NAME).get(0).getSegmentSet(serverSet.iterator().next()).size(), SEGMENT_COUNT);
}