KafkaItemWriter.write should not return until items are confirmed to have been written
Issue #3773
This commit is contained in:
committed by
Mahmoud Ben Hassine
parent
a6f283b4a8
commit
c886a60b4d
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user