Deliver jobs before stale broker messages, so a self-hosted job cannot stall

src/main/broker-proxy-service.test.ts

@@ -450,3 +450,268 @@ describe('extractGitHubJobInfo', () => {
     expect(info.githubWorkflow).toBe('integration');
   });
 });
+
+describe('message routing', () => {
+  interface Instance { sessionId?: string; runner: { agentName: string } }
+  interface RoutingInternals {
+    targets: Map<string, { target: Target; instances: Map<number, Instance> }>;
+    messageQueues: Map<string, string[]>;
+    pendingTargetAssignments: string[];
+    localSessions: Map<string, { targetId?: string; currentJobId?: string }>;
+    handleRequest(req: unknown, res: unknown): Promise<void>;
+    processMessage(state: unknown, instance: unknown, body: string): Promise<void>;
+  }
+
+  // Broker messages as GitHub delivers them: the inner body is a JSON string.
+  const jobMessage = JSON.stringify({
+    messageId: 2,
+    messageType: 'RunnerJobRequest',
+    body: JSON.stringify({ runner_request_id: 'req-1' }),
+  });
+  const refreshMessage = JSON.stringify({
+    messageId: 1,
+    messageType: 'RunnerRefreshConfig',
+    body: JSON.stringify({ config_type: 'runner' }),
+  });
+  const cancelMessage = JSON.stringify({
+    messageId: 3,
+    messageType: 'JobCancellation',
+    body: JSON.stringify({ jobId: 'req-1' }),
+  });
+
+  const fakeRequest = (method: string, url: string, body = '') => {
+    const req = new EventEmitter() as EventEmitter & {
+      method: string; url: string; headers: Record<string, string>;
+      [Symbol.asyncIterator]: () => AsyncGenerator<Buffer>;
+    };
+    req.method = method;
+    req.url = url;
+    req.headers = {};
+    req[Symbol.asyncIterator] = async function* () { if (body) yield Buffer.from(body); };
+    return req;
+  };
+
+  const fakeResponse = () => {
+    let ended = false;
+    const res = new EventEmitter() as EventEmitter & {
+      statusCode: number; body: string; writableEnded: boolean;
+      writeHead: (code: number) => unknown; end: (chunk?: unknown) => unknown;
+    };
+    Object.defineProperty(res, 'writableEnded', { get: () => ended });
+    res.writeHead = (code) => { res.statusCode = code; return res; };
+    res.end = (chunk) => { res.body = chunk ? String(chunk) : ''; ended = true; res.emit('close'); return res; };
+    return res;
+  };
+
+  let service: BrokerProxyService;
+  let internals: RoutingInternals;
+
+  const request = async (method: string, url: string, body?: string) => {
+    const res = fakeResponse();
+    await internals.handleRequest(fakeRequest(method, url, body), res);
+    return res;
+  };
+
+  const addTargetWithRunner = (id: string, agentName: string) => {
+    const target = createMockTarget({ id, displayName: id });
+    const cred = createMockInstanceCredentials(1);
+    service.addTarget(target, [{ ...cred, runner: { ...cred.runner, agentName } }]);
+    // An upstream session already exists, so /session doesn't try to create one.
+    internals.targets.get(id)!.instances.get(1)!.sessionId = `upstream-${id}`;
+    return target;
+  };
+
+  const createSession = async (body?: string): Promise<string> => {
+    const res = await request('POST', '/session', body);
+    expect(res.statusCode).toBe(201);
+    return JSON.parse(res.body).sessionId;
+  };
+
+  beforeEach(() => {
+    service = new BrokerProxyService(8787);
+    internals = service as unknown as RoutingInternals;
+    // Upstream calls (acknowledge) succeed silently.
+    mockHttpsRequest.mockImplementation((...args: unknown[]) => {
+      const callback = args[1] as (res: EventEmitter) => void;
+      const req = new EventEmitter() as EventEmitter & { setTimeout: () => void; write: () => void; end: () => void };
+      req.setTimeout = () => {};
+      req.write = () => {};
+      req.end = () => {
+        const res = new EventEmitter() as EventEmitter & { statusCode: number };
+        res.statusCode = 200;
+        callback(res);
+        res.emit('end');
+      };
+      return req;
+    });
+  });
+
+  it('delivers the queued job before a stale RunnerRefreshConfig', async () => {
+    // GitHub pushes RunnerRefreshConfig between jobs. If the next worker's first
+    // poll returns that instead of its job, the runner rewrites its config and
+    // restarts its session, and the job is never delivered.
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    internals.messageQueues.set(target.id, [refreshMessage, jobMessage]);
+    internals.pendingTargetAssignments.push(target.id);
+
+    const sessionId = await createSession();
+    const res = await request('GET', `/message?sessionId=${sessionId}`);
+
+    expect(res.statusCode).toBe(200);
+    expect(JSON.parse(res.body).messageType).toBe('RunnerJobRequest');
+  });
+
+  it('drops RunnerRefreshConfig instead of queuing it for the next worker', async () => {
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    const state = internals.targets.get(target.id)!;
+
+    await internals.processMessage(state, state.instances.get(1), refreshMessage);
+
+    expect(internals.messageQueues.get(target.id) ?? []).toEqual([]);
+  });
+
+  it('queues a JobCancellation behind the job it is for', async () => {
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    internals.messageQueues.set(target.id, [jobMessage]);
+    const state = internals.targets.get(target.id)!;
+
+    await internals.processMessage(state, state.instances.get(1), cancelMessage);
+
+    expect(internals.messageQueues.get(target.id)).toEqual([jobMessage, cancelMessage]);
+  });
+
+  it('queues a JobCancellation for a job a worker is running', async () => {
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    internals.messageQueues.set(target.id, [jobMessage]);
+    internals.pendingTargetAssignments.push(target.id);
+    const sessionId = await createSession();
+    await request('GET', `/message?sessionId=${sessionId}`);
+    const state = internals.targets.get(target.id)!;
+
+    await internals.processMessage(state, state.instances.get(1), cancelMessage);
+
+    expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]);
+  });
+
+  it('does not hand a cancellation to a worker that holds no job', async () => {
+    // With no job queued, taking the head of the queue gave a jobless worker a
+    // cancellation meant for whoever runs that job - and marked its session as
+    // holding a job it never had.
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    internals.messageQueues.set(target.id, [cancelMessage]);
+    internals.pendingTargetAssignments.push(target.id);
+    const sessionId = await createSession();
+
+    // Nothing is deliverable, so the handler long-polls rather than answering -
+    // which is the point. Let it poll, assert the cancellation stayed put, then
+    // end the poll the way it really ends, by a job arriving: that both cleans
+    // the request up and shows the cancellation was skipped rather than eaten.
+    const pending = request('GET', `/message?sessionId=${sessionId}`);
+    await new Promise((r) => setTimeout(r, 60));
+
+    expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]);
+    expect(internals.localSessions.get(sessionId)?.currentJobId).toBeUndefined();
+
+    internals.messageQueues.get(target.id)!.push(jobMessage);
+    const res = await pending;
+
+    expect(res.body).toBe(jobMessage);
+    expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]);
+  });
+
+  it('does hand the cancellation to the worker actually running that job', async () => {
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    internals.messageQueues.set(target.id, [jobMessage]);
+    internals.pendingTargetAssignments.push(target.id);
+    const sessionId = await createSession();
+    await request('GET', `/message?sessionId=${sessionId}`); // takes the job
+    internals.messageQueues.set(target.id, [cancelMessage]);
+
+    const res = await request('GET', `/message?sessionId=${sessionId}`);
+
+    expect(res.body).toContain('Cancel');
+    expect(internals.messageQueues.get(target.id)).toEqual([]);
+  });
+
+  it('drops a JobCancellation for a job that is neither queued nor running', async () => {
+    // GitHub redelivers a cancellation for a while after the job ends. Queued
+    // for a target with no worker, it would be handed to the next worker.
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    const state = internals.targets.get(target.id)!;
+
+    await internals.processMessage(state, state.instances.get(1), cancelMessage);
+
+    expect(internals.messageQueues.get(target.id) ?? []).toEqual([]);
+  });
+
+  it('does not let a queued cancellation keep itself alive across redeliveries', async () => {
+    // GitHub redelivers a cancellation while the job is unfinished. A queued
+    // cancellation names the job, so counting it as evidence the job is live
+    // made every redelivery queue another copy - unbounded, and the next
+    // worker to poll gets a cancellation instead of a job.
+    const target = addTargetWithRunner('target-a', 'runner-a.1');
+    const state = internals.targets.get(target.id)!;
+    internals.messageQueues.set(target.id, [cancelMessage]);
+
+    await internals.processMessage(state, state.instances.get(1), cancelMessage);
+
+    expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]);
+  });
+
+  it("answers the runner's own acknowledge locally", async () => {
+    const res = await request('POST', '/acknowledge?sessionId=abc', JSON.stringify({ runnerRequestId: 'req-1' }));
+
+    expect(res.statusCode).toBe(200);
+    expect(res.body).toBe('{}');
+  });
+
+  it('rejects an oversized session request instead of buffering it', async () => {
+    const res = await request('POST', '/session', 'x'.repeat(65 * 1024));
+
+    expect(res.statusCode).toBe(413);
+  });
+
+  describe('session to target binding', () => {
+    it('leaves a session unbound when its target has no job waiting', async () => {
+      // A listener spawned ahead of any job runs in the generic sandbox, not
+      // the repository's approved policy. Binding it would let it win the next
+      // job and fail it (seen live: cargo denied reading ~/.rustup).
+      addTargetWithRunner('target-a', 'runner-a.1');
+      addTargetWithRunner('target-b', 'runner-b.1');
+
+      const sessionId = await createSession(JSON.stringify({ agent: { name: 'runner-b.1' } }));
+
+      expect(internals.localSessions.get(sessionId)?.targetId).toBeUndefined();
+    });
+
+    it("never hands a named session another target's pending assignment", async () => {
+      const targetA = addTargetWithRunner('target-a', 'runner-a.1');
+      addTargetWithRunner('target-b', 'runner-b.1');
+      internals.pendingTargetAssignments.push(targetA.id);
+
+      const sessionId = await createSession(JSON.stringify({ agent: { name: 'runner-b.1' } }));
+
+      expect(internals.localSessions.get(sessionId)?.targetId).toBeUndefined();
+      expect(internals.pendingTargetAssignments).toEqual([targetA.id]);
+    });
+
+    it('consumes the matching pending assignment, not the first one', async () => {
+      const targetA = addTargetWithRunner('target-a', 'runner-a.1');
+      const targetB = addTargetWithRunner('target-b', 'runner-b.1');
+      internals.pendingTargetAssignments.push(targetA.id, targetB.id);
+
+      await createSession(JSON.stringify({ agent: { name: 'runner-b.1' } }));
+
+      expect(internals.pendingTargetAssignments).toEqual([targetA.id]);
+    });
+
+    it('falls back to the pending assignment when the request names no runner', async () => {
+      const targetA = addTargetWithRunner('target-a', 'runner-a.1');
+      internals.pendingTargetAssignments.push(targetA.id);
+
+      const sessionId = await createSession();
+
+      expect(internals.localSessions.get(sessionId)?.targetId).toBe(targetA.id);
+    });
+  });
+});