[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

3 files changed, 130 insertions, 104 deletions

diff --git a/src/state/persisted/index.ts b/src/state/persisted/index.ts
index 639e4e47f..95f814850 100644
--- a/src/state/persisted/index.ts
+++ b/src/state/persisted/index.ts
@@ -1,7 +1,12 @@
 import AsyncStorage from '@react-native-async-storage/async-storage'
 
 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'
@@ -12,16 +17,9 @@ const BSKY_STORAGE = 'BSKY_STORAGE'
 let _state: Schema = defaults
 
 export async function init() {
-  try {
-    const stored = await readFromStorage()
-    if (!stored) {
-      await writeToStorage(defaults)
-    }
-    _state = stored || defaults
-  } catch (e) {
-    logger.error('persisted state: failed to load root state from storage', {
-      message: e,
-    })
+  const stored = await readFromStorage()
+  if (stored) {
+    _state = stored
   }
 }
 init satisfies PersistedApi['init']
@@ -35,14 +33,11 @@ export async function write<K extends keyof Schema>(
   key: K,
   value: Schema[K],
 ): Promise<void> {
-  try {
-    _state[key] = value
-    await writeToStorage(_state)
-  } catch (e) {
-    logger.error(`persisted state: failed writing root state to storage`, {
-      message: e,
-    })
+  _state = {
+    ..._state,
+    [key]: value,
   }
+  await writeToStorage(_state)
 }
 write satisfies PersistedApi['write']
 
@@ -61,31 +56,28 @@ export async function clearStorage() {
 clearStorage satisfies PersistedApi['clearStorage']
 
 async function writeToStorage(value: Schema) {
-  schema.parse(value)
-  await AsyncStorage.setItem(BSKY_STORAGE, JSON.stringify(value))
+  const rawData = tryStringify(value)
+  if (rawData) {
+    try {
+      await AsyncStorage.setItem(BSKY_STORAGE, rawData)
+    } catch (e) {
+      logger.error(`persisted state: failed writing root state to storage`, {
+        message: e,
+      })
+    }
+  }
 }
 
 async function readFromStorage(): Promise<Schema | undefined> {
-  const rawData = await AsyncStorage.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 = await AsyncStorage.getItem(BSKY_STORAGE)
+  } catch (e) {
+    logger.error(`persisted state: failed reading root state from storage`, {
+      message: e,
+    })
+  }
+  if (rawData) {
+    return tryParse(rawData)
   }
 }
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
+    }
   }
 }
diff --git a/src/state/persisted/schema.ts b/src/state/persisted/schema.ts
index 399a7e793..0b652a1f0 100644
--- a/src/state/persisted/schema.ts
+++ b/src/state/persisted/schema.ts
@@ -1,5 +1,6 @@
 import {z} from 'zod'
 
+import {logger} from '#/logger'
 import {deviceLocales} from '#/platform/detection'
 import {PlatformInfo} from '../../../modules/expo-bluesky-swiss-army'
 
@@ -43,7 +44,7 @@ const currentAccountSchema = accountSchema.extend({
 })
 export type PersistedCurrentAccount = z.infer<typeof currentAccountSchema>
 
-export const schema = z.object({
+const schema = z.object({
   colorMode: z.enum(['system', 'light', 'dark']),
   darkTheme: z.enum(['dim', 'dark']).optional(),
   session: z.object({
@@ -133,3 +134,43 @@ export const defaults: Schema = {
   kawaii: false,
   hasCheckedForStarterPack: false,
 }
+
+export function tryParse(rawData: string): Schema | undefined {
+  let objData
+  try {
+    objData = JSON.parse(rawData)
+  } catch (e) {
+    logger.error('persisted state: failed to parse root state from storage', {
+      message: e,
+    })
+  }
+  if (!objData) {
+    return undefined
+  }
+  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
+  }
+}
+
+export function tryStringify(value: Schema): string | undefined {
+  try {
+    schema.parse(value)
+    return JSON.stringify(value)
+  } catch (e) {
+    logger.error(`persisted state: failed stringifying root state`, {
+      message: e,
+    })
+    return undefined
+  }
+}