feat(grpc): provide an option to configure gRPC channel directly

src/main/java/io/weaviate/client6/v1/api/Config.java

@@ -1,12 +1,15 @@
 package io.weaviate.client6.v1.api;
 
+import static java.util.Objects.requireNonNull;
+
 import java.net.URI;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
 
 import javax.net.ssl.TrustManagerFactory;
 
+import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
 import io.weaviate.client6.v1.internal.BuildInfo;
 import io.weaviate.client6.v1.internal.ObjectBuilder;
 import io.weaviate.client6.v1.internal.Proxy;
@@ -26,7 +29,8 @@ public record Config(
     Authentication authentication,
     TrustManagerFactory trustManagerFactory,
     Timeout timeout,
-    Proxy proxy) {
+    Proxy proxy,
+    Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder) {
 
   public static Config of(Function<Custom, ObjectBuilder<Config>> fn) {
     return fn.apply(new Custom()).build();
@@ -43,23 +47,26 @@ private Config(Builder<?> builder) {
         builder.authentication,
         builder.trustManagerFactory,
         builder.timeout,
-        builder.proxy);
+        builder.proxy,
+        builder.channelBuilder);
   }
 
   RestTransportOptions restTransportOptions() {
     return restTransportOptions(null);
   }
 
   RestTransportOptions restTransportOptions(TokenProvider tokenProvider) {
-    return new RestTransportOptions(scheme, httpHost, httpPort, headers, tokenProvider, trustManagerFactory, timeout, proxy);
+    return new RestTransportOptions(scheme, httpHost, httpPort, headers, tokenProvider, trustManagerFactory, timeout,
+        proxy);
   }
 
   GrpcChannelOptions grpcTransportOptions() {
     return grpcTransportOptions(null);
   }
 
   GrpcChannelOptions grpcTransportOptions(TokenProvider tokenProvider) {
-    return new GrpcChannelOptions(scheme, grpcHost, grpcPort, headers, tokenProvider, trustManagerFactory, timeout, proxy);
+    return new GrpcChannelOptions(scheme, grpcHost, grpcPort, headers, tokenProvider, trustManagerFactory, timeout,
+        proxy, channelBuilder);
   }
 
   private abstract static class Builder<SelfT extends Builder<SelfT>> implements ObjectBuilder<Config> {
@@ -74,6 +81,7 @@ private abstract static class Builder<SelfT extends Builder<SelfT>> implements O
     protected Timeout timeout = new Timeout();
     protected Map<String, String> headers = new HashMap<>();
     protected Proxy proxy;
+    protected Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder;
 
     /**
      * Set URL scheme. Subclasses may increase the visibility of this method to
@@ -188,6 +196,22 @@ public SelfT proxy(Proxy proxy) {
       return (SelfT) this;
     }
 
+    /**
+     * Override or extend any of the configurations previously set for the gRPC
+     * channel.
+     *
+     * <br>
+     * This is best used to access the parameters specific to gRPC, which the client
+     * does not expose directly, e.g. keepalive configuration. For other
+     * configurations, like timeout and proxy, prefer using the APIs exposed by this
+     * class.
+     */
+    @SuppressWarnings("unchecked")
+    public SelfT grpcChannelOptions(Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder) {
+      this.channelBuilder = requireNonNull(channelBuilder, "channelBuilder");
+      return (SelfT) this;
+    }
+
     /**
      * Weaviate will use the URL in this header to call Weaviate Embeddings
      * Service if an appropriate vectorizer is configured for collection.

src/main/java/io/weaviate/client6/v1/internal/grpc/DefaultGrpcTransport.java

@@ -2,6 +2,8 @@
 
 import static java.util.Objects.requireNonNull;
 
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.util.OptionalInt;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
@@ -12,6 +14,7 @@
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+
 import io.grpc.HttpConnectProxiedSocketAddress;
 import io.grpc.ManagedChannel;
 import io.grpc.StatusRuntimeException;
@@ -29,12 +32,6 @@
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoBatch.BatchStreamReply;
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoBatch.BatchStreamRequest;
 
-import javax.net.ssl.SSLException;
-import java.net.InetSocketAddress;
-import java.net.SocketAddress;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.TimeUnit;
-
 public final class DefaultGrpcTransport implements GrpcTransport {
   /**
    * ListenableFuture callbacks are executed
@@ -175,6 +172,11 @@ private static ManagedChannel buildChannel(GrpcChannelOptions transportOptions)
     }
 
     channel.intercept(MetadataUtils.newAttachHeadersInterceptor(transportOptions.headers()));
+
+    if (transportOptions.channelBuilder() != null) {
+      channel = transportOptions.channelBuilder().apply(channel);
+    }
+
     return channel.build();
   }
 

src/main/java/io/weaviate/client6/v1/internal/grpc/GrpcChannelOptions.java

@@ -2,38 +2,88 @@
 
 import java.util.Map;
 import java.util.OptionalInt;
+import java.util.function.Function;
 
 import javax.net.ssl.TrustManagerFactory;
 
 import io.grpc.Metadata;
+import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
 import io.weaviate.client6.v1.internal.Proxy;
 import io.weaviate.client6.v1.internal.Timeout;
 import io.weaviate.client6.v1.internal.TokenProvider;
 import io.weaviate.client6.v1.internal.TransportOptions;
 
 public class GrpcChannelOptions extends TransportOptions<Metadata> {
   private final OptionalInt maxMessageSize;
+  private final Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder;
 
-  public GrpcChannelOptions(String scheme, String host, int port, Map<String, String> headers,
-      TokenProvider tokenProvider, TrustManagerFactory tmf, Timeout timeout, Proxy proxy) {
-    this(scheme, host, port, buildMetadata(headers), tokenProvider, tmf, OptionalInt.empty(), timeout, proxy);
+  public GrpcChannelOptions(
+      String scheme,
+      String host,
+      int port,
+      Map<String, String> headers,
+      TokenProvider tokenProvider,
+      TrustManagerFactory tmf,
+      Timeout timeout,
+      Proxy proxy,
+      Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder) {
+    this(
+        scheme,
+        host,
+        port,
+        buildMetadata(headers),
+        tokenProvider,
+        tmf,
+        OptionalInt.empty(),
+        timeout,
+        proxy,
+        channelBuilder);
   }
 
-  private GrpcChannelOptions(String scheme, String host, int port, Metadata headers,
-      TokenProvider tokenProvider, TrustManagerFactory tmf, OptionalInt maxMessageSize, Timeout timeout, Proxy proxy) {
-    super(scheme, host, port, headers, tokenProvider, tmf, timeout, proxy);
+  private GrpcChannelOptions(String scheme,
+      String host,
+      int port,
+      Metadata headers,
+      TokenProvider tokenProvider,
+      TrustManagerFactory tmf,
+      OptionalInt maxMessageSize,
+      Timeout timeout,
+      Proxy proxy,
+      Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder) {
+    super(
+        scheme,
+        host,
+        port,
+        headers,
+        tokenProvider,
+        tmf,
+        timeout,
+        proxy);
     this.maxMessageSize = maxMessageSize;
+    this.channelBuilder = channelBuilder;
   }
 
   public GrpcChannelOptions withMaxMessageSize(int maxMessageSize) {
-    return new GrpcChannelOptions(scheme, host, port, headers, tokenProvider, trustManagerFactory,
-        OptionalInt.of(maxMessageSize), timeout, proxy);
+    return new GrpcChannelOptions(scheme,
+        host,
+        port,
+        headers,
+        tokenProvider,
+        trustManagerFactory,
+        OptionalInt.of(maxMessageSize),
+        timeout,
+        proxy,
+        channelBuilder);
   }
 
   public OptionalInt maxMessageSize() {
     return maxMessageSize;
   }
 
+  public Function<NettyChannelBuilder, NettyChannelBuilder> channelBuilder() {
+    return channelBuilder;
+  }
+
   private static final Metadata buildMetadata(Map<String, String> headers) {
     var metadata = new Metadata();
     for (var header : headers.entrySet()) {