Stream query results to reduce memory usage

[jemiahw]

Previously, ReadEntities loaded all matching documents into memory before returning them to the API layer, which could cause OOM errors on large result sets.

This commit replaces ReadEntities with two methods:

• StreamEntities: Returns a channel that yields results as they're read from MongoDB, allowing the API to encode and flush each record incrementally
• ReadEntity: Fetches a single entity by ID (no streaming needed)

The API layer now writes JSON array elements one at a time, enabling gzip compression to begin after the first result arrives rather than waiting for the full dataset.

Limitations:

• Errors mid-stream will produce malformed JSON since we've already written partial output. This is rare and fixing it requires breaking API changes (deferred to a future v2 API).
• The "db" instrumentation now measures query setup time, not total query duration.

[Copilot]

Pull request overview

This pull request refactors the entity reading functionality to stream results instead of loading all matching documents into memory. The primary goal is to reduce memory usage for large result sets and enable incremental JSON encoding with compression.

Changes:

• Replaced ReadEntities with StreamEntities (returns a channel for streaming results) and ReadEntity (fetches a single entity by ID)
• Modified API layer to write JSON array elements incrementally using the streaming channel
• Updated all tests and mocks to support the new streaming interface

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
entity/store.go	Adds ReadEntity for single entity lookup and StreamEntities for streaming multiple entities via channel; includes EntityResult type for channel communication
api/api_gomux.go	Updates getEntitiesHandler to stream JSON incrementally, modifies getEntityHandler and getLabelsHandler to use new ReadEntity method
test/mock/entity.go	Updates mock to support new interface with ReadEntityFunc, StreamEntitiesFunc, and helper DoStreamEntities
entity/store_test.go	Adds readStream helper and updates all tests to use new streaming interface; adds tests for ReadEntity
entity/v09_test.go	Updates test to use streaming interface
api/single_entity_read_test.go	Updates tests to mock ReadEntity instead of ReadEntities
api/query_test.go	Updates tests to mock StreamEntities with channel-based interface
api/api_test.go	Updates tests to mock both ReadEntity and StreamEntities for context propagation testing

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zyin-squareup]

While it looks like this will reduce the memory usage, I am hesitant to approve this.
I feel like the multiple calls (possibly 10k+) to DocDB for a single query will result in slower query latency, higher CPU demand for DocDB and Block-etre, as well as (possibly) higher memory pressure on DocDB.

I believe this needs careful load testing and metrics from both components

[qian-squareup]

stamp to unblock. As discussed please

1. Clearly stating it in the release notes/ API docs
2. Guard it with a major version change

[zyin-squareup]

Update:

I dug into the code and made some timing script.
Overall, I see why it is probably not an issue.
The cursor.Next(ctx) was bothering me, but it seems to be returning chunks of results in the underlying.

Based on my tests, it only seems to be 15% slower.

New
Total time for 10 runs: 11.71810 s
Average time per run:       1.17181 s

Old
Total time for 10 runs: 10.31881 s
Average time per run:       1.03188 sZi Ning Yin  [2:55 PM]
2 things that are worth considering.

The bot raises a good point about line 165 blocking forever.
It is worth considering if we want to increase the channel size to 1k or 2k. This might improve the response speed slightly.

[jemiahw]

While it looks like this will reduce the memory usage, I am hesitant to approve this. I feel like the multiple calls (possibly 10k+) to DocDB for a single query will result in slower query latency, higher CPU demand for DocDB and Block-etre, as well as (possibly) higher memory pressure on DocDB.

I believe this needs careful load testing and metrics from both components

If you look at the code for the MongoDB driver, this is the same as calling cursor.All (the previous behavior). Either way the mongo client makes round trips to the server, fetching a batch of records at a time. The batch size is set in our configuration file.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 11 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jemiahw]

It is worth considering if we want to increase the channel size to 1k or 2k. This might improve the response speed slightly.

That's a good idea. I set the channel size to 2*BatchSize, where BatchSize is the number of records fetched from mongo in each round trip (currently configured to 5000). Will see if it helps the speed, but it makes sense it would because now we fetch 5000 records from the DB (slow), write them to the channel (fast), then while the API is encoding those and writing them out, the DB can start it's next fetch.

[jemiahw]

Based on my tests, it only seems to be 15% slower.

I ran a speed test with the updated channel size.

Fetching aurora_parameter _id from staging: Avg 2.721s
Fetching integration_test.payload1 _id from staging: Avg 2.735s

Fetching aurora_parameter _id from playpen: Avg 2.639s
Fetching integration_test.payload1 _id from playpen: Avg 2.086s

The new version might be slightly faster than the existing code now, but it's probably within the margin of error.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.