Skip to content

Commit ee0ffd1

Browse files
author
jruaux
committed
feat: Added retries for timeout exceptions in sink and source tasks
1 parent 62781b2 commit ee0ffd1

File tree

3 files changed

+155
-149
lines changed

3 files changed

+155
-149
lines changed

core/redis-kafka-connect/src/main/java/com/redis/kafka/connect/sink/RedisSinkTask.java

Lines changed: 66 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -15,31 +15,33 @@
1515
*/
1616
package com.redis.kafka.connect.sink;
1717

18-
import java.io.IOException;
1918
import java.util.ArrayList;
2019
import java.util.Collection;
2120
import java.util.Collections;
2221
import java.util.HashMap;
2322
import java.util.LinkedHashMap;
2423
import java.util.List;
2524
import java.util.Map;
26-
import java.util.concurrent.ConcurrentHashMap;
25+
import java.util.Map.Entry;
26+
import java.util.stream.Collector;
2727
import java.util.stream.Collectors;
2828

29+
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
2930
import org.apache.kafka.common.TopicPartition;
3031
import org.apache.kafka.common.config.ConfigException;
3132
import org.apache.kafka.connect.data.Field;
3233
import org.apache.kafka.connect.data.Struct;
3334
import org.apache.kafka.connect.errors.ConnectException;
3435
import org.apache.kafka.connect.errors.DataException;
36+
import org.apache.kafka.connect.errors.RetriableException;
3537
import org.apache.kafka.connect.json.JsonConverter;
3638
import org.apache.kafka.connect.sink.SinkRecord;
3739
import org.apache.kafka.connect.sink.SinkTask;
3840
import org.apache.kafka.connect.storage.Converter;
3941
import org.slf4j.Logger;
4042
import org.slf4j.LoggerFactory;
4143
import org.springframework.batch.item.ExecutionContext;
42-
import org.springframework.util.Assert;
44+
import org.springframework.util.CollectionUtils;
4345

4446
import com.fasterxml.jackson.annotation.JsonInclude;
4547
import com.fasterxml.jackson.core.JsonProcessingException;
@@ -65,8 +67,9 @@
6567

6668
import io.lettuce.core.AbstractRedisClient;
6769
import io.lettuce.core.KeyValue;
70+
import io.lettuce.core.RedisCommandTimeoutException;
71+
import io.lettuce.core.RedisConnectionException;
6872
import io.lettuce.core.codec.ByteArrayCodec;
69-
import io.netty.util.internal.StringUtil;
7073

