Adding retries to update the metadata store instead of failure

extensions-contrib/materialized-view-maintenance/src/main/java/org/apache/druid/indexing/materializedview/DerivativeDataSourceMetadata.java

@@ -85,6 +85,12 @@ public boolean matches(DataSourceMetadata other)
     return equals(other);
   }
 
+  @Override
+  public int compareTo(DataSourceMetadata o)
+  {
+    return 0;
+  }
+
   @Override
   public DataSourceMetadata plus(DataSourceMetadata other)
   {

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadata.java

@@ -26,8 +26,11 @@
 import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata;
 import org.apache.druid.indexing.seekablestream.SeekableStreamEndSequenceNumbers;
 import org.apache.druid.indexing.seekablestream.SeekableStreamSequenceNumbers;
+import org.apache.druid.java.util.common.IAE;
+
+import java.util.Comparator;
 
 public class KafkaDataSourceMetadata extends SeekableStreamDataSourceMetadata<KafkaTopicPartition, Long>
 {
 
   @JsonCreator
@@ -58,4 +61,19 @@
   {
     return new KafkaDataSourceMetadata(seekableStreamSequenceNumbers);
   }
+
+  @Override
+  public int compareTo(DataSourceMetadata other)
+  {
+    if (!getClass().equals(other.getClass())) {
+      throw new IAE(
+          "Expected instance of %s, got %s",
+          this.getClass().getName(),
+          other.getClass().getName()
+      );
+    }
+    final SeekableStreamDataSourceMetadata<KafkaTopicPartition, Long> that = (SeekableStreamDataSourceMetadata<KafkaTopicPartition, Long>) other;
+
+    return getSeekableStreamSequenceNumbers().compareTo(that.getSeekableStreamSequenceNumbers(), Comparator.naturalOrder());
+  }
 }

extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadata.java

@@ -26,7 +26,7 @@
 import org.apache.druid.indexing.seekablestream.SeekableStreamEndSequenceNumbers;
 import org.apache.druid.indexing.seekablestream.SeekableStreamSequenceNumbers;
 
 public class KinesisDataSourceMetadata extends SeekableStreamDataSourceMetadata<String, String>
 {
   @JsonCreator
   public KinesisDataSourceMetadata(
@@ -56,4 +56,9 @@
   {
     return new KinesisDataSourceMetadata(seekableStreamSequenceNumbers);
   }
+
+  @Override
+  public int compareTo(DataSourceMetadata o) {
+    return 0;
+  }
 }

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamEndSequenceNumbers.java

@@ -25,10 +25,12 @@
 import org.apache.druid.java.util.common.IAE;
 
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Represents the end sequenceNumber per partition of a sequence. Note that end sequenceNumbers are always
@@ -147,6 +149,37 @@ public SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> plus(
     }
   }
 
+  @Override
+  public int compareTo(SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> other, Comparator<SequenceOffsetType> comparator)
+  {
+    if (this.getClass() != other.getClass()) {
+      throw new IAE(
+          "Expected instance of %s, got %s",
+          this.getClass().getName(),
+          other.getClass().getName()
+      );
+    }
+
+    final SeekableStreamEndSequenceNumbers<PartitionIdType, SequenceOffsetType> otherStart =
+        (SeekableStreamEndSequenceNumbers<PartitionIdType, SequenceOffsetType>) other;
+
+    if (stream.equals(otherStart.stream)) {
+      //Same stream, compare the offset
+      AtomicReference<Boolean> res = new AtomicReference<>(false);
+      partitionSequenceNumberMap.forEach(
+          (partitionId, sequenceOffset) -> {
+            if (otherStart.partitionSequenceNumberMap.get(partitionId) != null && comparator.compare(sequenceOffset, otherStart.partitionSequenceNumberMap.get(partitionId)) > 0) {
+              res.set(true);
+            }
+          }
+      );
+      if (res.get()) {
+        return 1;
+      }
+    }
+    return 0;
+  }
+
   @Override
   public SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> minus(
       SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> other

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamSequenceNumbers.java

@@ -25,6 +25,7 @@
 import com.fasterxml.jackson.annotation.JsonTypeInfo.Id;
 import org.apache.druid.indexing.overlord.DataSourceMetadata;
 
+import java.util.Comparator;
 import java.util.Map;
 
 @JsonTypeInfo(use = Id.NAME, property = "type", defaultImpl = SeekableStreamEndSequenceNumbers.class)
@@ -61,4 +62,11 @@ SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> plus(
   SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> minus(
       SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> other
   );
+
+  /**
+   * Compare this and the other sequence offsets using comparator.
+   * Returns 1, if this sequence is ahead of the other.
+   * otherwise, Return 0
+   */
+  int compareTo(SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> seekableStreamSequenceNumbers, Comparator<SequenceOffsetType> comparator);
 }

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamStartSequenceNumbers.java

@@ -26,12 +26,14 @@
 
 import javax.annotation.Nullable;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Represents the start sequenceNumber per partition of a sequence. This class keeps an additional set of
@@ -161,6 +163,37 @@ public SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> plus(
     }
   }
 
+  @Override
+  public int compareTo(SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> other, Comparator<SequenceOffsetType> comparator)
+  {
+    if (this.getClass() != other.getClass()) {
+      throw new IAE(
+          "Expected instance of %s, got %s",
+          this.getClass().getName(),
+          other.getClass().getName()
+      );
+    }
+
+    final SeekableStreamStartSequenceNumbers<PartitionIdType, SequenceOffsetType> otherStart =
+        (SeekableStreamStartSequenceNumbers<PartitionIdType, SequenceOffsetType>) other;
+
+    if (stream.equals(otherStart.stream)) {
+      //Same stream, compare the offset
+      AtomicReference<Boolean> res = new AtomicReference<>(false);
+      partitionSequenceNumberMap.forEach(
+          (partitionId, sequenceOffset) -> {
+            if (otherStart.partitionSequenceNumberMap.get(partitionId) != null && comparator.compare(sequenceOffset, otherStart.partitionSequenceNumberMap.get(partitionId)) > 0) {
+              res.set(true);
+            }
+          }
+      );
+      if (res.get()) {
+        return 1;
+      }
+    }
+    return 0;
+  }
+
   @Override
   public SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> minus(
       SeekableStreamSequenceNumbers<PartitionIdType, SequenceOffsetType> other

indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertActionTest.java

@@ -29,17 +29,24 @@
 import org.apache.druid.indexing.overlord.ObjectMetadata;
 import org.apache.druid.indexing.overlord.SegmentPublishResult;
 import org.apache.druid.indexing.overlord.Segments;
+import org.apache.druid.indexing.overlord.TaskLockbox;
 import org.apache.druid.indexing.overlord.TimeChunkLockRequest;
+import org.apache.druid.indexing.overlord.config.TaskLockConfig;
+import org.apache.druid.indexing.overlord.supervisor.SupervisorManager;
 import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.concurrent.ScheduledExecutors;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.partition.LinearShardSpec;
 import org.assertj.core.api.Assertions;
+import org.easymock.EasyMock;
 import org.hamcrest.CoreMatchers;
 import org.joda.time.Interval;
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 
 public class SegmentTransactionalInsertActionTest
 {
@@ -134,6 +141,105 @@ public void testTransactionalUpdateDataSourceMetadata() throws Exception
     );
   }
 
+  @Test
+  public void testTransactionalUpdateDataSourceMetadataWithRecoverFromMetadataMismatch() throws Exception
+  {
+    final Task task1 = NoopTask.create();
+    final TaskLockbox taskLockbox1 = new TaskLockbox(actionTestKit.getTaskStorage(), actionTestKit.getMetadataStorageCoordinator());
+    taskLockbox1.add(task1);
+    taskLockbox1.lock(task1, new TimeChunkLockRequest(TaskLockType.EXCLUSIVE, task1, INTERVAL, null), 5000);
+
+    final Task task2 = NoopTask.create();
+    final TaskLockbox taskLockbox2 = new TaskLockbox(actionTestKit.getTaskStorage(), actionTestKit.getMetadataStorageCoordinator());
+    taskLockbox2.add(task2);
+    taskLockbox2.lock(task2, new TimeChunkLockRequest(TaskLockType.EXCLUSIVE, task2, INTERVAL, null), 5000);
+
+    final TaskLockConfig taskLockConfig = new TaskLockConfig()
+    {
+      @Override
+      public boolean isBatchSegmentAllocation()
+      {
+        return true;
+      }
+
+      @Override
+      public long getBatchAllocationWaitTime()
+      {
+        return 10L;
+      }
+    };
+    TaskActionToolbox taskActionToolbox = actionTestKit.getTaskActionToolbox();
+
+    // Task1 and Task2 tries to publish segment1 and segment2 for same partition at around same time.
+    // With different start and end offsets. Segment2 -> {1 - 2}, Segment1 -> {null - 1}
+    Future<SegmentPublishResult> result2Future = CompletableFuture.supplyAsync(() -> {
+      return SegmentTransactionalInsertAction.appendAction(
+          ImmutableSet.of(SEGMENT2),
+          new ObjectMetadata(ImmutableList.of(1)),
+          new ObjectMetadata(ImmutableList.of(2))
+      ).perform(
+          task2,
+          new TaskActionToolbox(
+              taskLockbox2,
+              taskActionToolbox.getTaskStorage(),
+              taskActionToolbox.getIndexerMetadataStorageCoordinator(),
+              new SegmentAllocationQueue(
+                  taskLockbox2,
+                  taskLockConfig,
+                  taskActionToolbox.getIndexerMetadataStorageCoordinator(),
+                  taskActionToolbox.getEmitter(),
+                  ScheduledExecutors::fixed
+              ),
+              taskActionToolbox.getEmitter(),
+              EasyMock.createMock(SupervisorManager.class),
+              taskActionToolbox.getJsonMapper()
+          )
+      );
+    });
+
+
+    Future<SegmentPublishResult> result1Future = CompletableFuture.supplyAsync(() -> {
+      return SegmentTransactionalInsertAction.appendAction(
+          ImmutableSet.of(SEGMENT1),
+          new ObjectMetadata(null),
+          new ObjectMetadata(ImmutableList.of(1))
+      ).perform(
+          task1,
+          new TaskActionToolbox(
+              taskLockbox1,
+              taskActionToolbox.getTaskStorage(),
+              taskActionToolbox.getIndexerMetadataStorageCoordinator(),
+              new SegmentAllocationQueue(
+                  taskLockbox1,
+                  taskLockConfig,
+                  taskActionToolbox.getIndexerMetadataStorageCoordinator(),
+                  taskActionToolbox.getEmitter(),
+                  ScheduledExecutors::fixed
+              ),
+              taskActionToolbox.getEmitter(),
+              EasyMock.createMock(SupervisorManager.class),
+              taskActionToolbox.getJsonMapper()
+          )
+      );
+    });
+
+    SegmentPublishResult result2 = result2Future.get();
+    SegmentPublishResult result1 = result1Future.get();
+
+    Assert.assertEquals(SegmentPublishResult.ok(ImmutableSet.of(SEGMENT1)), result1);
+    Assert.assertEquals(SegmentPublishResult.ok(ImmutableSet.of(SEGMENT2)), result2);
+
+    Assertions.assertThat(
+        actionTestKit.getMetadataStorageCoordinator()
+            .retrieveUsedSegmentsForInterval(DATA_SOURCE, INTERVAL, Segments.ONLY_VISIBLE)
+    ).containsExactlyInAnyOrder(SEGMENT1, SEGMENT2);
+
+    Assert.assertEquals(
+        new ObjectMetadata(ImmutableList.of(2)),
+        actionTestKit.getMetadataStorageCoordinator().retrieveDataSourceMetadata(DATA_SOURCE)
+    );
+  }
+
   @Test
   public void testFailTransactionalUpdateDataSourceMetadata() throws Exception
   {
@@ -152,9 +258,8 @@ public void testFailTransactionalUpdateDataSourceMetadata() throws Exception
 
     Assert.assertEquals(
         SegmentPublishResult.fail(
-          "java.lang.RuntimeException: Inconsistent metadata state. " +
-          "This can happen if you update input topic in a spec without changing the supervisor name. " +
-          "Stored state: [null], Target state: [ObjectMetadata{theObject=[1]}]."
+          "org.apache.druid.metadata.RetryTransactionException: Failed to update the metadata Store. " +
+          "The new start metadata: [ObjectMetadata{theObject=[1]}] is ahead of last commited end state: [null]."
         ),
         result
     );

indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionTestKit.java

@@ -59,6 +59,11 @@ public TaskLockbox getTaskLockbox()
     return taskLockbox;
   }
 
+  public TaskStorage getTaskStorage()
+  {
+    return taskStorage;
+  }
+
   public IndexerMetadataStorageCoordinator getMetadataStorageCoordinator()
   {
     return metadataStorageCoordinator;

indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/TestSeekableStreamDataSourceMetadata.java

@@ -23,7 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.indexing.overlord.DataSourceMetadata;
 
 public class TestSeekableStreamDataSourceMetadata extends SeekableStreamDataSourceMetadata<String, String>
 {
   @JsonCreator
   public TestSeekableStreamDataSourceMetadata(
@@ -45,4 +45,10 @@
   {
     return null;
   }
+
+  @Override
+  public int compareTo(DataSourceMetadata o)
+  {
+    return 0;
+  }
 }

indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java

@@ -2414,7 +2414,7 @@
     }
   }
 
   private static class TestSeekableStreamDataSourceMetadata extends SeekableStreamDataSourceMetadata<String, String>
   {
 
     @JsonCreator
@@ -2445,6 +2445,12 @@
     {
       return new TestSeekableStreamDataSourceMetadata(seekableStreamSequenceNumbers);
     }
+
+    @Override
+    public int compareTo(DataSourceMetadata o)
+    {
+      return 0;
+    }
   }
 
   private static SeekableStreamIndexTaskIOConfig createTaskIoConfigExt(

server/src/main/java/org/apache/druid/indexing/overlord/DataSourceMetadata.java

@@ -36,7 +36,7 @@
 @JsonSubTypes(value = {
     @JsonSubTypes.Type(name = "object", value = ObjectMetadata.class)
 })
-public interface DataSourceMetadata
+public interface DataSourceMetadata extends Comparable<DataSourceMetadata>
 {
   /**
    * Returns true if this instance should be considered a valid starting point for a new dataSource that has

server/src/main/java/org/apache/druid/indexing/overlord/ObjectMetadata.java

@@ -60,6 +60,12 @@ public boolean matches(DataSourceMetadata other)
     return equals(other);
   }
 
+  @Override
+  public int compareTo(DataSourceMetadata o)
+  {
+    return 0;
+  }
+
   @Override
   public DataSourceMetadata plus(DataSourceMetadata other)
   {