Java Code Examples for org.apache.helix.HelixManager#disconnect()

@Test (groups = {"disabledOnTravis"})
public void testTaskAssignmentAfterHelixConnectionRetry()
    throws Exception {
  Config jobConfigOverrides = ClusterIntegrationTestUtils.buildSleepingJob(JOB_ID, TASK_STATE_FILE);
  this.suite = new TaskAssignmentAfterConnectionRetry(jobConfigOverrides);

  String zkConnectString = suite.getManagerConfig().getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  String clusterName = suite.getManagerConfig().getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);
  //A test manager instance for observing the state of the cluster
  HelixManager helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "TestManager", InstanceType.SPECTATOR, zkConnectString);

  suite.startCluster();

  helixManager.connect();

  //Ensure that Helix has created a workflow
  AssertWithBackoff.create().maxSleepMs(1000).backoffFactor(1).
      assertTrue(ClusterIntegrationTest.isTaskStarted(helixManager, JOB_ID), "Waiting for the job to start...");

  //Ensure that the SleepingTask is running
  AssertWithBackoff.create().maxSleepMs(100).timeoutMs(2000).backoffFactor(1).
      assertTrue(ClusterIntegrationTest.isTaskRunning(TASK_STATE_FILE),"Waiting for the task to enter running state");

  helixManager.disconnect();
}
 

public static void main(String[] args) throws Exception {
  CommandLine cmd = processCommandLineArgs(args);
  String zkAddress = cmd.getOptionValue(zkAddr);
  String clusterName = cmd.getOptionValue(cluster);
  String instance = cmd.getOptionValue(instanceName);
  String stateModelName = cmd.getOptionValue(stateModel);

  HelixManager manager =
      new ZKHelixManager(clusterName, instance, InstanceType.PARTICIPANT, zkAddress);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory(stateModelName, new AgentStateModelFactory());

  Runtime.getRuntime().addShutdownHook(new HelixAgentShutdownHook(manager));

  try {
    manager.connect();
    Thread.currentThread().join();
  } catch (Exception e) {
    LOG.error(e.toString());
  } finally {
    if (manager != null && manager.isConnected()) {
      manager.disconnect();
    }
  }
}
 

@Test()
public void testControllerLeaderHistory() throws Exception {
  HelixManager manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  manager.connect();

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
  PropertyKey propertyKey = keyBuilder.controllerLeaderHistory();
  ControllerHistory controllerHistory = manager.getHelixDataAccessor().getProperty(propertyKey);
  Assert.assertNotNull(controllerHistory);
  List<String> list = controllerHistory.getRecord().getListField("HISTORY");
  Assert.assertEquals(list.size(), 1);

  for (int i = 0; i <= 12; i++) {
    _controller.syncStop();
    _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, "Controller-" + i);
    _controller.syncStart();
  }

  controllerHistory = manager.getHelixDataAccessor().getProperty(propertyKey);
  Assert.assertNotNull(controllerHistory);
  list = controllerHistory.getRecord().getListField("HISTORY");
  Assert.assertEquals(list.size(), 10);
  manager.disconnect();
}
 

@Test
public void testParticipant() throws Exception {
  Map<String, Integer> stateReachedCounts = Maps.newHashMap();
  HelixManager participant =
      HelixManagerFactory.getZKHelixManager(_clusterName, "localhost_12918",
          InstanceType.PARTICIPANT, ZK_ADDR);
  participant.getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new MyStateModelFactory(stateReachedCounts));
  participant.connect();

  Thread.sleep(1000);

  // Ensure that the external view coalesces
  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          _clusterName));
  Assert.assertTrue(result);

  // Ensure that there was only one state transition
  Assert.assertEquals(stateReachedCounts.size(), 1);
  Assert.assertTrue(stateReachedCounts.containsKey("ONLINE"));
  Assert.assertEquals(stateReachedCounts.get("ONLINE").intValue(), 1);

  // Now let's stop the ZK server; this should do nothing
  TestHelper.stopZkServer(_zkServer);
  Thread.sleep(1000);

  // Verify no change
  Assert.assertEquals(stateReachedCounts.size(), 1);
  Assert.assertTrue(stateReachedCounts.containsKey("ONLINE"));
  Assert.assertEquals(stateReachedCounts.get("ONLINE").intValue(), 1);

  if (participant.isConnected()) {
    participant.disconnect();
  }
}
 

