...

import {
  DataFrame,
  DataFrameView,
  DataQueryResponse,
  FieldColorModeId,
  FieldDTO,
  MutableDataFrame,
  NodeGraphDataFrameFieldNames as Fields,
  TimeRange,
  FieldType,
  toDataFrame,
} from '@grafana/data';

export const secondsMetric = 'traces_service_graph_request_server_seconds_sum';
export const totalsMetric = 'traces_service_graph_request_total';
export const failedMetric = 'traces_service_graph_request_failed_total';
export const histogramMetric = 'traces_service_graph_request_server_seconds_bucket';
export const nativeHistogramMetric = 'traces_service_graph_request_server_seconds';

export const rateMetric = {
  expr: 'sum(rate(traces_spanmetrics_calls_total{}[$__range])) by (span_name)',
  topk: 5,
  params: [],
};
export const errorRateMetric = {
  expr: 'sum(rate(traces_spanmetrics_calls_total{}[$__range])) by (span_name)',
  topk: 5,
  params: ['status_code="STATUS_CODE_ERROR"'],
};
export const durationMetric = {
  expr: 'histogram_quantile(.9, sum(rate(traces_spanmetrics_latency_bucket{}[$__range])) by (le))',
  params: [],
};
export const nativeHistogramDurationMetric = {
  expr: 'histogram_quantile(.9, sum(rate(traces_spanmetrics_latency{}[$__range])) by (le))',
  params: [],
};
export const defaultTableFilter = 'span_kind="SPAN_KIND_SERVER"';

export const serviceMapMetrics = [
  secondsMetric,
  totalsMetric,
  failedMetric,
  histogramMetric,
  // These are used for debugging the tempo collection so probably not useful for service map right now.
  // 'traces_service_graph_unpaired_spans_total',
  // 'traces_service_graph_untagged_spans_total',
];

/**
 * Map response from multiple prometheus metrics into a node graph data frames with nodes and edges.
 * @param responses
 * @param range
 */
export function mapPromMetricsToServiceMap(
  responses: DataQueryResponse[],
  range: TimeRange
): { nodes: DataFrame; edges: DataFrame } {
  const frames = getMetricFrames(responses);

// First just collect data from the metrics into a map with nodes and edges as keys
  const nodesMap: Record<string, NodeObject> = {};
  const edgesMap: Record<string, EdgeObject> = {};
  // At this moment we don't have any error/success or other counts so we just use these 2
  collectMetricData(frames[totalsMetric], 'total', totalsMetric, nodesMap, edgesMap);
  collectMetricData(frames[secondsMetric], 'seconds', secondsMetric, nodesMap, edgesMap);
  collectMetricData(frames[failedMetric], 'failed', failedMetric, nodesMap, edgesMap);

collectIsInstrumented(frames[`${totalsMetric}_labels`], nodesMap);
  collectIsInstrumented(frames[`${secondsMetric}_labels`], nodesMap);
  collectIsInstrumented(frames[`${failedMetric}_labels`], nodesMap);

return convertToDataFrames(nodesMap, edgesMap, range);
}

function createServiceMapDataFrames() {
  function createDF(name: string, fields: FieldDTO[]) {
    return new MutableDataFrame({ name, fields, meta: { preferredVisualisationType: 'nodeGraph' } });
  }

const nodes = createDF('Nodes', [
    { name: Fields.id, type: FieldType.string, values: [] },
    { name: Fields.title, type: FieldType.string, config: { displayName: 'Service name' }, values: [] },
    { name: Fields.subTitle, type: FieldType.string, config: { displayName: 'Service namespace' }, values: [] },
    {
      name: Fields.mainStat,
      type: FieldType.number,
      config: { unit: 'ms/r', displayName: 'Average response time' },
      values: [],
    },
    {
      name: Fields.secondaryStat,
      type: FieldType.number,
      config: { unit: 'r/sec', displayName: 'Requests per second' },
      values: [],
    },
    {
      name: Fields.arc + 'success',
      type: FieldType.number,
      config: { displayName: 'Success', color: { fixedColor: 'green', mode: FieldColorModeId.Fixed } },
      values: [],
    },
    {
      name: Fields.arc + 'failed',
      type: FieldType.number,
      config: { displayName: 'Failed', color: { fixedColor: 'red', mode: FieldColorModeId.Fixed } },
      values: [],
    },
    {
      name: Fields.isInstrumented,
      type: FieldType.boolean,
      values: [],
    },
  ]);
  const edges = createDF('Edges', [
    { name: Fields.id, type: FieldType.string, values: [] },
    { name: Fields.source, type: FieldType.string, values: [] },
    { name: AdditionalEdgeFields.sourceName, type: FieldType.string, values: [] },
    { name: AdditionalEdgeFields.sourceNamespace, type: FieldType.string, values: [] },
    { name: Fields.target, type: FieldType.string, values: [] },
    { name: AdditionalEdgeFields.targetName, type: FieldType.string, values: [] },
    { name: AdditionalEdgeFields.targetNamespace, type: FieldType.string, values: [] },
    {
      name: Fields.mainStat,
      type: FieldType.number,
      config: { unit: 'ms/r', displayName: 'Average response time' },
      values: [],
    },
    {
      name: Fields.secondaryStat,
      type: FieldType.number,
      config: { unit: 'r/sec', displayName: 'Requests per second' },
      values: [],
    },
  ]);

return [nodes, edges];
}

/**
 * Group frames from response based on ref id which is set the same as the metric name so we know which metric is where
 * and also put it into DataFrameView so it's easier to work with.
 * @param responses
 */
