Examples with IdealState org.apache.helix.model.IdealState used on opensource projects

Example 51 with IdealState

use of org.apache.helix.model.IdealState in project pinot by linkedin.

the class PinotLLCRealtimeSegmentManagerTest method testCompleteCommittingSegments.

@Test
public void testCompleteCommittingSegments() throws Exception {
    // Run multiple times randomizing the situation.
    for (int i = 0; i < 100; i++) {
        final List<ZNRecord> existingSegmentMetadata = new ArrayList<>(64);
        final int nPartitions = 16;
        final long seed = new Random().nextLong();
        Random random = new Random(seed);
        final int maxSeq = 10;
        final long now = System.currentTimeMillis();
        final String tableName = "table";
        final String realtimeTableName = TableNameBuilder.REALTIME_TABLE_NAME_BUILDER.forTable(tableName);
        final IdealState idealState = PinotTableIdealStateBuilder.buildEmptyKafkaConsumerRealtimeIdealStateFor(realtimeTableName, 19);
        int nIncompleteCommits = 0;
        final String topic = "someTopic";
        final int nInstances = 5;
        final int nReplicas = 3;
        List<String> instances = getInstanceList(nInstances);
        final String startOffset = KAFKA_OFFSET;
        FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager(false, null);
        segmentManager.setupHelixEntries(topic, realtimeTableName, nPartitions, instances, nReplicas, startOffset, DUMMY_HOST, idealState, false, 10000);
        ZNRecord partitionAssignment = segmentManager.getKafkaPartitionAssignment(realtimeTableName);
        for (int p = 0; p < nPartitions; p++) {
            // Current segment sequence ID for that partition
            int curSeq = random.nextInt(maxSeq);
            if (curSeq == 0) {
                curSeq++;
            }
            boolean incomplete = false;
            if (random.nextBoolean()) {
                incomplete = true;
            }
            for (int s = 0; s < curSeq; s++) {
                LLCSegmentName segmentName = new LLCSegmentName(tableName, p, s, now);
                String segNameStr = segmentName.getSegmentName();
                String state = PinotHelixSegmentOnlineOfflineStateModelGenerator.ONLINE_STATE;
                CommonConstants.Segment.Realtime.Status status = CommonConstants.Segment.Realtime.Status.DONE;
                if (s == curSeq - 1) {
                    state = PinotHelixSegmentOnlineOfflineStateModelGenerator.CONSUMING_STATE;
                    if (!incomplete) {
                        status = CommonConstants.Segment.Realtime.Status.IN_PROGRESS;
                    }
                }
                List<String> instancesForThisSeg = partitionAssignment.getListField(Integer.toString(p));
                for (String instance : instancesForThisSeg) {
                    idealState.setPartitionState(segNameStr, instance, state);
                }
                LLCRealtimeSegmentZKMetadata metadata = new LLCRealtimeSegmentZKMetadata();
                metadata.setSegmentName(segNameStr);
                metadata.setStatus(status);
                existingSegmentMetadata.add(metadata.toZNRecord());
            }
            // Add an incomplete commit to some of them
            if (incomplete) {
                nIncompleteCommits++;
                LLCSegmentName segmentName = new LLCSegmentName(tableName, p, curSeq, now);
                LLCRealtimeSegmentZKMetadata metadata = new LLCRealtimeSegmentZKMetadata();
                metadata.setSegmentName(segmentName.getSegmentName());
                metadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE);
                existingSegmentMetadata.add(metadata.toZNRecord());
            }
        }
        segmentManager._tableIdealState = idealState;
        segmentManager._existingSegmentMetadata = existingSegmentMetadata;
        segmentManager.completeCommittingSegments(TableNameBuilder.REALTIME_TABLE_NAME_BUILDER.forTable(tableName));
        Assert.assertEquals(segmentManager._nCallsToUpdateHelix, nIncompleteCommits, "Failed with seed " + seed);
    }
}

Example 52 with IdealState

use of org.apache.helix.model.IdealState in project pinot by linkedin.

the class SegmentDeletionManagerTest method makeHelixAdmin.

