change mongo scans to filter client-side (#347)

Author: agavra

kafka-client/src/main/java/dev/responsive/kafka/internal/db/CassandraClient.java

@@ -36,6 +36,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.function.BooleanSupplier;
+import org.apache.kafka.common.utils.Bytes;
 
 /**
  * {@code CassandraClient} wraps a {@link CqlSession} with utility methods
@@ -49,10 +50,10 @@ public class CassandraClient {
   private final CqlSession session;
 
   private final ResponsiveConfig config;
-  private final TableCache<RemoteKVTable<BoundStatement>> kvFactory;
-  private final TableCache<RemoteKVTable<BoundStatement>> factFactory;
+  private final TableCache<Bytes, Integer, RemoteKVTable<BoundStatement>> kvFactory;
+  private final TableCache<Bytes, Integer, RemoteKVTable<BoundStatement>> factFactory;
   private final WindowedTableCache<RemoteWindowedTable<BoundStatement>> windowedFactory;
-  private final TableCache<CassandraFactTable> globalFactory;
+  private final TableCache<Bytes, Integer, CassandraFactTable> globalFactory;
 
   /**
    * @param session the Cassandra session, expected to be initialized
@@ -178,15 +179,15 @@ public void shutdown() {
     session.close();
   }
 
-  public TableCache<CassandraFactTable> globalFactory() {
+  public TableCache<Bytes, Integer, CassandraFactTable> globalFactory() {
     return globalFactory;
   }
 
-  public TableCache<RemoteKVTable<BoundStatement>> kvFactory() {
+  public TableCache<Bytes, Integer, RemoteKVTable<BoundStatement>> kvFactory() {
     return kvFactory;
   }
 
-  public TableCache<RemoteKVTable<BoundStatement>> factFactory() {
+  public TableCache<Bytes, Integer, RemoteKVTable<BoundStatement>> factFactory() {
     return factFactory;
   }
 

kafka-client/src/main/java/dev/responsive/kafka/internal/db/CassandraFactFlushManager.java

@@ -34,13 +34,14 @@ public class CassandraFactFlushManager extends KVFlushManager {
   public CassandraFactFlushManager(
       final CassandraFactTable table,
       final CassandraClient client,
-      final int kafkaPartition
+      final int kafkaPartition,
+      final TablePartitioner<Bytes, Integer> partitioner
   ) {
     this.table = table;
     this.client = client;
     this.kafkaPartition = kafkaPartition;
+    this.partitioner = partitioner;
 
-    partitioner = TablePartitioner.defaultPartitioner();
     logPrefix = String.format("%s[%d] fact-store", table.name(), kafkaPartition);
   }
 

kafka-client/src/main/java/dev/responsive/kafka/internal/db/CassandraFactTable.java

@@ -32,6 +32,7 @@
 import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
 import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
 import com.datastax.oss.driver.api.querybuilder.schema.CreateTableWithOptions;
+import dev.responsive.kafka.internal.db.partitioning.TablePartitioner;
 import dev.responsive.kafka.internal.db.spec.RemoteTableSpec;
 import java.nio.ByteBuffer;
 import java.time.Instant;
@@ -56,6 +57,7 @@ public class CassandraFactTable implements RemoteKVTable<BoundStatement> {
   private final PreparedStatement delete;
   private final PreparedStatement fetchOffset;
   private final PreparedStatement setOffset;
+  private final TablePartitioner<Bytes, Integer> partitioner;
 
   public CassandraFactTable(
       final String name,
@@ -64,7 +66,8 @@ public CassandraFactTable(
       final PreparedStatement insert,
       final PreparedStatement delete,
       final PreparedStatement fetchOffset,
-      final PreparedStatement setOffset
+      final PreparedStatement setOffset,
+      final TablePartitioner<Bytes, Integer> partitioner
   ) {
     this.name = name;
     this.client = client;
@@ -73,10 +76,11 @@ public CassandraFactTable(
     this.delete = delete;
     this.fetchOffset = fetchOffset;
     this.setOffset = setOffset;
+    this.partitioner = partitioner;
   }
 
   public static CassandraFactTable create(
-      final RemoteTableSpec spec,
+      final RemoteTableSpec<Bytes, Integer> spec,
       final CassandraClient client
   ) {
     final String name = spec.tableName();
@@ -161,7 +165,8 @@ public static CassandraFactTable create(
         insert,
         delete,
         fetchOffset,
-        setOffset
+        setOffset,
+        spec.partitioner()
     );
   }
 
@@ -193,7 +198,7 @@ public CassandraFactFlushManager init(
             .build()
     );
 
-    return new CassandraFactFlushManager(this, client, kafkaPartition);
+    return new CassandraFactFlushManager(this, client, kafkaPartition, partitioner);
   }
 
   @Override

kafka-client/src/main/java/dev/responsive/kafka/internal/db/CassandraKeyValueTable.java

@@ -77,7 +77,7 @@ public class CassandraKeyValueTable implements RemoteKVTable<BoundStatement> {
   private final PreparedStatement ensureEpoch;
 
   public static CassandraKeyValueTable create(
-      final RemoteTableSpec spec,
+      final RemoteTableSpec<Bytes, Integer> spec,
       final CassandraClient client
   ) throws InterruptedException, TimeoutException {
     final String name = spec.tableName();

kafka-client/src/main/java/dev/responsive/kafka/internal/db/CassandraWindowedTable.java

@@ -97,7 +97,7 @@ public class CassandraWindowedTable implements RemoteWindowedTable<BoundStatemen
   private final PreparedStatement ensureEpoch;
 
   public static CassandraWindowedTable create(
-      final RemoteTableSpec spec,
+      final RemoteTableSpec<WindowedKey, Segmenter.SegmentPartition> spec,
       final CassandraClient client,
       final WindowSegmentPartitioner partitioner
   ) throws InterruptedException, TimeoutException {

kafka-client/src/main/java/dev/responsive/kafka/internal/db/MongoKVFlushManager.java

@@ -44,13 +44,14 @@ public class MongoKVFlushManager extends KVFlushManager {
   public MongoKVFlushManager(
       final MongoKVTable table,
       final MongoCollection<KVDoc> kvDocs,
-      final int kafkaPartition
+      final int kafkaPartition,
+      final TablePartitioner<Bytes, Integer> partitioner
   ) {
     this.table = table;
     this.kvDocs = kvDocs;
     this.kafkaPartition = kafkaPartition;
+    this.partitioner = partitioner;
 
-    partitioner = TablePartitioner.defaultPartitioner();
     logPrefix = String.format("%s[%d] kv-store {epoch=%d} ",
                               table.name(), kafkaPartition, table.localEpoch(kafkaPartition));
     log = new LogContext(logPrefix).logger(MongoKVFlushManager.class);

kafka-client/src/main/java/dev/responsive/kafka/internal/db/RemoteTableSpecFactory.java

@@ -40,35 +40,36 @@
  */
 public class RemoteTableSpecFactory {
 
-  public static RemoteTableSpec globalSpec(
+  public static RemoteTableSpec<Bytes, Integer> globalSpec(
       final ResponsiveKeyValueParams params,
       final TablePartitioner<Bytes, Integer> partitioner
   ) {
-    return new GlobalTableSpec(new BaseTableSpec(params.name().tableName(), partitioner));
+    return new GlobalTableSpec<>(new BaseTableSpec<>(params.name().tableName(), partitioner));
   }
 
-  public static RemoteTableSpec fromKVParams(
+  public static RemoteTableSpec<Bytes, Integer> fromKVParams(
       final ResponsiveKeyValueParams params,
       final TablePartitioner<Bytes, Integer> partitioner
   ) {
-    RemoteTableSpec spec = new BaseTableSpec(params.name().tableName(), partitioner);
+    RemoteTableSpec<Bytes, Integer> spec =
+        new BaseTableSpec<>(params.name().tableName(), partitioner);
 
     if (params.timeToLive().isPresent()) {
-      spec = new TtlTableSpec(spec, params.timeToLive().get());
+      spec = new TtlTableSpec<>(spec, params.timeToLive().get());
     }
 
     if (params.schemaType() == SchemaTypes.KVSchema.FACT) {
-      spec = new TimeWindowedCompactionTableSpec(spec);
+      spec = new TimeWindowedCompactionTableSpec<>(spec);
     }
 
     return spec;
   }
 
-  public static RemoteTableSpec fromWindowParams(
+  public static RemoteTableSpec<WindowedKey, SegmentPartition> fromWindowParams(
       final ResponsiveWindowParams params,
       final TablePartitioner<WindowedKey, SegmentPartition> partitioner
   ) {
-    return new BaseTableSpec(params.name().tableName(), partitioner);
+    return new BaseTableSpec<>(params.name().tableName(), partitioner);
   }
 
 }

kafka-client/src/main/java/dev/responsive/kafka/internal/db/SessionTableCache.java

@@ -16,8 +16,10 @@
 
 package dev.responsive.kafka.internal.db;
 
+import dev.responsive.kafka.internal.db.partitioning.Segmenter;
 import dev.responsive.kafka.internal.db.partitioning.SessionSegmentPartitioner;
 import dev.responsive.kafka.internal.db.spec.RemoteTableSpec;
+import dev.responsive.kafka.internal.utils.SessionKey;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeoutException;
@@ -28,8 +30,10 @@ public class SessionTableCache<T extends RemoteTable<?, ?>> {
 
   @FunctionalInterface
   public interface Factory<T> {
-    T create(final RemoteTableSpec spec, SessionSegmentPartitioner partitioner)
-        throws InterruptedException, TimeoutException;
+    T create(
+        final RemoteTableSpec<SessionKey, Segmenter.SegmentPartition> spec,
+        SessionSegmentPartitioner partitioner
+    ) throws InterruptedException, TimeoutException;
   }
 
   private final Map<String, T> tables = new HashMap<>();
@@ -43,7 +47,10 @@ public SessionTableCache(final Factory<T> factory) {
    * Creates a table with the supplied {@code tableName} with the
    * desired schema.
    */
-  public synchronized T create(RemoteTableSpec spec, SessionSegmentPartitioner partitioner)
+  public synchronized T create(
+      RemoteTableSpec<SessionKey, Segmenter.SegmentPartition> spec,
+      SessionSegmentPartitioner partitioner
+  )
       throws InterruptedException, TimeoutException {
     final T existing = tables.get(spec.tableName());
     if (existing != null) {

kafka-client/src/main/java/dev/responsive/kafka/internal/db/TableCache.java

@@ -28,26 +28,26 @@
  * table are only prepared once during the lifetime of the application.
  */
 @ThreadSafe
-public class TableCache<T extends RemoteTable<?, ?>> {
+public class TableCache<K, V, T extends RemoteTable<?, ?>> {
 
   @FunctionalInterface
-  public interface Factory<T> {
-    T create(final RemoteTableSpec spec)
+  public interface Factory<K, V, T> {
+    T create(final RemoteTableSpec<K, V> spec)
         throws InterruptedException, TimeoutException;
   }
 
   private final Map<String, T> tables = new HashMap<>();
-  private final Factory<T> factory;
+  private final Factory<K, V, T> factory;
 
-  public TableCache(final Factory<T> factory) {
+  public TableCache(final Factory<K, V, T> factory) {
     this.factory = factory;
   }
 
   /**
    * Creates a table with the supplied {@code tableName} with the
    * desired schema.
    */
-  public synchronized T create(RemoteTableSpec spec)
+  public synchronized T create(RemoteTableSpec<K, V> spec)
       throws InterruptedException, TimeoutException {
     final T existing = tables.get(spec.tableName());
     if (existing != null) {

kafka-client/src/main/java/dev/responsive/kafka/internal/db/WindowedTableCache.java

@@ -16,8 +16,10 @@
 
 package dev.responsive.kafka.internal.db;
 
+import dev.responsive.kafka.internal.db.partitioning.Segmenter;
 import dev.responsive.kafka.internal.db.partitioning.WindowSegmentPartitioner;
 import dev.responsive.kafka.internal.db.spec.RemoteTableSpec;
+import dev.responsive.kafka.internal.utils.WindowedKey;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeoutException;
@@ -28,7 +30,10 @@ public class WindowedTableCache<T extends RemoteTable<?, ?>> {
 
   @FunctionalInterface
   public interface Factory<T> {
-    T create(final RemoteTableSpec spec, WindowSegmentPartitioner partitioner)
+    T create(
+        final RemoteTableSpec<WindowedKey, Segmenter.SegmentPartition> spec,
+        WindowSegmentPartitioner partitioner
+    )
         throws InterruptedException, TimeoutException;
   }
 
@@ -43,7 +48,10 @@ public WindowedTableCache(final Factory<T> factory) {
    * Creates a table with the supplied {@code tableName} with the
    * desired schema.
    */
-  public synchronized T create(RemoteTableSpec spec, WindowSegmentPartitioner partitioner)
+  public synchronized T create(
+      RemoteTableSpec<WindowedKey, Segmenter.SegmentPartition> spec,
+      WindowSegmentPartitioner partitioner
+  )
       throws InterruptedException, TimeoutException {
     final T existing = tables.get(spec.tableName());
     if (existing != null) {