Merge branch 'master' into vmalert-proxy

Author: AndrewChubatiuk

app/vlinsert/elasticsearch/elasticsearch.go

@@ -187,7 +187,10 @@ func readBulkLine(lr *insertutil.LineReader, timeFields, msgFields []string, lmp
 		// Continue parsing next lines.
 		return true, nil
 	}
+
 	p := logstorage.GetJSONParser()
+	defer logstorage.PutJSONParser(p)
+
 	if err := p.ParseLogMessage(line); err != nil {
 		return false, fmt.Errorf("cannot parse json-encoded log entry: %w", err)
 	}
@@ -201,7 +204,6 @@ func readBulkLine(lr *insertutil.LineReader, timeFields, msgFields []string, lmp
 	}
 	logstorage.RenameField(p.Fields, msgFields, "_msg")
 	lmp.AddRow(ts, p.Fields, nil)
-	logstorage.PutJSONParser(p)
 
 	return true, nil
 }

app/vlinsert/syslog/syslog.go

@@ -100,6 +100,10 @@ func processQueryRequest(ctx context.Context, w http.ResponseWriter, r *http.Req
 
 		bb := bufs.Get(workerID)
 
+		// Write the marker of a regular data block.
+		bb.B = append(bb.B, 0)
+
+		// Marshal the data block.
 		bb.B = db.Marshal(bb.B)
 
 		if len(bb.B) < 1024*1024 {
@@ -117,7 +121,10 @@ func processQueryRequest(ctx context.Context, w http.ResponseWriter, r *http.Req
 		}
 	}
 
-	if err := vlstorage.RunQuery(ctx, cp.TenantIDs, cp.Query, writeBlock); err != nil {
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	if err := vlstorage.RunQuery(qctx, writeBlock); err != nil {
 		return err
 	}
 	if errGlobal != nil {
@@ -131,7 +138,13 @@ func processQueryRequest(ctx context.Context, w http.ResponseWriter, r *http.Req
 		}
 	}
 
-	return nil
+	// Send the query stats block.
+	bb := bufs.Get(0)
+	// Write the marker of query stats block.
+	bb.B = append(bb.B, 1)
+	// Marshal the block itself
+	bb.B = marshalQueryStatsBlock(bb.B, qctx)
+	return sendBuf(bb)
 }
 
 func processFieldNamesRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -140,12 +153,15 @@ func processFieldNamesRequest(ctx context.Context, w http.ResponseWriter, r *htt
 		return err
 	}
 
-	fieldNames, err := vlstorage.GetFieldNames(ctx, cp.TenantIDs, cp.Query)
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	fieldNames, err := vlstorage.GetFieldNames(qctx)
 	if err != nil {
 		return fmt.Errorf("cannot obtain field names: %w", err)
 	}
 
-	return writeValuesWithHits(w, fieldNames, cp.DisableCompression)
+	return writeValuesWithHits(w, qctx, fieldNames, cp.DisableCompression)
 }
 
 func processFieldValuesRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -161,12 +177,15 @@ func processFieldValuesRequest(ctx context.Context, w http.ResponseWriter, r *ht
 		return err
 	}
 
-	fieldValues, err := vlstorage.GetFieldValues(ctx, cp.TenantIDs, cp.Query, fieldName, uint64(limit))
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	fieldValues, err := vlstorage.GetFieldValues(qctx, fieldName, uint64(limit))
 	if err != nil {
 		return fmt.Errorf("cannot obtain field values: %w", err)
 	}
 
-	return writeValuesWithHits(w, fieldValues, cp.DisableCompression)
+	return writeValuesWithHits(w, qctx, fieldValues, cp.DisableCompression)
 }
 
 func processStreamFieldNamesRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -175,12 +194,15 @@ func processStreamFieldNamesRequest(ctx context.Context, w http.ResponseWriter,
 		return err
 	}
 
-	fieldNames, err := vlstorage.GetStreamFieldNames(ctx, cp.TenantIDs, cp.Query)
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	fieldNames, err := vlstorage.GetStreamFieldNames(qctx)
 	if err != nil {
 		return fmt.Errorf("cannot obtain stream field names: %w", err)
 	}
 
-	return writeValuesWithHits(w, fieldNames, cp.DisableCompression)
+	return writeValuesWithHits(w, qctx, fieldNames, cp.DisableCompression)
 }
 
 func processStreamFieldValuesRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -196,12 +218,15 @@ func processStreamFieldValuesRequest(ctx context.Context, w http.ResponseWriter,
 		return err
 	}
 
