github.com/benchkram/bob@v0.0.0-20240314204020-b7a57f2f9be9/pkg/nix/shell_cache.go (about)

     1  package nix
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/hex"
     6  	"fmt"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  
    11  	"github.com/benchkram/bob/bobtask/hash"
    12  	"github.com/benchkram/bob/pkg/file"
    13  	"github.com/benchkram/bob/pkg/filehash"
    14  	"github.com/benchkram/errz"
    15  )
    16  
    17  // ShellCache caches the output of nix-shell command
    18  type ShellCache struct {
    19  	// dir is the root directory where the cache files are stored
    20  	dir string
    21  	// dataCache caches the file contents, so we don't need to call os.ReadFile on second read of the same file
    22  	dataCache map[string][]byte
    23  }
    24  
    25  // NewShellCache creates a new instance of ShellCache
    26  func NewShellCache(dir string) *ShellCache {
    27  	dataCache := make(map[string][]byte)
    28  
    29  	return &ShellCache{dir, dataCache}
    30  }
    31  
    32  // Save caches the output inside a file named by the key cache
    33  func (c *ShellCache) Save(key string, output []byte) (err error) {
    34  	defer errz.Recover(&err)
    35  
    36  	err = os.MkdirAll(c.dir, 0775)
    37  	errz.Fatal(err)
    38  
    39  	err = os.WriteFile(filepath.Join(c.dir, key), output, 0644)
    40  	errz.Fatal(err)
    41  
    42  	return nil
    43  }
    44  
    45  // Get the data by cache key
    46  // If Reading the file returns an error, empty data is returned
    47  func (c *ShellCache) Get(key string) ([]byte, bool) {
    48  	if i, ok := c.dataCache[filepath.Join(c.dir, key)]; ok {
    49  		return i, true
    50  	}
    51  
    52  	if !file.Exists(filepath.Join(c.dir, key)) {
    53  		return []byte{}, false
    54  	}
    55  	data, err := os.ReadFile(filepath.Join(c.dir, key))
    56  	if err != nil {
    57  		return []byte{}, false
    58  	}
    59  
    60  	c.dataCache[filepath.Join(c.dir, key)] = data
    61  	return data, true
    62  }
    63  
    64  // GenerateKey generates key for the cache based on a list of Dependency and nix-shell command
    65  //
    66  // if Dependency it's a .nix file it will hash the nixpkgs + file contents
    67  // if Dependency it's a package name will hash the packageName:nixpkgs content
    68  func (c *ShellCache) GenerateKey(deps []Dependency, nixShellCmd string) (_ string, err error) {
    69  	defer errz.Recover(&err)
    70  	h := filehash.New()
    71  
    72  	for _, dependency := range deps {
    73  		if strings.HasSuffix(dependency.Name, ".nix") {
    74  			err = h.AddBytes(bytes.NewBufferString(dependency.Nixpkgs))
    75  			errz.Fatal(err)
    76  
    77  			err = h.AddFile(dependency.Name)
    78  			errz.Fatal(err)
    79  		} else {
    80  			toHash := fmt.Sprintf("%s:%s", dependency.Name, dependency.Nixpkgs)
    81  			err = h.AddBytes(bytes.NewBufferString(toHash))
    82  			errz.Fatal(err)
    83  		}
    84  	}
    85  
    86  	err = h.AddBytes(bytes.NewBufferString(nixShellCmd))
    87  	errz.Fatal(err)
    88  
    89  	hashIn := hash.In(hex.EncodeToString(h.Sum()))
    90  
    91  	return hashIn.String(), nil
    92  }
    93  
    94  // Clean will clean entire shell env cache
    95  func (c *ShellCache) Clean() (err error) {
    96  	defer errz.Recover(&err)
    97  
    98  	if !file.Exists(c.dir) {
    99  		return nil
   100  	}
   101  
   102  	entries, err := os.ReadDir(c.dir)
   103  	errz.Fatal(err)
   104  
   105  	for _, entry := range entries {
   106  		err = os.Remove(filepath.Join(c.dir, entry.Name()))
   107  		errz.Fatal(err)
   108  	}
   109  	return nil
   110  }