org.apache.helix.provisioning.tools.UpdateProvisionerConfig Maven / Gradle / Ivy

Go to download
package org.apache.helix.provisioning.tools;

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.OptionGroup;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.helix.HelixConnection;
import org.apache.helix.api.Resource;
import org.apache.helix.api.accessor.ClusterAccessor;
import org.apache.helix.api.config.ResourceConfig;
import org.apache.helix.api.id.ClusterId;
import org.apache.helix.api.id.ResourceId;
import org.apache.helix.manager.zk.ZkHelixConnection;
import org.apache.helix.provisioning.yarn.YarnProvisionerConfig;
import org.apache.log4j.Logger;

/**
 * Update the provisioner config
 */
public class UpdateProvisionerConfig {
  private static Logger LOG = Logger.getLogger(UpdateProvisionerConfig.class);
  private static String updateContainerCount = "updateContainerCount";
  private HelixConnection _connection;

  public UpdateProvisionerConfig(String zkAddress) {
    _connection = new ZkHelixConnection(zkAddress);
    _connection.connect();
  }

  public void setNumContainers(String appName, String serviceName, int numContainers) {
    ResourceId resourceId = ResourceId.from(serviceName);

    ClusterAccessor clusterAccessor = _connection.createClusterAccessor(ClusterId.from(appName));
    Resource resource = clusterAccessor.readResource(resourceId);
    LOG.info("Current provisioner config:" + resource.getProvisionerConfig());

    ResourceConfig.Delta delta = new ResourceConfig.Delta(resourceId);
    YarnProvisionerConfig config = new YarnProvisionerConfig(resourceId);
    config.setNumContainers(numContainers);
    delta.setProvisionerConfig(config);
    delta.setIdealState(resource.getIdealState());
    ResourceConfig updatedResourceConfig = clusterAccessor.updateResource(resourceId, delta);
    LOG.info("Update provisioner config:" + updatedResourceConfig.getProvisionerConfig());

  }

  @SuppressWarnings("static-access")
  public static void main(String[] args) throws ParseException {
    Option zkServerOption =
        OptionBuilder.withLongOpt("zookeeperAddress").withDescription("Provide zookeeper address")
            .create();
    zkServerOption.setArgs(1);
    zkServerOption.setRequired(true);
    zkServerOption.setArgName("zookeeperAddress(Required)");

    OptionGroup group = new OptionGroup();
    group.setRequired(true);

    // update container count per service
    Option updateContainerCountOption =
        OptionBuilder.withLongOpt(updateContainerCount)
            .withDescription("appName serviceName numContainers").create();
    updateContainerCountOption.setArgs(3);
    updateContainerCountOption.setRequired(false);
    updateContainerCountOption.setArgName("appName serviceName numContainers");

    group.addOption(updateContainerCountOption);

    Options options = new Options();
    options.addOption(zkServerOption);
    options.addOptionGroup(group);
    CommandLine cliParser = new GnuParser().parse(options, args);

    String zkAddress = cliParser.getOptionValue("zookeeperAddress");
    UpdateProvisionerConfig updater = new UpdateProvisionerConfig(zkAddress);

    if (cliParser.hasOption(updateContainerCount)) {
      String appName = cliParser.getOptionValues(updateContainerCount)[0];
      String serviceName = cliParser.getOptionValues(updateContainerCount)[1];
      int numContainers = Integer.parseInt(cliParser.getOptionValues(updateContainerCount)[2]);
      updater.setNumContainers(appName, serviceName, numContainers);
    }

  }
}