From ab11840dcfa7a72861c6c166bc81ee2aa7b41077 Mon Sep 17 00:00:00 2001
From: Benoit 'BoD' Lubek <bod@jraf.org>
Date: Mon, 29 Apr 2024 11:22:07 +0200
Subject: [PATCH] Call through to loadRecords from the MemoryCache to the SQL
 one (#5848)

(cherry picked from commit 7d5c7c7c8c66f3ce982a6a412fe50057f726bf98)
---
 .../cache/normalized/api/MemoryCache.kt       | 32 +++++++++++------
 .../cache/normalized/api/MemoryCache.kt       | 34 +++++++++++++------
 2 files changed, 45 insertions(+), 21 deletions(-)

diff --git a/libraries/apollo-normalized-cache-api-incubating/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt b/libraries/apollo-normalized-cache-api-incubating/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt
index 1bf0e933b3..8ca3c9d5c5 100644
--- a/libraries/apollo-normalized-cache-api-incubating/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt
+++ b/libraries/apollo-normalized-cache-api-incubating/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt
@@ -39,22 +39,34 @@ class MemoryCache(
     get() = lruCache.size()
 
   override fun loadRecord(key: String, cacheHeaders: CacheHeaders): Record? = lock.lock {
-    val cacheEntry = lruCache[key]?.also { cacheEntry ->
-      if (cacheEntry.isExpired || cacheHeaders.hasHeader(ApolloCacheHeaders.EVICT_AFTER_READ)) {
-        lruCache.remove(key)
-      }
+    val record = internalLoadRecord(key, cacheHeaders)
+    record ?: nextCache?.loadRecord(key, cacheHeaders)?.also { nextCachedRecord ->
+      lruCache[key] = CacheEntry(
+        record = nextCachedRecord,
+        expireAfterMillis = expireAfterMillis
+      )
     }
+  }
 
-    cacheEntry?.takeUnless { it.isExpired }?.record ?: nextCache?.loadRecord(key, cacheHeaders)?.also { nextCachedRecord ->
-      lruCache[key] = CacheEntry(
-          record = nextCachedRecord,
-          expireAfterMillis = expireAfterMillis
+  override fun loadRecords(keys: Collection<String>, cacheHeaders: CacheHeaders): Collection<Record> = lock.lock {
+    val recordsByKey: Map<String, Record?> = keys.associateWith { key -> internalLoadRecord(key, cacheHeaders) }
+    val missingKeys = recordsByKey.filterValues { it == null }.keys
+    val nextCachedRecords = nextCache?.loadRecords(missingKeys, cacheHeaders).orEmpty()
+    for (record in nextCachedRecords) {
+      lruCache[record.key] = CacheEntry(
+        record = record,
+        expireAfterMillis = expireAfterMillis
       )
     }
+    recordsByKey.values.filterNotNull() + nextCachedRecords
   }
 
-  override fun loadRecords(keys: Collection<String>, cacheHeaders: CacheHeaders): Collection<Record> {
-    return keys.mapNotNull { key -> loadRecord(key, cacheHeaders) }
+  private fun internalLoadRecord(key: String, cacheHeaders: CacheHeaders): Record? {
+    return lruCache[key]?.also { cacheEntry ->
+      if (cacheEntry.isExpired || cacheHeaders.hasHeader(ApolloCacheHeaders.EVICT_AFTER_READ)) {
+        lruCache.remove(key)
+      }
+    }?.takeUnless { it.isExpired }?.record
   }
 
   override fun clearAll() {
diff --git a/libraries/apollo-normalized-cache-api/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt b/libraries/apollo-normalized-cache-api/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt
index fc19783e86..2abd968b65 100644
--- a/libraries/apollo-normalized-cache-api/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt
+++ b/libraries/apollo-normalized-cache-api/src/commonMain/kotlin/com/apollographql/apollo3/cache/normalized/api/MemoryCache.kt
@@ -38,13 +38,8 @@ class MemoryCache(
     get() = lruCache.size()
 
   override fun loadRecord(key: String, cacheHeaders: CacheHeaders): Record? = lock.lock {
-    val cacheEntry = lruCache[key]?.also { cacheEntry ->
-      if (cacheEntry.isExpired || cacheHeaders.hasHeader(ApolloCacheHeaders.EVICT_AFTER_READ)) {
-        lruCache.remove(key)
-      }
-    }
-
-    cacheEntry?.takeUnless { it.isExpired }?.record ?: nextCache?.loadRecord(key, cacheHeaders)?.also { nextCachedRecord ->
+    val record = internalLoadRecord(key, cacheHeaders)
+    record ?: nextCache?.loadRecord(key, cacheHeaders)?.also { nextCachedRecord ->
       lruCache[key] = CacheEntry(
           record = nextCachedRecord,
           expireAfterMillis = expireAfterMillis
@@ -52,8 +47,25 @@ class MemoryCache(
     }
   }
 
-  override fun loadRecords(keys: Collection<String>, cacheHeaders: CacheHeaders): Collection<Record> {
-    return keys.mapNotNull { key -> loadRecord(key, cacheHeaders) }
+  override fun loadRecords(keys: Collection<String>, cacheHeaders: CacheHeaders): Collection<Record> = lock.lock {
+    val recordsByKey: Map<String, Record?> = keys.associateWith { key -> internalLoadRecord(key, cacheHeaders) }
+    val missingKeys = recordsByKey.filterValues { it == null }.keys
+    val nextCachedRecords = nextCache?.loadRecords(missingKeys, cacheHeaders).orEmpty()
+    for (record in nextCachedRecords) {
+      lruCache[record.key] = CacheEntry(
+          record = record,
+          expireAfterMillis = expireAfterMillis
+      )
+    }
+    recordsByKey.values.filterNotNull() + nextCachedRecords
+  }
+
+  private fun internalLoadRecord(key: String, cacheHeaders: CacheHeaders): Record? {
+    return lruCache[key]?.also { cacheEntry ->
+      if (cacheEntry.isExpired || cacheHeaders.hasHeader(ApolloCacheHeaders.EVICT_AFTER_READ)) {
+        lruCache.remove(key)
+      }
+    }?.takeUnless { it.isExpired }?.record
   }
 
   override fun clearAll() {
@@ -79,7 +91,7 @@ class MemoryCache(
     var total = 0
     val keys = HashSet(lruCache.keys()) // local copy to avoid concurrent modification
     keys.forEach {
-      if (regex.matches(it)){
+      if (regex.matches(it)) {
         lruCache.remove(it)
         total++
       }
@@ -132,7 +144,7 @@ class MemoryCache(
 
   private class CacheEntry(
       val record: Record,
-      val expireAfterMillis: Long
+      val expireAfterMillis: Long,
   ) {
     val cachedAtMillis: Long = currentTimeMillis()