Implement stacked graphs (dirty) and improve React wrapping

Signed-off-by: Julius Volz <julius.volz@gmail.com>

web/ui/mantine-ui/src/pages/query/Graph.tsx

@@ -2,7 +2,7 @@ import { FC, useEffect, useId, useState } from "react";
 import { Alert, Skeleton, Box, LoadingOverlay } from "@mantine/core";
 import { IconAlertTriangle, IconInfoCircle } from "@tabler/icons-react";
 import { InstantQueryResult } from "../../api/responseTypes/query";
-import { useAPIQuery } from "../../api/api";
+import { SuccessAPIResponse, useAPIQuery } from "../../api/api";
 import classes from "./Graph.module.css";
 import {
   GraphDisplayMode,
@@ -55,37 +55,45 @@ const Graph: FC<GraphProps> = ({
       enabled: expr !== "",
     });
 
-  // Keep the displayed chart range separate from the actual query range, so that
-  // the chart will keep displaying the old range while a query for a new range
-  // is still in progress.
-  const [displayedChartRange, setDisplayedChartRange] =
-    useState<UPlotChartRange>({
-      startTime: startTime,
-      endTime: effectiveEndTime,
-      resolution: effectiveResolution,
-    });
+  // Bundle the chart data and the displayed range together. This has two purposes:
+  // 1. If we update them separately, we cause unnecessary rerenders of the uPlot chart itself.
+  // 2. We want to keep displaying the old range in the chart while a query for a new range
+  //    is still in progress.
+  const [dataAndRange, setDataAndRange] = useState<{
+    data: SuccessAPIResponse<InstantQueryResult>;
+    range: UPlotChartRange;
+  } | null>(null);
 
   useEffect(() => {
-    setDisplayedChartRange({
-      startTime: startTime,
-      endTime: effectiveEndTime,
-      resolution: effectiveResolution,
-    });
-    // We actually want to update the displayed range only once the new data is there.
+    if (data !== undefined) {
+      setDataAndRange({
+        data: data,
+        range: {
+          startTime: startTime,
+          endTime: effectiveEndTime,
+          resolution: effectiveResolution,
+        },
+      });
+    }
+    // We actually want to update the displayed range only once the new data is there,
+    // so we don't want to include any of the range-related parameters in the dependencies.
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [data]);
 
+  // Re-execute the query when the user presses Enter (or hits the Execute button).
   useEffect(() => {
     expr !== "" && refetch();
   }, [retriggerIdx, refetch, expr, endTime, range, resolution]);
 
+  // The useElementSize hook above only gets a valid size on the second render, so this
+  // is a workaround to make the component render twice after mount.
   useEffect(() => {
-    if (data !== undefined && rerender) {
+    if (dataAndRange !== null && rerender) {
       setRerender(false);
     }
-  }, [data, rerender, setRerender]);
+  }, [dataAndRange, rerender, setRerender]);
 
-  // TODO: Share all the loading/error/empty data notices with the DataTable.
+  // TODO: Share all the loading/error/empty data notices with the DataTable?
 
   // Show a skeleton only on the first load, not on subsequent ones.
   if (isLoading) {
@@ -110,11 +118,11 @@ const Graph: FC<GraphProps> = ({
     );
   }
 
-  if (data === undefined) {
+  if (dataAndRange === null) {
     return <Alert variant="transparent">No data queried yet</Alert>;
   }
 
-  const { result, resultType } = data.data;
+  const { result, resultType } = dataAndRange.data.data;
 
   if (resultType !== "matrix") {
     return (
@@ -150,8 +158,8 @@ const Graph: FC<GraphProps> = ({
         // styles={{ loader: { width: "100%", height: "100%" } }}
       />
       <UPlotChart
-        data={data.data.result}
-        range={displayedChartRange}
+        data={dataAndRange.data.data.result}
+        range={dataAndRange.range}
         width={width}
         showExemplars={showExemplars}
         displayMode={displayMode}

web/ui/mantine-ui/src/pages/query/QueryPanel.tsx

@@ -14,7 +14,7 @@ import {
   IconGraph,
   IconTable,
 } from "@tabler/icons-react";
-import { FC, useState } from "react";
+import { FC, useCallback, useState } from "react";
 import { useAppDispatch, useAppSelector } from "../../state/hooks";
 import {
   GraphDisplayMode,
@@ -47,6 +47,24 @@ const QueryPanel: FC<PanelProps> = ({ idx, metricNames }) => {
   const panel = useAppSelector((state) => state.queryPage.panels[idx]);
   const dispatch = useAppDispatch();
 
+  const onSelectRange = useCallback(
+    (start: number, end: number) =>
+      dispatch(
+        setVisualizer({
+          idx,
+          visualizer: {
+            ...panel.visualizer,
+            range: (end - start) * 1000,
+            endTime: end * 1000,
+          },
+        })
+      ),
+    // TODO: How to have panel.visualizer in the dependencies, but not re-create
+    // the callback every time it changes by the callback's own update? This leads
+    // to extra renders of the plot further down.
+    [dispatch, idx, panel.visualizer]
+  );
+
   return (
     <Stack gap="lg">
       <ExpressionInput
@@ -220,18 +238,7 @@ const QueryPanel: FC<PanelProps> = ({ idx, metricNames }) => {
             showExemplars={panel.visualizer.showExemplars}
             displayMode={panel.visualizer.displayMode}
             retriggerIdx={retriggerIdx}
-            onSelectRange={(start: number, end: number) =>
-              dispatch(
-                setVisualizer({
-                  idx,
-                  visualizer: {
-                    ...panel.visualizer,
-                    range: (end - start) * 1000,
-                    endTime: end * 1000,
-                  },
-                })
-              )
-            }
+            onSelectRange={onSelectRange}
           />
         </Tabs.Panel>
       </Tabs>

web/ui/mantine-ui/src/pages/query/UPlotChart.tsx

@@ -10,6 +10,7 @@ import { useComputedColorScheme } from "@mantine/core";
 import "uplot/dist/uPlot.min.css";
 import "./uplot.css";
 import { getUPlotData, getUPlotOptions } from "./uPlotChartHelpers";
+import { setStackedOpts } from "./uPlotStackHelpers";
 
 export interface UPlotChartRange {
   startTime: number;
@@ -26,13 +27,19 @@ export interface UPlotChartProps {
   onSelectRange: (start: number, end: number) => void;
 }
 
+// This wrapper component translates the incoming Prometheus RangeSamples[] data to the
+// uPlot format and sets up the uPlot options object depending on the UI settings.
 const UPlotChart: FC<UPlotChartProps> = ({
   data,
   range: { startTime, endTime, resolution },
   width,
+  displayMode,
   onSelectRange,
 }) => {
   const [options, setOptions] = useState<uPlot.Options | null>(null);
+  const [processedData, setProcessedData] = useState<uPlot.AlignedData | null>(
+    null
+  );
   const { useLocalTime } = useSettings();
   const theme = useComputedColorScheme();
 
@@ -41,32 +48,49 @@ const UPlotChart: FC<UPlotChartProps> = ({
       return;
     }
 
-    setOptions(
-      getUPlotOptions(
-        width,
-        data,
-        useLocalTime,
-        theme === "light",
-        onSelectRange
-      )
+    const seriesData: uPlot.AlignedData = getUPlotData(
+      data,
+      startTime,
+      endTime,
+      resolution
     );
-  }, [width, data, useLocalTime, theme, onSelectRange]);
 
-  const seriesData: uPlot.AlignedData = getUPlotData(
+    const opts = getUPlotOptions(
+      seriesData,
+      width,
+      data,
+      useLocalTime,
+      theme === "light",
+      onSelectRange
+    );
+
+    if (displayMode === GraphDisplayMode.Stacked) {
+      setProcessedData(setStackedOpts(opts, seriesData).data);
+    } else {
+      setProcessedData(seriesData);
+    }
+
+    setOptions(opts);
+  }, [
+    width,
     data,
+    displayMode,
     startTime,
     endTime,
-    resolution
-  );
+    resolution,
+    useLocalTime,
+    theme,
+    onSelectRange,
+  ]);
 
-  if (options === null) {
+  if (options === null || processedData === null) {
     return;
   }
 
   return (
     <UplotReact
       options={options}
-      data={seriesData}
+      data={processedData}
       className={classes.uplotChart}
     />
   );

web/ui/mantine-ui/src/pages/query/uPlotChartHelpers.ts

@@ -3,7 +3,7 @@ import { formatSeries } from "../../lib/formatSeries";
 import { formatTimestamp } from "../../lib/formatTime";
 import { getSeriesColor } from "./colorPool";
 import { computePosition, shift, flip, offset } from "@floating-ui/dom";
-import uPlot, { Series } from "uplot";
+import uPlot, { AlignedData, Series } from "uplot";
 
 const formatYAxisTickValue = (y: number | null): string => {
   if (y === null) {
@@ -81,7 +81,7 @@ const formatLabels = (labels: { [key: string]: string }): string => `
                 .join("")}
             </div>`;
 
-const tooltipPlugin = (useLocalTime: boolean) => {
+const tooltipPlugin = (useLocalTime: boolean, data: AlignedData) => {
   let over: HTMLDivElement;
   let boundingLeft: number;
   let boundingTop: number;
@@ -141,7 +141,7 @@ const tooltipPlugin = (useLocalTime: boolean) => {
         }
 
         const ts = u.data[0][idx];
-        const value = u.data[selectedSeriesIdx][idx];
+        const value = data[selectedSeriesIdx][idx];
         const series = u.series[selectedSeriesIdx];
         // @ts-expect-error - uPlot doesn't have a field for labels, but we just attach some anyway.
         const labels = series.labels;
@@ -286,6 +286,7 @@ const onlyDrawPointsForDisconnectedSamplesFilter = (
 };
 
 export const getUPlotOptions = (
+  data: AlignedData,
   width: number,
   result: RangeSamples[],
   useLocalTime: boolean,
@@ -309,7 +310,7 @@ export const getUPlotOptions = (
   tzDate: useLocalTime
     ? undefined
     : (ts) => uPlot.tzDate(new Date(ts * 1e3), "Etc/UTC"),
-  plugins: [tooltipPlugin(useLocalTime)],
+  plugins: [tooltipPlugin(useLocalTime, data)],
   legend: {
     show: true,
     live: false,
@@ -408,15 +409,15 @@ export const getUPlotData = (
   }
 
   const values = inputData.map(({ values, histograms }) => {
-    // Insert nulls for all missing steps.
     const data: (number | null)[] = [];
     let valuePos = 0;
     let histogramPos = 0;
 
     for (let t = startTime; t <= endTime; t += resolution) {
-      // Allow for floating point inaccuracy.
       const currentValue = values && values[valuePos];
       const currentHistogram = histograms && histograms[histogramPos];
+
+      // Allow for floating point inaccuracy.
       if (
         currentValue &&
         values.length > valuePos &&
@@ -432,6 +433,7 @@ export const getUPlotData = (
         data.push(parseValue(currentHistogram[1].sum));
         histogramPos++;
       } else {
+        // Insert nulls for all missing steps.
         data.push(null);
       }
     }

web/ui/mantine-ui/src/pages/query/uPlotStackHelpers.ts

@@ -0,0 +1,96 @@
+import { lighten } from "@mantine/core";
+import uPlot, { AlignedData, TypedArray } from "uplot";
+
+// Stacking code adapted from https://leeoniya.github.io/uPlot/demos/stack.js
+function stack(
+  data: uPlot.AlignedData,
+  omit: (i: number) => boolean
+): { data: uPlot.AlignedData; bands: uPlot.Band[] } {
+  const data2: uPlot.AlignedData = [];
+  let bands: uPlot.Band[] = [];
+  const d0Len = data[0].length;
+  const accum = Array(d0Len);
+
+  for (let i = 0; i < d0Len; i++) {
+    accum[i] = 0;
+  }
+
+  for (let i = 1; i < data.length; i++) {
+    data2.push(
+      (omit(i)
+        ? data[i]
+        : data[i].map((v, i) => (accum[i] += +(v || 0)))) as TypedArray
+    );
+  }
+
+  for (let i = 1; i < data.length; i++) {
+    !omit(i) &&
+      bands.push({
+        series: [data.findIndex((_s, j) => j > i && !omit(j)), i],
+      });
+  }
+
+  bands = bands.filter((b) => b.series[1] > -1);
+
+  return {
+    data: [data[0]].concat(data2) as AlignedData,
+    bands,
+  };
+}
+
+export function setStackedOpts(opts: uPlot.Options, data: uPlot.AlignedData) {
+  const stacked = stack(data, (_i) => false);
+  opts.bands = stacked.bands;
+
+  opts.cursor = opts.cursor || {};
+  opts.cursor.dataIdx = (_u, seriesIdx, closestIdx, _xValue) =>
+    data[seriesIdx][closestIdx] == null ? null : closestIdx;
+
+  opts.series.forEach((s) => {
+    // s.value = (u, v, si, i) => data[si][i];
+
+    s.points = s.points || {};
+
+    if (s.stroke) {
+      s.fill = lighten(s.stroke as string, 0.6);
+    }
+
+    // scan raw unstacked data to return only real points
+    s.points.filter = (
+      _self: uPlot,
+      seriesIdx: number,
+      show: boolean,
+      _gaps?: null | number[][]
+    ): number[] | null => {
+      if (show) {
+        const pts: number[] = [];
+        data[seriesIdx].forEach((v, i) => {
+          v != null && pts.push(i);
+        });
+        return pts;
+      }
+      return null;
+    };
+  });
+
+  // force 0 to be the sum minimum this instead of the bottom series
+  opts.scales = opts.scales || {};
+  opts.scales.y = {
+    range: (_u, _min, max) => {
+      const minMax = uPlot.rangeNum(0, max, 0.1, true);
+      return [0, minMax[1]];
+    },
+  };
+
+  // restack on toggle
+  opts.hooks = opts.hooks || {};
+  opts.hooks.setSeries = opts.hooks.setSeries || [];
+  opts.hooks.setSeries.push((u, _i) => {
+    const stacked = stack(data, (i) => !u.series[i].show);
+    u.delBand(null);
+    stacked.bands.forEach((b) => u.addBand(b));
+    u.setData(stacked.data);
+  });
+
+  return { opts, data: stacked.data };
+}