server: Update health check to use plugin status

The health check now supports a `?plugin` option which will make the
response depend on whether or not all configured plugins are in an OK
state.

The `bundle` parameter will now use the bundle *and* discovery plugin
statuses to determine if the bundles are ready. This corrects an issue
where discovery bundles, and bundles defined by the discovery dynamic
config, were not included with `/health?bundle=true` checks.

The URL parameter parsing has also changed to allow for omitting the
value for the `bundle` option. It will default to `true` so that
`/health?bundle=true` can be shortened to `/health?bundle`.

Fixes: #2010
Fixes: #2015
Signed-off-by: Patrick East <east.patrick@gmail.com>

server/server.go

@@ -110,8 +110,6 @@ type Server struct {
 	pprofEnabled        bool
 	runtime             *ast.Term
 	httpListeners       []httpListener
-	bundleStatuses      map[string]*bundlePlugin.Status
-	bundleStatusMtx     sync.RWMutex
 	metrics             Metrics
 	defaultDecisionPath string
 }
@@ -180,18 +178,6 @@ func (s *Server) Init(ctx context.Context) (*Server, error) {
 	s.partials = map[string]rego.PartialResult{}
 	s.preparedEvalQueries = newCache(pqMaxCacheSize)
 
-	bp := bundlePlugin.Lookup(s.manager)
-	if bp != nil {
-
-		// initialize statuses to empty defaults for server /health check
-		s.bundleStatuses = map[string]*bundlePlugin.Status{}
-		for bundleName := range bp.Config().Bundles {
-			s.bundleStatuses[bundleName] = &bundlePlugin.Status{Name: bundleName}
-		}
-
-		bp.RegisterBulkListener("REST API Server", s.updateBundleStatus)
-	}
-
 	// 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 != "" {
@@ -871,12 +857,6 @@ func (s *Server) getCachedPreparedEvalQuery(key string, m metrics.Metrics) (*reg
 	return nil, false
 }
 
-func (s *Server) updateBundleStatus(status map[string]*bundlePlugin.Status) {
-	s.bundleStatusMtx.Lock()
-	defer s.bundleStatusMtx.Unlock()
-	s.bundleStatuses = status
-}
-
 func (s *Server) canEval(ctx context.Context) bool {
 	// Create very simple query that binds a single variable.
 	eval := rego.New(rego.Compiler(s.getCompiler()),
@@ -886,7 +866,7 @@ func (s *Server) canEval(ctx context.Context) bool {
 	if err != nil {
 		return false
 	}
-	type emptyObject struct{}
+
 	v, ok := rs[0].Bindings["x"]
 	if ok {
 		jsonNumber, ok := v.(json.Number)
@@ -897,15 +877,22 @@ func (s *Server) canEval(ctx context.Context) bool {
 	return false
 }
 
-func (s *Server) bundlesActivated() bool {
-	s.bundleStatusMtx.RLock()
-	defer s.bundleStatusMtx.RUnlock()
+func (s *Server) bundlesReady(pluginStatuses map[string]*plugins.Status) bool {
 
-	for _, status := range s.bundleStatuses {
-		// Ensure that all of the bundle statuses have an activation time set on them
-		if (status.LastSuccessfulActivation == time.Time{}) {
-			return false
-		}
+	// Look for a discovery plugin first, if it exists and isn't ready
+	// then don't bother with the others.
+	// Note: use "discovery" instead of `discovery.Name` to avoid import
+	// cycle problems..
+	dpStatus, ok := pluginStatuses["discovery"]
+	if ok && dpStatus != nil && (dpStatus.State != plugins.StateOK) {
+		return false
+	}
+
+	// The bundle plugin won't return "OK" until the first activation
+	// of each configured bundle.
+	bpStatus, ok := pluginStatuses[bundlePlugin.Name]
+	if ok && bpStatus != nil && (bpStatus.State != plugins.StateOK) {
+		return false
 	}
 
 	return true
@@ -913,23 +900,50 @@ func (s *Server) bundlesActivated() bool {
 
 func (s *Server) unversionedGetHealth(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
-	includeBundleStatus := getBoolParam(r.URL, types.ParamBundleActivationV1, false)
+	includeBundleStatus := getBoolParam(r.URL, types.ParamBundleActivationV1, true)
+	includePluginStatus := getBoolParam(r.URL, types.ParamPluginsV1, true)
 
 	// Ensure the server can evaluate a simple query
-	type emptyObject struct{}
 	if !s.canEval(ctx) {
-		writer.JSON(w, http.StatusInternalServerError, emptyObject{}, false)
+		writeHealthResponse(w, errors.New("unable to perform evaluation"), struct{}{})
 		return
 	}
 
+	pluginStatuses := s.manager.PluginStatus()
+
 	// Ensure that bundles (if configured, and requested to be included in the result)
-	// have been activated successfully.
-	if includeBundleStatus && s.hasBundle() && !s.bundlesActivated() {
-		writer.JSON(w, http.StatusInternalServerError, emptyObject{}, false)
+	// have been activated successfully. This will include discovery bundles as well as
+	// normal bundles that are configured.
+	if includeBundleStatus && !s.bundlesReady(pluginStatuses) {
+		// For backwards compatibility we don't return a payload with statuses for the bundle endpoint
+		writeHealthResponse(w, errors.New("not all configured bundles have been activated"), struct{}{})
 		return
 	}
 
-	writer.JSON(w, http.StatusOK, emptyObject{}, false)
+	if includePluginStatus {
+		// Ensure that all plugins (if requested to be included in the result) have an OK status.
+		hasErr := false
+		for _, status := range pluginStatuses {
+			if status.State != plugins.StateOK {
+				hasErr = true
+				break
+			}
+		}
+		if hasErr {
+			writeHealthResponse(w, errors.New("not all plugins in OK state"), struct{}{})
+			return
+		}
+	}
+	writeHealthResponse(w, nil, struct{}{})
+}
+
+func writeHealthResponse(w http.ResponseWriter, err error, payload interface{}) {
+	status := http.StatusOK
+	if err != nil {
+		status = http.StatusInternalServerError
+	}
+
+	writer.JSON(w, status, payload, false)
 }
 
 func (s *Server) v1CompilePost(w http.ResponseWriter, r *http.Request) {
@@ -2128,10 +2142,6 @@ func (s *Server) getProvenance() *types.ProvenanceV1 {
 	return p
 }
 
-func (s *Server) hasBundle() bool {
-	return bundlePlugin.Lookup(s.manager) != nil || s.legacyRevision != ""
-}
-
 func (s *Server) hasLegacyBundle() bool {
 	bp := bundlePlugin.Lookup(s.manager)
 	return s.legacyRevision != "" || (bp != nil && !bp.Config().IsMultiBundle())

server/server_test.go

@@ -73,17 +73,13 @@ func TestUnversionedGetHealthBundleNoBundleSet(t *testing.T) {
 	}
 }
 
-func TestUnversionedGetHealthCheckBundleActivationSingle(t *testing.T) {
+func TestUnversionedGetHealthCheckOnlyBundlePlugin(t *testing.T) {
 
 	f := newFixture(t)
-	bundleName := "test-bundle"
 
 	// Initialize the server as if a bundle plugin was
 	// configured on the manager.
-	f.server.manager.Register(pluginBundle.Name, &pluginBundle.Plugin{})
-	f.server.bundleStatuses = map[string]*pluginBundle.Status{
-		bundleName: &pluginBundle.Status{Name: bundleName},
-	}
+	f.server.manager.UpdatePluginStatus("bundle", &plugins.Status{State: plugins.StateNotReady})
 
 	// The bundle hasn't been activated yet, expect the health check to fail
 	req := newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
@@ -92,11 +88,40 @@ func TestUnversionedGetHealthCheckBundleActivationSingle(t *testing.T) {
 	}
 
 	// Set the bundle to be activated.
-	status := map[string]*pluginBundle.Status{
-		bundleName: &pluginBundle.Status{},
+	f.server.manager.UpdatePluginStatus("bundle", &plugins.Status{State: plugins.StateOK})
+
+	// The heath check should now respond as healthy
+	req = newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
+	if err := f.executeRequest(req, 200, `{}`); err != nil {
+		t.Fatal(err)
 	}
-	status[bundleName].SetActivateSuccess("")
-	f.server.updateBundleStatus(status)
+}
+
+func TestUnversionedGetHealthCheckDiscoveryWithBundle(t *testing.T) {
+
+	f := newFixture(t)
+
+	// Initialize the server as if a discovery bundle is configured
+	f.server.manager.UpdatePluginStatus("discovery", &plugins.Status{State: plugins.StateNotReady})
+
+	// The discovery bundle hasn't been activated yet, expect the health check to fail
+	req := newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
+	if err := f.executeRequest(req, 500, `{}`); err != nil {
+		t.Fatal(err)
+	}
+
+	// Set the bundle to be not ready (plugin configured and created, but hasn't activated all bundles yet).
+	f.server.manager.UpdatePluginStatus("discovery", &plugins.Status{State: plugins.StateOK})
+	f.server.manager.UpdatePluginStatus("bundle", &plugins.Status{State: plugins.StateNotReady})
+
+	// The discovery bundle is OK, but the newly configured bundle hasn't been activated yet, expect the health check to fail
+	req = newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
+	if err := f.executeRequest(req, 500, `{}`); err != nil {
+		t.Fatal(err)
+	}
+
+	// Set the bundle to be activated.
+	f.server.manager.UpdatePluginStatus("bundle", &plugins.Status{State: plugins.StateOK})
 
 	// The heath check should now respond as healthy
 	req = newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
@@ -113,6 +138,12 @@ func TestUnversionedGetHealthCheckBundleActivationSingleLegacy(t *testing.T) {
 
 	ctx := context.Background()
 
+	// The server doesn't know about any bundles, so return a healthy status
+	req := newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
+	if err := f.executeRequest(req, 200, `{}`); err != nil {
+		t.Fatal(err)
+	}
+
 	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: "a",
@@ -123,130 +154,305 @@ func TestUnversionedGetHealthCheckBundleActivationSingleLegacy(t *testing.T) {
 		t.Fatalf("Unexpected error: %s", err)
 	}
 
-	// The heath check should now respond as healthy
-	req := newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
+	// The heath check still respond as healthy with a legacy bundle found in storage
+	req = newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
 	if err := f.executeRequest(req, 200, `{}`); err != nil {
 		t.Fatal(err)
 	}
 }
 
-func TestUnversionedGetHealthCheckBundleActivationMulti(t *testing.T) {
+func TestBundlesReady(t *testing.T) {
 
+	cases := []struct {
+		note   string
+		status map[string]*plugins.Status
+		ready  bool
+	}{
+		{
+			note:   "nil status",
+			status: nil,
+			ready:  true,
+		},
+		{
+			note:   "empty status",
+			status: map[string]*plugins.Status{},
+			ready:  true,
+		},
+		{
+			note: "discovery not ready - bundle missing",
+			status: map[string]*plugins.Status{
+				"discovery": {State: plugins.StateNotReady},
+			},
+			ready: false,
+		},
+		{
+			note: "discovery ok - bundle missing",
+			status: map[string]*plugins.Status{
+				"discovery": {State: plugins.StateOK},
+			},
+			ready: true, // bundles aren't enabled, only discovery plugin configured
+		},
+		{
+			note: "discovery missing - bundle not ready",
+			status: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateNotReady},
+			},
+			ready: false,
+		},
+		{
+			note: "discovery missing - bundle ok",
+			status: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateOK},
+			},
+			ready: true, // discovery isn't enabled, only bundle plugin configured
+		},
+		{
+			note: "discovery not ready - bundle not ready",
+			status: map[string]*plugins.Status{
+				"discovery": {State: plugins.StateNotReady},
+				"bundle":    {State: plugins.StateNotReady},
+			},
+			ready: false,
+		},
+		{
+			note: "discovery ok - bundle not ready",
+			status: map[string]*plugins.Status{
+				"discovery": {State: plugins.StateOK},
+				"bundle":    {State: plugins.StateNotReady},
+			},
+			ready: false,
+		},
+		{
+			note: "discovery not ready - bundle ok",
+			status: map[string]*plugins.Status{
+				"discovery": {State: plugins.StateNotReady},
+				"bundle":    {State: plugins.StateOK},
+			},
+			ready: false,
+		},
+		{
+			note: "discovery ok - bundle ok",
+			status: map[string]*plugins.Status{
+				"discovery": {State: plugins.StateOK},
+				"bundle":    {State: plugins.StateOK},
+			},
+			ready: true,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			f := newFixture(t)
+
+			actual := f.server.bundlesReady(tc.status)
+			if actual != tc.ready {
+				t.Errorf("Expected %t got %t", tc.ready, actual)
+			}
+		})
+	}
+}
+
+func TestUnversionedGetHealthCheckDiscoveryWithPlugins(t *testing.T) {
+
+	// Use the same server through the cases, the status updates apply incrementally to it.
 	f := newFixture(t)
 
-	// Initialize the server as if a bundle plugin was
-	// configured on the manager.
-	bp := pluginBundle.New(&pluginBundle.Config{Bundles: map[string]*pluginBundle.Source{
-		"b1": {Service: "s1", Resource: "bundle.tar.gz"},
-		"b2": {Service: "s2", Resource: "bundle.tar.gz"},
-		"b3": {Service: "s3", Resource: "bundle.tar.gz"},
-	}}, f.server.manager)
-	f.server.manager.Register(pluginBundle.Name, bp)
-	f.server.bundleStatuses = map[string]*pluginBundle.Status{
-		"b1": {Name: "b1"},
-		"b2": {Name: "b2"},
-		"b3": {Name: "b3"},
+	cases := []struct {
+		note          string
+		statusUpdates map[string]*plugins.Status
+		exp           int
+	}{
+		{
+			note:          "no plugins configured",
+			statusUpdates: nil,
+			exp:           200,
+		},
+		{
+			note: "one plugin configured - not ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "one plugin configured - ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
+		{
+			note: "one plugin configured - error state",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateErr},
+			},
+			exp: 500,
+		},
+		{
+			note: "one plugin configured - recovered from error",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
+		{
+			note: "add second plugin - not ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+				"p2": {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "add third plugin - not ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+				"p2": {State: plugins.StateNotReady},
+				"p3": {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "mixed states - not ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+				"p2": {State: plugins.StateErr},
+				"p3": {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "mixed states - still not ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+				"p2": {State: plugins.StateErr},
+				"p3": {State: plugins.StateOK},
+			},
+			exp: 500,
+		},
+		{
+			note: "all plugins ready",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+				"p2": {State: plugins.StateOK},
+				"p3": {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
+		{
+			note: "one plugins fails",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateErr},
+				"p2": {State: plugins.StateOK},
+				"p3": {State: plugins.StateOK},
+			},
+			exp: 500,
+		},
+
+		{
+			note: "all plugins ready - recovery",
+			statusUpdates: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+				"p2": {State: plugins.StateOK},
+				"p3": {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
 	}
 
-	// No bundle has been activated yet, expect the health check to fail
-	req := newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
-	if err := f.executeRequest(req, 500, `{}`); err != nil {
-		t.Fatal(err)
-	}
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			for name, status := range tc.statusUpdates {
+				f.server.manager.UpdatePluginStatus(name, status)
+			}
 
-	// Set one bundle to be activated
-	update := map[string]*pluginBundle.Status{
-		"b1": {Name: "b1"},
-		"b2": {Name: "b2"},
-		"b3": {Name: "b3"},
-	}
-	update["b2"].SetActivateSuccess("A")
-	f.server.updateBundleStatus(update)
-
-	// The heath check should still respond as unhealthy
-	req = newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
-	if err := f.executeRequest(req, 500, `{}`); err != nil {
-		t.Fatal(err)
-	}
-
-	// Activate all the bundles
-	update["b1"].SetActivateSuccess("B")
-	update["b3"].SetActivateSuccess("C")
-	f.server.updateBundleStatus(update)
-
-	// The heath check should succeed now
-	req = newReqUnversioned(http.MethodGet, "/health?bundle=true", "")
-	if err := f.executeRequest(req, 200, `{}`); err != nil {
-		t.Fatal(err)
+			req := newReqUnversioned(http.MethodGet, "/health?plugins", "")
+			if err := f.executeRequest(req, tc.exp, `{}`); err != nil {
+				t.Fatal(err)
+			}
+		})
 	}
 }
 
-func TestInitWithBundlePlugin(t *testing.T) {
-	store := inmem.New()
-	m, err := plugins.New([]byte{}, "test", store)
-	if err != nil {
-		t.Fatalf("Unexpected error creating plugin manager: %s", err.Error())
+func TestUnversionedGetHealthCheckBundleAndPlugins(t *testing.T) {
+
+	cases := []struct {
+		note     string
+		statuses map[string]*plugins.Status
+		exp      int
+	}{
+		{
+			note:     "no plugins configured",
+			statuses: nil,
+			exp:      200,
+		},
+		{
+			note: "only bundle plugin configured - not ready",
+			statuses: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "only bundle plugin configured - ok",
+			statuses: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
+		{
+			note: "only custom plugin configured - not ready",
+			statuses: map[string]*plugins.Status{
+				"p1": {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "only custom plugin configured - ok",
+			statuses: map[string]*plugins.Status{
+				"p1": {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
+		{
+			note: "both configured - bundle not ready",
+			statuses: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateNotReady},
+				"p1":     {State: plugins.StateOK},
+			},
+			exp: 500,
+		},
+		{
+			note: "both configured - custom plugin not ready",
+			statuses: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateOK},
+				"p1":     {State: plugins.StateNotReady},
+			},
+			exp: 500,
+		},
+		{
+			note: "both configured - both ready",
+			statuses: map[string]*plugins.Status{
+				"bundle": {State: plugins.StateOK},
+				"p1":     {State: plugins.StateOK},
+			},
+			exp: 200,
+		},
 	}
 
-	bundleName := "test-bundle"
-	bundleConf := &pluginBundle.Config{
-		Name:    bundleName,
-		Service: "s1",
-		Bundles: map[string]*pluginBundle.Source{"b1": {}},
-	}
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			f := newFixture(t)
 
-	m.Register(pluginBundle.Name, pluginBundle.New(bundleConf, m))
+			for name, status := range tc.statuses {
+				f.server.manager.UpdatePluginStatus(name, status)
+			}
 
-	server, err := New().
-		WithStore(store).
-		WithManager(m).
-		Init(context.Background())
-
-	if err != nil {
-		t.Fatalf("Unexpected error initializing server: %s", err.Error())
-	}
-
-	if !server.hasBundle() {
-		t.Error("server.hasBundle should be true")
-	}
-
-	isActivated := server.bundlesActivated()
-	if isActivated {
-		t.Error("bundle should not be initialized to activated status")
-	}
-}
-
-func TestInitWithBundlePluginMultiBundle(t *testing.T) {
-	store := inmem.New()
-	m, err := plugins.New([]byte{}, "test", store)
-	if err != nil {
-		t.Fatalf("Unexpected error creating plugin manager: %s", err.Error())
-	}
-
-	bundleConf := &pluginBundle.Config{Bundles: map[string]*pluginBundle.Source{
-		"b1": {},
-		"b2": {},
-		"b3": {},
-	}}
-
-	m.Register(pluginBundle.Name, pluginBundle.New(bundleConf, m))
-
-	server, err := New().
-		WithStore(store).
-		WithManager(m).
-		Init(context.Background())
-
-	if err != nil {
-		t.Fatalf("Unexpected error initializing server: %s", err.Error())
-	}
-
-	if !server.hasBundle() {
-		t.Error("server.hasBundle should be true")
-	}
-
-	isActivated := server.bundlesActivated()
-	if isActivated {
-		t.Error("bundle should not be initialized to activated")
+			req := newReqUnversioned(http.MethodGet, "/health?plugins&bundle", "")
+			if err := f.executeRequest(req, tc.exp, `{}`); err != nil {
+				t.Fatal(err)
+			}
+		})
 	}
 }
 
@@ -1680,9 +1886,6 @@ func TestDataProvenanceSingleBundle(t *testing.T) {
 	// Initialize as if a bundle plugin is running
 	bp := pluginBundle.New(&pluginBundle.Config{Name: "b1"}, f.server.manager)
 	f.server.manager.Register(pluginBundle.Name, bp)
-	f.server.bundleStatuses = map[string]*pluginBundle.Status{
-		"b1": {Name: "b1"},
-	}
 
 	req := newReqV1(http.MethodPost, "/data?provenance", "")
 	f.reset()
@@ -1791,11 +1994,6 @@ func TestDataProvenanceMultiBundle(t *testing.T) {
 	}}, f.server.manager)
 	f.server.manager.Register(pluginBundle.Name, bp)
 
-	f.server.bundleStatuses = map[string]*pluginBundle.Status{
-		"b1": {Name: "b1"},
-		"b2": {Name: "b2"},
-	}
-
 	req := newReqV1(http.MethodPost, "/data?provenance", "")
 	f.reset()
 	f.server.Handler.ServeHTTP(f.recorder, req)