HelixAdmin makeHelixAdmin() {
    HelixAdmin admin = mock(HelixAdmin.class);
    ExternalView ev = mock(ExternalView.class);
    IdealState is = mock(IdealState.class);
    when(admin.getResourceExternalView(clusterName, tableName)).thenReturn(ev);
    when(admin.getResourceIdealState(clusterName, tableName)).thenReturn(is);
    List<String> segmentsInIs = segmentsInIdealStateOrExtView();
    Map<String, String> dummy = new HashMap<>(1);
    dummy.put("someHost", "ONLINE");
    for (String segment : segmentsInIs) {
        when(is.getInstanceStateMap(segment)).thenReturn(dummy);
    }
    when(ev.getStateMap(anyString())).thenReturn(null);
    return admin;
}

Example 53 with IdealState

use of org.apache.helix.model.IdealState in project pinot by linkedin.

the class ValidationManagerTest method testRebuildBrokerResourceWhenBrokerAdded.

@Test
public void testRebuildBrokerResourceWhenBrokerAdded() throws Exception {
    // Check that the first table we added doesn't need to be rebuilt(case where ideal state brokers and brokers in broker resource are the same.
    String partitionName = _offlineTableConfig.getTableName();
    HelixAdmin helixAdmin = _helixManager.getClusterManagmentTool();
    IdealState idealState = HelixHelper.getBrokerIdealStates(helixAdmin, HELIX_CLUSTER_NAME);
    // Ensure that the broker resource is not rebuilt.
    Assert.assertTrue(idealState.getInstanceSet(partitionName).equals(_pinotHelixResourceManager.getAllInstancesForBrokerTenant(ControllerTenantNameBuilder.DEFAULT_TENANT_NAME)));
    _pinotHelixResourceManager.rebuildBrokerResourceFromHelixTags(partitionName);
    // Add another table that needs to be rebuilt
    String offlineTableTwoConfigJson = ControllerRequestBuilderUtil.buildCreateOfflineTableJSON(TEST_TABLE_TWO, null, null, 1).toString();
    AbstractTableConfig offlineTableConfigTwo = AbstractTableConfig.init(offlineTableTwoConfigJson);
    _pinotHelixResourceManager.addTable(offlineTableConfigTwo);
    String partitionNameTwo = offlineTableConfigTwo.getTableName();
    // Add a new broker manually such that the ideal state is not updated and ensure that rebuild broker resource is called
    final String brokerId = "Broker_localhost_2";
    InstanceConfig instanceConfig = new InstanceConfig(brokerId);
    instanceConfig.setInstanceEnabled(true);
    instanceConfig.setHostName("Broker_localhost");
    instanceConfig.setPort("2");
    helixAdmin.addInstance(HELIX_CLUSTER_NAME, instanceConfig);
    helixAdmin.addInstanceTag(HELIX_CLUSTER_NAME, instanceConfig.getInstanceName(), ControllerTenantNameBuilder.getBrokerTenantNameForTenant(ControllerTenantNameBuilder.DEFAULT_TENANT_NAME));
    idealState = HelixHelper.getBrokerIdealStates(helixAdmin, HELIX_CLUSTER_NAME);
    // Assert that the two don't equal before the call to rebuild the broker resource.
    Assert.assertTrue(!idealState.getInstanceSet(partitionNameTwo).equals(_pinotHelixResourceManager.getAllInstancesForBrokerTenant(ControllerTenantNameBuilder.DEFAULT_TENANT_NAME)));
    _pinotHelixResourceManager.rebuildBrokerResourceFromHelixTags(partitionNameTwo);
    idealState = HelixHelper.getBrokerIdealStates(helixAdmin, HELIX_CLUSTER_NAME);
    // Assert that the two do equal after being rebuilt.
    Assert.assertTrue(idealState.getInstanceSet(partitionNameTwo).equals(_pinotHelixResourceManager.getAllInstancesForBrokerTenant(ControllerTenantNameBuilder.DEFAULT_TENANT_NAME)));
}

Example 54 with IdealState

use of org.apache.helix.model.IdealState in project pinot by linkedin.

the class BalanceNumSegmentAssignmentStrategy method getAssignedInstances.

