绘制 "火焰图" 总结

type Span = {
  span_id: number
  parent_id: number

  start_time: number
  duration: number

  event: string // 调用的方法名
}

{
  "trace_id": 5796316316865205225,
  "span_sets": [
    {
      "node_type": "TiKV",
      "spans": [
        {
          "span_id": 393276,
          "parent_id": 393275,
          "begin_unix_time_ns": 1607658272409814199,
          "duration_ns": 302332,
          "event": "Endpoint::parse_and_handle_unary_request"
        },
        {
          "span_id": 917515,
          "parent_id": 393278,
          "begin_unix_time_ns": 1607658272410116531,
          "duration_ns": 134483,
          "event": "RaftKv::async_snapshot"
        },
        {
          "span_id": 917516,
          "parent_id": 917515,
          "begin_unix_time_ns": 1607658272410116531,
          "duration_ns": 134483,
          "event": "LocalReader::propose_raft_command"
        },
        ...
}

// 这个类型声明是由代码生成器 (openapi-generator) 根据后端的 model 自动生成的
// 默认所有属性都是可选的 (虽然实际情况它们都是有值的)
interface TraceSpan {
  span_id?: number
  parent_id?: number

  begin_unix_time_ns?: number
  duration_ns?: number

  event?: string
}

// 这是我们自己扩展的类型声明
interface IFullSpan extends TraceSpan {
  node_type: string // 区分 span 对应的方法是在 tidb 还是 tikv 中执行的

  children: IFullSpan[]
  parent?: IFullSpan

  relative_begin_unix_time_ns: number
  relative_end_unix_time_ns: number
  max_relative_end_time_ns: number // include children span

  depth: number // which layer it should be drawed in, rootSpan is 0
  max_child_depth: number
}

const allSpans: IFullSpan[] = []
source.span_sets?.forEach((spanSet) => {
  spanSet.spans?.forEach((span) => {
    allSpans.push({
      ...span,

      node_type: spanSet.node_type!,
      children: [],

      relative_begin_unix_time_ns: 0,
      relative_end_unix_time_ns: 0,
      max_relative_end_time_ns: 0,
      depth: 0,
      max_child_depth: 0,
    })
  })
}

const rootSpan = allSpans.find((span) => span.parent_id === 0)!
const startTime = rootSpan.begin_unix_time_ns!
allSpans.forEach((span) => {
  span.relative_begin_unix_time_ns = span.begin_unix_time_ns! - startTime
  span.relative_end_unix_time_ns =
    span.relative_begin_unix_time_ns + span.duration_ns!
  span.max_relative_end_time_ns = span.relative_end_unix_time_ns
})

function findChildren(parentSpan: IFullSpan, allSpans: IFullSpan[]) {
  parentSpan.children = allSpans.filter(
    (span) => span.parent_id === parentSpan.span_id
  )
  parentSpan.children.forEach((child) => {
    child.parent = parentSpan
    findChildren(child, allSpans)
  })
}

findChildren(rootSpan, allSpans)

export type FullSpanMap = Record<string, IFullSpan>

function buildTree(allSpans: IFullSpan[]): FullSpanMap {
  // convert arr to map
  let spansObj = allSpans.reduce((accu, cur) => {
    accu[cur.span_id!] = cur
    return accu
  }, {} as FullSpanMap)

  // set children and parent
  Object.values(spansObj).forEach((span) => {
    const parent = spansObj[span.parent_id!]
    span.parent = parent
    // the root span has no parent
    if (parent) {
      parent.children.push(span)
    }
  })

  // sort children
  Object.values(spansObj).forEach((span) => {
    span.children.sort((a, b) => {
      let delta = a.relative_begin_unix_time_ns - b.relative_begin_unix_time_ns
      if (delta === 0) {
        // make the span with longer duration in the front when they have the same begin time
        // so we can draw the span with shorter duration first
        // to make them closer to the parent span
        delta = b.duration_ns! - a.duration_ns!
      }
      return delta
    })
  })
  return spansObj
}

function calcMaxEndTime(span: IFullSpan) {
  // return condition
  if (span.children.length === 0) {
    span.max_end_time_ns = span.end_unix_time_ns
    return span.end_unix_time_ns
  }

  const childrenTime = span.children
    .map((childSpan) => calcMaxEndTime(childSpan))
    .concat(span.end_unix_time_ns)
  const maxTime = Math.max(...childrenTime)
  span.max_end_time_ns = maxTime
  return maxTime
}

calcMaxEndTime(rootSpan)

function calcMaxEndTime(spansObj: FullSpanMap) {
  Object.values(spansObj)
    .filter((span) => span.children.length === 0) // find leaf spans
    .forEach(calcParentMaxEndTime)
}

// from bottom to top
function calcParentMaxEndTime(span: IFullSpan) {
  const parent = span.parent
  if (parent === undefined) return

  if (span.max_relative_end_time_ns > parent.max_relative_end_time_ns) {
    parent.max_relative_end_time_ns = span.max_relative_end_time_ns
  }
  calcParentMaxEndTime(parent)
}

calcMaxEndTime(spansObj)

// from top to bottom
function calcDepth(parentSpan: IFullSpan) {
  const childrenMaxIdx = parentSpan.children.length - 1
  // keep the same logic as datadog
  // compare the spans from right to left
  for (let i = childrenMaxIdx; i >= 0; i--) {
    const curSpan = parentSpan.children[i]

    if (i === childrenMaxIdx) {
      curSpan.depth = parentSpan.depth + 1
    } else {
      const lastSpan = parentSpan.children[i + 1]
      if (
        curSpan.max_relative_end_time_ns >
          lastSpan.relative_begin_unix_time_ns ||
        curSpan.relative_begin_unix_time_ns ===
          lastSpan.relative_begin_unix_time_ns
      ) {
        if (lastSpan.max_child_depth === lastSpan.depth) {
          // lastSpan has no children
          curSpan.depth = lastSpan.max_child_depth + 1
        } else {
          // keep the same logic as datadog
          // add a more empty layer
          curSpan.depth = lastSpan.max_child_depth + 2
        }
      } else {
        curSpan.depth = parentSpan.depth + 1
        // equal
        // curSpan.depth = lastSpan.depth
      }
    }
    curSpan.max_child_depth = curSpan.depth
    updateParentChildDepth(curSpan)
    calcDepth(curSpan)
  }
}

// from bottom to top
function updateParentChildDepth(span: IFullSpan) {
  const parent = span.parent
  if (parent === undefined) return

  if (span.max_child_depth > parent.max_child_depth) {
    parent.max_child_depth = span.max_child_depth
    updateParentChildDepth(parent)
  }
}

calcDepth(rootSpan)

setTimeLenScale() {
  this.timeLenScale = scaleLinear()
    .domain([0, this.timeDuration])
    .range([0, this.width])
}

windowToTimeRange(window: Window): TimeRange {
  return {
    start: this.timeLenScale.invert(window.left),
    end: this.timeLenScale.invert(window.right),
  }
}

fixPixelRatio() {
  // https://developer.mozilla.org/zh-CN/docs/Web/API/Window/devicePixelRatio
  const dpr = window.devicePixelRatio || 1

  this.context.canvas.style.width = this.width + 'px'
  this.context.canvas.style.height = this.height + 'px'
  this.context.canvas.width = this.width * dpr
  this.context.canvas.height = this.height * dpr
  this.context.scale(dpr, dpr)
}

import { getValueFormat } from '@baurine/grafana-value-formats'

getValueFormat('ns')(13387520, 2) // 得到 "13.39 ms"

windowToCanvasLoc(windowX: number, windowY: number) {
  const canvasBox = this.context.canvas.getBoundingClientRect()
  return {
    x: windowX - canvasBox.left,
    y: windowY - canvasBox.top,
  }
}

draw() {
  // 重绘之前先清空画布
  this.context.clearRect(0, 0, this.width, this.height)

  this.drawTimePointsAndVerticalLines()
  this.drawFlameGraph()
  this.drawWindow()
  this.drawSelectedWindow()
  this.drawMoveVerticalLine()
}

// setup
constructor(container: HTMLDivElement, flameGraph: IFlameGraph) {
  // ...

  this.drawOffscreenCanvas()
  this.draw()
  // ...
}

//////////////
// offscreen canvas
drawOffscreenCanvas() {
  this.offscreenContext.save()
  this.drawSpan(this.flameGraph.rootSpan, this.offscreenContext)
  this.offscreenContext.restore()
}

drawSpan(span: IFullSpan, ctx: CanvasRenderingContext2D) {
  if (span.node_type === 'TiDB') {
    ctx.fillStyle = '#aab254'
  } else {
    ctx.fillStyle = '#507359'
  }

  // 绘制矩形
  const x = this.timeLenScale(span.relative_begin_unix_time_ns)
  const y = span.depth * TimelineOverviewChart.OFFSCREEN_CANVAS_LAYER_HEIGHT
  let width = Math.max(
    this.timeLenScale(span.duration_ns!),
    TimelineOverviewChart.OFFSCREEN_CANVAS_SPAN_WIDTH
  )
  const height = TimelineOverviewChart.OFFSCREEN_CANVAS_LAYER_HEIGHT - 1
  ctx.fillRect(x, y, width, height)

  // 绘制竖线
  const deltaDepth = span.depth - (span.parent?.depth || 0)
  if (deltaDepth > 1) {
    ctx.strokeStyle = ctx.fillStyle
    ctx.lineWidth = 0.5
    ctx.beginPath()
    ctx.moveTo(x, y)
    ctx.lineTo(
      x,
      y -
        (deltaDepth - 1) * TimelineOverviewChart.OFFSCREEN_CANVAS_LAYER_HEIGHT
    )
    ctx.stroke()
  }

  // 继续绘制子 span
  span.children.forEach((s) => this.drawSpan(s, ctx))
}

// 重绘时从离屏 canvas 整体拷贝
drawFlameGraph() {
  this.context.save()
  this.context.drawImage(
    this.offscreenContext.canvas,
    0,
    0,
    this.width,
    this.offscreenCanvasHeight,
    0,
    16,
    this.width,
    this.height - 16
  )
  this.context.restore()
}

drawWindow() {
  const { left, right } = this.curWindow

  this.context.save()

  // draw unselected window area
  this.context.globalAlpha = TimelineOverviewChart.UNSELECTED_WINDOW_ALPHA
  this.context.fillStyle = TimelineOverviewChart.UNSELECTED_WINDOW_FILL_STYLE
  this.context.fillRect(0, 0, left, this.height)
  this.context.fillRect(right, 0, this.width, this.height)o

  // draw window left and right borders
  // ...

  // draw resize area
  // ...
}

drawSelectedWindow() {
  if (this.mouseDownPos === null || this.action !== Action.SelectWindow) {
    return
  }

  this.context.save()
  this.context.globalAlpha = TimelineOverviewChart.SELECTED_WINDOW_ALPHA
  this.context.fillStyle = TimelineOverviewChart.SELECTED_WINDOW_FILL_STYLE
  if (this.curMousePos.x > this.mouseDownPos.x) {
    this.context.fillRect(
      this.mouseDownPos.x,
      0,
      this.curMousePos.x - this.mouseDownPos.x,
      this.height
    )
  } else {
    this.context.fillRect(
      this.curMousePos.x,
      0,
      this.mouseDownPos.x - this.curMousePos.x,
      this.height
    )
  }
  this.context.restore()
}

drawMoveVerticalLine() {
  // not draw it when mouse move outside the canvas
  // to keep same as the chrome dev tool
  if (
    this.action !== Action.SelectWindow ||
    this.mouseOutsideCanvas(this.curMousePos)
  ) {
    return
  }

  this.context.save()
  this.context.strokeStyle =
    TimelineOverviewChart.MOVED_VERTICAL_LINE_STROKE_STYLE
  this.context.lineWidth = TimelineOverviewChart.MOVED_VERTICAL_LINE_WIDTH
  this.context.beginPath()
  this.context.moveTo(this.curMousePos.x, 0)
  this.context.lineTo(this.curMousePos.x, this.height)
  this.context.stroke()
  this.context.restore()
}

/////////////////////////////////////
// setup
constructor(container: HTMLDivElement, flameGraph: IFlameGraph) {
  // ...

  this.registerHanlers()
}

/////////////////////////////////////
// event handlers: mousedown, mousemove, mouseup, mousewheel, resize
registerHanlers() {
  window.addEventListener('resize', this.onResize)
  // https://developer.mozilla.org/en-US/docs/Web/API/Element/wheel_event
  this.context.canvas.addEventListener('wheel', this.onMouseWheel)
  this.context.canvas.addEventListener('mousedown', this.onMouseDown)
  this.context.canvas.addEventListener('mousemove', this.onCanvasMouseMove)
  this.context.canvas.addEventListener('mouseout', this.onCanvasMouseOut)
  window.addEventListener('mousemove', this.onWindowMouseMove)
  window.addEventListener('mouseup', this.onMouseUp)
}

updateAction(loc: Pos) {
  // only change it when mouse isn't down
  if (this.mouseDownPos) return

  const { left, right } = this.curWindow
  if (this.mouseOutsideCanvas(loc)) {
    this.action = Action.None
  } else if (loc.y > this.dragAreaHeight) {
    this.action = Action.SelectWindow
  } else if (
    loc.x > left - TimelineOverviewChart.WINDOW_RESIZE_LINE_WIDTH_HALF &&
    loc.x < left + TimelineOverviewChart.WINDOW_RESIZE_LINE_WIDTH_HALF
  ) {
    this.action = Action.MoveWindowLeft
  } else if (
    loc.x > right - TimelineOverviewChart.WINDOW_RESIZE_LINE_WIDTH_HALF &&
    loc.x < right + TimelineOverviewChart.WINDOW_RESIZE_LINE_WIDTH_HALF
  ) {
    this.action = Action.MoveWindowRight
  } else {
    this.action = Action.MoveWindow
  }
  this.updateCursor()
}

getSpanInPos(span: IFullSpan, pos: Pos): IFullSpan | null {
  const { x, y } = pos
  const x1 = this.timeLenScale(span.relative_begin_unix_time_ns)
  let x2 = this.timeLenScale(span.relative_end_unix_time_ns)
  if (x2 === x1) {
    x2 = x1 + TimelineDetailChart.MIN_SPAN_WIDTH
  }

  const y1 = span.depth * TimelineDetailChart.LAYER_HEIGHT
  const y2 = y1 + TimelineDetailChart.LAYER_HEIGHT - 1
  if (x <= x2 && x >= x1 && y <= y2 && y >= y1) {
    return span
  }
  if (span.children.length === 0) {
    return null
  }

  // traverse children
  for (let i = 0; i < span.children.length; i++) {
    const targetSpan = this.getSpanInPos(span.children[i], pos)
    if (targetSpan) {
      return targetSpan
    }
  }
  return null
}

onCanvasMouseMove = (event) => {
  //...

  const loc = this.windowToCanvasLoc(event.clientX, event.clientY)
  this.hoverSpan = this.getSpanInPos(this.flameGraph.rootSpan, loc)
  this.showTooltip({ x: event.clientX, y: event.clientY })
  this.draw()
}

onMouseUp = (event) => {
  //...
  const loc = this.windowToCanvasLoc(event.clientX, event.clientY)

  // handle click
  if (loc.x === this.mouseDownPos?.x && loc.y === this.mouseDownPos?.y) {
    this.clickedSpan = this.getSpanInPos(this.flameGraph.rootSpan, loc)
    //...
  }

  // release mouse
  this.mouseDownPos = null
  this.draw()
}

draw() {
  this.context.clearRect(0, 0, this.width, this.height)

  this.drawFlameGraph() // 包括了显示文字的操作
  this.drawClickedSpan()
}

showTooltip(windowPos: Pos) {
  if (this.tooltipDomElement === null) return

  if (this.hoverSpan === null) {
    this.tooltipDomElement.style.opacity = '0.0'
  } else {
    this.tooltipDomElement.style.opacity = '1.0'
    this.tooltipDomElement.style.transform = `translate(${
      windowPos.x + 8
    }px, ${windowPos.y + 8}px)`
    this.tooltipDomElement.innerHTML = `<span>${getValueFormat('ns')(
      this.hoverSpan.duration_ns!,
      2
    )}</span>&nbsp;&nbsp;${this.hoverSpan.event!}`
  }
}

drawFlameGraph() {
  this.context.save()
  this.drawSpan(this.flameGraph.rootSpan)
  this.context.restore()
}

drawSpan(span: IFullSpan) {
  const { start, end } = this.selectedTimeRange
  const inside =
    span.relative_end_unix_time_ns > start ||
    span.relative_begin_unix_time_ns < end

  // 只绘制区间范围内的 span
  if (inside) {
    // 对 hover 的 span 使用不同的透明度
    if (span === this.hoverSpan) {
      this.context.globalAlpha = 1.0
    } else {
      this.context.globalAlpha = 0.9
    }
    if (span.node_type === 'TiDB') {
      this.context.fillStyle = '#aab254'
    } else {
      this.context.fillStyle = '#507359'
    }
    // 计算 x, y, width, height
    let x = this.timeLenScale(span.relative_begin_unix_time_ns)
    if (x < 0) {
      x = 0
    }
    const y = span.depth * 20
    let width = Math.max(
      this.timeLenScale(span.relative_end_unix_time_ns) - x,
      TimelineDetailChart.MIN_SPAN_WIDTH
    )
    if (x + width > this.width) {
      width = this.width - x
    }
    const height = TimelineDetailChart.LAYER_HEIGHT - 1
    // 绘制矩形
    this.context.fillRect(x, y, width, height)

    // 绘制和父 span 之间的竖线
    const deltaDepth = span.depth - (span.parent?.depth || 0)
    if (deltaDepth > 1) {
      this.context.strokeStyle = this.context.fillStyle
      this.context.lineWidth = 0.5
      this.context.beginPath()
      this.context.moveTo(x, y)
      this.context.lineTo(
        x,
        y - (deltaDepth - 1) * TimelineDetailChart.LAYER_HEIGHT
      )
      this.context.stroke()
    }

    // 绘制文字
    const durationStr = getValueFormat('ns')(span.duration_ns!, 2)
    const fullStr = `${span.event!} ${durationStr}`
    const fullStrWidth = this.context.measureText(fullStr).width
    const eventStrWidth = this.context.measureText(span.event!).width
    const singleCharWidth = this.context.measureText('n').width
    this.context.textAlign = 'start'
    this.context.textBaseline = 'middle'
    this.context.fillStyle = 'white'
    this.context.globalAlpha = 1.0
    if (width >= fullStrWidth + 4) {
      // 显示完整的信息，在左边显示 event，右边显示时间
      this.context.fillText(span.event!, x + 2, y + 10)
      this.context.textAlign = 'end'
      this.context.fillText(durationStr, x + width - 2, y + 10)
    } else if (width >= eventStrWidth + 2) {
      // 只显示 event
      this.context.fillText(span.event!, x + 2, y + 10)
    } else {
      // truncate event
      const charCount = Math.floor((width - 10) / singleCharWidth)
      if (charCount > 1) {
        const str = `${span.event!.slice(0, charCount)}...`
        this.context.fillText(str, x + 2, y + 10)
      }
    }
  }

  span.children.forEach((s) => this.drawSpan(s))
}

function setupCharts(flameGraph: IFlameGraph) {
  //...

  if (overviewChartRef.current) {
    overviewChart.current = new TimelineOverviewChart(
      overviewChartRef.current!,
      flameGraph!
    )
    overviewChart.current.addTimeRangeListener((newTimeRange) => {
      detailChart.current?.setTimeRange(newTimeRange)
    })
  }
  if (detailChartRef.current) {
    detailChart.current = new TimelineDetailChart(
      detailChartRef.current!,
      flameGraph
    )
    detailChart.current.addTimeRangeListener((newTimeRange) => {
      overviewChart.current?.setTimeRange(newTimeRange)
    })
    //...
  }
}