diff --git a/src/main/java/io/qdrant/client/QdrantGrpcClient.java b/src/main/java/io/qdrant/client/QdrantGrpcClient.java index e0ecf758..9ec85bf1 100644 --- a/src/main/java/io/qdrant/client/QdrantGrpcClient.java +++ b/src/main/java/io/qdrant/client/QdrantGrpcClient.java @@ -269,6 +269,7 @@ private static ManagedChannel createChannel( } channelBuilder.userAgent(userAgent); + channelBuilder.intercept(new ResourceExhaustedInterceptor()); return channelBuilder.build(); } diff --git a/src/main/java/io/qdrant/client/ResourceExhaustedInterceptor.java b/src/main/java/io/qdrant/client/ResourceExhaustedInterceptor.java new file mode 100644 index 00000000..f2f27f55 --- /dev/null +++ b/src/main/java/io/qdrant/client/ResourceExhaustedInterceptor.java @@ -0,0 +1,49 @@ +package io.qdrant.client; + +import io.grpc.*; + +/** An Interceptor that handles Resource Exhausted errors */ +public class ResourceExhaustedInterceptor implements ClientInterceptor { + /** Default constructor for {@link ResourceExhaustedInterceptor} */ + public ResourceExhaustedInterceptor() {} + + @Override + public ClientCall interceptCall( + MethodDescriptor method, CallOptions callOptions, Channel next) { + return new ForwardingClientCall.SimpleForwardingClientCall( + next.newCall(method, callOptions)) { + + @Override + public void start(Listener responseListener, Metadata headers) { + super.start( + new ForwardingClientCallListener.SimpleForwardingClientCallListener( + responseListener) { + @Override + public void onClose(Status status, Metadata trailers) { + if (status.getCode() == Status.Code.RESOURCE_EXHAUSTED) { + String retryAfter = + trailers.get( + Metadata.Key.of("Retry-After", Metadata.ASCII_STRING_MARSHALLER)); + if (retryAfter != null) { + try { + int retryAfterSeconds = Integer.parseInt(retryAfter); + status = + status.withCause( + new ResourceExhaustedResponse( + "Too many requests", retryAfterSeconds)); + } catch (NumberFormatException e) { + throw new QdrantException( + "Retry-After header value is not a valid integer: " + retryAfter); + } + } else { + super.onClose(status, trailers); + } + } + super.onClose(status, trailers); + } + }, + headers); + } + }; + } +} diff --git a/src/main/java/io/qdrant/client/ResourceExhaustedResponse.java b/src/main/java/io/qdrant/client/ResourceExhaustedResponse.java new file mode 100644 index 00000000..f61f02c3 --- /dev/null +++ b/src/main/java/io/qdrant/client/ResourceExhaustedResponse.java @@ -0,0 +1,27 @@ +package io.qdrant.client; + +/** An exception indicating that rate limit is reached */ +public class ResourceExhaustedResponse extends RuntimeException { + /** The number of seconds to wait before retrying */ + private final int retryAfterSeconds; + + /** + * Instantiates a new instance of {@link ResourceExhaustedResponse} + * + * @param message The message to display + * @param retryAfterSeconds The number of seconds to wait before retrying + */ + public ResourceExhaustedResponse(String message, int retryAfterSeconds) { + super(message); + this.retryAfterSeconds = retryAfterSeconds; + } + + /** + * Gets the number of seconds to wait before retrying + * + * @return the number of seconds to wait before retrying + */ + public int getRetryAfterSeconds() { + return retryAfterSeconds; + } +}