com.redis.spring.batch.step.FlushingChunkProvider Maven / Gradle / Ivy
The newest version!
package com.redis.spring.batch.step;
import java.time.Duration;
import java.util.concurrent.TimeUnit;
import org.springframework.batch.core.StepContribution;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.observability.BatchMetrics;
import org.springframework.batch.core.step.item.SimpleChunkProvider;
import org.springframework.batch.core.step.item.SkipOverflowException;
import org.springframework.batch.core.step.skip.SkipListenerFailedException;
import org.springframework.batch.item.Chunk;
import org.springframework.batch.item.ItemReader;
import org.springframework.batch.repeat.RepeatOperations;
import org.springframework.batch.repeat.RepeatStatus;
import org.springframework.util.Assert;
import com.redis.spring.batch.item.PollableItemReader;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Timer;
import io.micrometer.core.instrument.Timer.Sample;
/**
* Fault-tolerant implementation of the ChunkProvider interface, that allows for
* skipping or retry of items that cause exceptions, as well as incomplete
* chunks when timeout is reached.
*/
public class FlushingChunkProvider extends SimpleChunkProvider {
public static final Duration DEFAULT_FLUSH_INTERVAL = Duration.ofMillis(50);
public static final Duration DEFAULT_IDLE_TIMEOUT = Duration.ofMillis(Long.MAX_VALUE); // no idle timeout by default
private final RepeatOperations repeatOperations;
private long flushInterval = DEFAULT_FLUSH_INTERVAL.toMillis();
private long idleTimeout = DEFAULT_IDLE_TIMEOUT.toMillis();
private long lastActivity = 0;
public FlushingChunkProvider(ItemReader itemReader, RepeatOperations repeatOperations) {
super(itemReader, repeatOperations);
Assert.isTrue(itemReader instanceof PollableItemReader, "Reader must extend PollableItemReader");
this.repeatOperations = repeatOperations;
}
public void setFlushInterval(Duration interval) {
this.flushInterval = interval.toMillis();
}
public void setIdleTimeout(Duration idleTimeout) {
this.idleTimeout = idleTimeout.toMillis();
}
private void stopFlushTimer(Timer.Sample sample, StepExecution stepExecution, String status) {
sample.stop(BatchMetrics.createTimer(Metrics.globalRegistry, "item.read", "Item reading duration",
Tag.of("job.name", stepExecution.getJobExecution().getJobInstance().getJobName()),
Tag.of("step.name", stepExecution.getStepName()), Tag.of("status", status)));
}
@Override
public Chunk provide(StepContribution contribution) {
long start = System.currentTimeMillis();
if (lastActivity == 0) {
lastActivity = start;
}
final Chunk inputs = new Chunk<>();
repeatOperations.iterate(context -> {
long pollingTimeout = flushInterval - millisSince(start);
if (pollingTimeout < 0) {
return RepeatStatus.FINISHED;
}
Sample sample = Timer.start(Metrics.globalRegistry);
I item;
try {
item = read(contribution, inputs, pollingTimeout);
} catch (SkipOverflowException e) {
// read() tells us about an excess of skips by throwing an exception
stopFlushTimer(sample, contribution.getStepExecution(), BatchMetrics.STATUS_FAILURE);
return RepeatStatus.FINISHED;
}
if (item == null) {
if (millisSince(lastActivity) > idleTimeout) {
inputs.setEnd();
}
return RepeatStatus.CONTINUABLE;
}
stopFlushTimer(sample, contribution.getStepExecution(), BatchMetrics.STATUS_SUCCESS);
inputs.add(item);
contribution.incrementReadCount();
lastActivity = System.currentTimeMillis();
return RepeatStatus.CONTINUABLE;
});
return inputs;
}
private long millisSince(long start) {
return System.currentTimeMillis() - start;
}
protected I read(StepContribution contribution, Chunk chunk, long timeout)
throws Exception, InterruptedException {
while (true) {
try {
return doRead(timeout);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw e;
}
}
}
@SuppressWarnings("unchecked")
protected final I doRead(long timeout) throws Exception, InterruptedException {
try {
getListener().beforeRead();
I item = ((PollableItemReader) itemReader).poll(timeout, TimeUnit.MILLISECONDS);
if (item != null) {
getListener().afterRead(item);
}
return item;
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw e;
} catch (Exception e) {
getListener().onReadError(e);
throw e;
}
}
@Override
public void postProcess(StepContribution contribution, Chunk chunk) {
for (Exception e : chunk.getErrors()) {
try {
getListener().onSkipInRead(e);
} catch (RuntimeException ex) {
throw new SkipListenerFailedException("Fatal exception in SkipListener.", ex, e);
}
}
}
}