Java Code Examples for org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker#currentRestriction()

@ProcessElement
public void processElement(
    @Element Read read,
    OutputReceiver<Result> out,
    RestrictionTracker<ByteKeyRange, ByteKey> tracker)
    throws Exception {
  Connection connection = ConnectionFactory.createConnection(read.getConfiguration());
  TableName tableName = TableName.valueOf(read.getTableId());
  Table table = connection.getTable(tableName);
  final ByteKeyRange range = tracker.currentRestriction();
  try (ResultScanner scanner =
      table.getScanner(HBaseUtils.newScanInRange(read.getScan(), range))) {
    for (Result result : scanner) {
      ByteKey key = ByteKey.copyFrom(result.getRow());
      if (!tracker.tryClaim(key)) {
        return;
      }
      out.output(result);
    }
    tracker.tryClaim(ByteKey.EMPTY);
  }
}
 

@ProcessElement
public ProcessContinuation processElement(
    @Element SequenceDefinition srcElement,
    OutputReceiver<Instant> out,
    RestrictionTracker<OffsetRange, Long> restrictionTracker) {

  OffsetRange restriction = restrictionTracker.currentRestriction();
  Long interval = srcElement.durationMilliSec;
  Long nextOutput = restriction.getFrom() + interval;

  boolean claimSuccess = true;

  while (claimSuccess && Instant.ofEpochMilli(nextOutput).isBeforeNow()) {
    claimSuccess = restrictionTracker.tryClaim(nextOutput);
    if (claimSuccess) {
      Instant output = Instant.ofEpochMilli(nextOutput);
      out.outputWithTimestamp(output, output);
      nextOutput = nextOutput + interval;
    }
  }

  ProcessContinuation continuation = ProcessContinuation.stop();
  if (claimSuccess) {
    Duration offset = new Duration(Instant.now(), Instant.ofEpochMilli(nextOutput));
    continuation = ProcessContinuation.resume().withResumeDelay(offset);
  }
  return continuation;
}
 

/**
 * List messages.
 *
 * @param hl7v2Store the HL7v2 store to list messages from
 * @throws IOException the io exception
 */
@ProcessElement
public void listMessages(
    @Element String hl7v2Store,
    RestrictionTracker<OffsetRange, Long> tracker,
    OutputReceiver<HL7v2Message> outputReceiver)
    throws IOException {
  OffsetRange currentRestriction = (OffsetRange) tracker.currentRestriction();
  Instant startRestriction = Instant.ofEpochMilli(currentRestriction.getFrom());
  Instant endRestriction = Instant.ofEpochMilli(currentRestriction.getTo());
  HttpHealthcareApiClient.HL7v2MessagePages pages =
      new HttpHealthcareApiClient.HL7v2MessagePages(
          client, hl7v2Store, startRestriction, endRestriction, filter.get(), "sendTime");
  Instant cursor;
  long lastClaimedMilliSecond = startRestriction.getMillis() - 1;
  for (HL7v2Message msg : FluentIterable.concat(pages)) {
    cursor = Instant.parse(msg.getSendTime());
    if (cursor.getMillis() > lastClaimedMilliSecond) {
      // Return early after the first claim failure preventing us from iterating
      // through the remaining messages.
      if (!tracker.tryClaim(cursor.getMillis())) {
        return;
      }
      lastClaimedMilliSecond = cursor.getMillis();
    }

    outputReceiver.output(msg);
  }

  // We've paginated through all messages for this restriction but the last message may be
  // before the end of the restriction
  tracker.tryClaim(currentRestriction.getTo());
}
 

@ProcessElement
public ProcessContinuation process(
    ProcessContext c,
    RestrictionTracker<GrowthState, KV<Growth.PollResult<OutputT>, TerminationStateT>> tracker,
    ManualWatermarkEstimator<Instant> watermarkEstimator)
    throws Exception {

  GrowthState currentRestriction = tracker.currentRestriction();
  if (currentRestriction instanceof NonPollingGrowthState) {
    Growth.PollResult<OutputT> priorPoll =
        ((NonPollingGrowthState<OutputT>) currentRestriction).getPending();
    if (tracker.tryClaim(KV.of(priorPoll, null))) {
      if (!priorPoll.getOutputs().isEmpty()) {
        LOG.info(
            "{} - re-emitting output of prior poll containing {} results.",
            c.element(),
            priorPoll.getOutputs().size());
        c.output(KV.of(c.element(), priorPoll.getOutputs()));
      }
      watermarkEstimator.setWatermark(priorPoll.getWatermark());
    }
    return stop();
  }

  // Poll for additional elements.
  Instant now = Instant.now();
  Growth.PollResult<OutputT> res =
      spec.getPollFn().getClosure().apply(c.element(), wrapProcessContext(c));

  PollingGrowthState<TerminationStateT> pollingRestriction =
      (PollingGrowthState<TerminationStateT>) currentRestriction;
  // Produce a poll result that only contains never seen before results.
  Growth.PollResult<OutputT> newResults =
      computeNeverSeenBeforeResults(pollingRestriction, res);

  // If we had zero new results, attempt to update the watermark if the poll result
  // provided a watermark. Otherwise attempt to claim all pending outputs.
  LOG.info(
      "{} - current round of polling took {} ms and returned {} results, "
          + "of which {} were new.",
      c.element(),
      new Duration(now, Instant.now()).getMillis(),
      res.getOutputs().size(),
      newResults.getOutputs().size());

  TerminationStateT terminationState = pollingRestriction.getTerminationState();
  if (!newResults.getOutputs().isEmpty()) {
    terminationState =
        getTerminationCondition().onSeenNewOutput(Instant.now(), terminationState);
  }

  if (!tracker.tryClaim(KV.of(newResults, terminationState))) {
    LOG.info("{} - will not emit poll result tryClaim failed.", c.element());
    return stop();
  }

  if (!newResults.getOutputs().isEmpty()) {
    c.output(KV.of(c.element(), newResults.getOutputs()));
  }

  if (newResults.getWatermark() != null) {
    watermarkEstimator.setWatermark(newResults.getWatermark());
  }

  Instant currentTime = Instant.now();
  if (getTerminationCondition().canStopPolling(currentTime, terminationState)) {
    LOG.info(
        "{} - told to stop polling by polling function at {} with termination state {}.",
        c.element(),
        currentTime,
        getTerminationCondition().toString(terminationState));
    return stop();
  }

  if (BoundedWindow.TIMESTAMP_MAX_VALUE.equals(newResults.getWatermark())) {
    LOG.info("{} - will stop polling, reached max timestamp.", c.element());
    return stop();
  }

  LOG.info(
      "{} - will resume polling in {} ms.", c.element(), spec.getPollInterval().getMillis());
  return resume().withResumeDelay(spec.getPollInterval());
}