Examples with StepContext org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StepContext used on opensource projects

Example 1 with StepContext

use of org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StepContext in project beam by apache.

the class StreamingModeExecutionContext method flushState.

public Map<Long, Runnable> flushState() {
    Map<Long, Runnable> callbacks = new HashMap<>();
    for (StepContext stepContext : getAllStepContexts()) {
        stepContext.flushState();
    }
    if (activeReader != null) {
        Windmill.SourceState.Builder sourceStateBuilder = outputBuilder.getSourceStateUpdatesBuilder();
        final UnboundedSource.CheckpointMark checkpointMark = activeReader.getCheckpointMark();
        final Instant watermark = activeReader.getWatermark();
        long id = ThreadLocalRandom.current().nextLong();
        sourceStateBuilder.addFinalizeIds(id);
        callbacks.put(id, () -> {
            try {
                checkpointMark.finalizeCheckpoint();
            } catch (IOException e) {
                throw new RuntimeException("Exception while finalizing checkpoint", e);
            }
        });
        @SuppressWarnings("unchecked") Coder<UnboundedSource.CheckpointMark> checkpointCoder = ((UnboundedSource<?, UnboundedSource.CheckpointMark>) activeReader.getCurrentSource()).getCheckpointMarkCoder();
        if (checkpointCoder != null) {
            ByteString.Output stream = ByteString.newOutput();
            try {
                checkpointCoder.encode(checkpointMark, stream, Coder.Context.OUTER);
            } catch (IOException e) {
                throw new RuntimeException("Exception while encoding checkpoint", e);
            }
            sourceStateBuilder.setState(stream.toByteString());
        }
        outputBuilder.setSourceWatermark(WindmillTimeUtils.harnessToWindmillTimestamp(watermark));
        backlogBytes = activeReader.getSplitBacklogBytes();
        if (backlogBytes == UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN && WorkerCustomSources.isFirstUnboundedSourceSplit(getSerializedKey())) {
            // Only call getTotalBacklogBytes() on the first split.
            backlogBytes = activeReader.getTotalBacklogBytes();
        }
        outputBuilder.setSourceBacklogBytes(backlogBytes);
        readerCache.cacheReader(getComputationKey(), getWork().getCacheToken(), getWork().getWorkToken(), activeReader);
        activeReader = null;
    }
    return callbacks;
}

Example 2 with StepContext

use of org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StepContext in project beam by apache.

the class StreamingModeExecutionContext method start.

public void start(@Nullable Object key, Windmill.WorkItem work, Instant inputDataWatermark, @Nullable Instant outputDataWatermark, @Nullable Instant synchronizedProcessingTime, WindmillStateReader stateReader, StateFetcher stateFetcher, Windmill.WorkItemCommitRequest.Builder outputBuilder) {
    this.key = key;
    this.work = work;
    this.computationKey = WindmillComputationKey.create(computationId, work.getKey(), work.getShardingKey());
    this.stateFetcher = stateFetcher;
    this.outputBuilder = outputBuilder;
    this.sideInputCache.clear();
    clearSinkFullHint();
    Instant processingTime = Instant.now();
    // timers.  Otherwise a trigger could ignore the timer and orphan the window.
    for (Windmill.Timer timer : work.getTimers().getTimersList()) {
        if (timer.getType() == Windmill.Timer.Type.REALTIME) {
            Instant inferredFiringTime = WindmillTimeUtils.windmillToHarnessTimestamp(timer.getTimestamp()).plus(Duration.millis(1));
            if (inferredFiringTime.isAfter(processingTime)) {
                processingTime = inferredFiringTime;
            }
        }
    }
    Collection<? extends StepContext> stepContexts = getAllStepContexts();
    if (!stepContexts.isEmpty()) {
        // This must be only created once for the workItem as token validation will fail if the same
        // work token is reused.
        WindmillStateCache.ForKey cacheForKey = stateCache.forKey(getComputationKey(), getWork().getCacheToken(), getWorkToken());
        for (StepContext stepContext : stepContexts) {
            stepContext.start(stateReader, inputDataWatermark, processingTime, cacheForKey, outputDataWatermark, synchronizedProcessingTime);
        }
    }
}