protected void stopCurrentLeader(HelixZkClient zkClient, String clusterName,
    Map<String, Thread> threadMap, Map<String, HelixManager> managerMap) {
  String leader = getCurrentLeader(zkClient, clusterName);
  Assert.assertTrue(leader != null);
  System.out.println("stop leader:" + leader + " in " + clusterName);
  Assert.assertTrue(leader != null);

  HelixManager manager = managerMap.remove(leader);
  Assert.assertTrue(manager != null);
  manager.disconnect();

  Thread thread = threadMap.remove(leader);
  Assert.assertTrue(thread != null);
  thread.interrupt();

  boolean isNewLeaderElected = false;
  try {
    // Thread.sleep(2000);
    for (int i = 0; i < 5; i++) {
      Thread.sleep(1000);
      String newLeader = getCurrentLeader(zkClient, clusterName);
      if (!newLeader.equals(leader)) {
        isNewLeaderElected = true;
        System.out.println("new leader elected: " + newLeader + " in " + clusterName);
        break;
      }
    }
  } catch (InterruptedException e) {
    e.printStackTrace();
  }
  if (!isNewLeaderElected) {
    System.out.println("fail to elect a new leader elected in " + clusterName);
  }
  AssertJUnit.assertTrue(isNewLeaderElected);
}
 

@Test
public void testDelayedRebalancer() throws Exception {
  System.out.println("START testDelayedRebalancer at " + new Date(System.currentTimeMillis()));
  HelixManager manager =
      HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, null, InstanceType.SPECTATOR, ZK_ADDR);
  manager.connect();
  manager.addExternalViewChangeListener(this);
  manager.addIdealStateChangeListener(this);
  enablePersistBestPossibleAssignment(_gZkClient, CLUSTER_NAME, true);

  // Start half number of nodes.
  int i = 0;
  for (; i < NUM_NODE / 2; i++) {
    _participants.get(i).syncStart();
  }

  int replica = 3;
  int partition = 30;
  for (String stateModel : TestStateModels) {
    String db = "Test-DB-" + stateModel;
    createResourceWithDelayedRebalance(CLUSTER_NAME, db, stateModel, partition, replica, replica,
        0);
  }
  Assert.assertTrue(_clusterVerifier.verifyByPolling(50000L, 100L));

  _startListen = true;
  DelayedTransition.setDelay(5);

  // add the other half of nodes.
  for (; i < NUM_NODE; i++) {
    _participants.get(i).syncStart();
  }
  Assert.assertTrue(_clusterVerifier.verify(70000L));
  Assert.assertTrue(_testSuccess);

  if (manager.isConnected()) {
    manager.disconnect();
  }
  System.out.println("END testDelayedRebalancer at " + new Date(System.currentTimeMillis()));
}
 

@SuppressWarnings("deprecation")
@Test
public void testSpectator() throws Exception {
  Map<String, Integer> stateReachedCounts = Maps.newHashMap();
  HelixManager participant =
      HelixManagerFactory.getZKHelixManager(_clusterName, "localhost_12918",
          InstanceType.PARTICIPANT, ZK_ADDR);
  participant.getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new MyStateModelFactory(stateReachedCounts));
  participant.connect();

  RoutingTableProvider routingTableProvider = new RoutingTableProvider();
  try {
    HelixManager spectator = HelixManagerFactory
        .getZKHelixManager(_clusterName, "spectator", InstanceType.SPECTATOR, ZK_ADDR);
    spectator.connect();
    spectator.addConfigChangeListener(routingTableProvider);
    spectator.addExternalViewChangeListener(routingTableProvider);
    Thread.sleep(1000);

    // Now let's stop the ZK server; this should do nothing
    TestHelper.stopZkServer(_zkServer);
    Thread.sleep(1000);

    // Verify routing table still works
    Assert.assertEquals(routingTableProvider.getInstances("resource0", "ONLINE").size(), 1);
    Assert.assertEquals(routingTableProvider.getInstances("resource0", "OFFLINE").size(), 0);
  } finally {
    routingTableProvider.shutdown();
    if (participant.isConnected()) {
      participant.disconnect();
    }
  }
}
 

