Stream query results to reduce memory usage

api/api_gomux.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log"
 	"math/rand"
 	"net/http"
@@ -576,26 +577,96 @@ 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 {
+
+	rc.inst.Start("encode-response")
+
+	// 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 {
+			api.readError(rc, w, e.Err)
+			return
+		}
+		// Handle context timeouts
+		if err := ctx.Err(); err != nil {
+			api.readError(rc, w, err)
+			return
+		}
+
+		// 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)
+			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
+		// 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)
+		} 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
+		}
+
+		// 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()
+		}
+	}
+
+	// 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
 	}
-	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)
+	// Clean up the JSON array
+	if count == 0 {
+		// 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
+	}
+
+	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")
 }
 
@@ -804,18 +875,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 +907,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())
 }
 
 // --------------------------------------------------------------------------

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

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)

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{}