org.apache.cassandra.streaming.StreamState Java Examples

private Set<StreamState> parseCompositeData(Set<CompositeData> payload) {
  Set<StreamState> result = Sets.newHashSet();

  for (CompositeData compositeData : payload) {

    try {
      // start by trying to parse with classes coming from Reaper's C* dependency
      StreamState streamState = StreamStateCompositeData.fromCompositeData(compositeData);
      result.add(streamState);
    } catch (AssertionError | InvalidKeyException e) {
      try {
        // if that fails, try the old version
        StreamState olderStreamState = parseStreamStatePre2_1(compositeData);
        result.add(olderStreamState);
      } catch (InvalidKeyException ie) {
        // and if even that fails, try the new version
        // if that still fails, exception goes up
        StreamState newerStreamState = parseStreamState4_0_0(compositeData);
        result.add(newerStreamState);
      }
    }
  }
  return result;
}
 

public static StreamState fromCompositeData(CompositeData cd)
{
    assert cd.getCompositeType().equals(COMPOSITE_TYPE);
    Object[] values = cd.getAll(ITEM_NAMES);
    UUID planId = UUID.fromString((String) values[0]);
    String description = (String) values[1];
    Set<SessionInfo> sessions = Sets.newHashSet(Iterables.transform(Arrays.asList((CompositeData[]) values[2]),
                                                                    new Function<CompositeData, SessionInfo>()
                                                                    {
                                                                        public SessionInfo apply(CompositeData input)
                                                                        {
                                                                            return SessionInfoCompositeData.fromCompositeData(input);
                                                                        }
                                                                    }));
    return new StreamState(planId, description, sessions);
}
 

private StreamState parseStreamStatePre2_1(CompositeData compositeData) {
  UUID planId = UUID.fromString((String)compositeData.get("planId"));
  String description = (String) compositeData.get("description");

  CompositeData[] sessionCompositeData = (CompositeData[]) compositeData.get("sessions");

  Set<SessionInfo> sessions = Arrays.stream(sessionCompositeData)
      .map(this::parseSessionInfoPre2_1)
      .collect(Collectors.toSet());

  return new StreamState(planId, description, sessions);
}
 

private StreamState parseStreamState4_0_0(CompositeData compositeData) {
  UUID planId = UUID.fromString((String)compositeData.get("planId"));
  String description = (String) compositeData.get("description");

  CompositeData[] sessionCompositeData = (CompositeData[]) compositeData.get("sessions");

  Set<SessionInfo> sessions = Arrays.stream(sessionCompositeData)
      .map(this::parseSessionInfo4_0_0)
      .collect(Collectors.toSet());

  return new StreamState(planId, description, sessions);
}
 

/**
 * This is called when stream info is explicitly pulled from a node.
 *
 * @param host hostname of the node the streams are pulled from.
 * @param streamState parsed payload received from Cassandra.
 * @return stream session.
 */
public static StreamSession fromStreamState(String host, StreamState streamState) {

  UUID planId = streamState.planId;

  Map<String, Stream> streams = streamState.sessions.stream()
      .map(sessionInfo -> StreamFactory.newStream(host, sessionInfo, System.currentTimeMillis()))
      .collect(Collectors.toMap(Stream::getId, stream -> stream));

  return StreamSession.builder()
      .withPlanId(planId.toString())
      .withStreams(streams)
      .build();
}
 

