@Test
  public void nonLeaderTest() throws Exception {
    final String tableName = "myTable";
    List<String> allTableNames = new ArrayList<String>();
    allTableNames.add(tableName);

    HelixAdmin helixAdmin;
    {
      helixAdmin = mock(HelixAdmin.class);
    }
    {
      helixResourceManager = mock(PinotHelixResourceManager.class);
      when(helixResourceManager.isLeader()).thenReturn(false);
      when(helixResourceManager.getAllPinotTableNames()).thenReturn(allTableNames);
      when(helixResourceManager.getHelixClusterName()).thenReturn("StatusChecker");
      when(helixResourceManager.getHelixAdmin()).thenReturn(helixAdmin);
    }
    {
      config = mock(ControllerConf.class);
      when(config.getStatusControllerFrequencyInSeconds()).thenReturn(300);
    }
    metricsRegistry = new MetricsRegistry();
    controllerMetrics = new ControllerMetrics(metricsRegistry);
    segmentStatusChecker = new SegmentStatusChecker(helixResourceManager, config);
    segmentStatusChecker.setMetricsRegistry(controllerMetrics);
    segmentStatusChecker.runSegmentMetrics();
    Assert.assertEquals(
        controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.SEGMENTS_IN_ERROR_STATE),
        0);
    Assert.assertEquals(
        controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.NUMBER_OF_REPLICAS), 0);
    segmentStatusChecker.stop();
  }
 /*
  * For tenant creation
  */
 @Override
 @Post("json")
 public Representation post(Representation entity) {
   StringRepresentation presentation = null;
   try {
     PinotResourceManagerResponse response = null;
     final Tenant tenant = _objectMapper.readValue(entity.getText(), Tenant.class);
     switch (tenant.getTenantRole()) {
       case BROKER:
         response = _pinotHelixResourceMananger.createBrokerTenant(tenant);
         presentation = new StringRepresentation(response.toString());
         break;
       case SERVER:
         response = _pinotHelixResourceMananger.createServerTenant(tenant);
         presentation = new StringRepresentation(response.toString());
         break;
       default:
         throw new RuntimeException("Not a valid tenant creation call");
     }
   } catch (final Exception e) {
     presentation = exceptionToStringRepresentation(e);
     LOGGER.error("Caught exception while processing put request", e);
     setStatus(Status.SERVER_ERROR_INTERNAL);
   }
   return presentation;
 }
  @Test
  public void missingEVPartitionTest() throws Exception {
    final String tableName = "myTable";
    List<String> allTableNames = new ArrayList<String>();
    allTableNames.add(tableName);
    IdealState idealState = new IdealState(tableName);
    idealState.setPartitionState("myTable_0", "pinot1", "ONLINE");
    idealState.setPartitionState("myTable_0", "pinot2", "ONLINE");
    idealState.setPartitionState("myTable_0", "pinot3", "ONLINE");
    idealState.setPartitionState("myTable_1", "pinot1", "ONLINE");
    idealState.setPartitionState("myTable_1", "pinot2", "ONLINE");
    idealState.setPartitionState("myTable_1", "pinot3", "ONLINE");
    idealState.setPartitionState("myTable_2", "pinot3", "OFFLINE");
    idealState.setPartitionState("myTable_3", "pinot3", "ONLINE");
    idealState.setReplicas("2");
    idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);

    ExternalView externalView = new ExternalView(tableName);
    externalView.setState("myTable_0", "pinot1", "ONLINE");
    externalView.setState("myTable_0", "pinot2", "ONLINE");
    externalView.setState("myTable_1", "pinot1", "ERROR");
    externalView.setState("myTable_1", "pinot2", "ONLINE");

    HelixAdmin helixAdmin;
    {
      helixAdmin = mock(HelixAdmin.class);
      when(helixAdmin.getResourceIdealState("StatusChecker", "myTable")).thenReturn(idealState);
      when(helixAdmin.getResourceExternalView("StatusChecker", "myTable")).thenReturn(externalView);
    }
    {
      helixResourceManager = mock(PinotHelixResourceManager.class);
      when(helixResourceManager.isLeader()).thenReturn(true);
      when(helixResourceManager.getAllPinotTableNames()).thenReturn(allTableNames);
      when(helixResourceManager.getHelixClusterName()).thenReturn("StatusChecker");
      when(helixResourceManager.getHelixAdmin()).thenReturn(helixAdmin);
    }
    {
      config = mock(ControllerConf.class);
      when(config.getStatusControllerFrequencyInSeconds()).thenReturn(300);
    }
    metricsRegistry = new MetricsRegistry();
    controllerMetrics = new ControllerMetrics(metricsRegistry);
    segmentStatusChecker = new SegmentStatusChecker(helixResourceManager, config);
    segmentStatusChecker.setMetricsRegistry(controllerMetrics);
    segmentStatusChecker.runSegmentMetrics();
    Assert.assertEquals(
        controllerMetrics.getValueOfTableGauge(
            externalView.getId(), ControllerGauge.SEGMENTS_IN_ERROR_STATE),
        1);
    Assert.assertEquals(
        controllerMetrics.getValueOfTableGauge(
            externalView.getId(), ControllerGauge.NUMBER_OF_REPLICAS),
        0);
    segmentStatusChecker.stop();
  }
 @Override
 @Delete
 public Representation delete() {
   StringRepresentation presentation = null;
   try {
     final String tenantName = (String) getRequest().getAttributes().get(TENANT_NAME);
     final String type = getReference().getQueryAsForm().getValues("type");
     if (type == null) {
       presentation =
           new StringRepresentation(
               "Not specify the type for the tenant name. Please try to append:"
                   + "/?type=SERVER or /?type=BROKER ");
     } else {
       TenantRole tenantRole = TenantRole.valueOf(type.toUpperCase());
       PinotResourceManagerResponse res = null;
       switch (tenantRole) {
         case BROKER:
           if (_pinotHelixResourceMananger.isBrokerTenantDeletable(tenantName)) {
             res = _pinotHelixResourceMananger.deleteBrokerTenantFor(tenantName);
           } else {
             res = new PinotResourceManagerResponse();
             res.status = STATUS.failure;
             res.errorMessage = "Broker Tenant is not null, cannot delete it.";
           }
           break;
         case SERVER:
           if (_pinotHelixResourceMananger.isServerTenantDeletable(tenantName)) {
             res = _pinotHelixResourceMananger.deleteOfflineServerTenantFor(tenantName);
             if (res.isSuccessfull()) {
               res = _pinotHelixResourceMananger.deleteRealtimeServerTenantFor(tenantName);
             }
           } else {
             res = new PinotResourceManagerResponse();
             res.status = STATUS.failure;
             res.errorMessage = "Server Tenant is not null, cannot delete it.";
           }
           break;
         default:
           break;
       }
       presentation = new StringRepresentation(res.toString());
     }
   } catch (final Exception e) {
     presentation = exceptionToStringRepresentation(e);
     LOGGER.error("Caught exception while processing delete request", e);
     setStatus(Status.SERVER_ERROR_INTERNAL);
   }
   return presentation;
 }
  /**
   * Helper method to perform idempotent operation to refresh all watches (related to real-time
   * segments): - Data change listener for all existing real-time tables. - Child creation listener
   * for all existing real-time tables. - Data change listener for all existing real-time segments
   *
   * @param path
   */
  private void refreshWatchers(String path) {
    LOGGER.info("Received change notification for path: {}", path);
    List<Stat> stats = new ArrayList<>();
    List<ZNRecord> tableConfigs =
        _pinotHelixResourceManager.getPropertyStore().getChildren(TABLE_CONFIG, stats, 0);

    if (tableConfigs == null) {
      return;
    }

    for (ZNRecord tableConfig : tableConfigs) {
      try {
        AbstractTableConfig abstractTableConfig = AbstractTableConfig.fromZnRecord(tableConfig);
        if (abstractTableConfig.isRealTime()) {
          String realtimeTable = abstractTableConfig.getTableName();
          String realtimeSegmentsPathForTable =
              _propertyStorePath + SEGMENTS_PATH + "/" + realtimeTable;

          LOGGER.info("Setting data/child changes watch for real-time table '{}'", realtimeTable);
          _zkClient.subscribeDataChanges(realtimeSegmentsPathForTable, this);
          _zkClient.subscribeChildChanges(realtimeSegmentsPathForTable, this);

          List<String> childNames =
              _pinotHelixResourceManager
                  .getPropertyStore()
                  .getChildNames(SEGMENTS_PATH + "/" + realtimeTable, 0);

          if (childNames != null && !childNames.isEmpty()) {
            for (String segmentName : childNames) {
              String segmentPath = realtimeSegmentsPathForTable + "/" + segmentName;
              LOGGER.info("Setting data change watch for real-time segment: {}", segmentPath);
              _zkClient.subscribeDataChanges(segmentPath, this);
            }
          }
        }
      } catch (JSONException e) {
        LOGGER.error("Caught exception while reading table config", e);
      } catch (IOException e) {
        LOGGER.error(
            "Caught exception while setting change listeners for realtime tables/segments", e);
      }
    }
  }
  public void start() {
    LOGGER.info(
        "Starting realtime segments manager, adding a listener on the property store table configs path.");
    String zkUrl = _pinotHelixResourceManager.getHelixZkURL();
    _zkClient =
        new ZkClient(zkUrl, ZkClient.DEFAULT_SESSION_TIMEOUT, ZkClient.DEFAULT_CONNECTION_TIMEOUT);
    _zkClient.setZkSerializer(new ZNRecordSerializer());
    _zkClient.waitUntilConnected();

    // Subscribe to any data/child changes to property
    _zkClient.subscribeChildChanges(_tableConfigPath, this);
    _zkClient.subscribeDataChanges(_tableConfigPath, this);

    // Setup change listeners for already existing tables, if any.
    processPropertyStoreChange(_tableConfigPath);
  }
  /**
   * called with optional resourceName if resourceName is not present then it sends back a list of
   *
   * @return
   */
  @Override
  @Get
  public Representation get() {
    StringRepresentation presentation = null;
    try {
      final String tenantName = (String) getRequest().getAttributes().get(TENANT_NAME);
      if (tenantName == null) {
        // Return all the tags.
        final JSONObject ret = new JSONObject();
        final String type = getReference().getQueryAsForm().getValues("type");
        if (type == null || type.equals("server")) {
          ret.put("SERVER_TENANTS", _pinotHelixResourceMananger.getAllServerTenantNames());
        }
        if (type == null || type.equals("broker")) {
          ret.put("BROKER_TENANTS", _pinotHelixResourceMananger.getAllBrokerTenantNames());
        }
        presentation = new StringRepresentation(ret.toString(), MediaType.APPLICATION_JSON);
      } else {
        // Return instances related to given tenant name.
        final String type = getReference().getQueryAsForm().getValues("type");

        JSONObject resourceGetRet = new JSONObject();
        if (type == null) {
          resourceGetRet.put(
              "ServerInstances",
              _pinotHelixResourceMananger.getAllInstancesForServerTenant(tenantName));
          resourceGetRet.put(
              "BrokerInstances",
              _pinotHelixResourceMananger.getAllInstancesForBrokerTenant(tenantName));
        } else {
          if (type.equals("server")) {
            resourceGetRet.put(
                "ServerInstances",
                _pinotHelixResourceMananger.getAllInstancesForServerTenant(tenantName));
          }
          if (type.equals("broker")) {
            resourceGetRet.put(
                "BrokerInstances",
                _pinotHelixResourceMananger.getAllInstancesForBrokerTenant(tenantName));
          }
        }
        resourceGetRet.put(TENANT_NAME, tenantName);

        presentation =
            new StringRepresentation(resourceGetRet.toString(), MediaType.APPLICATION_JSON);
      }
    } catch (final Exception e) {
      presentation = exceptionToStringRepresentation(e);
      LOGGER.error("Caught exception while processing get request", e);
      setStatus(Status.SERVER_ERROR_INTERNAL);
    }
    return presentation;
  }
  private synchronized void assignRealtimeSegmentsToServerInstancesIfNecessary()
      throws JSONException, IOException {
    // Fetch current ideal state snapshot
    Map<String, IdealState> idealStateMap = new HashMap<String, IdealState>();

    for (String resource : _pinotHelixResourceManager.getAllRealtimeTables()) {
      idealStateMap.put(
          resource,
          _pinotHelixResourceManager
              .getHelixAdmin()
              .getResourceIdealState(_pinotHelixResourceManager.getHelixClusterName(), resource));
    }

    List<String> listOfSegmentsToAdd = new ArrayList<String>();

    for (String resource : idealStateMap.keySet()) {
      IdealState state = idealStateMap.get(resource);

      // Are there any partitions?
      if (state.getPartitionSet().size() == 0) {
        // No, this is a brand new ideal state, so we will add one new segment to every partition
        // and replica
        List<String> instancesInResource = new ArrayList<String>();
        try {
          instancesInResource.addAll(
              _pinotHelixResourceManager.getServerInstancesForTable(resource, TableType.REALTIME));
        } catch (Exception e) {
          LOGGER.error("Caught exception while fetching instances for resource {}", resource, e);
        }

        // Assign a new segment to all server instances
        for (String instanceId : instancesInResource) {
          InstanceZKMetadata instanceZKMetadata =
              _pinotHelixResourceManager.getInstanceZKMetadata(instanceId);
          String groupId = instanceZKMetadata.getGroupId(resource);
          String partitionId = instanceZKMetadata.getPartition(resource);
          listOfSegmentsToAdd.add(
              SegmentNameBuilder.Realtime.build(
                  resource,
                  instanceId,
                  groupId,
                  partitionId,
                  String.valueOf(System.currentTimeMillis())));
        }
      } else {
        // Add all server instances to the list of instances for which to assign a realtime segment
        Set<String> instancesToAssignRealtimeSegment = new HashSet<String>();
        instancesToAssignRealtimeSegment.addAll(
            _pinotHelixResourceManager.getServerInstancesForTable(resource, TableType.REALTIME));

        // Remove server instances that are currently processing a segment
        for (String partition : state.getPartitionSet()) {
          RealtimeSegmentZKMetadata realtimeSegmentZKMetadata =
              ZKMetadataProvider.getRealtimeSegmentZKMetadata(
                  _pinotHelixResourceManager.getPropertyStore(),
                  SegmentNameBuilder.Realtime.extractTableName(partition),
                  partition);
          if (realtimeSegmentZKMetadata.getStatus() == Status.IN_PROGRESS) {
            String instanceName = SegmentNameBuilder.Realtime.extractInstanceName(partition);
            instancesToAssignRealtimeSegment.remove(instanceName);
          }
        }

        // Assign a new segment to the server instances not currently processing this segment
        for (String instanceId : instancesToAssignRealtimeSegment) {
          InstanceZKMetadata instanceZKMetadata =
              _pinotHelixResourceManager.getInstanceZKMetadata(instanceId);
          String groupId = instanceZKMetadata.getGroupId(resource);
          String partitionId = instanceZKMetadata.getPartition(resource);
          listOfSegmentsToAdd.add(
              SegmentNameBuilder.Realtime.build(
                  resource,
                  instanceId,
                  groupId,
                  partitionId,
                  String.valueOf(System.currentTimeMillis())));
        }
      }
    }

    LOGGER.info(
        "Computed list of new segments to add : " + Arrays.toString(listOfSegmentsToAdd.toArray()));

    // Add the new segments to the server instances
    for (String segmentId : listOfSegmentsToAdd) {
      String resourceName = SegmentNameBuilder.Realtime.extractTableName(segmentId);
      String instanceName = SegmentNameBuilder.Realtime.extractInstanceName(segmentId);

      // Does the ideal state already contain this segment?
      if (!idealStateMap.get(resourceName).getPartitionSet().contains(segmentId)) {
        // No, add it
        // Create the realtime segment metadata
        RealtimeSegmentZKMetadata realtimeSegmentMetadataToAdd = new RealtimeSegmentZKMetadata();
        realtimeSegmentMetadataToAdd.setTableName(
            TableNameBuilder.extractRawTableName(resourceName));
        realtimeSegmentMetadataToAdd.setSegmentType(SegmentType.REALTIME);
        realtimeSegmentMetadataToAdd.setStatus(Status.IN_PROGRESS);
        realtimeSegmentMetadataToAdd.setSegmentName(segmentId);

        // Add the new metadata to the property store
        ZKMetadataProvider.setRealtimeSegmentZKMetadata(
            _pinotHelixResourceManager.getPropertyStore(), realtimeSegmentMetadataToAdd);

        // Update the ideal state to add the new realtime segment
        HelixHelper.updateIdealState(
            _pinotHelixResourceManager.getHelixZkManager(),
            resourceName,
            idealState ->
                PinotTableIdealStateBuilder.addNewRealtimeSegmentToIdealState(
                    segmentId, idealState, instanceName),
            RetryPolicies.exponentialBackoffRetryPolicy(5, 500L, 2.0f));
      }
    }
  }
 public void stop() {
   LOGGER.info("Stopping realtime segments manager, stopping property store.");
   _pinotHelixResourceManager.getPropertyStore().stop();
 }
 public PinotRealtimeSegmentManager(PinotHelixResourceManager pinotManager) {
   _pinotHelixResourceManager = pinotManager;
   String clusterName = _pinotHelixResourceManager.getHelixClusterName();
   _propertyStorePath = PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, clusterName);
   _tableConfigPath = _propertyStorePath + TABLE_CONFIG;
 }
 private boolean isLeader() {
   return _pinotHelixResourceManager.isLeader();
 }