Clean up policy a tad

packages/net2/src/policy.ts

@@ -1,4 +1,4 @@
-import {on, sleep, spec, ago, now} from "@welshman/lib"
+import {on, call, sleep, spec, ago, now} from "@welshman/lib"
 import {AUTH_JOIN} from "@welshman/util"
 import {
   ClientMessage,
@@ -16,15 +16,17 @@ import {AuthState, AuthStatus, AuthStateEventType} from "./auth.js"
 
 // Pause sending messages when the socket isn't open
 export const socketPolicySendWhenOpen = (socket: Socket) => {
-  const unsubscribe = on(socket, SocketEventType.Status, (newStatus: SocketStatus) => {
+  const unsubscribers = [
+    on(socket, SocketEventType.Status, (newStatus: SocketStatus) => {
       if (newStatus === SocketStatus.Open) {
         socket._sendQueue.start()
       } else {
         socket._sendQueue.stop()
       }
-  })
+    }),
+  ]
 
-  return unsubscribe
+  return () => unsubscribers.forEach(call)
 }
 
 export const socketPolicyDeferOnAuth = (socket: Socket) => {
@@ -32,8 +34,9 @@ export const socketPolicyDeferOnAuth = (socket: Socket) => {
   const authState = new AuthState(socket)
   const okStatuses = [AuthStatus.None, AuthStatus.Ok]
 
+  const unsubscribers = [
     // Pause sending certain messages when we're not authenticated
-  const unsubscribeEnqueue = on(socket, SocketEventType.Enqueue, (message: ClientMessage) => {
+    on(socket, SocketEventType.Enqueue, (message: ClientMessage) => {
       // If we're closing a request, but it never got sent, remove both from the queue
       // Otherwise, always send CLOSE
       if (isClientClose(message)) {
@@ -58,20 +61,19 @@ export const socketPolicyDeferOnAuth = (socket: Socket) => {
         buffer.push(message)
         socket._sendQueue.remove(message)
       }
-  })
+    }),
-
     // Send buffered messages when we get successful auth
-  const unsubscribeAuthStatus = on(authState, AuthStateEventType.Status, (status: AuthStatus) => {
+    on(authState, AuthStateEventType.Status, (status: AuthStatus) => {
       if (okStatuses.includes(status) && buffer.length > 0) {
         for (const message of buffer.splice(0)) {
           socket.send(message)
         }
       }
-  })
+    }),
+  ]
 
   return () => {
-    unsubscribeAuthStatus()
+    unsubscribers.forEach(call)
-    unsubscribeEnqueue()
     authState.cleanup()
   }
 }
@@ -80,8 +82,9 @@ export const socketPolicyRetryAuthRequired = (socket: Socket) => {
   const retried = new Set<string>()
   const pending = new Map<string, ClientMessage>()
 
+  const unsubscribers = [
     // Watch outgoing events and requests and keep a copy
-  const unsubscribeSend = on(socket, SocketEventType.Send, (message: ClientMessage) => {
+    on(socket, SocketEventType.Send, (message: ClientMessage) => {
       if (isClientEvent(message)) {
         const [_, event] = message
 
@@ -97,10 +100,9 @@ export const socketPolicyRetryAuthRequired = (socket: Socket) => {
           pending.set(id, message)
         }
       }
-  })
+    }),
-
     // If a message is rejected with auth-required, re-enqueue it one time
-  const unsubscribeReceive = on(socket, SocketEventType.Receive, (message: RelayMessage) => {
+    on(socket, SocketEventType.Receive, (message: RelayMessage) => {
       if (isRelayOk(message)) {
         const [_, id, ok, detail] = message
         const pendingMessage = pending.get(id)
@@ -124,19 +126,18 @@ export const socketPolicyRetryAuthRequired = (socket: Socket) => {
 
         pending.delete(id)
       }
-  })
+    }),
+  ]
 
-  return () => {
+  return () => unsubscribers.forEach(call)
-    unsubscribeSend()
-    unsubscribeReceive()
-  }
 }
 
 export const socketPolicyConnectOnSend = (socket: Socket) => {
   let lastError = 0
   let currentStatus = SocketStatus.Closed
 
-  const unsubscribeStatus = on(socket, SocketEventType.Status, (newStatus: SocketStatus) => {
+  const unsubscribers = [
+    on(socket, SocketEventType.Status, (newStatus: SocketStatus) => {
       // Keep track of the most recent error
       if (newStatus === SocketStatus.Error) {
         lastError = now()
@@ -144,31 +145,29 @@ export const socketPolicyConnectOnSend = (socket: Socket) => {
 
       // Keep track of the current status
       currentStatus = newStatus
-  })
+    }),
-
+    on(socket, SocketEventType.Send, (message: ClientMessage) => {
-  const unsubscribeSend = on(socket, SocketEventType.Send, (message: ClientMessage) => {
       // When a new message is sent, make sure the socket is open (unless there was a recent error)
       if (currentStatus === SocketStatus.Closed && now() - lastError < ago(30)) {
         socket.open()
       }
-  })
+    }),
+  ]
 
-  return () => {
+  return () => unsubscribers.forEach(call)
-    unsubscribeStatus()
-    unsubscribeSend()
-  }
 }
 
 export const socketPolicyCloseOnTimeout = (socket: Socket) => {
   let lastActivity = 0
 
-  const unsubscribeSend = on(socket, SocketEventType.Send, (message: ClientMessage) => {
+  const unsubscribers = [
+    on(socket, SocketEventType.Send, (message: ClientMessage) => {
       lastActivity = now()
-  })
+    }),
-
+    on(socket, SocketEventType.Receive, (message: RelayMessage) => {
-  const unsubscribeReceive = on(socket, SocketEventType.Receive, (message: RelayMessage) => {
       lastActivity = now()
-  })
+    }),
+  ]
 
   const interval = setInterval(() => {
     if (lastActivity < ago(30)) {
@@ -177,8 +176,7 @@ export const socketPolicyCloseOnTimeout = (socket: Socket) => {
   }, 3000)
 
   return () => {
-    unsubscribeSend()
+    unsubscribers.forEach(call)
-    unsubscribeReceive()
     clearInterval(interval)
   }
 }
@@ -188,7 +186,8 @@ export const socketPolicyReopenActive = (socket: Socket) => {
 
   let lastOpen = 0
 
-  const unsubscribeStatus = on(socket, SocketEventType.Status, (newStatus: SocketStatus) => {
+  const unsubscribers = [
+    on(socket, SocketEventType.Status, (newStatus: SocketStatus) => {
       // Keep track of the most recent error
       if (newStatus === SocketStatus.Open) {
         lastOpen = Date.now()
@@ -202,9 +201,8 @@ export const socketPolicyReopenActive = (socket: Socket) => {
           }
         })
       }
-  })
+    }),
-
+    on(socket, SocketEventType.Send, (message: ClientMessage) => {
-  const unsubscribeSend = on(socket, SocketEventType.Send, (message: ClientMessage) => {
       if (isClientEvent(message)) {
         pending.set(message[1].id, message)
       }
@@ -216,19 +214,15 @@ export const socketPolicyReopenActive = (socket: Socket) => {
       if (isClientClose(message)) {
         pending.delete(message[1])
       }
-  })
+    }),
-
+    on(socket, SocketEventType.Receive, (message: RelayMessage) => {
-  const unsubscribeReceive = on(socket, SocketEventType.Receive, (message: RelayMessage) => {
       if (isRelayClosed(message) || isRelayOk(message)) {
         pending.delete(message[1])
       }
-  })
+    }),
+  ]
 
-  return () => {
+  return () => unsubscribers.forEach(call)
-    unsubscribeStatus()
-    unsubscribeSend()
-    unsubscribeReceive()
-  }
 }
 
 export const defaultSocketPolicies = [