From e9cafcc84e1e7ad6bd160d66ec44aad3175f0ad6 Mon Sep 17 00:00:00 2001
From: Ernest Iliiasov <ernest@hyperdx.io>
Date: Tue, 16 Jul 2024 13:49:16 -0600
Subject: [PATCH] feat: Heatmap

---
 packages/app/src/SearchPage.tsx               | 198 ++++++++++++++----
 .../components/Heatmap/Heatmap.module.scss    |  50 +++++
 .../components/Heatmap/Heatmap.stories.tsx    |  18 ++
 .../app/src/components/Heatmap/Heatmap.tsx    | 127 +++++++++++
 4 files changed, 353 insertions(+), 40 deletions(-)
 create mode 100644 packages/app/src/components/Heatmap/Heatmap.module.scss
 create mode 100644 packages/app/src/components/Heatmap/Heatmap.stories.tsx
 create mode 100644 packages/app/src/components/Heatmap/Heatmap.tsx

diff --git a/packages/app/src/SearchPage.tsx b/packages/app/src/SearchPage.tsx
index ef8ab2b84..fae2a6c88 100644
--- a/packages/app/src/SearchPage.tsx
+++ b/packages/app/src/SearchPage.tsx
@@ -13,6 +13,8 @@ import Link from 'next/link';
 import { useRouter } from 'next/router';
 import cx from 'classnames';
 import { clamp, sub } from 'date-fns';
