From 19d79bea3bccfcf7738ea3b04d75c9efcdd18b1b Mon Sep 17 00:00:00 2001
From: Kamil Kisiela <kamil.kisiela@gmail.com>
Date: Mon, 30 Jan 2023 18:32:07 +0100
Subject: [PATCH] Extra read from KV storage (#1185)

---
 packages/services/cdn-worker/src/handler.ts   |  7 ++++-
 .../services/cdn-worker/tests/cdn.spec.ts     | 26 +++++++++----------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/packages/services/cdn-worker/src/handler.ts b/packages/services/cdn-worker/src/handler.ts
index 24ee2275d..61f76e9c3 100644
--- a/packages/services/cdn-worker/src/handler.ts
+++ b/packages/services/cdn-worker/src/handler.ts
@@ -193,7 +193,12 @@ export const createRequestHandler = (deps: RequestHandlerDependencies) => {
     analytics.track({ type: 'artifact', value: artifactType, version: 'v0' }, targetId);
 
     const kvStorageKey = `target:${targetId}:${storageKeyType}`;
-    const rawValue = await deps.getRawStoreValue(kvStorageKey);
+    const rawValue = await deps.getRawStoreValue(kvStorageKey).catch(() => {
+      // Do an extra attempt to read the value from the store.
+      // If we see that a single retry does not help, we should do a proper retry logic here.
+      // Why not now? Because we do have a new implementation that is based on R2 storage and this change is simple enough.
+      return deps.getRawStoreValue(kvStorageKey);
+    });
 
     if (rawValue) {
       const etag = await createETag(`${kvStorageKey}|${rawValue}`);
diff --git a/packages/services/cdn-worker/tests/cdn.spec.ts b/packages/services/cdn-worker/tests/cdn.spec.ts
index da26564b3..44a1f78eb 100644
--- a/packages/services/cdn-worker/tests/cdn.spec.ts
+++ b/packages/services/cdn-worker/tests/cdn.spec.ts
@@ -46,7 +46,7 @@ describe('CDN Worker', () => {
           },
         },
       }),
-      getRawStoreValue: (key: string) => map.get(key),
+      getRawStoreValue: async (key: string) => map.get(key),
     });
 
     const schemaRequest = new Request(`https://fake-worker.com/${targetId}/schema`, {
@@ -92,7 +92,7 @@ describe('CDN Worker', () => {
           },
         },
       }),
-      getRawStoreValue: (key: string) => map.get(key),
+      getRawStoreValue: async (key: string) => map.get(key),
     });
 
     const token = createToken(SECRET, targetId);
@@ -160,7 +160,7 @@ describe('CDN Worker', () => {
           },
         },
       }),
-      getRawStoreValue: (key: string) => map.get(key),
+      getRawStoreValue: async (key: string) => map.get(key),
     });
 
     const firstRequest = new Request(`https://fake-worker.com/${targetId}/supergraph`, {
@@ -221,7 +221,7 @@ describe('CDN Worker', () => {
           },
         },
       }),
-      getRawStoreValue: (key: string) => map.get(key),
+      getRawStoreValue: async (key: string) => map.get(key),
     });
 
     const token = createToken(SECRET, targetId);
@@ -284,7 +284,7 @@ describe('CDN Worker', () => {
           },
         },
       }),
-      getRawStoreValue: (key: string) => map.get(key),
+      getRawStoreValue: async (key: string) => map.get(key),
     });
 
     const token = createToken(SECRET, targetId);
@@ -352,7 +352,7 @@ describe('CDN Worker', () => {
           },
         },
       }),
-      getRawStoreValue: (key: string) => map.get(key),
+      getRawStoreValue: async (key: string) => map.get(key),
     });
 
     const token = createToken(SECRET, targetId);
@@ -397,7 +397,7 @@ describe('CDN Worker', () => {
     it('Should throw when target id is missing', async () => {
       const handleRequest = createRequestHandler({
         isKeyValid: KeyValidators.AlwaysTrue,
-        getRawStoreValue: (_key: string) => Promise.resolve(null),
+        getRawStoreValue: async (_key: string) => Promise.resolve(null),
       });
 
       const request = new Request('https://fake-worker.com/', {});
@@ -410,7 +410,7 @@ describe('CDN Worker', () => {
     it('Should throw when requested resource is not valid', async () => {
       const handleRequest = createRequestHandler({
         isKeyValid: KeyValidators.AlwaysTrue,
-        getRawStoreValue: (_key: string) => Promise.resolve(null),
+        getRawStoreValue: async (_key: string) => Promise.resolve(null),
       });
 
       const request = new Request('https://fake-worker.com/fake-target-id/error', {});
@@ -423,7 +423,7 @@ describe('CDN Worker', () => {
     it('Should throw when auth key is missing', async () => {
       const handleRequest = createRequestHandler({
         isKeyValid: KeyValidators.AlwaysTrue,
-        getRawStoreValue: (_key: string) => Promise.resolve(null),
+        getRawStoreValue: async (_key: string) => Promise.resolve(null),
       });
 
       const request = new Request('https://fake-worker.com/fake-target-id/sdl', {});
@@ -436,7 +436,7 @@ describe('CDN Worker', () => {
     it('Should throw when key validation function fails', async () => {
       const handleRequest = createRequestHandler({
         isKeyValid: KeyValidators.AlwaysFalse,
-        getRawStoreValue: (_key: string) => Promise.resolve(null),
+        getRawStoreValue: async (_key: string) => Promise.resolve(null),
       });
 
       const request = new Request('https://fake-worker.com/fake-target-id/sdl', {
@@ -474,7 +474,7 @@ describe('CDN Worker', () => {
             },
           },
         }),
-        getRawStoreValue: (key: string) => map.get(key),
+        getRawStoreValue: async (key: string) => map.get(key),
       });
 
       const token = createToken(SECRET, targetId);
@@ -509,7 +509,7 @@ describe('CDN Worker', () => {
             },
           },
         }),
-        getRawStoreValue: (key: string) => map.get(key),
+        getRawStoreValue: async (key: string) => map.get(key),
       });
 
       const token = createToken(SECRET, 'fake-target-id');
@@ -542,7 +542,7 @@ describe('CDN Worker', () => {
             },
           },
         }),
-        getRawStoreValue: (key: string) => new Map().get(key),
+        getRawStoreValue: async (key: string) => new Map().get(key),
       });
 
       const request = new Request(`https://fake-worker.com/some-target/sdl`, {