Collector Run will listen for context done (#4954)

* Collector Run will listen for context done Signed-off-by: Corbin Phelps <[email protected]> * Updated changelog with PR number Signed-off-by: Corbin Phelps <[email protected]>
open-telemetry · Mar 3, 2022 · 93272a5 · 93272a5
1 parent c257aba
commit 93272a5
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,10 @@
 
 - Remove `Type` funcs in pdata (#4933)
 
+## 🧰 Bug fixes 🧰
+
+- Collector `Run` will now exit when a context cancels (#4954)
+
 ## v0.46.0 Beta
 
 ### 🛑 Breaking changes 🛑

diff --git a/service/collector.go b/service/collector.go
@@ -174,6 +174,10 @@ LOOP:
 		case <-col.shutdownChan:
 			col.logger.Info("Received shutdown request")
 			break LOOP
+		case <-ctx.Done():
+			col.logger.Error("Context done, terminating process", zap.Error(ctx.Err()))
+			// Call shutdown with background context as the passed in context has been canceled
+			return col.shutdown(context.Background())
 		}
 	}
 	return col.shutdown(ctx)

diff --git a/service/collector_test.go b/service/collector_test.go
@@ -253,6 +253,43 @@ func TestCollector_ReportError(t *testing.T) {
 	assert.Equal(t, Closed, col.GetState())
 }
 
+// TestCollector_ContextCancel tests that the collector gracefully exits on context cancel
+func TestCollector_ContextCancel(t *testing.T) {
+	// use a mock AppTelemetry struct to return an error on shutdown
+	preservedAppTelemetry := collectorTelemetry
+	collectorTelemetry = &colTelemetry{}
+	defer func() { collectorTelemetry = preservedAppTelemetry }()
+
+	factories, err := testcomponents.NewDefaultFactories()
+	require.NoError(t, err)
+
+	set := CollectorSettings{
+		BuildInfo: component.NewDefaultBuildInfo(),
+		Factories: factories,
+		ConfigProvider: MustNewDefaultConfigProvider([]string{filepath.Join("testdata", "otelcol-config.yaml")},
+			[]string{"service.telemetry.metrics.address=localhost:" + strconv.FormatUint(uint64(testutil.GetAvailablePort(t)), 10)}),
+	}
+	col, err := New(set)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	colDone := make(chan struct{})
+	go func() {
+		defer close(colDone)
+		require.NoError(t, col.Run(ctx))
+	}()
+
+	assert.Eventually(t, func() bool {
+		return Running == col.GetState()
+	}, 2*time.Second, 200*time.Millisecond)
+
+	cancel()
+
+	<-colDone
+	assert.Equal(t, Closed, col.GetState())
+}
+
 func assertMetrics(t *testing.T, metricsPort uint16, mandatoryLabels []string) {
 	client := &http.Client{}
 	resp, err := client.Get(fmt.Sprintf("http://localhost:%d/metrics", metricsPort))