Pass the ?pool query param through to the backend. (#2246)

Pass the ?pool query param through to the backend.

This will slice off the ?pool= param and pass it to the WebSocket
request, which requests that the Zoo API use a particular pool of
engines. This isn't something any users of the zoo api require; but it's
needed for the internal engine Zoo development workflow. This may be
used in the future, but for now this'll be always enabled. Passing any
value in the production servers will result in a "no backend" error for
now.

src/components/ModelingMachineProvider.tsx

@@ -56,6 +56,7 @@ import toast from 'react-hot-toast'
 import { EditorSelection } from '@uiw/react-codemirror'
 import { CoreDumpManager } from 'lib/coredump'
 import { useHotkeys } from 'react-hotkeys-hook'
+import { useSearchParams } from 'react-router-dom'
 import { letEngineAnimateAndSyncCamAfter } from 'clientSideScene/CameraControls'
 
 type MachineContext<T extends AnyStateMachine> = {
@@ -84,7 +85,12 @@ export const ModelingMachineProvider = ({
   } = useSettingsAuthContext()
   const token = auth?.context?.token
   const streamRef = useRef<HTMLDivElement>(null)
+
+  let [searchParams] = useSearchParams()
+  const pool = searchParams.get('pool')
+
   useSetupEngineManager(streamRef, token, {
+    pool: pool,
     theme: theme.current,
     highlightEdges: highlightEdges.current,
     enableSSAO: enableSSAO.current,

src/hooks/useSetupEngineManager.ts

@@ -9,10 +9,12 @@ export function useSetupEngineManager(
   streamRef: React.RefObject<HTMLDivElement>,
   token?: string,
   settings = {
+    pool: null,
     theme: Themes.System,
     highlightEdges: true,
     enableSSAO: true,
   } as {
+    pool: string | null
     theme: Themes
     highlightEdges: boolean
     enableSSAO: boolean
@@ -35,6 +37,12 @@ export function useSetupEngineManager(
 
   const hasSetNonZeroDimensions = useRef<boolean>(false)
 
+  if (settings.pool) {
+    // override the pool param (?pool=) to request a specific engine instance
+    // from a particular pool.
+    engineCommandManager.pool = settings.pool
+  }
+
   useLayoutEffect(() => {
     // Load the engine command manager once with the initial width and height,
     // then we do not want to reload it.

src/lang/std/engineConnection.ts

@@ -888,6 +888,7 @@ export class EngineCommandManager {
   sceneCommandArtifacts: ArtifactMap = {}
   outSequence = 1
   inSequence = 1
+  pool?: string
   engineConnection?: EngineConnection
   defaultPlanes: DefaultPlanes | null = null
   commandLogs: CommandLog[] = []
@@ -914,8 +915,9 @@ export class EngineCommandManager {
   callbacksEngineStateConnection: ((state: EngineConnectionState) => void)[] =
     []
 
-  constructor() {
+  constructor(pool?: string) {
     this.engineConnection = undefined
+    this.pool = pool
   }
 
   private _camControlsCameraChange = () => {}
@@ -972,7 +974,8 @@ export class EngineCommandManager {
     }
 
     const additionalSettings = settings.enableSSAO ? '&post_effect=ssao' : ''
-    const url = `${VITE_KC_API_WS_MODELING_URL}?video_res_width=${width}&video_res_height=${height}${additionalSettings}`
+    const pool = this.pool == undefined ? '' : `&pool=${this.pool}`
+    const url = `${VITE_KC_API_WS_MODELING_URL}?video_res_width=${width}&video_res_height=${height}${additionalSettings}${pool}`
     this.engineConnection = new EngineConnection({
       engineCommandManager: this,
       url,

src/lib/coredump.ts

@@ -49,6 +49,11 @@ export class CoreDumpManager {
     return APP_VERSION
   }
 
+  // Get the backend pool we've requested.
+  pool(): string {
+    return this.engineCommandManager.pool || ''
+  }
+
   // Get the os information.
   getOsInfo(): Promise<string> {
     if (this.isTauri()) {

src/wasm-lib/kcl/src/coredump/local.rs

@@ -33,6 +33,10 @@ impl CoreDump for CoreDumper {
         Ok(env!("CARGO_PKG_VERSION").to_string())
     }
 
+    fn pool(&self) -> Result<String> {
+        Ok("".to_owned())
+    }
+
     async fn os(&self) -> Result<crate::coredump::OsInfo> {
         Ok(crate::coredump::OsInfo {
             platform: Some(std::env::consts::OS.to_string()),

src/wasm-lib/kcl/src/coredump/mod.rs

@@ -19,6 +19,8 @@ pub trait CoreDump: Clone {
 
     fn version(&self) -> Result<String>;
 
+    fn pool(&self) -> Result<String>;
+
     async fn os(&self) -> Result<OsInfo>;
 
     fn is_tauri(&self) -> Result<bool>;
@@ -71,6 +73,7 @@ pub trait CoreDump: Clone {
             os,
             webrtc_stats,
             github_issue_url: None,
+            pool: self.pool()?,
         };
         app_info.set_github_issue_url(&screenshot_url)?;
 
@@ -103,6 +106,9 @@ pub struct AppInfo {
     /// This gets prepoulated with all the core dump info.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub github_issue_url: Option<String>,
+
+    /// Engine pool the client is connected to.
+    pub pool: String,
 }
 
 impl AppInfo {

src/wasm-lib/kcl/src/coredump/wasm.rs

@@ -16,6 +16,9 @@ extern "C" {
     #[wasm_bindgen(method, js_name = baseApiUrl, catch)]
     fn baseApiUrl(this: &CoreDumpManager) -> Result<String, js_sys::Error>;
 
+    #[wasm_bindgen(method, js_name = pool, catch)]
+    fn pool(this: &CoreDumpManager) -> Result<String, js_sys::Error>;
+
     #[wasm_bindgen(method, js_name = version, catch)]
     fn version(this: &CoreDumpManager) -> Result<String, js_sys::Error>;
 
@@ -66,6 +69,12 @@ impl CoreDump for CoreDumper {
             .map_err(|e| anyhow::anyhow!("Failed to get response from version: {:?}", e))
     }
 
+    fn pool(&self) -> Result<String> {
+        self.manager
+            .pool()
+            .map_err(|e| anyhow::anyhow!("Failed to get response from pool: {:?}", e))
+    }
+
     async fn os(&self) -> Result<crate::coredump::OsInfo> {
         let promise = self
             .manager