Optimize query latency by moving the read count incrementing to a background task

backend/server/server.go

@@ -216,9 +216,6 @@ func apiQueryHandler(ctx context.Context, w http.ResponseWriter, r *http.Request
 	userId := getRequiredQueryParam(r, "user_id")
 	deviceId := getRequiredQueryParam(r, "device_id")
 	updateUsageData(ctx, r, userId, deviceId, 0, true)
-	// Increment the count
-	// TODO(optimization): Have this run in a background thread to avoid slowing down the entire response
-	checkGormResult(GLOBAL_DB.WithContext(ctx).Exec("UPDATE enc_history_entries SET read_count = read_count + 1 WHERE device_id = ?", deviceId))
 
 	// Delete any entries that match a pending deletion request
 	var deletionRequests []*shared.DeletionRequest
@@ -230,7 +227,7 @@ func apiQueryHandler(ctx context.Context, w http.ResponseWriter, r *http.Request
 		}
 	}
 
-	// Then retrieve, to avoid a race condition
+	// Then retrieve
 	tx := GLOBAL_DB.WithContext(ctx).Where("device_id = ? AND read_count < 5", deviceId)
 	var historyEntries []*shared.EncHistoryEntry
 	checkGormResult(tx.Find(&historyEntries))
@@ -240,9 +237,29 @@ func apiQueryHandler(ctx context.Context, w http.ResponseWriter, r *http.Request
 		panic(err)
 	}
 	w.Write(resp)
+
+	// And finally, kick off a background goroutine that will increment the read count. Doing it in the background avoids
+	// blocking the entire response. This does have a potential race condition, but that is fine.
+	if isProductionEnvironment() {
+		go func() {
+			span, ctx := tracer.StartSpanFromContext(ctx, "apiQueryHandler.incrementReadCount")
+			err = incrementReadCounts(ctx, deviceId)
+			span.Finish(tracer.WithError(err))
+		}()
+	} else {
+		err = incrementReadCounts(ctx, deviceId)
+		if err != nil {
+			panic(fmt.Sprintf("failed to increment read counts"))
+		}
+	}
+
 	GLOBAL_STATSD.Incr("hishtory.query", []string{}, 1.0)
 }
 
+func incrementReadCounts(ctx context.Context, deviceId string) error {
+	return GLOBAL_DB.WithContext(ctx).Exec("UPDATE enc_history_entries SET read_count = read_count + 1 WHERE device_id = ?", deviceId).Error
+}
+
 func getRemoteAddr(r *http.Request) string {
 	addr, ok := r.Header["X-Real-Ip"]
 	if !ok || len(addr) == 0 {

backend/server/server_test.go

@@ -64,7 +64,7 @@ func TestESubmitThenQuery(t *testing.T) {
 	if dbEntry.UserId != data.UserId("key") {
 		t.Fatalf("Response contains an incorrect device ID: %#v", *dbEntry)
 	}
-	if dbEntry.ReadCount != 1 {
+	if dbEntry.ReadCount != 0 {
 		t.Fatalf("db.ReadCount should have been 1, was %v", dbEntry.ReadCount)
 	}
 	decEntry, err := data.DecryptHistoryEntry("key", *dbEntry)
@@ -92,7 +92,7 @@ func TestESubmitThenQuery(t *testing.T) {
 	if dbEntry.UserId != data.UserId("key") {
 		t.Fatalf("Response contains an incorrect device ID: %#v", *dbEntry)
 	}
-	if dbEntry.ReadCount != 1 {
+	if dbEntry.ReadCount != 0 {
 		t.Fatalf("db.ReadCount should have been 1, was %v", dbEntry.ReadCount)
 	}
 	decEntry, err = data.DecryptHistoryEntry("key", *dbEntry)
@@ -270,7 +270,7 @@ func TestDumpRequestAndResponse(t *testing.T) {
 		if dbEntry.UserId != userId {
 			t.Fatalf("Response contains an incorrect user ID: %#v", *dbEntry)
 		}
-		if dbEntry.ReadCount != 1 {
+		if dbEntry.ReadCount != 0 {
 			t.Fatalf("db.ReadCount should have been 1, was %v", dbEntry.ReadCount)
 		}
 		decEntry, err := data.DecryptHistoryEntry("dkey", *dbEntry)
@@ -381,7 +381,7 @@ func TestDeletionRequests(t *testing.T) {
 		if dbEntry.UserId != data.UserId("dkey") {
 			t.Fatalf("Response contains an incorrect device ID: %#v", *dbEntry)
 		}
-		if dbEntry.ReadCount != 1 {
+		if dbEntry.ReadCount != 0 {
 			t.Fatalf("db.ReadCount should have been 1, was %v", dbEntry.ReadCount)
 		}
 		decEntry, err := data.DecryptHistoryEntry("dkey", *dbEntry)
@@ -406,6 +406,7 @@ func TestDeletionRequests(t *testing.T) {
 	addDeletionRequestHandler(context.Background(), nil, req)
 
 	// Query again for device id 1 and get a single result
+	time.Sleep(10 * time.Millisecond)
 	w = httptest.NewRecorder()
 	searchReq = httptest.NewRequest(http.MethodGet, "/?device_id="+devId1+"&user_id="+userId, nil)
 	apiQueryHandler(context.Background(), w, searchReq)
@@ -424,7 +425,7 @@ func TestDeletionRequests(t *testing.T) {
 	if dbEntry.UserId != data.UserId("dkey") {
 		t.Fatalf("Response contains an incorrect device ID: %#v", *dbEntry)
 	}
-	if dbEntry.ReadCount != 2 {
+	if dbEntry.ReadCount != 1 {
 		t.Fatalf("db.ReadCount should have been 1, was %v", dbEntry.ReadCount)
 	}
 	decEntry, err := data.DecryptHistoryEntry("dkey", *dbEntry)
@@ -452,7 +453,7 @@ func TestDeletionRequests(t *testing.T) {
 	if dbEntry.UserId != data.UserId("dOtherkey") {
 		t.Fatalf("Response contains an incorrect device ID: %#v", *dbEntry)
 	}
-	if dbEntry.ReadCount != 1 {
+	if dbEntry.ReadCount != 0 {
 		t.Fatalf("db.ReadCount should have been 1, was %v", dbEntry.ReadCount)
 	}
 	decEntry, err = data.DecryptHistoryEntry("dOtherkey", *dbEntry)