Add FileEntry cache to TagCache for selfupdate.

cipd/client/cipd/internal/tagcache.go

Index: cipd/client/cipd/internal/tagcache.go
diff --git a/cipd/client/cipd/internal/tagcache.go b/cipd/client/cipd/internal/tagcache.go
index cc9556195347a79d2292902b275fc702f570f61f..750142f3563b745786251fcfe16697ba96dda974 100644
--- a/cipd/client/cipd/internal/tagcache.go
+++ b/cipd/client/cipd/internal/tagcache.go
@@ -9,6 +9,7 @@ import (
 	"io/ioutil"
 	"os"
 	"sort"
+	"strings"
 	"sync"
 
 	"golang.org/x/net/context"
@@ -20,8 +21,9 @@ import (
 )
 
 const (
-	tagCacheMaxSize  = 300
-	tagCacheFilename = "tagcache.db"
+	tagCacheMaxSize    = 300
+	tagCacheMaxExeSize = 20
+	tagCacheFilename   = "tagcache.db"
 )
 
 // TagCache provides a mapping (package name, tag) -> instance ID.
@@ -34,20 +36,28 @@ const (
 // cache is to avoid round trips to the service to increase reliability of
 // 'cipd ensure' calls that use only tags to specify versions. It happens to be
 // the most common case of 'cipd ensure' usage by far.
+//
+// Additionally, this TagCache stores a mapping of (pin, file_name) ->
+// hex(sha1(executable)) to assist in the `selfupdate` flow. Whenever selfupdate
+// resolves what instance ID it SHOULD be at, it checks the SHA1 of its own
+// binary to see if it's actually already at that instance ID (since instance ID
+// is of the whole package, and not of the actual executable inside the
+// package).
 type TagCache struct {
 	fs   local.FileSystem
 	lock sync.Mutex
 
-	cache *messages.TagCache                  // the last loaded state, if not nil.
-	added map[string]*messages.TagCache_Entry // entries added by AddTag
+	cache      *messages.TagCache                      // the last loaded state, if not nil.
+	addedTags  map[string]*messages.TagCache_Entry     // entries added by AddTag
+	addedFiles map[string]*messages.TagCache_FileEntry // entries added by AddFile
 }
 
-// tagMapKey constructs a key for TagCache.added map.
+// stringKey constructs keys for the TagCache.added* maps.
 //
 // We use string math instead of a struct to avoid reimplementing sorting
 // interface.
-func tagMapKey(pkgName, tag string) string {
-	return pkgName + ":" + tag
+func stringKey(strs ...string) string {
+	return strings.Join(strs, ":")
 }
 
 // NewTagCache initializes TagCache.
@@ -57,39 +67,43 @@ func NewTagCache(fs local.FileSystem) *TagCache {
 	return &TagCache{fs: fs}
 }
 
+func (c *TagCache) lazyLoadLocked(ctx context.Context) (err error) {
+	// Lazy-load the cache the first time it is used. We reload it again in Save
+	// right before overwriting the file. We don't reload it anywhere else though,
+	// so the implementation essentially assumes Save is called relatively soon
+	// after ResolveTag call. If it's not the case, cache updates made by other
+	// processes will be "invisible" to the TagCache. It still tries not to
+	// overwrite them in Save, so it's fine.
+	if c.cache == nil {
+		c.cache, err = c.loadFromDisk(ctx)
+	}
+	return
+}
+
 // ResolveTag returns cached tag or empty Pin{} if such tag is not in the cache.
 //
 // Returns error if the cache can't be read.
-func (c *TagCache) ResolveTag(ctx context.Context, pkg, tag string) (common.Pin, error) {
-	if err := common.ValidatePackageName(pkg); err != nil {
-		return common.Pin{}, err
+func (c *TagCache) ResolveTag(ctx context.Context, pkg, tag string) (pin common.Pin, err error) {
+	if err = common.ValidatePackageName(pkg); err != nil {
+		return
 	}
-	if err := common.ValidateInstanceTag(tag); err != nil {
-		return common.Pin{}, err
+	if err = common.ValidateInstanceTag(tag); err != nil {
+		return
 	}
 
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
 	// Already added with AddTag recently?
-	if e := c.added[tagMapKey(pkg, tag)]; e != nil {
+	if e := c.addedTags[stringKey(pkg, tag)]; e != nil {
 		return common.Pin{
 			PackageName: e.Package,
 			InstanceID:  e.InstanceId,
 		}, nil
 	}
 
-	// Lazy-load the cache the first time it is used. We reload it again in Save
-	// right before overwriting the file. We don't reload it anywhere else though,
-	// so the implementation essentially assumes Save is called relatively soon
-	// after ResolveTag call. If it's not the case, cache updates made by other
-	// processes will be "invisible" to the TagCache. It still tries not to
-	// overwrite them in Save, so it's fine.
-	if c.cache == nil {
-		var err error
-		if c.cache, err = c.loadFromDisk(ctx); err != nil {
-			return common.Pin{}, err
-		}
+	if err = c.lazyLoadLocked(ctx); err != nil {
+		return
 	}
 
 	// Most recently used tags are usually at the end, search in reverse as a
@@ -104,17 +118,46 @@ func (c *TagCache) ResolveTag(ctx context.Context, pkg, tag string) (common.Pin,
 		}
 	}
 
-	return common.Pin{}, nil
+	return
+}
+
+// ResolveFile returns file hash or "" if that file is not in the cache.
+//
+// Returns error if the cache can't be read.
+func (c *TagCache) ResolveFile(ctx context.Context, pin common.Pin, fileName string) (hash string, err error) {
+	if err = common.ValidatePin(pin); err != nil {
+		return
+	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	key := stringKey(pin.PackageName, pin.InstanceID, fileName)
+	if e := c.addedFiles[key]; e != nil {
+		return e.Hash, nil
+	}
+
+	if err = c.lazyLoadLocked(ctx); err != nil {
+		return
+	}
+
+	// Most recently used tags are usually at the end, search in reverse as a
+	// silly optimization.
+	for i := len(c.cache.FileEntries) - 1; i >= 0; i-- {
+		e := c.cache.FileEntries[i]
+		if e.Package == pin.PackageName && e.InstanceId == pin.InstanceID && e.FileName == fileName {
+			hash = e.Hash
+			return
+		}
+	}
+	return
 }
 
 // AddTag records that (pin.PackageName, tag) maps to pin.InstanceID.
 //
 // Call 'Save' later to persist these changes to the cache file on disk.
 func (c *TagCache) AddTag(ctx context.Context, pin common.Pin, tag string) error {
-	if err := common.ValidatePackageName(pin.PackageName); err != nil {
-		return err
-	}
-	if err := common.ValidateInstanceID(pin.InstanceID); err != nil {
+	if err := common.ValidatePin(pin); err != nil {
 		return err
 	}
 	if err := common.ValidateInstanceTag(tag); err != nil {
@@ -124,12 +167,12 @@ func (c *TagCache) AddTag(ctx context.Context, pin common.Pin, tag string) error
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
-	if c.added == nil {
-		c.added = make(map[string]*messages.TagCache_Entry, 1)
+	if c.addedTags == nil {
+		c.addedTags = make(map[string]*messages.TagCache_Entry, 1)
 	}
 
 	// 'Save' will merge this into 'c.cache' before dumping to disk.
-	c.added[tagMapKey(pin.PackageName, tag)] = &messages.TagCache_Entry{
+	c.addedTags[stringKey(pin.PackageName, tag)] = &messages.TagCache_Entry{
 		Package:    pin.PackageName,
 		Tag:        tag,
 		InstanceId: pin.InstanceID,
@@ -138,6 +181,33 @@ func (c *TagCache) AddTag(ctx context.Context, pin common.Pin, tag string) error
 	return nil
 }
 
+// AddFile records that (pin, fileName) maps to hash (where hash is the
+// hex-encoded sha1 of the named file).
+//
+// Call 'Save' later to persist these changes to the cache file on disk.
+func (c *TagCache) AddFile(ctx context.Context, pin common.Pin, fileName, hash string) error {
+	if err := common.ValidatePin(pin); err != nil {
+		return err
+	}
+	if err := common.ValidateFileHash(hash); err != nil {
+		return err
+	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	if c.addedFiles == nil {
+		c.addedFiles = make(map[string]*messages.TagCache_FileEntry, 1)
+	}
+	c.addedFiles[stringKey(pin.PackageName, pin.InstanceID, fileName)] = &messages.TagCache_FileEntry{
+		Package:    pin.PackageName,
+		InstanceId: pin.InstanceID,
+		FileName:   fileName,
+		Hash:       hash,
+	}
+	return nil
+}
+
 // Save stores all pending cache updates to the file system.
 //
 // It effectively resets the object to the initial state.
@@ -147,17 +217,22 @@ func (c *TagCache) Save(ctx context.Context) error {
 
 	// Nothing to store? Just clean the state, so that ResolveTag can fetch the
 	// up-to-date cache from disk later if needed.
-	if len(c.added) == 0 {
+	if len(c.addedTags) == 0 && len(c.addedFiles) == 0 {
 		c.cache = nil
 		return nil
 	}
 
 	// Sort all new entries, for consistency.
-	sorted := make([]string, 0, len(c.added))
-	for k := range c.added {
-		sorted = append(sorted, k)
+	sortedTags := make([]string, 0, len(c.addedTags))
+	for k := range c.addedTags {
+		sortedTags = append(sortedTags, k)
+	}
+	sort.Strings(sortedTags)
+	sortedFiles := make([]string, 0, len(c.addedFiles))
+	for k := range c.addedFiles {
+		sortedFiles = append(sortedFiles, k)
 	}
-	sort.Strings(sorted)
+	sort.Strings(sortedFiles)
 
 	// Load the most recent data to avoid overwriting it.
 	recent, err := c.loadFromDisk(ctx)
@@ -166,34 +241,49 @@ func (c *TagCache) Save(ctx context.Context) error {
 	}
 
 	// Copy all existing entries, except the ones we are moving to the tail.
-	merged := make([]*messages.TagCache_Entry, 0, len(recent.Entries)+len(c.added))
+	mergedTags := make([]*messages.TagCache_Entry, 0, len(recent.Entries)+len(c.addedTags))
 	for _, e := range recent.Entries {
-		if c.added[tagMapKey(e.Package, e.Tag)] == nil {
-			merged = append(merged, e)
+		if c.addedTags[stringKey(e.Package, e.Tag)] == nil {
+			mergedTags = append(mergedTags, e)
 		}
 	}
 
 	// Add new entries to the tail.
-	for _, k := range sorted {
-		merged = append(merged, c.added[k])
+	for _, k := range sortedTags {
+		mergedTags = append(mergedTags, c.addedTags[k])
 	}
 
 	// Trim the end result, discard the head: it's where old items are.
-	if len(merged) > tagCacheMaxSize {
-		merged = merged[len(merged)-tagCacheMaxSize:]
+	if len(mergedTags) > tagCacheMaxSize {
+		mergedTags = mergedTags[len(mergedTags)-tagCacheMaxSize:]
+	}
+
+	// Do the same for file entries
+	mergedFiles := make([]*messages.TagCache_FileEntry, 0, len(recent.FileEntries)+len(c.addedFiles))
+	for _, e := range recent.FileEntries {
+		if c.addedFiles[e.Hash] == nil {
+			mergedFiles = append(mergedFiles, e)
+		}
+	}
+	for _, k := range sortedFiles {
+		mergedFiles = append(mergedFiles, c.addedFiles[k])
+	}
+	if len(mergedFiles) > tagCacheMaxExeSize {
+		mergedFiles = mergedFiles[len(mergedFiles)-tagCacheMaxExeSize:]
 	}
 
 	// Serialize and write to disk. We still can accidentally replace someone
 	// else's changes, but the probability should be relatively low. It can happen
 	// only if two processes call 'Save' at the exact same time.
-	updated := &messages.TagCache{Entries: merged}
+	updated := &messages.TagCache{Entries: mergedTags, FileEntries: mergedFiles}
 	if err := c.dumpToDisk(ctx, updated); err != nil {
 		return err
 	}
 
 	// The state is persisted now.
 	c.cache = updated
-	c.added = nil
+	c.addedTags = nil
+	c.addedFiles = nil
 
 	return nil
 }