Skip to content

Commit

Permalink
feat: Added retries for timeout exceptions in sink and source tasks
Browse files Browse the repository at this point in the history
  • Loading branch information
jruaux committed Oct 19, 2023
1 parent 62781b2 commit ee0ffd1
Show file tree
Hide file tree
Showing 3 changed files with 155 additions and 149 deletions.
134 changes: 66 additions & 68 deletions core/redis-kafka-connect/src/main/java/com/redis/kafka/connect/sink/RedisSinkTask.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,31 +15,33 @@
*/
package com.redis.kafka.connect.sink;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map.Entry;
import java.util.stream.Collector;
import java.util.stream.Collectors;

import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.config.ConfigException;
import org.apache.kafka.connect.data.Field;
import org.apache.kafka.connect.data.Struct;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.errors.DataException;
import org.apache.kafka.connect.errors.RetriableException;
import org.apache.kafka.connect.json.JsonConverter;
import org.apache.kafka.connect.sink.SinkRecord;
import org.apache.kafka.connect.sink.SinkTask;
import org.apache.kafka.connect.storage.Converter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
Expand All @@ -65,8 +67,9 @@

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.KeyValue;
import io.lettuce.core.RedisCommandTimeoutException;
import io.lettuce.core.RedisConnectionException;
import io.lettuce.core.codec.ByteArrayCodec;
import io.netty.util.internal.StringUtil;