public static void main(String[] args) {
  if (args.length < 3) {
    System.err.println("USAGE: DummyParticipant zkAddress clusterName instanceName");
    System.exit(1);
  }

  String zkAddr = args[0];
  String clusterName = args[1];
  String instanceName = args[2];

  HelixManager manager = null;
  try {
    manager =
        HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
            InstanceType.PARTICIPANT, zkAddr);

    StateMachineEngine stateMach = manager.getStateMachineEngine();
    DummyMSModelFactory msModelFactory = new DummyMSModelFactory();
    stateMach.registerStateModelFactory("MasterSlave", msModelFactory);

    manager.connect();

    Thread.currentThread().join();
  } catch (Exception e) {
    // TODO Auto-generated catch block
    e.printStackTrace();
  } finally {
    if (manager != null) {
      manager.disconnect();
    }
  }
}
 

@Test
public void simpleIntegrationTest() throws Exception {
  int n = 1;

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      4, // partitions per resource
      n, // number of nodes
      1, // replicas
      "MasterSlave", true); // do rebalance

  HelixManager participant =
      new ZKHelixManager(clusterName, "localhost_12918", InstanceType.PARTICIPANT, ZK_ADDR);
  participant.getStateMachineEngine().registerStateModelFactory("MasterSlave",
      new MockMSModelFactory());
  participant.connect();

  HelixManager controller =
      new ZKHelixManager(clusterName, "controller_0", InstanceType.CONTROLLER, ZK_ADDR);
  controller.connect();

  verifyHelixManagerMetrics(InstanceType.PARTICIPANT, MonitorLevel.DEFAULT,
      participant.getInstanceName());
  verifyHelixManagerMetrics(InstanceType.CONTROLLER, MonitorLevel.DEFAULT,
      controller.getInstanceName());

  BestPossibleExternalViewVerifier verifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkClient(_gZkClient)
          .setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling());

  // cleanup
  controller.disconnect();
  participant.disconnect();

  // verify all live-instances and leader nodes are gone
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_gZkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  Assert.assertNull(accessor.getProperty(keyBuilder.liveInstance("localhost_12918")));
  Assert.assertNull(accessor.getProperty(keyBuilder.controllerLeader()));
}
 

