@Override public void onFinishWithKeys(List<? extends PartitionKey> partitionKeys, boolean succeeded) { ConsumerWorkingSet workingSet = readState(); doFinish(workingSet, partitionKeys, succeeded); statePersistor.persistState(workingSet.toBytes()); }
@Override public void untake(ConsumerWorkingSet workingSet, List<? extends PartitionKey> partitionKeys) { doExpiry(workingSet); for (PartitionKey key : partitionKeys) { ConsumablePartition consumablePartition = workingSet.lookup(key); // don't need to assertInProgress because untake() already does that consumablePartition.untake(); } }
@Override public PartitionConsumerResult doConsume(ConsumerWorkingSet workingSet, PartitionAcceptor acceptor) { doExpiry(workingSet); workingSet.populate(getPartitionedFileSet(), getConfiguration()); List<PartitionDetail> toConsume = selectPartitions(acceptor, workingSet); return new PartitionConsumerResult(toConsume, removeDiscardedPartitions(workingSet)); }
@Override public void apply() throws Exception { // using the same PartitionConsumer (which remembers the PartitionConsumerState) to consume additional // partitions results in only the newly added partitions (corresponding to partitionKeys2) to be returned Assert.assertEquals(partitionKeys2, toKeys(partitionConsumer.consumePartitions().getPartitions())); } });
@Override public PartitionConsumerResult consumePartitions(PartitionAcceptor acceptor) { ConsumerWorkingSet workingSet = readState(); PartitionConsumerResult partitionsResult = doConsume(workingSet, acceptor); statePersistor.persistState(workingSet.toBytes()); return partitionsResult; }
@Override public void untakeWithKeys(List<? extends PartitionKey> partitionKeys) { ConsumerWorkingSet workingSet = readState(); untake(workingSet, partitionKeys); statePersistor.persistState(workingSet.toBytes()); }
private ConsumerWorkingSet readState() { byte[] bytes = statePersistor.readState(); return bytes == null ? new ConsumerWorkingSet() : ConsumerWorkingSet.fromBytes(bytes); } }
/** * Removes the given partition keys from the working set, as they have been successfully processed. */ protected void commit(ConsumerWorkingSet workingSet, List<? extends PartitionKey> partitionKeys) { for (PartitionKey key : partitionKeys) { ConsumablePartition consumablePartition = workingSet.lookup(key); assertInProgress(consumablePartition); workingSet.remove(key); } }
@Override public PartitionConsumerResult consumePartitions(final int limit) { return consumePartitions(new PartitionAcceptor.Limit(limit)); }
private PartitionConsumer getPartitionConsumer(DatasetContext context) { PartitionedFileSet lines = context.getDataset(partitionedFileSetName); return new ConcurrentPartitionConsumer(lines, new DelegatingStatePersistor(context, statePersistor), consumerConfiguration); } }
@Override public void onFinishWithKeys(final List<? extends PartitionKey> partitionKeys, boolean succeeded) { Transactionals.execute(transactional, context -> { getPartitionConsumer(context).onFinishWithKeys(partitionKeys, succeeded); }); }
@Override public PartitionConsumerResult consumePartitions(int limit) { return consumePartitions(new PartitionAcceptor.Limit(limit)); }
/** * See {@link #setInput(MapReduceContext, String, DatasetStatePersistor, ConsumerConfiguration)}, * but using {@link ConsumerConfiguration#DEFAULT}. */ public static BatchPartitionCommitter setInput(MapReduceContext mapreduceContext, String partitionedFileSetName, DatasetStatePersistor statePersistor) { return setInput(mapreduceContext, partitionedFileSetName, statePersistor, ConsumerConfiguration.DEFAULT); }
@Override public void take() { assertState(ProcessState.AVAILABLE); processState = ProcessState.IN_PROGRESS; }
/** * @return a {@link Builder} instance to build an instance of a ConsumerConfiguration. */ public static Builder builder() { return new Builder(); }
/** * Create a ConsumerConfiguration from this builder, using the private ConsumerConfiguration * constructor. */ public ConsumerConfiguration build() { return new ConsumerConfiguration(partitionPredicate, maxWorkingSetSize, timeout, maxRetries); } }
@Override public void retry() { untake(); numFailures++; }
@Override public void apply() throws Exception { // the consumed partition keys should correspond to partitionKeys2, and not include the dropped, but unconsumed // partitions added before them Assert.assertEquals(partitionKeys2, toKeys(partitionConsumer.consumePartitions().getPartitions())); } });
@Override public void complete() { assertState(ProcessState.IN_PROGRESS); processState = ProcessState.COMPLETED; }