fix for grafana profiles plugin support

metrico · Aug 28, 2024 · c4d58bc · c4d58bc
1 parent 4a65580
commit c4d58bc
Show file tree

Hide file tree

Showing 2 changed files with 97 additions and 71 deletions.
diff --git a/pyroscope/render.js b/pyroscope/render.js
@@ -1,8 +1,7 @@
-const { parseTypeId } = require('./shared')
+const { parseQuery } = require('./shared')
 const { mergeStackTraces } = require('./merge_stack_traces')
 const querierMessages = require('./querier_pb')
 const { selectSeriesImpl } = require('./select_series')
-const types = require('./types/v1/types_pb')
 
 const render = async (req, res) => {
   const query = req.query.query
@@ -52,28 +51,50 @@ const render = async (req, res) => {
   const [bMergeStackTrace, selectSeries] =
     await Promise.all(promises)
   const mergeStackTrace = querierMessages.SelectMergeStacktracesResponse.deserializeBinary(bMergeStackTrace)
-  let series = new types.Series()
-  if (selectSeries.getSeriesList().length === 1) {
-    series = selectSeries.getSeriesList()[0]
+  let pTimeline = null
+  for (const series of selectSeries.getSeriesList()) {
+    if (!pTimeline) {
+      pTimeline = timeline(series,
+        fromTimeSec * 1000,
+        toTimeSec * 1000,
+        timelineStep)
+      continue
+    }
+    const _timeline = timeline(series,
+      fromTimeSec * 1000,
+      toTimeSec * 1000,
+      timelineStep)
+    pTimeline.samples = pTimeline.samples.map((v, i) => v + _timeline.samples[i])
   }
   const fb = toFlamebearer(mergeStackTrace.getFlamegraph(), parsedQuery.profileType)
-  fb.flamebearerProfileV1.timeline = timeline(series,
-    fromTimeSec * 1000,
-    toTimeSec * 1000,
-    timelineStep)
+  fb.flamebearerProfileV1.timeline = pTimeline
 
   if (groupBy.length > 0) {
+    const pGroupedTimelines = {}
     fb.flamebearerProfileV1.groups = {}
-    let key = '*'
-    series.getSeriesList().forEach((_series) => {
-      _series.getLabelsList().forEach((label) => {
-        key = label.getName() === groupBy[0] ? label.getValue() : key
-      })
-    })
-    fb.flamebearerProfileV1.groups[key] = timeline(series,
-      fromTimeSec * 1000,
-      toTimeSec * 1000,
-      timelineStep)
+    for (const series of selectSeries.getSeriesList()) {
+      const _key = {}
+      for (const label of series.getLabelsList()) {
+        if (groupBy.includes(label.getName())) {
+          _key[label.getName()] = label.getValue()
+        }
+      }
+      const key = '{' + Object.entries(_key).map(e => `${e[0]}=${JSON.stringify(e[1])}`)
+        .sort().join(', ') + '}'
+      if (!pGroupedTimelines[key]) {
+        pGroupedTimelines[key] = timeline(series,
+          fromTimeSec * 1000,
+          toTimeSec * 1000,
+          timelineStep)
+      } else {
+        const _timeline = timeline(series,
+          fromTimeSec * 1000,
+          toTimeSec * 1000,
+          timelineStep)
+        pGroupedTimelines[key].samples = pGroupedTimelines[key].samples.map((v, i) => v + _timeline.samples[i])
+      }
+    }
+    fb.flamebearerProfileV1.groups = pGroupedTimelines
   }
   res.code(200)
   res.headers({ 'Content-Type': 'application/json' })
@@ -208,43 +229,6 @@ function sizeToBackfill (startMs, endMs, stepSec) {
   return Math.floor((endMs - startMs) / (stepSec * 1000))
 }
 
-/**
- *
- * @param query {string}
- */
-const parseQuery = (query) => {
-  query = query.trim()
-  const match = query.match(/^([^{\s]+)\s*(\{(.*)})?$/)
-  if (!match) {
-    return null
-  }
-  const typeId = match[1]
-  const typeDesc = parseTypeId(typeId)
-  let strLabels = (match[3] || '').trim()
-  const labels = []
-  while (strLabels && strLabels !== '' && strLabels !== '}') {
-    const m = strLabels.match(/^(,)?\s*([A-Za-z0-9_]+)\s*(!=|!~|=~|=)\s*("([^"\\]|\\.)*")/)
-    if (!m) {
-      throw new Error('Invalid label selector')
-    }
-    labels.push([m[2], m[3], m[4]])
-    strLabels = strLabels.substring(m[0].length).trim()
-  }
-  const profileType = new types.ProfileType()
-  profileType.setId(typeId)
-  profileType.setName(typeDesc.type)
-  profileType.setSampleType(typeDesc.sampleType)
-  profileType.setSampleUnit(typeDesc.sampleUnit)
-  profileType.setPeriodType(typeDesc.periodType)
-  profileType.setPeriodUnit(typeDesc.periodUnit)
-  return {
-    typeId,
-    typeDesc,
-    labels,
-    labelSelector: strLabels,
-    profileType
-  }
-}
 
 const init = (fastify) => {
   fastify.get('/pyroscope/render', render)

diff --git a/pyroscope/shared.js b/pyroscope/shared.js
@@ -1,6 +1,6 @@
 const { QrynBadRequest } = require('../lib/handlers/errors')
 const Sql = require('@cloki/clickhouse-sql')
-const compiler = require('../parser/bnf')
+const types = require('./types/v1/types_pb')
 /**
  *
  * @param payload {ReadableStream}
@@ -77,14 +77,14 @@ const serviceNameSelectorQuery = (labelSelector) => {
   }
   const labelSelectorScript = parseLabelSelector(labelSelector)
   let conds = null
-  for (const rule of labelSelectorScript.Children('log_stream_selector_rule')) {
-    const label = rule.Child('label').value
+  for (const rule of labelSelectorScript) {
+    const label = rule[0]
     if (label !== 'service_name') {
       continue
     }
-    const val = JSON.parse(rule.Child('quoted_str').value)
+    const val = JSON.parse(rule[2])
     let valRul = null
-    switch (rule.Child('operator').value) {
+    switch (rule[1]) {
       case '=':
         valRul = Sql.Eq(new Sql.Raw('service_name'), Sql.val(val))
         break
@@ -102,12 +102,54 @@ const serviceNameSelectorQuery = (labelSelector) => {
   return conds || empty
 }
 
+/**
+ *
+ * @param query {string}
+ */
+const parseQuery = (query) => {
+  query = query.trim()
+  const match = query.match(/^([^{\s]+)\s*(\{(.*)})?$/)
+  if (!match) {
+    return null
+  }
+  const typeId = match[1]
+  const typeDesc = parseTypeId(typeId)
+  const strLabels = (match[3] || '').trim()
+  const labels = parseLabelSelector(strLabels)
+  const profileType = new types.ProfileType()
+  profileType.setId(typeId)
+  profileType.setName(typeDesc.type)
+  profileType.setSampleType(typeDesc.sampleType)
+  profileType.setSampleUnit(typeDesc.sampleUnit)
+  profileType.setPeriodType(typeDesc.periodType)
+  profileType.setPeriodUnit(typeDesc.periodUnit)
+  return {
+    typeId,
+    typeDesc,
+    labels,
+    labelSelector: strLabels,
+    profileType
+  }
+}
 
-const parseLabelSelector = (labelSelector) => {
-  if (labelSelector.endsWith(',}')) {
-    labelSelector = labelSelector.slice(0, -2) + '}'
+const parseLabelSelector = (strLabels) => {
+  strLabels = strLabels.trim()
+  if (strLabels.startsWith('{')) {
+    strLabels = strLabels.slice(1)
+  }
+  if (strLabels.endsWith('}')) {
+    strLabels = strLabels.slice(0, -1)
   }
-  return compiler.ParseScript(labelSelector).rootToken
+  const labels = []
+  while (strLabels && strLabels !== '' && strLabels !== '}' && strLabels !== ',') {
+    const m = strLabels.match(/^(,)?\s*([A-Za-z0-9_]+)\s*(!=|!~|=~|=)\s*("([^"\\]|\\.)*")/)
+    if (!m) {
+      throw new Error('Invalid label selector')
+    }
+    labels.push([m[2], m[3], m[4]])
+    strLabels = strLabels.substring(m[0].length).trim()
+  }
+  return labels
 }
 
 /**
@@ -128,7 +170,6 @@ const parseTypeId = (typeId) => {
   }
 }
 
-
 /**
  *
  * @param {Sql.Select} query
@@ -140,10 +181,10 @@ const labelSelectorQuery = (query, labelSelector) => {
   }
   const labelSelectorScript = parseLabelSelector(labelSelector)
   const labelsConds = []
-  for (const rule of labelSelectorScript.Children('log_stream_selector_rule')) {
-    const val = JSON.parse(rule.Child('quoted_str').value)
+  for (const rule of labelSelectorScript) {
+    const val = JSON.parse(rule[2])
     let valRul = null
-    switch (rule.Child('operator').value) {
+    switch (rule[1]) {
       case '=':
         valRul = Sql.Eq(new Sql.Raw('val'), Sql.val(val))
         break
@@ -157,7 +198,7 @@ const labelSelectorQuery = (query, labelSelector) => {
         valRul = Sql.Ne(new Sql.Raw(`match(val, ${Sql.quoteVal(val)})`), 1)
     }
     const labelSubCond = Sql.And(
-      Sql.Eq('key', Sql.val(rule.Child('label').value)),
+      Sql.Eq('key', Sql.val(rule[0])),
       valRul
     )
     labelsConds.push(labelSubCond)
@@ -183,5 +224,6 @@ module.exports = {
   serviceNameSelectorQuery,
   parseLabelSelector,
   labelSelectorQuery,
-  HISTORY_TIMESPAN
+  HISTORY_TIMESPAN,
+  parseQuery
 }