Run Orbit query streams in Rails without holding Puma threads

What does this MR do and why?

Rails now runs Orbit query streams itself, behind the orbit_query_in_rails flag. It opens the gRPC stream to Orbit, redacts in the same request, and writes the response. Workhorse only proxies the HTTP request. Orbit query changes then need Ruby only, not Go code or Workhorse review. This MR changes no Workhorse code.

Terms

  • Orbit wait: the time Rails waits for Orbit to run a query, often 1 to 5 seconds.
  • IO thread: a Puma thread marked with env['puma.mark_as_io_bound']. Puma 8 lets IO threads go over max_threads, up to max_io_threads.
  • Redaction: the step where Orbit asks Rails which result rows the user can see.

How it works

Before: Workhorse runs the stream. Rails returns a send-data header. Workhorse opens the stream and calls Rails again for redaction.

After: Rails runs the stream. Workhorse only proxies. Redaction runs in the same request.

1. The problem

The first Ruby version held a Puma thread for the full stream. A 4-second query used one of the 4 threads in a worker, and other requests waited. ADR 008 moved the stream to Workhorse to fix this.

2. Rails runs the stream

QueryRunner#run_stream takes the same parameters that Workhorse got in the send-data header. It uses GRPC::Core::Call#run_batch from one thread, so it needs no writer thread and no queue. A waiting run_batch releases the GVL and uses no CPU.

call.run_batch(OPS::SEND_INITIAL_METADATA => metadata, OPS::SEND_MESSAGE => initial_message.to_proto)

MAX_STREAM_MESSAGES.times do
  batch = call.run_batch(received ? recv_ops : recv_ops.merge(OPS::RECV_INITIAL_METADATA => nil))

3. The wait holds one thread, and production has threads to spare

Over 7 days of production logs, at most 18 Orbit queries were in flight at once, across about 1,100 api pods with 24 threads each. The worst case, peak rate times the 30-second timeout, is 55 in flight: under 0.3% of those threads.

Before the wait, the runner gives its primary database connections back and marks the thread as an IO thread. The mark does nothing until a fleet sets max_io_threads, so it is ready if Orbit traffic grows.

4. Puma IO threads work on Puma 8.0.2

Puma 8.0.2 does not lower its thread count when an extra IO thread exits. After a few requests, marking stops adding threads. The fix restores the count in the exit hook, which Puma calls with the pool mutex held.

def trigger_before_thread_exit_hooks
  @spawned = @processors.size if @spawned > @processors.size

  super
end

The initializer raises on any Puma version other than 8.0.2, so an upgrade must recheck the patch. Upstream fixed the bug in puma#3968, which is not released yet.

5. Redaction runs in the same request and fails closed

redact calls Authz::RedactionService with the request's user. Workhorse had to call Rails again with forwarded credentials and a client IP header, and that second request is gone. Any redaction error returns 502 and sends no rows.

6. Errors match Workhorse and reuse client_error?

  • Orbit's ExecuteQueryError codes map to the same HTTP statuses as Workhorse's gkgErrorToHTTPStatus.
  • A stream that ends with a client-class gRPC status reuses GrpcClient::ConnectionError#client_error? from gitlab!261313. It returns Orbit's text: a 400 for REST, and an isError result for MCP. Other codes return a generic 502.

7. Rails stops Orbit when the client leaves

DisconnectWatcher runs one thread for each process. It waits on the client sockets and cancels the gRPC call at EOF. Orbit stops the query, and Rails logs status 499.

8. Limits match the request

  • The gRPC deadline ends 2 seconds before Rack::Timeout, so Rack::Timeout never interrupts a running gRPC batch.
  • The response skips Rack::ETag, so Rails does not hash large results.

9. One helper serves every query entry point

POST /orbit/query, POST /orbit/query/:name, POST /orbit/agent/commands/query_graph, and the MCP query_graph tool call render_orbit_query!. With the flag off, the helper sends the same Workhorse header as before. GraphQL resolvers can call the same runner through context[:request].env.

