github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/internal/slicecache/sliceio.go (about)

     1  package slicecache
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/grailbio/base/backgroundcontext"
     8  	"github.com/grailbio/base/compress/zstd"
     9  	"github.com/grailbio/base/errors"
    10  	"github.com/grailbio/base/file"
    11  	"github.com/grailbio/base/log"
    12  	"github.com/grailbio/bigslice/frame"
    13  	"github.com/grailbio/bigslice/sliceio"
    14  )
    15  
    16  type fileReader struct {
    17  	sliceio.Reader
    18  	file file.File
    19  	zr   io.ReadCloser
    20  	path string
    21  }
    22  
    23  func (f *fileReader) Read(ctx context.Context, frame frame.Frame) (int, error) {
    24  	if f.file == nil {
    25  		var err error
    26  		f.file, err = file.Open(ctx, f.path)
    27  		if err != nil {
    28  			return 0, err
    29  		}
    30  		f.zr, err = zstd.NewReader(f.file.Reader(backgroundcontext.Get()))
    31  		if err != nil {
    32  			return 0, err
    33  		}
    34  		f.Reader = sliceio.NewDecodingReader(f.zr)
    35  	}
    36  	n, err := f.Reader.Read(ctx, frame)
    37  	if err != nil {
    38  		if closeErr := f.zr.Close(); closeErr != nil {
    39  			log.Error.Printf("%s: close zstd reader: %v", f.file.Name(), closeErr)
    40  		}
    41  		if closeErr := f.file.Close(ctx); closeErr != nil {
    42  			log.Error.Printf("%s: close file: %v", f.file.Name(), closeErr)
    43  		}
    44  	}
    45  	return n, err
    46  }
    47  
    48  func newFileReader(path string) sliceio.Reader {
    49  	return &fileReader{path: path}
    50  }
    51  
    52  type writethroughReader struct {
    53  	sliceio.Reader
    54  	path string
    55  	file file.File
    56  	zw   io.WriteCloser
    57  	enc  *sliceio.Encoder
    58  }
    59  
    60  func (r *writethroughReader) Read(ctx context.Context, frame frame.Frame) (int, error) {
    61  	if r.file == nil {
    62  		var err error
    63  		r.file, err = file.Create(ctx, r.path)
    64  		if err != nil {
    65  			return 0, err
    66  		}
    67  		// Ideally we'd use the underlying context for each op here,
    68  		// but the way encoder is set up, we can't (understandably)
    69  		// pass a new writer for each encode.
    70  		r.zw, err = zstd.NewWriter(r.file.Writer(backgroundcontext.Get()))
    71  		if err != nil {
    72  			return 0, err
    73  		}
    74  		r.enc = sliceio.NewEncodingWriter(r.zw)
    75  	}
    76  	n, err := r.Reader.Read(ctx, frame)
    77  	if err == nil || err == sliceio.EOF {
    78  		if writeErr := r.enc.Write(ctx, frame.Slice(0, n)); writeErr != nil {
    79  			return n, writeErr
    80  		}
    81  		if err == sliceio.EOF {
    82  			closeErr := r.zw.Close()
    83  			errors.CleanUpCtx(ctx, r.file.Close, &closeErr)
    84  			if closeErr != nil {
    85  				return n, closeErr
    86  			}
    87  		}
    88  	} else {
    89  		r.file.Discard(backgroundcontext.Get())
    90  	}
    91  	return n, err
    92  }
    93  
    94  func newWritethroughReader(reader sliceio.Reader, path string) sliceio.Reader {
    95  	return &writethroughReader{Reader: reader, path: path}
    96  }