Check the connection before connecting as lettuce does the retry automatically

Author: Lavkesh Lahngir

ingestion/src/main/java/feast/store/serving/redis/RedisCustomIO.java

@@ -189,11 +189,11 @@ public static class WriteDoFn extends DoFn<RedisMutation, FailedElement> {
       private final List<RedisMutation> mutations = new ArrayList<>();
       private int batchSize = DEFAULT_BATCH_SIZE;
       private int timeout = DEFAULT_TIMEOUT;
-      private RedisIngestionClient ingestionClient;
+      private RedisIngestionClient redisIngestionClient;
 
       WriteDoFn(StoreProto.Store store) {
         if (store.getType() == StoreProto.Store.StoreType.REDIS)
-          this.ingestionClient = new RedisStandaloneIngestionClient(store.getRedisConfig());
+          this.redisIngestionClient = new RedisStandaloneIngestionClient(store.getRedisConfig());
       }
 
       public WriteDoFn withBatchSize(int batchSize) {
@@ -212,38 +212,38 @@ public WriteDoFn withTimeout(int timeout) {
 
       @Setup
       public void setup() {
-        this.ingestionClient.setup();
+        this.redisIngestionClient.setup();
       }
 
       @StartBundle
       public void startBundle() {
         try {
-          ingestionClient.connect();
+          redisIngestionClient.connect();
         } catch (RedisConnectionException e) {
           log.error("Connection to redis cannot be established ", e);
         }
         mutations.clear();
       }
 
       private void executeBatch() throws Exception {
-        this.ingestionClient
+        this.redisIngestionClient
             .getBackOffExecutor()
             .execute(
                 new Retriable() {
                   @Override
                   public void execute() throws ExecutionException, InterruptedException {
-                    if (!ingestionClient.isConnected()) {
-                      ingestionClient.connect();
+                    if (!redisIngestionClient.isConnected()) {
+                      redisIngestionClient.connect();
                     }
                     mutations.forEach(
                         mutation -> {
                           writeRecord(mutation);
                           if (mutation.getExpiryMillis() != null
                               && mutation.getExpiryMillis() > 0) {
-                            ingestionClient.pexpire(mutation.getKey(), mutation.getExpiryMillis());
+                            redisIngestionClient.pexpire(mutation.getKey(), mutation.getExpiryMillis());
                           }
                         });
-                    ingestionClient.sync();
+                    redisIngestionClient.sync();
                     mutations.clear();
                   }
 
@@ -290,22 +290,22 @@ public void processElement(ProcessContext context) {
       private void writeRecord(RedisMutation mutation) {
         switch (mutation.getMethod()) {
           case APPEND:
-            ingestionClient.append(mutation.getKey(), mutation.getValue());
+            redisIngestionClient.append(mutation.getKey(), mutation.getValue());
             return;
           case SET:
-            ingestionClient.set(mutation.getKey(), mutation.getValue());
+            redisIngestionClient.set(mutation.getKey(), mutation.getValue());
             return;
           case LPUSH:
-            ingestionClient.lpush(mutation.getKey(), mutation.getValue());
+            redisIngestionClient.lpush(mutation.getKey(), mutation.getValue());
             return;
           case RPUSH:
-            ingestionClient.rpush(mutation.getKey(), mutation.getValue());
+            redisIngestionClient.rpush(mutation.getKey(), mutation.getValue());
             return;
           case SADD:
-            ingestionClient.sadd(mutation.getKey(), mutation.getValue());
+            redisIngestionClient.sadd(mutation.getKey(), mutation.getValue());
             return;
           case ZADD:
-            ingestionClient.zadd(mutation.getKey(), mutation.getScore(), mutation.getValue());
+            redisIngestionClient.zadd(mutation.getKey(), mutation.getScore(), mutation.getValue());
             return;
           default:
             throw new UnsupportedOperationException(
@@ -333,7 +333,7 @@ public void finishBundle(FinishBundleContext context)
 
       @Teardown
       public void teardown() {
-        ingestionClient.shutdown();
+        redisIngestionClient.shutdown();
       }
     }
   }

ingestion/src/main/java/feast/store/serving/redis/RedisStandaloneIngestionClient.java

@@ -63,8 +63,10 @@ public void shutdown() {
 
   @Override
   public void connect() {
-    this.connection = this.redisclient.connect(new ByteArrayCodec());
-    this.commands = connection.async();
+    if (!isConnected()) {
+      this.connection = this.redisclient.connect(new ByteArrayCodec());
+      this.commands = connection.async();
+    }
   }
 
   @Override
@@ -75,6 +77,7 @@ public boolean isConnected() {
   @Override
   public void sync() {
     // Wait for some time for futures to complete
+    // TODO: should this be configurable?
     LettuceFutures.awaitAll(60, TimeUnit.SECONDS, futures.toArray(new RedisFuture[0]));
     futures.clear();
   }