Move walWriter into manager and contain it there

Author: kalleep

internal/component/common/loki/client/manager.go

@@ -72,8 +72,9 @@ func (p watcherClientPair) Stop(drain bool) {
 // https://github.com/grafana/loki/issues/8197, this Manager will be
 // responsible for instantiating all client types: Logger, Multi and WAL.
 type Manager struct {
-	clients []Client
-	pairs   []watcherClientPair
+	clients   []Client
+	pairs     []watcherClientPair
+	walWriter *wal.Writer
 
 	entries chan loki.Entry
 	once    sync.Once
@@ -82,20 +83,34 @@ type Manager struct {
 }
 
 // NewManager creates a new Manager
-func NewManager(metrics *Metrics, logger log.Logger, reg prometheus.Registerer, walCfg wal.Config, notifier WriterEventsNotifier, clientCfgs ...Config) (*Manager, error) {
+func NewManager(metrics *Metrics, logger log.Logger, reg prometheus.Registerer, walCfg wal.Config, clientCfgs ...Config) (*Manager, error) {
 	var fake struct{}
 
-	walWatcherMetrics := wal.NewWatcherMetrics(reg)
-	walMarkerMetrics := internal.NewMarkerMetrics(reg)
-	queueClientMetrics := NewQueueClientMetrics(reg)
-
 	if len(clientCfgs) == 0 {
 		return nil, fmt.Errorf("at least one client config must be provided")
 	}
 
-	clientsCheck := make(map[string]struct{})
-	clients := make([]Client, 0, len(clientCfgs))
-	pairs := make([]watcherClientPair, 0, len(clientCfgs))
+	var walWriter *wal.Writer
+	if walCfg.Enabled {
+		var err error
+		walWriter, err = wal.NewWriter(walCfg, logger, reg)
+		if err != nil {
+			return nil, fmt.Errorf("error creating wal writer: %w", err)
+		}
+	}
+
+	var (
+		walWatcherMetrics  = wal.NewWatcherMetrics(reg)
+		walMarkerMetrics   = internal.NewMarkerMetrics(reg)
+		queueClientMetrics = NewQueueClientMetrics(reg)
+	)
+
+	var (
+		clientsCheck = make(map[string]struct{})
+		clients      = make([]Client, 0, len(clientCfgs))
+		pairs        = make([]watcherClientPair, 0, len(clientCfgs))
+	)
+
 	for _, cfg := range clientCfgs {
 		// Don't allow duplicate clients, we have client specific metrics that need at least one unique label value (name).
 		clientName := getClientName(cfg)
@@ -120,13 +135,13 @@ func NewManager(metrics *Metrics, logger log.Logger, reg prometheus.Registerer,
 				return nil, fmt.Errorf("error starting queue client: %w", err)
 			}
 
+			watcher := wal.NewWatcher(walCfg.Dir, clientName, walWatcherMetrics, queue, wlog, walCfg.WatchConfig, markerHandler)
+
 			// subscribe watcher's wal.WriteTo to writer events. This will make the writer trigger the cleanup of the wal.WriteTo
 			// series cache whenever a segment is deleted.
-			notifier.SubscribeCleanup(queue)
-
-			watcher := wal.NewWatcher(walCfg.Dir, clientName, walWatcherMetrics, queue, wlog, walCfg.WatchConfig, markerHandler)
+			walWriter.SubscribeCleanup(queue)
 			// subscribe watcher to wal write events
-			notifier.SubscribeWrite(watcher)
+			walWriter.SubscribeWrite(watcher)
 
 			level.Debug(logger).Log("msg", "starting WAL watcher for client", "client", clientName)
 			watcher.Start()
@@ -149,9 +164,10 @@ func NewManager(metrics *Metrics, logger log.Logger, reg prometheus.Registerer,
 		}
 	}
 	manager := &Manager{
-		clients: clients,
-		pairs:   pairs,
-		entries: make(chan loki.Entry),
+		clients:   clients,
+		pairs:     pairs,
+		walWriter: walWriter,
+		entries:   make(chan loki.Entry),
 	}
 
 	if walCfg.Enabled {
@@ -187,6 +203,9 @@ func (m *Manager) startWithForward() {
 }
 
 func (m *Manager) Chan() chan<- loki.Entry {
+	if m.walWriter != nil {
+		return m.walWriter.Chan()
+	}
 	return m.entries
 }
 
@@ -195,11 +214,15 @@ func (m *Manager) Stop() {
 	m.StopWithDrain(false)
 }
 
-// StopWithDrain will stop the manager, its Write-Ahead Log watchers, and clients accordingly. If drain is enabled,
+// StopWithDrain will stop the manager, its WalWriter, Write-Ahead Log watchers, and clients accordingly. If drain is enabled,
 // the Watchers will attempt to drain the WAL completely.
-// The shutdown procedure first stops the Watchers, allowing them to flush as much data into the clients as possible. Then
-// the clients are shut down accordingly.
+// The shutdown procedure first stops the WalWriter, Then Watchers, allowing them to flush as much data into the clients as possible.
+// Lastly the clients are shut down accordingly.
 func (m *Manager) StopWithDrain(drain bool) {
+	if m.walWriter != nil {
+		m.walWriter.Stop()
+	}
+
 	// first stop the receiving channel
 	m.once.Do(func() { close(m.entries) })
 	m.wg.Wait()

internal/component/common/loki/client/manager_test.go

@@ -37,7 +37,7 @@ func TestManager_NoDuplicateMetricsPanic(t *testing.T) {
 		for range 2 {
 			_, err := NewManager(metrics, log.NewNopLogger(), reg, wal.Config{
 				WatchConfig: wal.DefaultWatchConfig,
-			}, NilNotifier, Config{
+			}, Config{
 				URL: flagext.URLValue{URL: host},
 			})
 			require.NoError(t, err)
@@ -53,7 +53,7 @@ func TestManager_ErrorCreatingWhenNoClientConfigsProvided(t *testing.T) {
 				Dir:         walDir,
 				Enabled:     walEnabled,
 				WatchConfig: wal.DefaultWatchConfig,
-			}, NilNotifier)
+			})
 			require.Error(t, err)
 		})
 	}
@@ -74,7 +74,7 @@ func TestManager_ErrorCreatingWhenRepeatedConfigs(t *testing.T) {
 				Dir:         walDir,
 				Enabled:     walEnabled,
 				WatchConfig: wal.DefaultWatchConfig,
-			}, NilNotifier, config1, config1Copy)
+			}, config1, config1Copy)
 			require.Error(t, err)
 		})
 	}
