From d2e1de2343ddf9ce2abce7323eb072dc205fce14 Mon Sep 17 00:00:00 2001 From: Jemiah Westerman Date: Tue, 24 Feb 2026 19:57:21 -0800 Subject: [PATCH 1/3] Stream query results to reduce memory usage 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. --- api/api_gomux.go | 90 +++++++++++++++------ api/api_test.go | 12 ++- api/query_test.go | 24 +++--- api/single_entity_read_test.go | 36 ++++----- entity/store.go | 143 ++++++++++++++++++++++++--------- entity/store_test.go | 90 ++++++++++++++++++--- entity/v09_test.go | 2 +- test/mock/entity.go | 31 ++++++- 8 files changed, 313 insertions(+), 115 deletions(-) diff --git a/api/api_gomux.go b/api/api_gomux.go index 584b377..4b4ed7e 100644 --- a/api/api_gomux.go +++ b/api/api_gomux.go @@ -8,6 +8,7 @@ import ( "context" "encoding/json" "fmt" + "io" "log" "math/rand" "net/http" @@ -576,26 +577,71 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) { // Query data store (instrumented) rc.inst.Start("db") - entities, err := api.es.ReadEntities(ctx, rc.entityType, q, f) + entities := api.es.StreamEntities(ctx, rc.entityType, q, f) rc.inst.Stop("db") - if err != nil { - api.readError(rc, w, err) - return - } - rc.gm.Val(metrics.ReadMatch, int64(len(entities))) - // Success: return matching entities (possibly empty list) rc.inst.Start("encode-response") - if len(entities) > 0 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { - // use gzip compression if we have data to send back and the client accepts gzip - w.Header().Set("Content-Encoding", "gzip") - gzw := gzip.NewWriter(w) - defer gzw.Close() - json.NewEncoder(gzw).Encode(entities) + first := true + + // Stream the results + count := 0 + var finalWriter io.Writer + finalWriter = w + var encoder *json.Encoder + for e := range entities { + if e.Err != nil { + // Handle errors returned by the database + api.readError(rc, w, e.Err) + return + } + + if err := ctx.Err(); err != nil { + // Handle context timeouts + // api.readError will mangle the response, but we can't just return partial results so let it blow up + api.readError(rc, w, err) + return + } + + count++ + + // Enable compression if client accepts gzip and this is the first entity (we want to avoid compressing an empty response) + if first && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { + w.Header().Set("Content-Encoding", "gzip") + gzw := gzip.NewWriter(w) + defer gzw.Close() + finalWriter = gzw + } + + // Start the JSON array if this is the first record, or put a comma separator if not the first record + if first { + finalWriter.Write([]byte("[")) + encoder = json.NewEncoder(finalWriter) + first = false + } else { + finalWriter.Write([]byte(",")) + } + + // Write the record and handle the error. + // encoder.Encode will add a newline, which is fine (nice for human readable output) + err = encoder.Encode(e.Entity) + if err != nil { + // api.readError will mangle the response, but if the encoder failed then the response is already mangled and there's not much else we can do. + api.readError(rc, w, err) + log.Println("ERROR: Read error while encoding response: ", err) + return + } + } + + // Clean up the JSON array + if first { + // Never saw any data. Write an empty array. + finalWriter.Write([]byte("[]")) } else { - // no compression - json.NewEncoder(w).Encode(entities) + // We wrote some data, now we need to close the array + finalWriter.Write([]byte("]")) // end of JSON array } + + rc.gm.Val(metrics.ReadMatch, int64(count)) rc.inst.Stop("encode-response") } @@ -804,18 +850,17 @@ func (api *API) getEntityHandler(w http.ResponseWriter, r *http.Request) { } // Read the entity by ID - q, _ := query.Translate("_id=" + rc.entityId) - entities, err := api.es.ReadEntities(ctx, rc.entityType, q, f) + entity, err := api.es.ReadEntity(ctx, rc.entityType, rc.entityId, f) if err != nil { api.readError(rc, w, err) return } - if len(entities) == 0 { + if entity == nil { w.WriteHeader(http.StatusNotFound) return } - json.NewEncoder(w).Encode(entities[0]) + json.NewEncoder(w).Encode(entity) } // getLabelsHandler godoc @@ -837,18 +882,17 @@ func (api *API) getLabelsHandler(w http.ResponseWriter, r *http.Request) { rc.gm.Inc(metrics.ReadLabels, 1) // specific read type - q, _ := query.Translate("_id=" + rc.entityId) - entities, err := api.es.ReadEntities(ctx, rc.entityType, q, etre.QueryFilter{}) + entity, err := api.es.ReadEntity(ctx, rc.entityType, rc.entityId, etre.QueryFilter{}) if err != nil { api.readError(rc, w, err) return } - if len(entities) == 0 { + if entity == nil { w.WriteHeader(http.StatusNotFound) return } - json.NewEncoder(w).Encode(entities[0].Labels()) + json.NewEncoder(w).Encode(entity.Labels()) } // -------------------------------------------------------------------------- diff --git a/api/api_test.go b/api/api_test.go index 0b1dbd1..0cef138 100644 --- a/api/api_test.go +++ b/api/api_test.go @@ -198,9 +198,9 @@ func TestClientQueryTimeout(t *testing.T) { // and plumbed all the way down to the entity.Store context var gotCtx context.Context store := mock.EntityStore{} - store.ReadEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + store.ReadEntityFunc = func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { gotCtx = ctx - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil } server := setup(t, defaultConfig, store) defer server.ts.Close() @@ -242,9 +242,13 @@ func TestContextPropagation(t *testing.T) { var gotCtx context.Context store := mock.EntityStore{} // We're going to test all operations, so we need to set all of these funcs - store.ReadEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + store.ReadEntityFunc = func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { gotCtx = ctx - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil + } + store.StreamEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + gotCtx = ctx + return mock.DoStreamEntities(testEntitiesWithObjectIDs[0:1], nil) } store.CreateEntitiesFunc = func(ctx context.Context, op entity.WriteOp, entities []etre.Entity) ([]string, error) { gotCtx = ctx diff --git a/api/query_test.go b/api/query_test.go index c52f50e..fa9a87b 100644 --- a/api/query_test.go +++ b/api/query_test.go @@ -28,10 +28,10 @@ func TestQueryBasic(t *testing.T) { var gotQuery query.Query var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { gotQuery = q gotFilter = f - return testEntitiesWithObjectIDs, nil + return mock.DoStreamEntities(testEntitiesWithObjectIDs, nil) }, } server := setup(t, defaultConfig, store) @@ -171,9 +171,9 @@ func TestQueryNoMatches(t *testing.T) { // is still 200 OK in this case because there's no error. var gotQuery query.Query store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { gotQuery = q - return []etre.Entity{}, nil // no matching queries + return mock.DoStreamEntities(nil, nil) }, } server := setup(t, defaultConfig, store) @@ -220,8 +220,8 @@ func TestQueryErrorsDatabaseError(t *testing.T) { // Test that GET /entities/:type?query=Q handles a database error correctly. // Db errors (and only db errors return HTTP 503 "Service Unavailable". store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - return nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"} + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + return mock.DoStreamEntities(nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"}) }, } server := setup(t, defaultConfig, store) @@ -269,8 +269,8 @@ func TestQueryErrorsNoEntityType(t *testing.T) { // You can run "../test/coverage -test.run TestQueryErrorsNoEntityType" and // see that the handler is never called. store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - return nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"} + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + return mock.DoStreamEntities(nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"}) }, } server := setup(t, defaultConfig, store) @@ -310,11 +310,11 @@ func TestQueryErrorsTimeout(t *testing.T) { // Test that GET /entities/:type?query=Q handles a database timeout correctly. // Db errors (and only db errors return HTTP 503 "Service Unavailable". store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() <-ctx.Done() - return nil, entity.DbError{Err: ctx.Err(), Type: "db-read"} + return mock.DoStreamEntities(nil, entity.DbError{Err: ctx.Err(), Type: "db-read"}) }, } server := setup(t, defaultConfig, store) @@ -357,8 +357,8 @@ func TestQueryErrorsTimeout(t *testing.T) { func TestResponseCompression(t *testing.T) { // Stand up the server store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - return testEntitiesWithObjectIDs, nil + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + return mock.DoStreamEntities(testEntitiesWithObjectIDs, nil) }, } server := setup(t, defaultConfig, store) diff --git a/api/single_entity_read_test.go b/api/single_entity_read_test.go index 987c192..a233d9e 100644 --- a/api/single_entity_read_test.go +++ b/api/single_entity_read_test.go @@ -17,7 +17,6 @@ import ( "github.com/square/etre/auth" "github.com/square/etre/entity" "github.com/square/etre/metrics" - "github.com/square/etre/query" "github.com/square/etre/test" "github.com/square/etre/test/mock" ) @@ -27,15 +26,14 @@ import ( // ////////////////////////////////////////////////////////////////////////// func TestGetEntityBasic(t *testing.T) { - // Test the most basic GET /entity/:type/:id gets the entity. This is - // really a wrapper to call ReadEntitiesFunc() with _id=:id. - var gotQuery query.Query + // Test the most basic GET /entity/:type/:id gets the entity. + var gotEntityId string var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - gotQuery = q + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + gotEntityId = entityId gotFilter = f - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil }, } server := setup(t, defaultConfig, store) @@ -49,8 +47,7 @@ func TestGetEntityBasic(t *testing.T) { assert.Equal(t, http.StatusOK, statusCode, "response status = %d, expected %d", statusCode, http.StatusOK) // GET /entity/:type/:id = "_id=:id" - expectQuery, _ := query.Translate("_id=" + testEntityIds[0]) - assert.Equal(t, expectQuery, gotQuery) + assert.Equal(t, testEntityIds[0], gotEntityId) // No filter options provided in URL expectFilter := etre.QueryFilter{} @@ -84,9 +81,9 @@ func TestGetEntityReturnLabels(t *testing.T) { // that the URL param "labels=" is processed and passed along to the entity.Store. var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { gotFilter = f - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil }, } server := setup(t, defaultConfig, store) @@ -130,8 +127,8 @@ func TestGetEntityNotFound(t *testing.T) { // We simulate this by making ReadEntities() below return an empty list which // the real entity.Store() does when no entity exists with the given _id. store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - return []etre.Entity{}, nil + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + return nil, nil }, } server := setup(t, defaultConfig, store) @@ -164,7 +161,7 @@ func TestGetEntityErrors(t *testing.T) { read := false var dbError error store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { read = true return nil, dbError }, @@ -255,13 +252,13 @@ func TestGetEntityErrors(t *testing.T) { func TestGetEntityLabels(t *testing.T) { // Test that GET /entity/:type/:id/labels works - var gotQuery query.Query + var gotEntityId string var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - gotQuery = q + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + gotEntityId = entityId gotFilter = f - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil }, } server := setup(t, defaultConfig, store) @@ -275,8 +272,7 @@ func TestGetEntityLabels(t *testing.T) { assert.Equal(t, http.StatusOK, statusCode) // GET /entity/:type/:id = "_id=:id" - expectQuery, _ := query.Translate("_id=" + testEntityIds[0]) - assert.Equal(t, expectQuery, gotQuery) + assert.Equal(t, testEntityIds[0], gotEntityId) // No filter options provided in URL expectFilter := etre.QueryFilter{} diff --git a/entity/store.go b/entity/store.go index b4bace7..41cd77d 100644 --- a/entity/store.go +++ b/entity/store.go @@ -17,9 +17,13 @@ import ( "github.com/square/etre/query" ) +const ( + StreamChannelBufferSize = 100 +) + // Store interface has methods needed to do CRUD operations on entities. type Store interface { - ReadEntities(context.Context, string, query.Query, etre.QueryFilter) ([]etre.Entity, error) + ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) CreateEntities(context.Context, WriteOp, []etre.Entity) ([]string, error) @@ -28,6 +32,8 @@ type Store interface { DeleteEntities(context.Context, WriteOp, query.Query) ([]etre.Entity, error) DeleteLabel(context.Context, WriteOp, string) (etre.Entity, error) + + StreamEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan EntityResult } type store struct { @@ -45,40 +51,13 @@ func NewStore(entities map[string]*mongo.Collection, cdcStore cdc.Store, cfg con } } -// ReadEntities queries the db and returns a slice of Entity objects if -// something is found, a nil slice if nothing is found, and an error if one -// occurs. -func (s store) ReadEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { +// ReadEntity queries the db for a single entity. Returns nil if not found. +func (s store) ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { c, ok := s.coll[entityType] if !ok { panic("invalid entity type passed to ReadEntities: " + entityType) } - // Distinct optimizaiton: unique values for the one return label. For example, - // "es -u node.metacluster zone=pd" returns a list of unique metacluster names. - // This is 10x faster than "es node.metacluster zone=pd | sort -u". - if len(f.ReturnLabels) == 1 && f.Distinct { - dr := c.Distinct(ctx, f.ReturnLabels[0], Filter(q)) - if err := dr.Err(); err != nil { - nfe := mongo.ErrNoDocuments - if errors.Is(err, nfe) { - // No documents found, return empty slice - return []etre.Entity{}, nil - } - return nil, s.dbError(ctx, err, "db-read-distinct") - } - var values []string - err := dr.Decode(&values) - if err != nil { - return nil, s.dbError(ctx, err, "db-read-distinct") - } - entities := make([]etre.Entity, len(values)) - for i, v := range values { - entities[i] = etre.Entity{f.ReturnLabels[0]: v} - } - return entities, nil - } - // Find and return all matching entities p := bson.M{} if len(f.ReturnLabels) > 0 { @@ -93,17 +72,105 @@ func (s store) ReadEntities(ctx context.Context, entityType string, q query.Quer } } - // Set batch size and projection - opts := options.Find().SetProjection(p).SetBatchSize(int32(s.config.BatchSize)) - cursor, err := c.Find(ctx, Filter(q), opts) - if err != nil { - return nil, s.dbError(ctx, err, "db-query") + // Query, we should only get one row + q, _ := query.Translate("_id=" + entityId) + result := c.FindOne(ctx, Filter(q), options.FindOne().SetProjection(p)) + if err := result.Err(); err != nil { + nfe := mongo.ErrNoDocuments + if errors.Is(err, nfe) { + return nil, nil + } + return nil, s.dbError(ctx, result.Err(), "db-query") } - entities := []etre.Entity{} - if err := cursor.All(ctx, &entities); err != nil { + entity := etre.Entity{} + if err := result.Decode(&entity); err != nil { return nil, s.dbError(ctx, err, "db-read-cursor") } - return entities, nil + return entity, nil +} + +type EntityResult struct { + Entity etre.Entity + Err error +} + +func (s store) StreamEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan EntityResult { + c, ok := s.coll[entityType] + if !ok { + panic("invalid entity type passed to StreamEntities: " + entityType) + } + + ch := make(chan EntityResult, StreamChannelBufferSize) + go func() { + defer close(ch) + + // Distinct optimization: unique values for the one return label. For example, + // "es -u node.metacluster zone=pd" returns a list of unique metacluster names. + // This is 10x faster than "es node.metacluster zone=pd | sort -u". + if len(f.ReturnLabels) == 1 && f.Distinct { + dr := c.Distinct(ctx, f.ReturnLabels[0], Filter(q)) + if err := dr.Err(); err != nil { + nfe := mongo.ErrNoDocuments + if errors.Is(err, nfe) { + // No documents found, return to close the channel + return + } + ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-distinct")} + return + } + + var values []string + err := dr.Decode(&values) + if err != nil { + ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-distinct")} + return + } + // Distinct doesn't return a cursor, so we just have to loop and send the results + for _, v := range values { + ch <- EntityResult{Entity: etre.Entity{f.ReturnLabels[0]: v}} + } + return + } + + // Find and return all matching entities + p := bson.M{} + if len(f.ReturnLabels) > 0 { + for _, label := range f.ReturnLabels { + p[label] = 1 + } + // Only include _id if explicitly set in f.ReturnLabels. If not, + // ok is false and we must explicitly exlude it because MongoDB + // returns it by default. + if _, ok := p["_id"]; !ok { + p["_id"] = 0 + } + } + + // Run the query + opts := options.Find().SetProjection(p).SetBatchSize(int32(s.config.BatchSize)) + cursor, err := c.Find(ctx, Filter(q), opts) + if err != nil { + ch <- EntityResult{Err: s.dbError(ctx, err, "db-query")} + return + } + defer cursor.Close(ctx) + + // Stream results + for cursor.Next(ctx) { + var entity etre.Entity + if err := cursor.Decode(&entity); err != nil { + ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-cursor")} + return + } + ch <- EntityResult{Entity: entity} + } + // Check for errors from iterating over cursor + if err := cursor.Err(); err != nil { + ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-cursor")} + return + } + }() + return ch } // CreateEntities inserts many entities into DB. This method allows for partial diff --git a/entity/store_test.go b/entity/store_test.go index 84389a8..9fa7be5 100644 --- a/entity/store_test.go +++ b/entity/store_test.go @@ -94,7 +94,7 @@ func docs(entities []etre.Entity) []interface{} { // Read // -------------------------------------------------------------------------- -func TestReadEntitiesWithAllOperators(t *testing.T) { +func TestStreamEntitiesWithAllOperators(t *testing.T) { // Test all operators: in, notin, =, ==, !=, (has) y, (does not have) !foo, <, > // All values are set such that it only matches the first test node to make // testing easier and ensure we don't match the other entities. @@ -115,7 +115,7 @@ func TestReadEntitiesWithAllOperators(t *testing.T) { q, err := query.Translate(qs) require.NoError(t, err) - actual, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + actual, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) assert.Equal(t, expect, actual) } @@ -126,7 +126,7 @@ type readTest struct { expect []etre.Entity } -func TestReadEntitiesMatching(t *testing.T) { +func TestStreamEntitiesMatching(t *testing.T) { // Test various combinations of queries to ensure that we match and return // the correct entities. This is the fundamental job of Etre, so it should // be very thoroughly tested. @@ -157,13 +157,40 @@ func TestReadEntitiesMatching(t *testing.T) { q, err := query.Translate(rt.query) require.NoError(t, err) - got, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) assert.Equal(t, rt.expect, got) } } -func TestReadEntitiesFilterDistinct(t *testing.T) { +func TestReadEntityMatching(t *testing.T) { + // Test various combinations of queries to ensure that we match and return + // the correct entities. This is the fundamental job of Etre, so it should + // be very thoroughly tested. + store := setup(t, &mock.CDCStore{}) + readTests := []struct { + id string + expect etre.Entity + }{ + { + // found + id: testNodes[1]["_id"].(bson.ObjectID).Hex(), + expect: testNodes[1], + }, + { + // not found + id: "bogus", + expect: nil, + }, + } + for _, rt := range readTests { + got, err := store.ReadEntity(context.Background(), entityType, rt.id, etre.QueryFilter{}) + require.NoError(t, err) + assert.Equal(t, rt.expect, got) + } +} + +func TestStreamEntitiesFilterDistinct(t *testing.T) { // Test that etre.QueryFilter{Distinct: true} returns a list of unique values // for one label. The 1st test node has y=a and the 2nd and 3rd both have y=b, // so the unique values are [a,b]. @@ -175,7 +202,7 @@ func TestReadEntitiesFilterDistinct(t *testing.T) { ReturnLabels: []string{"y"}, // only works with 1 return label Distinct: true, } - got, err := store.ReadEntities(context.Background(), entityType, q, f) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, f)) require.NoError(t, err) expect := []etre.Entity{ @@ -185,7 +212,7 @@ func TestReadEntitiesFilterDistinct(t *testing.T) { assert.Equal(t, expect, got) } -func TestReadEntitiesFilterDistinctNoResult(t *testing.T) { +func TestStreamEntitiesFilterDistinctNoResult(t *testing.T) { // Test that etre.QueryFilter{Distinct: true} returns an empty list // if no rows match the query. store := setup(t, &mock.CDCStore{}) @@ -196,12 +223,12 @@ func TestReadEntitiesFilterDistinctNoResult(t *testing.T) { ReturnLabels: []string{"BOGUS"}, // only works with 1 return label Distinct: true, } - got, err := store.ReadEntities(context.Background(), entityType, q, f) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, f)) require.NoError(t, err) assert.Empty(t, got) } -func TestReadEntitiesFilterReturnLabels(t *testing.T) { +func TestStreamEntitiesFilterReturnLabels(t *testing.T) { // Test that etre.QueryFilter{ReturnLabels: []string{x}} returns only that // label and not the others (y, z, bar, foo). We'll select/match by label y // but return only label x. @@ -217,17 +244,33 @@ func TestReadEntitiesFilterReturnLabels(t *testing.T) { {"x": int64(4)}, {"x": int64(6)}, } - got, err := store.ReadEntities(context.Background(), entityType, q, f) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, f)) + require.NoError(t, err) + assert.Equal(t, expect, got) +} + +func TestReadEntityFilterReturnLabels(t *testing.T) { + // Test that etre.QueryFilter{ReturnLabels: []string{x}} returns only that + // label and not the others (y, z, bar, foo). We'll select/match by label y + // but return only label x. + store := setup(t, &mock.CDCStore{}) + + f := etre.QueryFilter{ + ReturnLabels: []string{"x"}, // testing this + } + expect := etre.Entity{"x": int64(2)} + + got, err := store.ReadEntity(context.Background(), entityType, testNodes[0]["_id"].(bson.ObjectID).Hex(), f) require.NoError(t, err) assert.Equal(t, expect, got) } -func TestReadEntitiesFilterReturnMetalabels(t *testing.T) { +func TestStreamEntitiesFilterReturnMetalabels(t *testing.T) { store := setup(t, &mock.CDCStore{}) q, err := query.Translate("y=a") require.NoError(t, err) - actual, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{ReturnLabels: []string{"_id", "_type", "_rev", "y", "_created", "_updated"}}) + actual, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{ReturnLabels: []string{"_id", "_type", "_rev", "y", "_created", "_updated"}})) require.NoError(t, err) expect := []etre.Entity{ @@ -237,6 +280,16 @@ func TestReadEntitiesFilterReturnMetalabels(t *testing.T) { assert.Equal(t, expect, actual) } +func TestReadEntityFilterReturnMetalabels(t *testing.T) { + store := setup(t, &mock.CDCStore{}) + + actual, err := store.ReadEntity(context.Background(), entityType, testNodes[0]["_id"].(bson.ObjectID).Hex(), etre.QueryFilter{ReturnLabels: []string{"_id", "_type", "_rev", "y", "_created", "_updated"}}) + require.NoError(t, err) + + expect := etre.Entity{"_id": testNodes[0]["_id"], "_type": entityType, "_rev": int64(0), "y": "a", "_created": testNodes[0]["_created"], "_updated": testNodes[0]["_updated"]} + assert.Equal(t, expect, actual) +} + // -------------------------------------------------------------------------- // Create // -------------------------------------------------------------------------- @@ -745,7 +798,7 @@ func TestDeleteLabel(t *testing.T) { // The foo label should no longer be set on the entity q, _ := query.Translate("y=a") - gotNew, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + gotNew, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) e := etre.Entity{} @@ -780,3 +833,14 @@ func TestDeleteLabel(t *testing.T) { } assert.Equal(t, expectEvent, gotEvents) } + +func readStream(ch <-chan entity.EntityResult) ([]etre.Entity, error) { + entities := []etre.Entity{} + for r := range ch { + if r.Err != nil { + return nil, r.Err + } + entities = append(entities, r.Entity) + } + return entities, nil +} diff --git a/entity/v09_test.go b/entity/v09_test.go index 1f8a9e6..8a11ed9 100644 --- a/entity/v09_test.go +++ b/entity/v09_test.go @@ -293,7 +293,7 @@ func TestV09DeleteLabel(t *testing.T) { // The foo label should no longer be set on the entity q, _ := query.Translate("x=a") - gotNew, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + gotNew, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) e := etre.Entity{} diff --git a/test/mock/entity.go b/test/mock/entity.go index edf39bd..43fa3a1 100644 --- a/test/mock/entity.go +++ b/test/mock/entity.go @@ -11,12 +11,13 @@ import ( ) type EntityStore struct { - ReadEntitiesFunc func(context.Context, string, query.Query, etre.QueryFilter) ([]etre.Entity, error) + ReadEntityFunc func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) DeleteEntityLabelFunc func(context.Context, entity.WriteOp, string) (etre.Entity, error) CreateEntitiesFunc func(context.Context, entity.WriteOp, []etre.Entity) ([]string, error) UpdateEntitiesFunc func(context.Context, entity.WriteOp, query.Query, etre.Entity) ([]etre.Entity, error) DeleteEntitiesFunc func(context.Context, entity.WriteOp, query.Query) ([]etre.Entity, error) DeleteLabelFunc func(context.Context, entity.WriteOp, string) (etre.Entity, error) + StreamEntitiesFunc func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult } func (s EntityStore) DeleteEntityLabel(ctx context.Context, wo entity.WriteOp, label string) (etre.Entity, error) { @@ -33,9 +34,9 @@ func (s EntityStore) CreateEntities(ctx context.Context, wo entity.WriteOp, enti return nil, nil } -func (s EntityStore) ReadEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - if s.ReadEntitiesFunc != nil { - return s.ReadEntitiesFunc(ctx, entityType, q, f) +func (s EntityStore) ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + if s.ReadEntityFunc != nil { + return s.ReadEntityFunc(ctx, entityType, entityId, f) } return nil, nil } @@ -60,3 +61,25 @@ func (s EntityStore) DeleteLabel(ctx context.Context, wo entity.WriteOp, label s } return etre.Entity{}, nil } + +func (s EntityStore) StreamEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + if s.StreamEntitiesFunc != nil { + return s.StreamEntitiesFunc(ctx, entityType, q, f) + } + return DoStreamEntities(nil, nil) +} + +func DoStreamEntities(entities []etre.Entity, err error) <-chan entity.EntityResult { + ch := make(chan entity.EntityResult) + go func() { + defer close(ch) + if err != nil { + ch <- entity.EntityResult{Err: err} + return + } + for _, e := range entities { + ch <- entity.EntityResult{Entity: e} + } + }() + return ch +} From ce3c8cac4d25db8c9cca9a4d9b58da6752881a8f Mon Sep 17 00:00:00 2001 From: Jemiah Westerman Date: Wed, 25 Feb 2026 09:58:21 -0800 Subject: [PATCH 2/3] Flush the gzip writer periodically to force chunking --- api/api_gomux.go | 48 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/api/api_gomux.go b/api/api_gomux.go index 4b4ed7e..fe631e3 100644 --- a/api/api_gomux.go +++ b/api/api_gomux.go @@ -581,42 +581,47 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) { rc.inst.Stop("db") rc.inst.Start("encode-response") - first := true - // Stream the results - count := 0 + // FinalWriter is where we will write the data. If we're using gzip compression, this will be the gzip writer. If not, it will just be the original response writer. var finalWriter io.Writer finalWriter = w + + // gzip writer, only initialized if client accepts gzip encoding. + var gzw *gzip.Writer + + // encoder is the JSON encoder writing to finalWriter. We initialize it after we know whether we're using gzip or not. var encoder *json.Encoder + + // Number of non-error records sent to the client + count := 0 for e := range entities { + + // Handle errors returned by the database if e.Err != nil { - // Handle errors returned by the database api.readError(rc, w, e.Err) return } - + // Handle context timeouts if err := ctx.Err(); err != nil { - // Handle context timeouts - // api.readError will mangle the response, but we can't just return partial results so let it blow up api.readError(rc, w, err) return } - count++ - - // Enable compression if client accepts gzip and this is the first entity (we want to avoid compressing an empty response) - if first && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { + // Initialize gzip writer and JSON encoder on the first record, after we know there is data to return to the client. + // We also check the client's Accept-Encoding header to see if gzip is supported. + if count == 0 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { w.Header().Set("Content-Encoding", "gzip") - gzw := gzip.NewWriter(w) + gzw = gzip.NewWriter(w) defer gzw.Close() finalWriter = gzw } // Start the JSON array if this is the first record, or put a comma separator if not the first record - if first { + // We don't do this before the for loop because we don't want to write the opening bracket if there is a database + // error in the first record returned. + if count == 0 { finalWriter.Write([]byte("[")) encoder = json.NewEncoder(finalWriter) - first = false } else { finalWriter.Write([]byte(",")) } @@ -630,10 +635,19 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) { log.Println("ERROR: Read error while encoding response: ", err) return } + + // Count the record + count++ + + // If we're compressing, flush every 100 records to ensure data is sent to the client in a timely manner and not buffered in memory. + // Otherwise, the gzip flusher may buffer all the data instead of chunking. + if gzw != nil && count%100 == 0 { + gzw.Flush() + } } // Clean up the JSON array - if first { + if count == 0 { // Never saw any data. Write an empty array. finalWriter.Write([]byte("[]")) } else { @@ -641,6 +655,10 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) { finalWriter.Write([]byte("]")) // end of JSON array } + if gzw != nil { + gzw.Flush() // flush any remaining compressed data to the client + } + rc.gm.Val(metrics.ReadMatch, int64(count)) rc.inst.Stop("encode-response") } From 690aee7e6a5c8f9956d51d5d447f1cfd5c64b810 Mon Sep 17 00:00:00 2001 From: Jemiah Westerman Date: Wed, 25 Feb 2026 13:41:49 -0800 Subject: [PATCH 3/3] Prevent goroutine leak on ctx timeout. --- api/api_gomux.go | 7 ++++++ entity/store.go | 60 ++++++++++++++++++++++++++++++++++-------------- 2 files changed, 50 insertions(+), 17 deletions(-) diff --git a/api/api_gomux.go b/api/api_gomux.go index fe631e3..a792a42 100644 --- a/api/api_gomux.go +++ b/api/api_gomux.go @@ -646,6 +646,13 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) { } } + // One final check of the context timeout. If the store closed the channel because the context was canceled, + // we want to return an error instead of partial results. + if err := ctx.Err(); err != nil { + api.readError(rc, w, err) + return + } + // Clean up the JSON array if count == 0 { // Never saw any data. Write an empty array. diff --git a/entity/store.go b/entity/store.go index 41cd77d..dbc23a6 100644 --- a/entity/store.go +++ b/entity/store.go @@ -17,10 +17,6 @@ import ( "github.com/square/etre/query" ) -const ( - StreamChannelBufferSize = 100 -) - // Store interface has methods needed to do CRUD operations on entities. type Store interface { ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) @@ -55,17 +51,17 @@ func NewStore(entities map[string]*mongo.Collection, cdcStore cdc.Store, cfg con func (s store) ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { c, ok := s.coll[entityType] if !ok { - panic("invalid entity type passed to ReadEntities: " + entityType) + panic("invalid entity type passed to ReadEntity: " + entityType) } - // Find and return all matching entities + // Find and return the matching entity p := bson.M{} if len(f.ReturnLabels) > 0 { for _, label := range f.ReturnLabels { p[label] = 1 } // Only include _id if explicitly set in f.ReturnLabels. If not, - // ok is false and we must explicitly exlude it because MongoDB + // ok is false and we must explicitly exclude it because MongoDB // returns it by default. if _, ok := p["_id"]; !ok { p["_id"] = 0 @@ -73,7 +69,11 @@ func (s store) ReadEntity(ctx context.Context, entityType string, entityId strin } // Query, we should only get one row - q, _ := query.Translate("_id=" + entityId) + q, err := query.Translate("_id=" + entityId) + if err != nil { + return nil, s.dbError(ctx, err, "invalid-query") + } + result := c.FindOne(ctx, Filter(q), options.FindOne().SetProjection(p)) if err := result.Err(); err != nil { nfe := mongo.ErrNoDocuments @@ -100,7 +100,7 @@ func (s store) StreamEntities(ctx context.Context, entityType string, q query.Qu panic("invalid entity type passed to StreamEntities: " + entityType) } - ch := make(chan EntityResult, StreamChannelBufferSize) + ch := make(chan EntityResult, 2*int32(s.config.BatchSize)) go func() { defer close(ch) @@ -115,19 +115,19 @@ func (s store) StreamEntities(ctx context.Context, entityType string, q query.Qu // No documents found, return to close the channel return } - ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-distinct")} + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-query-distinct")) return } var values []string err := dr.Decode(&values) if err != nil { - ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-distinct")} + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-query-distinct")) return } // Distinct doesn't return a cursor, so we just have to loop and send the results for _, v := range values { - ch <- EntityResult{Entity: etre.Entity{f.ReturnLabels[0]: v}} + s.writeEntityToChannel(ctx, ch, etre.Entity{f.ReturnLabels[0]: v}) } return } @@ -139,7 +139,7 @@ func (s store) StreamEntities(ctx context.Context, entityType string, q query.Qu p[label] = 1 } // Only include _id if explicitly set in f.ReturnLabels. If not, - // ok is false and we must explicitly exlude it because MongoDB + // ok is false and we must explicitly exclude it because MongoDB // returns it by default. if _, ok := p["_id"]; !ok { p["_id"] = 0 @@ -150,7 +150,7 @@ func (s store) StreamEntities(ctx context.Context, entityType string, q query.Qu opts := options.Find().SetProjection(p).SetBatchSize(int32(s.config.BatchSize)) cursor, err := c.Find(ctx, Filter(q), opts) if err != nil { - ch <- EntityResult{Err: s.dbError(ctx, err, "db-query")} + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-query")) return } defer cursor.Close(ctx) @@ -159,20 +159,46 @@ func (s store) StreamEntities(ctx context.Context, entityType string, q query.Qu for cursor.Next(ctx) { var entity etre.Entity if err := cursor.Decode(&entity); err != nil { - ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-cursor")} + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-read-cursor")) return } - ch <- EntityResult{Entity: entity} + s.writeEntityToChannel(ctx, ch, entity) } // Check for errors from iterating over cursor if err := cursor.Err(); err != nil { - ch <- EntityResult{Err: s.dbError(ctx, err, "db-read-cursor")} + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-read-cursor")) return } }() return ch } +func (s store) writeEntityToChannel(ctx context.Context, ch chan EntityResult, entity etre.Entity) { + select { + case <-ctx.Done(): + // The context was canceled or timed out. Bail out. + // We can't write the error to the channel because the receiver may have stopped listening. + // We depend on the receiver to see the ctx timeout as well. + return + case ch <- EntityResult{Entity: entity}: + // Successfully wrote to the channel + } + return +} + +func (s store) writeErrToChannel(ctx context.Context, ch chan EntityResult, err error) { + select { + case <-ctx.Done(): + // The context was canceled or timed out. Bail out. + // We can't write the error to the channel because the receiver may have stopped listening. + // We depend on the receiver to see the ctx timeout as well. + return + case ch <- EntityResult{Err: err}: + // Successfully wrote to the channel + } + return +} + // CreateEntities inserts many entities into DB. This method allows for partial // success and failure which means the return value and error are _not_ // mutually exclusive. Caller should check and handle both.