scheduler: add component shutdown timeout of 10m (#4577)

* scheduler: add component shutdown timeout of 10m

* Expose a flag to change deadline

Author: thampiotr

CHANGELOG.md

@@ -48,16 +48,19 @@ Main (unreleased)
 
 - Fix direction of arrows for pyroscope components in UI graph. (@dehaansa)
 
+- Fix an issue where component shutdown could block indefinitely by adding a warning log message and a deadline of 10 minutes. The deadline can be configured with the `--feature.component-shutdown-deadline` flag if the default is not suitable. (@thampiotr)
+
 v1.11.1
 -----------------
 
+### Bugfixes
+
 - Fix potential deadlock in `loki.source.journal` when stopping or reloading the component. (@thampiotr)
 
 - Honor sync timeout when waiting for network availability for prometheus.operator.* components. (@dehaansa)
 
 - Fix `prometheus.exporter.cloudwatch` to not always emit debug logs but respect debug property. (@kalleep)
 
-### Bugfixes
 
 v1.11.0
 -----------------

docs/sources/reference/cli/run.md

@@ -65,6 +65,7 @@ The following flags are supported:
 * `--config.extra-args`: Extra arguments from the original format used by the converter.
 * `--stability.level`: The minimum permitted stability level of functionality. Supported values: `experimental`, `public-preview`, and `generally-available` (default `"generally-available"`).
 * `--feature.community-components.enabled`: Enable community components (default `false`).
+* `--feature.component-shutdown-deadline`: Maximum duration to wait for a component to shut down before giving up and logging an error (default `"10m"`).
 * `--windows.priority`: The priority to set for the {{< param "PRODUCT_NAME" >}} process when running on Windows. This is only available on Windows. Supported values: `above_normal`, `below_normal`, `normal`, `high`, `idle`, or `realtime` (default `"normal"`).
 
 {{< admonition type="note" >}}

internal/alloycli/cmd_run.go

@@ -68,6 +68,7 @@ func runCommand() *cobra.Command {
 		clusterRejoinInterval: 60 * time.Second,
 		disableSupportBundle:  false,
 		windowsPriority:       windowspriority.PriorityNormal,
+		taskShutdownDeadline:  10 * time.Minute,
 	}
 
 	cmd := &cobra.Command{
@@ -166,6 +167,7 @@ depending on the nature of the reload error.
 	if runtime.GOOS == "windows" {
 		cmd.Flags().StringVar(&r.windowsPriority, "windows.priority", r.windowsPriority, fmt.Sprintf("Process priority to use when running on windows. This flag is currently in public preview. Supported values: %s", strings.Join(slices.Collect(windowspriority.PriorityValues()), ", ")))
 	}
+	cmd.Flags().DurationVar(&r.taskShutdownDeadline, "feature.component-shutdown-deadline", r.taskShutdownDeadline, "Maximum duration to wait for a component to shut down before giving up and logging an error")
 
 	addDeprecatedFlags(cmd)
 	return cmd
@@ -201,6 +203,7 @@ type alloyRun struct {
 	enableCommunityComps         bool
 	disableSupportBundle         bool
 	windowsPriority              string
+	taskShutdownDeadline         time.Duration
 }
 
 func (fr *alloyRun) Run(cmd *cobra.Command, configPath string) error {
@@ -384,6 +387,7 @@ func (fr *alloyRun) Run(cmd *cobra.Command, configPath string) error {
 			remoteCfgService,
 			uiService,
 		},
+		TaskShutdownDeadline: fr.taskShutdownDeadline,
 	})
 
 	ready = f.Ready

internal/runtime/alloy.go

@@ -114,6 +114,9 @@ type Options struct {
 
 	// EnableCommunityComps enables the use of community components.
 	EnableCommunityComps bool
+
+	// TaskShutdownDeadline is the maximum duration to wait for a component to shut down before giving up and logging an error.
+	TaskShutdownDeadline time.Duration
 }
 
 // Runtime is the Alloy system.
@@ -137,10 +140,11 @@ type Runtime struct {
 // New creates a new, unstarted Alloy controller. Call Run to run the controller.
 func New(o Options) *Runtime {
 	return newController(controllerOptions{
-		Options:        o,
-		ModuleRegistry: newModuleRegistry(),
-		IsModule:       false, // We are creating a new root controller.
-		WorkerPool:     worker.NewDefaultWorkerPool(),
+		Options:              o,
+		ModuleRegistry:       newModuleRegistry(),
+		IsModule:             false, // We are creating a new root controller.
+		WorkerPool:           worker.NewDefaultWorkerPool(),
+		TaskShutdownDeadline: o.TaskShutdownDeadline,
 	})
 }
 
@@ -154,6 +158,8 @@ type controllerOptions struct {
 	IsModule          bool               // Whether this controller is for a module.
 	// A worker pool to evaluate components asynchronously. A default one will be created if this is nil.
 	WorkerPool worker.Pool
+	// TaskShutdownDeadline is the maximum duration to wait for a component to shut down before giving up and logging an error.
+	TaskShutdownDeadline time.Duration
 }
 
 // newController creates a new, unstarted Alloy controller with a specific
@@ -186,7 +192,7 @@ func newController(o controllerOptions) *Runtime {
 		opts:   o,
 
 		updateQueue: controller.NewQueue(),
-		sched:       controller.NewScheduler(log),
+		sched:       controller.NewScheduler(log, o.TaskShutdownDeadline),
 
 		modules: o.ModuleRegistry,
 

internal/runtime/internal/controller/scheduler.go

@@ -4,12 +4,19 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/go-kit/log"
 
 	"github.com/grafana/alloy/internal/runtime/logging/level"
 )
 
+var (
+	// TaskShutdownWarningTimeout is the duration after which a warning is logged
+	// when a task is taking too long to shut down.
+	TaskShutdownWarningTimeout = time.Minute
+)
+
 // RunnableNode is any BlockNode which can also be run.
 type RunnableNode interface {
 	BlockNode
@@ -18,10 +25,11 @@ type RunnableNode interface {
 
 // Scheduler runs components.
 type Scheduler struct {
-	ctx     context.Context
-	cancel  context.CancelFunc
-	running sync.WaitGroup
-	logger  log.Logger
+	ctx                  context.Context
+	cancel               context.CancelFunc
+	running              sync.WaitGroup
+	logger               log.Logger
+	taskShutdownDeadline time.Duration
 
 	tasksMut sync.Mutex
 	tasks    map[string]*task
@@ -31,12 +39,13 @@ type Scheduler struct {
 // components which are running.
 //
 // Call Close to stop the Scheduler and all running components.
-func NewScheduler(logger log.Logger) *Scheduler {
+func NewScheduler(logger log.Logger, taskShutdownDeadline time.Duration) *Scheduler {
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Scheduler{
-		ctx:    ctx,
-		cancel: cancel,
-		logger: logger,
+		ctx:                  ctx,
+		cancel:               cancel,
+		logger:               logger,
+		taskShutdownDeadline: taskShutdownDeadline,
 
 		tasks: make(map[string]*task),
 	}
@@ -52,7 +61,6 @@ func NewScheduler(logger log.Logger) *Scheduler {
 // call to Synchronize.
 func (s *Scheduler) Synchronize(rr []RunnableNode) error {
 	s.tasksMut.Lock()
-	defer s.tasksMut.Unlock()
 
 	if s.ctx.Err() != nil {
 		return fmt.Errorf("Scheduler is closed")
@@ -89,9 +97,9 @@ func (s *Scheduler) Synchronize(rr []RunnableNode) error {
 		)
 
 		opts := taskOptions{
-			Context:  s.ctx,
-			Runnable: newRunnable,
-			OnDone: func(err error) {
+			context:  s.ctx,
+			runnable: newRunnable,
+			onDone: func(err error) {
 				defer s.running.Done()
 
 				if err != nil {
@@ -104,12 +112,16 @@ func (s *Scheduler) Synchronize(rr []RunnableNode) error {
 				defer s.tasksMut.Unlock()
 				delete(s.tasks, nodeID)
 			},
+			logger:               log.With(s.logger, "taskID", nodeID),
+			taskShutdownDeadline: s.taskShutdownDeadline,
 		}
 
 		s.running.Add(1)
 		s.tasks[nodeID] = newTask(opts)
 	}
 
+	// Unlock the tasks mutex so that Stop calls can complete.
+	s.tasksMut.Unlock()
 	// Wait for all stopping runnables to exit.
 	stopping.Wait()
 	return nil
@@ -125,36 +137,64 @@ func (s *Scheduler) Close() error {
 
 // task is a scheduled runnable.
 type task struct {
-	ctx    context.Context
-	cancel context.CancelFunc
-	exited chan struct{}
+	ctx      context.Context
+	cancel   context.CancelFunc
+	exited   chan struct{}
+	opts     taskOptions
+	doneOnce sync.Once
 }
 
 type taskOptions struct {
-	Context  context.Context
-	Runnable RunnableNode
-	OnDone   func(error)
+	context              context.Context
+	runnable             RunnableNode
+	onDone               func(error)
+	logger               log.Logger
+	taskShutdownDeadline time.Duration
 }
 
 // newTask creates and starts a new task.
 func newTask(opts taskOptions) *task {
-	ctx, cancel := context.WithCancel(opts.Context)
+	ctx, cancel := context.WithCancel(opts.context)
 
 	t := &task{
 		ctx:    ctx,
 		cancel: cancel,
 		exited: make(chan struct{}),
+		opts:   opts,
 	}
 
 	go func() {
-		err := opts.Runnable.Run(t.ctx)
+		err := opts.runnable.Run(t.ctx)
 		close(t.exited)
-		opts.OnDone(err)
+		t.doneOnce.Do(func() {
+			t.opts.onDone(err)
+		})
 	}()
 	return t
 }
 
 func (t *task) Stop() {
 	t.cancel()
-	<-t.exited
+
+	deadlineDuration := t.opts.taskShutdownDeadline
+	if deadlineDuration == 0 {
+		deadlineDuration = time.Hour * 24 * 365 * 100 // infinite timeout ~= 100 years
+	}
+
+	deadlineCtx, deadlineCancel := context.WithTimeout(context.Background(), deadlineDuration)
+	defer deadlineCancel()
+
+	for {
+		select {
+		case <-t.exited:
+			return // Task exited normally.
+		case <-time.After(TaskShutdownWarningTimeout):
+			level.Warn(t.opts.logger).Log("msg", "task shutdown is taking longer than expected")
+		case <-deadlineCtx.Done():
+			t.doneOnce.Do(func() {
+				t.opts.onDone(fmt.Errorf("task shutdown deadline exceeded"))
+			})
+			return // Task took too long to exit, don't wait.
+		}
+	}
 }

internal/runtime/internal/controller/scheduler_test.go

@@ -1,10 +1,12 @@
 package controller_test
 
 import (
+	"bytes"
 	"context"
 	"os"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/stretchr/testify/require"
@@ -30,7 +32,7 @@ func TestScheduler_Synchronize(t *testing.T) {
 			return nil
 		}
 
-		sched := controller.NewScheduler(logger)
+		sched := controller.NewScheduler(logger, 1*time.Minute)
 		sched.Synchronize([]controller.RunnableNode{
 			fakeRunnable{ID: "component-a", Component: mockComponent{RunFunc: runFunc}},
 			fakeRunnable{ID: "component-b", Component: mockComponent{RunFunc: runFunc}},
@@ -52,7 +54,7 @@ func TestScheduler_Synchronize(t *testing.T) {
 			return nil
 		}
 
-		sched := controller.NewScheduler(logger)
+		sched := controller.NewScheduler(logger, 1*time.Minute)
 
 		for i := 0; i < 10; i++ {
 			// If a new runnable is created, runFunc will panic since the WaitGroup
@@ -78,7 +80,7 @@ func TestScheduler_Synchronize(t *testing.T) {
 			return nil
 		}
 
-		sched := controller.NewScheduler(logger)
+		sched := controller.NewScheduler(logger, 1*time.Minute)
 
 		sched.Synchronize([]controller.RunnableNode{
 			fakeRunnable{ID: "component-a", Component: mockComponent{RunFunc: runFunc}},
@@ -114,3 +116,54 @@ var _ component.Component = (*mockComponent)(nil)
 
 func (mc mockComponent) Run(ctx context.Context) error              { return mc.RunFunc(ctx) }
 func (mc mockComponent) Update(newConfig component.Arguments) error { return mc.UpdateFunc(newConfig) }
+
+func TestScheduler_TaskTimeoutLogging(t *testing.T) {
+	// Temporarily modify timeout values for testing
+	originalWarningTimeout := controller.TaskShutdownWarningTimeout
+	controller.TaskShutdownWarningTimeout = 50 * time.Millisecond
+	defer func() {
+		controller.TaskShutdownWarningTimeout = originalWarningTimeout
+	}()
+
+	// Create a buffer to capture log output
+	var logBuffer bytes.Buffer
+	logger := log.NewLogfmtLogger(&logBuffer)
+
+	var started sync.WaitGroup
+	started.Add(1)
+
+	// Create a component that will block and not respond to context cancellation
+	runFunc := func(ctx context.Context) error {
+		started.Done()
+		// Block indefinitely, ignoring context cancellation
+		// Use a long sleep to simulate a component that doesn't respond to cancellation
+		time.Sleep(1 * time.Second)
+		return nil
+	}
+
+	sched := controller.NewScheduler(logger, 150*time.Millisecond)
+
+	// Start a component
+	err := sched.Synchronize([]controller.RunnableNode{
+		fakeRunnable{ID: "blocking-component", Component: mockComponent{RunFunc: runFunc}},
+	})
+	require.NoError(t, err)
+	started.Wait()
+
+	// Remove the component, which should trigger the timeout behavior. This will block until the component exits.
+	err = sched.Synchronize([]controller.RunnableNode{})
+	require.NoError(t, err)
+
+	logOutput := logBuffer.String()
+	t.Logf("actual log output:\n%s", logOutput)
+
+	// Should contain warning message
+	require.Contains(t, logOutput, "task shutdown is taking longer than expected")
+	require.Contains(t, logOutput, "level=warn")
+
+	// Should contain error message
+	require.Contains(t, logOutput, "task shutdown deadline exceeded")
+	require.Contains(t, logOutput, "level=error")
+
+	require.NoError(t, sched.Close())
+}