+import { useAtom } from 'jotai';
+import { atomWithStorage } from 'jotai/utils';
 import { Button } from 'react-bootstrap';
 import { useHotkeys } from 'react-hotkeys-hook';
 import {
@@ -30,13 +32,16 @@ import {
   useQueryParams,
   withDefault,
 } from 'use-query-params';
-import { ActionIcon, Indicator } from '@mantine/core';
+import { ActionIcon, Indicator, Tooltip as MTooltip } from '@mantine/core';
 import { notifications } from '@mantine/notifications';
 
 import { TimePicker } from '@/components/TimePicker';
 
 import { ErrorBoundary } from './components/ErrorBoundary';
-import api from './api';
+import { Heatmap } from './components/Heatmap/Heatmap';
+import { Icon } from './components/Icon';
+import api, { useMultiSeriesChartV2 } from './api';
+import { convertDateRangeToGranularityString } from './ChartUtils';
 import CreateLogAlertModal from './CreateLogAlertModal';
 import { withAppNav } from './layout';
 import LogSidePanel from './LogSidePanel';
@@ -48,12 +53,18 @@ import { SearchPageFilters, ToggleFilterButton } from './SearchPage.components';
 import SearchPageActionBar from './SearchPageActionBar';
 import { Tags } from './Tags';
 import { useTimeQuery } from './timeQuery';
+import type { TimeChartSeries } from './types';
 import { useDisplayedColumns } from './useDisplayedColumns';
 import { FormatTime, useFormatTime } from './useFormatTime';
 
 import 'react-modern-drawer/dist/index.css';
 import styles from '../styles/SearchPage.module.scss';
 
+const chartModeAtom = atomWithStorage<'heatmap' | 'histogram'>(
+  'hdx-search-page-chart-mode',
+  'histogram',
+);
+
 const HistogramBarChartTooltip = (props: any) => {
   const { active, payload, label } = props;
   if (active && payload && payload.length) {
@@ -119,7 +130,7 @@ const HDXHistogram = memo(
     const formatTime = useFormatTime();
 
     return isHistogramResultsLoading ? (
-      <div className="w-100 h-100 d-flex align-items-center justify-content-center">
+      <div className="w-100 h-100 fs-8 text-slate-300 d-flex align-items-center justify-content-center">
         Loading Graph...
       </div>
     ) : (
@@ -219,6 +230,79 @@ const HDXHistogram = memo(
   },
 );
 
+function genLogScale(total_intervals: number, start: number, end: number) {
+  const x = (Math.log(end) - Math.log(start)) / total_intervals;
+  const factor = Math.exp(x);
+  const result = [start];
+  let i;
+
+  for (i = 1; i < total_intervals; i++) {
+    result.push(result[result.length - 1] * factor);
+  }
+  result.push(end);
+  return result;
+}
+
+const HDXHeatmap = ({
+  config,
+  isLive,
+}: {
+  config: {
+    dateRange: [Date, Date];
+    where: string;
+  };
+  isLive: boolean;
+}) => {
+  const formatTime = useFormatTime();
+
+  const input = useMemo(() => {
+    const scale = genLogScale(14, 1, 30 * 60 * 1000); // ms
+    return {
+      startDate: config.dateRange[0],
+      endDate: config.dateRange[1],
+      where: config.where,
+      series: scale.map((v, i) => ({
+        type: 'time' as const,
+        table: 'logs' as const,
+        aggFn: 'count' as const,
+        where: `duration:>${scale[i - 1] || 0} AND duration:<${v} AND ${
+          config.where
+        }`,
+        groupBy: [],
+      })),
+      seriesReturnType: 'column' as const,
+      granularity: convertDateRangeToGranularityString(config.dateRange, 200),
+    };
+  }, [config]);
+
+  const { isFetching, isLoading, data } = api.useMultiSeriesChart(input, {
+    keepPreviousData: true,
+  });
+
+  const xLabels = useMemo(() => {
+    return [formatTime(config.dateRange[0]), formatTime(config.dateRange[1])];
+  }, [config.dateRange, formatTime]);
+
+  const yLabels = useMemo(() => ['0ms', '30m'], []);
+
+  if (isLoading) {
+    return (
+      <div className="w-100 fs-8 text-slate-300 h-100 d-flex align-items-center justify-content-center">
+        Loading Graph...
+      </div>
+    );
+  }
+
+  return (
+    <Heatmap
+      xLabels={xLabels}
+      yLabels={yLabels}
+      data={data}
+      isFetching={isFetching}
+    />
+  );
+};
+
 const HistogramResultCounter = ({
   config: { dateRange, where },
 }: {
@@ -681,6 +765,8 @@ function SearchPage() {
     [displayedSearchQuery, displayedTimeInputValue, doSearch],
   );
 
+  const [chartMode, setChartMode] = useAtom(chartModeAtom);
+
   return (
     <div style={{ height: '100vh' }}>
       <Head>
@@ -847,7 +933,31 @@ function SearchPage() {
           </ErrorBoundary>
           <div className="d-flex flex-column flex-grow-1">
             <div className="d-flex mx-4 mt-2 justify-content-between">
-              <div className="fs-8 text-muted">
+              <div className="fs-8 text-muted d-flex align-items-center gap-1">
+                <MTooltip label="Histogram" color="gray">
+                  <ActionIcon
+                    color="gray"
+                    variant={chartMode === 'histogram' ? 'filled' : 'subtle'}
+                    size="sm"
+                    onClick={() => setChartMode('histogram')}
+                  >
+                    <Icon
+                      name="bar-chart-line-fill"
+                      className="fs-8 text-success"
+                    />
+                  </ActionIcon>
+                </MTooltip>
+                <MTooltip color="gray" label="Heat map">
+                  <ActionIcon
+                    color="gray"
+                    size="sm"
+                    mr={4}
+                    variant={chartMode === 'heatmap' ? 'filled' : 'subtle'}
+                    onClick={() => setChartMode('heatmap')}
+                  >
+                    <Icon name="grid-fill" className="fs-8 text-success" />
+                  </ActionIcon>
+                </MTooltip>
                 {isReady ? (
                   <HistogramResultCounter
                     config={{
@@ -860,37 +970,41 @@ function SearchPage() {
                   />
                 ) : null}
               </div>
-              <div className="d-flex">
-                <Link
-                  href={generateSearchUrl(searchedQuery, [
-                    zoomOutFrom,
-                    zoomOutTo,
-                  ])}
-                  className="text-muted-hover text-decoration-none fs-8 me-3"
-                >
-                  <i className="bi bi-zoom-out me-1"></i>Zoom Out
-                </Link>
-                <Link
-                  href={generateSearchUrl(searchedQuery, [
-                    zoomInFrom,
-                    zoomInTo,
-                  ])}
-                  className="text-muted-hover text-decoration-none fs-8 me-3"
-                >
-                  <i className="bi bi-zoom-in me-1"></i>Zoom In
-                </Link>
-                <Link
-                  href={generateChartUrl({
-                    table: 'logs',
-                    aggFn: 'count',
-                    field: undefined,
-                    groupBy: ['level'],
-                  })}
-                  className="text-muted-hover text-decoration-none fs-8"
-                >
-                  <i className="bi bi-plus-circle me-1"></i>Create Chart
-                </Link>
-              </div>
+              {chartMode === 'histogram' ? (
+                <div className="d-flex">
+                  <Link
+                    href={generateSearchUrl(searchedQuery, [
+                      zoomOutFrom,
+                      zoomOutTo,
+                    ])}
+                    className="text-muted-hover text-decoration-none fs-8 me-3"
+                  >
+                    <i className="bi bi-zoom-out me-1"></i>Zoom Out
+                  </Link>
+                  <Link
+                    href={generateSearchUrl(searchedQuery, [
+                      zoomInFrom,
+                      zoomInTo,
+                    ])}
+                    className="text-muted-hover text-decoration-none fs-8 me-3"
+                  >
+                    <i className="bi bi-zoom-in me-1"></i>Zoom In
+                  </Link>
+                  <Link
+                    href={generateChartUrl({
+                      table: 'logs',
+                      aggFn: 'count',
+                      field: undefined,
+                      groupBy: ['level'],
+                    })}
+                    className="text-muted-hover text-decoration-none fs-8"
+                  >
+                    <i className="bi bi-plus-circle me-1"></i>Create Chart
+                  </Link>
+                </div>
+              ) : (
+                <div className="fs-8 text-slate-600">H Y P E R D X</div>
+              )}
             </div>
             <div style={{ height: 110 }} className="my-2 px-3 w-100">
               {/* Hack, recharts will release real fix soon https://github.com/recharts/recharts/issues/172 */}
@@ -911,11 +1025,15 @@ function SearchPage() {
                   }}
                 >
                   {isReady ? (
-                    <HDXHistogram
-                      config={chartsConfig}
-                      onTimeRangeSelect={onTimeRangeSelect}
-                      isLive={isLive}
-                    />
+                    chartMode === 'histogram' ? (
+                      <HDXHistogram
+                        config={chartsConfig}
+                        onTimeRangeSelect={onTimeRangeSelect}
+                        isLive={isLive}
+                      />
+                    ) : (
+                      <HDXHeatmap isLive={isLive} config={chartsConfig} />
+                    )
                   ) : null}
                 </div>
               </div>
diff --git a/packages/app/src/components/Heatmap/Heatmap.module.scss b/packages/app/src/components/Heatmap/Heatmap.module.scss
new file mode 100644
index 000000000..786b2309b
--- /dev/null
+++ b/packages/app/src/components/Heatmap/Heatmap.module.scss
@@ -0,0 +1,50 @@
+.wrapper {
+  position: relative;
+  height: 100%;
+}
+
+.xLabels,
+.yLabels {
+  font-size: 10px;
+  color: var(--mantine-color-gray-5);
+  line-height: 1;
+}
+
+.xLabels {
+  position: absolute;
+  bottom: 0;
+  left: 20px;
+  right: 0;
+  display: flex;
+  justify-content: space-between;
+}
+
+.yLabels {
+  position: absolute;
+  top: 0;
+  bottom: 16px;
+  left: 0;
+  display: flex;
+  flex-direction: column-reverse;
+  justify-content: space-between;
+  padding: 0 5px;
+
+  > div {
+    writing-mode: vertical-rl;
+  }
+}
+
+.heatmap {
+  position: absolute;
+  inset: 0 0 16px 22px;
+  display: grid;
+  grid-template-columns: repeat(100, 1fr);
+  grid-template-rows: repeat(10, 1fr);
+  place-items: stretch stretch;
+  gap: 1px;
+}
+
+.cell {
+  transition: background-color 100ms;
+  cursor: pointer;
+}
diff --git a/packages/app/src/components/Heatmap/Heatmap.stories.tsx b/packages/app/src/components/Heatmap/Heatmap.stories.tsx
new file mode 100644
index 000000000..28e33191b
--- /dev/null
+++ b/packages/app/src/components/Heatmap/Heatmap.stories.tsx
@@ -0,0 +1,18 @@
+import type { Meta, StoryObj } from '@storybook/react';
+
+import { Heatmap } from './Heatmap';
+
+const meta = {
+  component: Heatmap,
+} satisfies Meta<typeof Heatmap>;
+
+export default meta;
+
+type Story = StoryObj<typeof meta>;
+
+export const Default = () => (
+  <Heatmap
+    xLabels={['Jun 1 20:20:200', 'Jun 10 20:20:200']}
+    yLabels={['0ms', '30m']}
+  />
+);
diff --git a/packages/app/src/components/Heatmap/Heatmap.tsx b/packages/app/src/components/Heatmap/Heatmap.tsx
new file mode 100644
index 000000000..447561a21
--- /dev/null
+++ b/packages/app/src/components/Heatmap/Heatmap.tsx
@@ -0,0 +1,127 @@
+import React from 'react';
+import cx from 'classnames';
+
+import classes from './Heatmap.module.scss';
+
+function percentageToColor(percentage: number, warn = false) {
+  if (percentage === 0) {
+    return 'transparent';
+  }
+
+  if (warn) {
+    return `rgba(255, 200, 0, ${Math.max(percentage, 0.01)})`;
+  }
+
+  return `rgba(80, 250, 123, ${Math.max(percentage, 0.01)})`;
+}
+
+// ts_bucket, duration_bucket, count
+type HeatmapDataPoint = {
+  ts_bucket: number;
+  duration_bucket: number;
+  count: number;
+};
+type HeatmapData = HeatmapDataPoint[];
+
+const generateMockHeatmapData = (w = 100, h = 10): HeatmapData => {
+  return Array.from({ length: w }, (_, i) =>
+    Array.from({ length: h }, (_, j) => ({
+      ts_bucket: i,
+      duration_bucket: j,
+      count: Math.random() > 0.5 ? 0 : Math.random(),
+    })),
+  ).flat();
+};
+
+// const MOCK_DATA = generateMockHeatmapData(w, h);
+
+export const Heatmap = ({
+  xLabels,
+  yLabels,
+  data,
+  isFetching,
+}: {
+  data: any;
+  xLabels?: string[];
+  yLabels?: string[];
+  isFetching: boolean;
+}) => {
+  if (!data.data) {
+    return null;
+  }
+
+  const w = data.data.length;
+  const h = data.meta.filter(({ name }: { name: string }) =>
+    name.includes('series_'),
+  ).length;
+
+  const dataPoints: {
+    ts_bucket: number;
+    duration_bucket: number;
+    count: number;
+    series: string;
+  }[] = [];
+
+  let maxCount = 1;
+
+  for (const [index, point] of data.data.entries()) {
+    for (let j = 0; j <= h; j++) {
+      const count = point[`series_${j}.data`] || 0;
+      if (count > maxCount) {
+        maxCount = count;
+      }
+      // point[`series_${j}`] = point[`series_${j}`] || 0;
+      dataPoints.push({
+        ts_bucket: index,
+        duration_bucket: h - j - 4,
+        count,
+        series: `series_${j}`,
+      });
+    }
+  }
+
+  return (
+    <div className={cx(classes.wrapper, { 'effect-pulse': isFetching })}>
+      <div className={classes.yLabels}>
+        {yLabels?.map(label => (
+          <div key={label} className={classes.yLabel}>
+            {label}
+          </div>
+        ))}
+      </div>
+      <div className={classes.xLabels}>
+        {xLabels?.map(label => (
+          <div key={label} className={classes.xLabel}>
+            {label}
+          </div>
+        ))}
+      </div>
+      <div
+        className={classes.heatmap}
+        style={{
+          gridTemplateColumns: `repeat(${w}, 1fr)`,
+          gridTemplateRows: `repeat(${h}, 1fr)`,
+        }}
+      >
+        {dataPoints.map(({ ts_bucket, duration_bucket, count, series }) =>
+          count > 0 ? (
+            <div
+              key={`${ts_bucket + 1}-${duration_bucket + 1}`}
+              className={classes.heatmapCell}
+              style={{
+                backgroundColor: percentageToColor(
+                  Math.max(0.05, count / maxCount),
+                  duration_bucket < 5,
+                ),
+                gridArea: `${duration_bucket + 1} / ${
+                  ts_bucket + 1
+                } / span 1 / span 1`,
+              }}
+              title={`${count} event of ${maxCount} ${series}`}
+            />
+          ) : null,
+        )}
+      </div>
+    </div>
+  );
+};