golden/go/expstorage/fs_expstore/fs_expstore.go - buildbot - Git at Google

 // Package fs_expstore an ExpectationsStore based on Firestore. See FIRESTORE.md for the schema
 // and design rationale.
 package fs_expstore

 import (
 	"context"
 	"errors"
 	"fmt"
 	"math/rand"
 	"strings"
 	"time"

 	"cloud.google.com/go/firestore"
 	"github.com/cenkalti/backoff"
 	"go.skia.org/infra/go/eventbus"
 	ifirestore "go.skia.org/infra/go/firestore"
 	"go.skia.org/infra/go/metrics2"
 	"go.skia.org/infra/go/skerr"
 	"go.skia.org/infra/go/sklog"
 	"go.skia.org/infra/go/util"
 	"go.skia.org/infra/golden/go/expstorage"
 	"go.skia.org/infra/golden/go/fs_utils"
 	"go.skia.org/infra/golden/go/shared"
 	"go.skia.org/infra/golden/go/types"
 	"go.skia.org/infra/golden/go/types/expectations"
 	"golang.org/x/sync/errgroup"
 )

 // AccessMode indicates if this ExpectationsStore can update existing Expectations
 // in the backing store or if if can only read them.
 type AccessMode int

 const (
 	ReadOnly AccessMode = iota
 	ReadWrite
 )

 var (
 	ReadOnlyErr = errors.New("expectationStore is in read-only mode")
 )

 const (
 	// These are the collections in Firestore.
 	expectationsCollection  = "expstore_expectations_v2"
 	triageRecordsCollection = "expstore_triage_records_v2"
 	triageChangesCollection = "expstore_triage_changes_v2"

 	// Fields in the Collections we query by.
 	committedField = "committed"
 	digestField    = "digest"
 	groupingField  = "grouping"
 	crsCLIDField   = "crs_cl_id"
 	recordIDField  = "record_id"
 	tsField        = "ts"

 	maxOperationTime = 2 * time.Minute

 	// There will not be very many entries on ChangeLists, relative to the masterBranch, so
 	// we can get away with many fewer shards to avoid the overhead of so many
 	// simultaneous queries.
 	clShards = 4

 	// snapshotShards was determined empirically on a data set of about 550k expectationEntry
 	// The more shards here, the more overhead and contention with the masterShards,
 	// so we aim for the sweet spot, erring on the side of too few shards.
 	// Times are for the New() function (i.e. initial fetch)
 	// 1 shard -> ???
 	// 8 shards -> 49s
 	// 16 shards -> 25s
 	// 32 shards -> 17s
 	// 64 shards -> 15s
 	// 96 shards -> ???
 	// 128 shards -> ???
 	// 512 shards -> ???
 	snapshotShards = 32

 	masterBranch = ""

 	// recoverTime is the minimum amount of time to wait before recreating any QuerySnapshotIterator
 	// if it fails. A random amount of time should be added to this, proportional to recoverTime.
 	recoverTime = 30 * time.Second
 )

 // Store implements expstorage.ExpectationsStore backed by
 // Firestore. It has a write-through caching mechanism.
 type Store struct {
 	client     *ifirestore.Client
 	mode       AccessMode
 	crsAndCLID string // crs+"_"+id. Empty string means master branch.

 	// eventBus allows this Store to communicate with the outside world when
 	// expectations change.
 	eventBus eventbus.EventBus
 	// globalEvent keeps track whether we want to send events within this instance
 	// or on the global eventbus.
 	globalEvent bool
 	// eventExpChange keeps track of which event to fire when the expectations change.
 	// This will be for either the MasterExpectations or for an IssueExpectations.
 	eventExpChange string

 	cache *expectations.Expectations

 	masterQuerySnapshots []*firestore.QuerySnapshotIterator
 }

 // expectationEntry is the document type stored in the expectationsCollection.
 type expectationEntry struct {
 	Grouping   types.TestName     `firestore:"grouping"`
 	Digest     types.Digest       `firestore:"digest"`
 	Label      expectations.Label `firestore:"label"`
 	Updated    time.Time          `firestore:"updated"`
 	CRSAndCLID string             `firestore:"crs_cl_id"`
 }

 // ID returns the deterministic ID that lets us update existing entries.
 func (e *expectationEntry) ID() string {
 	s := string(e.Grouping) + "|" + string(e.Digest)
 	// firestore gets cranky if there are / in key names
 	return strings.Replace(s, "/", "-", -1)
 }

 // triageRecord is the document type stored in the triageRecordsCollection.
 type triageRecord struct {
 	UserName   string    `firestore:"user"`
 	TS         time.Time `firestore:"ts"`
 	CRSAndCLID string    `firestore:"crs_cl_id"`
 	Changes    int       `firestore:"changes"`
 	Committed  bool      `firestore:"committed"`
 }

 // triageChanges is the document type stored in the triageChangesCollection.
 type triageChanges struct {
 	RecordID    string             `firestore:"record_id"`
 	Grouping    types.TestName     `firestore:"grouping"`
 	Digest      types.Digest       `firestore:"digest"`
 	LabelBefore expectations.Label `firestore:"before"`
 	LabelAfter  expectations.Label `firestore:"after"`
 }

 // New returns a new Store using the given firestore client. The Store will track
 // masterBranch- see ForChangeList() for getting Stores that track ChangeLists.
 // The passed in context is used for the QuerySnapshots (in ReadOnly mode).
 func New(ctx context.Context, client *ifirestore.Client, eventBus eventbus.EventBus, mode AccessMode) (*Store, error) {
 	defer metrics2.FuncTimer().Stop()
 	defer shared.NewMetricsTimer("expstore_init").Stop()
 	f := &Store{
 		client:         client,
 		eventBus:       eventBus,
 		eventExpChange: expstorage.EV_EXPSTORAGE_CHANGED,
 		globalEvent:    true,
 		crsAndCLID:     masterBranch,
 		mode:           mode,
 		cache:          &expectations.Expectations{},
 	}

 	err := f.initQuerySnapshot(ctx)
 	if err != nil {
 		return nil, skerr.Wrapf(err, "could not get initial query snapshot")
 	}

 	sklog.Infof("Loaded %d master expectations for %d tests", f.cache.Len(), f.cache.NumTests())

 	// Starts several go routines to listen to the snapshots created earlier.
 	f.listenToQuerySnapshots(ctx)

 	return f, nil
 }

 // ForChangeList implements the ExpectationsStore interface.
 func (f *Store) ForChangeList(id, crs string) expstorage.ExpectationsStore {
 	if id == masterBranch {
 		// It is invalid to re-request the master branch
 		return nil
 	}
 	return &Store{
 		client:     f.client,
 		eventBus:   nil,
 		crsAndCLID: crs + "_" + id,
 		mode:       f.mode,
 		cache:      &expectations.Expectations{},
 	}
 }

 // Get implements the ExpectationsStore interface.
 func (f *Store) Get() (expectations.ReadOnly, error) {
 	if f.crsAndCLID == masterBranch {
 		defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "true"}).Stop()
 		return f.cache, nil
 	}
 	defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "false"}).Stop()
 	return f.getExpectationsForCL()
 }

 // GetCopy implements the ExpectationsStore interface.
 func (f *Store) GetCopy() (*expectations.Expectations, error) {
 	if f.crsAndCLID == masterBranch {
 		defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "true"}).Stop()
 		return f.cache.DeepCopy(), nil
 	}
 	defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "false"}).Stop()
 	return f.getExpectationsForCL()
 }

 // initQuerySnapshot creates many firestore.QuerySnapshotIterator objects based on a shard of
 // all expectations and does the first Next() on them (which will try to return all data
 // in those shards). This data is loaded into the cache. Without sharding the queries, this times
 //  out with many expectations because of the fact that the first call to Next() fetches all data
 // currently there.
 func (f *Store) initQuerySnapshot(ctx context.Context) error {
 	q := f.client.Collection(expectationsCollection).Where(crsCLIDField, "==", masterBranch)
 	queries := fs_utils.ShardQueryOnDigest(q, digestField, snapshotShards)

 	f.masterQuerySnapshots = make([]*firestore.QuerySnapshotIterator, snapshotShards)
 	es := make([][]expectationEntry, snapshotShards)
 	var eg errgroup.Group
 	for shard, q := range queries {
 		func(shard int, q firestore.Query) {
 			eg.Go(func() error {
 				snap := q.Snapshots(ctx)
 				qs, err := snap.Next()
 				if err != nil {
 					return skerr.Wrapf(err, "getting initial snapshot data")
 				}
 				es[shard] = extractExpectationEntries(qs)

 				f.masterQuerySnapshots[shard] = snap
 				return nil
 			})
 		}(shard, q)
 	}
 	err := eg.Wait()
 	if err != nil {
 		return skerr.Wrap(err)
 	}

 	for _, entries := range es {
 		for _, e := range entries {
 			f.cache.Set(e.Grouping, e.Digest, e.Label)
 		}
 	}

 	return nil
 }

 // listenToQuerySnapshots takes the f.masterQuerySnapshots from earlier and spins up N
 // go routines that listen to those snapshots. If they see new triages (i.e. expectationEntry),
 // they update the f.cache (which is protected by cacheMutex).
 func (f *Store) listenToQuerySnapshots(ctx context.Context) {
 	metrics2.GetCounter("stopped_expstore_shards").Reset()
 	for i := 0; i < snapshotShards; i++ {
 		go func(shard int) {
 			for {
 				if err := ctx.Err(); err != nil {
 					f.masterQuerySnapshots[shard].Stop()
 					sklog.Debugf("Stopping query of snapshots on shard %d due to context err: %s", shard, err)
 					metrics2.GetCounter("stopped_expstore_shards").Inc(1)
 					return
 				}
 				qs, err := f.masterQuerySnapshots[shard].Next()
 				if err != nil {
 					sklog.Errorf("reading query snapshot %d: %s", shard, err)
 					f.masterQuerySnapshots[shard].Stop()
 					// sleep and rebuild the snapshot query. Once a SnapshotQueryIterator returns
 					// an error, it seems to always return that error. We sleep for a
 					// semi-randomized amount of time to spread out the re-building of shards
 					// (as it is likely all the shards will fail at about the same time).
 					t := recoverTime + time.Duration(float32(recoverTime)*rand.Float32())
 					time.Sleep(t)
 					sklog.Infof("Trying to recreate query snapshot %d after having slept %s", shard, t)
 					q := f.client.Collection(expectationsCollection).Where(crsCLIDField, "==", masterBranch)
 					queries := fs_utils.ShardQueryOnDigest(q, digestField, snapshotShards)
 					// This will trigger a complete re-request of this shard's data, to catch any
 					// updates that happened while we were not listening.
 					f.masterQuerySnapshots[shard] = queries[shard].Snapshots(ctx)
 					continue
 				}
 				entries := extractExpectationEntries(qs)
 				func() {
 					for _, e := range entries {
 						f.cache.Set(e.Grouping, e.Digest, e.Label)
 					}
 				}()

 				if f.eventBus != nil {
 					for _, e := range entries {
 						f.eventBus.Publish(f.eventExpChange, &expstorage.EventExpectationChange{
 							ExpectationDelta: expstorage.Delta{
 								Grouping: e.Grouping,
 								Digest:   e.Digest,
 								Label:    e.Label,
 							},
 							CRSAndCLID: f.crsAndCLID,
 						}, f.globalEvent)
 					}
 				}
 			}
 		}(i)
 	}
 }

 // extractExpectationEntries retrieves all []expectationEntry from a given QuerySnapshot, logging
 // any errors (which should be exceedingly rare)
 func extractExpectationEntries(qs *firestore.QuerySnapshot) []expectationEntry {
 	var entries []expectationEntry
 	for _, dc := range qs.Changes {
 		if dc.Kind == firestore.DocumentRemoved {
 			sklog.Warningf("Unexpected DocumentRemoved event: %#v", dc)
 			continue // There will likely never be DocumentRemoved events
 		}
 		entry := expectationEntry{}
 		if err := dc.Doc.DataTo(&entry); err != nil {
 			id := dc.Doc.Ref.ID
 			sklog.Errorf("corrupt data in firestore, could not unmarshal expectationEntry with id %s", id)
 			continue
 		}
 		entries = append(entries, entry)
 	}
 	return entries
 }

 // getExpectationsForCL returns an Expectations object which is safe to mutate
 // that has all cl-specific Expectations.
 // It fetches everything from firestore every time, as there could be multiple
 // readers and writers and thus caching isn't safe.
 func (f *Store) getExpectationsForCL() (*expectations.Expectations, error) {
 	defer metrics2.FuncTimer().Stop()

 	q := f.client.Collection(expectationsCollection).Where(crsCLIDField, "==", f.crsAndCLID)

 	es := make([]*expectations.Expectations, clShards)
 	queries := fs_utils.ShardQueryOnDigest(q, digestField, clShards)

 	maxRetries := 3
 	err := f.client.IterDocsInParallel(context.TODO(), "loadExpectations", f.crsAndCLID, queries, maxRetries, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error {
 		if doc == nil {
 			return nil
 		}
 		entry := expectationEntry{}
 		if err := doc.DataTo(&entry); err != nil {
 			id := doc.Ref.ID
 			return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal expectationEntry with id %s", id)
 		}
 		if es[i] == nil {
 			es[i] = &expectations.Expectations{}
 		}
 		es[i].Set(entry.Grouping, entry.Digest, entry.Label)
 		return nil
 	})

 	if err != nil {
 		return nil, skerr.Wrapf(err, "fetching expectations for ChangeList %s", f.crsAndCLID)
 	}

 	e := expectations.Expectations{}
 	for _, ne := range es {
 		e.MergeExpectations(ne)
 	}
 	return &e, nil
 }

 // AddChange implements the ExpectationsStore interface.
 func (f *Store) AddChange(ctx context.Context, delta []expstorage.Delta, userID string) error {
 	defer metrics2.FuncTimer().Stop()
 	if f.mode == ReadOnly {
 		return ReadOnlyErr
 	}
 	// Create the entries that we want to write (using the previous values)
 	now := time.Now()
 	entries, changes := f.flatten(now, delta)

 	// Nothing to add
 	if len(entries) == 0 {
 		return nil
 	}

 	// firestore can do up to 500 writes at once, we have 2 writes per entry, plus 1 triageRecord
 	const batchSize = (ifirestore.MAX_TRANSACTION_DOCS / 2) - 1

 	b := f.client.Batch()

 	// First write the triage record, with Committed being false (i.e. in progress)
 	tr := f.client.Collection(triageRecordsCollection).NewDoc()
 	record := triageRecord{
 		UserName:   userID,
 		TS:         now,
 		CRSAndCLID: f.crsAndCLID,
 		Changes:    len(entries),
 		Committed:  false,
 	}
 	b.Set(tr, record)

 	err := util.ChunkIter(len(entries), batchSize, func(start, stop int) error {
 		sklog.Debugf("Storing new expectations [%d, %d]", start, stop)
 		for offset, entry := range entries[start:stop] {
 			e := f.client.Collection(expectationsCollection).Doc(entry.ID())
 			b.Set(e, entry)

 			tc := f.client.Collection(triageChangesCollection).NewDoc()
 			change := changes[start+offset]
 			change.RecordID = tr.ID
 			b.Set(tc, change)
 		}

 		exp := &backoff.ExponentialBackOff{
 			InitialInterval:     time.Second,
 			RandomizationFactor: 0.5,
 			Multiplier:          2,
 			MaxInterval:         maxOperationTime / 4,
 			MaxElapsedTime:      maxOperationTime,
 			Clock:               backoff.SystemClock,
 		}

 		o := func() error {
 			_, err := b.Commit(ctx)
 			return err
 		}

 		if err := backoff.Retry(o, exp); err != nil {
 			// We really hope this doesn't happen, as it may leave the data in a partially
 			// broken state.
 			return skerr.Wrapf(err, "writing entries with retry [%d, %d]", start, stop)
 		}
 		// Go on to the next batch, if needed.
 		if stop < len(entries) {
 			b = f.client.Batch()
 		}
 		return nil
 	})
 	if err != nil {
 		return skerr.Wrap(err)
 	}

 	// We have succeeded this potentially long write, so mark it completed.
 	update := map[string]interface{}{
 		committedField: true,
 	}
 	_, err = f.client.Set(ctx, tr, update, 10, maxOperationTime, firestore.MergeAll)
 	return err
 }

 // flatten creates the data for the Documents to be written for a given Expectations delta.
 // It requires that the f.cache is safe to read (i.e. the mutex is held), because
 // it needs to determine the previous values.
 func (f *Store) flatten(now time.Time, delta []expstorage.Delta) ([]expectationEntry, []triageChanges) {
 	var entries []expectationEntry
 	var changes []triageChanges

 	for _, d := range delta {
 		entries = append(entries, expectationEntry{
 			Grouping:   d.Grouping,
 			Digest:     d.Digest,
 			Label:      d.Label,
 			Updated:    now,
 			CRSAndCLID: f.crsAndCLID,
 		})

 		changes = append(changes, triageChanges{
 			// RecordID will be filled out later
 			Grouping:    d.Grouping,
 			Digest:      d.Digest,
 			LabelBefore: f.cache.Classification(d.Grouping, d.Digest),
 			LabelAfter:  d.Label,
 		})
 	}
 	return entries, changes
 }

 // QueryLog implements the ExpectationsStore interface.
 func (f *Store) QueryLog(ctx context.Context, offset, size int, details bool) ([]expstorage.TriageLogEntry, int, error) {
 	if offset < 0 || size <= 0 {
 		return nil, -1, skerr.Fmt("offset: %d and size: %d must be positive", offset, size)
 	}
 	defer metrics2.FuncTimer().Stop()

 	// Fetch the records, which have everything except the details.
 	q := f.client.Collection(triageRecordsCollection).OrderBy(tsField, firestore.Desc).Offset(offset).Limit(size)
 	q = q.Where(crsCLIDField, "==", f.crsAndCLID).Where(committedField, "==", true)
 	var rv []expstorage.TriageLogEntry
 	d := fmt.Sprintf("offset: %d, size %d", offset, size)
 	err := f.client.IterDocs(ctx, "query_log", d, q, 3, maxOperationTime, func(doc *firestore.DocumentSnapshot) error {
 		if doc == nil {
 			return nil
 		}
 		tr := triageRecord{}
 		if err := doc.DataTo(&tr); err != nil {
 			id := doc.Ref.ID
 			return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal triageRecord with id %s", id)
 		}
 		rv = append(rv, expstorage.TriageLogEntry{
 			ID:          doc.Ref.ID,
 			User:        tr.UserName,
 			TS:          tr.TS,
 			ChangeCount: tr.Changes,
 		})
 		return nil
 	})
 	if err != nil {
 		return nil, -1, skerr.Wrapf(err, "could not request triage records [%d: %d]", offset, size)
 	}

 	n := len(rv)
 	if n == size && n != 0 {
 		// We don't know how many there are and it might be too slow to count, so just give
 		// the "many" response.
 		n = expstorage.CountMany
 	} else {
 		// We know exactly either 1) how many there are (if n > 0) or 2) an upper bound on how many
 		// there are (if n == 0)
 		n += offset
 	}

 	if len(rv) == 0 || !details {
 		return rv, n, nil
 	}

 	// Make a query for each of the records to fetch the changes belonging to that record.
 	qs := make([]firestore.Query, 0, len(rv))
 	for _, r := range rv {
 		q := f.client.Collection(triageChangesCollection).Where(recordIDField, "==", r.ID)
 		// Sort them by grouping, then Digest for determinism
 		q = q.OrderBy(groupingField, firestore.Asc).OrderBy(digestField, firestore.Asc)
 		qs = append(qs, q)
 	}

 	// Then fire them all off in parallel.
 	err = f.client.IterDocsInParallel(ctx, "query_log_details", d, qs, 3, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error {
 		if doc == nil {
 			return nil
 		}
 		tc := triageChanges{}
 		if err := doc.DataTo(&tc); err != nil {
 			id := doc.Ref.ID
 			return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal triageChanges with id %s", id)
 		}
 		rv[i].Details = append(rv[i].Details, expstorage.Delta{
 			Grouping: tc.Grouping,
 			Digest:   tc.Digest,
 			Label:    tc.LabelAfter,
 		})
 		return nil
 	})
 	if err != nil {
 		return nil, -1, skerr.Wrapf(err, "could not query details")
 	}

 	return rv, n, nil
 }

 // UndoChange implements the ExpectationsStore interface.
 func (f *Store) UndoChange(ctx context.Context, changeID, userID string) error {
 	defer metrics2.FuncTimer().Stop()
 	if f.mode == ReadOnly {
 		return ReadOnlyErr
 	}
 	// Verify the original change id exists.
 	dr := f.client.Collection(triageRecordsCollection).Doc(changeID)
 	doc, err := f.client.Get(ctx, dr, 3, maxOperationTime)
 	if err != nil || !doc.Exists() {
 		return skerr.Wrapf(err, "could not find change to undo with id %s", changeID)
 	}

 	q := f.client.Collection(triageChangesCollection).Where(recordIDField, "==", changeID)
 	var delta []expstorage.Delta
 	err = f.client.IterDocs(ctx, "undo_query", changeID, q, 3, maxOperationTime, func(doc *firestore.DocumentSnapshot) error {
 		if doc == nil {
 			return nil
 		}
 		tc := triageChanges{}
 		if err := doc.DataTo(&tc); err != nil {
 			id := doc.Ref.ID
 			return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal triageChanges with id %s", id)
 		}
 		delta = append(delta, expstorage.Delta{
 			Grouping: tc.Grouping,
 			Digest:   tc.Digest,
 			Label:    tc.LabelBefore,
 		})
 		return nil
 	})
 	if err != nil {
 		return skerr.Wrapf(err, "could not get delta to undo %s", changeID)
 	}

 	if err = f.AddChange(ctx, delta, userID); err != nil {
 		return skerr.Wrapf(err, "could not apply delta to undo %s", changeID)
 	}

 	return nil
 }

 // Make sure Store fulfills the ExpectationsStore interface
 var _ expstorage.ExpectationsStore = (*Store)(nil)
	// Package fs_expstore an ExpectationsStore based on Firestore. See FIRESTORE.md for the schema
	// and design rationale.
	package fs_expstore

	import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"strings"
	"time"

	"cloud.google.com/go/firestore"
	"github.com/cenkalti/backoff"
	"go.skia.org/infra/go/eventbus"
	ifirestore "go.skia.org/infra/go/firestore"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/expstorage"
	"go.skia.org/infra/golden/go/fs_utils"
	"go.skia.org/infra/golden/go/shared"
	"go.skia.org/infra/golden/go/types"
	"go.skia.org/infra/golden/go/types/expectations"
	"golang.org/x/sync/errgroup"
	)

	// AccessMode indicates if this ExpectationsStore can update existing Expectations
	// in the backing store or if if can only read them.
	type AccessMode int

	const (
	ReadOnly AccessMode = iota
	ReadWrite
	)

	var (
	ReadOnlyErr = errors.New("expectationStore is in read-only mode")
	)

	const (
	// These are the collections in Firestore.
	expectationsCollection = "expstore_expectations_v2"
	triageRecordsCollection = "expstore_triage_records_v2"
	triageChangesCollection = "expstore_triage_changes_v2"

	// Fields in the Collections we query by.
	committedField = "committed"
	digestField = "digest"
	groupingField = "grouping"
	crsCLIDField = "crs_cl_id"
	recordIDField = "record_id"
	tsField = "ts"

	maxOperationTime = 2 * time.Minute

	// There will not be very many entries on ChangeLists, relative to the masterBranch, so
	// we can get away with many fewer shards to avoid the overhead of so many
	// simultaneous queries.
	clShards = 4

	// snapshotShards was determined empirically on a data set of about 550k expectationEntry
	// The more shards here, the more overhead and contention with the masterShards,
	// so we aim for the sweet spot, erring on the side of too few shards.
	// Times are for the New() function (i.e. initial fetch)
	// 1 shard -> ???
	// 8 shards -> 49s
	// 16 shards -> 25s
	// 32 shards -> 17s
	// 64 shards -> 15s
	// 96 shards -> ???
	// 128 shards -> ???
	// 512 shards -> ???
	snapshotShards = 32

	masterBranch = ""

	// recoverTime is the minimum amount of time to wait before recreating any QuerySnapshotIterator
	// if it fails. A random amount of time should be added to this, proportional to recoverTime.
	recoverTime = 30 * time.Second
	)

	// Store implements expstorage.ExpectationsStore backed by
	// Firestore. It has a write-through caching mechanism.
	type Store struct {
	client *ifirestore.Client
	mode AccessMode
	crsAndCLID string // crs+"_"+id. Empty string means master branch.

	// eventBus allows this Store to communicate with the outside world when
	// expectations change.
	eventBus eventbus.EventBus
	// globalEvent keeps track whether we want to send events within this instance
	// or on the global eventbus.
	globalEvent bool
	// eventExpChange keeps track of which event to fire when the expectations change.
	// This will be for either the MasterExpectations or for an IssueExpectations.
	eventExpChange string

	cache *expectations.Expectations

	masterQuerySnapshots []*firestore.QuerySnapshotIterator
	}

	// expectationEntry is the document type stored in the expectationsCollection.
	type expectationEntry struct {
	Grouping types.TestName `firestore:"grouping"`
	Digest types.Digest `firestore:"digest"`
	Label expectations.Label `firestore:"label"`
	Updated time.Time `firestore:"updated"`
	CRSAndCLID string `firestore:"crs_cl_id"`
	}

	// ID returns the deterministic ID that lets us update existing entries.
	func (e *expectationEntry) ID() string {
	s := string(e.Grouping) + "\|" + string(e.Digest)
	// firestore gets cranky if there are / in key names
	return strings.Replace(s, "/", "-", -1)
	}

	// triageRecord is the document type stored in the triageRecordsCollection.
	type triageRecord struct {
	UserName string `firestore:"user"`
	TS time.Time `firestore:"ts"`
	CRSAndCLID string `firestore:"crs_cl_id"`
	Changes int `firestore:"changes"`
	Committed bool `firestore:"committed"`
	}

	// triageChanges is the document type stored in the triageChangesCollection.
	type triageChanges struct {
	RecordID string `firestore:"record_id"`
	Grouping types.TestName `firestore:"grouping"`
	Digest types.Digest `firestore:"digest"`
	LabelBefore expectations.Label `firestore:"before"`
	LabelAfter expectations.Label `firestore:"after"`
	}

	// New returns a new Store using the given firestore client. The Store will track
	// masterBranch- see ForChangeList() for getting Stores that track ChangeLists.
	// The passed in context is used for the QuerySnapshots (in ReadOnly mode).
	func New(ctx context.Context, client ifirestore.Client, eventBus eventbus.EventBus, mode AccessMode) (Store, error) {
	defer metrics2.FuncTimer().Stop()
	defer shared.NewMetricsTimer("expstore_init").Stop()
	f := &Store{
	client: client,
	eventBus: eventBus,
	eventExpChange: expstorage.EV_EXPSTORAGE_CHANGED,
	globalEvent: true,
	crsAndCLID: masterBranch,
	mode: mode,
	cache: &expectations.Expectations{},
	}

	err := f.initQuerySnapshot(ctx)
	if err != nil {
	return nil, skerr.Wrapf(err, "could not get initial query snapshot")
	}

	sklog.Infof("Loaded %d master expectations for %d tests", f.cache.Len(), f.cache.NumTests())

	// Starts several go routines to listen to the snapshots created earlier.
	f.listenToQuerySnapshots(ctx)

	return f, nil
	}

	// ForChangeList implements the ExpectationsStore interface.
	func (f *Store) ForChangeList(id, crs string) expstorage.ExpectationsStore {
	if id == masterBranch {
	// It is invalid to re-request the master branch
	return nil
	}
	return &Store{
	client: f.client,
	eventBus: nil,
	crsAndCLID: crs + "_" + id,
	mode: f.mode,
	cache: &expectations.Expectations{},
	}
	}

	// Get implements the ExpectationsStore interface.
	func (f *Store) Get() (expectations.ReadOnly, error) {
	if f.crsAndCLID == masterBranch {
	defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "true"}).Stop()
	return f.cache, nil
	}
	defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "false"}).Stop()
	return f.getExpectationsForCL()
	}

	// GetCopy implements the ExpectationsStore interface.
	func (f Store) GetCopy() (expectations.Expectations, error) {
	if f.crsAndCLID == masterBranch {
	defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "true"}).Stop()
	return f.cache.DeepCopy(), nil
	}
	defer metrics2.NewTimer("gold_get_expectations", map[string]string{"master_branch": "false"}).Stop()
	return f.getExpectationsForCL()
	}

	// initQuerySnapshot creates many firestore.QuerySnapshotIterator objects based on a shard of
	// all expectations and does the first Next() on them (which will try to return all data
	// in those shards). This data is loaded into the cache. Without sharding the queries, this times
	// out with many expectations because of the fact that the first call to Next() fetches all data
	// currently there.
	func (f *Store) initQuerySnapshot(ctx context.Context) error {
	q := f.client.Collection(expectationsCollection).Where(crsCLIDField, "==", masterBranch)
	queries := fs_utils.ShardQueryOnDigest(q, digestField, snapshotShards)

	f.masterQuerySnapshots = make([]*firestore.QuerySnapshotIterator, snapshotShards)
	es := make([][]expectationEntry, snapshotShards)
	var eg errgroup.Group
	for shard, q := range queries {
	func(shard int, q firestore.Query) {
	eg.Go(func() error {
	snap := q.Snapshots(ctx)
	qs, err := snap.Next()
	if err != nil {
	return skerr.Wrapf(err, "getting initial snapshot data")
	}
	es[shard] = extractExpectationEntries(qs)

	f.masterQuerySnapshots[shard] = snap
	return nil
	})
	}(shard, q)
	}
	err := eg.Wait()
	if err != nil {
	return skerr.Wrap(err)
	}

	for _, entries := range es {
	for _, e := range entries {
	f.cache.Set(e.Grouping, e.Digest, e.Label)
	}
	}

	return nil
	}

	// listenToQuerySnapshots takes the f.masterQuerySnapshots from earlier and spins up N
	// go routines that listen to those snapshots. If they see new triages (i.e. expectationEntry),
	// they update the f.cache (which is protected by cacheMutex).
	func (f *Store) listenToQuerySnapshots(ctx context.Context) {
	metrics2.GetCounter("stopped_expstore_shards").Reset()
	for i := 0; i < snapshotShards; i++ {
	go func(shard int) {
	for {
	if err := ctx.Err(); err != nil {
	f.masterQuerySnapshots[shard].Stop()
	sklog.Debugf("Stopping query of snapshots on shard %d due to context err: %s", shard, err)
	metrics2.GetCounter("stopped_expstore_shards").Inc(1)
	return
	}
	qs, err := f.masterQuerySnapshots[shard].Next()
	if err != nil {
	sklog.Errorf("reading query snapshot %d: %s", shard, err)
	f.masterQuerySnapshots[shard].Stop()
	// sleep and rebuild the snapshot query. Once a SnapshotQueryIterator returns
	// an error, it seems to always return that error. We sleep for a
	// semi-randomized amount of time to spread out the re-building of shards
	// (as it is likely all the shards will fail at about the same time).
	t := recoverTime + time.Duration(float32(recoverTime)*rand.Float32())
	time.Sleep(t)
	sklog.Infof("Trying to recreate query snapshot %d after having slept %s", shard, t)
	q := f.client.Collection(expectationsCollection).Where(crsCLIDField, "==", masterBranch)
	queries := fs_utils.ShardQueryOnDigest(q, digestField, snapshotShards)
	// This will trigger a complete re-request of this shard's data, to catch any
	// updates that happened while we were not listening.
	f.masterQuerySnapshots[shard] = queries[shard].Snapshots(ctx)
	continue
	}
	entries := extractExpectationEntries(qs)
	func() {
	for _, e := range entries {
	f.cache.Set(e.Grouping, e.Digest, e.Label)
	}
	}()

	if f.eventBus != nil {
	for _, e := range entries {
	f.eventBus.Publish(f.eventExpChange, &expstorage.EventExpectationChange{
	ExpectationDelta: expstorage.Delta{
	Grouping: e.Grouping,
	Digest: e.Digest,
	Label: e.Label,
	},
	CRSAndCLID: f.crsAndCLID,
	}, f.globalEvent)
	}
	}
	}
	}(i)
	}
	}

	// extractExpectationEntries retrieves all []expectationEntry from a given QuerySnapshot, logging
	// any errors (which should be exceedingly rare)
	func extractExpectationEntries(qs *firestore.QuerySnapshot) []expectationEntry {
	var entries []expectationEntry
	for _, dc := range qs.Changes {
	if dc.Kind == firestore.DocumentRemoved {
	sklog.Warningf("Unexpected DocumentRemoved event: %#v", dc)
	continue // There will likely never be DocumentRemoved events
	}
	entry := expectationEntry{}
	if err := dc.Doc.DataTo(&entry); err != nil {
	id := dc.Doc.Ref.ID
	sklog.Errorf("corrupt data in firestore, could not unmarshal expectationEntry with id %s", id)
	continue
	}
	entries = append(entries, entry)
	}
	return entries
	}

	// getExpectationsForCL returns an Expectations object which is safe to mutate
	// that has all cl-specific Expectations.
	// It fetches everything from firestore every time, as there could be multiple
	// readers and writers and thus caching isn't safe.
	func (f Store) getExpectationsForCL() (expectations.Expectations, error) {
	defer metrics2.FuncTimer().Stop()

	q := f.client.Collection(expectationsCollection).Where(crsCLIDField, "==", f.crsAndCLID)

	es := make([]*expectations.Expectations, clShards)
	queries := fs_utils.ShardQueryOnDigest(q, digestField, clShards)

	maxRetries := 3
	err := f.client.IterDocsInParallel(context.TODO(), "loadExpectations", f.crsAndCLID, queries, maxRetries, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error {
	if doc == nil {
	return nil
	}
	entry := expectationEntry{}
	if err := doc.DataTo(&entry); err != nil {
	id := doc.Ref.ID
	return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal expectationEntry with id %s", id)
	}
	if es[i] == nil {
	es[i] = &expectations.Expectations{}
	}
	es[i].Set(entry.Grouping, entry.Digest, entry.Label)
	return nil
	})

	if err != nil {
	return nil, skerr.Wrapf(err, "fetching expectations for ChangeList %s", f.crsAndCLID)
	}

	e := expectations.Expectations{}
	for _, ne := range es {
	e.MergeExpectations(ne)
	}
	return &e, nil
	}

	// AddChange implements the ExpectationsStore interface.
	func (f *Store) AddChange(ctx context.Context, delta []expstorage.Delta, userID string) error {
	defer metrics2.FuncTimer().Stop()
	if f.mode == ReadOnly {
	return ReadOnlyErr
	}
	// Create the entries that we want to write (using the previous values)
	now := time.Now()
	entries, changes := f.flatten(now, delta)

	// Nothing to add
	if len(entries) == 0 {
	return nil
	}

	// firestore can do up to 500 writes at once, we have 2 writes per entry, plus 1 triageRecord
	const batchSize = (ifirestore.MAX_TRANSACTION_DOCS / 2) - 1

	b := f.client.Batch()

	// First write the triage record, with Committed being false (i.e. in progress)
	tr := f.client.Collection(triageRecordsCollection).NewDoc()
	record := triageRecord{
	UserName: userID,
	TS: now,
	CRSAndCLID: f.crsAndCLID,
	Changes: len(entries),
	Committed: false,
	}
	b.Set(tr, record)

	err := util.ChunkIter(len(entries), batchSize, func(start, stop int) error {
	sklog.Debugf("Storing new expectations [%d, %d]", start, stop)
	for offset, entry := range entries[start:stop] {
	e := f.client.Collection(expectationsCollection).Doc(entry.ID())
	b.Set(e, entry)

	tc := f.client.Collection(triageChangesCollection).NewDoc()
	change := changes[start+offset]
	change.RecordID = tr.ID
	b.Set(tc, change)
	}

	exp := &backoff.ExponentialBackOff{
	InitialInterval: time.Second,
	RandomizationFactor: 0.5,
	Multiplier: 2,
	MaxInterval: maxOperationTime / 4,
	MaxElapsedTime: maxOperationTime,
	Clock: backoff.SystemClock,
	}

	o := func() error {
	_, err := b.Commit(ctx)
	return err
	}

	if err := backoff.Retry(o, exp); err != nil {
	// We really hope this doesn't happen, as it may leave the data in a partially
	// broken state.
	return skerr.Wrapf(err, "writing entries with retry [%d, %d]", start, stop)
	}
	// Go on to the next batch, if needed.
	if stop < len(entries) {
	b = f.client.Batch()
	}
	return nil
	})
	if err != nil {
	return skerr.Wrap(err)
	}

	// We have succeeded this potentially long write, so mark it completed.
	update := map[string]interface{}{
	committedField: true,
	}
	_, err = f.client.Set(ctx, tr, update, 10, maxOperationTime, firestore.MergeAll)
	return err
	}

	// flatten creates the data for the Documents to be written for a given Expectations delta.
	// It requires that the f.cache is safe to read (i.e. the mutex is held), because
	// it needs to determine the previous values.
	func (f *Store) flatten(now time.Time, delta []expstorage.Delta) ([]expectationEntry, []triageChanges) {
	var entries []expectationEntry
	var changes []triageChanges

	for _, d := range delta {
	entries = append(entries, expectationEntry{
	Grouping: d.Grouping,
	Digest: d.Digest,
	Label: d.Label,
	Updated: now,
	CRSAndCLID: f.crsAndCLID,
	})

	changes = append(changes, triageChanges{
	// RecordID will be filled out later
	Grouping: d.Grouping,
	Digest: d.Digest,
	LabelBefore: f.cache.Classification(d.Grouping, d.Digest),
	LabelAfter: d.Label,
	})
	}
	return entries, changes
	}

	// QueryLog implements the ExpectationsStore interface.
	func (f *Store) QueryLog(ctx context.Context, offset, size int, details bool) ([]expstorage.TriageLogEntry, int, error) {
	if offset < 0 \|\| size <= 0 {
	return nil, -1, skerr.Fmt("offset: %d and size: %d must be positive", offset, size)
	}
	defer metrics2.FuncTimer().Stop()

	// Fetch the records, which have everything except the details.
	q := f.client.Collection(triageRecordsCollection).OrderBy(tsField, firestore.Desc).Offset(offset).Limit(size)
	q = q.Where(crsCLIDField, "==", f.crsAndCLID).Where(committedField, "==", true)
	var rv []expstorage.TriageLogEntry
	d := fmt.Sprintf("offset: %d, size %d", offset, size)
	err := f.client.IterDocs(ctx, "query_log", d, q, 3, maxOperationTime, func(doc *firestore.DocumentSnapshot) error {
	if doc == nil {
	return nil
	}
	tr := triageRecord{}
	if err := doc.DataTo(&tr); err != nil {
	id := doc.Ref.ID
	return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal triageRecord with id %s", id)
	}
	rv = append(rv, expstorage.TriageLogEntry{
	ID: doc.Ref.ID,
	User: tr.UserName,
	TS: tr.TS,
	ChangeCount: tr.Changes,
	})
	return nil
	})
	if err != nil {
	return nil, -1, skerr.Wrapf(err, "could not request triage records [%d: %d]", offset, size)
	}

	n := len(rv)
	if n == size && n != 0 {
	// We don't know how many there are and it might be too slow to count, so just give
	// the "many" response.
	n = expstorage.CountMany
	} else {
	// We know exactly either 1) how many there are (if n > 0) or 2) an upper bound on how many
	// there are (if n == 0)
	n += offset
	}

	if len(rv) == 0 \|\| !details {
	return rv, n, nil
	}

	// Make a query for each of the records to fetch the changes belonging to that record.
	qs := make([]firestore.Query, 0, len(rv))
	for _, r := range rv {
	q := f.client.Collection(triageChangesCollection).Where(recordIDField, "==", r.ID)
	// Sort them by grouping, then Digest for determinism
	q = q.OrderBy(groupingField, firestore.Asc).OrderBy(digestField, firestore.Asc)
	qs = append(qs, q)
	}

	// Then fire them all off in parallel.
	err = f.client.IterDocsInParallel(ctx, "query_log_details", d, qs, 3, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error {
	if doc == nil {
	return nil
	}
	tc := triageChanges{}
	if err := doc.DataTo(&tc); err != nil {
	id := doc.Ref.ID
	return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal triageChanges with id %s", id)
	}
	rv[i].Details = append(rv[i].Details, expstorage.Delta{
	Grouping: tc.Grouping,
	Digest: tc.Digest,
	Label: tc.LabelAfter,
	})
	return nil
	})
	if err != nil {
	return nil, -1, skerr.Wrapf(err, "could not query details")
	}

	return rv, n, nil
	}

	// UndoChange implements the ExpectationsStore interface.
	func (f *Store) UndoChange(ctx context.Context, changeID, userID string) error {
	defer metrics2.FuncTimer().Stop()
	if f.mode == ReadOnly {
	return ReadOnlyErr
	}
	// Verify the original change id exists.
	dr := f.client.Collection(triageRecordsCollection).Doc(changeID)
	doc, err := f.client.Get(ctx, dr, 3, maxOperationTime)
	if err != nil \|\| !doc.Exists() {
	return skerr.Wrapf(err, "could not find change to undo with id %s", changeID)
	}

	q := f.client.Collection(triageChangesCollection).Where(recordIDField, "==", changeID)
	var delta []expstorage.Delta
	err = f.client.IterDocs(ctx, "undo_query", changeID, q, 3, maxOperationTime, func(doc *firestore.DocumentSnapshot) error {
	if doc == nil {
	return nil
	}
	tc := triageChanges{}
	if err := doc.DataTo(&tc); err != nil {
	id := doc.Ref.ID
	return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal triageChanges with id %s", id)
	}
	delta = append(delta, expstorage.Delta{
	Grouping: tc.Grouping,
	Digest: tc.Digest,
	Label: tc.LabelBefore,
	})
	return nil
	})
	if err != nil {
	return skerr.Wrapf(err, "could not get delta to undo %s", changeID)
	}

	if err = f.AddChange(ctx, delta, userID); err != nil {
	return skerr.Wrapf(err, "could not apply delta to undo %s", changeID)
	}

	return nil
	}

	// Make sure Store fulfills the ExpectationsStore interface
	var _ expstorage.ExpectationsStore = (*Store)(nil)