Proxy client count change requests through BRAD to avoid CORS restric…

…tions (#485)

* Proxy requests through BRAD to avoid CORS restrictions

* Fix lint

* Fixes

* Fix lint, formatting

setup.py

@@ -68,6 +68,8 @@
     "fastapi",
     "uvicorn[standard]",
     "pydantic",
+    "requests",
+    "types-requests",
 ]
 
 KEYWORDS = []

src/brad/ui/manager_impl.py

@@ -3,10 +3,10 @@
 import logging
 import importlib.resources as pkg_resources
 import numpy as np
-from fastapi import FastAPI
+import requests
+from fastapi import FastAPI, HTTPException
 from fastapi.staticfiles import StaticFiles
 from typing import Optional, List
-from pydantic import BaseModel
 
 import brad.ui.static as brad_app
 from brad.blueprint import Blueprint
@@ -26,6 +26,8 @@
     VirtualInfrastructure,
     DisplayableTable,
     Status,
+    ClientState,
+    SetClientState,
 )
 from brad.daemon.front_end_metrics import FrontEndMetric
 from brad.daemon.system_event_logger import SystemEventLogger, SystemEventRecord
@@ -190,25 +192,44 @@ def get_system_state(filter_tables_for_demo: bool = False) -> SystemState:
     return system_state
 
 
-class ClientState(BaseModel):
-    max_clients: int
-    curr_clients: int
-
-
-class SetClientState(BaseModel):
-    curr_clients: int
-
-
-@app.get("/clients")
-def get_clients_dummy() -> ClientState:
-    # Used for debugging without starting the variable client runner.
-    return ClientState(max_clients=12, curr_clients=3)
-
-
-@app.post("/clients")
-def set_clients_dummy(clients: SetClientState) -> ClientState:
-    # Used for debugging without starting the variable client runner.
-    return ClientState(max_clients=12, curr_clients=clients.curr_clients)
+@app.get("/api/1/clients")
+def get_workload_clients(runner_port: Optional[int] = None) -> ClientState:
+    # This proxies the request to the runner, which runs as a different process
+    # and listens for requests on a different port. We require a proxy to avoid
+    # CORS restrictions.
+    if runner_port is None:
+        # Used for debugging without starting the variable client runner.
+        return ClientState(max_clients=12, curr_clients=3)
+    else:
+        try:
+            r = requests.get(f"http://localhost:{runner_port}/clients", timeout=2)
+            if r.status_code != 200:
+                raise HTTPException(r.status_code, r.reason)
+            return ClientState(**r.json())
+        except requests.ConnectionError as ex:
+            raise HTTPException(400, f"Unable to connect to port {runner_port}") from ex
+
+
+@app.post("/api/1/clients")
+def set_clients(clients: SetClientState) -> ClientState:
+    # This proxies the request to the runner, which runs as a different process
+    # and listens for requests on a different port. We require a proxy to avoid
+    # CORS restrictions.
+    if clients.runner_port is None:
+        # Used for debugging without starting the variable client runner.
+        return ClientState(max_clients=12, curr_clients=clients.curr_clients)
+    else:
+        try:
+            r = requests.post(
+                f"http://localhost:{clients.runner_port}/clients", timeout=2
+            )
+            if r.status_code != 200:
+                raise HTTPException(r.status_code, r.reason)
+            return ClientState(**r.json())
+        except requests.ConnectionError as ex:
+            raise HTTPException(
+                400, f"Unable to connect to port {clients.runner_port}"
+            ) from ex
 
 
 def _analytics_table_mapper_temp(table_name: str, blueprint: Blueprint) -> List[str]:

src/brad/ui/models.py

@@ -107,3 +107,13 @@ class SystemState(BaseModel):
     virtual_infra: VirtualInfrastructure
     blueprint: DisplayableBlueprint
     next_blueprint: Optional[DisplayableBlueprint]
+
+
+class ClientState(BaseModel):
+    max_clients: int
+    curr_clients: int
+
+
+class SetClientState(BaseModel):
+    runner_port: Optional[int] = None
+    curr_clients: int

ui/src/App.jsx

@@ -27,12 +27,10 @@ function App() {
       {
         host: "localhost",
         port: 8585,
-        // port: 7583,
       },
       {
         host: "localhost",
         port: 8586,
-        // port: 7583,
       },
     ],
   });
