KafkaItemWriter.write should not return until items are confirmed to have been written

Issue #3773
This commit is contained in:
Jacob Botuck
2021-01-05 11:35:08 -06:00
committed by Mahmoud Ben Hassine
parent a6f283b4a8
commit c886a60b4d
3 changed files with 31 additions and 5 deletions

View File

@@ -43,7 +43,9 @@ public abstract class KeyValueItemWriter<K, V> implements ItemWriter<V>, Initial
K key = itemKeyMapper.convert(item);
writeKeyValue(key, item);
}
flush();
}
protected void flush() throws Exception {}
/**
* Subclasses implement this method to write each item to key value store

View File

@@ -19,7 +19,12 @@ package org.springframework.batch.item.kafka;
import org.springframework.batch.item.ItemWriter;
import org.springframework.batch.item.KeyValueItemWriter;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.support.SendResult;
import org.springframework.util.Assert;
import org.springframework.util.concurrent.ListenableFuture;
import java.util.ArrayList;
import java.util.List;
/**
* <p>
@@ -34,16 +39,25 @@ import org.springframework.util.Assert;
public class KafkaItemWriter<K, T> extends KeyValueItemWriter<K, T> {
protected KafkaTemplate<K, T> kafkaTemplate;
private final List<ListenableFuture<SendResult<K, T>>> listenableFutures = new ArrayList<>();
@Override
protected void writeKeyValue(K key, T value) {
if (this.delete) {
this.kafkaTemplate.sendDefault(key, null);
listenableFutures.add(this.kafkaTemplate.sendDefault(key, null));
}
else {
this.kafkaTemplate.sendDefault(key, value);
listenableFutures.add(this.kafkaTemplate.sendDefault(key, value));
}
}
@Override
protected void flush() throws Exception{
kafkaTemplate.flush();
for(ListenableFuture<SendResult<K,T>> future: listenableFutures){
future.get();
}
listenableFutures.clear();
}
@Override
protected void init() {

View File

@@ -25,25 +25,31 @@ import org.mockito.MockitoAnnotations;
import org.springframework.core.convert.converter.Converter;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.support.SendResult;
import org.springframework.util.concurrent.ListenableFuture;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
public class KafkaItemWriterTests {
@Mock
private KafkaTemplate<String, String> kafkaTemplate;
@Mock
private ListenableFuture<SendResult<String, String>> future;
private KafkaItemKeyMapper itemKeyMapper;
private KafkaItemWriter<String, String> writer;
@Before
public void setUp() throws Exception {
MockitoAnnotations.initMocks(this);
MockitoAnnotations.openMocks(this);
when(this.kafkaTemplate.getDefaultTopic()).thenReturn("defaultTopic");
when(this.kafkaTemplate.sendDefault(any(), any())).thenReturn(future);
this.itemKeyMapper = new KafkaItemKeyMapper();
this.writer = new KafkaItemWriter<>();
this.writer.setKafkaTemplate(this.kafkaTemplate);
@@ -90,6 +96,8 @@ public class KafkaItemWriterTests {
verify(this.kafkaTemplate).sendDefault(items.get(0), items.get(0));
verify(this.kafkaTemplate).sendDefault(items.get(1), items.get(1));
verify(this.kafkaTemplate).flush();
verify(this.future, times(2)).get();
}
@Test
@@ -101,6 +109,8 @@ public class KafkaItemWriterTests {
verify(this.kafkaTemplate).sendDefault(items.get(0), null);
verify(this.kafkaTemplate).sendDefault(items.get(1), null);
verify(this.kafkaTemplate).flush();
verify(this.future, times(2)).get();
}
@Test