feat(core): use quick-lru for documents edit state

packages/sanity/src/core/store/_legacy/document/document-pair/editState.ts

@@ -1,14 +1,14 @@
 import {type SanityClient} from '@sanity/client'
 import {type SanityDocument, type Schema} from '@sanity/types'
-import {combineLatest, defer, merge, type Observable, of} from 'rxjs'
-import {finalize, map, publishReplay, refCount, startWith, switchMap, tap} from 'rxjs/operators'
+import {combineLatest, defer, type Observable, of} from 'rxjs'
+import {map, publishReplay, refCount, startWith, switchMap, tap} from 'rxjs/operators'
 
+import {type DocumentsStorage} from '../documentsStorage'
 import {type IdPair, type PendingMutationsEvent} from '../types'
 import {memoize} from '../utils/createMemoizer'
 import {memoizeKeyGen} from './memoizeKeyGen'
 import {snapshotPair} from './snapshotPair'
 import {isLiveEditEnabled} from './utils/isLiveEditEnabled'
-import {getPairFromLocalStorage, savePairToLocalStorage} from './utils/localStoragePOC'
 
 interface TransactionSyncLockState {
   enabled: boolean
@@ -36,81 +36,59 @@ export const editState = memoize(
       client: SanityClient
       schema: Schema
       serverActionsEnabled: Observable<boolean>
+      storage: DocumentsStorage
     },
     idPair: IdPair,
     typeName: string,
   ): Observable<EditStateFor> => {
     const liveEdit = isLiveEditEnabled(ctx.schema, typeName)
-    const INITIAL = {
-      id: idPair.publishedId,
-      type: typeName,
-      draft: null,
-      published: null,
-      liveEdit,
-      ready: false,
-      transactionSyncLock: null,
-    }
+    const storage = ctx.storage
 
-    let cachedDocumentPair: {
-      draft: SanityDocument | null
-      published: SanityDocument | null
-    } | null = null
-
-    function getCachedPair() {
-      // try first read it from memory
-      // if we haven't got it in memory, see if it's in localstorage
-      if (cachedDocumentPair) {
-        return cachedDocumentPair
-      }
-      return getPairFromLocalStorage(idPair)
-    }
-
-    return snapshotPair(ctx.client, idPair, typeName, ctx.serverActionsEnabled).pipe(
-      switchMap((versions) =>
-        combineLatest([
-          versions.draft.snapshots$,
-          versions.published.snapshots$,
-          versions.transactionsPendingEvents$.pipe(
-            // eslint-disable-next-line max-nested-callbacks
-            map((ev: PendingMutationsEvent) => (ev.phase === 'begin' ? LOCKED : NOT_LOCKED)),
-            startWith(NOT_LOCKED),
+    return defer(() => {
+      return of({
+        draft: storage.get(idPair.draftId) || null,
+        published: storage.get(idPair.publishedId) || null,
+      })
+    }).pipe(
+      switchMap((cachePair) => {
+        return snapshotPair(ctx.client, idPair, typeName, ctx.serverActionsEnabled).pipe(
+          switchMap((versions) =>
+            combineLatest([
+              versions.draft.snapshots$,
+              versions.published.snapshots$,
+              versions.transactionsPendingEvents$.pipe(
+                // eslint-disable-next-line max-nested-callbacks
+                map((ev: PendingMutationsEvent) => (ev.phase === 'begin' ? LOCKED : NOT_LOCKED)),
+                startWith(NOT_LOCKED),
+              ),
+            ]),
           ),
-        ]),
-      ),
-      tap(([draft, published]) => {
-        cachedDocumentPair = {draft, published}
-      }),
-      map(([draftSnapshot, publishedSnapshot, transactionSyncLock]) => ({
-        id: idPair.publishedId,
-        type: typeName,
-        draft: draftSnapshot,
-        published: publishedSnapshot,
-        liveEdit,
-        ready: true,
-        transactionSyncLock,
-      })),
-      // todo: turn this into a proper operator function - It's like startWith only that it takes a function that will be invoked upon subscription
-      (input$) => {
-        return defer(() => {
-          const cachedPair = getCachedPair()
-          return merge(
-            cachedPair
-              ? of({
-                  id: idPair.publishedId,
-                  type: typeName,
-                  draft: cachedPair.draft,
-                  published: cachedPair.published,
-                  liveEdit,
-                  ready: false,
-                  transactionSyncLock: null,
-                })
-              : of(INITIAL),
-            input$,
-          )
-        })
-      },
-      finalize(() => {
-        savePairToLocalStorage(cachedDocumentPair)
+          tap(([draftSnapshot, publishedSnapshot]) => {
+            if (draftSnapshot) {
+              storage.save(idPair.draftId, draftSnapshot)
+            } else if (publishedSnapshot) {
+              storage.save(idPair.publishedId, publishedSnapshot)
+            }
+          }),
+          map(([draftSnapshot, publishedSnapshot, transactionSyncLock]) => ({
+            id: idPair.publishedId,
+            type: typeName,
+            draft: draftSnapshot,
+            published: publishedSnapshot,
+            liveEdit,
+            ready: true,
+            transactionSyncLock,
+          })),
+          startWith({
+            id: idPair.publishedId,
+            type: typeName,
+            draft: cachePair.draft,
+            published: cachePair.published,
+            liveEdit,
+            ready: false,
+            transactionSyncLock: null,
+          }),
+        )
       }),
       publishReplay(1),
       refCount(),

packages/sanity/src/core/store/_legacy/document/document-pair/validation.ts

@@ -9,6 +9,7 @@ import {type SourceClientOptions} from '../../../../config'
 import {type LocaleSource} from '../../../../i18n'
 import {type DraftsModelDocumentAvailability} from '../../../../preview'
 import {validateDocumentWithReferences, type ValidationStatus} from '../../../../validation'
+import {type DocumentsStorage} from '../documentsStorage'
 import {type IdPair} from '../types'
 import {memoize} from '../utils/createMemoizer'
 import {editState} from './editState'
@@ -31,6 +32,7 @@ export const validation = memoize(
       schema: Schema
       i18n: LocaleSource
       serverActionsEnabled: Observable<boolean>
+      storage: DocumentsStorage
     },
     {draftId, publishedId}: IdPair,
     typeName: string,

packages/sanity/src/core/store/_legacy/document/document-store.ts

@@ -23,11 +23,11 @@ import {
 } from './document-pair/operationEvents'
 import {type OperationsAPI} from './document-pair/operations'
 import {validation} from './document-pair/validation'
+import {createDocumentsStorage} from './documentsStorage'
 import {getInitialValueStream, type InitialValueMsg, type InitialValueOptions} from './initialValue'
 import {listenQuery, type ListenQueryOptions} from './listenQuery'
 import {resolveTypeForDocument} from './resolveTypeForDocument'
 import {type IdPair} from './types'
-
 /**
  * @hidden
  * @beta */
@@ -104,7 +104,7 @@ export function createDocumentStore({
   // internal operations, and a `getClient` method that we expose to user-land
   // for things like validations
   const client = getClient(DEFAULT_STUDIO_CLIENT_OPTIONS)
-
+  const storage = createDocumentsStorage()
   const ctx = {
     client,
     getClient,
@@ -113,6 +113,7 @@ export function createDocumentStore({
     schema,
     i18n,
     serverActionsEnabled,
+    storage,
   }
 
   return {

packages/sanity/src/core/store/_legacy/document/documentsStorage.ts

@@ -0,0 +1,21 @@
+import {type SanityDocument} from '@sanity/types'
+import QuickLRU from 'quick-lru'
+
+export interface DocumentsStorage {
+  save: (id: string, doc: SanityDocument) => void
+  get: (id: string) => SanityDocument | null
+}
+
+export function createDocumentsStorage(): DocumentsStorage {
+  const documentsCache = new QuickLRU<string, SanityDocument | null>({
+    maxSize: 50,
+  })
+  return {
+    save(id, doc) {
+      documentsCache.set(id, doc)
+    },
+    get(id) {
+      return documentsCache.get(id) || null
+    },
+  }
+}