add requestTimeoutMs for Client, Producer, Reader (#166)

daleiz · web-flow · commit 1f237bd0bf98 · 2023-01-17T12:23:42.000+08:00
diff --git a/client/src/main/java/io/hstream/BufferedProducerBuilder.java b/client/src/main/java/io/hstream/BufferedProducerBuilder.java
@@ -31,5 +31,7 @@ public interface BufferedProducerBuilder {
    */
   BufferedProducerBuilder compressionType(CompressionType compressionType);
 
+  BufferedProducerBuilder requestTimeoutMs(long timeoutMs);
+
   BufferedProducer build();
 }
diff --git a/client/src/main/java/io/hstream/HStreamClientBuilder.java b/client/src/main/java/io/hstream/HStreamClientBuilder.java
@@ -51,5 +51,7 @@ public interface HStreamClientBuilder {
    */
   HStreamClientBuilder tlsCertPath(String certPath);
 
+  HStreamClientBuilder requestTimeoutMs(long timeoutMs);
+
   HStreamClient build();
 }
diff --git a/client/src/main/java/io/hstream/ProducerBuilder.java b/client/src/main/java/io/hstream/ProducerBuilder.java
@@ -5,5 +5,7 @@ public interface ProducerBuilder {
 
   ProducerBuilder stream(String streamName);
 
+  ProducerBuilder requestTimeoutMs(long timeoutMs);
+
   Producer build();
 }
diff --git a/client/src/main/java/io/hstream/ReaderBuilder.java b/client/src/main/java/io/hstream/ReaderBuilder.java
@@ -13,5 +13,7 @@ public interface ReaderBuilder {
 
   ReaderBuilder readerId(String readerId);
 
+  ReaderBuilder requestTimeoutMs(long timeoutMs);
+
   Reader build();
 }
diff --git a/client/src/main/java/io/hstream/impl/BufferedProducerBuilderImpl.java b/client/src/main/java/io/hstream/impl/BufferedProducerBuilderImpl.java
@@ -1,5 +1,6 @@
 package io.hstream.impl;
 
+import com.google.common.base.Preconditions;
 import io.hstream.BatchSetting;
 import io.hstream.BufferedProducer;
 import io.hstream.BufferedProducerBuilder;
@@ -15,6 +16,8 @@ public class BufferedProducerBuilderImpl implements BufferedProducerBuilder {
   private FlowControlSetting flowControlSetting = FlowControlSetting.newBuilder().build();
   private CompressionType compressionType = CompressionType.NONE;
 
+  private long requestTimeoutMs = DefaultSettings.GRPC_CALL_TIMEOUT_MS;
+
   public BufferedProducerBuilderImpl(HStreamClientKtImpl client) {
     this.client = client;
   }
@@ -43,6 +46,12 @@ public BufferedProducerBuilder compressionType(CompressionType compressionType)
     return this;
   }
 
+  @Override
+  public BufferedProducerBuilder requestTimeoutMs(long timeoutMs) {
+    this.requestTimeoutMs = timeoutMs;
+    return this;
+  }
+
   @Override
   public BufferedProducer build() {
     if (streamName == null) {
@@ -56,7 +65,8 @@ public BufferedProducer build() {
               "BatchSetting.ageLimit:[%d] should not be greater than flowControlSetting.bytesLimit:[%d]",
               batchBytes, flowBytes));
     }
+    Preconditions.checkArgument(requestTimeoutMs > 0);
     return new BufferedProducerKtImpl(
-        client, streamName, batchSetting, flowControlSetting, compressionType);
+        client, streamName, requestTimeoutMs, batchSetting, flowControlSetting, compressionType);
   }
 }
diff --git a/client/src/main/java/io/hstream/impl/DefaultSettings.java b/client/src/main/java/io/hstream/impl/DefaultSettings.java
@@ -8,5 +8,5 @@ public class DefaultSettings {
 
   public static final String DEFAULT_PARTITION_KEY = "";
 
-  public static final long GRPC_CALL_TIMEOUT_SECONDS = 5;
+  public static final long GRPC_CALL_TIMEOUT_MS = 5000;
 }
