Module id fail

CHANGELOG.md

@@ -19,6 +19,9 @@ Main (unreleased)
 
 - Fixed an issue where `loki.process` validation for stage `metric.counter` was 
   allowing invalid combination of configuration options. (@thampiotr)
+
+- Fixed issue where adding a module after initial start, that failed to load then subsequently resolving the issue would cause the module to
+  permanently fail to load with `id already exists` error. (@mattdurham)
 
 ### Enhancements
 

pkg/flow/internal/controller/loader_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/grafana/agent/component"
 	"github.com/grafana/agent/pkg/flow/internal/controller"
 	"github.com/grafana/agent/pkg/flow/internal/dag"
 	"github.com/grafana/agent/pkg/flow/logging"
@@ -209,7 +210,7 @@ func TestScopeWithFailingComponent(t *testing.T) {
 				OnComponentUpdate: func(cn *controller.ComponentNode) { /* no-op */ },
 				Registerer:        prometheus.NewRegistry(),
 				NewModuleController: func(id string, availableServices []string) controller.ModuleController {
-					return nil
+					return fakeModuleController{}
 				},
 			},
 		}
@@ -304,3 +305,16 @@ func requireGraph(t *testing.T, g *dag.Graph, expect graphDefinition) {
 	}
 	require.ElementsMatch(t, expect.OutEdges, actualEdges, "List of edges do not match")
 }
+
+type fakeModuleController struct{}
+
+func (f fakeModuleController) NewModule(id string, export component.ExportFunc) (component.Module, error) {
+	return nil, nil
+}
+
+func (f fakeModuleController) ModuleIDs() []string {
+	return nil
+}
+
+func (f fakeModuleController) ClearModuleIDs() {
+}

pkg/flow/internal/controller/module.go

@@ -9,4 +9,7 @@ type ModuleController interface {
 
 	// ModuleIDs returns the list of managed modules in unspecified order.
 	ModuleIDs() []string
+
+	// ClearModuleIDs removes all the module id references from the controller.
+	ClearModuleIDs()
 }

pkg/flow/internal/controller/node_component.go

