loki.source.docker: refactor to use shared scheduler

Author: kalleep

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

@@ -25,7 +25,7 @@ import (
 	"github.com/grafana/alloy/internal/component/common/loki"
 	alloy_relabel "github.com/grafana/alloy/internal/component/common/relabel"
 	"github.com/grafana/alloy/internal/component/discovery"
-	dt "github.com/grafana/alloy/internal/component/loki/source/docker/internal/dockertarget"
+	"github.com/grafana/alloy/internal/component/loki/source"
 	"github.com/grafana/alloy/internal/component/loki/source/internal/positions"
 	"github.com/grafana/alloy/internal/featuregate"
 	"github.com/grafana/alloy/internal/runtime/logging/level"
@@ -47,9 +47,11 @@ func init() {
 var userAgent = useragent.Get()
 
 const (
+	// See github.com/prometheus/prometheus/discovery/moby
 	dockerLabel                = model.MetaLabelPrefix + "docker_"
 	dockerLabelContainerPrefix = dockerLabel + "container_"
 	dockerLabelContainerID     = dockerLabelContainerPrefix + "id"
+	dockerLabelLogStream       = dockerLabelContainerPrefix + "log_stream"
 )
 
 // Arguments holds values which are used to configure the loki.source.docker
@@ -102,16 +104,15 @@ var (
 // Component implements the loki.source.file component.
 type Component struct {
 	opts    component.Options
-	metrics *dt.Metrics
+	metrics *metrics
 
-	mut           sync.RWMutex
-	args          Arguments
-	manager       *manager
-	lastOptions   *options
-	handler       loki.LogsReceiver
-	posFile       positions.Positions
-	rcs           []*relabel.Config
-	defaultLabels model.LabelSet
+	mut       sync.RWMutex
+	args      Arguments
+	scheduler *source.Scheduler[string]
+	client    client.APIClient
+	handler   loki.LogsReceiver
+	posFile   positions.Positions
+	rcs       []*relabel.Config
 
 	receiversMut sync.RWMutex
 	receivers    []loki.LogsReceiver
@@ -134,11 +135,10 @@ func New(o component.Options, args Arguments) (*Component, error) {
 	}
 
 	c := &Component{
-		opts:    o,
-		metrics: dt.NewMetrics(o.Registerer),
-
+		opts:      o,
+		metrics:   newMetrics(o.Registerer),
 		handler:   loki.NewLogsReceiver(),
-		manager:   newManager(o.Logger, nil),
+		scheduler: source.NewScheduler[string](),
 		receivers: args.ForwardTo,
 		posFile:   positionsFile,
 	}
@@ -159,11 +159,8 @@ func (c *Component) Run(ctx context.Context) error {
 		c.mut.Lock()
 		defer c.mut.Unlock()
 
-		// Guard for safety, but it's not possible for Run to be called without
-		// c.tailer being initialized.
-		if c.manager != nil {
-			c.manager.stop()
-		}
+		// FIXME: We need to drain here
+		c.scheduler.Stop()
 	}()
 
 	for {
@@ -175,7 +172,11 @@ func (c *Component) Run(ctx context.Context) error {
 			receivers := c.receivers
 			c.receiversMut.RUnlock()
 			for _, receiver := range receivers {
-				receiver.Chan() <- entry
+				select {
+				case <-ctx.Done():
+					return nil
+				case receiver.Chan() <- entry:
+				}
 			}
 		}
 	}
@@ -198,33 +199,22 @@ func (c *Component) Update(args component.Arguments) error {
 	c.mut.Lock()
 	defer c.mut.Unlock()
 
-	managerOpts, err := c.getManagerOptions(newArgs)
+	opts, err := c.getClient(newArgs)
 	if err != nil {
 		return err
 	}
 
-	if managerOpts != c.lastOptions {
-		// Options changed; pass it to the tailer.
-		// This will never fail because it only fails if the context gets canceled.
-		_ = c.manager.updateOptions(context.Background(), managerOpts)
-		c.lastOptions = managerOpts
+	if opts != c.client {
+		// Stop all tailers so all will be restarted
+		c.scheduler.Stop()
 	}
 
 	defaultLabels := make(model.LabelSet, len(newArgs.Labels))
 	for k, v := range newArgs.Labels {
 		defaultLabels[model.LabelName(k)] = model.LabelValue(v)
 	}
-	c.defaultLabels = defaultLabels
-
-	if len(newArgs.RelabelRules) > 0 {
-		c.rcs = alloy_relabel.ComponentToPromRelabelConfigs(newArgs.RelabelRules)
-	} else {
-		c.rcs = []*relabel.Config{}
-	}
 
-	// Convert input targets into targets to give to tailer.
-	targets := make([]*dt.Target, 0, len(newArgs.Targets))
-	seenTargets := make(map[string]struct{}, len(newArgs.Targets))
+	c.rcs = alloy_relabel.ComponentToPromRelabelConfigs(newArgs.RelabelRules)
 
 	promTargets := make([]promTarget, len(newArgs.Targets))
 	for i, target := range newArgs.Targets {
@@ -238,53 +228,67 @@ func (c *Component) Update(args component.Arguments) error {
 		return promTargets[i].fingerPrint < promTargets[j].fingerPrint
 	})
 
+	shouldRun := make(map[string]struct{}, len(newArgs.Targets))
 	for _, markedTarget := range promTargets {
 		containerID, ok := markedTarget.labels[dockerLabelContainerID]
 		if !ok {
 			level.Debug(c.opts.Logger).Log("msg", "docker target did not include container ID label:"+dockerLabelContainerID)
 			continue
 		}
-		if _, seen := seenTargets[string(containerID)]; seen {
+
+		if c.scheduler.Contains(string(containerID)) {
 			continue
 		}
-		seenTargets[string(containerID)] = struct{}{}
 
-		tgt, err := dt.NewTarget(
+		tailer, err := newTailer(
 			c.metrics,
 			log.With(c.opts.Logger, "target", fmt.Sprintf("docker/%s", containerID)),
-			c.manager.opts.handler,
-			c.manager.opts.positions,
+			c.handler,
+			c.posFile,
 			string(containerID),
-			markedTarget.labels.Merge(c.defaultLabels),
+			markedTarget.labels.Merge(defaultLabels),
 			c.rcs,
-			c.manager.opts.client,
+			opts,
+			5*time.Second,
 		)
+
 		if err != nil {
 			return err
 		}
-		targets = append(targets, tgt)
+
+		shouldRun[tailer.Key()] = struct{}{}
+		c.scheduler.ScheduleSource(tailer)
 	}
 
-	// This will never fail because it only fails if the context gets canceled.
-	_ = c.manager.syncTargets(context.Background(), targets)
+	// Avoid mutating the scheduler state during iteration. Collect sources to
+	// remove and stop them in a separate loop.
+	var toDelete []source.Source[string]
+	for source := range c.scheduler.Sources() {
+		if _, ok := shouldRun[source.Key()]; ok {
+			continue
+		}
+		toDelete = append(toDelete, source)
+	}
+
+	for _, s := range toDelete {
+		c.scheduler.StopSource(s) // stops without blocking
+	}
 
 	c.args = newArgs
 	return nil
 }
 
-// getTailerOptions gets tailer options from arguments. If args hasn't changed
-// from the last call to getTailerOptions, c.lastOptions is returned.
-// c.lastOptions must be updated by the caller.
-//
-// getTailerOptions must only be called when c.mut is held.
-func (c *Component) getManagerOptions(args Arguments) (*options, error) {
-	if reflect.DeepEqual(c.args.Host, args.Host) && c.lastOptions != nil {
-		return c.lastOptions, nil
+// getClient created client from args. If args hasn't changed
+// from the last call to getClient, c.client is returned.
+// getClient must only be called when c.mut is held.
+func (c *Component) getClient(args Arguments) (client.APIClient, error) {
+	if reflect.DeepEqual(c.args.Host, args.Host) && c.client != nil {
+		return c.client, nil
 	}
 
 	hostURL, err := url.Parse(args.Host)
 	if err != nil {
-		return c.lastOptions, err
+		return c.client, err
 	}
 
 	opts := []client.Opt{
@@ -298,7 +302,7 @@ func (c *Component) getManagerOptions(args Arguments) (*options, error) {
 	if hostURL.Scheme == "http" || hostURL.Scheme == "https" {
 		rt, err := config.NewRoundTripperFromConfig(*args.HTTPClientConfig.Convert(), "docker_sd")
 		if err != nil {
-			return c.lastOptions, err
+			return c.client, err
 		}
 		opts = append(opts,
 			client.WithHTTPClient(&http.Client{
@@ -315,41 +319,33 @@ func (c *Component) getManagerOptions(args Arguments) (*options, error) {
 	client, err := client.NewClientWithOpts(opts...)
 	if err != nil {
 		level.Error(c.opts.Logger).Log("msg", "could not create new Docker client", "err", err)
-		return c.lastOptions, fmt.Errorf("failed to build docker client: %w", err)
+		return c.client, fmt.Errorf("failed to build docker client: %w", err)
 	}
 
-	return &options{
-		client:                client,
-		handler:               loki.NewEntryHandler(c.handler.Chan(), func() {}),
-		positions:             c.posFile,
-		targetRestartInterval: 5 * time.Second,
-	}, nil
+	return client, nil
 }
 
 // DebugInfo returns information about the status of tailed targets.
 func (c *Component) DebugInfo() interface{} {
+	c.mut.RLock()
+	defer c.mut.RUnlock()
+
 	var res readerDebugInfo
-	for _, tgt := range c.manager.targets() {
-		details := tgt.Details()
-		res.TargetsInfo = append(res.TargetsInfo, targetInfo{
-			Labels:     tgt.LabelsStr(),
-			ID:         details["id"],
-			LastError:  details["error"],
-			IsRunning:  details["running"],
-			ReadOffset: details["position"],
-		})
+	for s := range c.scheduler.Sources() {
+		ds := s.(source.DebugSource[string])
+		res.TargetsInfo = append(res.TargetsInfo, ds.DebugInfo())
 	}
 	return res
 }
 
 type readerDebugInfo struct {
-	TargetsInfo []targetInfo `alloy:"targets_info,block"`
+	TargetsInfo []any `alloy:"targets_info,block"`
 }
 
-type targetInfo struct {
+type sourceInfo struct {
 	ID         string `alloy:"id,attr"`
 	LastError  string `alloy:"last_error,attr"`
 	Labels     string `alloy:"labels,attr"`
-	IsRunning  string `alloy:"is_running,attr"`
+	IsRunning  bool   `alloy:"is_running,attr"`
 	ReadOffset string `alloy:"read_offset,attr"`
 }