Stream Orbit query results through Rails in chunks

What does this MR do and why?

Orbit query results now stream through Rails in chunks of 512 KiB. Before, Rails kept the whole result in the Ruby heap several times. A 6.2 MB result cost 43 MB of malloc in one request, and results over 8 MiB failed with 502.

The status, content type, Cache-Control, and body stay the same. Only the transfer changes from Content-Length to chunked.

Terms:

  • Buffered path: Orbit sends one result message, and Rails returns one string. Old Orbit servers still use this path.
  • Stream path: Rails asks Orbit for chunks. A body object writes each chunk as it arrives. When Puma closes the body, it releases the gRPC call and records the metric.

This MR bumps gitlab-orbit-proto to 0.147.0, the first release with the Orbit chunk messages.

References

  • Orbit server MR, which adds the chunk messages.
  • Feature issue for running Orbit queries in Rails.
  • The Rails query path, which this MR follows up.
  • Production finding during the 75% rollout: one agent burst allocated up to 140 MB in each Rails request.

Screenshots or screen recordings

End to end on a GDK VM. The stack is Workhorse, Puma (2 workers, 4 threads), and a real Orbit server. The query returns 40 source files with content, a 6.2 MB result. Each row is the median of 3 alternating runs. CPU is Puma worker CPU for each successful request.

Clients CPU per request, before CPU per request, after Peak RSS growth, before Peak RSS growth, after
1 312 ms 318 ms 184 MB 150 MB
2 259 ms 256 ms 147 MB 130 MB
4 306 ms 323 ms 161 MB 86 MB
8 352 ms 455 ms 115 MB 69 MB
  • Memory is lower at every level, by up to 47%.
  • CPU does not change at 1 to 4 clients. In GDK development mode, most of the CPU goes to work that is not result handling.
  • At 8 clients, Orbit and Gitaly time out on both paths. Only 119 and 110 of 144 requests succeeded, so that CPU row is not a fair comparison.

Isolated response path. This test runs Puma 8.0.2, Rack 2.2.23, and Grape 3.2.1 with the same protobuf and JSON gems, and no Rails. It sends a 7.5 MB result 10 times per client.

Clients CPU, before CPU, after Peak memory, before Peak memory, after
1 0.47 s 0.30 s 175 MB 100 MB
8 1.51 s 0.71 s 410 MB 148 MB

Results over 8 MiB. An 80-file query with an 11 MB result returned 502 before. It returns 200 now.

Same response. llm results are byte-identical between the two paths. Raw results have the same length and the same JSON. Orbit's raw key order changes from run to run on either path.

How to set up and validate locally

  1. Run an Orbit server from main (v0.147.0 or later).

  2. Enable the Rails path for your user:

    Feature.enable(:orbit_query_in_rails, User.find_by_username('root'))
  3. Send a large query and check the headers:

    curl -sk -D - -o /tmp/result.json -H "PRIVATE-TOKEN: $TOKEN" -H 'Content-Type: application/json' \
      -X POST https://gdk.test:3443/api/v4/orbit/query \
      -d '{"query":"MATCH (f:File) WHERE f.size_bytes > 40000 RETURN f.path, f.content LIMIT 40","response_format":"raw"}'

    You see transfer-encoding: chunked, and the file is valid JSON.

  4. Run the same query against an Orbit server older than v0.147.0. You see content-length and the same JSON.

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.

Implementation notes
  • Same bytes: the head and tail come from the existing renderers. Rails renders a sentinel string and splits on it. For MCP, each chunk is JSON-escaped on its own. Chunks split on UTF-8 boundaries, so the joined output equals one escaped string.
  • Headers: Grape stream removes Content-Length and sets Cache-Control: no-cache. Rails' Rack::ETag, "no-cache" already sends that value for these responses, so only the transfer header changes.
  • Failure after the first byte: the body logs the cause and raises an IOError. Puma 8 turns that into a silent socket close (response.rb rescues IOError and raises ConnectionError, which client_error ignores). The client sees a cut-off response, not a 500 written into the body and not a complete 200. Workhorse passes the abort on with http.ErrAbortHandler.
  • Limits: the Orbit server owns the total size limit (grpc.max_result_bytes, 64 MiB by default). It rejects a larger result before the first chunk, and Rails returns that error as a 400. The gRPC deadline now also covers the transfer.
  • Observability: the body runs after app.call, so api_json duration_s and mem_bytes stop before the stream. The gRPC duration metric covers the full stream, and stream failures log with the request's Labkit context.
  • Review: three adversarial review passes. Findings fixed: error details on stream failure, Labkit context, the byte cap, a 500 written mid-body, the Puma constant dependency, and spec gaps.
  • Allocator: the isolated test ran on macOS malloc. The GDK test ran on Linux.
Edited by Michael Angelo Rivera

Merge request reports

Loading
Loading