Merge pull request #217 from weaviate/grpc/add-retry-middleware-for-unavailable

Use `nice-grpc-client-middleware-retry` for auto-retry of UNAVAILABLE grpc error

Author: tsmith023

package-lock.json

@@ -55,7 +55,7 @@
     "graphql-request": "^6.1.0",
     "long": "^5.2.3",
     "nice-grpc": "^2.1.9",
-    "nice-grpc-client-middleware-deadline": "^2.0.12",
+    "nice-grpc-client-middleware-retry": "^3.1.9",
     "uuid": "^9.0.1"
   },
   "devDependencies": {

package.json

@@ -6,7 +6,7 @@ import { InternalConnectionParams } from './http.js';
 import { ConsistencyLevel } from '../data/index.js';
 
 import { ChannelCredentials, ChannelOptions, createChannel, createClientFactory, Metadata } from 'nice-grpc';
-import { deadlineMiddleware } from 'nice-grpc-client-middleware-deadline';
+import { retryMiddleware } from 'nice-grpc-client-middleware-retry';
 
 import { HealthCheckResponse_ServingStatus, HealthDefinition } from '../proto/google/health/v1/health.js';
 import { WeaviateDefinition } from '../proto/v1/weaviate.js';
@@ -23,7 +23,7 @@ export interface GrpcConnectionParams extends InternalConnectionParams {
   grpcSecure: boolean;
 }
 
-const clientFactory = createClientFactory().use(deadlineMiddleware);
+const clientFactory = createClientFactory().use(retryMiddleware);
 
 const MAX_GRPC_MESSAGE_LENGTH = 104858000; // 10mb, needs to be synchronized with GRPC server
 

src/connection/grpc.ts

@@ -2,20 +2,21 @@ import { isAbortError } from 'abort-controller-x';
 import { ConsistencyLevel } from '../data/index.js';
 
 import { Metadata } from 'nice-grpc';
+import { RetryOptions } from 'nice-grpc-client-middleware-retry';
 import { WeaviateRequestTimeoutError } from '../errors.js';
 import { ConsistencyLevel as ConsistencyLevelGRPC } from '../proto/v1/base.js';
 import { WeaviateClient } from '../proto/v1/weaviate.js';
 
 export default class Base {
-  protected connection: WeaviateClient;
+  protected connection: WeaviateClient<RetryOptions>;
   protected collection: string;
   protected timeout: number;
   protected consistencyLevel?: ConsistencyLevelGRPC;
   protected tenant?: string;
   protected metadata?: Metadata;
 
   protected constructor(
-    connection: WeaviateClient,
+    connection: WeaviateClient<RetryOptions>,
     collection: string,
     metadata: Metadata,
     timeout: number,

src/grpc/base.ts

@@ -5,11 +5,14 @@ import { ConsistencyLevel } from '../data/index.js';
 import { BatchObject, BatchObjectsReply, BatchObjectsRequest } from '../proto/v1/batch.js';
 import { WeaviateClient } from '../proto/v1/weaviate.js';
 
+import { RetryOptions } from 'nice-grpc-client-middleware-retry';
 import { WeaviateBatchError, WeaviateDeleteManyError } from '../errors.js';
 import { Filters } from '../proto/v1/base.js';
 import { BatchDeleteReply, BatchDeleteRequest } from '../proto/v1/batch_delete.js';
 import Base from './base.js';
 
+import { retryOptions } from './retry.js';
+
 export interface Batch {
   withDelete: (args: BatchDeleteArgs) => Promise<BatchDeleteReply>;
   withObjects: (args: BatchObjectsArgs) => Promise<BatchObjectsReply>;
@@ -27,7 +30,7 @@ export interface BatchDeleteArgs {
 
 export default class Batcher extends Base implements Batch {
   public static use(
-    connection: WeaviateClient,
+    connection: WeaviateClient<RetryOptions>,
     collection: string,
     metadata: Metadata,
     timeout: number,
@@ -70,6 +73,7 @@ export default class Batcher extends Base implements Batch {
           {
             metadata: this.metadata,
             signal,
+            ...retryOptions,
           }
         )
         .catch((err) => {

src/grpc/batcher.ts

@@ -0,0 +1,11 @@
+import { ClientError, Status } from 'nice-grpc';
+import { RetryOptions } from 'nice-grpc-client-middleware-retry';
+
+export const retryOptions: RetryOptions = {
+  retry: true,
+  retryMaxAttempts: 5,
+  retryableStatuses: [Status.UNAVAILABLE],
+  onRetryableError(error: ClientError, attempt: number, delayMs: number) {
+    console.warn(error, `Attempt ${attempt} failed. Retrying in ${delayMs}ms.`);
+  },
+};

src/grpc/retry.ts

@@ -24,9 +24,11 @@ import {
 } from '../proto/v1/search_get.js';
 import { WeaviateClient } from '../proto/v1/weaviate.js';
 
+import { RetryOptions } from 'nice-grpc-client-middleware-retry';
 import { WeaviateQueryError } from '../errors.js';
 import { GenerativeSearch } from '../proto/v1/generative.js';
 import Base from './base.js';
+import { retryOptions } from './retry.js';
 
 export type SearchFetchArgs = {
   limit?: number;
@@ -113,7 +115,7 @@ export interface Search {
 
 export default class Searcher extends Base implements Search {
   public static use(
-    connection: WeaviateClient,
+    connection: WeaviateClient<RetryOptions>,
     collection: string,
     metadata: Metadata,
     timeout: number,
@@ -151,6 +153,7 @@ export default class Searcher extends Base implements Search {
           {
             metadata: this.metadata,
             signal,
+            ...retryOptions,
           }
         )
         .catch((err) => {

src/grpc/searcher.ts

@@ -1,7 +1,9 @@
 import { Metadata } from 'nice-grpc';
+import { RetryOptions } from 'nice-grpc-client-middleware-retry';
 import { TenantsGetReply, TenantsGetRequest } from '../proto/v1/tenants.js';
 import { WeaviateClient } from '../proto/v1/weaviate.js';
 import Base from './base.js';
+import { retryOptions } from './retry.js';
 
 export type TenantsGetArgs = {
   names?: string[];
@@ -11,9 +13,9 @@ export interface Tenants {
   withGet: (args: TenantsGetArgs) => Promise<TenantsGetReply>;
 }
 
-export default class TenantsManager extends Base implements TenantsManager {
+export default class TenantsManager extends Base implements Tenants {
   public static use(
-    connection: WeaviateClient,
+    connection: WeaviateClient<RetryOptions>,
     collection: string,
     metadata: Metadata,
     timeout: number
@@ -34,6 +36,7 @@ export default class TenantsManager extends Base implements TenantsManager {
         {
           metadata: this.metadata,
           signal,
+          ...retryOptions,
         }
       )
     );