diff --git a/client/src/main/java/io/hstream/impl/HStreamClientBuilderImpl.java b/client/src/main/java/io/hstream/impl/HStreamClientBuilderImpl.java
@@ -23,6 +23,8 @@ public class HStreamClientBuilderImpl implements HStreamClientBuilder {
   private String keyPath;
   private String certPath;
 
+  private long requestTimeoutMs = DefaultSettings.GRPC_CALL_TIMEOUT_MS;
+
   @Override
   public HStreamClientBuilder serviceUrl(String serviceUrl) {
     this.serviceUrl = serviceUrl;
@@ -59,9 +61,16 @@ public HStreamClientBuilder tlsCertPath(String certPath) {
     return this;
   }
 
+  @Override
+  public HStreamClientBuilder requestTimeoutMs(long timeoutMs) {
+    this.requestTimeoutMs = timeoutMs;
+    return this;
+  }
+
   @Override
   public HStreamClient build() {
     checkNotNull(serviceUrl);
+    checkArgument(requestTimeoutMs > 0);
     Pair<UrlSchema, List<String>> schemaHosts = parseServerUrls(serviceUrl);
     // FIXME: remove enableTls option
     if (schemaHosts.getKey().equals(UrlSchema.HSTREAMS) && !enableTls) {
@@ -74,12 +83,13 @@ public HStreamClient build() {
         if (enableTlsAuthentication) {
           credentialsBuilder = credentialsBuilder.keyManager(new File(certPath), new File(keyPath));
         }
-        return new HStreamClientKtImpl(schemaHosts.getRight(), credentialsBuilder.build());
+        return new HStreamClientKtImpl(
+            schemaHosts.getRight(), requestTimeoutMs, credentialsBuilder.build());
       } catch (IOException e) {
         throw new HStreamDBClientException(String.format("invalid tls options, %s", e));
       }
     }
-    return new HStreamClientKtImpl(schemaHosts.getRight(), null);
+    return new HStreamClientKtImpl(schemaHosts.getRight(), requestTimeoutMs, null);
   }
 
   private Pair<UrlSchema, List<String>> parseServerUrls(String url) {
diff --git a/client/src/main/java/io/hstream/impl/ProducerBuilderImpl.java b/client/src/main/java/io/hstream/impl/ProducerBuilderImpl.java
@@ -9,6 +9,7 @@ public class ProducerBuilderImpl implements ProducerBuilder {
 
   private final HStreamClientKtImpl client;
   private String streamName;
+  private long requestTimeoutMs = DefaultSettings.GRPC_CALL_TIMEOUT_MS;
 
   public ProducerBuilderImpl(HStreamClientKtImpl client) {
     this.client = client;
@@ -20,9 +21,16 @@ public ProducerBuilder stream(String streamName) {
     return this;
   }
 
+  @Override
+  public ProducerBuilder requestTimeoutMs(long timeoutMs) {
+    this.requestTimeoutMs = timeoutMs;
+    return this;
+  }
+
   @Override
   public Producer build() {
     checkNotNull(streamName);
-    return new ProducerKtImpl(client, streamName);
+    checkArgument(requestTimeoutMs > 0);
+    return new ProducerKtImpl(client, streamName, requestTimeoutMs);
   }
 }
diff --git a/client/src/main/java/io/hstream/impl/ReaderBuilderImpl.java b/client/src/main/java/io/hstream/impl/ReaderBuilderImpl.java
@@ -1,7 +1,6 @@
 package io.hstream.impl;
 
-import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Preconditions.*;
 
 import io.hstream.Reader;
 import io.hstream.ReaderBuilder;
@@ -19,6 +18,8 @@ public class ReaderBuilderImpl implements ReaderBuilder {
   private int timeoutMs = 0;
   private String readerId = UUID.randomUUID().toString();
 
+  private long requestTimeoutMs = DefaultSettings.GRPC_CALL_TIMEOUT_MS;
+
   public ReaderBuilderImpl(HStreamClientKtImpl client) {
     this.client = client;
   }
@@ -53,6 +54,12 @@ public ReaderBuilder readerId(String readerId) {
     return this;
   }
 
+  @Override
+  public ReaderBuilder requestTimeoutMs(long timeoutMs) {
+    this.requestTimeoutMs = timeoutMs;
+    return this;
+  }
+
   @Override
   public Reader build() {
     checkNotNull(client);
@@ -61,6 +68,8 @@ public Reader build() {
     checkNotNull(shardOffset);
     checkState(timeoutMs >= 0);
     checkNotNull(readerId);
-    return new ReaderKtImpl(client, streamName, shardId, shardOffset, timeoutMs, readerId);
+    checkArgument(requestTimeoutMs > 0);
+    return new ReaderKtImpl(
+        client, streamName, shardId, shardOffset, timeoutMs, readerId, requestTimeoutMs);
   }
 }
diff --git a/client/src/main/kotlin/io/hstream/impl/BufferedProducerKtImpl.kt b/client/src/main/kotlin/io/hstream/impl/BufferedProducerKtImpl.kt
@@ -26,10 +26,11 @@ typealias Futures = MutableList<CompletableFuture<String>>
 class BufferedProducerKtImpl(
     client: HStreamClientKtImpl,
     stream: String,
+    requestTimeoutMs: Long,
     private val batchSetting: BatchSetting,
     private val flowControlSetting: FlowControlSetting,
     private val compressionType: CompressionType,
-) : ProducerKtImpl(client, stream), BufferedProducer {
+) : ProducerKtImpl(client, stream, requestTimeoutMs), BufferedProducer {
     private var lock = ReentrantLock()
     private var shardAppendBuffer: HashMap<Long, Records> = HashMap()
     private var shardAppendFutures: HashMap<Long, Futures> = HashMap()
diff --git a/client/src/main/kotlin/io/hstream/impl/HStreamClientKtImpl.kt b/client/src/main/kotlin/io/hstream/impl/HStreamClientKtImpl.kt
@@ -41,24 +41,24 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
 import kotlin.streams.toList
 
-class HStreamClientKtImpl(bootstrapServerUrls: List<String>, credentials: ChannelCredentials? = null) : HStreamClient {
+class HStreamClientKtImpl(bootstrapServerUrls: List<String>, private val requestTimeoutMs: Long, credentials: ChannelCredentials? = null) : HStreamClient {
 
     private val logger = LoggerFactory.getLogger(HStreamClientKtImpl::class.java)
 
     val channelProvider = ChannelProvider(credentials)
     val clusterServerUrls: AtomicReference<List<String>> = AtomicReference(null)
 
     fun <Resp> unaryCallAsync(call: suspend (stub: HStreamApiGrpcKt.HStreamApiCoroutineStub) -> Resp): CompletableFuture<Resp> {
-        return unaryCallAsync(clusterServerUrls, channelProvider, call)
+        return unaryCallAsync(clusterServerUrls, channelProvider, requestTimeoutMs, call)
     }
 
     // warning: this method will block current thread. Do not call this in suspend functions, use unaryCallCoroutine instead!
     fun <Resp> unaryCallBlocked(call: suspend (stub: HStreamApiGrpcKt.HStreamApiCoroutineStub) -> Resp): Resp {
-        return unaryCallBlocked(clusterServerUrls, channelProvider, call)
+        return unaryCallBlocked(clusterServerUrls, channelProvider, requestTimeoutMs, call)
     }
 
     suspend fun <Resp> unaryCallCoroutine(call: suspend (stub: HStreamApiGrpcKt.HStreamApiCoroutineStub) -> Resp): Resp {
-        return unaryCallCoroutine(clusterServerUrls, channelProvider, call)
+        return unaryCallCoroutine(clusterServerUrls, channelProvider, requestTimeoutMs, call)
     }
 
     fun getCoroutineStub(url: String): HStreamApiGrpcKt.HStreamApiCoroutineStub {
diff --git a/client/src/main/kotlin/io/hstream/impl/ProducerKtImpl.kt b/client/src/main/kotlin/io/hstream/impl/ProducerKtImpl.kt
@@ -28,7 +28,11 @@ 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 {
+open class ProducerKtImpl(
+    private val client: HStreamClientKtImpl,
+    private val stream: String,
+    val requestTimeoutMs: Long
+) : Producer {
     private val serverUrls: HashMap<Long, String> = HashMap()
     private val serverUrlsLock: Mutex = Mutex()
     private val shards: List<Shard>
@@ -137,7 +141,7 @@ open class ProducerKtImpl(private val client: HStreamClientKtImpl, private val s
         val serverUrl = lookupServerUrl(shardId, forceUpdate)
         logger.debug("try append with serverUrl [{}], current left tryTimes is [{}]", serverUrl, tryTimes)
         return try {
-            client.getCoroutineStubWithTimeout(serverUrl, DefaultSettings.GRPC_CALL_TIMEOUT_SECONDS)
+            client.getCoroutineStubWithTimeout(serverUrl, requestTimeoutMs)
                 .append(appendRequest).recordIdsList.map(GrpcUtils::recordIdFromGrpc)
         } catch (e: StatusException) {
             handleGRPCException(serverUrl, e)
diff --git a/client/src/main/kotlin/io/hstream/impl/ReaderKtImpl.kt b/client/src/main/kotlin/io/hstream/impl/ReaderKtImpl.kt
@@ -25,7 +25,8 @@ class ReaderKtImpl(
     private val shardId: Long,
     private val shardOffset: StreamShardOffset,
     private val timeoutMs: Int,
-    private val readerId: String
+    private val readerId: String,
+    private val requestTimeoutMs: Long
 ) : Reader {
 
     private var serverUrl: String
@@ -50,7 +51,7 @@ class ReaderKtImpl(
         readerScope.launch {
             try {
                 val readShardRequest = ReadShardRequest.newBuilder().setReaderId(readerId).setMaxRecords(maxRecords).build()
-                val readShardResponse = client.getCoroutineStubWithTimeout(serverUrl, DefaultSettings.GRPC_CALL_TIMEOUT_SECONDS)
+                val readShardResponse = client.getCoroutineStubWithTimeout(serverUrl, requestTimeoutMs)
                     .readShard(readShardRequest)
                 val res = readShardResponse.receivedRecordsList.flatMap {
                     RecordUtils.decompress(it).map { receivedHStreamRecord ->
@@ -82,7 +83,7 @@ class ReaderKtImpl(
             .setReaderId(readerId)
             .build()
         runBlocking {
-            client.getCoroutineStubWithTimeout(serverUrl, DefaultSettings.GRPC_CALL_TIMEOUT_SECONDS).deleteShardReader(deleteShardReaderRequest)
+            client.getCoroutineStubWithTimeout(serverUrl, requestTimeoutMs).deleteShardReader(deleteShardReaderRequest)
         }
 
         logger.info("Reader [{}] closed", readerId)
diff --git a/client/src/main/kotlin/io/hstream/impl/Utils.kt b/client/src/main/kotlin/io/hstream/impl/Utils.kt
@@ -74,7 +74,12 @@ suspend fun refreshClusterInfo(serverUrls: List<String>, channelProvider: Channe
     }
 }
 
-suspend fun <Resp> unaryCallCoroutine(urlsRef: AtomicReference<List<String>>, channelProvider: ChannelProvider, call: suspend (stub: HStreamApiCoroutineStub) -> Resp): Resp {
+suspend fun <Resp> unaryCallCoroutine(
+    urlsRef: AtomicReference<List<String>>,
+    channelProvider: ChannelProvider,
+    timeoutMs: Long,
+    call: suspend (stub: HStreamApiCoroutineStub) -> Resp
+): Resp {
     // Note: A failed grpc call can throw both 'StatusException' and 'StatusRuntimeException'.
     //       This function is for handling them.
     suspend fun handleGRPCException(urls: List<String>, e: Throwable): Resp {
@@ -95,21 +100,38 @@ suspend fun <Resp> unaryCallCoroutine(urlsRef: AtomicReference<List<String>>, ch
     logger.debug("unary rpc with urls [{}]", urls)
 
     try {
-        return call(HStreamApiCoroutineStub(channelProvider.get(urls[0])).withDeadlineAfter(DefaultSettings.GRPC_CALL_TIMEOUT_SECONDS, TimeUnit.SECONDS))
+        return call(HStreamApiCoroutineStub(channelProvider.get(urls[0])).withDeadlineAfter(timeoutMs, TimeUnit.MILLISECONDS))
     } catch (e: StatusException) {
         return handleGRPCException(urls, e)
     } catch (e: StatusRuntimeException) {
         return handleGRPCException(urls, e)
     }
 }
 
-fun <Resp> unaryCallAsync(urlsRef: AtomicReference<List<String>>, channelProvider: ChannelProvider, call: suspend (stub: HStreamApiCoroutineStub) -> Resp): CompletableFuture<Resp> {
-    return futureForIO { unaryCallCoroutine(urlsRef, channelProvider, call) }
+fun <Resp> unaryCallAsync(
+    urlsRef: AtomicReference<List<String>>,
+    channelProvider: ChannelProvider,
+    timeoutMs: Long,
+    call: suspend (stub: HStreamApiCoroutineStub) -> Resp
+): CompletableFuture<Resp> {
+    return futureForIO { unaryCallCoroutine(urlsRef, channelProvider, timeoutMs, call) }
 }
 
 // warning: this method will block current thread. Do not call this in suspend functions, use unaryCallCoroutine instead!
-fun <Resp> unaryCallBlocked(urlsRef: AtomicReference<List<String>>, channelProvider: ChannelProvider, call: suspend (stub: HStreamApiCoroutineStub) -> Resp): Resp {
-    return futureForIO(MoreExecutors.directExecutor().asCoroutineDispatcher()) { unaryCallCoroutine(urlsRef, channelProvider, call) }.join()
+fun <Resp> unaryCallBlocked(
+    urlsRef: AtomicReference<List<String>>,
+    channelProvider: ChannelProvider,
+    timeoutMs: Long,
+    call: suspend (stub: HStreamApiCoroutineStub) -> Resp
+): Resp {
+    return futureForIO(MoreExecutors.directExecutor().asCoroutineDispatcher()) {
+        unaryCallCoroutine(
+            urlsRef,
+            channelProvider,
+            timeoutMs,
+            call
+        )
+    }.join()
 }
 
 fun <Resp> unaryCallWithCurrentUrlsAsync(urls: List<String>, channelProvider: ChannelProvider, call: suspend (stub: HStreamApiCoroutineStub) -> Resp): CompletableFuture<Resp> {

Original file line number	Diff line number	Diff line change
`@@ -31,5 +31,7 @@ public interface BufferedProducerBuilder {`
`31`	`31`	`*/`
`32`	`32`	`BufferedProducerBuilder compressionType(CompressionType compressionType);`
`33`	`33`
	`34`	`+ BufferedProducerBuilder requestTimeoutMs(long timeoutMs);`
	`35`	`+`
`34`	`36`	`BufferedProducer build();`
`35`	`37`	`}`
Original file line number	Diff line number	Diff line change
`@@ -51,5 +51,7 @@ public interface HStreamClientBuilder {`
`51`	`51`	`*/`
`52`	`52`	`HStreamClientBuilder tlsCertPath(String certPath);`
`53`	`53`
	`54`	`+ HStreamClientBuilder requestTimeoutMs(long timeoutMs);`
	`55`	`+`
`54`	`56`	`HStreamClient build();`
`55`	`57`	`}`
Original file line number	Diff line number	Diff line change
`@@ -5,5 +5,7 @@ public interface ProducerBuilder {`
`5`	`5`
`6`	`6`	`ProducerBuilder stream(String streamName);`
`7`	`7`
	`8`	`+ ProducerBuilder requestTimeoutMs(long timeoutMs);`
	`9`	`+`
`8`	`10`	`Producer build();`
`9`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,5 +13,7 @@ public interface ReaderBuilder {`
`13`	`13`
`14`	`14`	`ReaderBuilder readerId(String readerId);`
`15`	`15`
	`16`	`+ ReaderBuilder requestTimeoutMs(long timeoutMs);`
	`17`	`+`
`16`	`18`	`Reader build();`
`17`	`19`	`}`
Original file line number	Diff line number	Diff line change
`@@ -8,5 +8,5 @@ public class DefaultSettings {`
`8`	`8`
`9`	`9`	`public static final String DEFAULT_PARTITION_KEY = "";`
`10`	`10`
`11`		`- public static final long GRPC_CALL_TIMEOUT_SECONDS = 5;`
	`11`	`+ public static final long GRPC_CALL_TIMEOUT_MS = 5000;`
`12`	`12`	`}`