feat(batch): reconnect to GCP every 160 seconds

Author: bevzzz

src/main/java/io/weaviate/client6/v1/api/Config.java

@@ -11,6 +11,7 @@
 import io.weaviate.client6.v1.internal.ObjectBuilder;
 import io.weaviate.client6.v1.internal.Timeout;
 import io.weaviate.client6.v1.internal.TokenProvider;
+import io.weaviate.client6.v1.internal.TransportOptions;
 import io.weaviate.client6.v1.internal.grpc.GrpcChannelOptions;
 import io.weaviate.client6.v1.internal.rest.RestTransportOptions;
 
@@ -181,17 +182,6 @@ public SelfT timeout(int initSeconds, int querySeconds, int insertSeconds) {
     private static final String HEADER_X_WEAVIATE_CLUSTER_URL = "X-Weaviate-Cluster-URL";
     private static final String HEADER_X_WEAVIATE_CLIENT = "X-Weaviate-Client";
 
-    /**
-     * isWeaviateDomain returns true if the host matches weaviate.io,
-     * semi.technology, or weaviate.cloud domain.
-     */
-    private static boolean isWeaviateDomain(String host) {
-      var lower = host.toLowerCase();
-      return lower.contains("weaviate.io") ||
-          lower.contains("semi.technology") ||
-          lower.contains("weaviate.cloud");
-    }
-
     private static final String VERSION = "weaviate-client-java/"
         + ((!BuildInfo.TAGS.isBlank() && BuildInfo.TAGS != "unknown") ? BuildInfo.TAGS
             : (BuildInfo.BRANCH + "-" + BuildInfo.COMMIT_ID_ABBREV));
@@ -200,7 +190,7 @@ private static boolean isWeaviateDomain(String host) {
     public Config build() {
       // For clusters hosted on Weaviate Cloud, Weaviate Embedding Service
       // will be available under the same domain.
-      if (isWeaviateDomain(httpHost) && authentication != null) {
+      if (TransportOptions.isWeaviateDomain(httpHost) && authentication != null) {
         setHeader(HEADER_X_WEAVIATE_CLUSTER_URL, "https://" + httpHost + ":" + httpPort);
       }
       setHeader(HEADER_X_WEAVIATE_CLIENT, VERSION);

src/main/java/io/weaviate/client6/v1/api/collections/batch/BatchContext.java

@@ -174,6 +174,7 @@ void setClosing(Exception ex) {
     this.queue = new ArrayBlockingQueue<>(DEFAULT_QUEUE_SIZE);
     this.batch = new Batch(DEFAULT_BATCH_SIZE, maxSizeBytes);
     setState(CLOSED);
+
   }
 
   /** Add {@link WeaviateObject} to the batch. */
@@ -701,9 +702,6 @@ public void onEnter(State prev) {
    * Reconnecting state is entererd either by the server finishing a shutdown
    * and closing it's end of the stream or an unexpected stream hangup.
    *
-   * <p>
-   *
-   *
    * @see Recv#onCompleted graceful server shutdown
    * @see Recv#onError stream hangup
    */
@@ -781,7 +779,6 @@ private void reconnectAfter(long delaySeconds) {
         } catch (ExecutionException e) {
           onEvent(new Event.ClientError(e));
         }
-
       }, delaySeconds, TimeUnit.SECONDS);
     }
   }
@@ -844,4 +841,32 @@ public void onEnter(State prev) {
       closed = true;
     }
   };
+
+  // --------------------------------------------------------------------------
+
+  private final ScheduledExecutorService reconnectExec = Executors.newScheduledThreadPool(1);
+
+  void scheduleReconnect(int reconnectIntervalSeconds) {
+    reconnectExec.scheduleWithFixedDelay(() -> {
+      if (Thread.currentThread().isInterrupted()) {
+        onEvent(Event.SHUTTING_DOWN);
+      }
+      if (Thread.currentThread().isInterrupted()) {
+        onEvent(Event.EOF);
+      }
+
+      // We want to count down from the moment we re-opened the stream,
+      // not from the moment we initialited the sequence.
+      lock.lock();
+      try {
+        while (state != ACTIVE) {
+          stateChanged.await();
+        }
+      } catch (InterruptedException ignored) {
+        // Let the process exit normally.
+      } finally {
+        lock.unlock();
+      }
+    }, reconnectIntervalSeconds, reconnectIntervalSeconds, TimeUnit.SECONDS);
+  }
 }

src/main/java/io/weaviate/client6/v1/api/collections/batch/Event.java

@@ -5,6 +5,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
 
 import io.grpc.Status;
 import io.weaviate.client6.v1.api.collections.batch.Event.Acks;