public class RedisSinkTask extends SinkTask {

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

private static final ObjectMapper objectMapper = objectMapper();

private static final Collector<SinkOffsetState, ?, Map<String, String>> offsetCollector = Collectors
.toMap(RedisSinkTask::offsetKey, RedisSinkTask::offsetValue);

private RedisSinkConfig config;

private AbstractRedisClient client;
Expand Down Expand Up @@ -107,42 +113,45 @@ public void start(final Map<String, String> props) {
client = config.client();
connection = RedisModulesUtils.connection(client);
writer = new OperationItemWriter<>(client, ByteArrayCodec.INSTANCE, operation());
writer.setMultiExec(config.isMultiexec());
writer.setMultiExec(config.isMultiExec());
writer.setWaitReplicas(config.getWaitReplicas());
writer.setWaitTimeout(config.getWaitTimeout());
writer.setPoolSize(config.getPoolSize());
writer.open(new ExecutionContext());
final java.util.Set<TopicPartition> assignment = this.context.assignment();
if (!assignment.isEmpty()) {
Map<TopicPartition, Long> partitionOffsets = new HashMap<>(assignment.size());
for (SinkOffsetState state : offsetStates(assignment)) {
partitionOffsets.put(state.topicPartition(), state.offset());
log.info("Requesting offset {} for {}", state.offset(), state.topicPartition());
}
for (TopicPartition topicPartition : assignment) {
partitionOffsets.putIfAbsent(topicPartition, 0L);
}
this.context.offset(partitionOffsets);
java.util.Set<TopicPartition> assignment = this.context.assignment();
if (CollectionUtils.isEmpty(assignment)) {
return;
}
Map<TopicPartition, Long> partitionOffsets = new HashMap<>(assignment.size());
for (SinkOffsetState state : offsetStates(assignment)) {
partitionOffsets.put(state.topicPartition(), state.offset());
log.info("Requesting offset {} for {}", state.offset(), state.topicPartition());
}
for (TopicPartition topicPartition : assignment) {
partitionOffsets.putIfAbsent(topicPartition, 0L);
}
this.context.offset(partitionOffsets);
}

private Collection<SinkOffsetState> offsetStates(java.util.Set<TopicPartition> assignment) {
Collection<SinkOffsetState> offsetStates = new ArrayList<>();
String[] partitionKeys = assignment.stream().map(a -> offsetKey(a.topic(), a.partition())).toArray(String[]::new);
String[] partitionKeys = assignment.stream().map(this::offsetKey).toArray(String[]::new);
List<KeyValue<String, String>> values = connection.sync().mget(partitionKeys);
for (KeyValue<String, String> value : values) {
if (value.hasValue()) {
try {
offsetStates.add(objectMapper.readValue(value.getValue(), SinkOffsetState.class));
} catch (IOException e) {
throw new DataException(e);
}
}
return values.stream().filter(KeyValue::hasValue).map(this::offsetState).collect(Collectors.toList());
}

private String offsetKey(TopicPartition partition) {
return offsetKey(partition.topic(), partition.partition());
}

private SinkOffsetState offsetState(KeyValue<String, String> value) {
try {
return objectMapper.readValue(value.getValue(), SinkOffsetState.class);
} catch (JsonProcessingException e) {
throw new DataException("Could not parse sink offset state", e);
}
return offsetStates;
}

private String offsetKey(String topic, Integer partition) {
private static String offsetKey(String topic, Integer partition) {
return String.format(OFFSET_KEY_FORMAT, topic, partition);
}

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

private byte[] jsonValue(SinkRecord sinkRecord) {
Object value = sinkRecord.value();
if (value == null) {
return null;
}
if (value instanceof byte[]) {
return (byte[]) value;
}
Expand Down Expand Up @@ -262,9 +268,6 @@ private String keyspace(SinkRecord sinkRecord) {
}

private byte[] bytes(String source, Object input) {
if (input == null) {
return null;
}
if (input instanceof byte[]) {
return (byte[]) input;
}
Expand All @@ -283,9 +286,6 @@ private byte[] collectionKey(SinkRecord sinkRecord) {
@SuppressWarnings("unchecked")
private Map<byte[], byte[]> map(SinkRecord sinkRecord) {
Object value = sinkRecord.value();
if (value == null) {
return null;
}
if (value instanceof Struct) {
Map<byte[], byte[]> body = new LinkedHashMap<>();
Struct struct = (Struct) value;
Expand All @@ -311,16 +311,13 @@ private Map<byte[], byte[]> map(SinkRecord sinkRecord) {
public void stop() {
if (writer != null) {
writer.close();
writer = null;
}
if (connection != null) {
connection.close();
connection = null;
}
if (client != null) {
client.shutdown();
client.getResources().shutdown();
client = null;
}
}

Expand All @@ -330,36 +327,37 @@ public void put(final Collection<SinkRecord> records) {
try {
writer.write(new ArrayList<>(records));
log.info("Wrote {} records", records.size());
} catch (Exception e) {
log.warn("Could not write {} records", records.size(), e);
} catch (RedisConnectionException e) {
throw new RetriableException("Could not get connection to Redis", e);
} catch (RedisCommandTimeoutException e) {
throw new RetriableException("Timeout while writing sink records", e);
}
Map<TopicPartition, Long> data = new ConcurrentHashMap<>(100);
for (SinkRecord sinkRecord : records) {
Assert.isTrue(!StringUtil.isNullOrEmpty(sinkRecord.topic()), "topic cannot be null or empty.");
Assert.notNull(sinkRecord.kafkaPartition(), "partition cannot be null.");
Assert.isTrue(sinkRecord.kafkaOffset() >= 0, "offset must be greater than or equal 0.");
TopicPartition partition = new TopicPartition(sinkRecord.topic(), sinkRecord.kafkaPartition());
long current = data.getOrDefault(partition, Long.MIN_VALUE);
if (sinkRecord.kafkaOffset() > current) {
data.put(partition, sinkRecord.kafkaOffset());
}
}
List<SinkOffsetState> offsetData = data.entrySet().stream().map(e -> SinkOffsetState.of(e.getKey(), e.getValue()))
.collect(Collectors.toList());
if (!offsetData.isEmpty()) {
Map<String, String> offsets = new LinkedHashMap<>(offsetData.size());
for (SinkOffsetState e : offsetData) {
String key = offsetKey(e.topic(), e.partition());
String value;
try {
value = objectMapper.writeValueAsString(e);
} catch (JsonProcessingException e1) {
throw new DataException(e1);
}
offsets.put(key, value);
log.trace("put() - Setting offset: {}", e);
}
}

@Override
public void flush(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
Map<String, String> offsets = currentOffsets.entrySet().stream().map(this::offsetState).collect(offsetCollector);
log.trace("Writing offsets: {}", offsets);
try {
connection.sync().mset(offsets);
} catch (RedisCommandTimeoutException e) {
throw new RetriableException("Could not write offsets", e);
}
}

private SinkOffsetState offsetState(Entry<TopicPartition, OffsetAndMetadata> entry) {
return SinkOffsetState.of(entry.getKey(), entry.getValue().offset());
}

private static String offsetKey(SinkOffsetState state) {
return offsetKey(state.topic(), state.partition());
}

private static String offsetValue(SinkOffsetState state) {
try {
return objectMapper.writeValueAsString(state);
} catch (JsonProcessingException e) {
throw new DataException("Could not serialize sink offset state", e);
}
}

Expand Down
Loading

0 comments on commit ee0ffd1

Please sign in to comment.