Use keepalive for ETCD Lease Renewals (#1511)

Author: ermontross

astra/pom.xml

@@ -465,6 +465,11 @@
             <version>${jetcd.version}</version>
             <scope>test</scope>
         </dependency>
+        <!-- testcontainers version uses version from <dependencyManagement> -->
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>testcontainers</artifactId>
+        </dependency>
         <!-- Note: next time we upgrade this, check if mockito-core depends on 1.14.8+ version of bytebuddy -->
         <!-- if it does then we don't need to define bytebuddy explicitly anymore -->
         <!-- Run mvn dependency:tree to verify. armeria also has a runtime dependency on bytebuddy so we need to check that as well -->
@@ -586,6 +591,17 @@
 
     </dependencies>
 
+    <dependencyManagement>
+        <dependencies>
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>testcontainers</artifactId>
+            <version>2.0.2</version>
+            <scope>test</scope>
+        </dependency>
+        </dependencies>
+    </dependencyManagement>
+
     <build>
         <defaultGoal>clean install</defaultGoal>
         <extensions>

astra/src/main/java/com/slack/astra/chunkrollover/DiskOrMessageCountBasedRolloverStrategy.java

@@ -88,14 +88,16 @@ public DiskOrMessageCountBasedRolloverStrategy(
             if (dirSize > 0) {
               approximateDirectoryBytes.set(dirSize);
             }
-            if (!maxTimePerChunksMinsReached.get()
-                && Instant.now()
-                    .isAfter(rolloverStartTime.plus(maxTimePerChunksSeconds, ChronoUnit.SECONDS))) {
-              LOG.info(
-                  "Max time per chunk reached. chunkStartTime: {} currentTime: {}",
-                  rolloverStartTime,
-                  Instant.now());
-              maxTimePerChunksMinsReached.set(true);
+            // Only check max time if it's not Long.MAX_VALUE (which would cause overflow)
+            if (!maxTimePerChunksMinsReached.get() && maxTimePerChunksSeconds < Long.MAX_VALUE) {
+              if (Instant.now()
+                  .isAfter(rolloverStartTime.plus(maxTimePerChunksSeconds, ChronoUnit.SECONDS))) {
+                LOG.info(
+                    "Max time per chunk reached. chunkStartTime: {} currentTime: {}",
+                    rolloverStartTime,
+                    Instant.now());
+                maxTimePerChunksMinsReached.set(true);
+              }
             }
           } catch (Exception e) {
             LOG.error("Error calculating directory size", e);

astra/src/main/java/com/slack/astra/metadata/core/EtcdMetadataStore.java

@@ -8,10 +8,12 @@
 import io.etcd.jetcd.KeyValue;
 import io.etcd.jetcd.Watch.Watcher;
 import io.etcd.jetcd.kv.GetResponse;
+import io.etcd.jetcd.lease.LeaseKeepAliveResponse;
 import io.etcd.jetcd.options.GetOption;
 import io.etcd.jetcd.options.PutOption;
 import io.etcd.jetcd.options.WatchOption;
 import io.etcd.jetcd.watch.WatchEvent;
+import io.grpc.stub.StreamObserver;
 import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.MeterRegistry;
 import java.io.Closeable;
@@ -56,8 +58,8 @@ public class EtcdMetadataStore<T extends AstraMetadata> implements Closeable {
   /** Shared lease ID for all ephemeral nodes. Only valid if createMode is EPHEMERAL. */
   private volatile long sharedLeaseId = -1;
 
-  /** Used for refreshing ephemeral node leases. */
-  private final ScheduledExecutorService leaseRefreshExecutor;
+  /** Flag to track if the store is being closed to prevent keepalive restarts during shutdown. */
+  private volatile boolean isClosing = false;
 
   /** Used for watch retry operations with delays. */
   private final ScheduledExecutorService watchRetryExecutor;
@@ -68,8 +70,6 @@ public class EtcdMetadataStore<T extends AstraMetadata> implements Closeable {
   /** TTL in milliseconds for ephemeral nodes. */
   private final long ephemeralTtlMs;
 
-  private final int ephemeralMaxRetries;
-
   private static final Logger LOG = LoggerFactory.getLogger(EtcdMetadataStore.class);
 
   protected final String storeFolder;
@@ -140,7 +140,7 @@ public EtcdMetadataStore(
       MeterRegistry meterRegistry,
       MetadataSerializer<T> serializer,
       EtcdCreateMode createMode,
-      Client etcClient) {
+      Client etcdClient) {
     this.storeFolder = storeFolder;
     this.namespace = config.getNamespace();
     this.meterRegistry = meterRegistry;
@@ -149,7 +149,6 @@ public EtcdMetadataStore(
     this.watchers = new ConcurrentHashMap<>();
     this.createMode = createMode;
     this.ephemeralTtlMs = config.getEphemeralNodeTtlMs();
-    this.ephemeralMaxRetries = config.getEphemeralNodeMaxRetries();
     this.etcdOperationTimeoutMs = config.getOperationsTimeoutMs();
 
     // Store retry configuration for watch operations
@@ -170,15 +169,15 @@ public EtcdMetadataStore(
     this.leaseRefreshHandlerFired =
         this.meterRegistry.counter(ASTRA_ETCD_LEASE_REFRESH_HANDLER_FIRED, "store", store);
 
-    if (etcClient == null) {
+    if (etcdClient == null) {
       throw new IllegalArgumentException("External etcd client must be provided");
     }
 
     LOG.info(
         "Using provided external etcd client for store folder: {} with mode: {}",
         storeFolder,
         createMode);
-    this.etcdClient = etcClient;
+    this.etcdClient = etcdClient;
 
     // Initialize watch retry executor - scheduled cached thread pool that scales to 0
     this.watchRetryExecutor =
@@ -191,28 +190,16 @@ public EtcdMetadataStore(
               return t;
             });
 
-    // Initialize lease refresh executor if we're creating ephemeral nodes
     if (createMode == EtcdCreateMode.EPHEMERAL) {
-      this.leaseRefreshExecutor =
-          Executors.newSingleThreadScheduledExecutor(
-              r -> {
-                Thread t = new Thread(r);
-                t.setDaemon(true);
-                t.setName("etcd-lease-refresh-" + storeFolder);
-                return t;
-              });
-
-      // Calculate the refresh interval - currently 1/3 of the TTL
-      long refreshIntervalMs = ephemeralTtlMs / 3;
-
       // Create a single shared lease for all ephemeral nodes synchronously
       try {
         sharedLeaseId =
-            etcdClient
+            this.etcdClient
                 .getLeaseClient()
                 .grant(ephemeralTtlMs / 1000) // grant ttl is in seconds
                 .get(ephemeralTtlMs, TimeUnit.MILLISECONDS)
                 .getID();
+        startKeepAlive();
       } catch (InterruptedException | ExecutionException | TimeoutException e) {
         throw new RuntimeException(e);
       }
@@ -222,14 +209,6 @@ public EtcdMetadataStore(
           sharedLeaseId,
           Long.toHexString(sharedLeaseId),
           ephemeralTtlMs);
-
-      LOG.info("Starting lease refresh thread with interval: {} ms", refreshIntervalMs);
-
-      // Start the refresh task
-      leaseRefreshExecutor.scheduleWithFixedDelay(
-          this::refreshAllLeases, refreshIntervalMs, refreshIntervalMs, TimeUnit.MILLISECONDS);
-    } else {
-      leaseRefreshExecutor = null;
     }
 
     // Initialize cache if needed
@@ -247,6 +226,52 @@ public EtcdMetadataStore(
     }
   }
 
+  /** Creates KeepAlive GRPC connection and handles error and completed cases */
+  private void startKeepAlive() {
+    this.etcdClient
+        .getLeaseClient()
+        .keepAlive(
+            sharedLeaseId,
+            new StreamObserver<LeaseKeepAliveResponse>() {
+              @Override
+              public void onNext(LeaseKeepAliveResponse response) {
+                LOG.debug(
+                    "Received keepAlive response for lease {}, TTL: {}",
+                    response.getID(),
+                    response.getTTL());
+                leaseRefreshHandlerFired.increment();
+              }
+
+              @Override
+              public void onError(Throwable t) {
+                if (isClosing) {
+                  LOG.debug(
+                      "KeepAlive stream error during shutdown for lease {}, not restarting",
+                      sharedLeaseId);
+                  return;
+                }
+                LOG.error(
+                    "Error in keepAlive stream for shared lease {}: {}",
+                    sharedLeaseId,
+                    t.getMessage());
+                startKeepAlive();
+              }
+
+              @Override
+              public void
+                  onCompleted() { // not the same as a shutdown, we want to restart the connection
+                if (isClosing) {
+                  LOG.debug(
+                      "KeepAlive stream completed during shutdown for lease {}, not restarting",
+                      sharedLeaseId);
+                  return;
+                }
+                LOG.warn("KeepAlive stream completed for shared lease {}", sharedLeaseId);
+                startKeepAlive();
+              }
+            });
+  }
+
   /**
    * Converts a path string to an etcd ByteSequence key.
    *
@@ -1091,49 +1116,12 @@ private void populateInitialCache() {
     // as initialized on success or interruption, not on errors.
   }
 
-  /** Refreshes the shared lease for all ephemeral nodes to prevent them from expiring. */
-  private void refreshAllLeases() {
-    // Skip if lease hasn't been initialized (though with synchronous initialization this shouldn't
-    // happen)
-    if (sharedLeaseId == -1) {
-      LOG.debug("Skipping lease refresh - lease not yet initialized");
-      return;
-    }
-
-    LOG.debug("Refreshing shared lease {} for store {}", sharedLeaseId, storeFolder);
-
-    long retryTimeoutMs = ephemeralTtlMs / ephemeralMaxRetries;
-    int retryCounter = 0;
-    while (retryCounter <= ephemeralMaxRetries) {
-      try {
-        etcdClient
-            .getLeaseClient()
-            .keepAliveOnce(sharedLeaseId)
-            .get(retryTimeoutMs, TimeUnit.MILLISECONDS);
-        LOG.trace("Successfully refreshed shared lease {}", sharedLeaseId);
-        this.leaseRefreshHandlerFired.increment();
-        break;
-      } catch (InterruptedException e) {
-        LOG.warn("Interrupted while refreshing shared lease {}", sharedLeaseId, e);
-        Thread.currentThread().interrupt(); // Preserve interrupt status
-      } catch (Exception e) {
-        retryCounter++;
-        if (retryCounter >= ephemeralMaxRetries) {
-          LOG.error(
-              "Failed to refresh shared lease max times, fataling {}: {}",
-              sharedLeaseId,
-              e.getMessage());
-          // This is a critical error since it affects all ephemeral nodes
-          new RuntimeHalterImpl().handleFatal(e);
-        }
-        LOG.error("Failed to refresh shared lease, retrying {}: {}", sharedLeaseId, e.getMessage());
-      }
-    }
-  }
-
   @Override
   public void close() {
     LOG.info("Closing etcd clients and watchers");
+    // Set closing flag to prevent keepalive restarts during shutdown
+    isClosing = true;
+
     // Close all active watchers
     watchers.values().forEach(Watcher::close);
     watchers.clear();
@@ -1166,24 +1154,15 @@ public void close() {
       }
     }
 
-    // Shut down lease refresh executor if it exists
-    if (leaseRefreshExecutor != null) {
-      leaseRefreshExecutor.shutdownNow();
+    // Revoke the shared lease if we have one
+    if (sharedLeaseId != -1) {
       try {
-        leaseRefreshExecutor.awaitTermination(5, TimeUnit.SECONDS);
-      } catch (InterruptedException ignored) {
-      }
-
-      // Revoke the shared lease if we have one
-      if (sharedLeaseId != -1) {
-        try {
-          LOG.info("Revoking shared lease {}", sharedLeaseId);
-          etcdClient.getLeaseClient().revoke(sharedLeaseId).get(5, TimeUnit.SECONDS);
-        } catch (Exception e) {
-          LOG.warn("Failed to revoke shared lease {}: {}", sharedLeaseId, e.getMessage());
-        } finally {
-          sharedLeaseId = -1;
-        }
+        LOG.info("Revoking shared lease {}", sharedLeaseId);
+        etcdClient.getLeaseClient().revoke(sharedLeaseId).get(5, TimeUnit.SECONDS);
+      } catch (Exception e) {
+        LOG.warn("Failed to revoke shared lease {}: {}", sharedLeaseId, e.getMessage());
+      } finally {
+        sharedLeaseId = -1;
       }
     }
 

astra/src/test/java/com/slack/astra/server/BulkIngestApiTest.java

@@ -227,15 +227,15 @@ public void shutdownOpenSearchAPI() throws Exception {
       bulkIngestKafkaProducer.awaitTerminated(DEFAULT_START_STOP_DURATION);
     }
     kafkaServer.close();
-    if (etcdClient != null) {
-      etcdClient.close();
-    }
 
     datasetMetadataStore.close();
     preprocessorMetadataStore.close();
     curatorFramework.unwrap().close();
     zkServer.close();
     meterRegistry.close();
+    if (etcdClient != null) {
+      etcdClient.close();
+    }
   }
 
   public KafkaConsumer getTestKafkaConsumer() {