Dedupe did cache refreshes (#1773)

* dedupe refreshes to did cache * handle expired cache entries as well * apply in pds as well * changeset
bluesky-social · Oct 25, 2023 · bb039d8 · bb039d8
1 parent ce28725
commit bb039d8
Show file tree

Hide file tree

Showing 6 changed files with 88 additions and 46 deletions.
diff --git a/.changeset/new-snails-hope.md b/.changeset/new-snails-hope.md
@@ -0,0 +1,5 @@
+---
+'@atproto/identity': minor
+---
+
+Pass stale did doc into refresh cache functions
diff --git a/packages/bsky/src/did-cache.ts b/packages/bsky/src/did-cache.ts
@@ -17,28 +17,42 @@ export class DidSqlCache implements DidCache {
     this.pQueue = new PQueue()
   }
 
-  async cacheDid(did: string, doc: DidDocument): Promise<void> {
-    await this.db.db
-      .insertInto('did_cache')
-      .values({ did, doc, updatedAt: Date.now() })
-      .onConflict((oc) =>
-        oc.column('did').doUpdateSet({
-          doc: excluded(this.db.db, 'doc'),
-          updatedAt: excluded(this.db.db, 'updatedAt'),
-        }),
-      )
-      .executeTakeFirst()
+  async cacheDid(
+    did: string,
+    doc: DidDocument,
+    prevResult?: CacheResult,
+  ): Promise<void> {
+    if (prevResult) {
+      await this.db.db
+        .updateTable('did_cache')
+        .set({ doc, updatedAt: Date.now() })
+        .where('did', '=', did)
+        .where('updatedAt', '=', prevResult.updatedAt)
+        .execute()
+    } else {
+      await this.db.db
+        .insertInto('did_cache')
+        .values({ did, doc, updatedAt: Date.now() })
+        .onConflict((oc) =>
+          oc.column('did').doUpdateSet({
+            doc: excluded(this.db.db, 'doc'),
+            updatedAt: excluded(this.db.db, 'updatedAt'),
+          }),
+        )
+        .executeTakeFirst()
+    }
   }
 
   async refreshCache(
     did: string,
     getDoc: () => Promise<DidDocument | null>,
+    prevResult?: CacheResult,
   ): Promise<void> {
     this.pQueue?.add(async () => {
       try {
         const doc = await getDoc()
         if (doc) {
-          await this.cacheDid(did, doc)
+          await this.cacheDid(did, doc, prevResult)
         } else {
           await this.clearEntry(did)
         }
@@ -55,20 +69,17 @@ export class DidSqlCache implements DidCache {
       .selectAll()
       .executeTakeFirst()
     if (!res) return null
+
     const now = Date.now()
     const updatedAt = new Date(res.updatedAt).getTime()
-
     const expired = now > updatedAt + this.maxTTL
-    if (expired) {
-      return null
-    }
-
     const stale = now > updatedAt + this.staleTTL
     return {
       doc: res.doc,
       updatedAt,
       did,
       stale,
+      expired,
     }
   }
 

diff --git a/packages/identity/src/did/base-resolver.ts b/packages/identity/src/did/base-resolver.ts
@@ -1,6 +1,12 @@
 import * as crypto from '@atproto/crypto'
 import { check } from '@atproto/common-web'
-import { DidCache, AtprotoData, DidDocument, didDocument } from '../types'
+import {
+  DidCache,
+  AtprotoData,
+  DidDocument,
+  didDocument,
+  CacheResult,
+} from '../types'
 import * as atprotoData from './atproto-data'
 import { DidNotFoundError, PoorlyFormattedDidDocumentError } from '../errors'
 
@@ -25,20 +31,25 @@ export abstract class BaseResolver {
     return this.validateDidDoc(did, got)
   }
 
-  async refreshCache(did: string): Promise<void> {
-    await this.cache?.refreshCache(did, () => this.resolveNoCache(did))
+  async refreshCache(did: string, prevResult?: CacheResult): Promise<void> {
+    await this.cache?.refreshCache(
+      did,
+      () => this.resolveNoCache(did),
+      prevResult,
+    )
   }
 
   async resolve(
     did: string,
     forceRefresh = false,
   ): Promise<DidDocument | null> {
+    let fromCache: CacheResult | null = null
     if (this.cache && !forceRefresh) {
-      const fromCache = await this.cache.checkCache(did)
-      if (fromCache?.stale) {
-        await this.refreshCache(did)
-      }
-      if (fromCache) {
+      fromCache = await this.cache.checkCache(did)
+      if (fromCache && !fromCache.expired) {
+        if (fromCache?.stale) {
+          await this.refreshCache(did, fromCache)
+        }
         return fromCache.doc
       }
     }
@@ -48,7 +59,7 @@ export abstract class BaseResolver {
       await this.cache?.clearEntry(did)
       return null
     }
-    await this.cache?.cacheDid(did, got)
+    await this.cache?.cacheDid(did, got, fromCache ?? undefined)
     return got
   }
 

diff --git a/packages/identity/src/did/memory-cache.ts b/packages/identity/src/did/memory-cache.ts
@@ -35,13 +35,12 @@ export class MemoryCache implements DidCache {
     if (!val) return null
     const now = Date.now()
     const expired = now > val.updatedAt + this.maxTTL
-    if (expired) return null
-
     const stale = now > val.updatedAt + this.staleTTL
     return {
       ...val,
       did,
       stale,
+      expired,
     }
   }
 

diff --git a/packages/identity/src/types.ts b/packages/identity/src/types.ts
@@ -30,14 +30,20 @@ export type CacheResult = {
   doc: DidDocument
   updatedAt: number
   stale: boolean
+  expired: boolean
 }
 
 export interface DidCache {
-  cacheDid(did: string, doc: DidDocument): Promise<void>
+  cacheDid(
+    did: string,
+    doc: DidDocument,
+    prevResult?: CacheResult,
+  ): Promise<void>
   checkCache(did: string): Promise<CacheResult | null>
   refreshCache(
     did: string,
     getDoc: () => Promise<DidDocument | null>,
+    prevResult?: CacheResult,
   ): Promise<void>
   clearEntry(did: string): Promise<void>
   clear(): Promise<void>

diff --git a/packages/pds/src/did-cache.ts b/packages/pds/src/did-cache.ts
@@ -15,28 +15,42 @@ export class DidSqlCache implements DidCache {
     this.pQueue = new PQueue()
   }
 
-  async cacheDid(did: string, doc: DidDocument): Promise<void> {
-    await this.db.db
-      .insertInto('did_cache')
-      .values({ did, doc: JSON.stringify(doc), updatedAt: Date.now() })
-      .onConflict((oc) =>
-        oc.column('did').doUpdateSet({
-          doc: excluded(this.db.db, 'doc'),
-          updatedAt: excluded(this.db.db, 'updatedAt'),
-        }),
-      )
-      .executeTakeFirst()
+  async cacheDid(
+    did: string,
+    doc: DidDocument,
+    prevResult?: CacheResult,
+  ): Promise<void> {
+    if (prevResult) {
+      await this.db.db
+        .updateTable('did_cache')
+        .set({ doc: JSON.stringify(doc), updatedAt: Date.now() })
+        .where('did', '=', did)
+        .where('updatedAt', '=', prevResult.updatedAt)
+        .execute()
+    } else {
+      await this.db.db
+        .insertInto('did_cache')
+        .values({ did, doc: JSON.stringify(doc), updatedAt: Date.now() })
+        .onConflict((oc) =>
+          oc.column('did').doUpdateSet({
+            doc: excluded(this.db.db, 'doc'),
+            updatedAt: excluded(this.db.db, 'updatedAt'),
+          }),
+        )
+        .executeTakeFirst()
+    }
   }
 
   async refreshCache(
     did: string,
     getDoc: () => Promise<DidDocument | null>,
+    prevResult?: CacheResult,
   ): Promise<void> {
     this.pQueue?.add(async () => {
       try {
         const doc = await getDoc()
         if (doc) {
-          await this.cacheDid(did, doc)
+          await this.cacheDid(did, doc, prevResult)
         } else {
           await this.clearEntry(did)
         }
@@ -55,18 +69,14 @@ export class DidSqlCache implements DidCache {
     if (!res) return null
     const now = Date.now()
     const updatedAt = new Date(res.updatedAt).getTime()
-
     const expired = now > updatedAt + this.maxTTL
-    if (expired) {
-      return null
-    }
-
     const stale = now > updatedAt + this.staleTTL
     return {
       doc: JSON.parse(res.doc) as DidDocument,
       updatedAt,
       did,
       stale,
+      expired,
     }
   }