@Test
public void testMissingTopStateDurationMonitoring() throws Exception {
  String clusterName = "testCluster-TestControllerLeadershipChange";
  String instanceName = clusterName + "-participant";
  String resourceName = "testResource";
  int numPartition = 1;
  int numReplica = 1;
  String stateModel = "LeaderStandby";
  ObjectName resourceMBeanObjectName = getResourceMonitorObjectName(clusterName, resourceName);
  MBeanServer beanServer = ManagementFactory.getPlatformMBeanServer();

  // Create cluster
  _gSetupTool.addCluster(clusterName, true);

  // Create cluster verifier
  ZkHelixClusterVerifier clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkClient(_gZkClient).build();

  // Create participant
  _gSetupTool.addInstanceToCluster(clusterName, instanceName);
  MockParticipantManager participant =
      new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
  participant.syncStart();

  // Create controller, since this is the only controller, it will be the leader
  HelixManager manager1 = HelixManagerFactory
      .getZKHelixManager(clusterName, clusterName + "-manager1", InstanceType.CONTROLLER,
          ZK_ADDR);
  manager1.connect();
  Assert.assertTrue(manager1.isLeader());

  // Create resource
  _gSetupTool.addResourceToCluster(clusterName, resourceName, numPartition, stateModel,
      IdealState.RebalanceMode.SEMI_AUTO.name());

  // Rebalance Resource
  _gSetupTool.rebalanceResource(clusterName, resourceName, numReplica);

  // Wait for rebalance
  Assert.assertTrue(clusterVerifier.verifyByPolling());

  // Trigger missing top state in manager1
  participant.syncStop();

  Thread.sleep(1000);

  // Starting manager2
  HelixManager manager2 = HelixManagerFactory
      .getZKHelixManager(clusterName, clusterName + "-manager2", InstanceType.CONTROLLER,
          ZK_ADDR);
  manager2.connect();

  // Set leader to manager2
  setLeader(manager2);

  Assert.assertFalse(manager1.isLeader());
  Assert.assertTrue(manager2.isLeader());

  // Wait for rebalance
  Assert.assertTrue(clusterVerifier.verify());

  Thread.sleep(1000);
  setLeader(manager1);

  Assert.assertTrue(manager1.isLeader());
  Assert.assertFalse(manager2.isLeader());

  // Make resource top state to come back by restarting participant
  participant = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
  participant.syncStart();

  _gSetupTool.rebalanceResource(clusterName, resourceName, numReplica);

  Assert.assertTrue(clusterVerifier.verifyByPolling());

  // Resource lost top state, and manager1 lost leadership for 2000ms, because manager1 will
  // clean monitoring cache after re-gaining leadership, so max value of hand off duration should
  // not have such a large value
  Assert.assertTrue((long) beanServer
      .getAttribute(resourceMBeanObjectName, "PartitionTopStateHandoffDurationGauge.Max") < 500);

  participant.syncStop();
  manager1.disconnect();
  manager2.disconnect();
  deleteCluster(clusterName);
}
 

@Override
public Representation post(Representation entity) {
  try {
    String clusterName = (String) getRequest().getAttributes().get("clusterName");
    Form form = new Form(entity);

    // Get the workflow and submit it
    if (form.size() < 1) {
      throw new HelixException("yaml workflow is required!");
    }
    Parameter payload = form.get(0);
    String yamlPayload = payload.getName();
    if (yamlPayload == null) {
      throw new HelixException("yaml workflow is required!");
    }
    String zkAddr =
        (String) getContext().getAttributes().get(RestAdminApplication.ZKSERVERADDRESS);
    HelixManager manager =
        HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.ADMINISTRATOR,
            zkAddr);
    manager.connect();
    try {
      Workflow workflow = Workflow.parse(yamlPayload);
      TaskDriver driver = new TaskDriver(manager);
      driver.start(workflow);
    } finally {
      manager.disconnect();
    }

    getResponse().setEntity(getHostedEntitiesRepresentation(clusterName));
    getResponse().setStatus(Status.SUCCESS_OK);
  }

  catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in posting " + entity, e);
  }
  return null;
}
 

protected void stopFakeInstances() {
  for (HelixManager helixManager : _fakeInstanceHelixManagers) {
    helixManager.disconnect();
  }
  _fakeInstanceHelixManagers.clear();
}
 

/**
 * Parses the first argument as a driver command and the rest of the
 * arguments are parsed based on that command. Constructs a Helix
 * message and posts it to the controller
 */
