Adding DiskCache stats

Author: ola-rozenfeld

go/pkg/client/client.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"errors"
+
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/actas"
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/balancer"
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
@@ -244,6 +245,10 @@ func (c *Client) Close() error {
 	if c.casConnection != c.connection {
 		return c.casConnection.Close()
 	}
+	if c.diskCache != nil {
+		// Waits for local disk GC to complete.
+		c.diskCache.Shutdown()
+	}
 	return nil
 }
 
@@ -354,21 +359,12 @@ func (o *TreeSymlinkOpts) Apply(c *Client) {
 }
 
 type DiskCacheOpts struct {
-	Context       context.Context
-	Path          string
-	MaxCapacityGb float64
+	DiskCache *diskcache.DiskCache
 }
 
 // Apply sets the client's TreeSymlinkOpts.
 func (o *DiskCacheOpts) Apply(c *Client) {
-	if o.Path != "" {
-		capBytes := uint64(o.MaxCapacityGb * 1024 * 1024 * 1024)
-		var err error
-		// TODO(ola): propagate errors from Apply.
-		if c.diskCache, err = diskcache.New(o.Context, o.Path, capBytes); err != nil {
-			log.Errorf("Error initializing disk cache on %s: %v", o.Path, err)
-		}
-	}
+	c.diskCache = o.DiskCache
 }
 
 // MaxBatchDigests is maximum amount of digests to batch in upload and download operations.

go/pkg/diskcache/BUILD.bazel

@@ -28,7 +28,7 @@ go_test(
         "//go/pkg/testutil",
         "@com_github_bazelbuild_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
         "@com_github_google_go_cmp//cmp:go_default_library",
-        "@com_github_pborman_uuid//:go_default_library",
+        "@com_github_google_uuid//:uuid",
         "@org_golang_x_sync//errgroup:go_default_library",
     ],
 )

go/pkg/diskcache/diskcache.go

@@ -97,12 +97,27 @@ type DiskCache struct {
 	mu               sync.Mutex     // protects the queue.
 	store            sync.Map       // map of keys to qitems.
 	queue            *priorityQueue // keys by last accessed time.
-	sizeBytes        int64          // total size.
 	ctx              context.Context
 	shutdown         chan bool
+	shutdownOnce     sync.Once
 	gcTick           uint64
 	gcReq            chan uint64
-	testGcTicks      chan uint64
+	gcDone           chan bool
+	stats            *DiskCacheStats
+}
+
+type DiskCacheStats struct {
+	TotalSizeBytes         int64
+	TotalNumFiles          int64
+	NumFilesStored         int64
+	TotalStoredBytes       int64
+	NumFilesGCed           int64
+	TotalGCedSizeBytes     int64
+	NumCacheHits           int64
+	NumCacheMisses         int64
+	TotalCacheHitSizeBytes int64
+	InitTime               time.Duration
+	TotalGCTime            time.Duration
 }
 
 func New(ctx context.Context, root string, maxCapacityBytes uint64) (*DiskCache, error) {
@@ -115,7 +130,11 @@ func New(ctx context.Context, root string, maxCapacityBytes uint64) (*DiskCache,
 		},
 		gcReq:    make(chan uint64, maxConcurrentRequests),
 		shutdown: make(chan bool),
+		gcDone:   make(chan bool),
+		stats:    &DiskCacheStats{},
 	}
