Read packages from index in batch to insert them in the SQL database (#1326)

Updated the process to load the search index from the storage to
read in batches all the packages. Each batch of packages is inserted
directly in the database.
This allows us to not keep all the databases from the storage index
in memory.

Author: mrodm

internal/storage/index.go

@@ -23,6 +23,7 @@ type packageIndex struct {
 
 func loadSearchIndexAll(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, bucketName, rootStoragePath string, aCursor cursor) (*packages.Packages, error) {
 	span, ctx := apm.StartSpan(ctx, "LoadSearchIndexAll", "app")
+	span.Context.SetLabel("load.method", "full")
 	defer span.End()
 
 	indexFile := searchIndexAllFile
@@ -114,3 +115,112 @@ func LoadPackagesAndCursorFromIndex(ctx context.Context, logger *zap.Logger, sto
 	}
 	return anIndex, storageCursor.Current, nil
 }
+
+func loadSearchIndexAllBatches(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, bucketName, rootStoragePath string, aCursor cursor, batchSize int, process func(packages.Packages) error) error {
+	span, ctx := apm.StartSpan(ctx, "LoadSearchIndexAll", "app")
+	span.Context.SetLabel("load.method", "batches")
+	span.Context.SetLabel("load.batch.size", batchSize)
+	defer span.End()
+
+	indexFile := searchIndexAllFile
+
+	logger.Debug("load search-index-all index", zap.String("index.file", indexFile))
+
+	rootedIndexStoragePath := buildIndexStoragePath(rootStoragePath, aCursor, indexFile)
+	objectReader, err := storageClient.Bucket(bucketName).Object(rootedIndexStoragePath).NewReader(ctx)
+	if err != nil {
+		return fmt.Errorf("can't read the index file (path: %s): %w", rootedIndexStoragePath, err)
+	}
+	defer objectReader.Close()
+
+	// Using a decoder here as tokenizer to parse the list of packages as a stream
+	// instead of needing the whole document in memory at the same time. This helps
+	// reducing memory usage.
+	// Using `Unmarshal(doc, &sia)` would require to read the whole document.
+	// Using `dec.Decode(&sia)` would also make the decoder to keep the whole document
+	// in memory.
+	// `jsoniter` seemed to be slightly faster, but to use more memory for our use case,
+	// and we are looking to optimize for memory use.
+	dec := json.NewDecoder(objectReader)
+	count := 0
+	packages := make(packages.Packages, 0, batchSize)
+	for dec.More() {
+		// Read everything till the "packages" key in the map.
+		token, err := dec.Token()
+		if err != nil {
+			return fmt.Errorf("unexpected error while reading index file: %w", err)
+		}
+		if key, ok := token.(string); !ok || key != "packages" {
+			continue
+		}
+
+		// Read the opening array now.
+		token, err = dec.Token()
+		if err != nil {
+			return fmt.Errorf("unexpected error while reading index file: %w", err)
+		}
+		if delim, ok := token.(json.Delim); !ok || delim != '[' {
+			return fmt.Errorf("expected opening array, found %v", token)
+		}
+
+		// Read the array of packages one by one.
+		for dec.More() {
+			var p packageIndex
+			err = dec.Decode(&p)
+			if err != nil {
+				return fmt.Errorf("unexpected error parsing package from index file (token: %v): %w", token, err)
+			}
+			packages = append(packages, p.PackageManifest)
+			count++
+
+			if count >= batchSize {
+				err = process(packages)
+				if err != nil {
+					return fmt.Errorf("error processing batch of packages: %w", err)
+				}
+				count = 0
+				packages = packages[:0] // Reset the slice to reuse the memory
+			}
+		}
+
+		// Read the closing array delimiter.
+		token, err = dec.Token()
+		if err != nil {
+			return fmt.Errorf("unexpected error while reading index file: %w", err)
+		}
+		if delim, ok := token.(json.Delim); !ok || delim != ']' {
+			return fmt.Errorf("expected closing array, found %v", token)
+		}
+	}
+	if len(packages) > 0 {
+		err = process(packages)
+		if err != nil {
+			return fmt.Errorf("error processing final batch of packages: %w", err)
+		}
+	}
+	return nil
+}
+
+func LoadPackagesAndCursorFromIndexBatches(ctx context.Context, logger *zap.Logger, storageClient *storage.Client, storageBucketInternal, currentCursor string, batchSize int, process func(packages.Packages) error) (string, error) {
+	bucketName, rootStoragePath, err := extractBucketNameFromURL(storageBucketInternal)
+	if err != nil {
+		return "", fmt.Errorf("can't extract bucket name from URL (url: %s): %w", storageBucketInternal, err)
+	}
+
+	storageCursor, err := loadCursor(ctx, logger, storageClient, bucketName, rootStoragePath)
+	if err != nil {
+		return "", fmt.Errorf("can't load latest cursor: %w", err)
+	}
+
+	if storageCursor.Current == currentCursor {
+		logger.Info("cursor is up-to-date", zap.String("cursor.current", currentCursor))
+		return currentCursor, nil
+	}
+	logger.Info("cursor will be updated", zap.String("cursor.current", currentCursor), zap.String("cursor.next", storageCursor.Current))
+
+	err = loadSearchIndexAllBatches(ctx, logger, storageClient, bucketName, rootStoragePath, *storageCursor, batchSize, process)
+	if err != nil {
+		return "", fmt.Errorf("can't load the search-index-all index content: %w", err)
+	}
+	return storageCursor.Current, nil
+}

internal/storage/sqlindexer.go

@@ -199,45 +199,45 @@ func (i *SQLIndexer) updateIndex(ctx context.Context) error {
 		metrics.StorageIndexerUpdateIndexDurationSeconds.Observe(time.Since(start).Seconds())
 	}()
 
-	anIndex, currentCursor, err := LoadPackagesAndCursorFromIndex(ctx, i.logger, i.storageClient, i.options.PackageStorageBucketInternal, i.cursor)
+	defer func() {
+		startClean := time.Now()
+		if err := i.cleanBackupDatabase(ctx); err != nil {
+			i.logger.Error("Failed to clean backup database", zap.Error(err))
+		}
+		startCleanDuration := time.Since(startClean)
+		i.logger.Debug("Cleaned backup database", zap.Duration("elapsed.time", time.Since(startClean)), zap.String("elapsed.time.human", startCleanDuration.String()))
+	}()
+
+	numPackages := 0
+	currentCursor, err := LoadPackagesAndCursorFromIndexBatches(ctx, i.logger, i.storageClient, i.options.PackageStorageBucketInternal, i.cursor, i.maxBulkAddBatch, func(pkgs packages.Packages) error {
+		// This function is called for each batch of packages read from the index.
+		startUpdate := time.Now()
+		if err := i.updateDatabase(ctx, &pkgs); err != nil {
+			return fmt.Errorf("failed to update database: %w", err)
+		}
+		startDuration := time.Since(startUpdate)
+		numPackages += len(pkgs)
+		i.logger.Debug("Filled database with a batch of packages", zap.Duration("elapsed.time", startDuration), zap.String("elapsed.time.human", startDuration.String()), zap.Int("num.packages", len(pkgs)))
+		return nil
+	})
 	if err != nil {
 		metrics.StorageIndexerUpdateIndexErrorsTotal.Inc()
 		return fmt.Errorf("can't load the search-index-all index content: %w", err)
 	}
 	if i.cursor == currentCursor {
 		return nil
 	}
-	if anIndex == nil {
-		i.logger.Info("Downloaded new search-index-all index. No packages found.")
-		return nil
-	}
-	i.logger.Info("Downloaded new search-index-all index", zap.String("index.packages.size", fmt.Sprintf("%d", len(*anIndex))))
-
-	i.logger.Info("Filling database")
-	startUpdate := time.Now()
-	err = i.updateDatabase(ctx, anIndex)
-	if err != nil {
-		return fmt.Errorf("failed to update database: %w", err)
-	}
-	startDuration := time.Since(startUpdate)
-	i.logger.Info("Filled database with latest packages", zap.Duration("elapsed.time", startDuration), zap.String("elapsed.time.human", startDuration.String()))
+	i.logger.Info("Downloaded new search-index-all index", zap.String("index.packages.size", fmt.Sprintf("%d", numPackages)))
 
 	startLock := time.Now()
-	i.swapDatabases(ctx, currentCursor, len(*anIndex))
-	i.logger.Info("Elapsed time in lock for updating index database", zap.Duration("lock.duration", time.Since(startLock)))
+	i.swapDatabases(ctx, currentCursor, numPackages)
+	i.logger.Debug("Elapsed time in lock for updating index database", zap.Duration("lock.duration", time.Since(startLock)))
 
 	if err != nil {
 		metrics.StorageIndexerUpdateIndexErrorsTotal.Inc()
 		return err
 	}
 
-	startClean := time.Now()
-	err = i.cleanBackupDatabase(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to clean backup database: %w", err)
-	}
-	startCleanDuration := time.Since(startClean)
-	i.logger.Info("Cleaned backup database", zap.Duration("elapsed.time", time.Since(startClean)), zap.String("elapsed.time.human", startCleanDuration.String()))
 	return nil
 }