feat: add port reconnection mechanism (#575)

* add async just in case --------- Co-authored-by: L❤️ ☮️ ✋ <[email protected]>
PlasmoHQ · May 10, 2023 · c1eee1c · c1eee1c
1 parent e4cdb2e
commit c1eee1c
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 38 deletions.
diff --git a/packages/messaging/package.json b/packages/messaging/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@plasmohq/messaging",
-  "version": "0.2.1",
+  "version": "0.3.0",
   "description": "Type-safe, zero-config messaging library for modern browser extensions",
   "type": "module",
   "module": "./dist/index.js",
@@ -26,6 +26,11 @@
       "require": "./dist/background.cjs",
       "types": "./dist/background.d.ts"
     },
+    "./message": {
+      "import": "./dist/message.js",
+      "require": "./dist/message.cjs",
+      "types": "./dist/message.d.ts"
+    },
     ".": {
       "import": "./dist/index.js",
       "require": "./dist/index.cjs",
@@ -34,6 +39,9 @@
   },
   "typesVersions": {
     "*": {
+      "message": [
+        "./dist/message.d.ts"
+      ],
       "hook": [
         "./dist/hook.d.ts"
       ],
@@ -56,6 +64,7 @@
       "src/index.ts",
       "src/port.ts",
       "src/relay.ts",
+      "src/message.ts",
       "src/background.ts",
       "src/hook.ts"
     ],

diff --git a/packages/messaging/src/hook.ts b/packages/messaging/src/hook.ts
@@ -1,7 +1,8 @@
 import { useEffect, useRef, useState } from "react"
 
 import { type MessageName, type PlasmoMessaging, relayMessage } from "./index"
-import { getPort } from "./port"
+import { listen as messageListen } from "./message"
+import { listen as portListen } from "./port"
 import { relay } from "./relay"
 
 /**
@@ -12,30 +13,14 @@ export const useMessage = <RequestBody, ResponseBody>(
 ) => {
   const [data, setData] = useState<RequestBody>()
 
-  useEffect(() => {
-    const metaListener = async (req, sender, sendResponse) => {
-      setData(req.body)
-      await handler?.(
-        {
-          ...req,
-          sender
-        },
-        {
-          send: (p) => sendResponse(p)
-        }
-      )
-    }
-
-    const listener = (req, sender, sendResponse) => {
-      metaListener(req, sender, sendResponse)
-      return
-    }
-
-    chrome.runtime.onMessage.addListener(listener)
-    return () => {
-      chrome.runtime.onMessage.removeListener(listener)
-    }
-  }, [handler])
+  useEffect(
+    () =>
+      messageListen<RequestBody, ResponseBody>(async (req, res) => {
+        setData(req.body)
+        await handler(req, res)
+      }),
+    [handler]
+  )
 
   return {
     data
@@ -44,26 +29,30 @@ export const useMessage = <RequestBody, ResponseBody>(
 
 export const usePort: PlasmoMessaging.PortHook = (name) => {
   const portRef = useRef<chrome.runtime.Port>()
+  const reconnectRef = useRef(0)
   const [data, setData] = useState()
 
   useEffect(() => {
     if (!name) {
       return null
     }
 
-    const port = getPort(name)
-
-    function messageHandler(msg) {
-      setData(msg)
-    }
-
-    port.onMessage.addListener(messageHandler)
+    const { port, disconnect } = portListen(
+      name,
+      (msg) => {
+        setData(msg)
+      },
+      () => {
+        reconnectRef.current = reconnectRef.current + 1
+      }
+    )
 
     portRef.current = port
-    return () => {
-      port.onMessage.removeListener(messageHandler)
-    }
-  }, [name])
+    return disconnect
+  }, [
+    name,
+    reconnectRef.current // This is needed to force a new port ref
+  ])
 
   return {
     data,
@@ -72,7 +61,9 @@ export const usePort: PlasmoMessaging.PortHook = (name) => {
         name,
         body
       })
-    }
+    },
+    listen: <T = string>(handler: (msg: T) => void) =>
+      portListen<T>(name, handler)
   }
 }
 

diff --git a/packages/messaging/src/message.ts b/packages/messaging/src/message.ts
@@ -0,0 +1,27 @@
+import { type PlasmoMessaging } from "./index"
+
+export const listen = <RequestBody, ResponseBody>(
+  handler: PlasmoMessaging.Handler<string, RequestBody, ResponseBody>
+) => {
+  const metaListener = async (req, sender, sendResponse) => {
+    await handler?.(
+      {
+        ...req,
+        sender
+      },
+      {
+        send: (p) => sendResponse(p)
+      }
+    )
+  }
+
+  const listener = (req, sender, sendResponse) => {
+    metaListener(req, sender, sendResponse)
+    return // Syncronous return to indicate this is an async listener
+  }
+
+  chrome.runtime.onMessage.addListener(listener)
+  return () => {
+    chrome.runtime.onMessage.removeListener(listener)
+  }
+}
diff --git a/packages/messaging/src/port.ts b/packages/messaging/src/port.ts
@@ -11,3 +11,31 @@ export const getPort = (name: PortName) => {
   portMap.set(name, newPort)
   return newPort
 }
+
+export const removePort = (name: PortName) => {
+  portMap.delete(name)
+}
+
+export const listen = <ResponseBody = any>(
+  name: PortName,
+  handler: (msg: ResponseBody) => Promise<void> | void,
+  onReconnect?: () => void
+) => {
+  const port = getPort(name)
+
+  function reconnectHandler() {
+    removePort(name)
+    onReconnect?.()
+  }
+
+  port.onMessage.addListener(handler)
+  port.onDisconnect.addListener(reconnectHandler)
+
+  return {
+    port,
+    disconnect: () => {
+      port.onMessage.removeListener(handler)
+      port.onDisconnect.removeListener(reconnectHandler)
+    }
+  }
+}