@@ -20,12 +21,16 @@ sealed interface Event
     permits Started, Acks, Results, Backoff, Oom, TerminationEvent, StreamHangup, ClientError {
 
   final static Event STARTED = new Started();
-  final static Event OOM = TerminationEvent.OOM;
   final static Event SHUTTING_DOWN = TerminationEvent.SHUTTING_DOWN;
   final static Event EOF = TerminationEvent.EOF;
 
-  /** */
-  record Started() implements Event {
+  /**
+   * The server has acknowledged our Start message and is ready to receive data.
+   *
+   * @param reconnectAfterSeconds Delay in seconds after which
+   *                              the stream should be renewed.
+   */
+  record Started(OptionalInt reconnectAfterSeconds) implements Event {
   }
 
   /**

src/main/java/io/weaviate/client6/v1/api/collections/batch/WeaviateBatchClient.java

@@ -5,6 +5,7 @@
 import java.util.OptionalInt;
 
 import io.weaviate.client6.v1.api.collections.CollectionHandleDefaults;
+import io.weaviate.client6.v1.internal.TransportOptions;
 import io.weaviate.client6.v1.internal.grpc.GrpcTransport;
 import io.weaviate.client6.v1.internal.orm.CollectionDescriptor;
 
@@ -34,11 +35,24 @@ public BatchContext<PropertiesT> start() {
     if (maxSizeBytes.isEmpty()) {
       throw new IllegalStateException("Server must have grpcMaxMessageSize configured to use server-side batching");
     }
+
     StreamFactory<Message, Event> streamFactory = new TranslatingStreamFactory(grpcTransport::createStream);
-    return new BatchContext<>(
+    BatchContext<PropertiesT> context = new BatchContext<>(
         streamFactory,
         maxSizeBytes.getAsInt(),
         collectionDescriptor,
         defaults);
+
+    if (isWeaviateCloudOnGoogleCloud(grpcTransport.host())) {
+      context.scheduleReconnect(GCP_RECONNECT_INTERVAL_SECONDS);
+    }
+
+    return context;
+  }
+
+  private static final int GCP_RECONNECT_INTERVAL_SECONDS = 160;
+
+  private static boolean isWeaviateCloudOnGoogleCloud(String host) {
+    return TransportOptions.isWeaviateDomain(host) && TransportOptions.isGoogleCloudDomain(host);
   }
 }

src/main/java/io/weaviate/client6/v1/internal/TransportOptions.java

@@ -57,4 +57,20 @@ public H headers() {
   public TrustManagerFactory trustManagerFactory() {
     return this.trustManagerFactory;
   }
+
+  /**
+   * isWeaviateDomain returns true if the host matches weaviate.io,
+   * semi.technology, or weaviate.cloud domain.
+   */
+  public static boolean isWeaviateDomain(String host) {
+    var lower = host.toLowerCase();
+    return lower.contains("weaviate.io") ||
+        lower.contains("semi.technology") ||
+        lower.contains("weaviate.cloud");
+  }
+
+  public static boolean isGoogleCloudDomain(String host) {
+    var lower = host.toLowerCase();
+    return lower.contains("gcp");
+  }
 }

src/main/java/io/weaviate/client6/v1/internal/grpc/DefaultGrpcTransport.java

@@ -177,4 +177,9 @@ public void close() throws Exception {
       callCredentials.close();
     }
   }
+
+  @Override
+  public String host() {
+    return transportOptions.host();
+  }
 }

src/main/java/io/weaviate/client6/v1/internal/grpc/GrpcTransport.java

@@ -23,6 +23,8 @@ <RequestT, RequestM, ReplyM, ResponseT> CompletableFuture<ResponseT> performRequ
   StreamObserver<WeaviateProtoBatch.BatchStreamRequest> createStream(
       StreamObserver<WeaviateProtoBatch.BatchStreamReply> recv);
 
+  String host();
+
   /**
    * Maximum inbound/outbound message size supported by the underlying channel.
    */

src/test/java/io/weaviate/testutil/transport/MockGrpcTransport.java

@@ -17,6 +17,7 @@
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoBatch.BatchStreamRequest;
 
 public class MockGrpcTransport implements GrpcTransport {
+  private final String host = "example.com";
 
   @FunctionalInterface
   public interface AssertFunction {
@@ -73,4 +74,9 @@ public OptionalInt maxMessageSizeBytes() {
     // TODO(dyma): implement for tests
     throw new UnsupportedOperationException("Unimplemented method 'maxMessageSizeBytes'");
   }
+
+  @Override
+  public String host() {
+    return host;
+  }
 }