1 files changed, 75 insertions, 40 deletions
diff --git a/src/state/unstable-post-source.tsx b/src/state/unstable-post-source.tsx
index 43aac6f4d..ac126d79c 100644
--- a/src/state/unstable-post-source.tsx
+++ b/src/state/unstable-post-source.tsx
@@ -1,62 +1,97 @@
-import {createContext, useCallback, useContext, useRef, useState} from 'react'
-import {type AppBskyFeedDefs} from '@atproto/api'
+import {useEffect, useId, useState} from 'react'
+import {type AppBskyFeedDefs, AtUri} from '@atproto/api'
 
-import {type FeedDescriptor} from './queries/post-feed'
+import {Logger} from '#/logger'
+import {type FeedDescriptor} from '#/state/queries/post-feed'
 
 /**
- * For passing the source of the post (i.e. the original post, from the feed) to the threadview,
- * without using query params. Deliberately unstable to avoid using query params, use for FeedFeedback
- * and other ephemeral non-critical systems.
+ * Separate logger for better debugging
  */
+const logger = Logger.create(Logger.Context.PostSource)
 
-type Source = {
+export type PostSource = {
   post: AppBskyFeedDefs.FeedViewPost
   feed?: FeedDescriptor
 }
 
-const SetUnstablePostSourceContext = createContext<
-  (key: string, source: Source) => void
->(() => {})
-const ConsumeUnstablePostSourceContext = createContext<
-  (uri: string) => Source | undefined
->(() => undefined)
+/**
+ * A cache of sources that will be consumed by the post thread view. This is
+ * cleaned up any time a source is consumed.
+ */
+const transientSources = new Map<string, PostSource>()
 
-export function Provider({children}: {children: React.ReactNode}) {
-  const sourcesRef = useRef<Map<string, Source>>(new Map())
+/**
+ * A cache of sources that have been consumed by the post thread view. This is
+ * not cleaned up, but because we use a new ID for each post thread view that
+ * consumes a source, this is never reused unless a user navigates back to a
+ * post thread view that has not been dropped from memory.
+ */
+const consumedSources = new Map<string, PostSource>()
 
-  const setUnstablePostSource = useCallback((key: string, source: Source) => {
-    sourcesRef.current.set(key, source)
-  }, [])
+/**
+ * For stashing the feed that the user was browsing when they clicked on a post.
+ *
+ * Used for FeedFeedback and other ephemeral non-critical systems.
+ */
+export function setUnstablePostSource(key: string, source: PostSource) {
+  assertValid(
+    key,
+    `setUnstablePostSource key should be a URI containing a handle, received ${key} — use buildPostSourceKey`,
+  )
+  logger.debug('set', {key, source})
+  transientSources.set(key, source)
+}
 
-  const consumeUnstablePostSource = useCallback((uri: string) => {
-    const source = sourcesRef.current.get(uri)
+/**
+ * This hook is unstable and should only be used for FeedFeedback and other
+ * ephemeral non-critical systems. Views that use this hook will continue to
+ * return a reference to the same source until those views are dropped from
+ * memory.
+ */
+export function useUnstablePostSource(key: string) {
+  const id = useId()
+  const [source] = useState(() => {
+    assertValid(
+      key,
+      `consumeUnstablePostSource key should be a URI containing a handle, received ${key} — use buildPostSourceKey`,
+    )
+    const source = consumedSources.get(id) || transientSources.get(key)
     if (source) {
-      sourcesRef.current.delete(uri)
+      logger.debug('consume', {id, key, source})
+      transientSources.delete(key)
+      consumedSources.set(id, source)
     }
     return source
-  }, [])
-
-  return (
-    <SetUnstablePostSourceContext.Provider value={setUnstablePostSource}>
-      <ConsumeUnstablePostSourceContext.Provider
-        value={consumeUnstablePostSource}>
-        {children}
-      </ConsumeUnstablePostSourceContext.Provider>
-    </SetUnstablePostSourceContext.Provider>
-  )
-}
+  })
 
-export function useSetUnstablePostSource() {
-  return useContext(SetUnstablePostSourceContext)
+  useEffect(() => {
+    return () => {
+      consumedSources.delete(id)
+      logger.debug('cleanup', {id})
+    }
+  }, [id])
+
+  return source
 }
 
 /**
- * DANGER - This hook is unstable and should only be used for FeedFeedback
- * and other ephemeral non-critical systems. Does not change when the URI changes.
+ * Builds a post source key. This (atm) is a URI where the `host` is the post
+ * author's handle, not DID.
  */
-export function useUnstablePostSource(uri: string) {
-  const consume = useContext(ConsumeUnstablePostSourceContext)
+export function buildPostSourceKey(key: string, handle: string) {
+  const urip = new AtUri(key)
+  urip.host = handle
+  return urip.toString()
+}
 
-  const [source] = useState(() => consume(uri))
-  return source
+/**
+ * Just a lil dev helper
+ */
+function assertValid(key: string, message: string) {
+  if (__DEV__) {
+    const urip = new AtUri(key)
+    if (urip.host.startsWith('did:')) {
+      throw new Error(message)
+    }
+  }
 }