update Append (#132)

* update Append

* update proto

Author: daleiz

client/build.gradle.kts

@@ -65,6 +65,14 @@ tasks.test {
     }
 }
 
+sourceSets {
+    main {
+        proto {
+            exclude("api/*")
+        }
+    }
+}
+
 protobuf {
     protoc {
         artifact = "com.google.protobuf:protoc:3.19.2"

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

@@ -63,6 +63,7 @@ static HStreamClientBuilder builder() {
    * List shards in a stream.
    *
    * @param streamName the name of the stream
+   * @return a list of {@link Shard}s.
    */
   List<Shard> listShards(String streamName);
 

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

@@ -108,17 +108,17 @@ class BufferedProducerKtImpl(
             val job = shardAppendJobs[shardId]
             shardAppendJobs[shardId] = batchScope.launch {
                 job?.join()
-                writeSingleKeyHStreamRecords(records, futures)
+                writeShard(shardId, records, futures)
                 logger.info("wrote batch for shard:$shardId")
                 flowController?.release(recordsBytesSize)
             }
         }
     }
 
     // only can be called by flush()
-    private suspend fun writeSingleKeyHStreamRecords(records: Records, futures: Futures) {
+    private suspend fun writeShard(shardId: Long, records: Records, futures: Futures) {
         try {
-            val ids = super.writeHStreamRecords(records, records[0].header.key)
+            val ids = super.writeHStreamRecords(records, shardId)
             for (i in ids.indices) {
                 futures[i].complete(ids[i])
             }

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

@@ -86,7 +86,8 @@ open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val s
         val future = CompletableFuture<String>()
         writeRecordScope.launch {
             try {
-                val ids = writeHStreamRecords(listOf(hStreamRecord), hStreamRecord.header.key)
+                val shardId = calculateShardIdByPartitionKey(hStreamRecord.header.key)
+                val ids = writeHStreamRecords(listOf(hStreamRecord), shardId)
                 future.complete(ids[0])
             } catch (e: Throwable) {
                 future.completeExceptionally(e)
@@ -112,7 +113,7 @@ open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val s
 
     private suspend fun appendWithRetry(
         appendRequest: AppendRequest,
-        partitionKey: String,
+        shardId: Long,
         tryTimes: Int,
         forceUpdate: Boolean = false
     ): List<String> {
@@ -123,15 +124,14 @@ 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, partitionKey, tryTimes - 1, true)
+                return appendWithRetry(appendRequest, shardId, tryTimes - 1, true)
             } else {
                 throw HStreamDBClientException(e)
             }
         }
 
         check(tryTimes > 0)
 
-        val shardId = calculateShardIdByPartitionKey(partitionKey)
         val serverUrl = lookupServerUrl(shardId, forceUpdate)
         logger.info("try append with serverUrl [{}], current left tryTimes is [{}]", serverUrl, tryTimes)
         return try {
@@ -144,9 +144,9 @@ open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val s
         }
     }
 
-    protected suspend fun writeHStreamRecords(hStreamRecords: List<HStreamRecord>, key: String): List<String> {
-        val appendRequest = AppendRequest.newBuilder().setStreamName(stream).addAllRecords(hStreamRecords).build()
-        return appendWithRetry(appendRequest, key, DefaultSettings.APPEND_RETRY_MAX_TIMES)
+    protected suspend fun writeHStreamRecords(hStreamRecords: List<HStreamRecord>, shardId: Long): List<String> {
+        val appendRequest = AppendRequest.newBuilder().setStreamName(stream).setShardId(shardId).addAllRecords(hStreamRecords).build()
+        return appendWithRetry(appendRequest, shardId, DefaultSettings.APPEND_RETRY_MAX_TIMES)
     }
 
     companion object {