public static void main(String[] args) throws Exception {
  String[] cmdArgs = Arrays.copyOfRange(args, 1, args.length);
  CommandLine cl = parseOptions(cmdArgs, constructOptions(), args[0]);
  String zkAddr = cl.getOptionValue(ZK_ADDRESS);
  String clusterName = cl.getOptionValue(CLUSTER_NAME_OPTION);
  String workflow = cl.getOptionValue(RESOURCE_OPTION);

  if (zkAddr == null || clusterName == null || workflow == null) {
    printUsage(constructOptions(), "[cmd]");
    throw new IllegalArgumentException(
        "zk, cluster, and resource must all be non-null for all commands");
  }

  HelixManager helixMgr =
      HelixManagerFactory.getZKHelixManager(clusterName, "Admin", InstanceType.ADMINISTRATOR,
          zkAddr);
  helixMgr.connect();
  TaskDriver driver = new TaskDriver(helixMgr);
  try {
    TaskDriver.DriverCommand cmd = TaskDriver.DriverCommand.valueOf(args[0]);
    switch (cmd) {
    case start:
      if (cl.hasOption(WORKFLOW_FILE_OPTION)) {
        driver.start(Workflow.parse(new File(cl.getOptionValue(WORKFLOW_FILE_OPTION))));
      } else {
        throw new IllegalArgumentException("Workflow file is required to start flow.");
      }
      break;
    case stop:
      driver.stop(workflow);
      break;
    case resume:
      driver.resume(workflow);
      break;
    case delete:
      driver.delete(workflow);
      break;
    case list:
      list(driver, workflow);
      break;
    case flush:
      driver.flushQueue(workflow);
      break;
    case clean:
      driver.cleanupQueue(workflow);
      break;
    default:
      throw new IllegalArgumentException("Unknown command " + args[0]);
    }
  } catch (IllegalArgumentException e) {
    LOG.error("Unknown driver command " + args[0]);
    throw e;
  }

  helixMgr.disconnect();
}
 

/**
 * LockManagerDemo clusterName, numInstances, lockGroupName, numLocks
 * @param args
 * @throws Exception
 */
public static void main(String[] args) throws Exception {
  final String zkAddress = "localhost:2199";
  final String clusterName = "lock-manager-demo";
  final String lockGroupName = "lock-group";
  final int numInstances = 3;
  final int numPartitions = 12;
  final boolean startController = false;
  HelixManager controllerManager = null;
  Thread[] processArray;
  processArray = new Thread[numInstances];
  try {
    startLocalZookeeper(2199);
    HelixAdmin admin = new ZKHelixAdmin(zkAddress);
    admin.addCluster(clusterName, true);
    StateModelConfigGenerator generator = new StateModelConfigGenerator();
    admin.addStateModelDef(clusterName, "OnlineOffline",
        new StateModelDefinition(generator.generateConfigForOnlineOffline()));
    admin.addResource(clusterName, lockGroupName, numPartitions, "OnlineOffline",
        RebalanceMode.FULL_AUTO.toString());
    admin.rebalance(clusterName, lockGroupName, 1);
    for (int i = 0; i < numInstances; i++) {
      final String instanceName = "localhost_" + (12000 + i);
      processArray[i] = new Thread(new Runnable() {

        @Override
        public void run() {
          LockProcess lockProcess = null;

          try {
            lockProcess = new LockProcess(clusterName, zkAddress, instanceName, startController);
            lockProcess.start();
            Thread.currentThread().join();
          } catch (InterruptedException e) {
            System.out.println(instanceName + "Interrupted");
            if (lockProcess != null) {
              lockProcess.stop();
            }
          } catch (Exception e) {
            e.printStackTrace();
          }
        }

      });
      processArray[i].start();
    }
    Thread.sleep(3000);
    controllerManager =
        HelixControllerMain.startHelixController(zkAddress, clusterName, "controller",
            HelixControllerMain.STANDALONE);
    Thread.sleep(5000);
    printStatus(admin, clusterName, lockGroupName);
    System.out.println("Stopping localhost_12000");
    processArray[0].interrupt();
    Thread.sleep(3000);
    printStatus(admin, clusterName, lockGroupName);
    Thread.currentThread().join();
  } catch (Exception e) {
    e.printStackTrace();
  } finally {
    if (controllerManager != null) {
      controllerManager.disconnect();
    }
    for (Thread process : processArray) {
      if (process != null) {
        process.interrupt();
      }
    }
  }
}
 

