com.datastax.driver.core.TokenRange Java Examples

@Override
public List<Shard> call() throws Exception {
    ArrayList<Shard> splits = new ArrayList<>();

    Map<TokenRange, Long> subSplits = getSubSplits(
            this.tokenRange,
            this.splitPartitions,
            this.splitSize);
    for (Map.Entry<TokenRange, Long> entry : subSplits.entrySet()) {
        List<TokenRange> ranges = entry.getKey().unwrap();
        for (TokenRange subrange : ranges) {
            String start = !isPartitionerOpp() ?
                           subrange.getStart().toString() :
                           subrange.getStart().toString().substring(2);
            String end = !isPartitionerOpp() ?
                         subrange.getEnd().toString() :
                         subrange.getEnd().toString().substring(2);
            long length = entry.getValue();
            splits.add(new Shard(start, end, length));
        }
    }
    return splits;
}
 

private Observable.Transformer<BoundStatement, Integer> applyMicroBatching() {
    return tObservable -> tObservable
            .groupBy(b -> {
                ByteBuffer routingKey = b.getRoutingKey(ProtocolVersion.NEWEST_SUPPORTED,
                        codecRegistry);
                Token token = metadata.newToken(routingKey);
                for (TokenRange tokenRange : session.getCluster().getMetadata().getTokenRanges()) {
                    if (tokenRange.contains(token)) {
                        return tokenRange;
                    }
                }
                log.warn("Unable to find any Cassandra node to insert token " + token.toString());
                return session.getCluster().getMetadata().getTokenRanges().iterator().next();
            })
            .flatMap(g -> g.compose(new BoundBatchStatementTransformer()))
            .flatMap(batch -> rxSession
                    .execute(batch)
                    .compose(applyInsertRetryPolicy())
                    .map(resultSet -> batch.size())
            );
}
 

private Map<String,String> getMappingTable() {
    // You only want to prepare the query once. Best to do it when you initialize the session.
    PreparedStatement findMappings = session.prepare(
            "SELECT project_id, project_name " +
                    "FROM openshift_metrics.metrics_mappings " +
                    "WHERE token(project_id) > ? AND token(project_id) <= ?");

    Map<String,String> mappings = new HashMap<>();
    if (hasMappingTable()) {
        for (TokenRange tokenRange : getTokenRanges()) {
            BoundStatement boundStatement = findMappings.bind().setToken(0, tokenRange.getStart())
                    .setToken(1, tokenRange.getEnd());
            ResultSet resultSet = session.execute(boundStatement);
            resultSet.forEach(row -> mappings.put(row.getString(0), row.getString(1)));
        }
    }

    mappings.remove("%succeeded");
    return mappings;
}
 

public SplitCallable(TokenRange tokenRange,
                     long splitPartitions, long splitSize) {
    if (splitSize <= 0 && splitPartitions <= 0) {
        throw new IllegalArgumentException(String.format(
                  "The split-partitions must be > 0, but got %s",
                  splitPartitions));
    }

    if (splitSize > 0 && splitSize < MIN_SHARD_SIZE) {
        // splitSize should be at least 1M if passed
        throw new IllegalArgumentException(String.format(
                  "The split-size must be >= %s bytes, but got %s",
                  MIN_SHARD_SIZE, splitSize));
    }

    this.tokenRange = tokenRange;
    this.splitPartitions = splitPartitions;
    this.splitSize = splitSize;
}
 

public Set<Host> getReplicas(String caseSensitiveSchemaName, TokenRange tokenRange)
{
    requireNonNull(caseSensitiveSchemaName, "keyspace is null");
    requireNonNull(tokenRange, "tokenRange is null");
    return executeWithSession(session ->
            session.getCluster().getMetadata().getReplicas(validSchemaName(caseSensitiveSchemaName), tokenRange));
}
 

private Set<TokenRange> unwrap(Set<TokenRange> tokenRanges)
{
    ImmutableSet.Builder<TokenRange> result = ImmutableSet.builder();
    for (TokenRange range : tokenRanges) {
        result.addAll(range.unwrap());
    }
    return result.build();
}
 

private Set<TokenRange> getTokenRanges() {
    Set<TokenRange> tokenRanges = new HashSet<>();
    for (TokenRange tokenRange : metadata.getTokenRanges()) {
        tokenRanges.addAll(tokenRange.unwrap());
    }
    return tokenRanges;
}
 

private Set<TokenRange> getTokenRanges() {
    Set<TokenRange> tokenRanges = new HashSet<>();
    for (TokenRange tokenRange : session.getCluster().getMetadata().getTokenRanges()) {
        tokenRanges.addAll(tokenRange.unwrap());
    }
    return tokenRanges;
}
 

private List<String> getEndpoints(String keyspace, TokenRange tokenRange)
{
    Set<Host> endpoints = session.getReplicas(keyspace, tokenRange);
    return unmodifiableList(endpoints.stream()
            .map(Host::toString)
            .collect(toList()));
}
 

