Promtail WAL support: Implement writer side

clients/pkg/promtail/client/client_test.go

@@ -3,9 +3,7 @@ package client
 import (
 	"fmt"
 	"io"
-	"math"
 	"net/http"
-	"net/http/httptest"
 	"net/url"
 	"strings"
 	"testing"
@@ -24,7 +22,6 @@ import (
 	"github.com/grafana/loki/clients/pkg/promtail/api"
 
 	"github.com/grafana/loki/pkg/logproto"
-	"github.com/grafana/loki/pkg/util"
 	lokiflag "github.com/grafana/loki/pkg/util/flagext"
 )
 
@@ -38,11 +35,6 @@ var logEntries = []api.Entry{
 	{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(6, 0).UTC(), Line: "line0123456789"}},
 }
 
-type receivedReq struct {
-	tenantID string
-	pushReq  logproto.PushRequest
-}
-
 func TestClient_Handle(t *testing.T) {
 	tests := map[string]struct {
 		clientBatchSize           int
@@ -500,7 +492,7 @@ func TestClient_Handle(t *testing.T) {
 			receivedReqsChan := make(chan receivedReq, 10)
 
 			// Start a local HTTP server
-			server := httptest.NewServer(createServerHandler(receivedReqsChan, testData.serverResponseStatus))
+			server := newTestRemoteWriteServer(receivedReqsChan, testData.serverResponseStatus)
 			require.NotNil(t, server)
 			defer server.Close()
 
@@ -642,7 +634,7 @@ func TestClient_StopNow(t *testing.T) {
 			receivedReqsChan := make(chan receivedReq, 10)
 
 			// Start a local HTTP server
-			server := httptest.NewServer(createServerHandler(receivedReqsChan, c.serverResponseStatus))
+			server := newTestRemoteWriteServer(receivedReqsChan, c.serverResponseStatus)
 			require.NotNil(t, server)
 			defer server.Close()
 
@@ -710,24 +702,6 @@ func TestClient_StopNow(t *testing.T) {
 	}
 }
 
-func createServerHandler(receivedReqsChan chan receivedReq, status int) http.HandlerFunc {
-	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
-		// Parse the request
-		var pushReq logproto.PushRequest
-		if err := util.ParseProtoReader(req.Context(), req.Body, int(req.ContentLength), math.MaxInt32, &pushReq, util.RawSnappy); err != nil {
-			rw.WriteHeader(500)
-			return
-		}
-
-		receivedReqsChan <- receivedReq{
-			tenantID: req.Header.Get("X-Scope-OrgID"),
-			pushReq:  pushReq,
-		}
-
-		rw.WriteHeader(status)
-	})
-}
-
 type RoundTripperFunc func(*http.Request) (*http.Response, error)
 
 func (r RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {

clients/pkg/promtail/client/manager.go

@@ -0,0 +1,108 @@
+package client
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/go-kit/log"
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/grafana/loki/clients/pkg/promtail/api"
+	"github.com/grafana/loki/clients/pkg/promtail/wal"
+)
+
+// Manager manages remote write client instantiation, and connects the related components to orchestrate the flow of api.Entry
+// from the scrape targets, to the remote write clients themselves.
+//
+// Right now it just supports instantiating the WAL writer side of the future-to-be WAL enabled client. In follow-up
+// work, tracked in 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
+
+	entries chan api.Entry
+	once    sync.Once
+
+	wg sync.WaitGroup
+}
+
+// NewManager creates a new Manager
+func NewManager(metrics *Metrics, logger log.Logger, maxStreams, maxLineSize int, maxLineSizeTruncate bool, reg prometheus.Registerer, walCfg wal.Config, clientCfgs ...Config) (*Manager, error) {
+	// TODO: refactor this to instantiate all clients types
+	var fake struct{}
+
+	if len(clientCfgs) == 0 {
+		return nil, fmt.Errorf("at least one client config should be provided")
+	}
+	clientsCheck := make(map[string]struct{})
+	clients := make([]Client, 0, len(clientCfgs))
+	for _, cfg := range clientCfgs {
+		client, err := New(metrics, cfg, maxStreams, maxLineSize, maxLineSizeTruncate, logger)
+		if err != nil {
+			return nil, err
+		}
+
+		// Don't allow duplicate clients, we have client specific metrics that need at least one unique label value (name).
+		if _, ok := clientsCheck[client.Name()]; ok {
+			return nil, fmt.Errorf("duplicate client configs are not allowed, found duplicate for name: %s", cfg.Name)
+		}
+
+		clientsCheck[client.Name()] = fake
+		clients = append(clients, client)
+	}
+
+	manager := &Manager{
+		clients: clients,
+		entries: make(chan api.Entry),
+	}
+	manager.start()
+	return manager, nil
+}
+
+func (m *Manager) start() {
+	m.wg.Add(1)
+	go func() {
+		defer m.wg.Done()
+		// keep reading received entries
+		for e := range m.entries {
+			// then fanout to every remote write client
+			for _, c := range m.clients {
+				c.Chan() <- e
+			}
+		}
+	}()
+}
+
+func (m *Manager) StopNow() {
+	for _, c := range m.clients {
+		c.StopNow()
+	}
+}
+
+func (m *Manager) Name() string {
+	var sb strings.Builder
+	// name contains wal since manager is used as client only when WAL enabled for now
+	sb.WriteString("wal:")
+	for i, c := range m.clients {
+		sb.WriteString(c.Name())
+		if i != len(m.clients)-1 {
+			sb.WriteString(",")
+		}
+	}
+	return sb.String()
+}
+
+func (m *Manager) Chan() chan<- api.Entry {
+	return m.entries
+}
+
+func (m *Manager) Stop() {
+	// first stop the receiving channel
+	m.once.Do(func() { close(m.entries) })
+	m.wg.Wait()
+	// close clients
+	for _, c := range m.clients {
+		c.Stop()
+	}
+}

clients/pkg/promtail/client/manager_test.go

@@ -0,0 +1,108 @@
+package client
+
+import (
+	"net/http"
+	"net/url"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/grafana/dskit/backoff"
+	"github.com/grafana/dskit/flagext"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/clients/pkg/promtail/api"
+	"github.com/grafana/loki/clients/pkg/promtail/wal"
+	"github.com/grafana/loki/pkg/logproto"
+)
+
+func TestManager_ErrorCreatingWhenNoClientConfigsProvided(t *testing.T) {
+	walDir := t.TempDir()
+	_, err := NewManager(nil, log.NewLogfmtLogger(os.Stdout), 0, 0, false, prometheus.NewRegistry(), wal.Config{
+		Dir:     walDir,
+		Enabled: true,
+	})
+	require.Error(t, err)
+}
+
+type closer interface {
+	Close()
+}
+
+type closerFunc func()
+
+func (c closerFunc) Close() {
+	c()
+}
+
+func newServerAncClientConfig(t *testing.T) (Config, chan receivedReq, closer) {
+	receivedReqsChan := make(chan receivedReq, 10)
+
+	// Start a local HTTP server
+	server := newTestRemoteWriteServer(receivedReqsChan, http.StatusOK)
+	require.NotNil(t, server)
+
+	testClientURL, _ := url.Parse(server.URL)
+	testClientConfig := Config{
+		Name:      "test-client",
+		URL:       flagext.URLValue{URL: testClientURL},
+		Timeout:   time.Second * 2,
+		BatchSize: 1,
+		BackoffConfig: backoff.Config{
+			MaxRetries: 0,
+		},
+	}
+	return testClientConfig, receivedReqsChan, closerFunc(func() {
+		server.Close()
+		close(receivedReqsChan)
+	})
+}
+
+func TestManager_EntriesAreWrittenToClients(t *testing.T) {
+	walDir := t.TempDir()
+	reg := prometheus.NewRegistry()
+	testClientConfig, rwReceivedReqs, closeServer := newServerAncClientConfig(t)
+	clientMetrics := NewMetrics(reg)
+	manager, err := NewManager(clientMetrics, log.NewLogfmtLogger(os.Stdout), 0, 0, false, reg, wal.Config{
+		Dir:     walDir,
+		Enabled: true,
+	}, testClientConfig)
+	require.NoError(t, err)
+	require.Equal(t, "wal:test-client", manager.Name())
+
+	var testLabels = model.LabelSet{
+		"wal_enabled": "true",
+	}
+	var lines = []string{
+		"Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
+		"In eu nisl ac massa ultricies rutrum.",
+		"Sed eget felis at ipsum auctor congue.",
+	}
+	for _, line := range lines {
+		manager.Chan() <- api.Entry{
+			Labels: testLabels,
+			Entry: logproto.Entry{
+				Timestamp: time.Now(),
+				Line:      line,
+			},
+		}
+	}
+
+	// stop client to flush WAL, stop rw server
+	manager.Stop()
+	closeServer.Close()
+
+	// assert over rw client received entries
+	rwSeenEntriesCount := 0
+	for req := range rwReceivedReqs {
+		rwSeenEntriesCount++
+		require.Len(t, req.pushReq.Streams, 1, "expected 1 stream requests to be received")
+		require.Len(t, req.pushReq.Streams[0].Entries, 1, "expected 1 entry in the only stream received per request")
+		require.Equal(t, `{wal_enabled="true"}`, req.pushReq.Streams[0].Labels)
+		require.Contains(t, lines, req.pushReq.Streams[0].Entries[0].Line)
+	}
+	require.Equal(t, 3, rwSeenEntriesCount)
+}

clients/pkg/promtail/client/test_server.go

@@ -0,0 +1,40 @@
+package client
+
+import (
+	"math"
+	"net/http"
+	"net/http/httptest"
+
+	"github.com/grafana/loki/pkg/logproto"
+	"github.com/grafana/loki/pkg/util"
+)
+
+type receivedReq struct {
+	tenantID string
+	pushReq  logproto.PushRequest
+}
+
+// newTestMoreWriteServer creates a new httpserver.Server that can handle remote write request. When a request is handled,
+// the received entries are written to receivedChan, and status is responded.
+func newTestRemoteWriteServer(receivedChan chan receivedReq, status int) *httptest.Server {
+	server := httptest.NewServer(createServerHandler(receivedChan, status))
+	return server
+}
+
+func createServerHandler(receivedReqsChan chan receivedReq, receivedOKStatus int) http.HandlerFunc {
+	return func(rw http.ResponseWriter, req *http.Request) {
+		// Parse the request
+		var pushReq logproto.PushRequest
+		if err := util.ParseProtoReader(req.Context(), req.Body, int(req.ContentLength), math.MaxInt32, &pushReq, util.RawSnappy); err != nil {
+			rw.WriteHeader(500)
+			return
+		}
+
+		receivedReqsChan <- receivedReq{
+			tenantID: req.Header.Get("X-Scope-OrgID"),
+			pushReq:  pushReq,
+		}
+
+		rw.WriteHeader(receivedOKStatus)
+	}
+}

clients/pkg/promtail/config/config.go

@@ -7,9 +7,6 @@ import (
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	dskit_flagext "github.com/grafana/dskit/flagext"
-
-	"github.com/grafana/loki/pkg/tracing"
-
 	"gopkg.in/yaml.v2"
 
 	"github.com/grafana/loki/clients/pkg/promtail/client"
@@ -18,7 +15,8 @@ import (
 	"github.com/grafana/loki/clients/pkg/promtail/scrapeconfig"
 	"github.com/grafana/loki/clients/pkg/promtail/server"
 	"github.com/grafana/loki/clients/pkg/promtail/targets/file"
-
+	"github.com/grafana/loki/clients/pkg/promtail/wal"
+	"github.com/grafana/loki/pkg/tracing"
 	"github.com/grafana/loki/pkg/util/flagext"
 )
 
@@ -39,6 +37,7 @@ type Config struct {
 	LimitsConfig    limit.Config          `yaml:"limits_config,omitempty"`
 	Options         Options               `yaml:"options,omitempty"`
 	Tracing         tracing.Config        `yaml:"tracing"`
+	WAL             wal.Config            `yaml:"wal"`
 }
 
 // RegisterFlags with prefix registers flags where every name is prefixed by