@Nullable @Override public IdealState apply(@Nullable IdealState input) { Map<String, Map<String, String>> existingMapField = input.getRecord().getMapFields(); for (Map.Entry<String, Map<String, String>> segmentEntry : proposedIdealState.entrySet()) { existingMapField.put(segmentEntry.getKey(), segmentEntry.getValue()); } return input; } }, RetryPolicies.exponentialBackoffRetryPolicy(5, 500L, 2.0f));
public int getNumPartitionsFromIdealState(IdealState idealState) { Set<Integer> partitions = new HashSet<>(); Map<String, Map<String, String>> mapFields = idealState.getRecord().getMapFields(); for (Map.Entry<String, Map<String, String>> entry : mapFields.entrySet()) { String segmentName = entry.getKey(); if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); partitions.add(llcSegmentName.getPartitionId()); } } return partitions.size(); }
/** * Gets the next ideal state based on the target (rebalanced) state. If no downtime is desired, the next state * is set such that there is always atleast one common replica for each segment between current and next state. */ private IdealState getNextState(IdealState currentState, IdealState targetState, Configuration rebalanceUserConfig) { // make a copy of the ideal state so it can be updated IdealState idealStateCopy = HelixHelper.cloneIdealState(currentState); Map<String, Map<String, String>> currentMapFields = currentState.getRecord().getMapFields(); Map<String, Map<String, String>> targetMapFields = targetState.getRecord().getMapFields(); for (String segmentId : targetMapFields.keySet()) { updateSegmentIfNeeded(segmentId, currentMapFields.get(segmentId), targetMapFields.get(segmentId), idealStateCopy, rebalanceUserConfig); } return idealStateCopy; }
@Nullable public Map<String, Map<String, String>> getIdealState(@Nonnull String tableNameOptType, @Nullable CommonConstants.Helix.TableType tableType) { String tableNameWithType = getTableNameWithType(tableNameOptType, tableType); IdealState resourceIdealState = _pinotHelixResourceManager.getHelixAdmin() .getResourceIdealState(_pinotHelixResourceManager.getHelixClusterName(), tableNameWithType); return resourceIdealState == null ? null : resourceIdealState.getRecord().getMapFields(); }
/** * Gets stream partition assignment of a table by reading the segment assignment in ideal state */ public PartitionAssignment getStreamPartitionAssignmentFromIdealState(TableConfig tableConfig, IdealState idealState) { String tableNameWithType = tableConfig.getTableName(); // get latest segment in each partition Map<String, LLCSegmentName> partitionIdToLatestSegment = getPartitionToLatestSegments(idealState); // extract partition assignment from the latest segments PartitionAssignment partitionAssignment = new PartitionAssignment(tableNameWithType); Map<String, Map<String, String>> mapFields = idealState.getRecord().getMapFields(); for (Map.Entry<String, LLCSegmentName> entry : partitionIdToLatestSegment.entrySet()) { String segmentName = entry.getValue().getSegmentName(); Map<String, String> instanceStateMap = mapFields.get(segmentName); partitionAssignment.addPartition(entry.getKey(), Lists.newArrayList(instanceStateMap.keySet())); } return partitionAssignment; }
public IdealStateBuilderUtil clear() { _idealState.getRecord().getMapFields().clear(); return this; } }
public IdealStateBuilderUtil removeSegment(String segmentName) { _idealState.getRecord().getMapFields().remove(segmentName); return this; }
public static IdealState cloneIdealState(IdealState idealState) { return new IdealState( (ZNRecord) ZN_RECORD_SERIALIZER.deserialize(ZN_RECORD_SERIALIZER.serialize(idealState.getRecord()))); }
/** * Verifies that all entries in old ideal state are unchanged in the new ideal state * There could be new entries in the ideal state due to num partitions increase * @param oldMapFields * @param idealState */ private void verifyNoChangeToOldEntries(Map<String, Map<String, String>> oldMapFields, IdealState idealState) { Map<String, Map<String, String>> newMapFields = idealState.getRecord().getMapFields(); for (Map.Entry<String, Map<String, String>> oldMapFieldsEntry : oldMapFields.entrySet()) { String oldSegment = oldMapFieldsEntry.getKey(); Map<String, String> oldInstanceStateMap = oldMapFieldsEntry.getValue(); Assert.assertTrue(newMapFields.containsKey(oldSegment)); Assert.assertTrue(oldInstanceStateMap.equals(newMapFields.get(oldSegment))); } }
@Test public void noDowntimeUpdateWithNoCommonElements() { Map<String, String> targetMap = new HashMap<>(); targetMap.put("host4", "ONLINE"); targetMap.put("host3", "ONLINE"); Map<String, String> srcMap = current.getRecord().getMapField(segmentId); TableRebalancer updater = new TableRebalancer(null, null, null); updater.updateSegmentIfNeeded(segmentId, srcMap, targetMap, current, noDowntime); Map<String, String> tempMap = current.getRecord().getMapField(segmentId); Set<String> targetHosts = new HashSet<String>(Arrays.asList("host3", "host4")); Set<String> srcHosts = new HashSet<String>(Arrays.asList("host1", "host2")); Assert.assertEquals(tempMap.size(), targetHosts.size()); for (String instance : tempMap.keySet()) { Assert.assertTrue(targetHosts.contains(instance) || srcHosts.contains(instance)); } }
@Test public void downtimeUpdateWithCommonElements() { Map<String, String> targetMap = new HashMap<>(); targetMap.put("host1", "ONLINE"); targetMap.put("host3", "ONLINE"); Map<String, String> srcMap = current.getRecord().getMapField(segmentId); Assert.assertEquals(srcMap.size(), 2); TableRebalancer updater = new TableRebalancer(null, null, null); updater.updateSegmentIfNeeded(segmentId, srcMap, targetMap, current, downtime); Map<String, String> tempMap = current.getRecord().getMapField(segmentId); Assert.assertTrue(EqualityUtils.isEqual(tempMap, targetMap)); }
@Test public void noDowntimeUpdateWithCommonElements() { Map<String, String> targetMap = new HashMap<>(); targetMap.put("host1", "ONLINE"); targetMap.put("host3", "ONLINE"); Map<String, String> srcMap = current.getRecord().getMapField(segmentId); Assert.assertEquals(srcMap.size(), 2); TableRebalancer updater = new TableRebalancer(null, null, null); updater.updateSegmentIfNeeded(segmentId, srcMap, targetMap, current, noDowntime); Map<String, String> tempMap = current.getRecord().getMapField(segmentId); Assert.assertTrue(EqualityUtils.isEqual(tempMap, targetMap)); }
@Test public void downtimeUpdateWithNoCommonElements() { Map<String, String> targetMap = new HashMap<>(); targetMap.put("host4", "ONLINE"); targetMap.put("host3", "ONLINE"); Map<String, String> srcMap = current.getRecord().getMapField(segmentId); TableRebalancer updater = new TableRebalancer(null, null, null); updater.updateSegmentIfNeeded(segmentId, srcMap, targetMap, current, downtime); Map<String, String> tempMap = current.getRecord().getMapField(segmentId); Assert.assertTrue(EqualityUtils.isEqual(tempMap, targetMap)); } }
public IdealStateBuilderUtil disableIdealState() { _idealState.getRecord().setSimpleField(IdealState.IdealStateProperty.HELIX_ENABLED.name(), "false"); return this; }
private boolean validateNumSegments(int numSegments) { String tableNameWithType = TableNameBuilder.OFFLINE.tableNameWithType(TABLE_NAME); IdealState idealState = _helixAdmin.getResourceIdealState(getHelixClusterName(), tableNameWithType); return idealState.getRecord().getMapFields().keySet().size() == numSegments; }
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; }
public String getSegment(int partition, int seqNum) { for (String segmentName : _idealState.getRecord().getMapFields().keySet()) { if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); if (llcSegmentName.getPartitionId() == partition && llcSegmentName.getSequenceNumber() == seqNum) { return segmentName; } } } return null; }
private boolean validateTableLevelReplicaGroupRebalance() { TableConfig tableConfig = _helixResourceManager.getTableConfig(TABLE_NAME, CommonConstants.Helix.TableType.OFFLINE); String tableNameWithType = TableNameBuilder.OFFLINE.tableNameWithType(TABLE_NAME); ReplicaGroupPartitionAssignmentGenerator partitionAssignmentGenerator = new ReplicaGroupPartitionAssignmentGenerator(_propertyStore); ReplicaGroupPartitionAssignment replicaGroupMapping = partitionAssignmentGenerator.getReplicaGroupPartitionAssignment(tableNameWithType); IdealState idealState = _helixAdmin.getResourceIdealState(getHelixClusterName(), tableNameWithType); Map<String, Map<String, String>> segmentAssignment = idealState.getRecord().getMapFields(); Map<Integer, Set<String>> segmentsPerPartition = new HashMap<>(); segmentsPerPartition.put(0, segmentAssignment.keySet()); return ReplicaGroupTestUtils .validateReplicaGroupSegmentAssignment(tableConfig, replicaGroupMapping, segmentAssignment, segmentsPerPartition); }