From 7b1a91937201c4b7725922f6531f0c474735bb36 Mon Sep 17 00:00:00 2001
From: Benoit 'BoD' Lubek <BoD@JRAF.org>
Date: Mon, 29 Apr 2024 18:25:08 +0200
Subject: [PATCH] Backport #5848 to the v3 branch (#5863)

(cherry picked from commit 7d5c7c7c8c66f3ce982a6a412fe50057f726bf98)
---
 .../cache/normalized/api/MemoryCache.kt       | 30 +++++++++++-----
 .../cache/normalized/api/MemoryCache.kt       | 34 +++++++++++++------
 2 files changed, 44 insertions(+), 20 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 c47321ce559..bf59c0d057f 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,13 +39,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
@@ -53,8 +48,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() {
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 b2e689e8ffd..b6535a14ab6 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++
       }
@@ -137,7 +149,7 @@ class MemoryCache(
 
   private class CacheEntry(
       val record: Record,
-      val expireAfterMillis: Long
+      val expireAfterMillis: Long,
   ) {
     val cachedAtMillis: Long = currentTimeMillis()