+	start := time.Now()
+	defer func() { atomic.AddInt64((*int64)(&res.stats.InitTime), int64(time.Since(start))) }()
 	heap.Init(res.queue)
 	if err := os.MkdirAll(root, os.ModePerm); err != nil {
 		return nil, err
@@ -158,7 +177,8 @@ func New(ctx context.Context, root string, maxCapacityBytes uint64) (*DiskCache,
 					return fmt.Errorf("error getting file size of %s: %v", path, err)
 				}
 				res.store.Store(k, it)
-				atomic.AddInt64(&res.sizeBytes, size)
+				atomic.AddInt64(&res.stats.TotalSizeBytes, size)
+				atomic.AddInt64(&res.stats.TotalNumFiles, 1)
 				res.mu.Lock()
 				heap.Push(res.queue, it)
 				res.mu.Unlock()
@@ -180,39 +200,47 @@ func (d *DiskCache) getItemSize(k key) (int64, error) {
 	fname := d.getPath(k)
 	info, err := os.Stat(fname)
 	if err != nil {
-		return 0, fmt.Errorf("Error getting info for %s: %v", fname, err)
+		return 0, fmt.Errorf("error getting info for %s: %v", fname, err)
 	}
 	return info.Size(), nil
 }
 
-// Releases resources and terminates the GC daemon. Should be the last call to the DiskCache.
+// Terminates the GC daemon, waiting for it to complete. No further Store* calls to the DiskCache should be made.
 func (d *DiskCache) Shutdown() {
-	d.shutdown <- true
+	d.shutdownOnce.Do(func() {
+		d.shutdown <- true
+		<-d.gcDone
+		log.Infof("DiskCacheStats: %+v", d.stats)
+	})
 }
 
-func (d *DiskCache) TotalSizeBytes() uint64 {
-	return uint64(atomic.LoadInt64(&d.sizeBytes))
-}
-
-// This function is defined in https://pkg.go.dev/strings#CutSuffix
-// It is copy/pasted here as a hack, because I failed to upgrade the *Reclient* repo to the latest Go 1.20.7.
-func CutSuffix(s, suffix string) (before string, found bool) {
-	if !strings.HasSuffix(s, suffix) {
-		return s, false
+func (d *DiskCache) GetStats() *DiskCacheStats {
+	// Return a copy for safety.
+	return &DiskCacheStats{
+		TotalSizeBytes:         atomic.LoadInt64(&d.stats.TotalSizeBytes),
+		TotalNumFiles:          atomic.LoadInt64(&d.stats.TotalNumFiles),
+		NumFilesStored:         atomic.LoadInt64(&d.stats.NumFilesStored),
+		TotalStoredBytes:       atomic.LoadInt64(&d.stats.TotalStoredBytes),
+		NumFilesGCed:           atomic.LoadInt64(&d.stats.NumFilesGCed),
+		TotalGCedSizeBytes:     atomic.LoadInt64(&d.stats.TotalGCedSizeBytes),
+		NumCacheHits:           atomic.LoadInt64(&d.stats.NumCacheHits),
+		NumCacheMisses:         atomic.LoadInt64(&d.stats.NumCacheMisses),
+		TotalCacheHitSizeBytes: atomic.LoadInt64(&d.stats.TotalCacheHitSizeBytes),
+		InitTime:               time.Duration(atomic.LoadInt64((*int64)(&d.stats.InitTime))),
+		TotalGCTime:            time.Duration(atomic.LoadInt64((*int64)(&d.stats.TotalGCTime))),
 	}
-	return s[:len(s)-len(suffix)], true
 }
 
 func (d *DiskCache) getKeyFromFileName(fname string) (key, error) {
 	pair := strings.Split(fname, ".")
 	if len(pair) != 2 {
-		return key{}, fmt.Errorf("expected file name in the form [ac_]hash/size, got %s", fname)
+		return key{}, fmt.Errorf("expected file name in the form hash[_ac].size, got %s", fname)
 	}
 	size, err := strconv.ParseInt(pair[1], 10, 64)
 	if err != nil {
 		return key{}, fmt.Errorf("invalid size in digest %s: %s", fname, err)
 	}
-	hash, isAc := CutSuffix(pair[0], "ac_")
+	hash, isAc := strings.CutSuffix(pair[0], "_ac")
 	dg, err := digest.New(hash, size)
 	if err != nil {
 		return key{}, fmt.Errorf("invalid digest from file name %s: %v", fname, err)
@@ -248,7 +276,10 @@ func (d *DiskCache) StoreCas(dg digest.Digest, path string) error {
 	if err := copyFile(path, d.getPath(it.key), dg.Size); err != nil {
 		return err
 	}
-	newSize := uint64(atomic.AddInt64(&d.sizeBytes, dg.Size))
+	newSize := uint64(atomic.AddInt64(&d.stats.TotalSizeBytes, dg.Size))
+	atomic.AddInt64(&d.stats.TotalNumFiles, 1)
+	atomic.AddInt64(&d.stats.NumFilesStored, 1)
+	atomic.AddInt64(&d.stats.TotalStoredBytes, dg.Size)
 	if newSize > d.maxCapacityBytes {
 		select {
 		case d.gcReq <- atomic.AddUint64(&d.gcTick, 1):
@@ -281,7 +312,10 @@ func (d *DiskCache) StoreActionCache(dg digest.Digest, ar *repb.ActionResult) er
 	if err := os.WriteFile(d.getPath(it.key), bytes, 0644); err != nil {
 		return err
 	}
-	newSize := uint64(atomic.AddInt64(&d.sizeBytes, int64(size)))
+	newSize := uint64(atomic.AddInt64(&d.stats.TotalSizeBytes, int64(size)))
+	atomic.AddInt64(&d.stats.TotalNumFiles, 1)
+	atomic.AddInt64(&d.stats.NumFilesStored, 1)
+	atomic.AddInt64(&d.stats.TotalStoredBytes, int64(size))
 	if newSize > d.maxCapacityBytes {
 		select {
 		case d.gcReq <- atomic.AddUint64(&d.gcTick, 1):
@@ -292,15 +326,17 @@ func (d *DiskCache) StoreActionCache(dg digest.Digest, ar *repb.ActionResult) er
 }
 
 func (d *DiskCache) gc() {
+	defer func() { d.gcDone <- true }()
 	for {
 		select {
 		case <-d.shutdown:
 			return
 		case <-d.ctx.Done():
 			return
-		case t := <-d.gcReq:
+		case <-d.gcReq:
+			start := time.Now()
 			// Evict old entries until total size is below cap.
-			for uint64(atomic.LoadInt64(&d.sizeBytes)) > d.maxCapacityBytes {
+			for uint64(atomic.LoadInt64(&d.stats.TotalSizeBytes)) > d.maxCapacityBytes {
 				d.mu.Lock()
 				it := heap.Pop(d.queue).(*qitem)
 				d.mu.Unlock()
@@ -309,7 +345,10 @@ func (d *DiskCache) gc() {
 					log.Errorf("error getting item size for %v: %v", it.key, err)
 					size = 0
 				}
-				atomic.AddInt64(&d.sizeBytes, -size)
+				atomic.AddInt64(&d.stats.TotalSizeBytes, -size)
+				atomic.AddInt64(&d.stats.TotalNumFiles, -1)
+				atomic.AddInt64(&d.stats.TotalGCedSizeBytes, size)
+				atomic.AddInt64(&d.stats.NumFilesGCed, 1)
 				it.mu.Lock()
 				// We only delete the files, and not the prefix directories, because the prefixes are not worth worrying about.
 				if err := os.Remove(d.getPath(it.key)); err != nil {
@@ -318,12 +357,7 @@ func (d *DiskCache) gc() {
 				d.store.Delete(it.key)
 				it.mu.Unlock()
 			}
-			if d.testGcTicks != nil {
-				select {
-				case d.testGcTicks <- t:
-				default:
-				}
-			}
+			atomic.AddInt64((*int64)(&d.stats.TotalGCTime), int64(time.Since(start)))
 		}
 	}
 }
@@ -363,6 +397,7 @@ func (d *DiskCache) LoadCas(dg digest.Digest, path string) bool {
 	k := key{digest: dg, isCas: true}
 	iUntyped, loaded := d.store.Load(k)
 	if !loaded {
+		atomic.AddInt64(&d.stats.NumCacheMisses, 1)
 		return false
 	}
 	it := iUntyped.(*qitem)
@@ -371,52 +406,61 @@ func (d *DiskCache) LoadCas(dg digest.Digest, path string) bool {
 	it.mu.RUnlock()
 	if err != nil {
 		// It is not possible to prevent a race with GC; hence, we return false on copy errors.
+		atomic.AddInt64(&d.stats.NumCacheMisses, 1)
 		return false
 	}
 
 	d.mu.Lock()
 	d.queue.Bump(it)
 	d.mu.Unlock()
+	atomic.AddInt64(&d.stats.NumCacheHits, 1)
+	atomic.AddInt64(&d.stats.TotalCacheHitSizeBytes, dg.Size)
 	return true
 }
 
 func (d *DiskCache) LoadActionCache(dg digest.Digest) (ar *repb.ActionResult, loaded bool) {
 	k := key{digest: dg, isCas: false}
 	iUntyped, loaded := d.store.Load(k)
 	if !loaded {
+		atomic.AddInt64(&d.stats.NumCacheMisses, 1)
 		return nil, false
 	}
 	it := iUntyped.(*qitem)
 	it.mu.RLock()
 	ar = &repb.ActionResult{}
-	if err := d.loadActionResult(k, ar); err != nil {
+	size, err := d.loadActionResult(k, ar)
+	if err != nil {
 		// It is not possible to prevent a race with GC; hence, we return false on load errors.
 		it.mu.RUnlock()
+		atomic.AddInt64(&d.stats.NumCacheMisses, 1)
 		return nil, false
 	}
 	it.mu.RUnlock()
 
 	d.mu.Lock()
 	d.queue.Bump(it)
 	d.mu.Unlock()
+	atomic.AddInt64(&d.stats.NumCacheHits, 1)
+	atomic.AddInt64(&d.stats.TotalCacheHitSizeBytes, int64(size))
 	return ar, true
 }
 
-func (d *DiskCache) loadActionResult(k key, ar *repb.ActionResult) error {
+func (d *DiskCache) loadActionResult(k key, ar *repb.ActionResult) (int, error) {
 	bytes, err := os.ReadFile(d.getPath(k))
 	if err != nil {
-		return err
+		return 0, err
 	}
+	n := len(bytes)
 	// Required sanity check: sometimes the read pretends to succeed, but doesn't, if
 	// the file is being concurrently deleted. Empty ActionResult is advised against in
 	// the RE-API: https://github.com/bazelbuild/remote-apis/blob/main/build/bazel/remote/execution/v2/remote_execution.proto#L1052
-	if len(bytes) == 0 {
-		return fmt.Errorf("read empty ActionResult for %v", k.digest)
+	if n == 0 {
+		return n, fmt.Errorf("read empty ActionResult for %v", k.digest)
 	}
 	if err := proto.Unmarshal(bytes, ar); err != nil {
-		return fmt.Errorf("error unmarshalling %v as ActionResult: %v", bytes, err)
+		return n, fmt.Errorf("error unmarshalling %v as ActionResult: %v", bytes, err)
 	}
-	return nil
+	return n, nil
 }
 
 func getLastAccessTime(path string) (time.Time, error) {