Merge pull request #524 from alex268/master

Added a new way to use channel initializer

Author: alex268

core/src/main/java/tech/ydb/core/grpc/GrpcTransportBuilder.java

@@ -1,6 +1,8 @@
 package tech.ydb.core.grpc;
 
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
@@ -15,6 +17,7 @@
 import com.google.common.net.HostAndPort;
 import com.google.common.util.concurrent.MoreExecutors;
 import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
 
 import tech.ydb.auth.AuthRpcProvider;
 import tech.ydb.auth.NopAuthProvider;
@@ -69,6 +72,7 @@ public enum InitMode {
     private byte[] cert = null;
     private boolean useTLS = false;
     private ManagedChannelFactory.Builder channelFactoryBuilder = null;
+    private final List<Consumer<ManagedChannelBuilder<?>>> channelInitializers = new ArrayList<>();
     private Supplier<ScheduledExecutorService> schedulerFactory = YdbSchedulerFactory::createScheduler;
     private String localDc;
     private BalancingSettings balancingSettings;
@@ -183,6 +187,10 @@ public ManagedChannelFactory getManagedChannelFactory() {
         return channelFactoryBuilder.buildFactory(this);
     }
 
+    public List<Consumer<ManagedChannelBuilder<?>>> getChannelInitializers() {
+        return this.channelInitializers;
+    }
+
     /**
      * Set a custom factory of {@link ManagedChannel}. This option must be used only if you want to configure
      * grpc channels in a special way.
@@ -195,6 +203,17 @@ public GrpcTransportBuilder withChannelFactoryBuilder(ManagedChannelFactory.Buil
         return this;
     }
 
+    /**
+     * Add a custom initialization of {@link ManagedChannelBuilder}
+     *
+     * @param ci custom ManagedChannelBuilder initializer
+     * @return this
+     */
+    public GrpcTransportBuilder addChannelInitializer(Consumer<ManagedChannelBuilder<?>> ci) {
+        channelInitializers.add(ci);
+        return this;
+    }
+
     /**
      * Set a custom initialization of {@link io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder} <br>
      * This method is deprecated. Use

core/src/main/java/tech/ydb/core/impl/pool/NettyChannelFactory.java

@@ -1,13 +1,15 @@
 package tech.ydb.core.impl.pool;
 
 import java.io.ByteArrayInputStream;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 import javax.net.ssl.SSLException;
 
 import io.grpc.ClientInterceptor;
 import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
 import io.grpc.Metadata;
 import io.grpc.internal.DnsNameResolverProvider;
 import io.grpc.netty.GrpcSslContexts;
@@ -39,6 +41,7 @@ public class NettyChannelFactory implements ManagedChannelFactory {
     private final long connectTimeoutMs;
     private final boolean useDefaultGrpcResolver;
     private final Long grpcKeepAliveTimeMillis;
+    private final List<Consumer<ManagedChannelBuilder<?>>> initializers;
 
     private NettyChannelFactory(GrpcTransportBuilder builder) {
         this.database = builder.getDatabase();
@@ -49,6 +52,7 @@ private NettyChannelFactory(GrpcTransportBuilder builder) {
         this.connectTimeoutMs = builder.getConnectTimeoutMillis();
         this.useDefaultGrpcResolver = builder.useDefaultGrpcResolver();
         this.grpcKeepAliveTimeMillis = builder.getGrpcKeepAliveTimeMillis();
+        this.initializers = builder.getChannelInitializers();
     }
 
     @Override
@@ -97,6 +101,10 @@ public ManagedChannel newManagedChannel(String host, int port, String sslHostOve
             channelBuilder.disableRetry();
         }
 
+        for (Consumer<ManagedChannelBuilder<?>> initializer: initializers) {
+            initializer.accept(channelBuilder);
+        }
+
         configure(channelBuilder);
 
         return channelBuilder.build();

core/src/main/java/tech/ydb/core/impl/pool/ShadedNettyChannelFactory.java

@@ -1,13 +1,15 @@
 package tech.ydb.core.impl.pool;
 
 import java.io.ByteArrayInputStream;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 import javax.net.ssl.SSLException;
 
 import io.grpc.ClientInterceptor;
 import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
 import io.grpc.Metadata;
 import io.grpc.internal.DnsNameResolverProvider;
 import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
@@ -39,6 +41,7 @@ public class ShadedNettyChannelFactory implements ManagedChannelFactory {
     private final long connectTimeoutMs;
     private final boolean useDefaultGrpcResolver;
     private final Long grpcKeepAliveTimeMillis;
+    private final List<Consumer<ManagedChannelBuilder<?>>> initializers;
 
     public ShadedNettyChannelFactory(GrpcTransportBuilder builder) {
         this.database = builder.getDatabase();
@@ -49,6 +52,7 @@ public ShadedNettyChannelFactory(GrpcTransportBuilder builder) {
         this.connectTimeoutMs = builder.getConnectTimeoutMillis();
         this.useDefaultGrpcResolver = builder.useDefaultGrpcResolver();
         this.grpcKeepAliveTimeMillis = builder.getGrpcKeepAliveTimeMillis();
+        this.initializers = builder.getChannelInitializers();
     }
 
     @Override
@@ -97,6 +101,10 @@ public ManagedChannel newManagedChannel(String host, int port, String sslHostOve
             channelBuilder.disableRetry();
         }
 
+        for (Consumer<ManagedChannelBuilder<?>> initializer: initializers) {
+            initializer.accept(channelBuilder);
+        }
+
         configure(channelBuilder);
 
         return channelBuilder.build();

core/src/test/java/tech/ydb/core/impl/pool/DefaultChannelFactoryTest.java

@@ -149,6 +149,34 @@ public void customChannelInitializer() {
         Mockito.verify(channelBuilderMock, Mockito.times(1)).enableFullStreamDecompression();
     }
 
+    @Test
+    public void addChannelInitializerTest() {
+        GrpcTransportBuilder builder = GrpcTransport.forHost(MOCKED_HOST, MOCKED_PORT, "/Root")
+                .withUseDefaultGrpcResolver(true)
+                .addChannelInitializer(ci -> ci.usePlaintext())
+                .addChannelInitializer(ci -> ci.userAgent("TEST"));
+
+        ManagedChannelFactory factory = ShadedNettyChannelFactory.build()
+                .buildFactory(builder);
+
+        channelStaticMock.verify(FOR_ADDRESS, Mockito.times(0));
+
+        Assert.assertSame(channelMock, factory.newManagedChannel(MOCKED_HOST, MOCKED_PORT, null));
+
+        channelStaticMock.verify(FOR_ADDRESS, Mockito.times(1));
+
+        Mockito.verify(channelBuilderMock, Mockito.times(1)).negotiationType(NegotiationType.PLAINTEXT);
+        Mockito.verify(channelBuilderMock, Mockito.times(1))
+                .maxInboundMessageSize(ShadedNettyChannelFactory.INBOUND_MESSAGE_SIZE);
+        Mockito.verify(channelBuilderMock, Mockito.times(0))
+                .defaultLoadBalancingPolicy(ShadedNettyChannelFactory.DEFAULT_BALANCER_POLICY);
+        Mockito.verify(channelBuilderMock, Mockito.times(1))
+                .withOption(ChannelOption.ALLOCATOR, ByteBufAllocator.DEFAULT);
+        Mockito.verify(channelBuilderMock, Mockito.times(1)).withOption(ChannelOption.TCP_NODELAY, Boolean.TRUE);
+        Mockito.verify(channelBuilderMock, Mockito.times(1)).usePlaintext();
+        Mockito.verify(channelBuilderMock, Mockito.times(1)).userAgent("TEST");
+    }
+
     @Test
     public void customSslFactory() throws CertificateException, IOException {
         SelfSignedCertificate selfSignedCert = new SelfSignedCertificate(MOCKED_HOST);

pom.xml

@@ -238,6 +238,9 @@
             <id>jdk8-bootstrap</id>
             <activation>
                 <jdk>[9</jdk>
+                <property>
+                    <name>argLine</name>
+                </property>
             </activation>
             <build>
                 <plugins>