extend tests

Author: Jamie Chapman-Brown

docs/development/extensions-contrib/rabbit-stream-ingestion.md

@@ -49,7 +49,7 @@
 import java.util.TreeMap;
 
 /**
- * Kafka indexing task runner supporting incremental segments publishing
+ * RabbitStream indexing task runner supporting incremental segments publishing
  */
 public class IncrementalPublishingRabbitStreamIndexTaskRunner
     extends SeekableStreamIndexTaskRunner<String, Long, ByteEntity>
@@ -95,7 +95,7 @@ protected SeekableStreamEndSequenceNumbers<String, Long> deserializePartitionsFr
         SeekableStreamEndSequenceNumbers.class,
         SeekableStreamEndSequenceNumbers.class,
         String.class,
-        String.class));
+        Long.class));
   }
 
   @Override

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/IncrementalPublishingRabbitStreamIndexTaskRunner.java

@@ -32,6 +32,9 @@
 
 public class RabbitStreamIndexTaskModule implements DruidModule
 {
+
+  static final String PROPERTY_BASE = "druid.rabbit";
+
   @Override
   public List<? extends Module> getJacksonModules()
   {

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskModule.java

@@ -42,13 +42,13 @@ public class RabbitStreamIndexTaskTuningConfig extends SeekableStreamIndexTaskTu
   static final int ASSUMED_RECORD_SIZE = 10_000;
 
   /**
-   * Together with {@link KinesisIndexTaskIOConfig#RECORD_FETCH_MEMORY_MAX_HEAP_FRACTION}, don't take up more
+   * Together with {@link RabbitStreamIndexTaskIOConfig#RECORD_FETCH_MEMORY_MAX_HEAP_FRACTION}, don't take up more
    * than 15% of the heap per task.
    */
   private static final double RECORD_BUFFER_MEMORY_MAX_HEAP_FRACTION = 0.1;
 
   /**
-   * Together with {@link KinesisIndexTaskIOConfig#MAX_RECORD_FETCH_MEMORY}, don't take up more than 200MB per task.
+   * Together with {@link RabbitStreamIndexTaskIOConfig#MAX_RECORD_FETCH_MEMORY}, don't take up more than 200MB per task.
    */
   private static final int MAX_RECORD_BUFFER_MEMORY = 100_000_000;
 

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskTuningConfig.java

@@ -20,6 +20,7 @@
 package org.apache.druid.indexing.rabbitstream;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Queues;
@@ -32,8 +33,6 @@
 import com.rabbitmq.stream.OffsetSpecification;
 import com.rabbitmq.stream.impl.Client;
 import com.rabbitmq.stream.impl.Client.ClientParameters;
-import io.netty.handler.ssl.SslContext;
-import io.netty.handler.ssl.SslContextBuilder;
 import org.apache.druid.data.input.impl.ByteEntity;
 import org.apache.druid.indexing.rabbitstream.supervisor.RabbitStreamSupervisorIOConfig;
 import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord;
@@ -45,7 +44,6 @@
 
 import javax.annotation.Nonnull;
 
-import java.io.FileInputStream;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
@@ -86,18 +84,17 @@ public class RabbitStreamRecordSupplier implements RecordSupplier<String, Long,
 
   private String password;
   private String username;
-  private String certificateFilePath;
 
   public RabbitStreamRecordSupplier(
       Map<String, Object> consumerProperties,
       ObjectMapper mapper,
       String uri,
       int recordBufferSize,
       int recordBufferOfferTimeout,
-      int maxRecordsPerPoll)
+      int maxRecordsPerPoll
+  )
   {
     this.uri = uri;
-    this.env = this.getRabbitEnvironment();
     this.mapper = mapper;
 
     this.recordBufferSize = recordBufferSize;
@@ -118,7 +115,6 @@ public RabbitStreamRecordSupplier(
 
     this.password = null;
     this.username = null;
-    this.certificateFilePath = null;
 
     if (consumerProperties != null) {
 
@@ -137,38 +133,47 @@ public RabbitStreamRecordSupplier(
           if (e.getKey().equals("username")) {
             this.username = e.getValue();
           }
-          if (e.getKey().equals("cerfiticateFilePath")) {
-            this.certificateFilePath = e.getValue();
-          }
         }
       }
     }
+
+    this.env = this.getRabbitEnvironment();
+
   }
 
   public void startBackgroundFetch()
   {
-
     try {
-      stateSemaphore.acquire();
-      try {
-        if (this.isRunning != false) {
-          return;
-        }
-        for (Map.Entry<String, ConsumerBuilder> entry : streamBuilders.entrySet()) {
-          consumers.add(
-              entry.getValue().offset(
-                  offsetMap.get(entry.getKey())).build());
-        }
-        this.isRunning = true;
+      // aquire uninteruptibly to prevent state corruption issues
+      // on consumers and isRunning
+      stateSemaphore.acquireUninterruptibly();
+      if (this.isRunning != false) {
+        return;
       }
-      finally {
-        stateSemaphore.release();
+      for (Map.Entry<String, ConsumerBuilder> entry : streamBuilders.entrySet()) {
+        consumers.add(
+            entry.getValue().offset(
+                offsetMap.get(entry.getKey())).build());
       }
+      this.isRunning = true;
     }
-    catch (InterruptedException exc) {
+    finally {
+      stateSemaphore.release();
     }
   }
 
+  @VisibleForTesting
+  public int bufferSize()
+  {
+    return queue.size();
+  }
+
+  @VisibleForTesting
+  public boolean isRunning()
+  {
+    return this.isRunning;
+  }
+
   public void stopBackgroundFetch()
   {
     try {
@@ -192,9 +197,14 @@ public void stopBackgroundFetch()
 
   }
 
+  public EnvironmentBuilder getEnvBuilder()
+  {
+    return Environment.builder();
+  }
+
   public Environment getRabbitEnvironment()
   {
-    EnvironmentBuilder envBuilder = Environment.builder().uri(this.uri);
+    EnvironmentBuilder envBuilder = this.getEnvBuilder().uri(this.uri);
 
     if (this.password != null) {
       envBuilder = envBuilder.password(this.password);
@@ -203,22 +213,6 @@ public Environment getRabbitEnvironment()
       envBuilder = envBuilder.username(this.username);
     }
 
-    if (this.certificateFilePath != null) {
-      try {
-        FileInputStream inputStream =
-            new FileInputStream(this.certificateFilePath);
-        SslContext sslContext = SslContextBuilder
-            .forClient()
-            .trustManager(inputStream)  
-            .build();
-        envBuilder = envBuilder.tls().sslContext(sslContext).environmentBuilder();
-      }
-      catch (Exception exc) {
-        log.error("failed to configure ssl for cockroach");
-        log.error(exc, exc.getMessage());
-      }
-    }
-
     return envBuilder.build();
   }
 
@@ -229,6 +223,27 @@ public static String getStreamFromSubstream(String partionID)
     return String.join("-", res);
   }
 
+  private void removeOldAssignments(Set<StreamPartition<String>> streamPartitionstoKeep)
+  {
+    Iterator<Map.Entry<String, ConsumerBuilder>> streamBuilderIterator = streamBuilders.entrySet().iterator();
+    while (streamBuilderIterator.hasNext()) {
+      Map.Entry<String, ConsumerBuilder> entry = streamBuilderIterator.next();
+      StreamPartition<String> comparitor = new StreamPartition<String>(getStreamFromSubstream(entry.getKey()), entry.getKey());
+      if (!streamPartitionstoKeep.contains(comparitor)) {
+        streamBuilderIterator.remove();
+      }
+    }
+
+    Iterator<Map.Entry<String, OffsetSpecification>> offsetItterator = offsetMap.entrySet().iterator();
+    while (offsetItterator.hasNext()) {
+      Map.Entry<String, OffsetSpecification> entry = offsetItterator.next();
+      StreamPartition<String> comparitor = new StreamPartition<String>(getStreamFromSubstream(entry.getKey()), entry.getKey());
+      if (!streamPartitionstoKeep.contains(comparitor)) {
+        offsetItterator.remove();
+      }
+    }
+  }
+
   @Override
   public void assign(Set<StreamPartition<String>> streamPartitions)
   {
@@ -242,24 +257,7 @@ public void assign(Set<StreamPartition<String>> streamPartitions)
       this.superStream = part.getStream();
     }
 
-    Iterator<Map.Entry<String, ConsumerBuilder>> streamBuilderItterator = streamBuilders.entrySet().iterator();
-    while (streamBuilderItterator.hasNext()) {
-      Map.Entry<String, ConsumerBuilder> entry = streamBuilderItterator.next();
-      // I don't think this works....
-      StreamPartition<String> comparitor = new StreamPartition<String>(getStreamFromSubstream(entry.getKey()), entry.getKey());
-      if (!streamPartitions.contains(comparitor)) {
-        streamBuilderItterator.remove();
-      }
-    }
-
-    Iterator<Map.Entry<String, OffsetSpecification>> offsetItterator = offsetMap.entrySet().iterator();
-    while (offsetItterator.hasNext()) {
-      Map.Entry<String, OffsetSpecification> entry = offsetItterator.next();
-      StreamPartition<String> comparitor = new StreamPartition<String>(getStreamFromSubstream(entry.getKey()), entry.getKey());
-      if (!streamPartitions.contains(comparitor)) {
-        offsetItterator.remove();
-      }
-    }
+    removeOldAssignments(streamPartitions);
 
 
   }
@@ -346,9 +344,14 @@ public void handle(MessageHandler.Context context, Message message)
 
   }
 
+  /**
+   * optionalStartBackgroundFetch ensures that a background fetch is running
+   * if this.queue is running low on records. We want to minimize thrashing
+   * around starting/stopping the consumers.
+   */
   public void optionalStartBackgroundFetch()
   {
-    if (this.queue.size() < this.maxRecordsPerPoll) {
+    if (this.queue.size() < Math.min(this.maxRecordsPerPoll * 2, recordBufferSize / 2)) {
       this.startBackgroundFetch();
     }
   }
@@ -410,10 +413,21 @@ public Long getPosition(StreamPartition<String> partition)
     throw new UnsupportedOperationException("getPosition() is not supported in RabbitMQ streams");
   }
 
+
+  public ClientParameters getParameters()
+  {
+    return new ClientParameters();
+  }
+
+  public Client getClient(ClientParameters parameters)
+  {
+    return new Client(parameters);
+  }
+
   @Override
   public Set<String> getPartitionIds(String stream)
   {
-    ClientParameters parameters = new ClientParameters();
+    ClientParameters parameters = getParameters();
 
     try {
       URI parsedURI = new URI(uri);
@@ -431,7 +445,7 @@ public Set<String> getPartitionIds(String stream)
         parameters.username(username);
       }
 
-      Client client = new Client(parameters);
+      Client client = getClient(parameters);
 
       List<String> partitions = client.partitions(stream);
       client.close();

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamRecordSupplier.java

@@ -72,15 +72,15 @@
 import java.util.stream.Collectors;
 
 /**
- * Supervisor responsible for managing the KafkaIndexTasks for a single
+ * Supervisor responsible for managing the RabbitStreamIndexTasks for a single
  * dataSource. At a high level, the class accepts a
- * {@link RabbitStreamSupervisorSpec} which includes the Kafka topic and
+ * {@link RabbitStreamSupervisorSpec} which includes the rabbit super stream and
  * configuration as well as an ingestion spec which will be used to generate the
- * indexing tasks. The run loop periodically refreshes its view of the Kafka
- * topic's partitions and the list of running indexing tasks and ensures that
+ * indexing tasks. The run loop periodically refreshes its view of the super stream's
+ * partitions and the list of running indexing tasks and ensures that
  * all partitions are being read from and that there are enough tasks to satisfy
  * the desired number of replicas. As tasks complete, new tasks are queued to
- * process the next range of Kafka offsets.
+ * process the next range of rabbit stream offsets.
  */
 public class RabbitStreamSupervisor extends SeekableStreamSupervisor<String, Long, ByteEntity>
 {
@@ -377,6 +377,7 @@ protected void updatePartitionLagFromStream()
     }
     catch (Exception e) {
       log.warn("Could not fetch partitions for topic/stream [%s]", getIoConfig().getStream());
+      getRecordSupplierLock().unlock();
       throw new StreamException(e);
     }
 

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/supervisor/RabbitStreamSupervisor.java

@@ -83,8 +83,8 @@ public RabbitStreamSupervisorIOConfig(
         new IdleConfig(null, null));
 
     this.consumerProperties = consumerProperties;
+    Preconditions.checkNotNull(uri, "uri");
     this.uri = uri;
-    Preconditions.checkNotNull(stream, "uri");
 
     this.pollTimeout = pollTimeout != null ? pollTimeout : DEFAULT_POLL_TIMEOUT_MILLIS;
   }