Support time and bytes size based flush strategy (#61)

Author: s12f

client/src/main/java/io/hstream/BufferedProducer.java

@@ -0,0 +1,13 @@
+package io.hstream;
+
+import java.io.Closeable;
+
+/** The interface for the HStream BuffetedProducer. */
+public interface BufferedProducer extends Producer, Closeable {
+
+  /** explicitly flush buffered records. */
+  void flush();
+
+  /** closes the buffered producer, will call flush() first */
+  void close();
+}

client/src/main/java/io/hstream/BufferedProducerBuilder.java

@@ -0,0 +1,17 @@
+package io.hstream;
+
+/** A builder for {@link BufferedProducer}s. */
+public interface BufferedProducerBuilder {
+
+  BufferedProducerBuilder stream(String streamName);
+
+  BufferedProducerBuilder recordCountLimit(int recordCountLimit);
+
+  BufferedProducerBuilder flushIntervalMs(long flushIntervalMs);
+
+  BufferedProducerBuilder maxBytesSize(int maxBytesSize);
+
+  BufferedProducerBuilder throwExceptionIfFull(boolean throwExceptionIfFull);
+
+  BufferedProducer build();
+}

client/src/main/java/io/hstream/HStreamClient.java

@@ -14,6 +14,8 @@ static HStreamClientBuilder builder() {
   /** @return a {@link ProducerBuilder} */
   ProducerBuilder newProducer();
 
+  BufferedProducerBuilder newBufferedProducer();
+
   /** @return a {@link ConsumerBuilder} */
   ConsumerBuilder newConsumer();
 

client/src/main/java/io/hstream/ProducerBuilder.java

@@ -5,9 +5,5 @@ public interface ProducerBuilder {
 
   ProducerBuilder stream(String streamName);
 
-  ProducerBuilder enableBatch();
-
-  ProducerBuilder recordCountLimit(int recordCountLimit);
-
   Producer build();
 }

client/src/main/java/io/hstream/impl/BufferedProducerBuilderImpl.java

@@ -0,0 +1,64 @@
+package io.hstream.impl;
+
+import io.hstream.BufferedProducer;
+import io.hstream.BufferedProducerBuilder;
+import io.hstream.HStreamDBClientException;
+
+public class BufferedProducerBuilderImpl implements BufferedProducerBuilder {
+
+  private String streamName;
+  private int recordCountLimit = 100;
+  private long flushIntervalMs = 100;
+  private int maxBytesSize = 4096;
+  private boolean throwExceptionIfFull = false;
+
+  @Override
+  public BufferedProducerBuilder stream(String streamName) {
+    this.streamName = streamName;
+    return this;
+  }
+
+  /** @param recordCountLimit default value is 100, it can NOT be less than 1 */
+  @Override
+  public BufferedProducerBuilder recordCountLimit(int recordCountLimit) {
+    this.recordCountLimit = recordCountLimit;
+    return this;
+  }
+
+  /**
+   * @param flushIntervalMs default value is 100ms, if flushIntervalMs <= 0, disables timed based
+   *     flush strategy.
+   */
+  @Override
+  public BufferedProducerBuilder flushIntervalMs(long flushIntervalMs) {
+    this.flushIntervalMs = flushIntervalMs;
+    return this;
+  }
+
+  /**
+   * @param maxBytesSize default value is 4K(4096), if maxBytesSize <= 0, does not limit bytes size
+   */
+  @Override
+  public BufferedProducerBuilder maxBytesSize(int maxBytesSize) {
+    this.maxBytesSize = maxBytesSize;
+    return this;
+  }
+
+  @Override
+  public BufferedProducerBuilder throwExceptionIfFull(boolean throwExceptionIfFull) {
+    this.throwExceptionIfFull = throwExceptionIfFull;
+    return this;
+  }
+
+  @Override
+  public BufferedProducer build() {
+    if (recordCountLimit < 1) {
+      throw new HStreamDBClientException(
+          String.format(
+              "build buffedProducer failed, recordCountLimit(%d) can NOT be less than 1",
+              recordCountLimit));
+    }
+    return new BufferedProducerKtImpl(
+        streamName, recordCountLimit, flushIntervalMs, maxBytesSize, throwExceptionIfFull);
+  }
+}

client/src/main/java/io/hstream/impl/ProducerBuilderImpl.java

@@ -4,47 +4,20 @@
 
 import io.hstream.Producer;
 import io.hstream.ProducerBuilder;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicReference;
 
 public class ProducerBuilderImpl implements ProducerBuilder {
 
   private String streamName;
 
-  private boolean enableBatch = false;
-
-  private int recordCountLimit = 1;
-
-  private final AtomicReference<List<String>> serverUrls;
-  private final ChannelProvider channelProvider;
-
-  public ProducerBuilderImpl(
-      AtomicReference<List<String>> serverUrls, ChannelProvider channelProvider) {
-    this.serverUrls = serverUrls;
-    this.channelProvider = channelProvider;
-  }
-
   @Override
   public ProducerBuilder stream(String streamName) {
     this.streamName = streamName;
     return this;
   }
 
-  @Override
-  public ProducerBuilder enableBatch() {
-    this.enableBatch = true;
-    return this;
-  }
-
-  @Override
-  public ProducerBuilder recordCountLimit(int recordCountLimit) {
-    this.recordCountLimit = recordCountLimit;
-    return this;
-  }
-
   @Override
   public Producer build() {
     checkNotNull(streamName);
-    return new ProducerKtImpl(streamName, enableBatch, recordCountLimit);
+    return new ProducerKtImpl(streamName);
   }
 }

client/src/main/kotlin/io/hstream/impl/BufferedProducerKtImpl.kt

@@ -0,0 +1,119 @@
+package io.hstream.impl
+
+import io.hstream.BufferedProducer
+import io.hstream.HStreamDBClientException
+import io.hstream.RecordId
+import io.hstream.internal.HStreamRecord
+import org.slf4j.LoggerFactory
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Executors
+import java.util.concurrent.ScheduledFuture
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.collections.ArrayList
+import kotlin.concurrent.withLock
+
+class BufferedProducerKtImpl(
+    stream: String,
+    private val recordCountLimit: Int,
+    private val flushIntervalMs: Long,
+    private val maxBytesSize: Int,
+    private val throwExceptionIfFull: Boolean
+) : ProducerKtImpl(stream), BufferedProducer {
+    private var lock = ReentrantLock()
+    private var recordBuffer: MutableList<HStreamRecord> = ArrayList(recordCountLimit)
+    private var futures: MutableList<CompletableFuture<RecordId>> = ArrayList(recordCountLimit)
+    private var timerService: ScheduledFuture<*>? = null
+
+    @Volatile
+    private var closed: Boolean = false
+    private var bufferedBytesSize: Int = 0
+
+    @Volatile
+    private var isFull = false
+
+    init {
+        if (flushIntervalMs > 0) {
+            runTimer()
+        }
+    }
+
+    private fun runTimer() {
+        timerService = scheduler.scheduleAtFixedRate(
+            { flush() },
+            flushIntervalMs,
+            flushIntervalMs,
+            TimeUnit.MILLISECONDS
+        )
+    }
+
+    override fun writeInternal(hStreamRecord: HStreamRecord): CompletableFuture<RecordId> {
+        return addToBuffer(hStreamRecord)
+    }
+
+    private fun addToBuffer(hStreamRecord: HStreamRecord): CompletableFuture<RecordId> {
+        // fuzzy check
+        val recordFuture = CompletableFuture<RecordId>()
+        if (throwExceptionIfFull && isFull) {
+            recordFuture.completeExceptionally(HStreamDBClientException("buffer is full"))
+            return recordFuture
+        }
+        lock.withLock {
+            if (closed) {
+                throw HStreamDBClientException("BufferedProducer is closed")
+            }
+            // it is impossible that buffer is full after holding the lock,
+            // if buffer is full, there must exist another thread keeping the lock(flushing buffer).
+            recordBuffer.add(hStreamRecord)
+            futures.add(recordFuture)
+            bufferedBytesSize += hStreamRecord.payload.size()
+            if (isFull()) {
+                isFull = true
+                flush()
+            }
+            return recordFuture
+        }
+    }
+
+    private fun isFull(): Boolean {
+        return (recordBuffer.size == recordCountLimit) || maxBytesSize > 0 && bufferedBytesSize >= maxBytesSize
+    }
+
+    override fun flush() {
+        lock.withLock {
+            if (recordBuffer.isEmpty()) {
+                return
+            }
+            try {
+                val recordBufferCount = recordBuffer.size
+                logger.info("ready to flush recordBuffer, current buffer size is [{}]", recordBufferCount)
+                val ids = futureForIO { super.writeHStreamRecords(recordBuffer) }.join()
+                for (i in ids.indices) {
+                    futures[i].complete(ids[i])
+                }
+                logger.info("flush the record buffer successfully")
+            } catch (e: Throwable) {
+                for (i in futures.indices) {
+                    futures[i].completeExceptionally(e)
+                }
+            }
+            recordBuffer.clear()
+            futures.clear()
+            bufferedBytesSize = 0
+            isFull = false
+        }
+    }
+
+    override fun close() {
+        if (!closed) {
+            timerService?.cancel(false)
+            closed = true
+            flush()
+        }
+    }
+
+    companion object {
+        private val logger = LoggerFactory.getLogger(BufferedProducerKtImpl::class.java)
+        private val scheduler = Executors.newScheduledThreadPool(4)
+    }
+}

client/src/main/kotlin/io/hstream/impl/HStreamClientKtImpl.kt

@@ -1,6 +1,7 @@
 package io.hstream.impl
 
 import com.google.protobuf.Empty
+import io.hstream.BufferedProducerBuilder
 import io.hstream.ConsumerBuilder
 import io.hstream.HStreamClient
 import io.hstream.ProducerBuilder
@@ -57,7 +58,11 @@ class HStreamClientKtImpl(bootstrapServerUrls: List<String>) : HStreamClient {
     }
 
     override fun newProducer(): ProducerBuilder {
-        return ProducerBuilderImpl(clusterServerUrls, channelProvider)
+        return ProducerBuilderImpl()
+    }
+
+    override fun newBufferedProducer(): BufferedProducerBuilder {
+        return BufferedProducerBuilderImpl()
     }
 
     override fun newConsumer(): ConsumerBuilder {