@Test
public void testWagedRebalancer() throws Exception {
  System.out.println("START testWagedRebalancer at " + new Date(System.currentTimeMillis()));
  HelixManager manager =
      HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, null, InstanceType.SPECTATOR, ZK_ADDR);
  manager.connect();
  manager.addExternalViewChangeListener(this);
  manager.addIdealStateChangeListener(this);
  enablePersistBestPossibleAssignment(_gZkClient, CLUSTER_NAME, true);

  // Start half number of nodes.
  int i = 0;
  for (; i < NUM_NODE / 2; i++) {
    _participants.get(i).syncStart();
  }

  int replica = 3;
  int partition = 30;
  for (String stateModel : TestStateModels) {
    String db = "Test-DB-" + stateModel;
    createResourceWithWagedRebalance(CLUSTER_NAME, db, stateModel, partition, replica, replica);
  }
  // TODO remove this sleep after fix https://github.com/apache/helix/issues/526
  Thread.sleep(1000);
  Assert.assertTrue(_clusterVerifier.verifyByPolling(50000L, 100L));

  _startListen = true;
  DelayedTransition.setDelay(5);

  // add the other half of nodes.
  for (; i < NUM_NODE; i++) {
    _participants.get(i).syncStart();
  }
  Assert.assertTrue(_clusterVerifier.verify(70000L));
  Assert.assertTrue(_testSuccess);

  if (manager.isConnected()) {
    manager.disconnect();
  }
  System.out.println("END testWagedRebalancer at " + new Date(System.currentTimeMillis()));
}
 

@Test
public void testDataChange() throws Exception {
  // Create an admin and add LiveInstanceChange listener to it
  HelixManager adminManager =
      HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.ADMINISTRATOR,
          ZK_ADDR);
  adminManager.connect();
  final BlockingQueue<List<LiveInstance>> changeList =
      new LinkedBlockingQueue<List<LiveInstance>>();

  adminManager.addLiveInstanceChangeListener(new LiveInstanceChangeListener() {
    @Override
    public void onLiveInstanceChange(List<LiveInstance> liveInstances,
        NotificationContext changeContext) {
      // The queue is basically unbounded, so shouldn't throw exception when calling
      // "add".
      changeList.add(deepCopy(liveInstances));
    }
  });

  // Check the initial condition
  List<LiveInstance> instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertTrue(instances.isEmpty(), "Expecting an empty list of live instance");
  // Join as participant, should trigger a live instance change event
  HelixManager manager =
      HelixManagerFactory.getZKHelixManager(clusterName, "localhost_54321",
          InstanceType.PARTICIPANT, ZK_ADDR);
  manager.connect();
  instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertEquals(instances.size(), 1, "Expecting one live instance");
  Assert.assertEquals(instances.get(0).getInstanceName(), manager.getInstanceName());
  // Update data in the live instance node, should trigger another live instance change
  // event
  HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
  PropertyKey propertyKey =
      helixDataAccessor.keyBuilder().liveInstance(manager.getInstanceName());
  LiveInstance instance = helixDataAccessor.getProperty(propertyKey);

  Map<String, String> map = new TreeMap<String, String>();
  map.put("k1", "v1");
  instance.getRecord().setMapField("test", map);
  Assert.assertTrue(helixDataAccessor.updateProperty(propertyKey, instance),
      "Failed to update live instance node");

  instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertEquals(instances.get(0).getRecord().getMapField("test"), map, "Wrong map data.");
  manager.disconnect();
  Thread.sleep(1000); // wait for callback finish

  instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertTrue(instances.isEmpty(), "Expecting an empty list of live instance");

  adminManager.disconnect();

}
 

