From cc8f77726a3abd815457b7a53554eddafcc98a6f Mon Sep 17 00:00:00 2001
From: Jon Staab <shtaab@gmail.com>
Date: Tue, 5 Mar 2024 16:09:21 -0800
Subject: [PATCH] Fix fallbacks

---
 src/connect/Socket.ts |  5 ++++-
 src/util/Deferred.ts  |  4 ++--
 src/util/Router.ts    | 16 ++++++++--------
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/connect/Socket.ts b/src/connect/Socket.ts
index 2fe564d..18a8ffd 100644
--- a/src/connect/Socket.ts
+++ b/src/connect/Socket.ts
@@ -107,9 +107,12 @@ export class Socket {
       const currentWs = this.ws
 
       this.ready.then(() => currentWs.close())
-      this.ready.resolve(false)
+      this.ready = defer()
       this.opts.onClose()
       this.ws = undefined
+
+      // Resolve a different instance of the promise
+      this.ready.resolve(false)
     }
   }
 }
diff --git a/src/util/Deferred.ts b/src/util/Deferred.ts
index e8fb93f..dd7ba2b 100644
--- a/src/util/Deferred.ts
+++ b/src/util/Deferred.ts
@@ -3,12 +3,12 @@ export type Deferred<T> = Promise<T> & {
   reject: (arg: T) => void
 }
 
-export const defer = (): Deferred<any> => {
+export const defer = <T>(): Deferred<T> => {
   let resolve, reject
   const p = new Promise((resolve_, reject_) => {
     resolve = resolve_
     reject = reject_
   })
 
-  return Object.assign(p, {resolve, reject}) as any
+  return (Object.assign(p, {resolve, reject}) as unknown) as Deferred<T>
 }
diff --git a/src/util/Router.ts b/src/util/Router.ts
index 90893d8..2a5fb5c 100644
--- a/src/util/Router.ts
+++ b/src/util/Router.ts
@@ -27,7 +27,7 @@ export type RouterOptions = {
 
 export type RouteScores = Record<string, {score: number, count: number}>
 
-export type FallbackPolicy = (limit: number) => number
+export type FallbackPolicy = (urls: string[], limit: number) => number
 
 export class Router {
   constructor(readonly options: RouterOptions) {}
@@ -65,7 +65,6 @@ export class Router {
       }
     })
 
-    // Use log-sum-exp to get a a weighted sum
     for (const [url, score] of Object.entries(scores)) {
       const quality = this.options.getRelayQuality?.(url) || 1
 
@@ -182,11 +181,11 @@ export class Router {
 
   // Fallback policies
 
-  addNoFallbacks = (limit: number) => 0
+  addNoFallbacks = (urls: string[], limit: number) => 0
 
-  addMinimalFallbacks = (limit: number) => 1
+  addMinimalFallbacks = (urls: string[], limit: number) => Math.max(0, 1 - urls.length)
 
-  addMaximalFallbacks = (limit: number) => limit
+  addMaximalFallbacks = (urls: string[], limit: number) => Math.max(0, limit - urls.length)
 
   // Higher level utils that use hints
 
@@ -242,11 +241,12 @@ export class RouterScenario {
 
   getUrls = () => {
     const fallbackPolicy = this.getPolicy()
-    const limit = fallbackPolicy(this.getLimit())
     const urls = this.router.scoreGroups(this.groups).map(s => s.url)
+    const limit = this.getLimit()
+    const limitWithFallbacks = limit + fallbackPolicy(urls, limit)
 
     for (const url of this.getFallbackRelays()) {
-      if (urls.length >= limit) {
+      if (urls.length >= limitWithFallbacks) {
         break
       }
 
@@ -255,7 +255,7 @@ export class RouterScenario {
       }
     }
 
-    return urls.slice(0, limit)
+    return urls.slice(0, limitWithFallbacks)
   }
 
   getUrl = () => first(this.limit(1).getUrls())