@Test public void missingEVPartitionTest() throws Exception { final String tableName = "myTable"; List<String> allTableNames = new ArrayList<String>(); allTableNames.add(tableName); IdealState idealState = new IdealState(tableName); idealState.setPartitionState("myTable_0", "pinot1", "ONLINE"); idealState.setPartitionState("myTable_0", "pinot2", "ONLINE"); idealState.setPartitionState("myTable_0", "pinot3", "ONLINE"); idealState.setPartitionState("myTable_1", "pinot1", "ONLINE"); idealState.setPartitionState("myTable_1", "pinot2", "ONLINE"); idealState.setPartitionState("myTable_1", "pinot3", "ONLINE"); idealState.setPartitionState("myTable_2", "pinot3", "OFFLINE"); idealState.setPartitionState("myTable_3", "pinot3", "ONLINE"); idealState.setReplicas("2"); idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED); ExternalView externalView = new ExternalView(tableName); externalView.setState("myTable_0", "pinot1", "ONLINE"); externalView.setState("myTable_0", "pinot2", "ONLINE"); externalView.setState("myTable_1", "pinot1", "ERROR"); externalView.setState("myTable_1", "pinot2", "ONLINE"); HelixAdmin helixAdmin; { helixAdmin = mock(HelixAdmin.class); when(helixAdmin.getResourceIdealState("StatusChecker", "myTable")).thenReturn(idealState); when(helixAdmin.getResourceExternalView("StatusChecker", "myTable")).thenReturn(externalView); } { helixResourceManager = mock(PinotHelixResourceManager.class); when(helixResourceManager.isLeader()).thenReturn(true); when(helixResourceManager.getAllPinotTableNames()).thenReturn(allTableNames); when(helixResourceManager.getHelixClusterName()).thenReturn("StatusChecker"); when(helixResourceManager.getHelixAdmin()).thenReturn(helixAdmin); } { config = mock(ControllerConf.class); when(config.getStatusControllerFrequencyInSeconds()).thenReturn(300); } metricsRegistry = new MetricsRegistry(); controllerMetrics = new ControllerMetrics(metricsRegistry); segmentStatusChecker = new SegmentStatusChecker(helixResourceManager, config); segmentStatusChecker.setMetricsRegistry(controllerMetrics); segmentStatusChecker.runSegmentMetrics(); Assert.assertEquals( controllerMetrics.getValueOfTableGauge( externalView.getId(), ControllerGauge.SEGMENTS_IN_ERROR_STATE), 1); Assert.assertEquals( controllerMetrics.getValueOfTableGauge( externalView.getId(), ControllerGauge.NUMBER_OF_REPLICAS), 0); segmentStatusChecker.stop(); }
@Override public IdealState computeNewIdealState( String resourceName, IdealState currentIdealState, CurrentStateOutput currentStateOutput, ClusterDataCache clusterData) { testRebalancerInvoked = true; for (String partition : currentIdealState.getPartitionSet()) { String instance = currentIdealState.getPreferenceList(partition).get(0); currentIdealState.getPreferenceList(partition).clear(); currentIdealState.getPreferenceList(partition).add(instance); currentIdealState.getInstanceStateMap(partition).clear(); currentIdealState.getInstanceStateMap(partition).put(instance, "MASTER"); } currentIdealState.setReplicas("1"); return currentIdealState; }
@Test public void testSchemataSM() throws Exception { String className = TestHelper.getTestClassName(); String methodName = TestHelper.getTestMethodName(); String clusterName = className + "_" + methodName; int n = 5; MockParticipant[] participants = new MockParticipant[n]; System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis())); TestHelper.setupCluster( clusterName, _zkaddr, 12918, // participant start port "localhost", // participant name prefix "TestSchemata", // resource name prefix 1, // resources 1, // partitions per resource n, // number of nodes 0, // replicas "STORAGE_DEFAULT_SM_SCHEMATA", false); // don't rebalance // rebalance ideal-state to use ANY_LIVEINSTANCE for preference list ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor); PropertyKey.Builder keyBuilder = accessor.keyBuilder(); PropertyKey key = keyBuilder.idealStates("TestSchemata0"); IdealState idealState = accessor.getProperty(key); idealState.setReplicas(HelixConstants.StateModelToken.ANY_LIVEINSTANCE.toString()); idealState .getRecord() .setListField( "TestSchemata0_0", Arrays.asList(HelixConstants.StateModelToken.ANY_LIVEINSTANCE.toString())); accessor.setProperty(key, idealState); MockController controller = new MockController(_zkaddr, clusterName, "controller"); controller.syncStart(); // start n-1 participants for (int i = 1; i < n; i++) { String instanceName = "localhost_" + (12918 + i); participants[i] = new MockParticipant(_zkaddr, clusterName, instanceName); participants[i].syncStart(); } boolean result = ClusterStateVerifier.verifyByZkCallback( new BestPossAndExtViewZkVerifier(_zkaddr, clusterName)); Assert.assertTrue(result); // start the remaining 1 participant participants[0] = new MockParticipant(_zkaddr, clusterName, "localhost_12918"); participants[0].syncStart(); // make sure we have all participants in MASTER state result = ClusterStateVerifier.verifyByZkCallback( new BestPossAndExtViewZkVerifier(_zkaddr, clusterName)); Assert.assertTrue(result); key = keyBuilder.externalView("TestSchemata0"); ExternalView externalView = accessor.getProperty(key); Map<String, String> stateMap = externalView.getStateMap("TestSchemata0_0"); Assert.assertNotNull(stateMap); Assert.assertEquals(stateMap.size(), n, "all " + n + " participants should be in Master state"); for (int i = 0; i < n; i++) { String instanceName = "localhost_" + (12918 + i); Assert.assertNotNull(stateMap.get(instanceName)); Assert.assertEquals(stateMap.get(instanceName), "MASTER"); } // clean up controller.syncStop(); for (int i = 0; i < n; i++) { participants[i].syncStop(); } System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis())); }
@Test public void testZKReconnect() throws Exception { final AtomicReference<ZkServer> zkServerRef = new AtomicReference<ZkServer>(); final int zkPort = TestHelper.getRandomPort(); final String zkAddr = String.format("localhost:%d", zkPort); ZkServer zkServer = TestHelper.startZkServer(zkAddr); zkServerRef.set(zkServer); String className = TestHelper.getTestClassName(); String methodName = TestHelper.getTestMethodName(); String clusterName = className + "_" + methodName; // Setup cluster LOG.info("Setup clusters"); ClusterSetup clusterSetup = new ClusterSetup(zkAddr); clusterSetup.addCluster(clusterName, true); // Registers and starts controller LOG.info("Starts controller"); HelixManager controller = HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.CONTROLLER, zkAddr); controller.connect(); // Registers and starts participant LOG.info("Starts participant"); String hostname = "localhost"; String instanceId = String.format("%s_%d", hostname, 1); clusterSetup.addInstanceToCluster(clusterName, instanceId); HelixManager participant = HelixManagerFactory.getZKHelixManager( clusterName, instanceId, InstanceType.PARTICIPANT, zkAddr); participant.connect(); LOG.info("Register state machine"); final CountDownLatch latch = new CountDownLatch(1); participant .getStateMachineEngine() .registerStateModelFactory( "OnlineOffline", new StateModelFactory<StateModel>() { @Override public StateModel createNewStateModel(String stateUnitKey) { return new SimpleStateModel(latch); } }, "test"); String resourceName = "test-resource"; LOG.info("Ideal state assignment"); HelixAdmin helixAdmin = participant.getClusterManagmentTool(); helixAdmin.addResource( clusterName, resourceName, 1, "OnlineOffline", IdealState.RebalanceMode.CUSTOMIZED.toString()); IdealState idealState = helixAdmin.getResourceIdealState(clusterName, resourceName); idealState.setReplicas("1"); idealState.setStateModelFactoryName("test"); idealState.setPartitionState(resourceName + "_0", instanceId, "ONLINE"); LOG.info("Shutdown ZK server"); TestHelper.stopZkServer(zkServerRef.get()); Executors.newSingleThreadScheduledExecutor() .schedule( new Runnable() { @Override public void run() { try { LOG.info("Restart ZK server"); // zkServer.set(TestUtils.startZookeeper(zkDir, zkPort)); zkServerRef.set(TestHelper.startZkServer(zkAddr, null, false)); } catch (Exception e) { LOG.error(e.getMessage(), e); } } }, 2L, TimeUnit.SECONDS); // future.get(); LOG.info("Before update ideal state"); helixAdmin.setResourceIdealState(clusterName, resourceName, idealState); LOG.info("After update ideal state"); LOG.info("Wait for OFFLINE->ONLINE state transition"); try { Assert.assertTrue(latch.await(10, TimeUnit.SECONDS)); // wait until stable state boolean result = ClusterStateVerifier.verifyByZkCallback( new BestPossAndExtViewZkVerifier(zkAddr, clusterName)); Assert.assertTrue(result); } finally { participant.disconnect(); zkServerRef.get().shutdown(); } }
public static IdealState calculateRelayIdealState( List<String> partitions, List<String> instances, String resultRecordName, int replica, String firstValue, String restValue, String stateModelName) { Collections.sort(partitions); Collections.sort(instances); if (instances.size() % replica != 0) { throw new HelixException("Instances must be divided by replica"); } IdealState result = new IdealState(resultRecordName); result.setNumPartitions(partitions.size()); result.setReplicas("" + replica); result.setStateModelDefId(StateModelDefinitionId.from(stateModelName)); int groups = instances.size() / replica; int remainder = instances.size() % replica; int remainder2 = partitions.size() % groups; int storageNodeGroupSize = partitions.size() / groups; for (int i = 0; i < groups; i++) { int relayStart = 0, relayEnd = 0, storageNodeStart = 0, storageNodeEnd = 0; if (i < remainder) { relayStart = (replica + 1) * i; relayEnd = (replica + 1) * (i + 1); } else { relayStart = (replica + 1) * remainder + replica * (i - remainder); relayEnd = relayStart + replica; } // System.out.println("relay start :" + relayStart + " relayEnd:" + relayEnd); if (i < remainder2) { storageNodeStart = (storageNodeGroupSize + 1) * i; storageNodeEnd = (storageNodeGroupSize + 1) * (i + 1); } else { storageNodeStart = (storageNodeGroupSize + 1) * remainder2 + storageNodeGroupSize * (i - remainder2); storageNodeEnd = storageNodeStart + storageNodeGroupSize; } // System.out.println("storageNodeStart :" + storageNodeStart + " storageNodeEnd:" + // storageNodeEnd); List<String> snBatch = partitions.subList(storageNodeStart, storageNodeEnd); List<String> relayBatch = instances.subList(relayStart, relayEnd); Map<String, List<String>> sublistFields = calculateSubIdealState(snBatch, relayBatch, replica); result.getRecord().getListFields().putAll(sublistFields); } for (String snName : result.getRecord().getListFields().keySet()) { Map<String, String> mapField = new TreeMap<String, String>(); List<String> relayCandidates = result.getRecord().getListField(snName); mapField.put(relayCandidates.get(0), firstValue); for (int i = 1; i < relayCandidates.size(); i++) { mapField.put(relayCandidates.get(i), restValue); } result.getRecord().getMapFields().put(snName, mapField); } System.out.println(); return result; }