Java Code Examples for org.apache.helix.HelixDataAccessor#updateProperty()

/**
 * This function is deprecated. Please use RebalanceUtil.scheduleInstantPipeline method instead.
 * Trigger the controller to perform rebalance for a given resource.
 * @param accessor Helix data accessor
 * @param resource the name of the resource changed to triggering the execution
 */
@Deprecated
public static void invokeRebalance(HelixDataAccessor accessor, String resource) {
  LOG.info("invoke rebalance for " + resource);
  PropertyKey key = accessor.keyBuilder().idealStates(resource);
  IdealState is = accessor.getProperty(key);
  if (is != null) {
    // Here it uses the updateProperty function with no-op DataUpdater. Otherwise, it will use default
    // ZNRecordUpdater which will duplicate elements for listFields.
    if (!accessor.updateProperty(key, znRecord -> znRecord, is)) {
      LOG.warn("Failed to invoke rebalance on resource {}", resource);
    }
  } else {
    LOG.warn("Can't find ideal state for {}", resource);
  }
}
 

/**
 * This function is deprecated. Please use RebalanceUtil.scheduleInstantPipeline method instead.
 * Trigger the controller to perform rebalance for a given resource.
 * @param accessor Helix data accessor
 * @param resource the name of the resource changed to triggering the execution
 */
@Deprecated
public static void invokeRebalanceForResourceConfig(HelixDataAccessor accessor, String resource) {
  LOG.info("invoke rebalance for " + resource);
  PropertyKey key = accessor.keyBuilder().resourceConfig(resource);
  ResourceConfig cfg = accessor.getProperty(key);
  if (cfg != null) {
    // Here it uses the updateProperty function with no-op DataUpdater. Otherwise, it will use default
    // ZNRecordUpdater which will duplicate elements for listFields.
    if (!accessor.updateProperty(key, znRecord -> znRecord, cfg)) {
      LOG.warn("Failed to invoke rebalance on resource config {}", resource);
    }
  } else {
    LOG.warn("Can't find resource config for {}", resource);
  }
}
 

/**
 * Request a state change for a specific task.
 *
 * @param accessor  connected Helix data accessor
 * @param instance  the instance serving the task
 * @param sessionId the current session of the instance
 * @param resource  the job name
 * @param partition the task partition name
 * @param state     the requested state
 * @return true if the request was persisted, false otherwise
 */
private static boolean setRequestedState(HelixDataAccessor accessor, String instance,
    String sessionId, String resource, String partition, TaskPartitionState state) {
  LOG.debug(
      String.format("Requesting a state transition to %s for partition %s.", state, partition));
  try {
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();
    PropertyKey key = keyBuilder.currentState(instance, sessionId, resource);
    CurrentState currStateDelta = new CurrentState(resource);
    currStateDelta.setRequestedState(partition, state.name());

    return accessor.updateProperty(key, currStateDelta);
  } catch (Exception e) {
    LOG.error(String
        .format("Error when requesting a state transition to %s for partition %s.", state,
            partition), e);
    return false;
  }
}
 

public void postHandleMessage() {
  if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
    _batchMsgWrapper.end(_message, _notificationContext);
  }

  // update currentState
  HelixManager manager = _notificationContext.getManager();
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap =
      (ConcurrentHashMap<String, CurrentStateUpdate>) _notificationContext
          .get(MapKey.CURRENT_STATE_UPDATE.toString());

  if (csUpdateMap != null) {
    Map<PropertyKey, CurrentState> csUpdate = mergeCurStateUpdate(csUpdateMap);

    // TODO: change to use asyncSet
    for (PropertyKey key : csUpdate.keySet()) {
      // logger.info("updateCS: " + key);
      // System.out.println("\tupdateCS: " + key.getPath() + ", " +
      // curStateMap.get(key));
      if(!accessor.updateProperty(key, csUpdate.get(key))) {
        LOG.error(
            "Fails to persist current state to ZK for key " + key);
      }
    }
  }
}
 