@@ -80,18 +80,19 @@ type ComponentGlobals struct {
 // arguments and exports. ComponentNode manages the arguments for the component
 // from a River block.
 type ComponentNode struct {
-	id                ComponentID
-	globalID          string
-	label             string
-	componentName     string
-	nodeID            string // Cached from id.String() to avoid allocating new strings every time NodeID is called.
-	reg               component.Registration
-	managedOpts       component.Options
-	registry          *prometheus.Registry
-	exportsType       reflect.Type
-	moduleController  ModuleController
-	OnComponentUpdate func(cn *ComponentNode) // Informs controller that we need to reevaluate
-	lastUpdateTime    atomic.Time
+	id                 ComponentID
+	globalID           string
+	label              string
+	componentName      string
+	nodeID             string // Cached from id.String() to avoid allocating new strings every time NodeID is called.
+	reg                component.Registration
+	managedOpts        component.Options
+	registry           *prometheus.Registry
+	exportsType        reflect.Type
+	moduleFailureCheck sync.Once // This is used on first startup to see if module loaded correctly and if not then clean up registry.
+	moduleController   ModuleController
+	OnComponentUpdate  func(cn *ComponentNode) // Informs controller that we need to reevaluate
+	lastUpdateTime     atomic.Time
 
 	mut     sync.RWMutex
 	block   *ast.BlockStmt // Current River block to derive args from
@@ -257,6 +258,13 @@ func (cn *ComponentNode) Evaluate(scope *vm.Scope) error {
 		msg := fmt.Sprintf("component evaluation failed: %s", err)
 		cn.setEvalHealth(component.HealthTypeUnhealthy, msg)
 	}
+	// Check for failure on initial loading, if not then we need to remove any modules it the component may habe created.
+	// In the case where the component is not a module loader this is a noop.
+	cn.moduleFailureCheck.Do(func() {
+		if err != nil {
+			cn.moduleController.ClearModuleIDs()
+		}
+	})
 
 	return err
 }
@@ -312,6 +320,8 @@ func (cn *ComponentNode) Run(ctx context.Context) error {
 	cn.mut.RLock()
 	managed := cn.managed
 	cn.mut.RUnlock()
+	// When finished running clear the module ids.
+	defer cn.moduleController.ClearModuleIDs()
 
 	if managed == nil {
 		return ErrUnevaluated

pkg/flow/module.go

@@ -69,8 +69,19 @@ func (m *moduleController) removeID(id string) {
 	m.mut.Lock()
 	defer m.mut.Unlock()
 
-	delete(m.modules, id)
 	m.o.ModuleRegistry.Unregister(id)
+	delete(m.modules, id)
+}
+
+// ClearModuleIDs removes all ids from the registry.
+func (m *moduleController) ClearModuleIDs() {
+	m.mut.Lock()
+	defer m.mut.Unlock()
+
+	for id := range m.modules {
+		m.o.ModuleRegistry.Unregister(id)
+	}
+	m.modules = make(map[string]struct{})
 }
 
 // ModuleIDs implements [controller.ModuleController].

pkg/flow/module_fail_test.go

@@ -0,0 +1,105 @@
+package flow
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/grafana/agent/component"
+	mod "github.com/grafana/agent/component/module"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIDRemovalIfFailedToLoad(t *testing.T) {
+	f := New(testOptions(t))
+
+	fullContent := "test.fail.module \"t1\" { content = \"\" }"
+	fl, err := ParseSource("test", []byte(fullContent))
+	require.NoError(t, err)
+	err = f.LoadSource(fl, nil)
+	require.NoError(t, err)
+	ctx := context.Background()
+	ctx, cnc := context.WithTimeout(ctx, 600*time.Second)
+
+	go f.Run(ctx)
+	require.Eventually(t, func() bool {
+		t1 := f.loader.Components()[0].Component().(*testFailModule)
+		return t1 != nil
+	}, 10*time.Second, 100*time.Millisecond)
+	t1 := f.loader.Components()[0].Component().(*testFailModule)
+	badContent :=
+		`test.fail.module "int" {
+content=""
+fail=true
+}`
+	err = t1.updateContent(badContent)
+	// Because we have bad content this should fail, but the ids should be removed.
+	require.Error(t, err)
+	goodContent :=
+		`test.fail.module "int" { 
+content=""
+fail=false
+}`
+	err = t1.updateContent(goodContent)
+	require.NoError(t, err)
+	cnc()
+}
+
+func init() {
+	component.Register(component.Registration{
+		Name:    "test.fail.module",
+		Args:    TestFailArguments{},
+		Exports: mod.Exports{},
+
+		Build: func(opts component.Options, args component.Arguments) (component.Component, error) {
+			m, err := mod.NewModuleComponent(opts)
+			if err != nil {
+				return nil, err
+			}
+			if args.(TestFailArguments).Fail {
+				return nil, fmt.Errorf("module told to fail")
+			}
+			err = m.LoadFlowSource(nil, args.(TestFailArguments).Content)
+			if err != nil {
+				return nil, err
+			}
+			return &testFailModule{
+				mc:      m,
+				content: args.(TestFailArguments).Content,
+				opts:    opts,
+				fail:    args.(TestFailArguments).Fail,
+				ch:      make(chan error),
+			}, nil
+		},
+	})
+}
+
+type TestFailArguments struct {
+	Content string `river:"content,attr"`
+	Fail    bool   `river:"fail,attr,optional"`
+}
+
+type testFailModule struct {
+	content string
+	opts    component.Options
+	ch      chan error
+	mc      *mod.ModuleComponent
+	fail    bool
+}
+
+func (t *testFailModule) Run(ctx context.Context) error {
+	go t.mc.RunFlowController(ctx)
+	<-ctx.Done()
+	return nil
+}
+
+func (t *testFailModule) updateContent(content string) error {
+	t.content = content
+	err := t.mc.LoadFlowSource(nil, t.content)
+	return err
+}
+
+func (t *testFailModule) Update(_ component.Arguments) error {
+	return nil
+}

pkg/flow/module_test.go

@@ -207,24 +207,6 @@ func TestIDCollision(t *testing.T) {
 	require.Nil(t, m)
 }
 
-func TestIDRemoval(t *testing.T) {
-	defer verifyNoGoroutineLeaks(t)
-	opts := testModuleControllerOptions(t)
-	defer opts.WorkerPool.Stop()
-	opts.ID = "test"
-	nc := newModuleController(opts)
-	m, err := nc.NewModule("t1", func(exports map[string]any) {})
-	require.NoError(t, err)
-	err = m.LoadConfig([]byte(""), nil)
-	require.NoError(t, err)
-	require.NotNil(t, m)
-	ctx := context.Background()
-	ctx, cncl := context.WithTimeout(ctx, 1*time.Second)
-	defer cncl()
-	m.Run(ctx)
-	require.Len(t, nc.(*moduleController).modules, 0)
-}
-
 func testModuleControllerOptions(t *testing.T) *moduleControllerOptions {
 	t.Helper()