@Test
public void testCountProgressPerTableWithTable() throws UnknownHostException {
  UUID cfid = UUID.randomUUID();
  int files = 3;
  int totalSize = 2048;
  StreamSummary streamSummary = new StreamSummary(cfid, files, totalSize);

  InetAddress peer = InetAddress.getByName("127.0.0.1");
  int index = 0;
  InetAddress connecting = InetAddress.getByName("127.0.0.2");
  ImmutableSet<StreamSummary> receivingSummaries = ImmutableSet.of(streamSummary);
  ImmutableSet<StreamSummary> sendingSummaries = ImmutableSet.of();

  SessionInfo sessionInfo
      = new SessionInfo(peer, index, connecting, receivingSummaries, sendingSummaries, StreamSession.State.STREAMING);

  String file1 = "/cass/data/keyspace1/standard1-af5311f0633a11e89d71710c22f847e7/lb-4-big-Data.db";
  String file2 = "/cass/data/keyspace1/standard1-af5311f0633a11e89d71710c22f847e7/lb-5-big-Data.db";
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file1, ProgressInfo.Direction.OUT, 512, 1024));
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file2, ProgressInfo.Direction.OUT, 512, 1024));

  UUID planId = UUID.randomUUID();
  ImmutableSet<SessionInfo> sessionInfos = ImmutableSet.of(sessionInfo);
  StreamState streamState = new StreamState(planId, "descr", sessionInfos);

  io.cassandrareaper.core.StreamSession streamSession
      = StreamSessionFactory.fromStreamState(peer.toString(), streamState);

  assertEquals(1, streamSession.getStreams().size());

  List<Stream.TableProgress> progressSent = streamSession.getStreams().values().asList().get(0).getProgressSent();
  assertEquals(1, progressSent.size());

  Stream.TableProgress tableProgress = progressSent.get(0);
  assertEquals("keyspace1.standard1", tableProgress.getTable());
  assertEquals(Long.valueOf(1024), tableProgress.getCurrent());
  assertEquals(Long.valueOf(2048), tableProgress.getTotal());
}
 

@Test
public void testCountProgressPerTableWhenReceiving() throws UnknownHostException {
  UUID cfid = UUID.randomUUID();
  int files = 1;
  int totalSize = 1024;
  StreamSummary streamSummary = new StreamSummary(cfid, files, totalSize);

  InetAddress peer = InetAddress.getByName("127.0.0.1");
  int index = 0;
  InetAddress connecting = InetAddress.getByName("127.0.0.2");
  ImmutableSet<StreamSummary> receivingSummaries = ImmutableSet.of(streamSummary);
  ImmutableSet<StreamSummary> sendingSummaries = ImmutableSet.of();

  SessionInfo sessionInfo
      = new SessionInfo(peer, index, connecting, receivingSummaries, sendingSummaries, StreamSession.State.STREAMING);

  // this is the important part - when receiving, the absolute path is not known
  String file1 = "keyspace1/standard1";
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file1, ProgressInfo.Direction.IN, 512, 1024));

  UUID planId = UUID.randomUUID();
  ImmutableSet<SessionInfo> sessionInfos = ImmutableSet.of(sessionInfo);
  StreamState streamState = new StreamState(planId, "descr", sessionInfos);

  io.cassandrareaper.core.StreamSession streamSession
      = StreamSessionFactory.fromStreamState(peer.toString(), streamState);

  assertEquals(1, streamSession.getStreams().size());

  List<Stream.TableProgress> progressSent = streamSession.getStreams().values().asList().get(0).getProgressReceived();
  assertEquals(1, progressSent.size());

  Stream.TableProgress tableProgress = progressSent.get(0);
  assertEquals("keyspace1.standard1", tableProgress.getTable());
  assertEquals(Long.valueOf(512), tableProgress.getCurrent());
  assertEquals(Long.valueOf(1024), tableProgress.getTotal());
}
 

private void prepareWriter() throws IOException
{
    try
    {
        if (writer == null)
        {
            writer = CQLSSTableWriter.builder()
                .forTable(schema)
                .using(insertStatement)
                .withPartitioner(ConfigHelper.getOutputPartitioner(conf))
                .inDirectory(outputDir)
                .withBufferSizeInMB(Integer.parseInt(conf.get(BUFFER_SIZE_IN_MB, "64")))
                .build();
        }
        if (loader == null)
        {
            ExternalClient externalClient = new ExternalClient(conf);
            
            externalClient.addKnownCfs(keyspace, schema);

            this.loader = new SSTableLoader(outputDir, externalClient, new BulkRecordWriter.NullOutputHandler()) {
                @Override
                public void onSuccess(StreamState finalState)
                {
                    if (deleteSrc)
                        FileUtils.deleteRecursive(outputDir);
                }
            };
        }
    }
    catch (Exception e)
    {
        throw new IOException(e);
    }      
}
 

private void close() throws IOException
{
    if (writer != null)
    {
        writer.close();
        Future<StreamState> future = loader.stream();
        while (true)
        {
            try
            {
                future.get(1000, TimeUnit.MILLISECONDS);
                break;
            }
            catch (ExecutionException | TimeoutException te)
            {
                if (null != progress)
                    progress.progress();
                if (null != context)
                    HadoopCompat.progress(context);
            }
            catch (InterruptedException e)
            {
                throw new IOException(e);
            }
        }
        if (loader.getFailedHosts().size() > 0)
        {
            if (loader.getFailedHosts().size() > maxFailures)
                throw new IOException("Too many hosts failed: " + loader.getFailedHosts());
            else
                logger.warn("Some hosts failed: {}", loader.getFailedHosts());
        }
    }
}
 

