feat: add metadata setting api (#218)

Author: s12f

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

@@ -53,5 +53,7 @@ public interface HStreamClientBuilder {
 
   HStreamClientBuilder requestTimeoutMs(long timeoutMs);
 
+  HStreamClientBuilder withMetadata(String key, String value);
+
   HStreamClient build();
 }

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

@@ -1,10 +1,10 @@
 package io.hstream.impl;
 
 import com.google.common.util.concurrent.MoreExecutors;
-import io.grpc.ChannelCredentials;
-import io.grpc.Grpc;
-import io.grpc.ManagedChannel;
-import io.grpc.ManagedChannelBuilder;
+import io.grpc.*;
+import io.grpc.stub.MetadataUtils;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class ChannelProviderImpl implements ChannelProvider {
@@ -16,6 +16,8 @@ public class ChannelProviderImpl implements ChannelProvider {
   static final String userAgent =
       "hstreamdb-java/" + ChannelProvider.class.getPackage().getImplementationVersion();
 
+  Map<String, String> header = new HashMap<>();
+
   public ChannelProviderImpl(int size) {
     provider = new ConcurrentHashMap<>(size);
   }
@@ -25,31 +27,41 @@ public ChannelProviderImpl(ChannelCredentials credentials) {
     provider = new ConcurrentHashMap<>(DEFAULT_CHANNEL_PROVIDER_SIZE);
   }
 
+  public ChannelProviderImpl(ChannelCredentials credentials, Map<String, String> header) {
+    this.credentials = credentials;
+    provider = new ConcurrentHashMap<>(DEFAULT_CHANNEL_PROVIDER_SIZE);
+    if (header != null) {
+      this.header = header;
+    }
+  }
+
   public ChannelProviderImpl() {
     this(DEFAULT_CHANNEL_PROVIDER_SIZE);
   }
 
   @Override
   public ManagedChannel get(String serverUrl) {
+    return provider.computeIfAbsent(serverUrl, this::getInternal);
+  }
+
+  ManagedChannel getInternal(String url) {
+    ManagedChannelBuilder<?> builder;
     if (credentials == null) {
-      return provider.computeIfAbsent(
-          serverUrl,
-          url ->
-              ManagedChannelBuilder.forTarget(url)
-                  .disableRetry()
-                  .usePlaintext()
-                  .userAgent(userAgent)
-                  .executor(MoreExecutors.directExecutor())
-                  .build());
+      builder = ManagedChannelBuilder.forTarget(url).usePlaintext();
+    } else {
+      builder = Grpc.newChannelBuilder(url, credentials);
+    }
+    if (!header.isEmpty()) {
+      var metadata = new Metadata();
+      header.forEach(
+          (k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v));
+      builder.intercept(MetadataUtils.newAttachHeadersInterceptor(metadata));
     }
-    return provider.computeIfAbsent(
-        serverUrl,
-        url ->
-            Grpc.newChannelBuilder(url, credentials)
-                .disableRetry()
-                .userAgent(userAgent)
-                .executor(MoreExecutors.directExecutor())
-                .build());
+    return builder
+        .disableRetry()
+        .userAgent(userAgent)
+        .executor(MoreExecutors.directExecutor())
+        .build();
   }
 
   @Override

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

@@ -3,14 +3,17 @@
 import static com.google.common.base.Preconditions.*;
 import static io.hstream.util.UrlSchemaUtils.parseServerUrls;
 
+import io.grpc.ChannelCredentials;
 import io.grpc.TlsChannelCredentials;
 import io.hstream.HStreamClient;
 import io.hstream.HStreamClientBuilder;
 import io.hstream.HStreamDBClientException;
 import io.hstream.UrlSchema;
 import java.io.File;
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import org.apache.commons.lang3.tuple.Pair;
 
 public class HStreamClientBuilderImpl implements HStreamClientBuilder {
@@ -21,6 +24,7 @@ public class HStreamClientBuilderImpl implements HStreamClientBuilder {
   boolean enableTlsAuthentication;
   String keyPath;
   String certPath;
+  Map<String, String> metadata = new HashMap<>();
 
   long requestTimeoutMs = DefaultSettings.GRPC_CALL_TIMEOUT_MS;
 
@@ -72,6 +76,12 @@ public HStreamClientBuilder requestTimeoutMs(long timeoutMs) {
     return this;
   }
 
+  @Override
+  public HStreamClientBuilder withMetadata(String key, String value) {
+    metadata.put(key, value);
+    return this;
+  }
+
   @Override
   public HStreamClient build() {
     checkArgument(serviceUrl != null, "HStreamClientBuilder: `serviceUrl` should not be null");
@@ -81,6 +91,9 @@ public HStreamClient build() {
     if (schemaHosts.getKey().equals(UrlSchema.HSTREAMS) && !enableTls) {
       throw new HStreamDBClientException("hstreams url schema should enable tls");
     }
+
+    // tls
+    ChannelCredentials credentials = null;
     if (enableTls) {
       checkArgument(caPath != null, "when TLS is enabled, `caPath` should not be null");
       try {
@@ -94,12 +107,17 @@ public HStreamClient build() {
               keyPath != null, "when TLS authentication is enabled, `keyPath` should not be null");
           credentialsBuilder = credentialsBuilder.keyManager(new File(certPath), new File(keyPath));
         }
-        return new HStreamClientKtImpl(
-            schemaHosts.getRight(), requestTimeoutMs, credentialsBuilder.build(), channelProvider);
+        credentials = credentialsBuilder.build();
       } catch (IOException e) {
         throw new HStreamDBClientException(String.format("invalid tls options, %s", e));
       }
     }
-    return new HStreamClientKtImpl(schemaHosts.getRight(), requestTimeoutMs, null, channelProvider);
+
+    // channel provider
+    ChannelProvider provider = channelProvider;
+    if (provider == null) {
+      provider = new ChannelProviderImpl(credentials, metadata);
+    }
+    return new HStreamClientKtImpl(schemaHosts.getRight(), requestTimeoutMs, provider);
   }
 }

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

@@ -3,7 +3,6 @@ package io.hstream.impl
 import com.google.common.base.Preconditions.checkArgument
 import com.google.common.util.concurrent.MoreExecutors
 import com.google.protobuf.Empty
-import io.grpc.ChannelCredentials
 import io.hstream.BufferedProducerBuilder
 import io.hstream.Cluster
 import io.hstream.Connector
@@ -65,12 +64,10 @@ import kotlin.streams.toList
 class HStreamClientKtImpl(
     private val bootstrapServerUrls: List<String>,
     private val requestTimeoutMs: Long,
-    credentials: ChannelCredentials? = null,
-    channelProvider: ChannelProvider? = null,
+    private val channelProvider: ChannelProvider,
 ) : HStreamClient {
 
     private val logger = LoggerFactory.getLogger(HStreamClientKtImpl::class.java)
-    private var channelProvider: ChannelProvider
 
     private val clusterServerUrls: AtomicReference<List<String>> = AtomicReference(null)
     fun refreshClusterServerUrls() {
@@ -120,11 +117,6 @@ class HStreamClientKtImpl(
     }
 
     init {
-        if (channelProvider == null) {
-            this.channelProvider = ChannelProviderImpl(credentials)
-        } else {
-            this.channelProvider = channelProvider
-        }
         logger.info("client init with bootstrapServerUrls [{}]", bootstrapServerUrls)
         refreshClusterServerUrls()
     }