2
0
Эх сурвалжийг харах

Preconnect (#1047) (#1051)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot 8 сар өмнө
parent
commit
965b95cde4

+ 38 - 3
src/main/java/io/milvus/v2/client/MilvusClientV2.java

@@ -22,6 +22,9 @@ package io.milvus.v2.client;
 import io.grpc.ManagedChannel;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.milvus.grpc.ClientInfo;
+import io.milvus.grpc.ConnectRequest;
+import io.milvus.grpc.ConnectResponse;
 import io.milvus.grpc.MilvusServiceGrpc;
 import io.milvus.orm.iterator.QueryIterator;
 import io.milvus.orm.iterator.SearchIterator;
@@ -102,11 +105,11 @@ public class MilvusClientV2 {
         }
         channel = clientUtils.getChannel(connectConfig);
 
+        blockingStub = MilvusServiceGrpc.newBlockingStub(channel);
+        connect(connectConfig, blockingStub);
         if (connectConfig.getRpcDeadlineMs() > 0) {
-            blockingStub =  MilvusServiceGrpc.newBlockingStub(channel).withWaitForReady()
+            blockingStub = blockingStub.withWaitForReady()
                     .withDeadlineAfter(connectConfig.getRpcDeadlineMs(), TimeUnit.MILLISECONDS);
-        }else {
-            blockingStub = MilvusServiceGrpc.newBlockingStub(channel);
         }
 
         if (connectConfig.getDbName() != null) {
@@ -115,6 +118,38 @@ public class MilvusClientV2 {
         }
     }
 
+    /**
+     * This method is internal used, it calls a RPC Connect() to the remote server,
+     * and sends the client info to the server so that the server knows which client is interacting,
+     * especially for accesses log.
+     *
+     * The info includes:
+     * 1. username(if Authentication is enabled)
+     * 2. the client computer's name
+     * 3. sdk language type and version
+     * 4. the client's local time
+     */
+    private void connect(ConnectConfig connectConfig, MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub) {
+        String userName = connectConfig.getUsername();
+        if (userName == null) {
+            userName = ""; // ClientInfo.setUser() requires non-null value
+        }
+
+        ClientInfo info = ClientInfo.newBuilder()
+                .setSdkType("Java")
+                .setSdkVersion(clientUtils.getSDKVersion())
+                .setUser(userName)
+                .setHost(clientUtils.getHostName())
+                .setLocalTime(clientUtils.getLocalTimeStr())
+                .build();
+        ConnectRequest req = ConnectRequest.newBuilder().setClientInfo(info).build();
+        ConnectResponse resp = blockingStub.withDeadlineAfter(1, TimeUnit.SECONDS)
+                .connect(req);
+        if (resp.getStatus().getCode() != 0 || !resp.getStatus().getErrorCode().equals(io.milvus.grpc.ErrorCode.Success)) {
+            throw new RuntimeException("Failed to initialize connection. Error: " + resp.getStatus().getReason());
+        }
+    }
+
     public void retryConfig(RetryConfig retryConfig) {
         this.retryConfig = retryConfig;
     }

+ 28 - 0
src/main/java/io/milvus/v2/utils/ClientUtils.java

@@ -26,6 +26,7 @@ import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
 import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext;
 import io.grpc.stub.MetadataUtils;
+import io.milvus.client.MilvusServiceClient;
 import io.milvus.grpc.*;
 import io.milvus.v2.client.ConnectConfig;
 import org.apache.commons.lang3.StringUtils;
@@ -34,7 +35,10 @@ import org.slf4j.LoggerFactory;
 
 import java.io.File;
 import java.io.IOException;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
 import java.nio.charset.StandardCharsets;
+import java.time.LocalDateTime;
 import java.util.Base64;
 import java.util.concurrent.TimeUnit;
 
@@ -132,4 +136,28 @@ public class ClientUtils {
         rpcUtils.handleResponse("Get server version", response.getStatus());
         return response.getVersion();
     }
+
+    public String getHostName() {
+        try {
+            InetAddress address = InetAddress.getLocalHost();
+            return address.getHostName();
+        } catch (UnknownHostException e) {
+            logger.warn("Failed to get host name, error:{}\n", e.getMessage());
+            return "Unknown";
+        }
+    }
+
+    public String getLocalTimeStr() {
+        LocalDateTime now = LocalDateTime.now();
+        return now.toString();
+    }
+
+    public String getSDKVersion() {
+        Package pkg = MilvusServiceClient.class.getPackage();
+        String ver = pkg.getImplementationVersion();
+        if (ver == null) {
+            return "";
+        }
+        return ver;
+    }
 }

+ 1 - 0
src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java

@@ -1633,6 +1633,7 @@ class MilvusClientV2DockerTest {
         try {
             ConnectConfig connectConfig = ConnectConfig.builder()
                     .uri(milvus.getEndpoint())
+                    .rpcDeadlineMs(100L)
                     .build();
             PoolConfig poolConfig = PoolConfig.builder()
                     .build();