private Future<StreamState> streamHints()
{
    // StreamPlan will not fail if there are zero files to transfer, so flush anyway (need to get any in-memory hints, as well)
    ColumnFamilyStore hintsCF = Keyspace.open(Keyspace.SYSTEM_KS).getColumnFamilyStore(SystemKeyspace.HINTS_CF);
    FBUtilities.waitOnFuture(hintsCF.forceFlush());

    // gather all live nodes in the cluster that aren't also leaving
    List<InetAddress> candidates = new ArrayList<>(StorageService.instance.getTokenMetadata().cloneAfterAllLeft().getAllEndpoints());
    candidates.remove(FBUtilities.getBroadcastAddress());
    for (Iterator<InetAddress> iter = candidates.iterator(); iter.hasNext(); )
    {
        InetAddress address = iter.next();
        if (!FailureDetector.instance.isAlive(address))
            iter.remove();
    }

    if (candidates.isEmpty())
    {
        logger.warn("Unable to stream hints since no live endpoints seen");
        return Futures.immediateFuture(null);
    }
    else
    {
        // stream to the closest peer as chosen by the snitch
        DatabaseDescriptor.getEndpointSnitch().sortByProximity(FBUtilities.getBroadcastAddress(), candidates);
        InetAddress hintsDestinationHost = candidates.get(0);
        InetAddress preferred = SystemKeyspace.getPreferredIP(hintsDestinationHost);

        // stream all hints -- range list will be a singleton of "the entire ring"
        Token token = StorageService.getPartitioner().getMinimumToken();
        List<Range<Token>> ranges = Collections.singletonList(new Range<>(token, token));

        return new StreamPlan("Hints").transferRanges(hintsDestinationHost,
                                                      preferred,
                                                                  Keyspace.SYSTEM_KS,
                                                                  ranges,
                                                                  SystemKeyspace.HINTS_CF)
                                                  .execute();
    }
}
 

public Set<StreamState> getStreamStatus()
{
    return Sets.newHashSet(Iterables.transform(streamProxy.getCurrentStreams(), new Function<CompositeData, StreamState>()
    {
        public StreamState apply(CompositeData input)
        {
            return StreamStateCompositeData.fromCompositeData(input);
        }
    }));
}
 

private void close() throws IOException {
  LOG.info("SSTables built. Now starting streaming");
  heartbeat.startHeartbeat();
  try {
    if (writer != null) {
      writer.close();
      Future<StreamState> future =
          loader.stream(Collections.<InetAddress>emptySet(), new ProgressIndicator());
      try {
        StreamState streamState = Uninterruptibles.getUninterruptibly(future);
        if (streamState.hasFailedSession()) {
          LOG.warn("Some streaming sessions failed");
        } else {
          LOG.info("Streaming finished successfully");
        }
      } catch (ExecutionException e) {
        throw new RuntimeException("Streaming to the following hosts failed: " +
            loader.getFailedHosts(), e);
      }
    } else {
      LOG.info("SSTableWriter wasn't instantiated, no streaming happened.");
    }
  } finally {
    heartbeat.stopHeartbeat();
  }
  LOG.info("Successfully closed bulk record writer");
}
 

