🐛 remove dead providers from coordinator (#4010)

* 🐛 remove dead providers from coordinator Signed-off-by: Ivan Milchev <[email protected]> * log heartbeat errors Signed-off-by: Ivan Milchev <[email protected]> --------- Signed-off-by: Ivan Milchev <[email protected]>
mondoohq · May 16, 2024 · bdc7ec4 · bdc7ec4
1 parent ff2ea75
commit bdc7ec4
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 0 deletions.
diff --git a/providers/coordinator.go b/providers/coordinator.go
@@ -211,6 +211,14 @@ func (c *coordinator) RemoveRuntime(runtime *Runtime) {
 				log.Warn().Err(err).Str("provider", p.Name).Msg("failed to shut down provider")
 			}
 		}
+	} else {
+		// Check for killed/crashed providers and remove them from the list of running providers
+		for _, p := range c.runningByID {
+			if p.isCloseOrShutdown() {
+				log.Warn().Str("provider", p.Name).Msg("removing closed provider")
+				delete(c.runningByID, p.ID)
+			}
+		}
 	}
 
 	// If all providers have been killed, reset the connection IDs back to 0

diff --git a/providers/coordinator_test.go b/providers/coordinator_test.go
@@ -149,6 +149,60 @@ func TestRemoveRuntime_StopUnusedProvider(t *testing.T) {
 	assert.Empty(t, c.runningByID)
 }
 
+func TestRemoveRuntime_RemoveDeadProvider(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	// Setup 1 closed provider with 1 runtime
+	mockPlugin1 := NewMockProviderPlugin(ctrl)
+	p1 := &RunningProvider{
+		ID:       "provider1",
+		Plugin:   mockPlugin1,
+		isClosed: true,
+	}
+	r1 := &Runtime{
+		providers: map[string]*ConnectedProvider{
+			"provider1": {Instance: p1},
+		},
+		Provider: &ConnectedProvider{
+			Instance: p1,
+			Connection: &pp.ConnectRes{
+				Asset: &inventory.Asset{PlatformIds: []string{"platformId1"}},
+			},
+		},
+	}
+
+	// Setup another provider with another runtime
+	r2 := &Runtime{
+		providers: map[string]*ConnectedProvider{
+			"provider2": {Instance: p1},
+		},
+		Provider: &ConnectedProvider{
+			Instance: p1,
+			Connection: &pp.ConnectRes{
+				Asset: &inventory.Asset{PlatformIds: []string{"platformId2"}},
+			},
+		},
+	}
+
+	c := &coordinator{
+		runningByID: map[string]*RunningProvider{
+			"provider1": p1,
+		},
+		runtimes: map[string]*Runtime{
+			"platformId1": r1,
+			"platformId2": r2,
+		},
+	}
+
+	// Remove 1 runtime
+	c.RemoveRuntime(r1)
+
+	// Verify that the provider has been removed because it crashed
+	assert.Empty(t, c.runningByID)
+
+	c.RemoveRuntime(r2)
+}
+
 func TestRemoveRuntime_UsedProvider(t *testing.T) {
 	ctrl := gomock.NewController(t)
 

diff --git a/providers/running_provider.go b/providers/running_provider.go
@@ -38,6 +38,7 @@ type RunningProvider struct {
 // initialize the heartbeat with the provider
 func (p *RunningProvider) heartbeat() error {
 	if err := p.doOneHeartbeat(p.interval + p.gracePeriod); err != nil {
+		log.Error().Err(err).Str("plugin", p.Name).Msg("error in plugin heartbeat")
 		if err := p.Shutdown(); err != nil {
 			log.Error().Err(err).Str("plugin", p.Name).Msg("error in plugin shutdown")
 		}
@@ -47,6 +48,7 @@ func (p *RunningProvider) heartbeat() error {
 	go func() {
 		for !p.isCloseOrShutdown() {
 			if err := p.doOneHeartbeat(p.interval + p.gracePeriod); err != nil {
+				log.Error().Err(err).Str("plugin", p.Name).Msg("error in plugin heartbeat")
 				if err := p.Shutdown(); err != nil {
 					log.Error().Err(err).Str("plugin", p.Name).Msg("error in plugin shutdown")
 				}

diff --git a/providers/runtime.go b/providers/runtime.go
@@ -204,6 +204,11 @@ func (r *Runtime) Connect(req *plugin.ConnectReq) error {
 		runtime: r,
 	}
 
+	// TODO: we probably want to check here if the provider is dead and restart it
+	// if r.Provider.Instance.isCloseOrShutdown() {
+
+	// }
+
 	r.Provider.Connection, r.Provider.ConnectionError = r.Provider.Instance.Plugin.Connect(req, &callbacks)
 	if r.Provider.ConnectionError != nil {
 		return r.Provider.ConnectionError