feat: Workspace Proxy picker show latency to each proxy (#7486)

* chore: Add cors to workspace proxies to allow for latency checks
* Add latency check to wsproxy

Use performance API timings.
- Fix cors and timing headers
- Accept custom headers

site/src/components/AppLink/AppLink.stories.tsx

@@ -5,6 +5,7 @@ import {
   MockWorkspace,
   MockWorkspaceAgent,
   MockWorkspaceApp,
+  MockProxyLatencies,
 } from "testHelpers/entities"
 import { AppLink, AppLinkProps } from "./AppLink"
 import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext"
@@ -17,6 +18,7 @@ export default {
 const Template: Story<AppLinkProps> = (args) => (
   <ProxyContext.Provider
     value={{
+      proxyLatencies: MockProxyLatencies,
       proxy: getPreferredProxy(MockWorkspaceProxies, MockPrimaryWorkspaceProxy),
       proxies: MockWorkspaceProxies,
       isLoading: false,

site/src/components/Resources/AgentLatency.tsx

@@ -1,5 +1,6 @@
 import { useRef, useState, FC } from "react"
 import { makeStyles, useTheme } from "@mui/styles"
+import { Theme } from "@mui/material/styles"
 import {
   HelpTooltipText,
   HelpPopover,
@@ -7,7 +8,7 @@ import {
 } from "components/Tooltips/HelpTooltip"
 import { Stack } from "components/Stack/Stack"
 import { WorkspaceAgent, DERPRegion } from "api/typesGenerated"
-import { Theme } from "@mui/material/styles"
+import { getLatencyColor } from "utils/colors"
 
 const getDisplayLatency = (theme: Theme, agent: WorkspaceAgent) => {
   // Find the right latency to display
@@ -22,17 +23,9 @@ const getDisplayLatency = (theme: Theme, agent: WorkspaceAgent) => {
     return undefined
   }
 
-  // Get the color
-  let color = theme.palette.success.light
-  if (latency.latency_ms >= 150 && latency.latency_ms < 300) {
-    color = theme.palette.warning.light
-  } else if (latency.latency_ms >= 300) {
-    color = theme.palette.error.light
-  }
-
   return {
     ...latency,
-    color,
+    color: getLatencyColor(theme, latency.latency_ms),
   }
 }
 

site/src/components/Resources/AgentRow.stories.tsx

@@ -16,6 +16,7 @@ import {
   MockWorkspaceAgentStartTimeout,
   MockWorkspaceAgentTimeout,
   MockWorkspaceApp,
+  MockProxyLatencies,
 } from "testHelpers/entities"
 import { AgentRow, AgentRowProps } from "./AgentRow"
 import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext"
@@ -56,6 +57,7 @@ const TemplateFC = (
   return (
     <ProxyContext.Provider
       value={{
+        proxyLatencies: MockProxyLatencies,
         proxy: getPreferredProxy(proxies, selectedProxy),
         proxies: proxies,
         isLoading: false,

site/src/components/Resources/ResourceCard.stories.tsx

@@ -1,6 +1,10 @@
 import { action } from "@storybook/addon-actions"
 import { Story } from "@storybook/react"
-import { MockWorkspace, MockWorkspaceResource } from "testHelpers/entities"
+import {
+  MockProxyLatencies,
+  MockWorkspace,
+  MockWorkspaceResource,
+} from "testHelpers/entities"
 import { AgentRow } from "./AgentRow"
 import { ResourceCard, ResourceCardProps } from "./ResourceCard"
 import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext"
@@ -18,6 +22,7 @@ Example.args = {
   agentRow: (agent) => (
     <ProxyContext.Provider
       value={{
+        proxyLatencies: MockProxyLatencies,
         proxy: getPreferredProxy([], undefined),
         proxies: [],
         isLoading: false,
@@ -84,6 +89,7 @@ BunchOfMetadata.args = {
   agentRow: (agent) => (
     <ProxyContext.Provider
       value={{
+        proxyLatencies: MockProxyLatencies,
         proxy: getPreferredProxy([], undefined),
         proxies: [],
         isLoading: false,

site/src/components/Workspace/Workspace.stories.tsx

@@ -7,6 +7,7 @@ import { Workspace, WorkspaceErrors, WorkspaceProps } from "./Workspace"
 import { withReactContext } from "storybook-react-context"
 import EventSource from "eventsourcemock"
 import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext"
+import { MockProxyLatencies } from "../../testHelpers/entities"
 
 export default {
   title: "components/Workspace",
@@ -26,6 +27,7 @@ export default {
 const Template: Story<WorkspaceProps> = (args) => (
   <ProxyContext.Provider
     value={{
+      proxyLatencies: MockProxyLatencies,
       proxy: getPreferredProxy([], undefined),
       proxies: [],
       isLoading: false,

site/src/contexts/ProxyContext.tsx

@@ -9,10 +9,12 @@ import {
   useContext,
   useState,
 } from "react"
+import { ProxyLatencyReport, useProxyLatency } from "./useProxyLatency"
 
 interface ProxyContextValue {
   proxy: PreferredProxy
   proxies?: Region[]
+  proxyLatencies?: Record<string, ProxyLatencyReport>
   // isfetched is true when the proxy api call is complete.
   isFetched: boolean
   // isLoading is true if the proxy is in the process of being fetched.
@@ -72,6 +74,10 @@ export const ProxyProvider: FC<PropsWithChildren> = ({ children }) => {
     },
   })
 
+  // Everytime we get a new proxiesResponse, update the latency check
+  // to each workspace proxy.
+  const proxyLatencies = useProxyLatency(proxiesResp)
+
   const setAndSaveProxy = (
     selectedProxy?: Region,
     // By default the proxies come from the api call above.
@@ -95,6 +101,7 @@ export const ProxyProvider: FC<PropsWithChildren> = ({ children }) => {
   return (
     <ProxyContext.Provider
       value={{
+        proxyLatencies: proxyLatencies,
         proxy: experimentEnabled
           ? proxy
           : {

site/src/contexts/useProxyLatency.ts

@@ -0,0 +1,152 @@
+import { Region, RegionsResponse } from "api/typesGenerated"
+import { useEffect, useReducer } from "react"
+import PerformanceObserver from "@fastly/performance-observer-polyfill"
+import axios from "axios"
+import { generateRandomString } from "utils/random"
+
+export interface ProxyLatencyReport {
+  // accurate identifies if the latency was calculated using the
+  // PerformanceResourceTiming API. If this is false, then the
+  // latency is calculated using the total duration of the request
+  // and will be off by a good margin.
+  accurate: boolean
+  latencyMS: number
+  // at is when the latency was recorded.
+  at: Date
+}
+
+interface ProxyLatencyAction {
+  proxyID: string
+  report: ProxyLatencyReport
+}
+
+const proxyLatenciesReducer = (
+  state: Record<string, ProxyLatencyReport>,
+  action: ProxyLatencyAction,
+): Record<string, ProxyLatencyReport> => {
+  // Just overwrite any existing latency.
+  state[action.proxyID] = action.report
+  return state
+}
+
+export const useProxyLatency = (
+  proxies?: RegionsResponse,
+): Record<string, ProxyLatencyReport> => {
+  const [proxyLatencies, dispatchProxyLatencies] = useReducer(
+    proxyLatenciesReducer,
+    {},
+  )
+
+  // Only run latency updates when the proxies change.
+  useEffect(() => {
+    if (!proxies) {
+      return
+    }
+
+    // proxyMap is a map of the proxy path_app_url to the proxy object.
+    // This is for the observer to know which requests are important to
+    // record.
+    const proxyChecks = proxies.regions.reduce((acc, proxy) => {
+      // Only run the latency check on healthy proxies.
+      if (!proxy.healthy) {
+        return acc
+      }
+
+      // Add a random query param to the url to make sure we don't get a cached response.
+      // This is important in case there is some caching layer between us and the proxy.
+      const url = new URL(
+        `/latency-check?cache_bust=${generateRandomString(6)}`,
+        proxy.path_app_url,
+      )
+      acc[url.toString()] = proxy
+      return acc
+    }, {} as Record<string, Region>)
+
+    // dispatchProxyLatenciesGuarded will assign the latency to the proxy
+    // via the reducer. But it will only do so if the performance entry is
+    // a resource entry that we care about.
+    const dispatchProxyLatenciesGuarded = (entry: PerformanceEntry): void => {
+      if (entry.entryType !== "resource") {
+        // We should never get these, but just in case.
+        return
+      }
+
+      // The entry.name is the url of the request.
+      const check = proxyChecks[entry.name]
+      if (!check) {
+        // This is not a proxy request, so ignore it.
+        return
+      }
+
+      // These docs are super useful.
+      // https://developer.mozilla.org/en-US/docs/Web/API/Performance_API/Resource_timing
+      let latencyMS = 0
+      let accurate = false
+      if (
+        "requestStart" in entry &&
+        (entry as PerformanceResourceTiming).requestStart !== 0
+      ) {
+        // This is the preferred logic to get the latency.
+        const timingEntry = entry as PerformanceResourceTiming
+        latencyMS = timingEntry.responseStart - timingEntry.requestStart
+        accurate = true
+      } else {
+        // This is the total duration of the request and will be off by a good margin.
+        // This is a fallback if the better timing is not available.
+        // eslint-disable-next-line no-console -- We can remove this when we display the "accurate" bool on the UI
+        console.log(
+          `Using fallback latency calculation for "${entry.name}". Latency will be incorrect and larger then actual.`,
+        )
+        latencyMS = entry.duration
+      }
+      dispatchProxyLatencies({
+        proxyID: check.id,
+        report: {
+          latencyMS,
+          accurate,
+          at: new Date(),
+        },
+      })
+
+      return
+    }
+
+    // Start a new performance observer to record of all the requests
+    // to the proxies.
+    const observer = new PerformanceObserver((list) => {
+      // If we get entries via this callback, then dispatch the events to the latency reducer.
+      list.getEntries().forEach((entry) => {
+        dispatchProxyLatenciesGuarded(entry)
+      })
+    })
+
+    // The resource requests include xmlhttp requests.
+    observer.observe({ entryTypes: ["resource"] })
+
+    const proxyRequests = Object.keys(proxyChecks).map((latencyURL) => {
+      return axios.get(latencyURL, {
+        withCredentials: false,
+        // Must add a custom header to make the request not a "simple request"
+        // https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS#simple_requests
+        headers: { "X-LATENCY-CHECK": "true" },
+      })
+    })
+
+    // When all the proxy requests finish
+    Promise.all(proxyRequests)
+      // TODO: If there is an error on any request, we might want to store some indicator of that?
+      .finally(() => {
+        // takeRecords will return any entries that were not called via the callback yet.
+        // We want to call this before we disconnect the observer to make sure we get all the
+        // proxy requests recorded.
+        observer.takeRecords().forEach((entry) => {
+          dispatchProxyLatenciesGuarded(entry)
+        })
+        // At this point, we can be confident that all the proxy requests have been recorded
+        // via the performance observer. So we can disconnect the observer.
+        observer.disconnect()
+      })
+  }, [proxies])
+
+  return proxyLatencies
+}

site/src/pages/TerminalPage/TerminalPage.test.tsx

@@ -4,6 +4,7 @@ import WS from "jest-websocket-mock"
 import { rest } from "msw"
 import {
   MockPrimaryWorkspaceProxy,
+  MockProxyLatencies,
   MockWorkspace,
   MockWorkspaceAgent,
   MockWorkspaceProxies,
@@ -43,6 +44,7 @@ const renderTerminal = () => {
         element={
           <ProxyContext.Provider
             value={{
+              proxyLatencies: MockProxyLatencies,
               proxy: {
                 selectedProxy: MockPrimaryWorkspaceProxy,
                 preferredPathAppURL: "",

site/src/pages/UserSettingsPage/WorkspaceProxyPage/WorkspaceProxyPage.tsx

@@ -14,6 +14,7 @@ export const WorkspaceProxyPage: FC<PropsWithChildren<unknown>> = () => {
     "This selection only affects browser connections to your workspace."
 
   const {
+    proxyLatencies,
     proxies,
     error: proxiesError,
     isFetched: proxiesFetched,
@@ -30,6 +31,7 @@ export const WorkspaceProxyPage: FC<PropsWithChildren<unknown>> = () => {
       layout="fluid"
     >
       <WorkspaceProxyView
+        proxyLatencies={proxyLatencies}
         proxies={proxies}
         isLoading={proxiesLoading}
         hasLoaded={proxiesFetched}

site/src/pages/UserSettingsPage/WorkspaceProxyPage/WorkspaceProxyRow.tsx

@@ -9,15 +9,19 @@ import {
   HealthyBadge,
   NotHealthyBadge,
 } from "components/DeploySettingsLayout/Badges"
-import { makeStyles } from "@mui/styles"
+import { makeStyles, useTheme } from "@mui/styles"
 import { combineClasses } from "utils/combineClasses"
+import { ProxyLatencyReport } from "contexts/useProxyLatency"
+import { getLatencyColor } from "utils/colors"
 
 export const ProxyRow: FC<{
+  latency?: ProxyLatencyReport
   proxy: Region
   onSelectRegion: (proxy: Region) => void
   preferred: boolean
-}> = ({ proxy, onSelectRegion, preferred }) => {
+}> = ({ proxy, onSelectRegion, preferred, latency }) => {
   const styles = useStyles()
+  const theme = useTheme()
 
   const clickable = useClickableTableRow(() => {
     onSelectRegion(proxy)
@@ -53,6 +57,15 @@ export const ProxyRow: FC<{
       <TableCell>
         <ProxyStatus proxy={proxy} />
       </TableCell>
+      <TableCell>
+        <span
+          style={{
+            color: latency ? getLatencyColor(theme, latency.latencyMS) : "",
+          }}
+        >
+          {latency ? `${latency.latencyMS.toFixed(1)} ms` : "?"}
+        </span>
+      </TableCell>
     </TableRow>
   )
 }

site/src/pages/UserSettingsPage/WorkspaceProxyPage/WorkspaceProxyView.tsx

@@ -12,9 +12,11 @@ import { FC } from "react"
 import { AlertBanner } from "components/AlertBanner/AlertBanner"
 import { Region } from "api/typesGenerated"
 import { ProxyRow } from "./WorkspaceProxyRow"
+import { ProxyLatencyReport } from "contexts/useProxyLatency"
 
 export interface WorkspaceProxyViewProps {
   proxies?: Region[]
+  proxyLatencies?: Record<string, ProxyLatencyReport>
   getWorkspaceProxiesError?: Error | unknown
   isLoading: boolean
   hasLoaded: boolean
@@ -27,6 +29,7 @@ export const WorkspaceProxyView: FC<
   React.PropsWithChildren<WorkspaceProxyViewProps>
 > = ({
   proxies,
+  proxyLatencies,
   getWorkspaceProxiesError,
   isLoading,
   hasLoaded,
@@ -49,6 +52,7 @@ export const WorkspaceProxyView: FC<
               <TableCell width="40%">Proxy</TableCell>
               <TableCell width="30%">URL</TableCell>
               <TableCell width="10%">Status</TableCell>
+              <TableCell width="20%">Latency</TableCell>
             </TableRow>
           </TableHead>
           <TableBody>
@@ -62,6 +66,7 @@ export const WorkspaceProxyView: FC<
               <Cond>
                 {proxies?.map((proxy) => (
                   <ProxyRow
+                    latency={proxyLatencies?.[proxy.id]}
                     key={proxy.id}
                     proxy={proxy}
                     onSelectRegion={onSelect}

site/src/pages/UserSettingsPage/WorkspaceProxyPage/WorspaceProxyView.stories.tsx

@@ -4,6 +4,7 @@ import {
   MockWorkspaceProxies,
   MockPrimaryWorkspaceProxy,
   MockHealthyWildWorkspaceProxy,
+  MockProxyLatencies,
 } from "testHelpers/entities"
 import {
   WorkspaceProxyView,
@@ -27,6 +28,7 @@ PrimarySelected.args = {
   isLoading: false,
   hasLoaded: true,
   proxies: MockWorkspaceProxies,
+  proxyLatencies: MockProxyLatencies,
   preferredProxy: MockPrimaryWorkspaceProxy,
   onSelect: () => {
     return Promise.resolve()
@@ -38,6 +40,7 @@ Example.args = {
   isLoading: false,
   hasLoaded: true,
   proxies: MockWorkspaceProxies,
+  proxyLatencies: MockProxyLatencies,
   preferredProxy: MockHealthyWildWorkspaceProxy,
   onSelect: () => {
     return Promise.resolve()

site/src/testHelpers/entities.ts

@@ -7,6 +7,7 @@ import range from "lodash/range"
 import { Permissions } from "xServices/auth/authXService"
 import { TemplateVersionFiles } from "utils/templateVersion"
 import { FileTree } from "utils/filetree"
+import { ProxyLatencyReport } from "contexts/useProxyLatency"
 
 export const MockOrganization: TypesGen.Organization = {
   id: "fc0774ce-cc9e-48d4-80ae-88f7a4d4a8b0",
@@ -113,6 +114,34 @@ export const MockWorkspaceProxies: TypesGen.Region[] = [
   },
 ]
 
+export const MockProxyLatencies: Record<string, ProxyLatencyReport> = {
+  ...MockWorkspaceProxies.reduce((acc, proxy) => {
+    if (!proxy.healthy) {
+      return acc
+    }
+    acc[proxy.id] = {
+      // Make one of them inaccurate.
+      accurate: proxy.id !== "26e84c16-db24-4636-a62d-aa1a4232b858",
+      // This is a deterministic way to generate a latency to for each proxy.
+      // It will be the same for each run as long as the IDs don't change.
+      latencyMS:
+        (Number(
+          Array.from(proxy.id).reduce(
+            // Multiply each char code by some large prime number to increase the
+            // size of the number and allow use to get some decimal points.
+            (acc, char) => acc + char.charCodeAt(0) * 37,
+            0,
+          ),
+        ) /
+          // Cap at 250ms
+          100) %
+        250,
+      at: new Date(),
+    }
+    return acc
+  }, {} as Record<string, ProxyLatencyReport>),
+}
+
 export const MockBuildInfo: TypesGen.BuildInfoResponse = {
   external_url: "file:///mock-url",
   version: "v99.999.9999+c9cdf14",

site/src/utils/colors.ts

@@ -1,3 +1,5 @@
+import { Theme } from "@mui/material/styles"
+
 // Used to convert our theme colors to Hex since monaco theme only support hex colors
 // From https://www.jameslmilner.com/posts/converting-rgb-hex-hsl-colors/
 export function hslToHex(hsl: string): string {
@@ -21,3 +23,15 @@ export function hslToHex(hsl: string): string {
   }
   return `#${f(0)}${f(8)}${f(4)}`
 }
+
+// getLatencyColor is the text color to use for a given latency
+// in milliseconds.
+export const getLatencyColor = (theme: Theme, latency: number) => {
+  let color = theme.palette.success.light
+  if (latency >= 150 && latency < 300) {
+    color = theme.palette.warning.light
+  } else if (latency >= 300) {
+    color = theme.palette.error.light
+  }
+  return color
+}