coder · Emyrk · Jun 1, 2023 · Jun 1, 2023 · Jun 1, 2023 · Jun 1, 2023
diff --git a/site/jest.setup.ts b/site/jest.setup.ts
@@ -19,7 +19,7 @@ jest.mock("contexts/useProxyLatency", () => ({
   useProxyLatency: (proxies?: RegionsResponse) => {
     // Must use `useMemo` here to avoid infinite loop.
     // Mocking the hook with a hook.
-    const latencies = useMemo(() => {
+    const proxyLatencies = useMemo(() => {
       if (!proxies) {
         return {} as Record<string, ProxyLatencyReport>
       }
@@ -35,7 +35,7 @@ jest.mock("contexts/useProxyLatency", () => ({
       }, {} as Record<string, ProxyLatencyReport>)
     }, [proxies])
 
-    return latencies
+    return { proxyLatencies, refetch: jest.fn() }
   },
 }))
 

diff --git a/site/src/components/AppLink/AppLink.stories.tsx b/site/src/components/AppLink/AppLink.stories.tsx
@@ -29,6 +29,9 @@ const Template: Story<AppLinkProps> = (args) => (
       clearProxy: () => {
         return
       },
+      refetchProxyLatencies: () => {
+        return
+      },
     }}
   >
     <AppLink {...args} />

diff --git a/site/src/components/Navbar/NavbarView.test.tsx b/site/src/components/Navbar/NavbarView.test.tsx
@@ -19,6 +19,7 @@ const proxyContextValue: ProxyContextValue = {
   isFetched: true,
   setProxy: jest.fn(),
   clearProxy: action("clearProxy"),
+  refetchProxyLatencies: jest.fn(),
   proxyLatencies: {},
 }
 

diff --git a/site/src/components/Navbar/NavbarView.tsx b/site/src/components/Navbar/NavbarView.tsx
@@ -189,6 +189,7 @@ const ProxyMenu: FC<{ proxyContextValue: ProxyContextValue }> = ({
   const buttonRef = useRef<HTMLButtonElement>(null)
   const [isOpen, setIsOpen] = useState(false)
   const selectedProxy = proxyContextValue.proxy.proxy
+  const refreshLatencies = proxyContextValue.refetchProxyLatencies
   const closeMenu = () => setIsOpen(false)
   const navigate = useNavigate()
 
@@ -293,6 +294,9 @@ const ProxyMenu: FC<{ proxyContextValue: ProxyContextValue }> = ({
         >
           Proxy settings
         </MenuItem>
+        <MenuItem sx={{ fontSize: 14 }} onClick={refreshLatencies}>
+          Refresh Latencies
+        </MenuItem>
       </Menu>
     </>
   )

diff --git a/site/src/components/Resources/AgentRow.stories.tsx b/site/src/components/Resources/AgentRow.stories.tsx
@@ -68,6 +68,9 @@ const TemplateFC = (
         clearProxy: () => {
           return
         },
+        refetchProxyLatencies: () => {
+          return
+        },
       }}
     >
       <AgentRow {...args} />

diff --git a/site/src/components/Resources/ResourceCard.stories.tsx b/site/src/components/Resources/ResourceCard.stories.tsx
@@ -33,6 +33,9 @@ Example.args = {
         clearProxy: () => {
           return
         },
+        refetchProxyLatencies: () => {
+          return
+        },
       }}
     >
       <AgentRow
@@ -103,6 +106,9 @@ BunchOfMetadata.args = {
         clearProxy: () => {
           return
         },
+        refetchProxyLatencies: () => {
+          return
+        },
       }}
     >
       <AgentRow

diff --git a/site/src/components/Workspace/Workspace.stories.tsx b/site/src/components/Workspace/Workspace.stories.tsx
@@ -53,6 +53,9 @@ const Template: Story<WorkspaceProps> = (args) => (
         setProxy: () => {
           return
         },
+        refetchProxyLatencies: () => {
+          return
+        },
       }}
     >
       <Workspace {...args} />

diff --git a/site/src/contexts/ProxyContext.test.tsx b/site/src/contexts/ProxyContext.test.tsx
@@ -26,7 +26,7 @@ import userEvent from "@testing-library/user-event"
 // here and not inside a unit test.
 jest.mock("contexts/useProxyLatency", () => ({
   useProxyLatency: () => {
-    return hardCodedLatencies
+    return { proxyLatencies: hardCodedLatencies, refetch: jest.fn() }
   },
 }))
 

diff --git a/site/src/contexts/ProxyContext.tsx b/site/src/contexts/ProxyContext.tsx
@@ -45,6 +45,9 @@ export interface ProxyContextValue {
   // then the latency has not been fetched yet. Calculations happen async for each proxy in the list.
   // Refer to the returned report for a given proxy for more information.
   proxyLatencies: Record<string, ProxyLatencyReport>
+  // refetchProxyLatencies will trigger refreshing of the proxy latencies. By default the latencies
+  // are loaded once.
+  refetchProxyLatencies: () => void
   // setProxy is a function that sets the user's selected proxy. This function should
   // only be called if the user is manually selecting a proxy. This value is stored in local
   // storage and will persist across reloads and tabs.
@@ -101,7 +104,8 @@ export const ProxyProvider: FC<PropsWithChildren> = ({ children }) => {
 
   // Every time we get a new proxiesResponse, update the latency check
   // to each workspace proxy.
-  const proxyLatencies = useProxyLatency(proxiesResp)
+  const { proxyLatencies, refetch: refetchProxyLatencies } =
+    useProxyLatency(proxiesResp)
 
   // updateProxy is a helper function that when called will
   // update the proxy being used.
@@ -128,6 +132,7 @@ export const ProxyProvider: FC<PropsWithChildren> = ({ children }) => {
     <ProxyContext.Provider
       value={{
         proxyLatencies,
+        refetchProxyLatencies,
         userProxy: userSavedProxy,
         proxy: experimentEnabled
           ? proxy

diff --git a/site/src/contexts/useProxyLatency.ts b/site/src/contexts/useProxyLatency.ts
@@ -1,5 +1,5 @@
 import { Region, RegionsResponse } from "api/typesGenerated"
-import { useEffect, useReducer } from "react"
+import { useEffect, useReducer, useState } from "react"
 import PerformanceObserver from "@fastly/performance-observer-polyfill"
 import axios from "axios"
 import { generateRandomString } from "utils/random"
@@ -25,20 +25,31 @@ const proxyLatenciesReducer = (
   action: ProxyLatencyAction,
 ): Record<string, ProxyLatencyReport> => {
   // Just overwrite any existing latency.
-  return {
-    ...state,
-    [action.proxyID]: action.report,
-  }
+  state[action.proxyID] = action.report
+  return state
 }
 
 export const useProxyLatency = (
   proxies?: RegionsResponse,
-): Record<string, ProxyLatencyReport> => {
+): {
+  // Refetch can be called to refetch the proxy latencies.
+  // Until the new values are loaded, the old values will still be used.
+  refetch: () => void
+  proxyLatencies: Record<string, ProxyLatencyReport>
+} => {
   const [proxyLatencies, dispatchProxyLatencies] = useReducer(
     proxyLatenciesReducer,
     {},
   )
 
+  // This latestFetchRequest is used to trigger a refetch of the proxy latencies.
+  const [latestFetchRequest, setLatestFetchRequest] = useState(
+    new Date().toISOString(),
+  )
+  const refetch = () => {
+    setLatestFetchRequest(new Date().toISOString())
+  }
+
   // Only run latency updates when the proxies change.
   useEffect(() => {
     if (!proxies) {
@@ -148,7 +159,10 @@ export const useProxyLatency = (
         // via the performance observer. So we can disconnect the observer.
         observer.disconnect()
       })
-  }, [proxies])
+  }, [proxies, latestFetchRequest])
 
-  return proxyLatencies
+  return {
+    proxyLatencies,
+    refetch,
+  }
 }
diff --git a/site/src/pages/TerminalPage/TerminalPage.test.tsx b/site/src/pages/TerminalPage/TerminalPage.test.tsx
@@ -55,6 +55,7 @@ const renderTerminal = () => {
               isLoading: false,
               setProxy: jest.fn(),
               clearProxy: jest.fn(),
+              refetchProxyLatencies: jest.fn(),
             }}
           >
             <TerminalPage renderer="dom" />
-Original file line number
+Diff line change
@@ Expand Up / @@ -29,6 +29,9 @@ const Template: Story<AppLinkProps> = (args) => ( @@
           clearProxy: () => {
             return
           },
+          refetchProxyLatencies: () => {
+            return
+          },
         }}
       >
         <AppLink {...args} />
@@ Expand Down @@