Remove an in-memory buffer for LSIF transformation

changelogs/unreleased/284-remove-zip-buffer.yml

+---
+title: Remove an in-memory buffer for LSIF transformation
+merge_request: 586
+author:
+type: performance

internal/lsif_transformer/parser/parser.go

@@ -3,8 +3,8 @@ package parser
 import (
 	"archive/zip"
 	"bufio"
-	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
@@ -16,13 +16,15 @@ var (
 
 type Parser struct {
 	Docs *Docs
+
+	pr *io.PipeReader
 }
 
 type Config struct {
 	TempPath string
 }
 
-func NewParser(r io.Reader, config Config) (*Parser, error) {
+func NewParser(r io.Reader, config Config) (io.ReadCloser, error) {
 	docs, err := NewDocs(config)
 	if err != nil {
 		return nil, err
@@ -45,28 +47,44 @@ func NewParser(r io.Reader, config Config) (*Parser, error) {
 		}
 	}
 
-	return &Parser{Docs: docs}, nil
-}
-
-func (p *Parser) ZipReader() (io.Reader, error) {
-	buf := new(bytes.Buffer)
-	w := zip.NewWriter(buf)
-
-	if err := p.Docs.SerializeEntries(w); err != nil {
-		return nil, err
+	pr, pw := io.Pipe()
+	parser := &Parser{
+		Docs: docs,
+		pr:   pr,
 	}
 
-	if err := w.Close(); err != nil {
-		return nil, err
-	}
+	go parser.parse(pw)
+
+	return parser, nil
+}
 
-	return buf, nil
+func (p *Parser) Read(b []byte) (int, error) {
+	return p.pr.Read(b)
 }
 
 func (p *Parser) Close() error {
+	p.pr.Close()
+
 	return p.Docs.Close()
 }
 
+func (p *Parser) parse(pw *io.PipeWriter) {
+	zw := zip.NewWriter(pw)
+
+	if err := p.Docs.SerializeEntries(zw); err != nil {
+		zw.Close() // Free underlying resources only
+		pw.CloseWithError(fmt.Errorf("lsif parser: Docs.SerializeEntries: %v", err))
+		return
+	}
+
+	if err := zw.Close(); err != nil {
+		pw.CloseWithError(fmt.Errorf("lsif parser: ZipWriter.Close: %v", err))
+		return
+	}
+
+	pw.Close()
+}
+
 func openZipReader(reader io.Reader, tempDir string) (io.Reader, error) {
 	tempFile, err := ioutil.TempFile(tempDir, Lsif)
 	if err != nil {

internal/lsif_transformer/parser/parser_test.go

@@ -38,24 +38,21 @@ func verifyCorrectnessOf(t *testing.T, tmpDir, fileName string) {
 }
 
 func createFiles(t *testing.T, filePath, tmpDir string) {
+	t.Helper()
 	file, err := os.Open(filePath)
 	require.NoError(t, err)
 
-	p, err := NewParser(file, Config{})
+	parser, err := NewParser(file, Config{})
 	require.NoError(t, err)
 
-	r, err := p.ZipReader()
-	require.NoError(t, err)
-
-	require.NoError(t, p.Close())
-
 	zipFileName := tmpDir + ".zip"
 	w, err := os.Create(zipFileName)
 	require.NoError(t, err)
 	defer os.RemoveAll(zipFileName)
 
-	_, err = io.Copy(w, r)
+	_, err = io.Copy(w, parser)
 	require.NoError(t, err)
+	require.NoError(t, parser.Close())
 
 	extractZipFiles(t, tmpDir, zipFileName)
 }

internal/lsif_transformer/parser/performance_test.go

 package parser
 
 import (
+	"io"
+	"io/ioutil"
 	"os"
 	"runtime"
 	"testing"
@@ -19,12 +21,12 @@ func BenchmarkGenerate(b *testing.B) {
 			file, err := os.Open(filePath)
 			require.NoError(b, err)
 
-			p, err := NewParser(file, Config{})
+			parser, err := NewParser(file, Config{})
 			require.NoError(b, err)
 
-			_, err = p.ZipReader()
+			_, err = io.Copy(ioutil.Discard, parser)
 			require.NoError(b, err)
-			require.NoError(b, p.Close())
+			require.NoError(b, parser.Close())
 		})
 	}
 

internal/upload/rewrite.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"mime/multipart"
 	"net/http"
 	"strings"
@@ -124,7 +125,7 @@ func (rew *rewriter) handleFilePart(ctx context.Context, name string, p *multipa
 
 	opts.TempFilePrefix = filename
 
-	var inputReader io.Reader
+	var inputReader io.ReadCloser
 	var err error
 	switch {
 	case exif.IsExifFile(filename):
@@ -138,9 +139,11 @@ func (rew *rewriter) handleFilePart(ctx context.Context, name string, p *multipa
 			return err
 		}
 	default:
-		inputReader = p
+		inputReader = ioutil.NopCloser(p)
 	}
 
+	defer inputReader.Close()
+
 	fh, err := filestore.SaveFileFromReader(ctx, inputReader, -1, opts)
 	if err != nil {
 		switch err {
@@ -166,36 +169,25 @@ func (rew *rewriter) handleFilePart(ctx context.Context, name string, p *multipa
 	return rew.filter.ProcessFile(ctx, name, fh, rew.writer)
 }
 
-func handleExifUpload(ctx context.Context, r io.Reader, filename string) (io.Reader, error) {
+func handleExifUpload(ctx context.Context, r io.Reader, filename string) (io.ReadCloser, error) {
 	log.WithContextFields(ctx, log.Fields{
 		"filename": filename,
 	}).Print("running exiftool to remove any metadata")
 
-	return exif.NewCleaner(ctx, r)
-}
-
-func handleLsifUpload(ctx context.Context, reader io.Reader, tempPath, filename string, preauth *api.Response) (io.Reader, error) {
-	parserConfig := parser.Config{
-		TempPath: tempPath,
-	}
-
-	p, err := parser.NewParser(reader, parserConfig)
+	r, err := exif.NewCleaner(ctx, r)
 	if err != nil {
 		return nil, err
 	}
 
-	z, err := p.ZipReader()
-	if err != nil {
-		return nil, err
-	}
+	return ioutil.NopCloser(r), nil
+}
 
-	if err := p.Close(); err != nil {
-		log.WithContextFields(ctx, log.Fields{
-			"filename": filename,
-		}).Print("failed to close lsif parser: " + err.Error())
+func handleLsifUpload(ctx context.Context, reader io.Reader, tempPath, filename string, preauth *api.Response) (io.ReadCloser, error) {
+	parserConfig := parser.Config{
+		TempPath: tempPath,
 	}
 
-	return z, nil
+	return parser.NewParser(reader, parserConfig)
 }
 
 func (rew *rewriter) copyPart(ctx context.Context, name string, p *multipart.Part) error {