7174
public class RedisSinkTask extends SinkTask {
7275

@@ -76,6 +79,9 @@ public class RedisSinkTask extends SinkTask {
7679

7780
private static final ObjectMapper objectMapper = objectMapper();
7881

82+
private static final Collector<SinkOffsetState, ?, Map<String, String>> offsetCollector = Collectors
83+
.toMap(RedisSinkTask::offsetKey, RedisSinkTask::offsetValue);
84+
7985
private RedisSinkConfig config;
8086

8187
private AbstractRedisClient client;
@@ -107,42 +113,45 @@ public void start(final Map<String, String> props) {
107113
client = config.client();
108114
connection = RedisModulesUtils.connection(client);
109115
writer = new OperationItemWriter<>(client, ByteArrayCodec.INSTANCE, operation());
110-
writer.setMultiExec(config.isMultiexec());
116+
writer.setMultiExec(config.isMultiExec());
111117
writer.setWaitReplicas(config.getWaitReplicas());
112118
writer.setWaitTimeout(config.getWaitTimeout());
113119
writer.setPoolSize(config.getPoolSize());
114120
writer.open(new ExecutionContext());
115-
final java.util.Set<TopicPartition> assignment = this.context.assignment();
116-
if (!assignment.isEmpty()) {
117-
Map<TopicPartition, Long> partitionOffsets = new HashMap<>(assignment.size());
118-
for (SinkOffsetState state : offsetStates(assignment)) {
119-
partitionOffsets.put(state.topicPartition(), state.offset());
120-
log.info("Requesting offset {} for {}", state.offset(), state.topicPartition());
121-
}
122-
for (TopicPartition topicPartition : assignment) {
123-
partitionOffsets.putIfAbsent(topicPartition, 0L);
124-
}
125-
this.context.offset(partitionOffsets);
121+
java.util.Set<TopicPartition> assignment = this.context.assignment();
122+
if (CollectionUtils.isEmpty(assignment)) {
123+
return;
124+
}
125+
Map<TopicPartition, Long> partitionOffsets = new HashMap<>(assignment.size());
126+
for (SinkOffsetState state : offsetStates(assignment)) {
127+
partitionOffsets.put(state.topicPartition(), state.offset());
128+
log.info("Requesting offset {} for {}", state.offset(), state.topicPartition());
129+
}
130+
for (TopicPartition topicPartition : assignment) {
131+
partitionOffsets.putIfAbsent(topicPartition, 0L);
126132
}
133+
this.context.offset(partitionOffsets);
127134
}
128135

129136
private Collection<SinkOffsetState> offsetStates(java.util.Set<TopicPartition> assignment) {
130-
Collection<SinkOffsetState> offsetStates = new ArrayList<>();
131-
String[] partitionKeys = assignment.stream().map(a -> offsetKey(a.topic(), a.partition())).toArray(String[]::new);
137+
String[] partitionKeys = assignment.stream().map(this::offsetKey).toArray(String[]::new);
132138
List<KeyValue<String, String>> values = connection.sync().mget(partitionKeys);
133-
for (KeyValue<String, String> value : values) {
134-
if (value.hasValue()) {
135-
try {
136-
offsetStates.add(objectMapper.readValue(value.getValue(), SinkOffsetState.class));
137-
} catch (IOException e) {
138-
throw new DataException(e);
139-
}
140-
}
139+
return values.stream().filter(KeyValue::hasValue).map(this::offsetState).collect(Collectors.toList());
140+
}
141+
142+
private String offsetKey(TopicPartition partition) {
143+
return offsetKey(partition.topic(), partition.partition());
144+
}
145+
146+
private SinkOffsetState offsetState(KeyValue<String, String> value) {
147+
try {
148+
return objectMapper.readValue(value.getValue(), SinkOffsetState.class);
149+
} catch (JsonProcessingException e) {
150+
throw new DataException("Could not parse sink offset state", e);
141151
}
142-
return offsetStates;
143152
}
144153

145-
private String offsetKey(String topic, Integer partition) {
154+
private static String offsetKey(String topic, Integer partition) {
146155
return String.format(OFFSET_KEY_FORMAT, topic, partition);
147156
}
148157

@@ -208,9 +217,6 @@ private byte[] value(SinkRecord sinkRecord) {
208217

209218
private byte[] jsonValue(SinkRecord sinkRecord) {
210219
Object value = sinkRecord.value();
211-
if (value == null) {
212-
return null;
213-
}
214220
if (value instanceof byte[]) {
215221
return (byte[]) value;
216222
}
@@ -262,9 +268,6 @@ private String keyspace(SinkRecord sinkRecord) {
262268
}
263269

264270
private byte[] bytes(String source, Object input) {
265-
if (input == null) {
266-
return null;
267-
}
268271
if (input instanceof byte[]) {
269272
return (byte[]) input;
270273
}
@@ -283,9 +286,6 @@ private byte[] collectionKey(SinkRecord sinkRecord) {
283286
@SuppressWarnings("unchecked")
284287
private Map<byte[], byte[]> map(SinkRecord sinkRecord) {
285288
Object value = sinkRecord.value();
286-
if (value == null) {
287-
return null;
288-
}
289289
if (value instanceof Struct) {
290290
Map<byte[], byte[]> body = new LinkedHashMap<>();
291291
Struct struct = (Struct) value;
@@ -311,16 +311,13 @@ private Map<byte[], byte[]> map(SinkRecord sinkRecord) {
311311
public void stop() {
312312
if (writer != null) {
313313
writer.close();
314-
writer = null;
315314
}
316315
if (connection != null) {
317316
connection.close();
318-
connection = null;
319317
}
320318
if (client != null) {
321319
client.shutdown();
322320
client.getResources().shutdown();
323-
client = null;
324321
}
325322
}
326323

@@ -330,36 +327,37 @@ public void put(final Collection<SinkRecord> records) {
330327
try {
331328
writer.write(new ArrayList<>(records));
332329
log.info("Wrote {} records", records.size());
333-
} catch (Exception e) {
334-
log.warn("Could not write {} records", records.size(), e);
330+
} catch (RedisConnectionException e) {
331+
throw new RetriableException("Could not get connection to Redis", e);
332+
} catch (RedisCommandTimeoutException e) {
333+
throw new RetriableException("Timeout while writing sink records", e);
335334
}
336-
Map<TopicPartition, Long> data = new ConcurrentHashMap<>(100);
337-
for (SinkRecord sinkRecord : records) {
338-
Assert.isTrue(!StringUtil.isNullOrEmpty(sinkRecord.topic()), "topic cannot be null or empty.");
339-
Assert.notNull(sinkRecord.kafkaPartition(), "partition cannot be null.");
340-
Assert.isTrue(sinkRecord.kafkaOffset() >= 0, "offset must be greater than or equal 0.");
341-
TopicPartition partition = new TopicPartition(sinkRecord.topic(), sinkRecord.kafkaPartition());
342-
long current = data.getOrDefault(partition, Long.MIN_VALUE);
343-
if (sinkRecord.kafkaOffset() > current) {
344-
data.put(partition, sinkRecord.kafkaOffset());
345-
}
346-
}
347-
List<SinkOffsetState> offsetData = data.entrySet().stream().map(e -> SinkOffsetState.of(e.getKey(), e.getValue()))
348-
.collect(Collectors.toList());
349-
if (!offsetData.isEmpty()) {
350-
Map<String, String> offsets = new LinkedHashMap<>(offsetData.size());
351-
for (SinkOffsetState e : offsetData) {
352-
String key = offsetKey(e.topic(), e.partition());
353-
String value;
354-
try {
355-
value = objectMapper.writeValueAsString(e);
356-
} catch (JsonProcessingException e1) {
357-
throw new DataException(e1);
358-
}
359-
offsets.put(key, value);
360-
log.trace("put() - Setting offset: {}", e);
361-
}
335+
}
336+
337+
@Override
338+
public void flush(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
339+
Map<String, String> offsets = currentOffsets.entrySet().stream().map(this::offsetState).collect(offsetCollector);
340+
log.trace("Writing offsets: {}", offsets);
341+
try {
362342
connection.sync().mset(offsets);
343+
} catch (RedisCommandTimeoutException e) {
344+
throw new RetriableException("Could not write offsets", e);
345+
}
346+
}
347+
348+
private SinkOffsetState offsetState(Entry<TopicPartition, OffsetAndMetadata> entry) {
349+
return SinkOffsetState.of(entry.getKey(), entry.getValue().offset());
350+
}
351+
352+
private static String offsetKey(SinkOffsetState state) {
353+
return offsetKey(state.topic(), state.partition());
354+
}
355+
356+
private static String offsetValue(SinkOffsetState state) {
357+
try {
358+
return objectMapper.writeValueAsString(state);
359+
} catch (JsonProcessingException e) {
360+
throw new DataException("Could not serialize sink offset state", e);
363361
}
364362
}
365363

0 commit comments

Comments
 (0)