Enhance thread-safety in ClientSideCaching key retrieval redis#2402

Added per-key locking mechanism to ensure that only a single thread
fetches a value from Redis cache for a specific key. This prevents
redundant Redis server calls under high concurrent load.

Author: yybmion

src/main/java/io/lettuce/core/support/caching/ClientSideCaching.java

@@ -1,14 +1,15 @@
 package io.lettuce.core.support.caching;
 
-import java.util.List;
-import java.util.concurrent.Callable;
-import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.function.Consumer;
-
 import io.lettuce.core.StatefulRedisConnectionImpl;
 import io.lettuce.core.TrackingArgs;
 import io.lettuce.core.api.StatefulRedisConnection;
 import io.lettuce.core.codec.RedisCodec;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 
 /**
  * Utility to provide server-side assistance for client-side caches. This is a {@link CacheFrontend} that represents a two-level
@@ -41,6 +42,8 @@ public class ClientSideCaching<K, V> implements CacheFrontend<K, V> {
 
     private final List<Consumer<K>> invalidationListeners = new CopyOnWriteArrayList<>();
 
+    private final ConcurrentHashMap<K, ReentrantLock> keyLocks = new ConcurrentHashMap<>();
+
     private ClientSideCaching(CacheAccessor<K, V> cacheAccessor, RedisCache<K, V> redisCache) {
         this.cacheAccessor = cacheAccessor;
         this.redisCache = redisCache;
@@ -49,12 +52,12 @@ private ClientSideCaching(CacheAccessor<K, V> cacheAccessor, RedisCache<K, V> re
     /**
      * Enable server-assisted Client side caching for the given {@link CacheAccessor} and {@link StatefulRedisConnection}.
      * <p>
-     * Note that the {@link CacheFrontend} is associated with a Redis connection. Make sure to {@link CacheFrontend#close()
-     * close} the frontend object to release the Redis connection after use.
+     * Note that the {@link CacheFrontend} is associated with a Redis connection. Make sure to
+     * {@link CacheFrontend#close() close} the frontend object to release the Redis connection after use.
      *
      * @param cacheAccessor the accessor used to interact with the client-side cache.
      * @param connection the Redis connection to use. The connection will be associated with {@link CacheFrontend} and must be
-     *        closed through {@link CacheFrontend#close()}.
+     * closed through {@link CacheFrontend#close()}.
      * @param tracking the tracking parameters.
      * @param <K> Key type.
      * @param <V> Value type.
@@ -72,12 +75,12 @@ public static <K, V> CacheFrontend<K, V> enable(CacheAccessor<K, V> cacheAccesso
      * Create a server-assisted Client side caching for the given {@link CacheAccessor} and {@link StatefulRedisConnection}.
      * This method expects that client key tracking is already configured.
      * <p>
-     * Note that the {@link CacheFrontend} is associated with a Redis connection. Make sure to {@link CacheFrontend#close()
-     * close} the frontend object to release the Redis connection after use.
+     * Note that the {@link CacheFrontend} is associated with a Redis connection. Make sure to
+     * {@link CacheFrontend#close() close} the frontend object to release the Redis connection after use.
      *
      * @param cacheAccessor the accessor used to interact with the client-side cache.
      * @param connection the Redis connection to use. The connection will be associated with {@link CacheFrontend} and must be
-     *        closed through {@link CacheFrontend#close()}.
+     * closed through {@link CacheFrontend#close()}.
      * @param <K> Key type.
      * @param <V> Value type.
      * @return the {@link CacheFrontend} for value retrieval.
@@ -103,6 +106,7 @@ private static <K, V> CacheFrontend<K, V> create(CacheAccessor<K, V> cacheAccess
     }
 
     private void notifyInvalidate(K key) {
+        keyLocks.remove(key);
 
         for (java.util.function.Consumer<K> invalidationListener : invalidationListeners) {
             invalidationListener.accept(key);
@@ -111,6 +115,7 @@ private void notifyInvalidate(K key) {
 
     @Override
     public void close() {
+        keyLocks.clear();
         redisCache.close();
     }
 
@@ -124,10 +129,20 @@ public V get(K key) {
         V value = cacheAccessor.get(key);
 
         if (value == null) {
-            value = redisCache.get(key);
+            ReentrantLock keyLock = keyLocks.computeIfAbsent(key, k -> new ReentrantLock());
+            keyLock.lock();
+            try {
+                value = cacheAccessor.get(key);
+
+                if (value == null) {
+                    value = redisCache.get(key);
 
-            if (value != null) {
-                cacheAccessor.put(key, value);
+                    if (value != null) {
+                        cacheAccessor.put(key, value);
+                    }
+                }
+            } finally {
+                keyLock.unlock();
             }
         }
 
@@ -140,28 +155,38 @@ public V get(K key, Callable<V> valueLoader) {
         V value = cacheAccessor.get(key);
 
         if (value == null) {
-            value = redisCache.get(key);
-
-            if (value == null) {
+            ReentrantLock keyLock = keyLocks.computeIfAbsent(key, k -> new ReentrantLock());
+            keyLock.lock();
 
-                try {
-                    value = valueLoader.call();
-                } catch (Exception e) {
-                    throw new ValueRetrievalException(
-                            String.format("Value loader %s failed with an exception for key %s", valueLoader, key), e);
-                }
+            try {
+                value = cacheAccessor.get(key);
 
                 if (value == null) {
-                    throw new ValueRetrievalException(
-                            String.format("Value loader %s returned a null value for key %s", valueLoader, key));
-                }
-                redisCache.put(key, value);
+                    value = redisCache.get(key);
 
-                // register interest in key
-                redisCache.get(key);
-            }
+                    if (value == null) {
+                        try {
+                            value = valueLoader.call();
+                        } catch (Exception e) {
+                            throw new ValueRetrievalException(
+                                    String.format("Value loader %s failed with an exception for key %s", valueLoader, key), e);
+                        }
 
-            cacheAccessor.put(key, value);
+                        if (value == null) {
+                            throw new ValueRetrievalException(
+                                    String.format("Value loader %s returned a null value for key %s", valueLoader, key));
+                        }
+
+                        redisCache.put(key, value);
+
+                        redisCache.get(key);
+                    }
+
+                    cacheAccessor.put(key, value);
+                }
+            } finally {
+                keyLock.unlock();
+            }
         }
 
         return value;

src/test/java/io/lettuce/core/support/caching/ClientsideCachingIntegrationTests.java

@@ -3,19 +3,6 @@
 import static io.lettuce.TestTags.INTEGRATION_TEST;
 import static org.assertj.core.api.Assertions.assertThat;
 
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CopyOnWriteArrayList;
-
-import javax.inject.Inject;
-
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Tag;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
-
 import io.lettuce.core.ClientOptions;
 import io.lettuce.core.RedisClient;
 import io.lettuce.core.TestSupport;
@@ -29,11 +16,29 @@
 import io.lettuce.test.LettuceExtension;
 import io.lettuce.test.Wait;
 import io.lettuce.test.condition.EnabledOnCommand;
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import javax.inject.Inject;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 
 /**
  * Integration tests for server-side assisted cache invalidation.
  *
  * @author Mark Paluch
+ * @author Yoobin Yoon
  */
 @Tag(INTEGRATION_TEST)
 @ExtendWith(LettuceExtension.class)
