/** * Sets the idealstate for a given segment to the target mapping */ private void setTargetState(IdealState idealState, String segmentId, Map<String, String> targetMap) { idealState.getInstanceStateMap(segmentId).clear(); idealState.setInstanceStateMap(segmentId, targetMap); }
private IdealState updateIdealState(IdealState idealState, int newNumReplicas) { idealState.setReplicas(Integer.toString(newNumReplicas)); Set<String> segmentIds = idealState.getPartitionSet(); for (String segmentId : segmentIds) { Map<String, String> instanceStateMap = idealState.getInstanceStateMap(segmentId); if (instanceStateMap.size() > newNumReplicas) { Set<String> keys = instanceStateMap.keySet(); while (instanceStateMap.size() > newNumReplicas) { instanceStateMap.remove(keys.iterator().next()); } } else if (instanceStateMap.size() < newNumReplicas) { throw new RuntimeException( "Segment " + segmentId + " has " + instanceStateMap.size() + " replicas but want changed to " + newNumReplicas); } } return idealState; }
/** * Returns a map from server instance to list of segments it serves for the given table. */ public Map<String, List<String>> getServerToSegmentsMap(String tableNameWithType) { Map<String, List<String>> serverToSegmentsMap = new HashMap<>(); IdealState idealState = _helixAdmin.getResourceIdealState(_helixClusterName, tableNameWithType); if (idealState == null) { throw new IllegalStateException("Ideal state does not exist for table: " + tableNameWithType); } for (String segment : idealState.getPartitionSet()) { for (String server : idealState.getInstanceStateMap(segment).keySet()) { serverToSegmentsMap.computeIfAbsent(server, key -> new ArrayList<>()).add(segment); } } return serverToSegmentsMap; }
private boolean shouldDeleteInProgressLLCSegment(String segmentName, IdealState idealState, RealtimeSegmentZKMetadata realtimeSegmentZKMetadata) { if (idealState == null) { return false; } // delete a segment only if it is old enough (5 days) or else, // 1. latest segment could get deleted in the middle of repair by RealtimeSegmentValidationManager // 2. for a brand new segment, if this code kicks in after new metadata is created but ideal state entry is not yet created (between step 2 and 3), // the latest segment metadata could get marked for deletion if (System.currentTimeMillis() - realtimeSegmentZKMetadata.getCreationTime() <= OLD_LLC_SEGMENTS_RETENTION_IN_MILLIS) { return false; } Map<String, String> stateMap = idealState.getInstanceStateMap(segmentName); if (stateMap == null) { // Segment is in property store but not in ideal state, delete it return true; } else { // Delete segment if all of its replicas are OFFLINE Set<String> states = new HashSet<>(stateMap.values()); return states.size() == 1 && states .contains(CommonConstants.Helix.StateModel.SegmentOnlineOfflineStateModel.OFFLINE); } }
public IdealStateBuilderUtil setSegmentState(String segmentName, String state) { Map<String, String> instanceStateMap = _idealState.getInstanceStateMap(segmentName); for (Map.Entry<String, String> entry : instanceStateMap.entrySet()) { instanceStateMap.put(entry.getKey(), state); } return this; }
@Nullable @Override public IdealState apply(@Nullable IdealState input) { List<String> segmentNames = Lists.newArrayList(input.getPartitionSet()); Collections.sort(segmentNames); Map<String, String> instanceStateMap1 = input.getInstanceStateMap(segmentNames.get(0)); for (String instance : instanceStateMap1.keySet()) { instanceStateMap1.put(instance, "OFFLINE"); break; } return input; } }, RetryPolicies.fixedDelayRetryPolicy(2, 10));
private Map<String, Map<String, String>> getPrevAssignment(IdealState idealState) { Map<String, Map<String, String>> prevAssignment = new HashMap<>(1); for (String segment : idealState.getPartitionSet()) { Map<String, String> instanceMap = new HashMap<>(1); instanceMap.putAll(idealState.getInstanceStateMap(segment)); prevAssignment.put(segment, instanceMap); } return prevAssignment; }
Map<String, String> stateMap = idealState.getInstanceStateMap(newSegmentId); if (stateMap != null) { stateMap.clear();
private void validateIdealStateRealtime(IdealState rebalancedIdealState, int nSegmentsCompleted, int nSegmentsConsuming, int targetNumReplicas, List<String> instancesCompleted, List<String> instancesConsuming, Configuration rebalanceUserConfig) { Assert.assertEquals(rebalancedIdealState.getPartitionSet().size(), nSegmentsCompleted + nSegmentsConsuming); for (String segment : rebalancedIdealState.getPartitionSet()) { Map<String, String> instanceStateMap = rebalancedIdealState.getInstanceStateMap(segment); Assert.assertEquals(instanceStateMap.size(), targetNumReplicas); boolean rebalanceConsuming = rebalanceUserConfig.getBoolean(RebalanceUserConfigConstants.INCLUDE_CONSUMING); if (segment.contains("consuming")) { if (rebalanceConsuming) { Assert.assertTrue(instancesConsuming.containsAll(instanceStateMap.keySet())); } } else { Assert.assertTrue(instancesCompleted.containsAll(instanceStateMap.keySet())); } } }
private void verifyAssignmentIsFromLatest(StreamPartitionAssignmentGenerator streamPartitionAssignmentGenerator, IdealState idealState, Map<String, List<String>> assignment) { Map<String, LLCSegmentName> partitionToLatestSegments = streamPartitionAssignmentGenerator.getPartitionToLatestSegments(idealState); for (Map.Entry<String, List<String>> entry : assignment.entrySet()) { String segmentName = entry.getKey(); List<String> assignedInstances = entry.getValue(); LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); int partitionId = llcSegmentName.getPartitionId(); LLCSegmentName latestSegment = partitionToLatestSegments.get(String.valueOf(partitionId)); Set<String> instancesInIdealState = idealState.getInstanceStateMap(latestSegment.getSegmentName()).keySet(); Assert.assertEquals(assignedInstances.size(), instancesInIdealState.size()); Assert.assertTrue(assignedInstances.containsAll(instancesInIdealState)); } }
private void validateIdealState(IdealState rebalancedIdealState, int nSegments, int targetNumReplicas, List<String> instances, Map<String, Map<String, String>> prevAssignment, boolean changeExpected) { Assert.assertEquals(rebalancedIdealState.getPartitionSet().size(), nSegments); for (String segment : rebalancedIdealState.getPartitionSet()) { Map<String, String> instanceStateMap = rebalancedIdealState.getInstanceStateMap(segment); Assert.assertEquals(instanceStateMap.size(), targetNumReplicas); Assert.assertTrue(instances.containsAll(instanceStateMap.keySet())); } boolean changed = false; for (String segment : prevAssignment.keySet()) { Map<String, String> prevInstanceMap = prevAssignment.get(segment); Map<String, String> instanceStateMap = rebalancedIdealState.getInstanceStateMap(segment); if (!changeExpected) { Assert.assertTrue(prevInstanceMap.keySet().containsAll(instanceStateMap.keySet())); Assert.assertTrue(instanceStateMap.keySet().containsAll(prevInstanceMap.keySet())); } else { if (!prevInstanceMap.keySet().containsAll(instanceStateMap.keySet()) || !instanceStateMap.keySet() .containsAll(prevInstanceMap.keySet())) { changed = true; break; } } } Assert.assertEquals(changeExpected, changed); }
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; }
private ArrayNode getSegments(String tableName, String tableType) { ArrayNode segments = JsonUtils.newArrayNode(); String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(tableName); String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(tableName); String tableNameWithType; if (CommonConstants.Helix.TableType.valueOf(tableType).toString().equals("REALTIME")) { tableNameWithType = realtimeTableName; } else { tableNameWithType = offlineTableName; } List<String> segmentList = _pinotHelixResourceManager.getSegmentsFor(tableNameWithType); IdealState idealState = HelixHelper.getTableIdealState(_pinotHelixResourceManager.getHelixZkManager(), tableNameWithType); for (String segmentName : segmentList) { Map<String, String> map = idealState.getInstanceStateMap(segmentName); if (map == null) { continue; } if (!map.containsValue(PinotHelixSegmentOnlineOfflineStateModelGenerator.OFFLINE_STATE)) { segments.add(segmentName); } } return segments; }
private void verifySegmentAssignment(IdealState updatedIdealState, IdealState prevIdealState, List<String> completedInstanceList, int nReplicas, Map<String, Integer> segmentNameToNumCompletedInstances) { Assert.assertEquals(updatedIdealState.getPartitionSet().size(), prevIdealState.getPartitionSet().size()); Assert.assertTrue(prevIdealState.getPartitionSet().containsAll(updatedIdealState.getPartitionSet())); for (String segmentName : updatedIdealState.getPartitionSet()) { Map<String, String> newInstanceStateMap = updatedIdealState.getInstanceStateMap(segmentName); int onCompleted = 0; int notOnCompleted = 0; for (String instance : newInstanceStateMap.keySet()) { if (completedInstanceList.contains(instance)) { onCompleted++; } else { notOnCompleted++; } } int expectedOnCompletedServers = segmentNameToNumCompletedInstances.get(segmentName).intValue(); Assert.assertEquals(onCompleted, expectedOnCompletedServers); Assert.assertEquals(notOnCompleted, nReplicas - expectedOnCompletedServers); } }
/** * Given an ideal state find the segments that need to relocate a replica to completed servers, * and create a new instance state map for those segments * * @param realtimeTagConfig * @param idealState * @param completedServers * @param completedServersQueue * @return */ private void createNewIdealState(final RealtimeTagConfig realtimeTagConfig, IdealState idealState, final List<String> completedServers, MinMaxPriorityQueue<Map.Entry<String, Integer>> completedServersQueue) { // TODO: we are scanning the entire segments list every time. This is unnecessary because only the latest segments will need relocation // Can we do something to avoid this? // 1. Time boundary: scan only last day whereas runFrequency = hourly // 2. For each partition, scan in descending order, and stop when the first segment not needing relocation is found for (String segmentName : idealState.getPartitionSet()) { final Map<String, String> instanceStateMap = idealState.getInstanceStateMap(segmentName); Map<String, String> newInstanceStateMap = createNewInstanceStateMap(realtimeTagConfig, segmentName, instanceStateMap, completedServers, completedServersQueue); if (MapUtils.isNotEmpty(newInstanceStateMap)) { idealState.setInstanceStateMap(segmentName, newInstanceStateMap); } } }
private void verifyPartitionAssignmentFromIdealState(TableConfig tableConfig, IdealState idealState, int numPartitions) { TestStreamPartitionAssignmentGenerator partitionAssignmentGenerator = new TestStreamPartitionAssignmentGenerator(_mockHelixManager); PartitionAssignment partitionAssignmentFromIdealState = partitionAssignmentGenerator.getStreamPartitionAssignmentFromIdealState(tableConfig, idealState); Assert.assertEquals(tableConfig.getTableName(), partitionAssignmentFromIdealState.getTableName()); Assert.assertEquals(partitionAssignmentFromIdealState.getNumPartitions(), numPartitions); // check that latest segments are honoring partition assignment Map<String, LLCSegmentName> partitionIdToLatestLLCSegment = partitionAssignmentGenerator.getPartitionToLatestSegments(idealState); for (Map.Entry<String, LLCSegmentName> entry : partitionIdToLatestLLCSegment.entrySet()) { Set<String> idealStateInstances = idealState.getInstanceStateMap(entry.getValue().getSegmentName()).keySet(); List<String> partitionAssignmentInstances = partitionAssignmentFromIdealState.getInstancesListForPartition(entry.getKey()); Assert.assertEquals(idealStateInstances.size(), partitionAssignmentInstances.size()); Assert.assertTrue(idealStateInstances.containsAll(partitionAssignmentInstances)); } }
public List<String> getInstances(int partition, int seqNum) { List<String> instances = new ArrayList<>(); for (String segmentName : _idealState.getRecord().getMapFields().keySet()) { if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); if (llcSegmentName.getPartitionId() == partition && llcSegmentName.getSequenceNumber() == seqNum) { Map<String, String> instanceStateMap = _idealState.getInstanceStateMap(segmentName); instances = Lists.newArrayList(instanceStateMap.keySet()); break; } } } return instances; }
public IdealStateBuilderUtil setSegmentState(int partition, int seqNum, String state) { for (String segmentName : _idealState.getRecord().getMapFields().keySet()) { if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); if (llcSegmentName.getPartitionId() == partition && llcSegmentName.getSequenceNumber() == seqNum) { Map<String, String> instanceStateMap = _idealState.getInstanceStateMap(segmentName); for (Map.Entry<String, String> entry : instanceStateMap.entrySet()) { instanceStateMap.put(entry.getKey(), state); } break; } } } return this; }
public IdealStateBuilderUtil moveToServers(int partition, int seqNum, List<String> instances) { for (String segmentName : _idealState.getRecord().getMapFields().keySet()) { if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); if (llcSegmentName.getPartitionId() == partition && llcSegmentName.getSequenceNumber() == seqNum) { Map<String, String> instanceStateMap = _idealState.getInstanceStateMap(segmentName); Map<String, String> newInstanceStateMap = new HashMap<>(instanceStateMap.size()); int serverId = 0; for (Map.Entry<String, String> entry : instanceStateMap.entrySet()) { newInstanceStateMap.put(instances.get(serverId++), entry.getValue()); } _idealState.setInstanceStateMap(llcSegmentName.getSegmentName(), newInstanceStateMap); break; } } } return this; }
/** * Regression test for large ideal state updates failing silently */ @Test public void testWriteLargeIdealState() { final int numSegments = 20000; IdealState idealState = new IdealState(RESOURCE_NAME); idealState.setStateModelDefRef("OnlineOffline"); idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED); idealState.setReplicas("0"); String helixClusterName = getHelixClusterName(); _helixAdmin.addResource(helixClusterName, RESOURCE_NAME, idealState); HelixHelper.updateIdealState(_helixManager, RESOURCE_NAME, new Function<IdealState, IdealState>() { @Override public IdealState apply(@Nullable IdealState idealState) { Assert.assertNotNull(idealState); for (int i = 0; i < numSegments; i++) { idealState.setPartitionState("segment_" + i, INSTANCE_NAME, "ONLINE"); } return idealState; } }, RetryPolicies.noDelayRetryPolicy(1)); IdealState resourceIdealState = _helixAdmin.getResourceIdealState(helixClusterName, RESOURCE_NAME); for (int i = 0; i < numSegments; i++) { Assert.assertEquals(resourceIdealState.getInstanceStateMap("segment_" + i).get(INSTANCE_NAME), "ONLINE"); } }