void handleMessageUsingScheduledTaskQueue(Criteria recipientCriteria, Message messageTemplate,
    String controllerMsgId) {
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  String clusterName = recipientCriteria.getClusterName();
  if (clusterName != null && !clusterName.equals(_manager.getClusterName())) {
    throw new HelixException(String.format(
        "ScheduledTaskQueue cannot send message to another cluster. Local cluster name %s, remote cluster name %s.",
        _manager.getClusterName(), clusterName));
  }

  Map<String, String> sendSummary = new HashMap<String, String>();
  sendSummary.put("MessageCount", "0");
  Map<InstanceType, List<Message>> messages =
      _manager.getMessagingService().generateMessage(recipientCriteria, messageTemplate);

  // Calculate tasks, and put them into the idealState of the SCHEDULER_TASK_QUEUE resource.
  // List field are the destination node, while the Message parameters are stored in the
  // mapFields
  // task throttling can be done on SCHEDULER_TASK_QUEUE resource
  if (messages.size() > 0) {
    String taskQueueName = _message.getRecord().getSimpleField(SCHEDULER_TASK_QUEUE);
    if (taskQueueName == null) {
      throw new HelixException("SchedulerTaskMessage need to have " + SCHEDULER_TASK_QUEUE
          + " specified.");
    }
    IdealState newAddedScheduledTasks = new IdealState(taskQueueName);
    newAddedScheduledTasks.setBucketSize(TASKQUEUE_BUCKET_NUM);
    newAddedScheduledTasks.setStateModelDefRef(SCHEDULER_TASK_QUEUE);

    synchronized (_manager) {
      int existingTopPartitionId = 0;
      IdealState currentTaskQueue =
          _manager.getHelixDataAccessor().getProperty(
              accessor.keyBuilder().idealStates(newAddedScheduledTasks.getId()));
      if (currentTaskQueue != null) {
        existingTopPartitionId = findTopPartitionId(currentTaskQueue) + 1;
      }

      List<Message> taskMessages = (List<Message>) (messages.values().toArray()[0]);
      for (Message task : taskMessages) {
        String partitionId = taskQueueName + "_" + existingTopPartitionId;
        existingTopPartitionId++;
        String instanceName = task.getTgtName();
        newAddedScheduledTasks.setPartitionState(partitionId, instanceName, "COMPLETED");
        task.getRecord().setSimpleField(instanceName, "COMPLETED");
        task.getRecord().setSimpleField(CONTROLLER_MSG_ID, controllerMsgId);

        List<String> priorityList = new LinkedList<String>();
        priorityList.add(instanceName);
        newAddedScheduledTasks.getRecord().setListField(partitionId, priorityList);
        newAddedScheduledTasks.getRecord().setMapField(partitionId,
            task.getRecord().getSimpleFields());
        _logger.info("Scheduling for controllerMsg " + controllerMsgId + " , sending task "
            + partitionId + " " + task.getMsgId() + " to " + instanceName);

        if (_logger.isDebugEnabled()) {
          _logger.debug(task.getRecord().getSimpleFields().toString());
        }
      }
      _manager.getHelixDataAccessor().updateProperty(
          accessor.keyBuilder().idealStates(newAddedScheduledTasks.getId()),
          newAddedScheduledTasks);
      sendSummary.put("MessageCount", "" + taskMessages.size());
    }
  }
  // Record the number of messages sent into scheduler message status updates

  ZNRecord statusUpdate =
      accessor.getProperty(
          keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(),
              _message.getMsgId())).getRecord();

  statusUpdate.getMapFields().put("SentMessageCount", sendSummary);
  accessor.updateProperty(keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(),
      _message.getMsgId()), new StatusUpdate(statusUpdate));
}
 

