From 1570f972051381b7510dcecfa35da9f30dd8f38c Mon Sep 17 00:00:00 2001
From: William Dumont <william.dumont@grafana.com>
Date: Tue, 2 Apr 2024 17:06:37 +0200
Subject: [PATCH] Fix custom component registry access (#6811)

* lock before before evaluating to prevent using a not fully initialized custom component registry

* update changelog

* add test that was creating a panic
---
 CHANGELOG.md                                  |  2 +
 internal/flow/internal/controller/loader.go   | 33 +++++-----
 .../testdata/import_file/import_file_17.txtar | 65 +++++++++++++++++++
 3 files changed, 84 insertions(+), 16 deletions(-)
 create mode 100644 internal/flow/testdata/import_file/import_file_17.txtar

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2409fd3568bc..db1ccf5a9a39 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,6 +53,8 @@ Main (unreleased)
   whenever that argument is explicitly configured. This issue only affected a
   small subset of arguments across 15 components. (@erikbaranowski, @rfratto)
 
+- Fix a bug where a panic could occur when reloading custom components. (@wildum)
+
 ### Other changes
 
 - Clustering for Grafana Agent in Flow mode has graduated from beta to stable.
diff --git a/internal/flow/internal/controller/loader.go b/internal/flow/internal/controller/loader.go
index 8921d5ff1859..d9b328160497 100644
--- a/internal/flow/internal/controller/loader.go
+++ b/internal/flow/internal/controller/loader.go
@@ -37,21 +37,21 @@ type Loader struct {
 	// pool for evaluation in EvaluateDependants, because the queue is full. This is an unlikely scenario, but when
 	// it happens we should avoid retrying too often to give other goroutines a chance to progress. Having a backoff
 	// also prevents log spamming with errors.
-	backoffConfig        backoff.Config
+	backoffConfig backoff.Config
+
+	mut                  sync.RWMutex
+	graph                *dag.Graph
+	originalGraph        *dag.Graph
+	componentNodes       []ComponentNode
+	declareNodes         map[string]*DeclareNode
+	importConfigNodes    map[string]*ImportConfigNode
+	serviceNodes         []*ServiceNode
+	cache                *valueCache
+	blocks               []*ast.BlockStmt // Most recently loaded blocks, used for writing
+	cm                   *controllerMetrics
+	cc                   *controllerCollector
+	moduleExportIndex    int
 	componentNodeManager *ComponentNodeManager
-
-	mut               sync.RWMutex
-	graph             *dag.Graph
-	originalGraph     *dag.Graph
-	componentNodes    []ComponentNode
-	declareNodes      map[string]*DeclareNode
-	importConfigNodes map[string]*ImportConfigNode
-	serviceNodes      []*ServiceNode
-	cache             *valueCache
-	blocks            []*ast.BlockStmt // Most recently loaded blocks, used for writing
-	cm                *controllerMetrics
-	cc                *controllerCollector
-	moduleExportIndex int
 }
 
 // LoaderOptions holds options for creating a Loader.
@@ -787,10 +787,11 @@ func (l *Loader) concurrentEvalFn(n dag.Node, spanCtx context.Context, tracer tr
 	switch n := n.(type) {
 	case BlockNode:
 		ectx := l.cache.BuildContext()
-		evalErr := n.Evaluate(ectx)
 
-		// Only obtain loader lock after we have evaluated the node, allowing for concurrent evaluation.
+		// RLock before evaluate to prevent Evaluating while the config is being reloaded
 		l.mut.RLock()
+		evalErr := n.Evaluate(ectx)
+
 		err = l.postEvaluate(l.log, n, evalErr)
 
 		// Additional post-evaluation steps necessary for module exports.
diff --git a/internal/flow/testdata/import_file/import_file_17.txtar b/internal/flow/testdata/import_file/import_file_17.txtar
new file mode 100644
index 000000000000..8313fa642d33
--- /dev/null
+++ b/internal/flow/testdata/import_file/import_file_17.txtar
@@ -0,0 +1,65 @@
+This is a simple test with several levels of nesting.
+This test would often panic when re-evaluation could run concurrently with a config reload.
+Now that it is protected with mutex, this test should always pass.
+
+-- main.river --
+import.file "a_namespace" {
+	filename = "module.river"
+}
+
+a_namespace.a "default" {
+	a_argument = 47
+}
+
+testcomponents.summation "sum" {
+	input = 10
+}
+
+-- module.river --
+declare "a" {
+	import.file "b_namespace" {
+		filename = "nested_module.river"
+	}
+
+	argument "a_argument" {
+		comment = "Where to send collected metrics."
+	}
+
+	b_namespace.b "default" {
+		b_argument = 147
+	}
+
+	export "a_export" {
+		value = argument.a_argument.value + b_namespace.b.default.b_export + 1
+	}
+}
+
+-- nested_module.river --
+declare "b" {
+	import.file "c_namespace" {
+		filename = "other_nested_module.river"
+	}
+
+	argument "b_argument" {
+		optional = false
+	}
+
+	c_namespace.c "default" {
+		c_argument = 101
+	}
+
+	export "b_export" {
+		value = argument.b_argument.value + c_namespace.c.default.c_export + 1
+	}
+}
+
+-- other_nested_module.river --
+declare "c" {
+	argument "c_argument" {
+		optional = false
+	}
+
+	export "c_export" {
+		value = argument.c_argument.value + 1
+	}
+}
\ No newline at end of file