@@ -72,7 +70,7 @@ function App() {
     let timeoutId = null;
     const refreshData = async () => {
       const newSystemState = await fetchSystemState(
-        /*filterTablesForDemo=*/ true,
+        /*filterTablesForDemo=*/ false,
       );
       // TODO: Not the best way to check for equality.
       if (JSON.stringify(systemState) !== JSON.stringify(newSystemState)) {

ui/src/api.js

@@ -20,4 +20,24 @@ async function fetchSystemState(filterTablesForDemo) {
   return result.data;
 }
 
-export { fetchMetrics, fetchSystemState };
+async function fetchWorkloadClients(port) {
+  const args = port != null ? { params: { runner_port: port } } : {};
+  const result = await axios.get(`${API_PREFIX}/clients`, args);
+  return result.data;
+}
+
+async function setWorkloadClients(port, numClients) {
+  const args = { curr_clients: numClients };
+  if (port != null) {
+    args.runner_port = port;
+  }
+  const result = await axios.post(`${API_PREFIX}/clients`, args);
+  return result.data;
+}
+
+export {
+  fetchMetrics,
+  fetchSystemState,
+  fetchWorkloadClients,
+  setWorkloadClients,
+};

ui/src/components/PerfView.jsx

@@ -81,7 +81,7 @@ function PerfView({ virtualInfra }) {
   useEffect(() => {
     let timeoutId = null;
     const refreshData = async () => {
-      const rawMetrics = await fetchMetrics(60, /*useGenerated=*/ true);
+      const rawMetrics = await fetchMetrics(60, /*useGenerated=*/ false);
       const fetchedMetrics = parseMetrics(rawMetrics);
       const metricsManager = getMetricsManager();
       const addedNewMetrics = metricsManager.mergeInMetrics(fetchedMetrics);

ui/src/components/VdbeView.jsx

@@ -49,7 +49,7 @@ function VdbeView({
           max={workloadState.max_clients}
           value={workloadState.curr_clients}
           onChange={updateWorkloadNumClients}
-          debounceMs={2000}
+          debounceMs={800}
         />
       )}
       <DbCylinder color="green" onClick={toggleWorkloadAdjuster}>

ui/src/components/VirtualInfraView.jsx

@@ -1,12 +1,8 @@
-import axios from "axios";
 import Panel from "./Panel";
 import VdbeView from "./VdbeView";
 import "./styles/VirtualInfraView.css";
 import { useEffect, useState, useCallback } from "react";
-
-function baseEndpointFromObj({ host, port }) {
-  return `http://${host}:${port}`;
-}
+import { fetchWorkloadClients, setWorkloadClients } from "../api";
 
 function VirtualInfraView({
   virtualInfra,
@@ -25,11 +21,11 @@ function VirtualInfraView({
       ) {
         return;
       }
-      const baseEndpoint = baseEndpointFromObj(workloadRunners[vdbeIndex]);
-      const result = await axios.post(`${baseEndpoint}/clients`, {
-        curr_clients: numClients,
-      });
-      const newWorkloadState = result.data;
+      const endpoint = workloadRunners[vdbeIndex];
+      const newWorkloadState = await setWorkloadClients(
+        endpoint.port,
+        numClients,
+      );
 
       // Skip the state update if there was no change.
       const existingWorkloadState = workloadStates[vdbeIndex];
@@ -51,11 +47,11 @@ function VirtualInfraView({
 
   useEffect(async () => {
     const { workloadRunners } = endpoints;
-    const promises = workloadRunners
-      .map(baseEndpointFromObj)
-      .map((baseEndpoint) => axios.get(`${baseEndpoint}/clients`));
+    const promises = workloadRunners.map((endpoint) =>
+      fetchWorkloadClients(endpoint.port),
+    );
     const results = await Promise.all(promises);
-    setWorkloadStates(results.map(({ data }) => data));
+    setWorkloadStates(results);
   }, [endpoints]);
 
   return (

workloads/IMDB_extended/workload_utils/change_clients_api.py

@@ -1,20 +1,11 @@
 import uvicorn
 from fastapi import FastAPI
-from pydantic import BaseModel
 from typing import Optional
 
+from brad.ui.models import ClientState, SetClientState
 from .pause_controller import PauseController
 
 
-class ClientState(BaseModel):
-    max_clients: int
-    curr_clients: int
-
-
-class SetClientState(BaseModel):
-    curr_clients: int
-
-
 class Manager:
     def __init__(self, pc: PauseController) -> None:
         self.pc = pc