[Persisted] Fix the race condition causing clobbered writes between tabs (#4873)

* Broadcast the update in the same tick

The motivation for the original code is unclear. I was not able to reproduce the described behavior and have not seen it mentioned on the web. I'll assume that this was a misunderstanding.

* Remove defensive programming

The only places in this code that we can expect to throw are schema.parse(), JSON.parse(), JSON.stringify(), and localStorage.getItem/setItem/removeItem. Let's push try/catch'es where we expect them to be necessary.

* Don't write or clobber defaults

Writing defaults to local storage is unnecessary. We would write them as a part of next update anyway. So I'm removing that to reduce the number of moving pieces.

However, we do need to be wary of _state being set to defaults. Because _state gets mutated on write. We don't want to mutate the defaults object. To avoid having to think about this, let's copy on write. We don't write to this object very often.

* Refactor: extract tryParse

* Refactor: move string parsing into tryParse

* Extract tryStringify, split logging by platform

Shared data parsing/stringification errors are always logged. Storage errors are only logged on native because we trust the web APIs to work.

* Add a layer of caching to readFromStorage to web

We're going to be doing a read on every write so let's add a fast path that avoids parsing and validating.

* Fix the race condition causing clobbered writes between tabs

1 files changed, 55 insertions, 62 deletions

diff --git a/src/state/persisted/index.web.ts b/src/state/persisted/index.web.ts
index 50f28b6b8..d71b59096 100644
--- a/src/state/persisted/index.web.ts
+++ b/src/state/persisted/index.web.ts
@@ -2,7 +2,12 @@ import EventEmitter from 'eventemitter3'
 
 import BroadcastChannel from '#/lib/broadcast'
 import {logger} from '#/logger'
-import {defaults, Schema, schema} from '#/state/persisted/schema'
+import {
+  defaults,
+  Schema,
+  tryParse,
+  tryStringify,
+} from '#/state/persisted/schema'
 import {PersistedApi} from './types'
 
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
@@ -18,17 +23,9 @@ const _emitter = new EventEmitter()
 
 export async function init() {
   broadcast.onmessage = onBroadcastMessage
-
-  try {
-    const stored = readFromStorage()
-    if (!stored) {
-      writeToStorage(defaults)
-    }
-    _state = stored || defaults
-  } catch (e) {
-    logger.error('persisted state: failed to load root state from storage', {
-      message: e,
-    })
+  const stored = readFromStorage()
+  if (stored) {
+    _state = stored
   }
 }
 init satisfies PersistedApi['init']
@@ -42,16 +39,20 @@ export async function write<K extends keyof Schema>(
   key: K,
   value: Schema[K],
 ): Promise<void> {
-  try {
-    _state[key] = value
-    writeToStorage(_state)
-    // must happen on next tick, otherwise the tab will read stale storage data
-    setTimeout(() => broadcast.postMessage({event: UPDATE_EVENT}), 0)
-  } catch (e) {
-    logger.error(`persisted state: failed writing root state to storage`, {
-      message: e,
-    })
+  const next = readFromStorage()
+  if (next) {
+    // The storage could have been updated by a different tab before this tab is notified.
+    // Make sure this write is applied on top of the latest data in the storage as long as it's valid.
+    _state = next
+    // Don't fire the update listeners yet to avoid a loop.
+    // If there was a change, we'll receive the broadcast event soon enough which will do that.
   }
+  _state = {
+    ..._state,
+    [key]: value,
+  }
+  writeToStorage(_state)
+  broadcast.postMessage({event: UPDATE_EVENT})
 }
 write satisfies PersistedApi['write']
 
@@ -65,62 +66,54 @@ export async function clearStorage() {
   try {
     localStorage.removeItem(BSKY_STORAGE)
   } catch (e: any) {
-    logger.error(`persisted store: failed to clear`, {message: e.toString()})
+    // Expected on the web in private mode.
   }
 }
 clearStorage satisfies PersistedApi['clearStorage']
 
 async function onBroadcastMessage({data}: MessageEvent) {
   if (typeof data === 'object' && data.event === UPDATE_EVENT) {
-    try {
-      // read next state, possibly updated by another tab
-      const next = readFromStorage()
-
-      if (next) {
-        _state = next
-        _emitter.emit('update')
-      } else {
-        logger.error(
-          `persisted state: handled update update from broadcast channel, but found no data`,
-        )
-      }
-    } catch (e) {
+    // read next state, possibly updated by another tab
+    const next = readFromStorage()
+    if (next) {
+      _state = next
+      _emitter.emit('update')
+    } else {
       logger.error(
-        `persisted state: failed handling update from broadcast channel`,
-        {
-          message: e,
-        },
+        `persisted state: handled update update from broadcast channel, but found no data`,
       )
     }
   }
 }
 
 function writeToStorage(value: Schema) {
-  schema.parse(value)
-  localStorage.setItem(BSKY_STORAGE, JSON.stringify(value))
+  const rawData = tryStringify(value)
+  if (rawData) {
+    try {
+      localStorage.setItem(BSKY_STORAGE, rawData)
+    } catch (e) {
+      // Expected on the web in private mode.
+    }
+  }
 }
 
+let lastRawData: string | undefined
+let lastResult: Schema | undefined
 function readFromStorage(): Schema | undefined {
-  const rawData = localStorage.getItem(BSKY_STORAGE)
-  const objData = rawData ? JSON.parse(rawData) : undefined
-
-  // new user
-  if (!objData) return undefined
-
-  // existing user, validate
-  const parsed = schema.safeParse(objData)
-
-  if (parsed.success) {
-    return objData
-  } else {
-    const errors =
-      parsed.error?.errors?.map(e => ({
-        code: e.code,
-        // @ts-ignore exists on some types
-        expected: e?.expected,
-        path: e.path?.join('.'),
-      })) || []
-    logger.error(`persisted store: data failed validation on read`, {errors})
-    return undefined
+  let rawData: string | null = null
+  try {
+    rawData = localStorage.getItem(BSKY_STORAGE)
+  } catch (e) {
+    // Expected on the web in private mode.
+  }
+  if (rawData) {
+    if (rawData === lastRawData) {
+      return lastResult
+    } else {
+      const result = tryParse(rawData)
+      lastRawData = rawData
+      lastResult = result
+      return result
+    }
   }
 }