github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/exec/bigmachine_inv_disk_cache.go (about)

     1  package exec
     2  
     3  import (
     4  	"io"
     5  	"io/ioutil"
     6  	"os"
     7  	"path"
     8  	"strconv"
     9  	"sync"
    10  
    11  	"github.com/grailbio/base/compress/zstd"
    12  	"github.com/grailbio/base/errors"
    13  	"github.com/grailbio/base/fileio"
    14  	"github.com/grailbio/base/log"
    15  	"github.com/grailbio/base/must"
    16  )
    17  
    18  type invDiskCache struct {
    19  	mu       sync.Mutex
    20  	cacheDir string
    21  	// invPaths maps invocation indices to cache files on disk.
    22  	invPaths map[uint64]string
    23  }
    24  
    25  func newInvDiskCache() *invDiskCache {
    26  	return &invDiskCache{invPaths: make(map[uint64]string)}
    27  }
    28  
    29  // REQUIRES: Caller holds c.mu.
    30  func (c *invDiskCache) init() error {
    31  	if c.cacheDir != "" {
    32  		return nil
    33  	}
    34  	cacheDir, err := ioutil.TempDir("", "bigslice-inv-cache")
    35  	if err != nil {
    36  		return errors.E(err, "bigslice: could not create invocation disk cache")
    37  	}
    38  	c.cacheDir = cacheDir
    39  	return nil
    40  }
    41  
    42  func (c *invDiskCache) close() {
    43  	c.mu.Lock()
    44  	defer c.mu.Unlock()
    45  
    46  	must.Truef(c.invPaths != nil, "multiple close")
    47  
    48  	if err := os.RemoveAll(c.cacheDir); err != nil {
    49  		log.Printf("WARNING: error discarding bigslice invocation disk cache: %v", err)
    50  	}
    51  	c.invPaths = nil
    52  }
    53  
    54  func (c *invDiskCache) getOrCreate(invIndex uint64, create func(io.Writer) error) (io.ReadCloser, error) {
    55  	c.mu.Lock() // Note: cache access is serialized.
    56  	defer c.mu.Unlock()
    57  
    58  	must.Truef(c.invPaths != nil, "call after close")
    59  	if err := c.init(); err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	if _, ok := c.invPaths[invIndex]; !ok {
    64  		err := func() (err error) {
    65  			invPath := path.Join(c.cacheDir, strconv.Itoa(int(invIndex)))
    66  			f, err := os.OpenFile(invPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
    67  			if err != nil {
    68  				return err
    69  			}
    70  			defer fileio.CloseAndReport(f, &err)
    71  			zw, err := zstd.NewWriter(f)
    72  			if err != nil {
    73  				return err
    74  			}
    75  			defer fileio.CloseAndReport(zw, &err)
    76  			if err = create(zw); err != nil {
    77  				return err
    78  			}
    79  			c.invPaths[invIndex] = invPath
    80  			return nil
    81  		}()
    82  		if err != nil {
    83  			return nil, errors.E(err, "bigslice: could not create invocation disk cache entry")
    84  		}
    85  	}
    86  	f, err := os.Open(c.invPaths[invIndex])
    87  	if err != nil {
    88  		return nil, errors.E(err, "bigslice: could not open invocation disk cache entry")
    89  	}
    90  	zr, err := zstd.NewReader(f)
    91  	if err != nil {
    92  		err = errors.E(err, "bigslice: could not open (zstd) invocation disk cache entry")
    93  		fileio.CloseAndReport(f, &err)
    94  		return nil, err
    95  	}
    96  	return fileReadCloser{ReadCloser: zr, file: f}, nil
    97  }
    98  
    99  type fileReadCloser struct {
   100  	io.ReadCloser
   101  	file *os.File
   102  }
   103  
   104  func (f fileReadCloser) Close() (err error) {
   105  	fileio.CloseAndReport(f.ReadCloser, &err)
   106  	fileio.CloseAndReport(f.file, &err)
   107  	return
   108  }