server: Remove cached revisions (#3277)

This commit refactors how the server manages the bundle revisions that
are included in decision logs and provenance results for the
API. Previously the revisions were cached on the server struct,
outside of the store. The server would read the revisions from the
store on reload() to keep them consistent.

While it is more performant to keep the revisions cached outside of
the store, it requires that the server perform reads on the
transaction that has already been committed. The inmem store
implementation allows for this however it's not going to be possible
to support that with other transaction implementations in the future.

This commit updates the server to simply read the revisions out of the
store in the handlers that require them. This adds a small amount of
overhead to the handlers that wasn't present before however in
practice this is not a concern (the overhead measured on my machine
was approx. 5 microseconds compared to the entire server handler that
was taking approx. 75 microseconds.)

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

server/server.go

@@ -105,8 +105,6 @@ type Server struct {
 	store                  storage.Store
 	manager                *plugins.Manager
 	decisionIDFactory      func() string
-	revisions              map[string]string
-	legacyRevision         string
 	buffer                 Buffer
 	logger                 func(context.Context, *Info) error
 	errLimit               int
@@ -158,15 +156,7 @@ func (s *Server) Init(ctx context.Context) (*Server, error) {
 
 	s.partials = map[string]rego.PartialResult{}
 	s.preparedEvalQueries = newCache(pqMaxCacheSize)
-
-	// Check if there is a bundle revision available at the legacy storage path
-	rev, err := bundle.LegacyReadRevisionFromStore(ctx, s.store, txn)
-	if err == nil && rev != "" {
-		s.legacyRevision = rev
-	}
-
 	s.defaultDecisionPath = s.generateDefaultDecisionPath()
-
 	s.interQueryBuiltinCache = iCache.NewInterQueryCache(s.manager.InterQueryBuiltinCacheConfig())
 	s.manager.RegisterCacheTrigger(s.updateCacheConfig)
 
@@ -660,9 +650,9 @@ func (s *Server) instrumentHandler(handler func(http.ResponseWriter, *http.Reque
 	return http.HandlerFunc(handler)
 }
 
-func (s *Server) execQuery(ctx context.Context, r *http.Request, txn storage.Transaction, decisionID string, parsedQuery ast.Body, input ast.Value, m metrics.Metrics, explainMode types.ExplainModeV1, includeMetrics, includeInstrumentation, pretty bool) (results types.QueryResponseV1, err error) {
+func (s *Server) execQuery(ctx context.Context, r *http.Request, br bundleRevisions, txn storage.Transaction, decisionID string, parsedQuery ast.Body, input ast.Value, m metrics.Metrics, explainMode types.ExplainModeV1, includeMetrics, includeInstrumentation, pretty bool) (results types.QueryResponseV1, err error) {
 
-	logger := s.getDecisionLogger()
+	logger := s.getDecisionLogger(br)
 
 	var buf *topdown.BufferTracer
 	if explainMode != types.ExplainOffV1 {
@@ -768,7 +758,13 @@ func (s *Server) indexGet(w http.ResponseWriter, r *http.Request) {
 
 	defer s.store.Abort(ctx, txn)
 
-	results, err := s.execQuery(ctx, r, txn, decisionID, parsedQuery, input, nil, explainMode, false, false, true)
+	br, err := getRevisions(ctx, s.store, txn)
+	if err != nil {
+		writer.ErrorAuto(w, err)
+		return
+	}
+
+	results, err := s.execQuery(ctx, r, br, txn, decisionID, parsedQuery, input, nil, explainMode, false, false, true)
 	if err != nil {
 		renderQueryResult(w, nil, err, t0)
 		return
@@ -782,6 +778,40 @@ func (s *Server) registerHandler(router *mux.Router, version int, path string, m
 	router.Handle(prefix+path, h).Methods(method)
 }
 
+type bundleRevisions struct {
+	LegacyRevision string
+	Revisions      map[string]string
+}
+
+func getRevisions(ctx context.Context, store storage.Store, txn storage.Transaction) (bundleRevisions, error) {
+
+	var err error
+	var br bundleRevisions
+	br.Revisions = map[string]string{}
+
+	// Check if we still have a legacy bundle manifest in the store
+	br.LegacyRevision, err = bundle.LegacyReadRevisionFromStore(ctx, store, txn)
+	if err != nil && !storage.IsNotFound(err) {
+		return br, err
+	}
+
+	// read all bundle revisions from storage (if any exist)
+	names, err := bundle.ReadBundleNamesFromStore(ctx, store, txn)
+	if err != nil && !storage.IsNotFound(err) {
+		return br, err
+	}
+
+	for _, name := range names {
+		r, err := bundle.ReadBundleRevisionFromStore(ctx, store, txn, name)
+		if err != nil && !storage.IsNotFound(err) {
+			return br, err
+		}
+		br.Revisions[name] = r
+	}
+
+	return br, nil
+}
+
 func (s *Server) reload(ctx context.Context, txn storage.Transaction, event storage.TriggerEvent) {
 
 	// NOTE(tsandall): We currently rely on the storage txn to provide
@@ -793,29 +823,8 @@ func (s *Server) reload(ctx context.Context, txn storage.Transaction, event stor
 
 	// reset some cached info
 	s.partials = map[string]rego.PartialResult{}
-	s.revisions = map[string]string{}
 	s.preparedEvalQueries = newCache(pqMaxCacheSize)
 	s.defaultDecisionPath = s.generateDefaultDecisionPath()
-
-	// read all bundle revisions from storage (if any exist)
-	names, err := bundle.ReadBundleNamesFromStore(ctx, s.store, txn)
-	if err != nil && !storage.IsNotFound(err) {
-		panic(err)
-	}
-
-	for _, name := range names {
-		r, err := bundle.ReadBundleRevisionFromStore(ctx, s.store, txn, name)
-		if err != nil && !storage.IsNotFound(err) {
-			panic(err)
-		}
-		s.revisions[name] = r
-	}
-
-	// Check if we still have a legacy bundle manifest in the store
-	s.legacyRevision, err = bundle.LegacyReadRevisionFromStore(ctx, s.store, txn)
-	if err != nil && !storage.IsNotFound(err) {
-		panic(err)
-	}
 }
 
 func (s *Server) unversionedPost(w http.ResponseWriter, r *http.Request) {
@@ -858,11 +867,17 @@ func (s *Server) v0QueryPath(w http.ResponseWriter, r *http.Request, urlPath str
 
 	defer s.store.Abort(ctx, txn)
 
+	br, err := getRevisions(ctx, s.store, txn)
+	if err != nil {
+		writer.ErrorAuto(w, err)
+		return
+	}
+
 	if useDefaultDecisionPath {
 		urlPath = s.defaultDecisionPath
 	}
 
-	logger := s.getDecisionLogger()
+	logger := s.getDecisionLogger(br)
 
 	pqID := "v0QueryPath::" + urlPath
 	preparedQuery, ok := s.getCachedPreparedEvalQuery(pqID, m)
@@ -1180,7 +1195,13 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 	}
 	defer s.store.Abort(ctx, txn)
 
-	logger := s.getDecisionLogger()
+	br, err := getRevisions(ctx, s.store, txn)
+	if err != nil {
+		writer.ErrorAuto(w, err)
+		return
+	}
+
+	logger := s.getDecisionLogger(br)
 
 	var buf *topdown.BufferTracer
 
@@ -1257,7 +1278,7 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if provenance {
-		result.Provenance = s.getProvenance()
+		result.Provenance = s.getProvenance(br)
 	}
 
 	if len(rs) == 0 {
@@ -1380,9 +1401,16 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 		writer.ErrorAuto(w, err)
 		return
 	}
+
 	defer s.store.Abort(ctx, txn)
 
-	logger := s.getDecisionLogger()
+	br, err := getRevisions(ctx, s.store, txn)
+	if err != nil {
+		writer.ErrorAuto(w, err)
+		return
+	}
+
+	logger := s.getDecisionLogger(br)
 
 	var buf *topdown.BufferTracer
 
@@ -1464,7 +1492,7 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if provenance {
-		result.Provenance = s.getProvenance()
+		result.Provenance = s.getProvenance(br)
 	}
 
 	if len(rs) == 0 {
@@ -1897,7 +1925,13 @@ func (s *Server) v1QueryGet(w http.ResponseWriter, r *http.Request) {
 
 	defer s.store.Abort(ctx, txn)
 
-	results, err := s.execQuery(ctx, r, txn, decisionID, parsedQuery, nil, m, explainMode, includeMetrics, includeInstrumentation, pretty)
+	br, err := getRevisions(ctx, s.store, txn)
+	if err != nil {
+		writer.ErrorAuto(w, err)
+		return
+	}
+
+	results, err := s.execQuery(ctx, r, br, txn, decisionID, parsedQuery, nil, m, explainMode, includeMetrics, includeInstrumentation, pretty)
 	if err != nil {
 		switch err := err.(type) {
 		case ast.Errors:
@@ -1950,7 +1984,13 @@ func (s *Server) v1QueryPost(w http.ResponseWriter, r *http.Request) {
 
 	defer s.store.Abort(ctx, txn)
 
-	results, err := s.execQuery(ctx, r, txn, decisionID, parsedQuery, nil, m, explainMode, includeMetrics, includeInstrumentation, pretty)
+	br, err := getRevisions(ctx, s.store, txn)
+	if err != nil {
+		writer.ErrorAuto(w, err)
+		return
+	}
+
+	results, err := s.execQuery(ctx, r, br, txn, decisionID, parsedQuery, nil, m, explainMode, includeMetrics, includeInstrumentation, pretty)
 	if err != nil {
 		switch err := err.(type) {
 		case ast.Errors:
@@ -2032,12 +2072,12 @@ func (s *Server) checkPathScope(ctx context.Context, txn storage.Transaction, pa
 	return nil
 }
 
-func (s *Server) getDecisionLogger() (logger decisionLogger) {
+func (s *Server) getDecisionLogger(br bundleRevisions) (logger decisionLogger) {
 	// For backwards compatibility use `revision` as needed.
-	if s.hasLegacyBundle() {
-		logger.revision = s.legacyRevision
+	if s.hasLegacyBundle(br) {
+		logger.revision = br.LegacyRevision
 	} else {
-		logger.revisions = s.revisions
+		logger.revisions = br.Revisions
 	}
 	logger.logger = s.logger
 	logger.buffer = s.buffer
@@ -2201,7 +2241,7 @@ func (s *Server) generateDecisionID() string {
 	return ""
 }
 
-func (s *Server) getProvenance() *types.ProvenanceV1 {
+func (s *Server) getProvenance(br bundleRevisions) *types.ProvenanceV1 {
 
 	p := &types.ProvenanceV1{
 		Version:   version.Version,
@@ -2213,11 +2253,11 @@ func (s *Server) getProvenance() *types.ProvenanceV1 {
 	// For backwards compatibility, if the bundles are using the old
 	// style config we need to fill in the older `Revision` field.
 	// Otherwise use the newer `Bundles` keyword.
-	if s.hasLegacyBundle() {
-		p.Revision = s.legacyRevision
+	if s.hasLegacyBundle(br) {
+		p.Revision = br.LegacyRevision
 	} else {
 		p.Bundles = map[string]types.ProvenanceBundleV1{}
-		for name, revision := range s.revisions {
+		for name, revision := range br.Revisions {
 			p.Bundles[name] = types.ProvenanceBundleV1{Revision: revision}
 		}
 	}
@@ -2225,9 +2265,9 @@ func (s *Server) getProvenance() *types.ProvenanceV1 {
 	return p
 }
 
-func (s *Server) hasLegacyBundle() bool {
+func (s *Server) hasLegacyBundle(br bundleRevisions) bool {
 	bp := bundlePlugin.Lookup(s.manager)
-	return s.legacyRevision != "" || (bp != nil && !bp.Config().IsMultiBundle())
+	return br.LegacyRevision != "" || (bp != nil && !bp.Config().IsMultiBundle())
 }
 
 func (s *Server) generateDefaultDecisionPath() string {

server/server_test.go

@@ -1675,9 +1675,15 @@ func TestDataProvenanceSingleBundle(t *testing.T) {
 		t.Errorf("Unexpected provenance data: \n\n%+v\n\nExpected:\n%+v\n\n", result.Provenance, expectedProvenance)
 	}
 
+	ctx := context.Background()
+
 	// Update bundle revision and request again
-	f.server.revisions["b1"] = "r1"
-	f.server.legacyRevision = "r1"
+	err := storage.Txn(ctx, f.server.store, storage.WriteParams, func(txn storage.Transaction) error {
+		return bundle.LegacyWriteManifestToStore(ctx, f.server.store, txn, bundle.Manifest{Revision: "r1"})
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	req = newReqV1(http.MethodPost, "/data?provenance", "")
 	f.reset()
@@ -1711,7 +1717,14 @@ func TestDataProvenanceSingleFileBundle(t *testing.T) {
 	version.Hostname = "foo.bar.com"
 
 	// No bundle plugin initialized, just a legacy revision set
-	f.server.legacyRevision = "r1"
+	ctx := context.Background()
+
+	err := storage.Txn(ctx, f.server.store, storage.WriteParams, func(txn storage.Transaction) error {
+		return bundle.LegacyWriteManifestToStore(ctx, f.server.store, txn, bundle.Manifest{Revision: "r1"})
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	req := newReqV1(http.MethodPost, "/data?provenance", "")
 	f.reset()
@@ -1783,7 +1796,14 @@ func TestDataProvenanceMultiBundle(t *testing.T) {
 	}
 
 	// Update bundle revision for a single bundle and make the request again
-	f.server.revisions["b1"] = "r1"
+	ctx := context.Background()
+
+	err := storage.Txn(ctx, f.server.store, storage.WriteParams, func(txn storage.Transaction) error {
+		return bundle.WriteManifestToStore(ctx, f.server.store, txn, "b1", bundle.Manifest{Revision: "r1"})
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	req = newReqV1(http.MethodPost, "/data?provenance", "")
 	f.reset()
@@ -1808,8 +1828,16 @@ func TestDataProvenanceMultiBundle(t *testing.T) {
 	}
 
 	// Update both and check again
-	f.server.revisions["b1"] = "r2"
-	f.server.revisions["b2"] = "r1"
+	err = storage.Txn(ctx, f.server.store, storage.WriteParams, func(txn storage.Transaction) error {
+		err := bundle.WriteManifestToStore(ctx, f.server.store, txn, "b1", bundle.Manifest{Revision: "r2"})
+		if err != nil {
+			return err
+		}
+		return bundle.WriteManifestToStore(ctx, f.server.store, txn, "b2", bundle.Manifest{Revision: "r1"})
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	req = newReqV1(http.MethodPost, "/data?provenance", "")
 	f.reset()