@Test
public void testCachedCustomizedViews() throws Exception {
  String clusterName = "CLUSTER_" + TestHelper.getTestMethodName();

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  HelixManager manager = new DummyClusterManager(clusterName, accessor);

  // ideal state: node0 is MASTER, node1 is SLAVE
  // replica=2 means 1 master and 1 slave
  setupIdealState(clusterName, new int[] {
      0, 1
  }, new String[] {
      "TestDB"
  }, 1, 2);
  setupLiveInstances(clusterName, new int[] {
      0, 1
  });
  setupStateModel(clusterName);

  ClusterEvent event = new ClusterEvent(ClusterEventType.Unknown);
  ResourceControllerDataProvider cache = new ResourceControllerDataProvider(clusterName);
  event.addAttribute(AttributeName.helixmanager.name(), manager);
  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);

  CustomizedStateConfig config = new CustomizedStateConfig();
  List<String> aggregationEnabledTypes = new ArrayList<>();
  aggregationEnabledTypes.add(CUSTOMIZED_STATE_NAME);
  config.setAggregationEnabledTypes(aggregationEnabledTypes);

  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.customizedStateConfig(), config);

  CustomizedState customizedState = new CustomizedState(RESOURCE_NAME);
  customizedState.setState(PARTITION_NAME, "STARTED");
  accessor
      .setProperty(keyBuilder.customizedState(INSTANCE_NAME, "customizedState1", RESOURCE_NAME),
          customizedState);

  CustomizedViewAggregationStage customizedViewComputeStage =
      new CustomizedViewAggregationStage();
  Pipeline dataRefresh = new Pipeline();
  dataRefresh.addStage(new ReadClusterDataStage());
  runPipeline(event, dataRefresh);
  runStage(event, new ResourceComputationStage());
  runStage(event, new CustomizedStateComputationStage());
  runStage(event, customizedViewComputeStage);
  Assert.assertEquals(cache.getCustomizedViewCacheMap().size(),
      accessor.getChildNames(accessor.keyBuilder().customizedViews()).size());

  // Assure there is no customized view got updated when running the stage again
  List<CustomizedView> oldCustomizedViews =
      accessor.getChildValues(accessor.keyBuilder().customizedViews(), true);
  runStage(event, customizedViewComputeStage);
  List<CustomizedView> newCustomizedViews =
      accessor.getChildValues(accessor.keyBuilder().customizedViews(), true);
  Assert.assertEquals(oldCustomizedViews, newCustomizedViews);
  for (int i = 0; i < oldCustomizedViews.size(); i++) {
    Assert.assertEquals(oldCustomizedViews.get(i).getStat().getVersion(),
        newCustomizedViews.get(i).getStat().getVersion());
  }

  if (manager.isConnected()) {
    manager.disconnect(); // For DummyClusterManager, this is not necessary
  }
  deleteLiveInstances(clusterName);
  deleteCluster(clusterName);
}
 

@Test
public void testCachedExternalViews() throws Exception {
  String clusterName = "CLUSTER_" + TestHelper.getTestMethodName();

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  HelixManager manager = new DummyClusterManager(clusterName, accessor);

  // ideal state: node0 is MASTER, node1 is SLAVE
  // replica=2 means 1 master and 1 slave
  setupIdealState(clusterName, new int[] {
      0, 1
  }, new String[] {
      "TestDB"
  }, 1, 2);
  setupLiveInstances(clusterName, new int[] {
      0, 1
  });
  setupStateModel(clusterName);

  ClusterEvent event = new ClusterEvent(ClusterEventType.Unknown);
  ResourceControllerDataProvider cache = new ResourceControllerDataProvider(clusterName);
  event.addAttribute(AttributeName.helixmanager.name(), manager);
  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);

  ExternalViewComputeStage externalViewComputeStage = new ExternalViewComputeStage();
  Pipeline dataRefresh = new Pipeline();
  dataRefresh.addStage(new ReadClusterDataStage());
  runPipeline(event, dataRefresh);
  runStage(event, new ResourceComputationStage());
  runStage(event, new CurrentStateComputationStage());
  runStage(event, externalViewComputeStage);
  Assert.assertEquals(cache.getExternalViews().values(),
      accessor.getChildValues(accessor.keyBuilder().externalViews(), true));

  // Assure there is no external got updated
  List<ExternalView> oldExternalViews =
      accessor.getChildValues(accessor.keyBuilder().externalViews(), true);
  runStage(event, externalViewComputeStage);
  List<ExternalView> newExternalViews =
      accessor.getChildValues(accessor.keyBuilder().externalViews(), true);
  Assert.assertEquals(oldExternalViews, newExternalViews);
  for (int i = 0; i < oldExternalViews.size(); i++) {
    Assert.assertEquals(oldExternalViews.get(i).getStat().getVersion(),
        newExternalViews.get(i).getStat().getVersion());
  }

  if (manager.isConnected()) {
    manager.disconnect(); // For DummyClusterManager, this is not necessary
  }
  deleteLiveInstances(clusterName);
  deleteCluster(clusterName);
}
 

