feat: Update Stream to accept a bufSize for tar CopyBuffer window

Author: devanbenz

pkg/tar/stream.go

@@ -16,7 +16,7 @@ import (
 // possibly writing each file to a tar writer stream.  By default StreamFile is used, which will result in all files
 // being written.  A custom writeFunc can be passed so that each file may be written, modified+written, or skipped
 // depending on the custom logic.
-func Stream(w io.Writer, dir, relativePath string, writeFunc func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error) error {
+func Stream(w io.Writer, dir, relativePath string, bufSize uint64, writeFunc func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer, bufSize uint64) error) error {
 	tw := tar.NewWriter(w)
 	defer tw.Close()
 
@@ -41,30 +41,31 @@ func Stream(w io.Writer, dir, relativePath string, writeFunc func(f os.FileInfo,
 			return err
 		}
 
-		return writeFunc(f, filepath.Join(relativePath, subDir), path, tw)
+		return writeFunc(f, filepath.Join(relativePath, subDir), path, tw, bufSize)
 	})
 }
 
 // Generates a filtering function for Stream that checks an incoming file, and only writes the file to the stream if
 // its mod time is later than since.  Example: to tar only files newer than a certain datetime, use
 // tar.Stream(w, dir, relativePath, SinceFilterTarFile(datetime))
-func SinceFilterTarFile(since time.Time) func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error {
-	return func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error {
+func SinceFilterTarFile(since time.Time) func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer, bufSize uint64) error {
+	return func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer, bufSize uint64) error {
 		if f.ModTime().After(since) {
-			return StreamFile(f, shardRelativePath, fullPath, tw)
+			return StreamFile(f, shardRelativePath, fullPath, tw, bufSize)
 		}
 		return nil
 	}
 }
 
 // stream a single file to tw, extending the header name using the shardRelativePath
-func StreamFile(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error {
-	return StreamRenameFile(f, f.Name(), shardRelativePath, fullPath, tw)
+func StreamFile(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer, bufSize uint64) error {
+	return StreamRenameFile(f, f.Name(), shardRelativePath, fullPath, tw, bufSize)
 }
 
 // / Stream a single file to tw, using tarHeaderFileName instead of the actual filename
 // e.g., when we want to write a *.tmp file using the original file's non-tmp name.
-func StreamRenameFile(f os.FileInfo, tarHeaderFileName, relativePath, fullPath string, tw *tar.Writer) error {
+func StreamRenameFile(f os.FileInfo, tarHeaderFileName, relativePath, fullPath string, tw *tar.Writer, bufSize uint64) error {
+	buf := make([]byte, bufSize)
 	h, err := tar.FileInfoHeader(f, f.Name())
 	if err != nil {
 		return err
@@ -86,7 +87,7 @@ func StreamRenameFile(f os.FileInfo, tarHeaderFileName, relativePath, fullPath s
 
 	defer fr.Close()
 
-	_, err = io.CopyN(tw, fr, h.Size)
+	_, err = io.CopyBuffer(tw, fr, buf)
 
 	return err
 }

tsdb/engine/tsm1/engine.go

@@ -227,6 +227,9 @@ type Engine struct {
 
 	// muDigest ensures only one goroutine can generate a digest at a time.
 	muDigest sync.RWMutex
+
+	// TarStreamBufferSize is the size used for our buffer windows within the tar package.
+	TarStreamBufferSize uint64
 }
 
 // NewEngine returns a new instance of Engine.
@@ -994,13 +997,13 @@ func (e *Engine) Backup(w io.Writer, basePath string, since time.Time) error {
 		}
 	}()
 
-	return intar.Stream(w, path, basePath, intar.SinceFilterTarFile(since))
+	return intar.Stream(w, path, basePath, e.TarStreamBufferSize, intar.SinceFilterTarFile(since))
 }
 
-func (e *Engine) timeStampFilterTarFile(start, end time.Time) func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error {
-	return func(fi os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error {
+func (e *Engine) timeStampFilterTarFile(start, end time.Time) func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer, bufSize uint64) error {
+	return func(fi os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer, bufSize uint64) error {
 		if !strings.HasSuffix(fi.Name(), ".tsm") {
-			return intar.StreamFile(fi, shardRelativePath, fullPath, tw)
+			return intar.StreamFile(fi, shardRelativePath, fullPath, tw, e.TarStreamBufferSize)
 		}
 
 		f, err := os.Open(fullPath)
@@ -1014,7 +1017,7 @@ func (e *Engine) timeStampFilterTarFile(start, end time.Time) func(f os.FileInfo
 
 		// Grab the tombstone file if one exists.
 		if ts := r.TombstoneStats(); ts.TombstoneExists {
-			return intar.StreamFile(fi, shardRelativePath, filepath.Base(ts.Path), tw)
+			return intar.StreamFile(fi, shardRelativePath, filepath.Base(ts.Path), tw, e.TarStreamBufferSize)
 		}
 
 		min, max := r.TimeRange()
@@ -1042,7 +1045,7 @@ func (e *Engine) timeStampFilterTarFile(start, end time.Time) func(f os.FileInfo
 
 		// the TSM file is 100% inside the range, so we can just write it without scanning each block
 		if min >= start.UnixNano() && max <= end.UnixNano() {
-			if err := intar.StreamFile(fi, shardRelativePath, fullPath, tw); err != nil {
+			if err := intar.StreamFile(fi, shardRelativePath, fullPath, tw, bufSize); err != nil {
 				return err
 			}
 		}
@@ -1062,7 +1065,7 @@ func (e *Engine) Export(w io.Writer, basePath string, start time.Time, end time.
 		}
 	}()
 
-	return intar.Stream(w, path, basePath, e.timeStampFilterTarFile(start, end))
+	return intar.Stream(w, path, basePath, e.TarStreamBufferSize, e.timeStampFilterTarFile(start, end))
 }
 
 func (e *Engine) filterFileToBackup(r *TSMReader, fi os.FileInfo, shardRelativePath, fullPath string, start, end int64, tw *tar.Writer) error {
@@ -1117,7 +1120,7 @@ func (e *Engine) filterFileToBackup(r *TSMReader, fi os.FileInfo, shardRelativeP
 		return err
 	}
 
-	return intar.StreamRenameFile(tmpFi, fi.Name(), shardRelativePath, path, tw)
+	return intar.StreamRenameFile(tmpFi, fi.Name(), shardRelativePath, path, tw, e.TarStreamBufferSize)
 }
 
 // Restore reads a tar archive generated by Backup().