KAFKA-5520: KIP-171 - Extend Consumer Group Reset Offset for Stream Application - Updated

checkstyle/import-control-core.xml

@@ -60,6 +60,7 @@
     <allow pkg="kafka.consumer" />
     <allow pkg="joptsimple" />
     <allow pkg="org.apache.kafka.clients.consumer" />
+    <allow class="javax.xml.datatype.Duration" />
   </subpackage>
 
   <subpackage name="coordinator">

core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala

@@ -529,8 +529,8 @@ object ConsumerGroupCommand extends Logging {
         case "Empty" | "Dead" =>
           val partitionsToReset = getPartitionsToReset(groupId)
           val preparedOffsets = prepareOffsetsToReset(groupId, partitionsToReset)
-          val execute = opts.options.has(opts.executeOpt)
-          if (execute)
+          val dryRun = opts.options.has(opts.dryRunOpt)
+          if (!dryRun)
             getConsumer().commitSync(preparedOffsets.asJava)
           preparedOffsets
         case currentState =>
@@ -727,7 +727,7 @@ object ConsumerGroupCommand extends Logging {
       "Has 3 execution options: (default) to plan which offsets to reset, --execute to execute the reset-offsets process, and --export to export the results to a CSV format." + nl +
       "Has the following scenarios to choose: --to-datetime, --by-period, --to-earliest, --to-latest, --shift-by, --from-file, --to-current. One scenario must be choose" + nl +
       "To define the scope use: --all-topics or --topic. . One scope must be choose, unless you use '--from-file' scenario"
-    val ExecuteDoc = "Execute operation. Supported operations: reset-offsets."
+    val DryRunDoc = "Only show results without executing changes on Consumer Groups. Supported operations: reset-offsets."
     val ExportDoc = "Export operation execution to a CSV file. Supported operations: reset-offsets."
     val ResetToOffsetDoc = "Reset offsets to a specific offset."
     val ResetFromFileDoc = "Reset offsets to values defined in CSV file."
@@ -770,7 +770,7 @@ object ConsumerGroupCommand extends Logging {
                                   .describedAs("command config property file")
                                   .ofType(classOf[String])
     val resetOffsetsOpt = parser.accepts("reset-offsets", ResetOffsetsDoc)
-    val executeOpt = parser.accepts("execute", ExecuteDoc)
+    val dryRunOpt = parser.accepts("dry-run", DryRunDoc)
     val exportOpt = parser.accepts("export", ExportDoc)
     val resetToOffsetOpt = parser.accepts("to-offset", ResetToOffsetDoc)
                            .withRequiredArg()

core/src/main/scala/kafka/tools/StreamsResetter.java

@@ -21,20 +21,27 @@
 import joptsimple.OptionSet;
 import joptsimple.OptionSpec;
 import joptsimple.OptionSpecBuilder;
+import kafka.utils.CommandLineUtils;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.DeleteTopicsResult;
 import org.apache.kafka.clients.admin.KafkaAdminClient;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.annotation.InterfaceStability;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
 
+import javax.xml.datatype.DatatypeFactory;
+import javax.xml.datatype.Duration;
 import java.io.IOException;
+import java.text.ParseException;
+import java.text.SimpleDateFormat;
 import java.util.ArrayList;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
@@ -74,13 +81,21 @@ public class StreamsResetter {
     private static OptionSpec<String> applicationIdOption;
     private static OptionSpec<String> inputTopicsOption;
     private static OptionSpec<String> intermediateTopicsOption;
+    private static OptionSpec<Long> toOffsetOption;
+    private static OptionSpec<String> toDatetimeOption;
+    private static OptionSpec<String> byDurationOption;
+    private static OptionSpecBuilder toEarliestOption;
+    private static OptionSpecBuilder toLatestOption;
+    private static OptionSpec<String> fromFileOption;
+    private static OptionSpec<Long> shiftByOption;
     private static OptionSpecBuilder dryRunOption;
     private static OptionSpec<String> commandConfigOption;
 
     private OptionSet options = null;
     private final List<String> allTopics = new LinkedList<>();
     private boolean dryRun = false;
 
+
     public int run(final String[] args) {
         return run(args, new Properties());
     }
@@ -169,6 +184,24 @@ private void parseArguments(final String[] args) throws IOException {
             .ofType(String.class)
             .withValuesSeparatedBy(',')
             .describedAs("list");
+        toOffsetOption = optionParser.accepts("to-offset", "Reset offsets to a specific offset.")
+            .withRequiredArg()
+            .ofType(Long.class);
+        toDatetimeOption = optionParser.accepts("to-datetime", "Reset offsets to offset from datetime. Format: 'YYYY-MM-DDTHH:mm:SS.sss'")
+            .withRequiredArg()
+            .ofType(String.class);
+        byDurationOption = optionParser.accepts("by-duration", "Reset offsets to offset by duration from current timestamp. Format: 'PnDTnHnMnS'")
+            .withRequiredArg()
+            .ofType(String.class);
+        toEarliestOption = optionParser.accepts("to-earliest", "Reset offsets to earliest offset.");
+        toLatestOption = optionParser.accepts("to-latest", "Reset offsets to latest offset.");
+        fromFileOption = optionParser.accepts("from-file", "Reset offsets to values defined in CSV file.")
+            .withRequiredArg()
+            .ofType(String.class);
+        shiftByOption = optionParser.accepts("shift-by", "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative")
+            .withRequiredArg()
+            .describedAs("number-of-offsets")
+            .ofType(Long.class);
         commandConfigOption = optionParser.accepts("config-file", "Property file containing configs to be passed to admin clients and embedded consumer.")
             .withRequiredArg()
             .ofType(String.class)
@@ -184,9 +217,26 @@ private void parseArguments(final String[] args) throws IOException {
             printHelp(optionParser);
             throw e;
         }
+
+        scala.collection.immutable.HashSet<OptionSpec<?>> allScenarioOptions = new scala.collection.immutable.HashSet<>();
+        allScenarioOptions.$plus(toOffsetOption);
+        allScenarioOptions.$plus(toDatetimeOption);
+        allScenarioOptions.$plus(byDurationOption);
+        allScenarioOptions.$plus(toEarliestOption);
+        allScenarioOptions.$plus(toLatestOption);
+        allScenarioOptions.$plus(fromFileOption);
+        allScenarioOptions.$plus(shiftByOption);
+
+        CommandLineUtils.checkInvalidArgs(optionParser, options, toOffsetOption, allScenarioOptions.$minus(toOffsetOption));
+        CommandLineUtils.checkInvalidArgs(optionParser, options, toDatetimeOption, allScenarioOptions.$minus(toDatetimeOption));
+        CommandLineUtils.checkInvalidArgs(optionParser, options, byDurationOption, allScenarioOptions.$minus(byDurationOption));
+        CommandLineUtils.checkInvalidArgs(optionParser, options, toEarliestOption, allScenarioOptions.$minus(toEarliestOption));
+        CommandLineUtils.checkInvalidArgs(optionParser, options, toLatestOption, allScenarioOptions.$minus(toLatestOption));
+        CommandLineUtils.checkInvalidArgs(optionParser, options, fromFileOption, allScenarioOptions.$minus(fromFileOption));
+        CommandLineUtils.checkInvalidArgs(optionParser, options, shiftByOption, allScenarioOptions.$minus(shiftByOption));
     }
 
-    private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consumerConfig) {
+    private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consumerConfig) throws Exception {
         final List<String> inputTopics = options.valuesOf(inputTopicsOption);
         final List<String> intermediateTopics = options.valuesOf(intermediateTopicsOption);
 
@@ -202,7 +252,7 @@ private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consu
         }
 
         if (inputTopics.size() != 0) {
-            System.out.println("Seek-to-beginning for input topics " + inputTopics);
+            System.out.println("Reset-offsets for input topics " + inputTopics);
         }
         if (intermediateTopics.size() != 0) {
             System.out.println("Seek-to-end for intermediate topics " + intermediateTopics);
@@ -249,7 +299,7 @@ private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consu
                 }
             }
 
-            maybeSeekToBeginning(client, inputTopicPartitions);
+            maybeReset(client, inputTopicPartitions);
 
             maybeSeekToEnd(client, intermediateTopicPartitions);
 
@@ -274,7 +324,7 @@ private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consu
                 }
             }
 
