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,28 @@ 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 {
+	// If the logLabels map is empty, labels.FromMap will allocate some empty slices.
+	// Since this code is executed for every log line we receive, as an optimization
+	// to avoid those allocations we'll call labels.FromMap only if the map is not empty.
+	var lbls labels.Labels
+	if len(logLabels) > 0 {
+		lbls = labels.FromMap(logLabels)
+	}
+	return c.pushLogLine(line, timestamp, lbls, extraLabelList...)
 }
 
 func formatTS(ts time.Time) string {
@@ -101,21 +117,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"
 )
 

integration/loki_micro_services_test.go

@@ -396,13 +396,13 @@ func TestMicroServicesIngestQueryOverMultipleBucketSingleProvider(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

@@ -8,6 +8,8 @@ import (
 	"github.com/buger/jsonparser"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/modern-go/reflect2"
+
+	"github.com/grafana/loki/pkg/logproto"
 )
 
 func init() {
@@ -18,6 +20,21 @@ func init() {
 type Entry struct {
 	Timestamp time.Time
 	Line      string
+	Labels    LabelSet
+}
+
+func (e Entry) ToProto() logproto.Entry {
+	// If there are no labels, we return empty string instead of '{}'.
+	var labels string
+	if len(e.Labels) > 0 {
+		labels = e.Labels.String()
+	}
+
+	return logproto.Entry{
+		Timestamp: e.Timestamp,
+		Line:      e.Line,
+		Labels:    labels,
+	}
 }
 
 func (e *Entry) UnmarshalJSON(data []byte) error {
@@ -27,25 +44,45 @@ func (e *Entry) UnmarshalJSON(data []byte) error {
 	)
 	_, err := jsonparser.ArrayEach(data, func(value []byte, t jsonparser.ValueType, _ int, _ error) {
 		// assert that both items in array are of type string
-		if t != jsonparser.String {
-			parseError = jsonparser.MalformedStringError
-			return
-		}
 		switch i {
 		case 0: // timestamp
+			if t != jsonparser.String {
+				parseError = jsonparser.MalformedStringError
+				return
+			}
 			ts, err := jsonparser.ParseInt(value)
 			if err != nil {
 				parseError = err
 				return
 			}
 			e.Timestamp = time.Unix(0, ts)
 		case 1: // value
+			if t != jsonparser.String {
+				parseError = jsonparser.MalformedStringError
+				return
+			}
 			v, err := jsonparser.ParseString(value)
 			if err != nil {
 				parseError = err
 				return
 			}
 			e.Line = v
+		case 2: // labels
+			if t != jsonparser.Object {
+				parseError = jsonparser.MalformedObjectError
+				return
+			}
+			e.Labels = make(LabelSet)
+			if err := jsonparser.ObjectEach(value, func(key []byte, value []byte, dataType jsonparser.ValueType, _ int) error {
+				if dataType != jsonparser.String {
+					return jsonparser.MalformedStringError
+				}
+				e.Labels[yoloString(key)] = yoloString(value)
+				return nil
+			}); err != nil {
+				parseError = err
+				return
+			}
 		}
 		i++
 	})
@@ -67,6 +104,7 @@ func (sliceEntryDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
 		i := 0
 		var ts time.Time
 		var line string
+		var labels LabelSet
 		ok := iter.ReadArrayCB(func(iter *jsoniter.Iterator) bool {
 			var ok bool
 			switch i {
@@ -81,15 +119,28 @@ func (sliceEntryDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
 					return false
 				}
 				return true
+			case 2:
+				labels = make(LabelSet)
+				iter.ReadMapCB(func(iter *jsoniter.Iterator, labelName string) bool {
+					labelValue := iter.ReadString()
+					labels[labelName] = labelValue
+					return true
+				})
+				i++
+				if iter.Error != nil {
+					return false
+				}
+				return true
 			default:
-				iter.ReportError("error reading entry", "array must contains 2 values")
+				iter.ReportError("error reading entry", "array must have at least 2 and up to 3 values")
 				return false
 			}
 		})
 		if ok {
 			*((*[]Entry)(ptr)) = append(*((*[]Entry)(ptr)), Entry{
 				Timestamp: ts,
 				Line:      line,
+				Labels:    labels,
 			})
 			return true
 		}
@@ -126,6 +177,20 @@ func (EntryEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
 	stream.WriteRaw(`"`)
 	stream.WriteMore()
 	stream.WriteStringWithHTMLEscaped(e.Line)
+	if len(e.Labels) > 0 {
+		stream.WriteMore()
+		stream.WriteObjectStart()
+		var idx int
+		for lName, lValue := range e.Labels {
+			if idx > 0 {
+				stream.WriteMore()
+			}
+			stream.WriteObjectField(lName)
+			stream.WriteString(lValue)
+			idx++
+		}
+		stream.WriteObjectEnd()
+	}
 	stream.WriteArrayEnd()
 }
 

pkg/loghttp/query.go

@@ -56,9 +56,126 @@ 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 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, dataType jsonparser.ValueType, _ int) error {
+				if dataType != jsonparser.String {
+					return jsonparser.MalformedStringError
+				}
+				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
+	)
+	if _, 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)
+	}); err != nil {
+		parseError = err
+	}
+
+	if parseError != 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, dataType jsonparser.ValueType, _ int) error {
+				if dataType != jsonparser.String {
+					return jsonparser.MalformedStringError
+				}
+				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
@@ -121,7 +238,10 @@ func (s Streams) ToProto() []logproto.Stream {
 	}
 	result := make([]logproto.Stream, 0, len(s))
 	for _, s := range s {
-		entries := *(*[]logproto.Entry)(unsafe.Pointer(&s.Entries))
+		entries := make([]logproto.Entry, len(s.Entries), len(s.Entries))
+		for i, e := range s.Entries {
+			entries[i] = e.ToProto()
+		}
 		result = append(result, logproto.Stream{Labels: s.Labels.String(), Entries: entries})
 	}
 	return result
@@ -431,3 +551,7 @@ func labelVolumeLimit(r *http.Request) error {
 
 	return nil
 }
+
+func yoloString(buf []byte) string {
+	return *((*string)(unsafe.Pointer(&buf)))
+}