Handle new error code and retry (#707)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

src/main/java/io/milvus/client/MilvusClient.java

@@ -24,6 +24,7 @@ import io.milvus.grpc.*;
 import io.milvus.param.LogLevel;
 import io.milvus.param.R;
 import io.milvus.param.RpcStatus;
+import io.milvus.param.RetryParam;
 import io.milvus.param.alias.*;
 import io.milvus.param.bulkinsert.*;
 import io.milvus.param.collection.*;
@@ -54,11 +55,19 @@ public interface MilvusClient {
      */
     MilvusClient withTimeout(long timeout, TimeUnit timeoutUnit);
 
+    /**
+     * Sets the parameters for retry.
+     *
+     * @param retryParam {@link RetryParam}
+     */
+    MilvusClient withRetry(RetryParam retryParam);
+
     /**
      * Number of retry attempts.
      *
      * @param retryTimes     number of retry attempts.
      */
+    @Deprecated
     MilvusClient withRetry(int retryTimes);
 
     /**
@@ -67,6 +76,7 @@ public interface MilvusClient {
      * @param interval     time interval between retry attempts.
      * @param timeUnit     time unit
      */
+    @Deprecated
     MilvusClient withRetryInterval(long interval, TimeUnit timeUnit);
 
     /**

src/main/java/io/milvus/client/MilvusMultiServiceClient.java

@@ -104,6 +104,11 @@ public class MilvusMultiServiceClient implements MilvusClient {
         return clusterFactory.getMaster().getClient().withTimeout(timeout, timeoutUnit);
     }
 
+    @Override
+    public MilvusClient withRetry(RetryParam retryParam) {
+        return clusterFactory.getMaster().getClient().withRetry(retryParam);
+    }
+
     @Override
     public MilvusClient withRetry(int retryTimes) {
         return clusterFactory.getMaster().getClient().withRetry(retryTimes);

src/main/java/io/milvus/client/MilvusServiceClient.java

@@ -20,13 +20,13 @@
 package io.milvus.client;
 
 import io.grpc.*;
+import io.grpc.Status;
 import io.grpc.stub.MetadataUtils;
+import io.milvus.exception.MilvusException;
+import io.milvus.exception.ServerException;
 import io.milvus.grpc.*;
-import io.milvus.param.ConnectParam;
+import io.milvus.param.*;
 
-import io.milvus.param.LogLevel;
-import io.milvus.param.R;
-import io.milvus.param.RpcStatus;
 import io.milvus.param.alias.AlterAliasParam;
 import io.milvus.param.alias.CreateAliasParam;
 import io.milvus.param.alias.DropAliasParam;
@@ -72,8 +72,7 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
     private final MilvusServiceGrpc.MilvusServiceFutureStub futureStub;
     private final long rpcDeadlineMs;
     private long timeoutMs = 0;
-    private int retryTimes = 0;
-    private long retryIntervalMs = 500L;
+    private RetryParam retryParam = RetryParam.newBuilder().build();
 
     public MilvusServiceClient(@NonNull ConnectParam connectParam) {
         this.rpcDeadlineMs = connectParam.getRpcDeadlineMs();
@@ -158,9 +157,8 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
         this.futureStub = src.futureStub;
         this.rpcDeadlineMs = src.rpcDeadlineMs;
         this.timeoutMs = src.timeoutMs;
-        this.retryTimes = src.retryTimes;
-        this.retryIntervalMs = src.retryIntervalMs;
         this.logLevel = src.logLevel;
+        this.retryParam = src.retryParam;
     }
 
     @Override
@@ -227,6 +225,13 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
         return newClient;
     }
 
+    @Override
+    public MilvusClient withRetry(RetryParam retryParam) {
+        MilvusServiceClient newClient = new MilvusServiceClient(this);
+        newClient.retryParam = retryParam;
+        return newClient;
+    }
+
     @Override
     public MilvusClient withRetry(int retryTimes) {
         if (retryTimes <= 0) {
@@ -234,7 +239,7 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
         }
 
         MilvusServiceClient newClient = new MilvusServiceClient(this);
-        newClient.retryTimes = retryTimes;
+        newClient.retryParam.setMaxRetryTimes(retryTimes);
         return newClient;
     }
 
@@ -244,14 +249,17 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
             return this;
         }
 
+        // to compatible with the old behavior
         MilvusServiceClient newClient = new MilvusServiceClient(this);
-        newClient.retryIntervalMs = timeUnit.toMillis(interval);
+        newClient.retryParam.setInitialBackOffMs(timeUnit.toMillis(interval));
+        newClient.retryParam.setMaxBackOffMs(timeUnit.toMillis(interval));
         return newClient;
     }
 
     private <T> R<T> retry(Callable<R<T>> callable) {
+        int maxRetryTimes = this.retryParam.getMaxRetryTimes();
         // no retry, direct call the method
-        if (this.retryTimes <= 1) {
+        if (maxRetryTimes <= 1) {
             try {
                 return callable.call();
             } catch (Exception e) {
@@ -261,36 +269,83 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
 
         // method to check timeout
         long begin = System.currentTimeMillis();
-        Callable<Void> timeoutChecker = ()->{
+        Callable<Boolean> timeoutChecker = ()->{
             long current = System.currentTimeMillis();
             long cost = (current - begin);
             if (this.timeoutMs > 0 && cost >= this.timeoutMs) {
-                String msg = String.format("Retry timeout: %dms", this.timeoutMs);
-                throw new RuntimeException(msg);
+                return Boolean.TRUE;
             }
-            return null;
+            return Boolean.FALSE;
         };
 
         // retry within timeout
-        for (int i = 0; i < this.retryTimes; i++) {
+        long retryIntervalMs = this.retryParam.getInitialBackOffMs();
+        for (int k = 1; k <= maxRetryTimes; k++) {
             try {
                 R<T> resp = callable.call();
                 if (resp.getStatus() == R.Status.Success.getCode()) {
                     return resp;
                 }
 
-                if (i != this.retryTimes-1) {
-                    timeoutChecker.call();
-                    TimeUnit.MILLISECONDS.sleep(this.retryIntervalMs);
-                    timeoutChecker.call();
-                    logInfo(String.format("Retry again after %dms...", this.retryIntervalMs));
+                Exception e = resp.getException();
+                if (e instanceof StatusRuntimeException) {
+                    // for rpc exception, some error cannot be retried
+                    StatusRuntimeException rpcException = (StatusRuntimeException)e;
+                    Status.Code code = rpcException.getStatus().getCode();
+                    if (code == Status.DEADLINE_EXCEEDED.getCode()
+                            || code == Status.PERMISSION_DENIED.getCode()
+                            || code == Status.UNAUTHENTICATED.getCode()
+                            || code == Status.INVALID_ARGUMENT.getCode()
+                            || code == Status.ALREADY_EXISTS.getCode()
+                            || code == Status.RESOURCE_EXHAUSTED.getCode()) {
+                        return resp;
+                    }
+
+                    if (timeoutChecker.call() == Boolean.TRUE) {
+                        String msg = String.format("Retry timeout: %dms, maxRetry:%d, retries: %d, reason: %s",
+                                this.timeoutMs, maxRetryTimes, k, e.getMessage());
+                        throw new MilvusException(msg, code.value());
+                    }
+                } else if (e instanceof ServerException) {
+                    ServerException serverException = (ServerException)e;
+                    if (timeoutChecker.call() == Boolean.TRUE) {
+                        String msg = String.format("Retry timeout: %dms, maxRetry:%d, retries: %d, reason: %s",
+                                this.timeoutMs, maxRetryTimes, k, e.getMessage());
+                        throw new MilvusException(msg, serverException.getStatus());
+                    }
+
+                    // for server-side returned error, only retry for rate limit
+                    // in new error codes of v2.3, rate limit error value is 8
+                    if (!(serverException.getCompatibleCode() == ErrorCode.RateLimit
+                            || serverException.getStatus() == 8)) {
+                        return resp;
+                    }
+                } else {
+                    return resp;
+                }
+
+                // print log, follow the pymilvus logic
+                if (k > 3) {
+                    logWarning(String.format("Retry(%d) with interval %dms. Reason: %s",
+                            k, retryIntervalMs, e.getMessage()));
+                }
+
+                // sleep for interval
+                if (k != maxRetryTimes) {
+                    TimeUnit.MILLISECONDS.sleep(retryIntervalMs);
+                }
+
+                // reset the next interval value
+                retryIntervalMs = retryIntervalMs*this.retryParam.getBackOffMultiplier();
+                if (retryIntervalMs > this.retryParam.getMaxBackOffMs()) {
+                    retryIntervalMs = this.retryParam.getMaxBackOffMs();
                 }
             } catch (Exception e) {
                 logError(e.getMessage());
                 return R.failed(e);
             }
         }
-        String msg = String.format("Retry run out of %d retry times", this.retryTimes);
+        String msg = String.format("Finish %d retry times, stop retry", maxRetryTimes);
         logError(msg);
         return R.failed(new RuntimeException(msg));
     }

src/main/java/io/milvus/exception/MilvusException.java

@@ -22,7 +22,7 @@ package io.milvus.exception;
 /**
  * Base class of Milvus exceptions.
  */
-public abstract class MilvusException extends RuntimeException {
+public class MilvusException extends RuntimeException {
     protected Integer status;
 
     public MilvusException(String msg, Integer status) {

src/main/java/io/milvus/exception/ServerException.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.milvus.exception;
+
+import io.milvus.grpc.ErrorCode;
+import io.milvus.param.R;
+
+/**
+ * Exception for error response from server side.
+ */
+public class ServerException extends MilvusException {
+    protected ErrorCode compatibleCode;
+    public ServerException(String msg, Integer code, ErrorCode compatibleCode) {
+        super(msg, code);
+        this.compatibleCode = compatibleCode;
+    }
+
+    public ErrorCode getCompatibleCode() {
+        return compatibleCode;
+    }
+}

src/main/java/io/milvus/param/R.java

@@ -73,7 +73,6 @@ public class R<T> {
             r.setStatus(e.getStatus());
         } else {
             r.setStatus(Status.Unknown.getCode());
-            r.exception = exception;
         }
         r.setException(exception);
         return r;
@@ -138,33 +137,33 @@ public class R<T> {
         // Server side error
         Success(0),
         UnexpectedError(1),
-        ConnectFailed(2),
-        PermissionDenied(3),
-        CollectionNotExists(4),
+        ConnectFailed(2),           // Deprecated
+        PermissionDenied(3),        // Deprecated
+        CollectionNotExists(4),     // Deprecated
         IllegalArgument(5),
-        IllegalDimension(7),
-        IllegalIndexType(8),
-        IllegalCollectionName(9),
-        IllegalTOPK(10),
-        IllegalRowRecord(11),
-        IllegalVectorID(12),
-        IllegalSearchResult(13),
-        FileNotFound(14),
-        MetaFailed(15),
-        CacheFailed(16),
-        CannotCreateFolder(17),
-        CannotCreateFile(18),
-        CannotDeleteFolder(19),
-        CannotDeleteFile(20),
-        BuildIndexError(21),
-        IllegalNLIST(22),
-        IllegalMetricType(23),
-        OutOfMemory(24),
-        IndexNotExist(25),
+        IllegalDimension(7),        // Deprecated
+        IllegalIndexType(8),        // Deprecated
+        IllegalCollectionName(9),   // Deprecated
+        IllegalTOPK(10),            // Deprecated
+        IllegalRowRecord(11),       // Deprecated
+        IllegalVectorID(12),        // Deprecated
+        IllegalSearchResult(13),    // Deprecated
+        FileNotFound(14),           // Deprecated
+        MetaFailed(15),             // Deprecated
+        CacheFailed(16),            // Deprecated
+        CannotCreateFolder(17),     // Deprecated
+        CannotCreateFile(18),       // Deprecated
+        CannotDeleteFolder(19),     // Deprecated
+        CannotDeleteFile(20),       // Deprecated
+        BuildIndexError(21),        // Deprecated
+        IllegalNLIST(22),           // Deprecated
+        IllegalMetricType(23),      // Deprecated
+        OutOfMemory(24),            // Deprecated
+        IndexNotExist(25),          // Deprecated
         EmptyCollection(26),
 
         // internal error code.
-        DDRequestRace(1000),
+        DDRequestRace(1000),        // Deprecated
 
         // Client side error
         RpcError(-1),

src/main/java/io/milvus/param/RetryParam.java

@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.milvus.param;
+
+import io.milvus.exception.ParamException;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.Setter;
+import lombok.ToString;
+
+
+/**
+ * Parameters for retry on failure.
+ */
+@Getter
+@Setter
+@ToString
+public class RetryParam {
+    private int maxRetryTimes;
+    private long initialBackOffMs;
+    private long maxBackOffMs;
+    private int backOffMultiplier;
+    private boolean retryOnRateLimie;
+
+    protected RetryParam(@NonNull Builder builder) {
+        this.maxRetryTimes = builder.maxRetryTimes;
+        this.initialBackOffMs = builder.initialBackOffMs;
+        this.maxBackOffMs = builder.maxBackOffMs;
+        this.backOffMultiplier = builder.backOffMultiplier;
+        this.retryOnRateLimie = builder.retryOnRateLimie;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Builder for {@link ConnectParam}
+     */
+    @Getter
+    public static class Builder {
+        private int maxRetryTimes = 75;
+        private long initialBackOffMs = 10;
+        private long maxBackOffMs = 3000;
+        private int backOffMultiplier = 3;
+        private boolean retryOnRateLimie = true;
+
+        protected Builder() {
+        }
+
+        /**
+         * Sets the max retry times on failure.Default value is 75.
+         *
+         * @param maxRetryTimes the maxinum times to retry
+         * @return <code>Builder</code>
+         */
+        public Builder withMaxRetryTimes(int maxRetryTimes) {
+            this.maxRetryTimes = maxRetryTimes;
+            return this;
+        }
+
+        /**
+         * Sets the first time interval between two retries, units: millisecond. Default value is 10ms.
+         *
+         * @param initialBackOffMs time interval value
+         * @return <code>Builder</code>
+         */
+        public Builder withInitialBackOffMs(long initialBackOffMs) {
+            this.initialBackOffMs = initialBackOffMs;
+            return this;
+        }
+
+        /**
+         * Sets the maximum time interval between two retries, units: millisecond. Default value is 3000ms.
+         *
+         * @param maxBackOffMs time interval value
+         * @return <code>Builder</code>
+         */
+        public Builder withMaxBackOffMs(long maxBackOffMs) {
+            this.maxBackOffMs = maxBackOffMs;
+            return this;
+        }
+
+        /**
+         * Sets multiplier to increase time interval after each retry. Default value is 3.
+         *
+         * @param backOffMultiplier the multiplier to increase time interval after each retry
+         * @return <code>Builder</code>
+         */
+        public Builder withBackOffMultiplier(int backOffMultiplier) {
+            this.backOffMultiplier = backOffMultiplier;
+            return this;
+        }
+
+        /**
+         * Sets whether to retry when the returned error is rate limit.Default value is true.
+         *
+         * @param retryOnRateLimie whether to retry when the returned error is rate limit
+         * @return <code>Builder</code>
+         */
+        public Builder withRetryOnRateLimie(boolean retryOnRateLimie) {
+            this.retryOnRateLimie = retryOnRateLimie;
+            return this;
+        }
+
+        /**
+         * Verifies parameters and creates a new {@link RetryParam} instance.
+         *
+         * @return {@link RetryParam}
+         */
+        public RetryParam build() throws ParamException {
+            if (maxRetryTimes <= 0L) {
+                throw new ParamException("Max retry time value must be positive!");
+            }
+
+            if (initialBackOffMs <= 0L) {
+                throw new ParamException("The initial time interval must be positive!");
+            }
+
+            if (maxBackOffMs <= 0L) {
+                throw new ParamException("The max time interval must be positive!");
+            }
+
+            if (backOffMultiplier <= 0L) {
+                throw new ParamException("The multiplier to increase time interval must be positive!");
+            }
+
+            return new RetryParam(this);
+        }
+    }
+}

src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -163,8 +163,13 @@ class MilvusClientDockerTest {
     public static void setUp() {
         startDockerContainer();
 
-        ConnectParam connectParam = connectParamBuilder().withAuthorization("root", "Milvus").build();
-        client = new MilvusServiceClient(connectParam);
+        ConnectParam connectParam = connectParamBuilder()
+                .withAuthorization("root", "Milvus")
+                .build();
+        RetryParam retryParam = RetryParam.newBuilder()
+                .withMaxRetryTimes(10)
+                .build();
+        client = new MilvusServiceClient(connectParam).withRetry(retryParam).withTimeout(10, TimeUnit.SECONDS);
         generator = new RandomStringGenerator.Builder().withinRange('a', 'z').build();
     }
 

src/test/java/io/milvus/client/MilvusServiceClientTest.java

@@ -60,17 +60,20 @@ class MilvusServiceClientTest {
         return mockServer;
     }
 
-    private MilvusServiceClient startClient() {
+    private MilvusClient startClient() {
         String testHost = "localhost";
         ConnectParam connectParam = ConnectParam.newBuilder()
                 .withHost(testHost)
                 .withPort(testPort)
                 .build();
-        return new MilvusServiceClient(connectParam);
+        RetryParam retryParam = RetryParam.newBuilder()
+                .withMaxRetryTimes(2)
+                .build();
+        return new MilvusServiceClient(connectParam).withRetry(retryParam);
     }
 
     @SuppressWarnings("unchecked")
-    private <T, P> void invokeFunc(Method testFunc, MilvusServiceClient client, T param, int ret, boolean equalRet) {
+    private <T, P> void invokeFunc(Method testFunc, MilvusClient client, T param, int ret, boolean equalRet) {
         try {
             R<P> resp = (R<P>) testFunc.invoke(client, param);
             if (equalRet) {
@@ -88,7 +91,7 @@ class MilvusServiceClientTest {
     private <T, P> void testFuncByName(String funcName, T param) {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
         try {
             Class<?> clientClass = MilvusServiceClient.class;
             Method testFunc = clientClass.getMethod(funcName, param.getClass());
@@ -119,7 +122,7 @@ class MilvusServiceClientTest {
     private <T, P> void testAsyncFuncByName(String funcName, T param) {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
 
         try {
             Class<?> clientClass = MilvusServiceClient.class;
@@ -498,7 +501,7 @@ class MilvusServiceClientTest {
     void getCollectionStatistics() {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
 
         try {
             final String collectionName = "collection1";
@@ -632,7 +635,7 @@ class MilvusServiceClientTest {
     void loadCollection() {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
 
         String collectionName = "collection1";
         LoadCollectionParam param = LoadCollectionParam.newBuilder()
@@ -1009,7 +1012,7 @@ class MilvusServiceClientTest {
     void loadPartitions() {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
 
         String collectionName = "collection1";
         String partitionName = "partition1";
@@ -1324,7 +1327,7 @@ class MilvusServiceClientTest {
     void createIndex() {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
 
         // createIndex() calls describeCollection() to check input
         CollectionSchema schema = CollectionSchema.newBuilder()
@@ -1513,7 +1516,7 @@ class MilvusServiceClientTest {
     void dropIndex() {
         // start mock server
         MockMilvusServer server = startServer();
-        MilvusServiceClient client = startClient();
+        MilvusClient client = startClient();
 
         DropIndexParam param = DropIndexParam.newBuilder()
                 .withCollectionName("collection1")
@@ -1720,7 +1723,7 @@ class MilvusServiceClientTest {
         {
             // start mock server
             MockMilvusServer server = startServer();
-            MilvusServiceClient client = startClient();
+            MilvusClient client = startClient();
 
             // test return ok with correct input
             mockServerImpl.setDescribeCollectionResponse(DescribeCollectionResponse.newBuilder()
@@ -1840,7 +1843,7 @@ class MilvusServiceClientTest {
         {
             // start mock server
             MockMilvusServer server = startServer();
-            MilvusServiceClient client = startClient();
+            MilvusClient client = startClient();
 
             // test return ok with insertAsync
             try {