[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

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()
-    const stored = await readFromStorage()
+  if (stored) {
-    if (!stored) {
+    _state = stored
-      await writeToStorage(defaults)
-    }
-    _state = stored || defaults
-  } catch (e) {
-    logger.error('persisted state: failed to load root state from storage', {
-      message: e,
-    })
   }
 }
 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 = {
-    _state[key] = value
+    ..._state,
-    await writeToStorage(_state)
+    [key]: value,
-  } catch (e) {
-    logger.error(`persisted state: failed writing root state to storage`, {
-      message: e,
-    })
   }
+  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)
+  const rawData = tryStringify(value)
-  await AsyncStorage.setItem(BSKY_STORAGE, JSON.stringify(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)
+  let rawData: string | null = null
-  const objData = rawData ? JSON.parse(rawData) : undefined
+  try {
-
+    rawData = await AsyncStorage.getItem(BSKY_STORAGE)
-  // new user
+  } catch (e) {
-  if (!objData) return undefined
+    logger.error(`persisted state: failed reading root state from storage`, {
-
+      message: e,
-  // existing user, validate
+    })
-  const parsed = schema.safeParse(objData)
+  }
-
+  if (rawData) {
-  if (parsed.success) {
+    return tryParse(rawData)
-    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
   }
 }

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

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
+  }
+}