function getMetricFrames(responses: DataQueryResponse[]): Record<string, DataFrameView> {
  return (responses[0]?.data || []).reduce<Record<string, DataFrameView>>((acc, frameDTO) => {
    const frame = toDataFrame(frameDTO);
    acc[frame.refId ?? 'A'] = new DataFrameView(frame);
    return acc;
  }, {});
}

type ServiceMapStatistics = {
  total?: number;
  seconds?: number;
  failed?: number;
};

type NodeObject = ServiceMapStatistics & {
  name: string;
  namespace?: string;
  isInstrumented?: boolean;
};

type EdgeObject = ServiceMapStatistics & {
  source: string;
  sourceName: string;
  sourceNamespace: string;
  target: string;
  targetName: string;
  targetNamespace: string;
};

// These fields are not necessary for rendering, so not available from the Fields enum
// Will be used for linking when namespace is present
enum AdditionalEdgeFields {
  sourceName = 'sourceName',
  sourceNamespace = 'sourceNamespace',
  targetName = 'targetName',
  targetNamespace = 'targetNamespace',
}

/**
 * Collect data from a metric into a map of nodes and edges. The metric data is modeled as counts of metric per edge
 * which is a pair of client-server nodes. This means we convert each row of the metric 1-1 to edges and than we assign
 * the metric also to server. We count the stats for server only as we show requests/transactions that particular node
 * processed not those which it generated and other stats like average transaction time then stem from that.
 * @param frame
 * @param stat
 * @param metric
 * @param nodesMap
 * @param edgesMap
 */
function collectMetricData(
  frame: DataFrameView | undefined,
  stat: keyof ServiceMapStatistics,
  metric: string,
  nodesMap: Record<string, NodeObject>,
  edgesMap: Record<string, EdgeObject>
) {
  if (!frame) {
    return;
  }

// The name of the value column is in this format
  // Improvement: figure out if it can be changed
  const valueName = `Value #${metric}`;

for (let i = 0; i < frame.length; i++) {
    const row = frame.get(i);
    const serverId = row.server_service_namespace ? `${row.server_service_namespace}/${row.server}` : row.server;
    const clientId = row.client_service_namespace ? `${row.client_service_namespace}/${row.client}` : row.client;

const edgeId = `${clientId}_${serverId}`;

if (!edgesMap[edgeId]) {
      // Create edge as it does not exist yet
      edgesMap[edgeId] = {
        target: serverId,
        targetName: row.server,
        targetNamespace: row.server_service_namespace,
        source: clientId,
        sourceName: row.client,
        sourceNamespace: row.client_service_namespace,
        [stat]: row[valueName],
      };
    } else {
      // Add stat to edge
      // We are adding the values if exists but that should not happen in general as there should be single row for
      // an edge.
      edgesMap[edgeId][stat] = (edgesMap[edgeId][stat] || 0) + row[valueName];
    }

if (!nodesMap[serverId]) {
      // Create node for server
      nodesMap[serverId] = {
        name: row.server,
        namespace: row.server_service_namespace,
        [stat]: row[valueName],
      };
    } else {
      // Add stat to server node. Sum up values if there are multiple edges targeting this server node.
      nodesMap[serverId][stat] = (nodesMap[serverId][stat] || 0) + row[valueName];
    }

if (!nodesMap[clientId]) {
      // Create the client node but don't add the stat as edge stats are attributed to the server node. This means for
      // example that the number of requests in a node show how many requests it handled not how many it generated.
      nodesMap[clientId] = {
        name: row.client,
        namespace: row.client_service_namespace,
        [stat]: 0,
      };
    }
  }
}

function collectIsInstrumented(frame: DataFrameView | undefined, nodesMap: Record<string, NodeObject>) {
  if (!frame) {
    return;
  }

for (let i = 0; i < frame.length; i++) {
    const row = frame.get(i);
    const serverId = row.server_service_namespace ? `${row.server_service_namespace}/${row.server}` : row.server;

if (nodesMap[serverId] && nodesMap[serverId].isInstrumented !== true) {
      nodesMap[serverId].isInstrumented = row.connection_type === '' || row.connection_type === 'messaging_system';
    }
  }
}

return { nodes, edges };
}

.	Edit
..	Edit
.eslintignore	Edit
CHANGELOG.md	Edit
CheatSheet.tsx	Edit
QueryField.tsx	Edit
README.md	Edit
SearchTraceQLEditor	Edit
ServiceGraphSection.tsx	Edit
VariableQueryEditor.test.tsx	Edit
VariableQueryEditor.tsx	Edit
_importedDependencies	Edit
configuration	Edit
dataquery.cue	Edit
dataquery.gen.ts	Edit
datasource.test.ts	Edit
datasource.ts	Edit
dist	Edit
graphTransform.test.ts	Edit
graphTransform.ts	Edit
img	Edit
jest-setup.js	Edit
jest.config.js	Edit
language_provider.test.ts	Edit
language_provider.ts	Edit
module.tsx	Edit
package.json	Edit
plugin.json	Edit
project.json	Edit
resultTransformer.test.ts	Edit
resultTransformer.ts	Edit
streaming.ts	Edit
test	Edit
traceql	Edit
tracking.test.ts	Edit
tracking.ts	Edit
tsconfig.json	Edit
types.ts	Edit
utils.test.ts	Edit
utils.ts	Edit
variables.test.ts	Edit
variables.ts	Edit
webpack.config.ts	Edit