@@ -227,4 +232,127 @@ void serverAssistedCachingShouldUseValueLoader() throws InterruptedException {
         frontend.close();
     }
 
+    @Test
+    void valueLoaderShouldBeInvokedOnceForConcurrentRequests() throws Exception {
+
+        Map<String, String> clientCache = new ConcurrentHashMap<>();
+
+        StatefulRedisConnection<String, String> connection = redisClient.connect();
+
+        final String testKey = "concurrent-loader-key";
+        connection.sync().del(testKey);
+
+        AtomicInteger loaderCallCount = new AtomicInteger(0);
+
+        CacheFrontend<String, String> frontend = ClientSideCaching.enable(CacheAccessor.forMap(clientCache), connection,
+                TrackingArgs.Builder.enabled());
+
+        try {
+            int threadCount = 10;
+            CountDownLatch startLatch = new CountDownLatch(1);
+            CountDownLatch finishLatch = new CountDownLatch(threadCount);
+            List<String> results = new CopyOnWriteArrayList<>();
+
+            ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+            for (int i = 0; i < threadCount; i++) {
+                executor.submit(() -> {
+                    try {
+                        startLatch.await();
+
+                        String result = frontend.get(testKey, () -> {
+                            loaderCallCount.incrementAndGet();
+
+                            try {
+                                Thread.sleep(100);
+                            } catch (InterruptedException e) {
+                                Thread.currentThread().interrupt();
+                            }
+
+                            return "loaded-value";
+                        });
+
+                        results.add(result);
+                    } catch (Exception e) {
+                        e.printStackTrace();
+                    } finally {
+                        finishLatch.countDown();
+                    }
+                });
+            }
+
+            startLatch.countDown();
+
+            finishLatch.await(5, TimeUnit.SECONDS);
+            executor.shutdown();
+
+            assertThat(loaderCallCount.get()).isEqualTo(1);
+
+            assertThat(results).hasSize(threadCount);
+            assertThat(results).containsOnly("loaded-value");
+
+            assertThat(connection.sync().get(testKey)).isEqualTo("loaded-value");
+
+            assertThat(clientCache).containsEntry(testKey, "loaded-value");
+        } finally {
+            frontend.close();
+            connection.close();
+        }
+    }
+
+    @Test
+    void locksShouldBeProperlyCleanedUp() throws Exception {
+
+        Map<String, String> clientCache = new ConcurrentHashMap<>();
+
+        StatefulRedisConnection<String, String> connection = redisClient.connect();
+        StatefulRedisConnection<String, String> otherClient = redisClient.connect();
+
+        final String testKey1 = "lock-test-key1";
+        final String testKey2 = "lock-test-key2";
+        final String initialValue = "initial-value";
+        final String updatedValue = "updated-value";
+
+        connection.sync().del(testKey1, testKey2);
+        connection.sync().set(testKey1, initialValue);
+        connection.sync().set(testKey2, initialValue);
+
+        ClientSideCaching<String, String> frontend = (ClientSideCaching<String, String>) ClientSideCaching.enable(
+                CacheAccessor.forMap(clientCache), connection, TrackingArgs.Builder.enabled());
+
+        Field keyLocksField = ClientSideCaching.class.getDeclaredField("keyLocks");
+        keyLocksField.setAccessible(true);
+        ConcurrentHashMap<String, ReentrantLock> keyLocks = (ConcurrentHashMap<String, ReentrantLock>) keyLocksField.get(
+                frontend);
+
+        try {
+            frontend.get(testKey1);
+            frontend.get(testKey2);
+
+            assertThat(keyLocks).containsKey(testKey1);
+            assertThat(keyLocks).containsKey(testKey2);
+            assertThat(keyLocks).hasSize(2);
+
+            otherClient.sync().set(testKey1, updatedValue);
+
+            Thread.sleep(200);
+
+            assertThat(keyLocks).doesNotContainKey(testKey1);
+            assertThat(keyLocks).containsKey(testKey2);
+            assertThat(keyLocks).hasSize(1);
+
+            frontend.get(testKey1);
+
+            assertThat(keyLocks).containsKey(testKey1);
+            assertThat(keyLocks).hasSize(2);
+
+            frontend.close();
+
+            assertThat(keyLocks).isEmpty();
+
+        } finally {
+            connection.close();
+            otherClient.close();
+        }
+    }
+
 }