Verification

  • End to end: all surfaces with the flag on: REST, named queries, agent commands, MCP, the data explorer, Duo Workflow chat, and GraphQL. TLS and client disconnect behave the same as on Workhorse.
  • Redaction: a matrix across 29 entity types for two users found 0 leaked rows and identical results on both paths.
  • Fail closed: 7 injected faults, and none of them returned denied rows.
  • Parity: response bodies match Workhorse for 17 request types.

Performance

GDK on gitlab-work-vm-3: 2 Puma workers with 4 threads, queue_requests false, as in production. Orbit had a fixed delay to act like a slow query. A probe client called GET /api/v4/version during each run. These loads are far above production, which peaks at about 0.003 Orbit queries in flight per worker.

Load Workhorse Rails, plain threads Rails, IO 8 Rails, IO 64
No delay, 4 clients: queries/s 11.4 14.2 14.2 14.1
No delay, 4 clients: Puma CPU per query 172 ms 138 ms 137 ms 138 ms
2 s delay, 32 clients: queries/s 10.8 3.3 9.4 11.6
2 s delay, 32 clients: probe p50 / p99 143 / 964 ms 7.0 / 9.5 s 1.2 / 2.8 s 96 / 808 ms
4 s delay, 64 clients: queries/s 11.1 1.8 5.1 10.7
4 s delay, 64 clients: probe p50 / p99 112 / 1,481 ms 21.7 / 24.8 s 6.8 / 9.7 s 41 / 1,780 ms

"IO 8" and "IO 64" are max_io_threads values, on stock Puma 8.0.2 with the fix in step 4.

  • Each query uses about 20% less Puma CPU than the Workhorse path, because the redaction callback request is gone.
  • With enough IO threads, Rails keeps Puma as free as Workhorse does. Without the fix in step 4, the same run falls to 2.6 queries/s.
  • max_io_threads must cover the peak number of Orbit queries in flight per worker.

Rollout plan

  1. Enable orbit_query_in_rails by percentage on GitLab.com. Watch Puma saturation and the api apdex, because the Rails request now lasts as long as the Orbit query.
  2. If Orbit traffic grows, set max_io_threads on the api fleet. CNG and Omnibus need a setting that passes it to Puma.
  3. When all traffic uses the Rails path, remove /api/v4/internal/orbit/redaction and workhorse/internal/orbit.

References

Screenshots or screen recordings

Not a UI change.

How to set up and validate locally

  1. Use a GDK with Orbit running.
  2. Enable the flag:
    Feature.enable(:orbit_query_in_rails)
  3. Optional: add max_io_threads 8 to config/puma.rb and restart rails-web.
  4. Run a query and check that the response has rows:
    curl --header "PRIVATE-TOKEN: <token>" --header "Content-Type: application/json" \
      --data '{"query":{"query_type":"traversal","nodes":[{"id":"p","entity":"Project","columns":["name"]}],"limit":5}}' \
      "https://gdk.test:3443/api/v4/orbit/query"
Agent context
  • This MR deletes the old bidi GrpcClient#execute_query. It had no callers, and it failed open on redaction errors.
  • orbit_query_params now resolves DNS for private_address? only for endpoints without TLS. Before, every query did a DNS lookup on GitLab.com.
  • During the wait, the request keeps a leased replica connection if it used one. max_io_threads bounds this per worker. A value of 8 or less stays inside the default pool headroom of 10, so this MR changes no database code.
  • The Puma master gate for queue_requests false (puma#3969) counts IO threads as busy. It matters only for a future Puma upgrade, and the version guard forces that check.
  • Production traffic numbers come from 7 days of Kibana Rails and Workhorse logs, 2026-10-03 to 2026-10-10.
  • A dedicated Orbit fleet and a full Rack hijack with a thread pool were both tested and dropped. The fleet needs new infrastructure. The hijack skips Rails middleware and cannot cover GraphQL.
  • The bench harness and raw results are not in this MR. Ask the author for them.

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Michael Angelo Rivera

Merge request reports

Loading
Loading