@Test (groups = {"disabledOnTravis"})
//Test disabled on Travis because cluster integration tests are generally flaky on Travis.
public void testExecute() throws Exception {
  suite.startCluster();

  //Connect to the previously started Helix cluster
  helixManager.connect();

  //Ensure that Helix has created a workflow
  AssertWithBackoff.create().maxSleepMs(1000).backoffFactor(1).
      assertTrue(ClusterIntegrationTest.isTaskStarted(helixManager, JOB_ID), "Waiting for the job to start...");

  //Instantiate config for HelixAssignedParticipantCheck
  String helixJobId = Joiner.on("_").join(JOB_ID, JOB_ID);
  helixConfig = helixConfig.withValue(GobblinClusterConfigurationKeys.HELIX_INSTANCE_NAME_KEY,
      ConfigValueFactory.fromAnyRef(IntegrationBasicSuite.WORKER_INSTANCE_0))
      .withValue(GobblinClusterConfigurationKeys.HELIX_JOB_ID_KEY, ConfigValueFactory.fromAnyRef(helixJobId))
      .withValue(GobblinClusterConfigurationKeys.HELIX_PARTITION_ID_KEY, ConfigValueFactory.fromAnyRef(0));
  HelixAssignedParticipantCheck check = new HelixAssignedParticipantCheck(helixConfig);

  //Ensure that the SleepingTask is running
  AssertWithBackoff.create().maxSleepMs(100).timeoutMs(2000).backoffFactor(1).
      assertTrue(ClusterIntegrationTest.isTaskRunning(TASK_STATE_FILE),"Waiting for the task to enter running state");

  //Run the check. Ensure that the configured Helix instance is indeed the assigned participant
  // (i.e. no exceptions thrown).
  check.execute();

  //Disconnect the helixmanager used to check the assigned participant to force an Exception on the first attempt.
  //The test should succeed on the following attempt.
  HelixManager helixManagerOriginal = HelixAssignedParticipantCheck.getHelixManager();
  helixManagerOriginal.disconnect();
  check.execute();
  //Ensure that a new HelixManager instance is created.
  Assert.assertTrue(HelixAssignedParticipantCheck.getHelixManager() != helixManagerOriginal);

  //Create Helix config with invalid partition num. Ensure HelixAssignedParticipantCheck fails.
  helixConfig = helixConfig.withValue(GobblinClusterConfigurationKeys.HELIX_PARTITION_ID_KEY, ConfigValueFactory.fromAnyRef(1));
  check = new HelixAssignedParticipantCheck(helixConfig);

  try {
    check.execute();
    Assert.fail("Expected to throw CommitStepException");
  } catch (CommitStepException e) {
    //Expected to throw CommitStepException
    Assert.assertTrue(e.getClass().equals(CommitStepException.class));
  }
}
 

public void deregister(final String serviceId) {
  HelixManager helixManager = serviceMap.get(serviceId);
  if (helixManager != null && helixManager.isConnected()) {
    helixManager.disconnect();
  }
}