Extract common function to drain handler

Author: kalleep

internal/component/loki/source/docker/docker.go

@@ -156,11 +156,14 @@ func (c *Component) Run(ctx context.Context) error {
 	defer c.posFile.Stop()
 
 	defer func() {
-		c.mut.Lock()
-		defer c.mut.Unlock()
-
-		// FIXME: We need to drain here
-		c.scheduler.Stop()
+		c.posFile.Stop()
+
+		// Start black hole drain routine to prevent deadlock when we call c.scheduler.Stop().
+		source.Drain(c.handler, func() {
+			c.mut.Lock()
+			defer c.mut.Unlock()
+			c.scheduler.Stop()
+		})
 	}()
 
 	for {

internal/component/loki/source/drain.go

@@ -0,0 +1,33 @@
+package source
+
+import (
+	"context"
+
+	"github.com/grafana/alloy/internal/component/common/loki"
+)
+
+// Drain consumes log entries from recv in a background goroutine while f executes.
+// This prevents deadlocks that can occur when stopping components that may still be
+// sending entries to the receiver channel. The draining goroutine will continue
+// consuming entries until f returns, at which point the context is cancelled and
+// the goroutine exits.
+//
+// This is typically used during component shutdown to drain any remaining entries
+// from a receiver channel while performing cleanup operations (e.g., stopping
+// sources, closing channels).
+func Drain(recv loki.LogsReceiver, f func()) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-recv.Chan():
+				// Consume and discard entries to prevent channel blocking
+			}
+		}
+	}()
+
+	f()
+}

internal/component/loki/source/file/file.go

@@ -203,27 +203,18 @@ func New(o component.Options, args Arguments) (*Component, error) {
 func (c *Component) Run(ctx context.Context) error {
 	defer func() {
 		level.Info(c.opts.Logger).Log("msg", "loki.source.file component shutting down, stopping sources and positions file")
-
 		// We need to stop posFile first so we don't record entries we are draining
 		c.posFile.Stop()
-		// Start black hole drain routine to prevent deadlock when we call c.t.Stop().
-		drainCtx, cancelDrain := context.WithCancel(context.Background())
-		defer cancelDrain()
-		go func() {
-			for {
-				select {
-				case <-drainCtx.Done():
-					return
-				case <-c.handler.Chan(): // Ignore the remaining entries
-				}
-			}
-		}()
-		c.mut.Lock()
-		c.stopping.Store(true)
-		c.watcher.Stop()
-		c.scheduler.Stop()
-		close(c.handler.Chan())
-		c.mut.Unlock()
+
+		// Start black hole drain routine to prevent deadlock when we call c.scheduler.Stop().
+		source.Drain(c.handler, func() {
+			c.mut.Lock()
+			c.stopping.Store(true)
+			c.watcher.Stop()
+			c.scheduler.Stop()
+			close(c.handler.Chan())
+			c.mut.Unlock()
+		})
 	}()
 
 	var wg sync.WaitGroup