private void persistAssignment(final Resource resource, final ResourceControllerDataProvider cache,
    final ClusterEvent event, final BestPossibleStateOutput bestPossibleAssignment,
    final ClusterConfig clusterConfig, final HelixDataAccessor accessor,
    final PropertyKey.Builder keyBuilder) {
  String resourceId = resource.getResourceName();
  if (resource != null) {
    final IdealState idealState = cache.getIdealState(resourceId);
    if (idealState == null) {
      LogUtil.logWarn(LOG, event.getEventId(), "IdealState not found for resource " + resourceId);
      return;
    }
    IdealState.RebalanceMode mode = idealState.getRebalanceMode();
    if (!mode.equals(IdealState.RebalanceMode.SEMI_AUTO) && !mode
        .equals(IdealState.RebalanceMode.FULL_AUTO)) {
      // do not persist assignment for resource in neither semi or full auto.
      return;
    }

    boolean needPersist = false;
    if (mode.equals(IdealState.RebalanceMode.FULL_AUTO)) {
      // persist preference list in ful-auto mode.
      Map<String, List<String>> newLists = bestPossibleAssignment.getPreferenceLists(resourceId);
      if (newLists != null && hasPreferenceListChanged(newLists, idealState)) {
        idealState.setPreferenceLists(newLists);
        needPersist = true;
      }
    }

    PartitionStateMap partitionStateMap = bestPossibleAssignment.getPartitionStateMap(resourceId);
    if (clusterConfig.isPersistIntermediateAssignment()) {
      IntermediateStateOutput intermediateAssignment =
          event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());
      partitionStateMap = intermediateAssignment.getPartitionStateMap(resourceId);
    }

    //TODO: temporary solution for Espresso/Dbus backcompatible, should remove this.
    Map<Partition, Map<String, String>> assignmentToPersist =
        convertAssignmentPersisted(resource, idealState, partitionStateMap.getStateMap());

    if (assignmentToPersist != null && hasInstanceMapChanged(assignmentToPersist, idealState)) {
      for (Partition partition : assignmentToPersist.keySet()) {
        Map<String, String> instanceMap = assignmentToPersist.get(partition);
        idealState.setInstanceStateMap(partition.getPartitionName(), instanceMap);
      }
      needPersist = true;
    }

    if (needPersist) {
      // Update instead of set to ensure any intermediate changes that the controller does not update are kept.
      accessor.updateProperty(keyBuilder.idealStates(resourceId), new DataUpdater<ZNRecord>() {
        @Override
        public ZNRecord update(ZNRecord current) {
          if (current != null) {
            // Overwrite MapFields and ListFields items with the same key.
            // Note that default merge will keep old values in the maps or lists unchanged, which is not desired.
            current.getMapFields().clear();
            current.getMapFields().putAll(idealState.getRecord().getMapFields());
            current.getListFields().putAll(idealState.getRecord().getListFields());
          }
          return current;
        }
      }, idealState);
    }
  }
}
 

/**
 * Write an error record to zookeeper to the zookeeper store.
 * @param record
 *          the status update record
 * @param instanceName
 *          the instance name
 * @param updateSubPath
 *          the error update sub path
 * @param updateKey
 *          the error update key
 * @param sessionId
 *          the session id
 * @param accessor
 *          the zookeeper data accessor that writes the status update to zookeeper
 * @param isController
 *          if the error log is for a controller instance or not
 */
void publishErrorRecord(ZNRecord record, String instanceName, String updateSubPath,
    String updateKey, String sessionId, HelixDataAccessor accessor, boolean isController) {
  Builder keyBuilder = accessor.keyBuilder();
  if (isController) {
    // TODO need to fix: ERRORS_CONTROLLER doesn't have a form of
    // ../{sessionId}/{subPath}
    accessor.setProperty(keyBuilder.controllerTaskError(updateSubPath), new Error(record));
  } else {
    accessor.updateProperty(keyBuilder.stateTransitionError(instanceName, sessionId,
        updateSubPath, updateKey), new Error(record));
  }
}