Add metadata to push payload

integration/client/client.go

@@ -13,6 +13,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/prometheus/prometheus/model/labels"
 	"github.com/weaveworks/common/user"
 )
 
@@ -86,13 +87,21 @@ func New(instanceID, token, baseURL string, opts ...Option) *Client {
 
 // PushLogLine creates a new logline with the current time as timestamp
 func (c *Client) PushLogLine(line string, extraLabels ...map[string]string) error {
-	return c.pushLogLine(line, c.Now, extraLabels...)
+	return c.pushLogLine(line, c.Now, nil, extraLabels...)
+}
+
+func (c *Client) PushLogLineWithMetadata(line string, logLabels map[string]string, extraLabels ...map[string]string) error {
+	return c.PushLogLineWithTimestampAndMetadata(line, c.Now, logLabels, extraLabels...)
 }
 
 // PushLogLineWithTimestamp creates a new logline at the given timestamp
 // The timestamp has to be a Unix timestamp (epoch seconds)
-func (c *Client) PushLogLineWithTimestamp(line string, timestamp time.Time, extraLabelList ...map[string]string) error {
-	return c.pushLogLine(line, timestamp, extraLabelList...)
+func (c *Client) PushLogLineWithTimestamp(line string, timestamp time.Time, extraLabels ...map[string]string) error {
+	return c.pushLogLine(line, timestamp, nil, extraLabels...)
+}
+
+func (c *Client) PushLogLineWithTimestampAndMetadata(line string, timestamp time.Time, logLabels map[string]string, extraLabelList ...map[string]string) error {
+	return c.pushLogLine(line, timestamp, labels.FromMap(logLabels), extraLabelList...)
 }
 
 func formatTS(ts time.Time) string {
@@ -101,21 +110,22 @@ func formatTS(ts time.Time) string {
 
 type stream struct {
 	Stream map[string]string `json:"stream"`
-	Values [][]string        `json:"values"`
+	Values [][]any           `json:"values"`
 }
 
 // pushLogLine creates a new logline
-func (c *Client) pushLogLine(line string, timestamp time.Time, extraLabelList ...map[string]string) error {
+func (c *Client) pushLogLine(line string, timestamp time.Time, logLabels labels.Labels, extraLabelList ...map[string]string) error {
 	apiEndpoint := fmt.Sprintf("%s/loki/api/v1/push", c.baseURL)
 
 	s := stream{
 		Stream: map[string]string{
 			"job": "varlog",
 		},
-		Values: [][]string{
+		Values: [][]any{
 			{
 				formatTS(timestamp),
 				line,
+				logLabels,
 			},
 		},
 	}

integration/loki_micro_services_delete_test.go

@@ -12,7 +12,6 @@ import (
 
 	"github.com/grafana/loki/integration/client"
 	"github.com/grafana/loki/integration/cluster"
-
 	"github.com/grafana/loki/pkg/storage"
 )
 
@@ -100,18 +99,22 @@ func TestMicroServicesDeleteRequest(t *testing.T) {
 				{
 					strconv.FormatInt(now.Add(-48*time.Hour).UnixNano(), 10),
 					"lineA",
+					"",
 				},
 				{
 					strconv.FormatInt(now.Add(-48*time.Hour).UnixNano(), 10),
 					"lineB",
+					"",
 				},
 				{
 					strconv.FormatInt(now.Add(-time.Minute).UnixNano(), 10),
 					"lineC",
+					"",
 				},
 				{
 					strconv.FormatInt(now.Add(-time.Minute).UnixNano(), 10),
 					"lineD",
+					"",
 				},
 			},
 		})

integration/loki_micro_services_test.go

@@ -207,13 +207,13 @@ func TestMicroServicesMultipleBucketSingleProvider(t *testing.T) {
 			cliQueryFrontend.Now = now
 
 			t.Run("ingest-logs", func(t *testing.T) {
-				// ingest logs to the previous period
-				require.NoError(t, cliDistributor.PushLogLineWithTimestamp("lineA", time.Now().Add(-48*time.Hour), map[string]string{"job": "fake"}))
-				require.NoError(t, cliDistributor.PushLogLineWithTimestamp("lineB", time.Now().Add(-36*time.Hour), map[string]string{"job": "fake"}))
+				require.NoError(t, cliDistributor.PushLogLineWithTimestampAndMetadata("lineA", time.Now().Add(-48*time.Hour), map[string]string{"traceID": "123"}, map[string]string{"job": "fake"}))
+				require.NoError(t, cliDistributor.PushLogLineWithTimestampAndMetadata("lineB", time.Now().Add(-36*time.Hour), map[string]string{"traceID": "456"}, map[string]string{"job": "fake"}))
 
 				// ingest logs to the current period
-				require.NoError(t, cliDistributor.PushLogLine("lineC", map[string]string{"job": "fake"}))
-				require.NoError(t, cliDistributor.PushLogLine("lineD", map[string]string{"job": "fake"}))
+				require.NoError(t, cliDistributor.PushLogLineWithMetadata("lineC", map[string]string{"traceID": "789"}, map[string]string{"job": "fake"}))
+				require.NoError(t, cliDistributor.PushLogLineWithMetadata("lineD", map[string]string{"traceID": "123"}, map[string]string{"job": "fake"}))
+
 			})
 
 			t.Run("query-lookback-default", func(t *testing.T) {

pkg/loghttp/entry.go

@@ -18,6 +18,7 @@ func init() {
 type Entry struct {
 	Timestamp time.Time
 	Line      string
+	Labels    string
 }
 
 func (e *Entry) UnmarshalJSON(data []byte) error {
@@ -46,6 +47,13 @@ func (e *Entry) UnmarshalJSON(data []byte) error {
 				return
 			}
 			e.Line = v
+		case 2: // labels
+			il, err := jsonparser.ParseString(value)
+			if err != nil {
+				parseError = err
+				return
+			}
+			e.Labels = il
 		}
 		i++
 	})
@@ -67,6 +75,7 @@ func (sliceEntryDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
 		i := 0
 		var ts time.Time
 		var line string
+		var labels string
 		ok := iter.ReadArrayCB(func(iter *jsoniter.Iterator) bool {
 			var ok bool
 			switch i {
@@ -81,6 +90,13 @@ func (sliceEntryDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
 					return false
 				}
 				return true
+			case 2:
+				labels = iter.ReadString()
+				i++
+				if iter.Error != nil {
+					return false
+				}
+				return true
 			default:
 				iter.ReportError("error reading entry", "array must contains 2 values")
 				return false
@@ -90,6 +106,7 @@ func (sliceEntryDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
 			*((*[]Entry)(ptr)) = append(*((*[]Entry)(ptr)), Entry{
 				Timestamp: ts,
 				Line:      line,
+				Labels:    labels,
 			})
 			return true
 		}
@@ -126,6 +143,8 @@ func (EntryEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
 	stream.WriteRaw(`"`)
 	stream.WriteMore()
 	stream.WriteStringWithHTMLEscaped(e.Line)
+	stream.WriteMore()
+	stream.WriteString(e.Labels)
 	stream.WriteArrayEnd()
 }
 

pkg/loghttp/query.go

@@ -55,9 +55,122 @@ func (q *QueryResponse) UnmarshalJSON(data []byte) error {
 	})
 }
 
-// PushRequest models a log stream push
+// PushRequest models a log stream push but is unmarshalled to proto push format.
 type PushRequest struct {
-	Streams []*Stream `json:"streams"`
+	Streams []LogProtoStream `json:"streams"`
+}
+
+// LogProtoStream helps with unmarshalling of each log stream for push request.
+// This might look un-necessary but without it the CPU usage in benchmarks was increasing by ~25% :shrug:
+type LogProtoStream struct {
+	logproto.Stream
+}
+
+func (s *LogProtoStream) UnmarshalJSON(data []byte) error {
+	err := jsonparser.ObjectEach(data, func(key, val []byte, ty jsonparser.ValueType, _ int) error {
+		switch string(key) {
+		case "stream":
+			labels := make(LabelSet)
+			err := jsonparser.ObjectEach(val, func(key, val []byte, _ jsonparser.ValueType, _ int) error {
+				labels[yoloString(key)] = yoloString(val)
+				return nil
+			})
+			if err != nil {
+				return err
+			}
+			s.Labels = labels.String()
+		case "values":
+			if ty == jsonparser.Null {
+				return nil
+			}
+			entries, err := unmarshalHTTPToLogProtoEntries(val)
+			if err != nil {
+				return err
+			}
+			s.Entries = entries
+		}
+		return nil
+	})
+	return err
+}
+
+func unmarshalHTTPToLogProtoEntries(data []byte) ([]logproto.Entry, error) {
+	var (
+		entries    []logproto.Entry
+		parseError error
+	)
+	_, err := jsonparser.ArrayEach(data, func(value []byte, ty jsonparser.ValueType, _ int, err error) {
+		if err != nil || parseError != nil {
+			return
+		}
+		if ty == jsonparser.Null {
+			return
+		}
+		e, err := unmarshalHTTPToLogProtoEntry(value)
+		if err != nil {
+			parseError = err
+			return
+		}
+		entries = append(entries, e)
+	})
+	if parseError != nil {
+		return nil, parseError
+	}
+	if err != nil {
+		return nil, parseError
+	}
+
+	return entries, nil
+}
+
+func unmarshalHTTPToLogProtoEntry(data []byte) (logproto.Entry, error) {
+	var (
+		i          int
+		parseError error
+		e          logproto.Entry
+	)
+	_, err := jsonparser.ArrayEach(data, func(value []byte, t jsonparser.ValueType, _ int, _ error) {
+		// assert that both items in array are of type string
+		if (i == 0 || i == 1) && t != jsonparser.String {
+			parseError = jsonparser.MalformedStringError
+			return
+		} else if i == 2 && t != jsonparser.Object {
+			parseError = jsonparser.MalformedObjectError
+			return
+		}
+		switch i {
+		case 0: // timestamp
+			ts, err := jsonparser.ParseInt(value)
+			if err != nil {
+				parseError = err
+				return
+			}
+			e.Timestamp = time.Unix(0, ts)
+		case 1: // value
+			v, err := jsonparser.ParseString(value)
+			if err != nil {
+				parseError = err
+				return
+			}
+			e.Line = v
+		case 2: // labels
+			labels := make(LabelSet)
+			err := jsonparser.ObjectEach(value, func(key, val []byte, _ jsonparser.ValueType, _ int) error {
+				labels[yoloString(key)] = yoloString(val)
+				return nil
+			})
+			if err != nil {
+				parseError = err
+				return
+			}
+			e.Labels = labels.String()
+		}
+		i++
+	})
+	if parseError != nil {
+		return e, parseError
+	}
+	return e, err
 }
 
 // ResultType holds the type of the result
@@ -377,3 +490,7 @@ func labelVolumeLimit(r *http.Request) error {
 
 	return nil
 }
+
+func yoloString(buf []byte) string {
+	return *((*string)(unsafe.Pointer(&buf)))
+}