feat: improve bundle performance [IDE-937] (#87)

Author: andrewrobinsonhodges-snyk

internal/bundle/bundle.go

@@ -30,6 +30,7 @@ type Bundle interface {
 	UploadBatch(ctx context.Context, requestId string, batch *Batch) error
 	GetBundleHash() string
 	GetFiles() map[string]deepcode.BundleFile
+	ClearFiles()
 	GetMissingFiles() []string
 	GetLimitToFiles() []string
 	GetRootPath() string
@@ -81,6 +82,10 @@ func (b *deepCodeBundle) GetFiles() map[string]deepcode.BundleFile {
 	return b.files
 }
 
+func (b *deepCodeBundle) ClearFiles() {
+	b.files = make(map[string]deepcode.BundleFile)
+}
+
 func (b *deepCodeBundle) GetMissingFiles() []string {
 	return b.missingFiles
 }
@@ -136,15 +141,15 @@ func NewBatch(documents map[string]deepcode.BundleFile) *Batch {
 
 // todo simplify the size computation
 // maybe consider an addFile / canFitFile interface with proper error handling
-func (b *Batch) canFitFile(uri string, content []byte) bool {
-	docPayloadSize := b.getTotalDocPayloadSize(uri, content)
+func (b *Batch) canFitFile(uri string, contentSize int) bool {
+	docPayloadSize := b.getTotalDocPayloadSize(uri, contentSize)
 	newSize := docPayloadSize + b.getSize()
 	b.size += docPayloadSize
 	return newSize < maxUploadBatchSize
 }
 
-func (b *Batch) getTotalDocPayloadSize(documentURI string, content []byte) int {
-	return len(jsonHashSizePerFile) + len(jsonOverheadPerFile) + len([]byte(documentURI)) + len(content)
+func (b *Batch) getTotalDocPayloadSize(documentURI string, contentSize int) int {
+	return len(jsonHashSizePerFile) + len(jsonOverheadPerFile) + len([]byte(documentURI)) + contentSize
 }
 
 func (b *Batch) getSize() int {

internal/bundle/bundle_manager.go

@@ -105,20 +105,25 @@ func (b *bundleManager) Create(ctx context.Context,
 		if !supported {
 			continue
 		}
-		var fileContent []byte
-		fileContent, err = os.ReadFile(absoluteFilePath)
-		if err != nil {
-			b.logger.Error().Err(err).Str("filePath", absoluteFilePath).Msg("could not load content of file")
+
+		fileInfo, fileErr := os.Stat(absoluteFilePath)
+		if fileErr != nil {
+			b.logger.Error().Err(err).Str("filePath", absoluteFilePath).Msg("Failed to read file info")
 			continue
 		}
 
-		if !(len(fileContent) > 0 && len(fileContent) <= maxFileSize) {
+		if fileInfo.Size() == 0 || fileInfo.Size() > maxFileSize {
 			continue
 		}
 
-		var relativePath string
-		relativePath, err = util.ToRelativeUnixPath(rootPath, absoluteFilePath)
-		if err != nil {
+		fileContent, fileErr := os.ReadFile(absoluteFilePath)
+		if fileErr != nil {
+			b.logger.Error().Err(err).Str("filePath", absoluteFilePath).Msg("Failed to load content of file")
+			continue
+		}
+
+		relativePath, fileErr := util.ToRelativeUnixPath(rootPath, absoluteFilePath)
+		if fileErr != nil {
 			b.errorReporter.CaptureError(err, observability.ErrorReporterOptions{ErrorDiagnosticPath: rootPath})
 		}
 		relativePath = util.EncodePath(relativePath)
@@ -181,16 +186,44 @@ func (b *bundleManager) Upload(
 			if err := ctx.Err(); err != nil {
 				return bundle, err
 			}
+			b.enrichBatchWithFileContent(batch, bundle.GetRootPath())
 			err := bundle.UploadBatch(s.Context(), requestId, batch)
 			if err != nil {
 				return bundle, err
 			}
+			batch.documents = make(map[string]deepcode.BundleFile)
 		}
 	}
 
+	// bundle doesn't need file map anymore since they are already grouped and uploaded
+	bundle.ClearFiles()
 	return bundle, nil
 }
 
+func (b *bundleManager) enrichBatchWithFileContent(batch *Batch, rootPath string) {
+	for filePath, bundleFile := range batch.documents {
+		absPath, err := util.DecodePath(util.ToAbsolutePath(rootPath, filePath))
+		if err != nil {
+			b.logger.Error().Err(err).Str("file", filePath).Msg("Failed to decode Path")
+			continue
+		}
+		content, err := os.ReadFile(absPath)
+		if err != nil {
+			b.logger.Error().Err(err).Str("file", filePath).Msg("Failed to read bundle file")
+			continue
+		}
+
+		utf8Content, err := util.ConvertToUTF8(content)
+		if err != nil {
+			b.logger.Error().Err(err).Str("file", filePath).Msg("Failed to convert bundle file to UTF-8")
+			continue
+		}
+
+		bundleFile.Content = string(utf8Content)
+		batch.documents[filePath] = bundleFile
+	}
+}
+
 func (b *bundleManager) groupInBatches(
 	ctx context.Context,
 	bundle Bundle,
@@ -212,12 +245,11 @@ func (b *bundleManager) groupInBatches(
 		}
 
 		file := files[filePath]
-		var fileContent = []byte(file.Content)
-		if batch.canFitFile(filePath, fileContent) {
-			b.logger.Trace().Str("path", filePath).Int("size", len(fileContent)).Msgf("added to deepCodeBundle #%v", len(batches))
+		if batch.canFitFile(filePath, file.ContentSize) {
+			b.logger.Trace().Str("path", filePath).Int("size", file.ContentSize).Msgf("added to deepCodeBundle #%v", len(batches))
 			batch.documents[filePath] = file
 		} else {
-			b.logger.Trace().Str("path", filePath).Int("size", len(fileContent)).Msgf("created new deepCodeBundle - %v bundles in this upload so far", len(batches))
+			b.logger.Trace().Str("path", filePath).Int("size", file.ContentSize).Msgf("created new deepCodeBundle - %v bundles in this upload so far", len(batches))
 			newUploadBatch := NewBatch(map[string]deepcode.BundleFile{})
 			newUploadBatch.documents[filePath] = file
 			batches = append(batches, newUploadBatch)

internal/bundle/bundle_manager_test.go

@@ -444,7 +444,7 @@ func createTempFileInDir(t *testing.T, name string, size int, temporaryDir strin
 	t.Helper()
 
 	documentURI, fileContent := createFileOfSize(t, name, size, temporaryDir)
-	return documentURI, deepcode.BundleFile{Hash: util.Hash(fileContent), Content: string(fileContent)}
+	return documentURI, deepcode.BundleFile{Hash: util.Hash(fileContent), ContentSize: size}
 }
 
 func Test_IsSupported_Extensions(t *testing.T) {

internal/bundle/mocks/bundle.go

@@ -20,14 +20,16 @@ import (
 )
 
 type BundleFile struct {
-	Hash    string `json:"hash"`
-	Content string `json:"content"`
+	Hash        string `json:"hash"`
+	Content     string `json:"content"`
+	ContentSize int    `json:"size"`
 }
 
 func BundleFileFrom(content []byte) BundleFile {
 	file := BundleFile{
-		Hash:    util.Hash(content),
-		Content: string(content),
+		Hash:        util.Hash(content),
+		Content:     "", // We create the bundleFile empty, and enrich  with content later.
+		ContentSize: len(content),
 	}
 	return file
 }

internal/deepcode/helpers.go

@@ -26,13 +26,18 @@ import (
 )
 
 func Hash(content []byte) string {
-	byteReader := bytes.NewReader(content)
-	reader, _ := charset.NewReaderLabel("UTF-8", byteReader)
-	utf8content, err := io.ReadAll(reader)
+	utf8content, err := ConvertToUTF8(content)
 	if err != nil {
 		utf8content = content
 	}
 	b := sha256.Sum256(utf8content)
 	sum256 := hex.EncodeToString(b[:])
 	return sum256
 }
+
+func ConvertToUTF8(content []byte) ([]byte, error) {
+	byteReader := bytes.NewReader(content)
+	reader, _ := charset.NewReaderLabel("UTF-8", byteReader)
+	utf8content, err := io.ReadAll(reader)
+	return utf8content, err
+}