github.com/pachyderm/pachyderm@v1.13.4/src/server/worker/cache/cache.go (about)

     1  package cache
     2  
     3  import (
     4  	"os"
     5  	"path/filepath"
     6  	"sync"
     7  
     8  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
     9  	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
    10  )
    11  
    12  // WorkerCache is an interface for managing hashtree MergeCaches for multiple
    13  // concurrent jobs
    14  type WorkerCache interface {
    15  	GetOrCreateCache(jobID string) (*hashtree.MergeCache, error)
    16  	GetCache(jobID string) *hashtree.MergeCache
    17  	RemoveCache(jobID string) error
    18  }
    19  
    20  type workerCache struct {
    21  	mutex sync.Mutex
    22  
    23  	// hashtreeStorage is the where we store on disk hashtrees
    24  	hashtreeStorage string
    25  
    26  	caches map[string]*hashtree.MergeCache
    27  }
    28  
    29  // NewWorkerCache constructs a WorkerCache for maintaining hashtree caches for
    30  // multiple concurrent jobs
    31  func NewWorkerCache(hashtreeStorage string) WorkerCache {
    32  	return &workerCache{
    33  		hashtreeStorage: hashtreeStorage,
    34  		caches:          make(map[string]*hashtree.MergeCache),
    35  	}
    36  }
    37  
    38  func (wc *workerCache) GetOrCreateCache(jobID string) (*hashtree.MergeCache, error) {
    39  	wc.mutex.Lock()
    40  	defer wc.mutex.Unlock()
    41  
    42  	if cache, ok := wc.caches[jobID]; ok {
    43  		return cache, nil
    44  	}
    45  
    46  	newCache, err := hashtree.NewMergeCache(filepath.Join(wc.hashtreeStorage, jobID))
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  
    51  	wc.caches[jobID] = newCache
    52  	return newCache, nil
    53  }
    54  
    55  func (wc *workerCache) RemoveCache(jobID string) error {
    56  	wc.mutex.Lock()
    57  	defer wc.mutex.Unlock()
    58  
    59  	if _, ok := wc.caches[jobID]; ok {
    60  		if err := os.RemoveAll(filepath.Join(wc.hashtreeStorage, jobID)); err != nil {
    61  			return errors.EnsureStack(err)
    62  		}
    63  		delete(wc.caches, jobID)
    64  	}
    65  	return nil
    66  }
    67  
    68  func (wc *workerCache) GetCache(jobID string) *hashtree.MergeCache {
    69  	wc.mutex.Lock()
    70  	defer wc.mutex.Unlock()
    71  
    72  	return wc.caches[jobID]
    73  }
    74  
    75  func (wc *workerCache) Close() error {
    76  	wc.mutex.Lock()
    77  	defer wc.mutex.Unlock()
    78  
    79  	wc.caches = make(map[string]*hashtree.MergeCache)
    80  	return errors.EnsureStack(os.RemoveAll(wc.hashtreeStorage))
    81  }