@Override
public List<String> getAssignedInstances(HelixAdmin helixAdmin, String helixClusterName, SegmentMetadata segmentMetadata, int numReplicas, String tenantName) {
    String serverTenantName;
    String tableName;
    if ("realtime".equalsIgnoreCase(segmentMetadata.getIndexType())) {
        tableName = TableNameBuilder.REALTIME_TABLE_NAME_BUILDER.forTable(segmentMetadata.getTableName());
        serverTenantName = ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(tenantName);
    } else {
        tableName = TableNameBuilder.OFFLINE_TABLE_NAME_BUILDER.forTable(segmentMetadata.getTableName());
        serverTenantName = ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName);
    }
    List<String> selectedInstances = new ArrayList<String>();
    Map<String, Integer> currentNumSegmentsPerInstanceMap = new HashMap<String, Integer>();
    List<String> allTaggedInstances = HelixHelper.getEnabledInstancesWithTag(helixAdmin, helixClusterName, serverTenantName);
    for (String instance : allTaggedInstances) {
        currentNumSegmentsPerInstanceMap.put(instance, 0);
    }
    // Count number of segments assigned to each instance
    IdealState idealState = helixAdmin.getResourceIdealState(helixClusterName, tableName);
    if (idealState != null) {
        for (String partitionName : idealState.getPartitionSet()) {
            Map<String, String> instanceToStateMap = idealState.getInstanceStateMap(partitionName);
            if (instanceToStateMap != null) {
                for (String instanceName : instanceToStateMap.keySet()) {
                    if (currentNumSegmentsPerInstanceMap.containsKey(instanceName)) {
                        currentNumSegmentsPerInstanceMap.put(instanceName, currentNumSegmentsPerInstanceMap.get(instanceName) + 1);
                    }
                // else, ignore. Do not add servers, that are not tagged, to the map
                // By this approach, new segments will not be allotted to the server if tags changed
                }
            }
        }
    }
    // Select up to numReplicas instances with the fewest segments assigned
    PriorityQueue<Number2ObjectPair<String>> priorityQueue = new PriorityQueue<Number2ObjectPair<String>>(numReplicas, Pairs.getDescendingnumber2ObjectPairComparator());
    for (String key : currentNumSegmentsPerInstanceMap.keySet()) {
        priorityQueue.add(new Number2ObjectPair<String>(currentNumSegmentsPerInstanceMap.get(key), key));
        if (priorityQueue.size() > numReplicas) {
            priorityQueue.poll();
        }
    }
    while (!priorityQueue.isEmpty()) {
        selectedInstances.add(priorityQueue.poll().getB());
    }
    LOGGER.info("Segment assignment result for : " + segmentMetadata.getName() + ", in resource : " + segmentMetadata.getTableName() + ", selected instances: " + Arrays.toString(selectedInstances.toArray()));
    return selectedInstances;
}

Example 55 with IdealState

use of org.apache.helix.model.IdealState in project pinot by linkedin.

the class PinotLLCRealtimeSegmentManager method cleanupLLC.

// Remove all trace of LLC for this table.
public void cleanupLLC(final String realtimeTableName) {
    // Start by removing the kafka partition assigment znode. This will prevent any new segments being created.
    ZKMetadataProvider.removeKafkaPartitionAssignmentFromPropertyStore(_propertyStore, realtimeTableName);
    LOGGER.info("Removed Kafka partition assignemnt (if any) record for {}", realtimeTableName);
    // If there are any completions in the pipeline we let them commit.
    Uninterruptibles.sleepUninterruptibly(1, TimeUnit.SECONDS);
    IdealState idealState = HelixHelper.getTableIdealState(_helixManager, realtimeTableName);
    final List<String> segmentsToRemove = new ArrayList<String>();
    Set<String> allSegments = idealState.getPartitionSet();
    int removeCount = 0;
    for (String segmentName : allSegments) {
        if (SegmentName.isLowLevelConsumerSegmentName(segmentName)) {
            segmentsToRemove.add(segmentName);
            removeCount++;
        }
    }
    LOGGER.info("Attempting to remove {} LLC segments of table {}", removeCount, realtimeTableName);
    _helixResourceManager.deleteSegments(realtimeTableName, segmentsToRemove);
}