-        } catch (final RuntimeException e) {
+        } catch (final Exception e) {
             System.err.println("ERROR: Resetting offsets failed.");
             throw e;
         }
@@ -300,23 +350,160 @@ private void maybeSeekToEnd(final KafkaConsumer<byte[], byte[]> client,
         }
     }
 
-    private void maybeSeekToBeginning(final KafkaConsumer<byte[], byte[]> client,
-                                      final Set<TopicPartition> inputTopicPartitions) {
+    private void maybeReset(final KafkaConsumer<byte[], byte[]> client,
+                            final Set<TopicPartition> inputTopicPartitions)
+        throws Exception {
 
-        final List<String> inputTopics = options.valuesOf(inputTopicsOption);
         final String groupId = options.valueOf(applicationIdOption);
 
         if (inputTopicPartitions.size() > 0) {
-            System.out.println("Following input topics offsets will be reset to beginning (for consumer group " + groupId + ")");
-            for (final String topic : inputTopics) {
-                if (allTopics.contains(topic)) {
-                    System.out.println("Topic: " + topic);
+            System.out.println("Following input topics offsets will be reset to (for consumer group " + groupId + ")");
+            if (options.has(toOffsetOption)) {
+                final Long offset = options.valueOf(toOffsetOption);
+                final Map<TopicPartition, Long> endOffsets = client.endOffsets(inputTopicPartitions);
+                final Map<TopicPartition, Long> beginningOffsets = client.beginningOffsets(inputTopicPartitions);
+
+                final Map<TopicPartition, Long> topicPartitionsAndOffset = new HashMap<>(inputTopicPartitions.size());
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    topicPartitionsAndOffset.put(topicPartition, offset);
                 }
-            }
-            if (!dryRun) {
+
+                final Map<TopicPartition, Long> validatedTopicPartitionsAndOffset =
+                    checkOffsetRange(topicPartitionsAndOffset, beginningOffsets, endOffsets);
+
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    client.seek(topicPartition, validatedTopicPartitionsAndOffset.get(topicPartition));
+                }
+            } else if (options.has(toEarliestOption)) {
                 client.seekToBeginning(inputTopicPartitions);
+            } else if (options.has(toLatestOption)) {
+                client.seekToEnd(inputTopicPartitions);
+            } else if (options.has(shiftByOption)) {
+                final Long shiftBy = options.valueOf(shiftByOption);
+                final Map<TopicPartition, Long> endOffsets = client.endOffsets(inputTopicPartitions);
+                final Map<TopicPartition, Long> beginningOffsets = client.beginningOffsets(inputTopicPartitions);
+
+                final Map<TopicPartition, Long> topicPartitionsAndOffset = new HashMap<>(inputTopicPartitions.size());
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    final Long position = client.position(topicPartition);
+                    final Long offset = position + shiftBy;
+                    topicPartitionsAndOffset.put(topicPartition, offset);
+                }
+
+                final Map<TopicPartition, Long> validatedTopicPartitionsAndOffset =
+                    checkOffsetRange(topicPartitionsAndOffset, beginningOffsets, endOffsets);
+
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    client.seek(topicPartition, validatedTopicPartitionsAndOffset.get(topicPartition));
+                }
+            } else if (options.has(toDatetimeOption)) {
+                final String ts = options.valueOf(toDatetimeOption);
+                final Long timestamp = getDateTime(ts);
+
+                final Map<TopicPartition, Long> topicPartitionsAndTimes = new HashMap<>(inputTopicPartitions.size());
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    topicPartitionsAndTimes.put(topicPartition, timestamp);
+                }
+
+                final Map<TopicPartition, OffsetAndTimestamp> topicPartitionsAndOffset = client.offsetsForTimes(topicPartitionsAndTimes);
+
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    final Long offset = topicPartitionsAndOffset.get(topicPartition).offset();
+                    client.seek(topicPartition, offset);
+                }
+            } else if (options.has(byDurationOption)) {
+                final String duration = options.valueOf(byDurationOption);
+                final Date now = new Date();
+                final Duration durationParsed = DatatypeFactory.newInstance().newDuration(duration);
+                durationParsed.negate().addTo(now);
+                final Long timestamp = now.getTime();
+
+                final Map<TopicPartition, Long> topicPartitionsAndTimes = new HashMap<>(inputTopicPartitions.size());
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    topicPartitionsAndTimes.put(topicPartition, timestamp);
+                }
+
+                final Map<TopicPartition, OffsetAndTimestamp> topicPartitionsAndOffset = client.offsetsForTimes(topicPartitionsAndTimes);
+
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    final Long offset = topicPartitionsAndOffset.get(topicPartition).offset();
+                    client.seek(topicPartition, offset);
+                }
+            } else if (options.has(fromFileOption)) {
+                final Map<TopicPartition, Long> endOffsets = client.endOffsets(inputTopicPartitions);
+                final Map<TopicPartition, Long> beginningOffsets = client.beginningOffsets(inputTopicPartitions);
+
+                final String resetPlanPath = options.valueOf(fromFileOption);
+                final String resetPlanCsv = Utils.readFileAsString(resetPlanPath);
+                final Map<TopicPartition, Long> topicPartitionsAndOffset = parseResetPlan(resetPlanCsv);
+
+                final Map<TopicPartition, Long> validatedTopicPartitionsAndOffset =
+                    checkOffsetRange(topicPartitionsAndOffset, beginningOffsets, endOffsets);
+
+                for (final TopicPartition topicPartition : inputTopicPartitions) {
+                    final Long offset = validatedTopicPartitionsAndOffset.get(topicPartition);
+                    client.seek(topicPartition, offset);
+                }
+            } else {
+                client.seekToBeginning(inputTopicPartitions);
+            }
+
+            for (final TopicPartition p : inputTopicPartitions) {
+                final Long position = client.position(p);
+                System.out.println("Topic: " + p.topic() + " Partition: " + p.partition() + " Offset: " + position);
+            }
+        }
+    }
+
+    public Long getDateTime(String ts) throws ParseException {
+        if (!(ts.split("T")[1].contains("+") || ts.split("T")[1].contains("-") || ts.split("T")[1].contains("Z"))) {
+            ts = ts + "Z";
+        }
+
+        Date date;
+        try {
+            date = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX").parse(ts);
+        } catch (ParseException e) {
+            date = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSX").parse(ts);
+        }
+
+        return date.getTime();
+    }
+
+    private Map<TopicPartition, Long> parseResetPlan(String resetPlanCsv) {
+        final Map<TopicPartition, Long> topicPartitionAndOffset = new HashMap<>();
+        for (final String line : resetPlanCsv.split("\n")) {
+            final String[] parts = line.split(",");
+            final String topic = parts[0];
+            final int partition = Integer.parseInt(parts[1]);
+            final long offset = Long.parseLong(parts[2]);
+            final TopicPartition topicPartition = new TopicPartition(topic, partition);
+            topicPartitionAndOffset.put(topicPartition, offset);
+        }
+        return topicPartitionAndOffset;
+    }
+
+    private Map<TopicPartition, Long> checkOffsetRange(Map<TopicPartition, Long> inputTopicPartitionsAndOffset,
+                                                       Map<TopicPartition, Long> beginningOffsets,
+                                                       Map<TopicPartition, Long> endOffsets) {
+        Map<TopicPartition, Long> validatedTopicPartitionsOffsets = new HashMap<>();
+        for (final TopicPartition topicPartition : inputTopicPartitionsAndOffset.keySet()) {
+            final Long endOffset = endOffsets.get(topicPartition);
+            final Long offset = inputTopicPartitionsAndOffset.get(topicPartition);
+            if (offset < endOffset) {
+                final Long beginningOffset = beginningOffsets.get(topicPartition);
+                if (offset > beginningOffset) {
+                    validatedTopicPartitionsOffsets.put(topicPartition, offset);
+                } else {
+                    System.out.println("New offset (" + offset + ") is lower than earliest offset. Value will be set to " + beginningOffset);
+                    validatedTopicPartitionsOffsets.put(topicPartition, beginningOffset);
+                }
+            } else {
+                System.out.println("New offset (" + offset + ") is higher than latest offset. Value will be set to " + endOffset);
+                validatedTopicPartitionsOffsets.put(topicPartition, endOffset);
             }
         }
+        return validatedTopicPartitionsOffsets;
     }
 
     private boolean isInputTopic(final String topic) {