Add chunked uploading for initial bootstrapping for new clients to improve performance for #298

.github/workflows/go-test.yml

@@ -24,8 +24,8 @@ jobs:
     needs: extra-delay
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-14]
-        test_shard: ["0", "1", "2", "3", "4"]
+        os: [ubuntu-latest]
+        test_shard: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
       fail-fast: false
     steps:
     - uses: actions/checkout@v4
@@ -75,7 +75,7 @@ jobs:
         OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
       run: |
           go install gotest.tools/gotestsum@bc98120
-          NUM_TEST_SHARDS=5 CURRENT_SHARD_NUM=${{ matrix.test_shard }} make test
+          NUM_TEST_SHARDS=10 CURRENT_SHARD_NUM=${{ matrix.test_shard }} make test
     - name: Extra Delay
       run: |
 
@@ -105,11 +105,11 @@ jobs:
       with:
         name: goldens-used-${{ matrix.os }}-${{ matrix.test_shard }}
         path: /tmp/goldens-used.txt
-    # - name: Setup tmate session
-    #   if: ${{ failure() }}
-    #   uses: mxschmitt/action-tmate@v3
-    #   with:
-    #     limit-access-to-actor: true
+    - name: Setup tmate session
+      if: ${{ failure() }}
+      uses: mxschmitt/action-tmate@v3
+      with:
+        limit-access-to-actor: true
   check-goldens:
     runs-on: ubuntu-latest
     needs: test

backend/server/internal/server/api_handlers.go

@@ -142,6 +142,7 @@ func (s *Server) apiSubmitDumpHandler(w http.ResponseWriter, r *http.Request) {
 	userId := getRequiredQueryParam(r, "user_id")
 	srcDeviceId := getRequiredQueryParam(r, "source_device_id")
 	requestingDeviceId := getRequiredQueryParam(r, "requesting_device_id")
+	isChunk := getOptionalQueryParam(r, "is_chunk", s.isTestEnvironment) == "true"
 	var entries []*shared.EncHistoryEntry
 	err := json.NewDecoder(r.Body).Decode(&entries)
 	if err != nil {
@@ -159,8 +160,10 @@ func (s *Server) apiSubmitDumpHandler(w http.ResponseWriter, r *http.Request) {
 
 	err = s.db.AddHistoryEntries(r.Context(), entries...)
 	checkGormError(err)
-	err = s.db.DumpRequestDeleteForUserAndDevice(r.Context(), userId, requestingDeviceId)
-	checkGormError(err)
+	if !isChunk {
+		err = s.db.DumpRequestDeleteForUserAndDevice(r.Context(), userId, requestingDeviceId)
+		checkGormError(err)
+	}
 
 	version := getHishtoryVersion(r)
 	remoteIPAddr := getRemoteAddr(r)

client/cmd/saveHistoryEntry.go

@@ -306,10 +306,21 @@ func deletePresavedEntries(ctx context.Context, entry *data.HistoryEntry, isRetr
 func handleDumpRequests(ctx context.Context, dumpRequests []*shared.DumpRequest) error {
 	db := hctx.GetDb(ctx)
 	config := hctx.GetConf(ctx)
-	if len(dumpRequests) > 0 {
-		lib.CheckFatalError(lib.RetrieveAdditionalEntriesFromRemote(ctx, "newclient"))
-		entries, err := lib.Search(ctx, db, "", 0)
+	if len(dumpRequests) == 0 {
+		return nil
+	}
+	if config.IsOffline {
+		return nil
+	}
+	lib.CheckFatalError(lib.RetrieveAdditionalEntriesFromRemote(ctx, "newclient"))
+	chunkSize := 1000
+	offset := 0
+	for {
+		entries, err := lib.SearchWithOffset(ctx, db, "", chunkSize, offset)
 		lib.CheckFatalError(err)
+		if len(entries) == 0 {
+			break
+		}
 		var encEntries []*shared.EncHistoryEntry
 		for _, entry := range entries {
 			enc, err := data.EncryptHistoryEntry(config.UserSecret, *entry)
@@ -319,13 +330,17 @@ func handleDumpRequests(ctx context.Context, dumpRequests []*shared.DumpRequest)
 		reqBody, err := json.Marshal(encEntries)
 		lib.CheckFatalError(err)
 		for _, dumpRequest := range dumpRequests {
-			if !config.IsOffline {
-				// TODO: Test whether this fails if the data is extremely large? It may need to be chunked
-				_, err := lib.ApiPost(ctx, "/api/v1/submit-dump?user_id="+dumpRequest.UserId+"&requesting_device_id="+dumpRequest.RequestingDeviceId+"&source_device_id="+config.DeviceId, "application/json", reqBody)
-				lib.CheckFatalError(err)
-			}
+			_, err := lib.ApiPost(ctx, "/api/v1/submit-dump?user_id="+dumpRequest.UserId+"&requesting_device_id="+dumpRequest.RequestingDeviceId+"&source_device_id="+config.DeviceId+"&is_chunk=true", "application/json", reqBody)
+			lib.CheckFatalError(err)
 		}
 	}
+	// Send a final dump response without the `is_chunk` param so that the backend knows we're done
+	reqBody, err := json.Marshal([]*shared.EncHistoryEntry{})
+	lib.CheckFatalError(err)
+	for _, dumpRequest := range dumpRequests {
+		_, err := lib.ApiPost(ctx, "/api/v1/submit-dump?user_id="+dumpRequest.UserId+"&requesting_device_id="+dumpRequest.RequestingDeviceId+"&source_device_id="+config.DeviceId+"&is_chunk=true", "application/json", reqBody)
+		lib.CheckFatalError(err)
+	}
 	return nil
 }