@@ -131,10 +131,7 @@ func TestManager_WALEnabled(t *testing.T) {
 	testClientConfig, rwReceivedReqs, closeServer := newServerAndClientConfig(t)
 	clientMetrics := NewMetrics(reg)
 
-	// start writer and manager
-	writer, err := wal.NewWriter(walConfig, logger, reg)
-	require.NoError(t, err)
-	manager, err := NewManager(clientMetrics, logger, prometheus.NewRegistry(), walConfig, writer, testClientConfig)
+	manager, err := NewManager(clientMetrics, logger, prometheus.NewRegistry(), walConfig, testClientConfig)
 	require.NoError(t, err)
 
 	receivedRequests := utils.NewSyncSlice[utils.RemoteWriteRequest]()
@@ -145,7 +142,6 @@ func TestManager_WALEnabled(t *testing.T) {
 	}()
 
 	defer func() {
-		writer.Stop()
 		manager.Stop()
 		closeServer.Close()
 	}()
@@ -155,7 +151,7 @@ func TestManager_WALEnabled(t *testing.T) {
 	}
 	var totalLines = 100
 	for i := range totalLines {
-		writer.Chan() <- loki.Entry{
+		manager.Chan() <- loki.Entry{
 			Labels: testLabels,
 			Entry: push.Entry{
 				Timestamp: time.Now(),
@@ -189,7 +185,7 @@ func TestManager_WALDisabled(t *testing.T) {
 	clientMetrics := NewMetrics(reg)
 
 	// start writer and manager
-	manager, err := NewManager(clientMetrics, logger, prometheus.NewRegistry(), walConfig, NilNotifier, testClientConfig)
+	manager, err := NewManager(clientMetrics, logger, prometheus.NewRegistry(), walConfig, testClientConfig)
 	require.NoError(t, err)
 
 	receivedRequests := utils.NewSyncSlice[utils.RemoteWriteRequest]()
@@ -247,7 +243,7 @@ func TestManager_WALDisabled_MultipleConfigs(t *testing.T) {
 	clientMetrics := NewMetrics(reg)
 
 	// start writer and manager
-	manager, err := NewManager(clientMetrics, logger, prometheus.NewRegistry(), walConfig, NilNotifier, testClientConfig, testClientConfig2)
+	manager, err := NewManager(clientMetrics, logger, prometheus.NewRegistry(), walConfig, testClientConfig, testClientConfig2)
 	require.NoError(t, err)
 
 	receivedRequests := utils.NewSyncSlice[utils.RemoteWriteRequest]()

internal/component/common/loki/client/shards.go

@@ -233,6 +233,8 @@ type shards struct {
 // Each shard gets its own queue and a dedicated worker that processes batches
 // from that queue. The number of shards determines the parallelism level.
 func (s *shards) start(n int) {
+	n = max(n, 1)
+
 	s.mut.Lock()
 	defer s.mut.Unlock()
 

internal/component/loki/write/write.go

@@ -118,17 +118,11 @@ func New(o component.Options, args Arguments) (*Component, error) {
 // Run implements component.Component.
 func (c *Component) Run(ctx context.Context) error {
 	defer func() {
-		// First we need to stop the sink, this is either wrapped clientManger or walWriter.
-		// Stopping the sink will not stop the inner handler
+		// First we need to stop the sink. Stopping the sink will not stop the wrapped handler.
 		if c.sink != nil {
 			c.sink.Stop()
 		}
 
-		// when exiting Run, proceed to shut down first the writer component, and then
-		// the client manager, with the WAL and remote-write client inside
-		if c.walWriter != nil {
-			c.walWriter.Stop()
-		}
 		if c.clientManger != nil {
 			// drain, since the component is shutting down. That means Alloy is shutting down as well
 			c.clientManger.StopWithDrain(true)
@@ -164,10 +158,6 @@ func (c *Component) Update(args component.Arguments) error {
 		c.sink.Stop()
 	}
 
-	if c.walWriter != nil {
-		c.walWriter.Stop()
-	}
-
 	if c.clientManger != nil {
 		// only drain on component shutdown
 		c.clientManger.Stop()
@@ -199,19 +189,7 @@ func (c *Component) Update(args component.Arguments) error {
 	walCfg.Dir = filepath.Join(c.opts.DataPath, "wal")
 
 	var err error
-	var notifier client.WriterEventsNotifier = client.NilNotifier
-	// nil-out wal writer in case WAL was disabled
-	c.walWriter = nil
-	// only configure WAL Writer if enabled
-	if walCfg.Enabled {
-		c.walWriter, err = wal.NewWriter(walCfg, c.opts.Logger, c.opts.Registerer)
-		if err != nil {
-			return fmt.Errorf("error creating wal writer: %w", err)
-		}
-		notifier = c.walWriter
-	}
-
-	c.clientManger, err = client.NewManager(c.metrics, c.opts.Logger, c.opts.Registerer, walCfg, notifier, cfgs...)
+	c.clientManger, err = client.NewManager(c.metrics, c.opts.Logger, c.opts.Registerer, walCfg, cfgs...)
 	if err != nil {
 		return fmt.Errorf("failed to create client manager: %w", err)
 	}