@Test
public void testCountProgressPerTableWithMultipleTables() throws UnknownHostException {
  UUID cfid = UUID.randomUUID();
  int files = 3;
  int totalSize = 2048;
  StreamSummary streamSummary = new StreamSummary(cfid, files, totalSize);

  InetAddress peer = InetAddress.getByName("127.0.0.1");
  int index = 0;
  InetAddress connecting = InetAddress.getByName("127.0.0.2");
  ImmutableSet<StreamSummary> receivingSummaries = ImmutableSet.of(streamSummary);
  ImmutableSet<StreamSummary> sendingSummaries = ImmutableSet.of();

  SessionInfo sessionInfo
      = new SessionInfo(peer, index, connecting, receivingSummaries, sendingSummaries, StreamSession.State.STREAMING);

  String file1 = "/cass/data/keyspace1/standard1-af5311f0633a11e89d71710c22f847e7/lb-4-big-Data.db";
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file1, ProgressInfo.Direction.OUT, 32, 1024));
  String file2 = "/cass/data/keyspace1/standard1-af5311f0633a11e89d71710c22f847e7/lb-5-big-Data.db";
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file2, ProgressInfo.Direction.OUT, 64, 1024));
  String file3 = "/cass/data/keyspace1/counter1-af5311f0633a11e89d71710c22f847e7/lb-4-big-Data.db";
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file3, ProgressInfo.Direction.OUT, 512, 512));
  String file4 = "/cass/data/keyspace1/counter1-af5311f0633a11e89d71710c22f847e7/lb-5-big-Data.db";
  sessionInfo.updateProgress(new ProgressInfo(peer, index, file4, ProgressInfo.Direction.OUT, 128, 512));

  UUID planId = UUID.randomUUID();
  ImmutableSet<SessionInfo> sessionInfos = ImmutableSet.of(sessionInfo);
  StreamState streamState = new StreamState(planId, "descr", sessionInfos);

  io.cassandrareaper.core.StreamSession streamSession
      = StreamSessionFactory.fromStreamState(peer.toString(), streamState);

  assertEquals(1, streamSession.getStreams().size());

  List<Stream.TableProgress> progressSent = streamSession.getStreams().values().asList().get(0).getProgressSent();
  assertEquals(2, progressSent.size());

  Stream.TableProgress standardTableProgess = progressSent.stream()
      .filter(ps -> ps.getTable().equals("keyspace1.standard1"))
      .findFirst()
      .get();
  assertEquals(Long.valueOf(96), standardTableProgess.getCurrent());
  assertEquals(Long.valueOf(2048), standardTableProgess.getTotal());

  Stream.TableProgress counterTableProgress = progressSent.stream()
      .filter(ps -> ps.getTable().equals("keyspace1.counter1"))
      .findFirst()
      .get();
  assertEquals(Long.valueOf(640), counterTableProgress.getCurrent());
  assertEquals(Long.valueOf(1024), counterTableProgress.getTotal());
}
 

public Future<StreamState> stream()
{
    return streamPlan.execute();
}
 

public static CompositeData toCompositeData(final StreamState streamState)
{
    Map<String, Object> valueMap = new HashMap<>();
    valueMap.put(ITEM_NAMES[0], streamState.planId.toString());
    valueMap.put(ITEM_NAMES[1], streamState.description);

    CompositeData[] sessions = new CompositeData[streamState.sessions.size()];
    Lists.newArrayList(Iterables.transform(streamState.sessions, new Function<SessionInfo, CompositeData>()
    {
        public CompositeData apply(SessionInfo input)
        {
            return SessionInfoCompositeData.toCompositeData(streamState.planId, input);
        }
    })).toArray(sessions);
    valueMap.put(ITEM_NAMES[2], sessions);

    long currentRxBytes = 0;
    long totalRxBytes = 0;
    long currentTxBytes = 0;
    long totalTxBytes = 0;
    for (SessionInfo sessInfo : streamState.sessions)
    {
        currentRxBytes += sessInfo.getTotalSizeReceived();
        totalRxBytes += sessInfo.getTotalSizeToReceive();
        currentTxBytes += sessInfo.getTotalSizeSent();
        totalTxBytes += sessInfo.getTotalSizeToSend();
    }
    double rxPercentage = (totalRxBytes == 0 ? 100L : currentRxBytes * 100L / totalRxBytes);
    double txPercentage = (totalTxBytes == 0 ? 100L : currentTxBytes * 100L / totalTxBytes);

    valueMap.put(ITEM_NAMES[3], currentRxBytes);
    valueMap.put(ITEM_NAMES[4], totalRxBytes);
    valueMap.put(ITEM_NAMES[5], rxPercentage);
    valueMap.put(ITEM_NAMES[6], currentTxBytes);
    valueMap.put(ITEM_NAMES[7], totalTxBytes);
    valueMap.put(ITEM_NAMES[8], txPercentage);

    try
    {
        return new CompositeDataSupport(COMPOSITE_TYPE, valueMap);
    }
    catch (OpenDataException e)
    {
        throw Throwables.propagate(e);
    }
}
 

@Override
public void onSuccess(StreamState finalState) {
}