private List<BoundStatement> rangeQuery(PreparedStatement rangeStmt, TokenRange range) {
  List<BoundStatement> res = Lists.newArrayList();
  for (TokenRange subRange : range.unwrap()) {
    res.add(rangeStmt.bind(subRange.getStart(), subRange.getEnd()));
  }
  return res;
}
 

private List<TokenRange> buildTokenRanges() {
    List<TokenRange> result = new ArrayList<>();
    Metadata metadata = session.getCluster().getMetadata();
    for (TokenRange range : metadata.getTokenRanges()) {
        for (TokenRange split : range.splitEvenly(split)) {
            result.addAll(split.unwrap());
        }
    }
    logger.info("Configured with {} token ranges, and {} splits", metadata.getTokenRanges(), result.size());
    return result;
}
 

private static TokenSplit createSplit(TokenRange range, List<String> endpoints)
{
    checkArgument(!range.isEmpty(), "tokenRange must not be empty");
    String startToken = range.getStart().toString();
    String endToken = range.getEnd().toString();
    return new TokenSplit(startToken, endToken, endpoints);
}
 

private Map<TokenRange, Long> getSubSplits(TokenRange tokenRange,
                                           long splitPartitions,
                                           long splitSize) {
    try {
        return describeSplits(this.session, this.keyspace, this.table,
                              splitPartitions, splitSize, tokenRange);
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}
 

private TokenRange rangeToTokenRange(Range<Token> range) {
    TokenFactory tokenFactory = this.partitioner.getTokenFactory();
    Metadata metadata = this.session.metadata();
    return metadata.newTokenRange(
                    metadata.newToken(tokenFactory.toString(range.left)),
                    metadata.newToken(tokenFactory.toString(range.right)));
}
 

public List<TokenSplit> getSplits(String keyspace, String table, Optional<Long> sessionSplitsPerNode)
{
    Set<TokenRange> tokenRanges = session.getTokenRanges();

    if (tokenRanges.isEmpty()) {
        throw new PrestoException(CASSANDRA_METADATA_ERROR, "The cluster metadata is not available. " +
                "Please make sure that the Cassandra cluster is up and running, " +
                "and that the contact points are specified correctly.");
    }

    if (tokenRanges.stream().anyMatch(TokenRange::isWrappedAround)) {
        tokenRanges = unwrap(tokenRanges);
    }

    Optional<TokenRing> tokenRing = createForPartitioner(session.getPartitioner());
    long totalPartitionsCount = getTotalPartitionsCount(keyspace, table, sessionSplitsPerNode);

    List<TokenSplit> splits = new ArrayList<>();
    for (TokenRange tokenRange : tokenRanges) {
        if (tokenRange.isEmpty()) {
            continue;
        }
        checkState(!tokenRange.isWrappedAround(), "all token ranges must be unwrapped at this step");

        List<String> endpoints = getEndpoints(keyspace, tokenRange);
        checkState(!endpoints.isEmpty(), "endpoints is empty for token range: %s", tokenRange);

        if (tokenRing.isEmpty()) {
            checkState(!tokenRange.isWrappedAround(), "all token ranges must be unwrapped at this step");
            splits.add(createSplit(tokenRange, endpoints));
            continue;
        }

        double tokenRangeRingFraction = tokenRing.get().getRingFraction(tokenRange.getStart().toString(), tokenRange.getEnd().toString());
        long partitionsCountEstimate = round(totalPartitionsCount * tokenRangeRingFraction);
        checkState(partitionsCountEstimate >= 0, "unexpected partitions count estimate: %s", partitionsCountEstimate);
        int subSplitCount = max(toIntExact(partitionsCountEstimate / splitSize), 1);
        List<TokenRange> subRanges = tokenRange.splitEvenly(subSplitCount);

        for (TokenRange subRange : subRanges) {
            if (subRange.isEmpty()) {
                continue;
            }
            checkState(!subRange.isWrappedAround(), "all token ranges must be unwrapped at this step");
            splits.add(createSplit(subRange, endpoints));
        }
    }
    shuffle(splits, ThreadLocalRandom.current());
    return unmodifiableList(splits);
}
 

public List<TokenRange> getTokenRanges() {
    return meta.getRanges();
}
 

public List<TokenRange> getRanges() {
    return ranges;
}
 

private void loadMetricIdCache(ExecutorService executor) {
  final AtomicInteger tasks = new AtomicInteger(0);
  logger.info("Found token ranges: " + cluster.getMetadata().getTokenRanges().size());
  for (TokenRange range : cluster.getMetadata().getTokenRanges()) {
    List<BoundStatement> queries = rangeQuery(retrieveMetricIdStmt, range);
    for (BoundStatement query : queries) {
      tasks.incrementAndGet();
      logger.info("adding a metric id reading task, total: " + tasks.get());

      ResultSetFuture future = metricsSession.executeAsync(query);

      Futures.addCallback(future, new FutureCallback<ResultSet>() {
        @Override
        public void onSuccess(ResultSet result) {
          for (Row row : result) {
            String id = Bytes.toHexString(row.getBytes(METRIC_ID));
            if (id != null) {
              //remove '0x'
              metricIdCache.put(id.substring(2), Boolean.TRUE);
            }
          }

          tasks.decrementAndGet();

          logger.info("completed a metric id read task. Remaining tasks: " + tasks.get());
        }

        @Override
        public void onFailure(Throwable t) {
          logger.error("Failed to execute query to load metric id cache.", t);

          tasks.decrementAndGet();

          logger.info("Failed a metric id read task. Remaining tasks: " + tasks.get());
        }
      }, executor);

    }
  }

  while (tasks.get() > 0) {
    logger.debug("waiting for more metric id load tasks: " + tasks.get());

    try {
      Thread.sleep(3000);
    } catch (InterruptedException e) {
      logger.warn("load metric cache was interrupted", e);
    }
  }

  logger.info("loaded metric id cache from database: " + metricIdCache.size());
}
 

private static Map<TokenRange, Long> describeSplits(
        CassandraSessionPool.Session session,
        String keyspace,
        String table,
        long splitPartitions,
        long splitSize,
        TokenRange tokenRange) {

    String query = String.format(
            "SELECT mean_partition_size, partitions_count FROM %s.%s " +
            "WHERE keyspace_name = ? AND table_name = ? AND " +
            "range_start = ? AND range_end = ?",
            SchemaConstants.SYSTEM_KEYSPACE_NAME,
            SystemKeyspace.SIZE_ESTIMATES);

    ResultSet resultSet = session.execute(query, keyspace, table,
                                          tokenRange.getStart().toString(),
                                          tokenRange.getEnd().toString());
    Row row = resultSet.one();

    long meanPartitionSize = 0L;
    long partitionsCount = 0L;
    long splitCount = 0L;

    if (row != null) {
        meanPartitionSize = row.getLong("mean_partition_size");
        partitionsCount = row.getLong("partitions_count");
        assert splitSize <= 0 || splitSize >= MIN_SHARD_SIZE;
        splitCount = splitSize > 0 ?
                     (meanPartitionSize * partitionsCount / splitSize) :
                     (partitionsCount / splitPartitions);
    }

    /*
     * If we have no data on this split or the size estimate is 0,
     * return the full split i.e., do not sub-split
     * Assume smallest granularity of partition count available from
     * CASSANDRA-7688.
     */
    if (splitCount == 0) {
        return ImmutableMap.of(tokenRange, (long) 128);
    }

    List<TokenRange> ranges = tokenRange.splitEvenly((int) splitCount);
    Map<TokenRange, Long> rangesWithLength = new HashMap<>();
    for (TokenRange range : ranges) {
        // Add a sub-range (with its partitions count per sub-range)
        rangesWithLength.put(range, partitionsCount / splitCount);
    }
    return rangesWithLength;
}
 

private void loadMetricDimensionCache(ExecutorService executor) {

    final AtomicInteger tasks = new AtomicInteger(0);

    for (TokenRange range : cluster.getMetadata().getTokenRanges()) {
      List<BoundStatement> queries = rangeQuery(retrieveMetricDimensionStmt, range);
      for (BoundStatement query : queries) {
        tasks.incrementAndGet();

        logger.info("Adding a metric dimnesion read task, total: " + tasks.get());

        ResultSetFuture future = metricsSession.executeAsync(query);

        Futures.addCallback(future, new FutureCallback<ResultSet>() {
          @Override
          public void onSuccess(ResultSet result) {
            for (Row row : result) {
              String key = getMetricDimnesionEntryKey(row.getString(REGION),
                  row.getString(TENANT_ID_COLUMN), row.getString(METRIC_NAME),
                  row.getString(DIMENSION_NAME), row.getString(DIMENSION_VALUE));
              metricDimensionCache.put(key, Boolean.TRUE);
            }

            tasks.decrementAndGet();

            logger.info("Completed a metric dimension read task. Remaining tasks: " + tasks.get());
          }

          @Override
          public void onFailure(Throwable t) {
            logger.error("Failed to execute query to load metric id cache.", t);

            tasks.decrementAndGet();

            logger.info("Failed a metric dimension read task. Remaining tasks: " + tasks.get());
          }
        }, executor);

      }
    }

    while (tasks.get() > 0) {

      logger.debug("waiting for metric dimension cache to load ...");

      try {
        Thread.sleep(1000);
      } catch (InterruptedException e) {
        logger.warn("load metric dimension cache was interrupted", e);
      }
    }

    logger.info("loaded metric dimension cache from database: " + metricDimensionCache.size());
  }
 

private Map<TokenRange, Set<Host>> getRangeMap() {
    Metadata metadata = this.session.metadata();
    return metadata.getTokenRanges().stream().collect(Collectors.toMap(
            p -> p,
            p -> metadata.getReplicas('"' + this.keyspace + '"', p)));
}
 

public Set<TokenRange> getTokenRanges()
{
    return executeWithSession(session -> session.getCluster().getMetadata().getTokenRanges());
}