github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/fileutil/writer_reader.go (about)

     1  package fileutil
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"os"
     7  )
     8  
     9  // WriterThenReader writes data to storage, then allows reading it.  It is suitable for
    10  // repeatedly processing large volumes of data.
    11  type WriterThenReader interface {
    12  	io.Writer
    13  	// StartReading stops writing and returns a RewindableReader that will allow repeatedly
    14  	// reading the data and the total length of data.  The WriterThenReader should be
    15  	// discarded; calls to Write() after StartReading() will fail.
    16  	StartReading() (RewindableReader, int64, error)
    17  	// Name returns a user-visible name for underlying storage.  It may help debug some issues.
    18  	Name() string
    19  }
    20  
    21  // RewindableReader allows repeatedly reading the same stream.
    22  type RewindableReader interface {
    23  	io.ReadSeeker
    24  	// Rewind allows sets RewindableReader to start re-reading the same data.
    25  	Rewind() error
    26  	// Name returns a user-visible name for underlying storage.  It may help debug some issues.
    27  	Name() string
    28  }
    29  
    30  type fileWriterThenReader struct{ file *os.File }
    31  
    32  func NewFileWriterThenReader(basename string) (WriterThenReader, error) {
    33  	file, err := os.CreateTemp("", basename)
    34  	if err != nil {
    35  		return nil, fmt.Errorf("creating temporary file: %w", err)
    36  	}
    37  	if err = os.Remove(file.Name()); err != nil {
    38  		return nil, fmt.Errorf("removing file %s from directory: %w", file.Name(), err)
    39  	}
    40  	return &fileWriterThenReader{file: file}, nil
    41  }
    42  
    43  func (f fileWriterThenReader) Write(p []byte) (int, error) {
    44  	return f.file.Write(p)
    45  }
    46  
    47  func (f *fileWriterThenReader) StartReading() (RewindableReader, int64, error) {
    48  	offset, err := f.file.Seek(0, io.SeekEnd)
    49  	if err != nil {
    50  		return nil, -1, err
    51  	}
    52  	_, err = f.file.Seek(0, io.SeekStart)
    53  	ret := &fileRewindableReader{file: f.file}
    54  	// Break future attempts to use f
    55  	f.file = nil
    56  	return ret, offset, err
    57  }
    58  
    59  func (f fileWriterThenReader) Name() string { return f.file.Name() }
    60  
    61  type fileRewindableReader struct{ file *os.File }
    62  
    63  func (f fileRewindableReader) Read(p []byte) (n int, err error) {
    64  	return f.file.Read(p)
    65  }
    66  
    67  func (f fileRewindableReader) Seek(offset int64, whence int) (int64, error) {
    68  	return f.file.Seek(offset, whence)
    69  }
    70  
    71  func (f fileRewindableReader) Rewind() error {
    72  	_, err := f.file.Seek(0, io.SeekStart)
    73  	return err
    74  }
    75  
    76  func (f fileRewindableReader) Name() string { return f.file.Name() }