add support to repair offsets on restore (#253)

Author: agavra

kafka-client/src/main/java/dev/responsive/kafka/api/ResponsiveKafkaStreams.java

@@ -421,6 +421,7 @@ public Params withTime(final Time time) {
     public Params build() {
       this.responsiveKafkaClientSupplier = new ResponsiveKafkaClientSupplier(
           clientSupplier,
+          responsiveConfig,
           streamsConfig,
           storeRegistry,
           metrics,

kafka-client/src/main/java/dev/responsive/kafka/api/config/ResponsiveConfig.java

@@ -231,6 +231,15 @@ public class ResponsiveConfig extends AbstractConfig {
       + "runs in. When set to RUN, runs the Kafka Streams app. When set to MIGRATE, runs app"
       + " migration.";
 
+  // ------------------ Misc functional overrides ----------------------
+  public static final String RESTORE_OFFSET_REPAIR_ENABLED_CONFIG = "responsive.restore.offset.repair.enabled";
+  public static final boolean RESTORE_OFFSET_REPAIR_ENABLED_DEFAULT = false;
+  public static final String RESTORE_OFFSET_REPAIR_ENABLED_DOC = "When set to 'true', " + RESTORE_OFFSET_REPAIR_ENABLED_CONFIG
+      + " will ignore OffsetOutOfRangeException and instead seek to the earliest available offset. This exception "
+      + "should only happen in situations where there is truncation/retention on the changelog topic and restoring from the latest "
+      + "committed offset in the remote store is no longer possible. Note that in some situations this may cause data "
+      + "loss, use this configuration with caution";
+
   // ------------------ StreamsConfig overrides ----------------------
 
   // These configuration values are required by Responsive, and a ConfigException will
@@ -452,6 +461,12 @@ public class ResponsiveConfig extends AbstractConfig {
           ),
           Importance.LOW,
           RESPONSIVE_MODE_DOC
+      ).define(
+          RESTORE_OFFSET_REPAIR_ENABLED_CONFIG,
+          Type.BOOLEAN,
+          RESTORE_OFFSET_REPAIR_ENABLED_DEFAULT,
+          Importance.LOW,
+          RESTORE_OFFSET_REPAIR_ENABLED_DOC
       );
 
   /**

kafka-client/src/main/java/dev/responsive/kafka/internal/clients/ResponsiveKafkaClientSupplier.java

@@ -19,6 +19,7 @@
 import static org.apache.kafka.streams.StreamsConfig.AT_LEAST_ONCE;
 
 import dev.responsive.kafka.api.config.CompatibilityMode;
+import dev.responsive.kafka.api.config.ResponsiveConfig;
 import dev.responsive.kafka.internal.metrics.EndOffsetsPoller;
 import dev.responsive.kafka.internal.metrics.MetricPublishingCommitListener;
 import dev.responsive.kafka.internal.metrics.ResponsiveMetrics;
@@ -67,15 +68,25 @@ public final class ResponsiveKafkaClientSupplier implements KafkaClientSupplier
   private final String applicationId;
   private final boolean eos;
   private final CompatibilityMode compatibilityMode;
+  private final boolean repairRestoreOffsetOutOfRange;
 
   public ResponsiveKafkaClientSupplier(
       final KafkaClientSupplier clientSupplier,
+      final ResponsiveConfig responsiveConfig,
       final StreamsConfig configs,
       final ResponsiveStoreRegistry storeRegistry,
       final ResponsiveMetrics metrics,
       final CompatibilityMode compatibilityMode
   ) {
-    this(new Factories() {}, clientSupplier, configs, storeRegistry, metrics, compatibilityMode);
+    this(
+        new Factories() {},
+        clientSupplier,
+        configs,
+        storeRegistry,
+        metrics,
+        compatibilityMode,
+        responsiveConfig.getBoolean(ResponsiveConfig.RESTORE_OFFSET_REPAIR_ENABLED_CONFIG)
+    );
   }
 
   ResponsiveKafkaClientSupplier(
@@ -84,13 +95,15 @@ public ResponsiveKafkaClientSupplier(
       final StreamsConfig configs,
       final ResponsiveStoreRegistry storeRegistry,
       final ResponsiveMetrics metrics,
-      final CompatibilityMode compatibilityMode
+      final CompatibilityMode compatibilityMode,
+      final boolean repairRestoreOffsetOutOfRange
   ) {
     this.factories = factories;
     this.wrapped = wrapped;
     this.storeRegistry = storeRegistry;
     this.metrics = metrics;
     this.compatibilityMode = compatibilityMode;
+    this.repairRestoreOffsetOutOfRange = repairRestoreOffsetOutOfRange;
 
     eos = !(AT_LEAST_ONCE.equals(
         configs.getString(StreamsConfig.PROCESSING_GUARANTEE_CONFIG)));
@@ -174,7 +187,8 @@ public Consumer<byte[], byte[]> getRestoreConsumer(final Map<String, Object> con
     return factories.createRestoreConsumer(
         clientId,
         wrapped.getRestoreConsumer(config),
-        storeRegistry::getCommittedOffset
+        storeRegistry::getCommittedOffset,
+        repairRestoreOffsetOutOfRange
     );
 
   }
@@ -403,12 +417,14 @@ default MetricPublishingCommitListener createMetricsPublishingCommitListener(
     default ResponsiveRestoreConsumer<byte[], byte[]> createRestoreConsumer(
         final String clientId,
         final Consumer<byte[], byte[]> restoreConsumer,
-        final Function<TopicPartition, OptionalLong> getCommittedOffset
+        final Function<TopicPartition, OptionalLong> getCommittedOffset,
+        final boolean repairRestoreOffsetOutOfRange
     ) {
       return new ResponsiveRestoreConsumer<>(
           clientId,
           restoreConsumer,
-          getCommittedOffset
+          getCommittedOffset,
+          repairRestoreOffsetOutOfRange
       );
     }
   }

kafka-client/src/main/java/dev/responsive/kafka/internal/clients/ResponsiveRestoreConsumer.java

@@ -1,5 +1,6 @@
 package dev.responsive.kafka.internal.clients;
 
+import dev.responsive.kafka.api.config.ResponsiveConfig;
 import java.time.Duration;
 import java.util.Collection;
 import java.util.Collections;
@@ -12,6 +13,7 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.OffsetOutOfRangeException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.utils.LogContext;
 import org.slf4j.Logger;
@@ -48,17 +50,20 @@ public class ResponsiveRestoreConsumer<K, V> extends DelegatingConsumer<K, V> {
 
   private final Function<TopicPartition, OptionalLong> startOffsets;
   private final Set<TopicPartition> uninitializedOffsets = new HashSet<>();
+  private final boolean repairRestoreOffsetOutOfRange;
 
   public ResponsiveRestoreConsumer(
       final String clientId,
       final Consumer<K, V> delegate,
-      final Function<TopicPartition, OptionalLong> startOffsets
+      final Function<TopicPartition, OptionalLong> startOffsets,
+      final boolean repairRestoreOffsetOutOfRange
   ) {
     super(delegate);
     this.startOffsets = Objects.requireNonNull(startOffsets);
     this.log = new LogContext(
         String.format("responsive-restore-consumer [%s]", Objects.requireNonNull(clientId))
     ).logger(ResponsiveConsumer.class);
+    this.repairRestoreOffsetOutOfRange = repairRestoreOffsetOutOfRange;
   }
 
   private Set<TopicPartition> initializeOffsets(final Collection<TopicPartition> partitions) {
@@ -101,7 +106,23 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
       log.error("Found uninitialized changelog partitions during poll: {}", uninitializedOffsets);
       throw new IllegalStateException("Restore consumer invoked poll without initializing offsets");
     }
-    return super.poll(timeout);
+
+    try {
+      return super.poll(timeout);
+    } catch (final OffsetOutOfRangeException e) {
+      // e.partitions() should never be empty, but we check it because accidentally
+      // passing in empty partitions is catastrophic (it will seek all partitions to
+      // beginning)
+      if (repairRestoreOffsetOutOfRange && !e.partitions().isEmpty()) {
+        log.warn("The restore consumer attempted to seek to offsets that are no longer in range. "
+            + ResponsiveConfig.RESTORE_OFFSET_REPAIR_ENABLED_CONFIG + " is enabled and will "
+            + "automatically seek " + e.offsetOutOfRangePartitions() + " to their earliest offset "
+            + "and continue restoration.", e);
+        super.seekToBeginning(e.partitions());
+        return super.poll(timeout);
+      }
+      throw e;
+    }
   }
 
   @Override

kafka-client/src/test/java/dev/responsive/kafka/integration/ResponsiveKeyValueStoreRestoreIntegrationTest.java

@@ -104,6 +104,7 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.IsolationLevel;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.TopicConfig;
@@ -226,6 +227,76 @@ private static Stream<Arguments> shouldFlushTestParams() {
         );
   }
 
+  @ParameterizedTest
+  @EnumSource(KVSchema.class)
+  public void shouldRepairOffsetsIfOutOfRangeAndConfigured(final KVSchema type) throws Exception {
+    // Given:
+    final Map<String, Object> properties = getMutableProperties();
+    properties.put(ResponsiveConfig.RESTORE_OFFSET_REPAIR_ENABLED_CONFIG, true);
+    final KafkaProducer<Long, Long> producer = new KafkaProducer<>(properties);
+    final KafkaClientSupplier defaultClientSupplier = new DefaultKafkaClientSupplier();
+    final CassandraClientFactory defaultFactory = new DefaultCassandraClientFactory();
+    final TopicPartition input = new TopicPartition(inputTopic(), 0);
+    final TopicPartition changelog = new TopicPartition(name + "-" + aggName() + "-changelog", 0);
+
+    // When:
+    final long clOffset;
+    try (final ResponsiveKafkaStreams streams
+             = buildAggregatorApp(properties, defaultClientSupplier, defaultFactory, type, false)) {
+      IntegrationTestUtils.startAppAndAwaitRunning(Duration.ofSeconds(30), streams);
+      // Send some data through
+      pipeInput(
+          inputTopic(),
+          1,
+          producer,
+          System::currentTimeMillis,
+          0,
+          1,
+          LongStream.range(0, 100).toArray()
+      );
+      // Wait for it to be processed
+      waitTillFullyConsumed(input, Duration.ofSeconds(120));
+
+      final List<ConsumerRecord<Long, Long>> changelogRecords
+          = slurpPartition(changelog, properties);
+      clOffset = changelogRecords.get(changelogRecords.size() - 1).offset();
+    }
+
+    // produce some data so we can truncate the data that has been committed
+    final RecordMetadata recordMetadata =
+        producer.send(new ProducerRecord<>(changelog.topic(), changelog.partition(), -1L, -1L))
+            .get();
+
+    // truncate the offset that exists in remote
+    admin.deleteRecords(
+        Map.of(changelog, RecordsToDelete.beforeOffset(recordMetadata.offset()))
+    ).all().get();
+
+    // run another application
+    try (final ResponsiveKafkaStreams streams
+             = buildAggregatorApp(properties, defaultClientSupplier, defaultFactory, type, false)) {
+      IntegrationTestUtils.startAppAndAwaitRunning(Duration.ofSeconds(30), streams);
+      // Send some data through
+      pipeInput(
+          inputTopic(),
+          1,
+          producer,
+          System::currentTimeMillis,
+          0,
+          1,
+          LongStream.range(0, 100).toArray()
+      );
+      // Wait for it to be processed
+      waitTillFullyConsumed(input, Duration.ofSeconds(120));
+
+      // Verify it made progress
+      final List<ConsumerRecord<Long, Long>> changelogRecords
+          = slurpPartition(changelog, properties);
+      final long last = changelogRecords.get(changelogRecords.size() - 1).offset();
+      assertThat(last, greaterThan(clOffset));
+    }
+  }
+
   @ParameterizedTest
   @EnumSource(KVSchema.class)
   public void shouldRestoreUnflushedChangelog(final KVSchema type) throws Exception {

kafka-client/src/test/java/dev/responsive/kafka/internal/clients/ResponsiveKafkaClientSupplierTest.java

@@ -247,7 +247,12 @@ public void shouldCreateGlobalAndRestoreConsumerInFullCompatibilityMode() {
 
     // then:
     verify(factories, Mockito.atLeastOnce()).createGlobalConsumer(any(), any(), any());
-    verify(factories, Mockito.atLeastOnce()).createRestoreConsumer(any(), any(), any());
+    verify(factories, Mockito.atLeastOnce()).createRestoreConsumer(
+        any(),
+        any(),
+        any(),
+        anyBoolean()
+    );
   }
 
   @Test
@@ -261,7 +266,7 @@ public void shouldNotCreateGlobalOrRestoreConsumerInMetricsCompatibilityMode() {
 
     // then:
     verify(factories, Mockito.never()).createGlobalConsumer(any(), any(), any());
-    verify(factories, Mockito.never()).createRestoreConsumer(any(), any(), any());
+    verify(factories, Mockito.never()).createRestoreConsumer(any(), any(), any(), anyBoolean());
   }
 
   @NotNull
@@ -275,7 +280,8 @@ private ResponsiveKafkaClientSupplier supplier(
         new StreamsConfig(CONFIGS),
         storeRegistry,
         metrics,
-        compat
+        compat,
+        false
     );
   }
 

kafka-client/src/test/java/dev/responsive/kafka/internal/clients/ResponsiveRestoreConsumerTest.java

@@ -1,25 +1,31 @@
 package dev.responsive.kafka.internal.clients;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
-import dev.responsive.kafka.internal.clients.ResponsiveRestoreConsumer;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.OptionalLong;
 import java.util.Set;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.OffsetOutOfRangeException;
 import org.apache.kafka.common.TopicPartition;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
 
 @ExtendWith(MockitoExtension.class)
@@ -35,7 +41,11 @@ public class ResponsiveRestoreConsumerTest {
 
   @BeforeEach
   public void setup() {
-    restoreConsumer = new ResponsiveRestoreConsumer<>(
+    restoreConsumer = getRestoreConsumer(false);
+  }
+
+  private ResponsiveRestoreConsumer<?, ?> getRestoreConsumer(boolean repairOffsets) {
+    return  new ResponsiveRestoreConsumer<>(
         "restore-consumer",
         wrapped,
         tp -> {
@@ -46,7 +56,9 @@ public void setup() {
             return OptionalLong.of(456L);
           }
           return OptionalLong.empty();
-        });
+        },
+        repairOffsets
+    );
   }
 
   @Test
@@ -193,6 +205,43 @@ public void shouldClearUninitializedPartitionsWhenUnsubscribed() {
     restoreConsumer.poll(Duration.ofMillis(100));
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  public void shouldSeekToBeginningOnOffsetOutOfRangeIfRepairOffsetConfigured() {
+    // Given:
+    restoreConsumer = getRestoreConsumer(true);
+    final ArgumentCaptor<Collection<TopicPartition>> seekedTps
+        = ArgumentCaptor.forClass(Collection.class);
+
+    when(wrapped.poll(any()))
+        .thenThrow(new OffsetOutOfRangeException(Map.of(TOPIC_PARTITION1, 1L)))
+        .thenReturn(null);
+    doNothing().when(wrapped).seekToBeginning(seekedTps.capture());
+
+    // When:
+    restoreConsumer.assign(List.of(TOPIC_PARTITION1));
+    restoreConsumer.seekToBeginning(List.of(TOPIC_PARTITION1));
+    restoreConsumer.poll(Duration.ofMillis(100));
+
+    // Then:
+    verify(wrapped, Mockito.times(2)).poll(any());
+    verify(wrapped).seekToBeginning(Set.of(TOPIC_PARTITION1));
+  }
+
+  @Test
+  public void shouldThrowOffsetOutOfRangeExceptionIfNoRepairConfigured() {
+    // Given:
+    when(wrapped.poll(any()))
+        .thenThrow(new OffsetOutOfRangeException(Map.of(TOPIC_PARTITION1, 1L)))
+        .thenReturn(null);
+
+    // Then:
+    Assertions.assertThrows(
+        OffsetOutOfRangeException.class,
+        () -> restoreConsumer.poll(Duration.ofMillis(100))
+    );
+  }
+
   @SuppressWarnings("unchecked")
   private void clearConsumerInvocations() {
     clearInvocations(wrapped);