update Producers with the new sharding model (#126)

Author: daleiz

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

@@ -29,10 +29,10 @@ class BufferedProducerKtImpl(
     private val flowControlSetting: FlowControlSetting,
 ) : ProducerKtImpl(client, stream), BufferedProducer {
     private var lock = ReentrantLock()
-    private var orderingBuffer: HashMap<String, Records> = HashMap()
-    private var orderingFutures: HashMap<String, Futures> = HashMap()
-    private var orderingBytesSize: HashMap<String, Int> = HashMap()
-    private var orderingJobs: HashMap<String, Job> = HashMap()
+    private var shardAppendBuffer: HashMap<Long, Records> = HashMap()
+    private var shardAppendFutures: HashMap<Long, Futures> = HashMap()
+    private var shardAppendBytesSize: HashMap<Long, Int> = HashMap()
+    private var shardAppendJobs: HashMap<Long, Job> = HashMap()
     private var batchScope: CoroutineScope = CoroutineScope(Dispatchers.Default)
 
     private val flowController: FlowController? = if (flowControlSetting.bytesLimit > 0) FlowController(flowControlSetting.bytesLimit) else null
@@ -41,7 +41,7 @@ class BufferedProducerKtImpl(
     private var closed: Boolean = false
 
     private val scheduler = Executors.newScheduledThreadPool(1)
-    private var timerServices: HashMap<String, ScheduledFuture<*>> = HashMap()
+    private var timerServices: HashMap<Long, ScheduledFuture<*>> = HashMap()
 
     override fun writeInternal(hStreamRecord: HStreamRecord): CompletableFuture<String> {
         if (closed) {
@@ -59,56 +59,57 @@ class BufferedProducerKtImpl(
             }
 
             val recordFuture = CompletableFuture<String>()
-            val key = hStreamRecord.header.key
-            if (!orderingBuffer.containsKey(key)) {
-                orderingBuffer[key] = LinkedList()
-                orderingFutures[key] = LinkedList()
-                orderingBytesSize[key] = 0
+            val partitionKey = hStreamRecord.header.key
+            val shardId = calculateShardIdByPartitionKey(partitionKey)
+            if (!shardAppendBuffer.containsKey(shardId)) {
+                shardAppendBuffer[shardId] = LinkedList()
+                shardAppendFutures[shardId] = LinkedList()
+                shardAppendBytesSize[shardId] = 0
                 if (batchSetting.ageLimit > 0) {
-                    timerServices[key] =
-                        scheduler.schedule({ flushForKey(key) }, batchSetting.ageLimit, TimeUnit.MILLISECONDS)
+                    timerServices[shardId] =
+                        scheduler.schedule({ flushForShard(shardId) }, batchSetting.ageLimit, TimeUnit.MILLISECONDS)
                 }
             }
-            orderingBuffer[key]!!.add(hStreamRecord)
-            orderingFutures[key]!!.add(recordFuture)
-            orderingBytesSize[key] = orderingBytesSize[key]!! + hStreamRecord.payload.size()
-            if (isFull(key)) {
-                flushForKey(key)
+            shardAppendBuffer[shardId]!!.add(hStreamRecord)
+            shardAppendFutures[shardId]!!.add(recordFuture)
+            shardAppendBytesSize[shardId] = shardAppendBytesSize[shardId]!! + hStreamRecord.payload.size()
+            if (isFull(shardId)) {
+                flushForShard(shardId)
             }
             return recordFuture
         }
     }
 
-    private fun isFull(key: String): Boolean {
-        val recordCount = orderingBuffer[key]!!.size
-        val bytesSize = orderingBytesSize[key]!!
+    private fun isFull(shardId: Long): Boolean {
+        val recordCount = shardAppendBuffer[shardId]!!.size
+        val bytesSize = shardAppendBytesSize[shardId]!!
         return batchSetting.recordCountLimit in 1..recordCount || batchSetting.bytesLimit in 1..bytesSize
     }
 
     override fun flush() {
         lock.withLock {
-            for (key in orderingBuffer.keys.toList()) {
-                flushForKey(key)
+            for (shard in shardAppendBuffer.keys.toList()) {
+                flushForShard(shard)
             }
         }
     }
 
-    private fun flushForKey(key: String) {
+    private fun flushForShard(shardId: Long) {
         lock.withLock {
-            val records = orderingBuffer[key]!!
-            val futures = orderingFutures[key]!!
-            val recordsBytesSize = orderingBytesSize[key]!!
-            logger.info("ready to flush recordBuffer for key:$key, current buffer size is [{}]", records.size)
-            orderingBuffer.remove(key)
-            orderingFutures.remove(key)
-            orderingBytesSize.remove(key)
-            timerServices[key]?.cancel(true)
-            timerServices.remove(key)
-            val job = orderingJobs[key]
-            orderingJobs[key] = batchScope.launch {
+            val records = shardAppendBuffer[shardId]!!
+            val futures = shardAppendFutures[shardId]!!
+            val recordsBytesSize = shardAppendBytesSize[shardId]!!
+            logger.info("ready to flush recordBuffer for shard:$shardId, current buffer size is [{}]", records.size)
+            shardAppendBuffer.remove(shardId)
+            shardAppendFutures.remove(shardId)
+            shardAppendBytesSize.remove(shardId)
+            timerServices[shardId]?.cancel(true)
+            timerServices.remove(shardId)
+            val job = shardAppendJobs[shardId]
+            shardAppendJobs[shardId] = batchScope.launch {
                 job?.join()
                 writeSingleKeyHStreamRecords(records, futures)
-                logger.info("wrote batch for key:$key")
+                logger.info("wrote batch for shard:$shardId")
                 flowController?.release(recordsBytesSize)
             }
         }

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

@@ -9,7 +9,9 @@ import io.hstream.Producer
 import io.hstream.Record
 import io.hstream.internal.AppendRequest
 import io.hstream.internal.HStreamRecord
-import io.hstream.internal.LookupStreamRequest
+import io.hstream.internal.ListShardsRequest
+import io.hstream.internal.LookupShardRequest
+import io.hstream.internal.Shard
 import io.hstream.util.GrpcUtils
 import io.hstream.util.RecordUtils
 import kotlinx.coroutines.CoroutineScope
@@ -19,39 +21,49 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import org.slf4j.LoggerFactory
+import java.math.BigInteger
+import java.nio.charset.StandardCharsets
 import java.util.concurrent.CompletableFuture
 import kotlin.collections.HashMap
 
 open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val stream: String) : Producer {
-    private val serverUrls: HashMap<String, String> = HashMap()
+    private val serverUrls: HashMap<Long, String> = HashMap()
     private val serverUrlsLock: Mutex = Mutex()
+    private val shards: List<Shard>
 
-    private suspend fun lookupServerUrl(orderingKey: String, forceUpdate: Boolean = false): String {
+    init {
+        val listShardRequest = ListShardsRequest.newBuilder()
+            .setStreamName(stream)
+            .build()
+        val listShardResponse = client.unaryCallBlocked { it.listShards(listShardRequest) }
+        shards = listShardResponse.shardsList
+    }
+
+    private suspend fun lookupServerUrl(shardId: Long, forceUpdate: Boolean = false): String {
         if (forceUpdate) {
-            return updateServerUrl(orderingKey)
+            return updateServerUrl(shardId)
         }
         val server: String? = serverUrlsLock.withLock {
-            return@withLock serverUrls[orderingKey]
+            return@withLock serverUrls[shardId]
         }
         if (server != null) {
             return server
         }
-        return updateServerUrl(orderingKey)
+        return updateServerUrl(shardId)
     }
 
-    private suspend fun updateServerUrl(orderingKey: String): String {
-        val req = LookupStreamRequest.newBuilder()
-            .setStreamName(stream)
-            .setOrderingKey(orderingKey)
+    private suspend fun updateServerUrl(shardId: Long): String {
+        val req = LookupShardRequest.newBuilder()
+            .setShardId(shardId)
             .build()
         val server = client.unaryCallCoroutine {
-            val serverNode = it.lookupStream(req).serverNode
+            val serverNode = it.lookupShard(req).serverNode
             return@unaryCallCoroutine "${serverNode.host}:${serverNode.port}"
         }
         serverUrlsLock.withLock {
-            serverUrls[orderingKey] = server
+            serverUrls[shardId] = server
         }
-        logger.debug("updateServerUrl, key:$orderingKey, server:$server")
+        logger.debug("updateServerUrl, key:$shardId, server:$server")
         return server
     }
 
@@ -83,9 +95,24 @@ open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val s
         return future
     }
 
+    protected fun calculateShardIdByPartitionKey(partitionKey: String): Long {
+        val hashcode = com.google.common.hash.Hashing.md5().hashString(partitionKey, StandardCharsets.UTF_8)
+        val hashValue = BigInteger(hashcode.toString(), 16)
+        for (shard in shards) {
+            val start = BigInteger(shard.startHashRangeKey)
+            val end = BigInteger(shard.endHashRangeKey)
+            if (hashValue.compareTo(start) >= 0 && hashValue.compareTo(end) <= 0) {
+                return shard.shardId
+            }
+        }
+
+        check(false)
+        return -1
+    }
+
     private suspend fun appendWithRetry(
         appendRequest: AppendRequest,
-        orderingKey: String,
+        partitionKey: String,
         tryTimes: Int,
         forceUpdate: Boolean = false
     ): List<String> {
@@ -96,22 +123,24 @@ open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val s
             val status = Status.fromThrowable(e)
             if (status.code == Status.UNAVAILABLE.code && tryTimes > 1) {
                 delay(DefaultSettings.REQUEST_RETRY_INTERVAL_SECONDS * 1000)
-                return appendWithRetry(appendRequest, orderingKey, tryTimes - 1, true)
+                return appendWithRetry(appendRequest, partitionKey, tryTimes - 1, true)
             } else {
                 throw HStreamDBClientException(e)
             }
         }
 
         check(tryTimes > 0)
-        val serverUrl = lookupServerUrl(orderingKey, forceUpdate)
+
+        val shardId = calculateShardIdByPartitionKey(partitionKey)
+        val serverUrl = lookupServerUrl(shardId, forceUpdate)
         logger.info("try append with serverUrl [{}], current left tryTimes is [{}]", serverUrl, tryTimes)
-        try {
-            return client.getCoroutineStub(serverUrl)
+        return try {
+            client.getCoroutineStub(serverUrl)
                 .append(appendRequest).recordIdsList.map(GrpcUtils::recordIdFromGrpc)
         } catch (e: StatusException) {
-            return handleGRPCException(serverUrl, e)
+            handleGRPCException(serverUrl, e)
         } catch (e: StatusRuntimeException) {
-            return handleGRPCException(serverUrl, e)
+            handleGRPCException(serverUrl, e)
         }
     }