github.com/bazelbuild/remote-apis-sdks@v0.0.0-20240425170053-8a36686a6350/go/pkg/fakes/ac.go (about)

     1  // Package fakes contains configurable test fakes for an RE service.
     2  package fakes
     3  
     4  import (
     5  	"context"
     6  	"fmt"
     7  	"sync"
     8  
     9  	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
    10  	"google.golang.org/grpc/codes"
    11  	"google.golang.org/grpc/status"
    12  
    13  	repb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
    14  )
    15  
    16  // ActionCache implements the RE ActionCache interface, storing fixed results.
    17  type ActionCache struct {
    18  	mu      sync.RWMutex
    19  	results map[digest.Digest]*repb.ActionResult
    20  	reads   map[digest.Digest]int
    21  	writes  map[digest.Digest]int
    22  }
    23  
    24  // NewActionCache returns a new empty ActionCache.
    25  func NewActionCache() *ActionCache {
    26  	c := &ActionCache{}
    27  	c.Clear()
    28  	return c
    29  }
    30  
    31  // Clear removes all results from the cache.
    32  func (c *ActionCache) Clear() {
    33  	c.mu.Lock()
    34  	defer c.mu.Unlock()
    35  	c.results = make(map[digest.Digest]*repb.ActionResult)
    36  	c.reads = make(map[digest.Digest]int)
    37  	c.writes = make(map[digest.Digest]int)
    38  }
    39  
    40  // PutAction sets a fake result for a given action, and returns the action digest.
    41  func (c *ActionCache) PutAction(ac *repb.Action, res *repb.ActionResult) digest.Digest {
    42  	d := digest.TestNewFromMessage(ac)
    43  	c.Put(d, res)
    44  	return d
    45  }
    46  
    47  // Put sets a fake result for a given action digest.
    48  func (c *ActionCache) Put(d digest.Digest, res *repb.ActionResult) {
    49  	c.mu.Lock()
    50  	defer c.mu.Unlock()
    51  	c.results[d] = res
    52  }
    53  
    54  // Get returns a previously saved fake result for the given action digest.
    55  func (c *ActionCache) Get(d digest.Digest) *repb.ActionResult {
    56  	c.mu.RLock()
    57  	defer c.mu.RUnlock()
    58  	res, ok := c.results[d]
    59  	if !ok {
    60  		return nil
    61  	}
    62  	return res
    63  }
    64  
    65  // Reads returns the number of times GetActionResult was called for a given action digest.
    66  // These include both successful and unsuccessful reads.
    67  func (c *ActionCache) Reads(d digest.Digest) int {
    68  	return c.reads[d]
    69  }
    70  
    71  // Writes returns the number of times UpdateActionResult was called for a given action digest.
    72  func (c *ActionCache) Writes(d digest.Digest) int {
    73  	return c.writes[d]
    74  }
    75  
    76  // GetActionResult returns a stored result, if it was found.
    77  func (c *ActionCache) GetActionResult(ctx context.Context, req *repb.GetActionResultRequest) (res *repb.ActionResult, err error) {
    78  	c.mu.Lock()
    79  	defer c.mu.Unlock()
    80  	dg, err := digest.NewFromProto(req.ActionDigest)
    81  	if err != nil {
    82  		return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("invalid digest received: %v", req.ActionDigest))
    83  	}
    84  	c.reads[dg]++
    85  	if res, ok := c.results[dg]; ok {
    86  		return res, nil
    87  	}
    88  	return nil, status.Error(codes.NotFound, "")
    89  }
    90  
    91  // UpdateActionResult sets/updates a given result.
    92  func (c *ActionCache) UpdateActionResult(ctx context.Context, req *repb.UpdateActionResultRequest) (res *repb.ActionResult, err error) {
    93  	c.mu.Lock()
    94  	defer c.mu.Unlock()
    95  	dg, err := digest.NewFromProto(req.ActionDigest)
    96  	if err != nil {
    97  		return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("invalid digest received: %v", req.ActionDigest))
    98  	}
    99  	if req.ActionResult == nil {
   100  		return nil, status.Error(codes.InvalidArgument, "no action result received")
   101  	}
   102  	c.results[dg] = req.ActionResult
   103  	c.writes[dg]++
   104  	return req.ActionResult, nil
   105  }