-	fieldValues, err := vlstorage.GetStreamFieldValues(ctx, cp.TenantIDs, cp.Query, fieldName, uint64(limit))
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	fieldValues, err := vlstorage.GetStreamFieldValues(qctx, fieldName, uint64(limit))
 	if err != nil {
 		return fmt.Errorf("cannot obtain stream field values: %w", err)
 	}
 
-	return writeValuesWithHits(w, fieldValues, cp.DisableCompression)
+	return writeValuesWithHits(w, qctx, fieldValues, cp.DisableCompression)
 }
 
 func processStreamsRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -215,12 +240,15 @@ func processStreamsRequest(ctx context.Context, w http.ResponseWriter, r *http.R
 		return err
 	}
 
-	streams, err := vlstorage.GetStreams(ctx, cp.TenantIDs, cp.Query, uint64(limit))
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	streams, err := vlstorage.GetStreams(qctx, uint64(limit))
 	if err != nil {
 		return fmt.Errorf("cannot obtain streams: %w", err)
 	}
 
-	return writeValuesWithHits(w, streams, cp.DisableCompression)
+	return writeValuesWithHits(w, qctx, streams, cp.DisableCompression)
 }
 
 func processStreamIDsRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -234,19 +262,33 @@ func processStreamIDsRequest(ctx context.Context, w http.ResponseWriter, r *http
 		return err
 	}
 
-	streamIDs, err := vlstorage.GetStreamIDs(ctx, cp.TenantIDs, cp.Query, uint64(limit))
+	qctx := cp.NewQueryContext(ctx)
+	defer cp.UpdatePerQueryStatsMetrics()
+
+	streamIDs, err := vlstorage.GetStreamIDs(qctx, uint64(limit))
 	if err != nil {
 		return fmt.Errorf("cannot obtain streams: %w", err)
 	}
 
-	return writeValuesWithHits(w, streamIDs, cp.DisableCompression)
+	return writeValuesWithHits(w, qctx, streamIDs, cp.DisableCompression)
 }
 
 type commonParams struct {
 	TenantIDs []logstorage.TenantID
 	Query     *logstorage.Query
 
 	DisableCompression bool
+
+	// qs contains execution statistics for the Query.
+	qs logstorage.QueryStats
+}
+
+func (cp *commonParams) NewQueryContext(ctx context.Context) *logstorage.QueryContext {
+	return logstorage.NewQueryContext(ctx, &cp.qs, cp.TenantIDs, cp.Query)
+}
+
+func (cp *commonParams) UpdatePerQueryStatsMetrics() {
+	vlstorage.UpdatePerQueryStatsMetrics(&cp.qs)
 }
 
 func getCommonParams(r *http.Request, expectedProtocolVersion string) (*commonParams, error) {
@@ -287,12 +329,18 @@ func getCommonParams(r *http.Request, expectedProtocolVersion string) (*commonPa
 	return cp, nil
 }
 
-func writeValuesWithHits(w http.ResponseWriter, vhs []logstorage.ValueWithHits, disableCompression bool) error {
+func writeValuesWithHits(w http.ResponseWriter, qctx *logstorage.QueryContext, vhs []logstorage.ValueWithHits, disableCompression bool) error {
 	var b []byte
+
+	// Marshal vhs at first
+	b = encoding.MarshalUint64(b, uint64(len(vhs)))
 	for i := range vhs {
 		b = vhs[i].Marshal(b)
 	}
 
+	// Marshal query stats block after that
+	b = marshalQueryStatsBlock(b, qctx)
+
 	if !disableCompression {
 		b = zstd.CompressLevel(nil, b, 1)
 	}
@@ -306,6 +354,13 @@ func writeValuesWithHits(w http.ResponseWriter, vhs []logstorage.ValueWithHits,
 	return nil
 }
 
+func marshalQueryStatsBlock(dst []byte, qctx *logstorage.QueryContext) []byte {
+	queryDurationNsecs := qctx.QueryDurationNsecs()
+	db := qctx.QueryStats.CreateDataBlock(queryDurationNsecs)
+	dst = db.Marshal(dst)
+	return dst
+}
+
 func getInt64FromRequest(r *http.Request, argName string) (int64, error) {
 	s := r.FormValue(argName)
 	n, err := strconv.ParseInt(s, 10, 64)