perf(relay): reuse host data attachment owner without inventory copies (#20219)

cloud/apps/relay/src/host-session-client-accept.test.ts

@@ -424,6 +424,8 @@ describe('successful client accept timing', () => {
       ) as { connId: string; connTicket: string }
       // The desktop's data leg is the attach window this is meant to expose.
       now += 23
+      const session = h.registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })!
+      const ownerProbe = vi.spyOn(session.pendingConns, 'has')
       const accepted = await h.registry.acceptHostData(
         hostData as unknown as WebSocket,
         connOpen.connId,
@@ -432,6 +434,7 @@ describe('successful client accept timing', () => {
       )
 
       expect(accepted).toBe(true)
+      expect(ownerProbe).toHaveBeenCalledOnce()
       expect(h.observer.recordClientAcceptCompleted).toHaveBeenCalledWith({
         totalMs: 49,
         stageMs: { assignment: 5, credential: 7, activity: 11, attach: 23, basis: 3 }

cloud/apps/relay/src/host-session-owner-scan.test.ts

@@ -0,0 +1,44 @@
+import { describe, expect, it, vi } from 'vitest'
+import type WebSocket from 'ws'
+import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract'
+import { HostSessionRegistry, type HostSession } from './host-session-registry.js'
+
+describe('host data owner lookup', () => {
+  it.each([0, 299, 599, -1])('stops at the first owner (position %s)', async (ownerIndex) => {
+    const recordAuth = vi.fn()
+    const registry = new HostSessionRegistry(
+      ...([{}, vi.fn(), {}, {}, {}, { recordAuth }] as unknown as ConstructorParameters<
+        typeof HostSessionRegistry
+      >)
+    )
+    const sessions = (registry as unknown as { sessions: Map<string, HostSession> }).sessions
+    for (let index = 0; index < 600; index++) {
+      sessions.set(`host-${index}`, {
+        relayHostId: `host-${index}`,
+        pendingConns: new Map(index === ownerIndex ? [['connection', { connTicket: 'secret' }]] : [])
+      } as unknown as HostSession)
+    }
+    let visited = 0
+    const values = sessions.values.bind(sessions)
+    vi.spyOn(sessions, 'values').mockImplementation(() => {
+      const iterator = values()
+      const next = iterator.next.bind(iterator)
+      iterator.next = () => {
+        const result = next()
+        if (!result.done) visited++
+        return result
+      }
+      return iterator
+    })
+    const close = vi.fn()
+    expect(
+      await registry.acceptHostData({ close } as unknown as WebSocket, 'connection', 'wrong', 1)
+    ).toBe(false)
+    expect(close).toHaveBeenCalledWith(
+      RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL,
+      'invalid host data ticket'
+    )
+    expect(recordAuth).toHaveBeenCalledExactlyOnceWith(false)
+    expect(visited).toBe(ownerIndex === -1 ? 600 : ownerIndex + 1)
+  })
+})

cloud/apps/relay/src/host-session-registry.ts

@@ -511,16 +511,20 @@ export class HostSessionRegistry {
     connTicket: string,
     generation: number
   ): Promise<boolean> {
-    const owner = [...this.sessions.values()].find((candidate) =>
-      candidate.pendingConns.has(connId)
-    )
+    let owner: HostSession | undefined
+    for (const candidate of this.sessions.values()) {
+      if (candidate.pendingConns.has(connId)) {
+        owner = candidate
+        break
+      }
+    }
     const release = owner ? this.beginIdleWork(owner.relayHostId) : () => {}
     if (!release) {
       socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'idle cutover in progress')
       return false
     }
     try {
-      return await this.acceptHostDataUnfenced(socket, connId, connTicket, generation)
+      return await this.acceptHostDataUnfenced(socket, connId, connTicket, generation, owner)
     } finally {
       release()
     }
@@ -530,11 +534,9 @@ export class HostSessionRegistry {
     socket: WebSocket,
     connId: string,
     connTicket: string,
-    generation: number
+    generation: number,
+    session: HostSession | undefined
   ): Promise<boolean> {
-    const session = [...this.sessions.values()].find((candidate) =>
-      candidate.pendingConns.has(connId)
-    )
     const pending = session?.pendingConns.get(connId)
     if (
       !session ||