From f30debf7ad2d4aed001dbf8c1ed0e4a0eae27393 Mon Sep 17 00:00:00 2001
From: Erin Noe-Payne <erin.noe.payne@gmail.com>
Date: Fri, 3 Dec 2021 05:09:55 -0500
Subject: [PATCH] TraceView: Fixes virtualized scrolling when trace view is
 opened in right pane in Explore (#42480)

---
 .../CustomScrollbar/CustomScrollbar.tsx          |  8 ++++----
 public/app/features/explore/Explore.tsx          | 16 ++++++++++++++--
 .../app/features/explore/TraceView/TraceView.tsx |  4 ++--
 .../explore/TraceView/TraceViewContainer.tsx     | 10 ++++++++--
 .../explore/__snapshots__/Explore.test.tsx.snap  |  1 +
 5 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/packages/grafana-ui/src/components/CustomScrollbar/CustomScrollbar.tsx b/packages/grafana-ui/src/components/CustomScrollbar/CustomScrollbar.tsx
index a15e07057af..fa228b87877 100644
--- a/packages/grafana-ui/src/components/CustomScrollbar/CustomScrollbar.tsx
+++ b/packages/grafana-ui/src/components/CustomScrollbar/CustomScrollbar.tsx
@@ -1,10 +1,10 @@
-import React, { FC, RefCallback, useCallback, useEffect, useRef } from 'react';
-import { isNil } from 'lodash';
-import classNames from 'classnames';
 import { css } from '@emotion/css';
+import { GrafanaTheme2 } from '@grafana/data';
+import classNames from 'classnames';
+import { isNil } from 'lodash';
+import React, { FC, RefCallback, useCallback, useEffect, useRef } from 'react';
 import Scrollbars, { positionValues } from 'react-custom-scrollbars-2';
 import { useStyles2 } from '../../themes';
-import { GrafanaTheme2 } from '@grafana/data';
 
 export type ScrollbarPosition = positionValues;
 
diff --git a/public/app/features/explore/Explore.tsx b/public/app/features/explore/Explore.tsx
index 3f341a8298e..cea1890551c 100644
--- a/public/app/features/explore/Explore.tsx
+++ b/public/app/features/explore/Explore.tsx
@@ -96,6 +96,8 @@ export type Props = ExploreProps & ConnectedProps<typeof connector>;
  * `format`, to indicate eventual transformations by the datasources' result transformers.
  */
 export class Explore extends React.PureComponent<Props, ExploreState> {
+  scrollElement: HTMLDivElement | undefined;
+
   constructor(props: Props) {
     super(props);
     this.state = {
@@ -286,7 +288,14 @@ export class Explore extends React.PureComponent<Props, ExploreState> {
 
     return (
       // If there is no data (like 404) we show a separate error so no need to show anything here
-      dataFrames.length && <TraceViewContainer exploreId={exploreId} dataFrames={dataFrames} splitOpenFn={splitOpen} />
+      dataFrames.length && (
+        <TraceViewContainer
+          exploreId={exploreId}
+          dataFrames={dataFrames}
+          splitOpenFn={splitOpen}
+          scrollElement={this.scrollElement}
+        />
+      )
     );
   }
 
@@ -313,7 +322,10 @@ export class Explore extends React.PureComponent<Props, ExploreState> {
     const showQueryInspector = openDrawer === ExploreDrawer.QueryInspector;
 
     return (
-      <CustomScrollbar autoHeightMin={'100%'}>
+      <CustomScrollbar
+        autoHeightMin={'100%'}
+        scrollRefCallback={(scrollElement) => (this.scrollElement = scrollElement || undefined)}
+      >
         <ExploreToolbar exploreId={exploreId} onChangeTime={this.onChangeTime} />
         {datasourceMissing ? this.renderEmptyState() : null}
         {datasourceInstance && (
diff --git a/public/app/features/explore/TraceView/TraceView.tsx b/public/app/features/explore/TraceView/TraceView.tsx
index d4427bfb457..f4f105ed3de 100644
--- a/public/app/features/explore/TraceView/TraceView.tsx
+++ b/public/app/features/explore/TraceView/TraceView.tsx
@@ -36,6 +36,7 @@ type Props = {
   dataFrames: DataFrame[];
   splitOpenFn: SplitOpen;
   exploreId: ExploreId;
+  scrollElement?: Element;
 };
 
 export function TraceView(props: Props) {
@@ -105,7 +106,6 @@ export function TraceView(props: Props) {
     () => createSpanLinkFactory({ splitOpenFn: props.splitOpenFn, traceToLogsOptions, dataFrame: frame }),
     [props.splitOpenFn, traceToLogsOptions, frame]
   );
-  const scrollElement = document.getElementsByClassName('scrollbar-view')[0];
   const onSlimViewClicked = useCallback(() => setSlim(!slim), [slim]);
 
   if (!props.dataFrames?.length || !traceProp) {
@@ -170,7 +170,7 @@ export function TraceView(props: Props) {
           linksGetter={noop as any}
           uiFind={search}
           createSpanLink={createSpanLink}
-          scrollElement={scrollElement}
+          scrollElement={props.scrollElement}
         />
       </UIElementsContext.Provider>
     </ThemeProvider>
diff --git a/public/app/features/explore/TraceView/TraceViewContainer.tsx b/public/app/features/explore/TraceView/TraceViewContainer.tsx
index c9f60b39f21..740d8d4374a 100644
--- a/public/app/features/explore/TraceView/TraceViewContainer.tsx
+++ b/public/app/features/explore/TraceView/TraceViewContainer.tsx
@@ -8,13 +8,19 @@ interface Props {
   dataFrames: DataFrame[];
   splitOpenFn: SplitOpen;
   exploreId: ExploreId;
+  scrollElement?: Element;
 }
 export function TraceViewContainer(props: Props) {
-  const { dataFrames, splitOpenFn, exploreId } = props;
+  const { dataFrames, splitOpenFn, exploreId, scrollElement } = props;
 
   return (
     <Collapse label="Trace View" isOpen>
-      <TraceView exploreId={exploreId} dataFrames={dataFrames} splitOpenFn={splitOpenFn} />
+      <TraceView
+        exploreId={exploreId}
+        dataFrames={dataFrames}
+        splitOpenFn={splitOpenFn}
+        scrollElement={scrollElement}
+      />
     </Collapse>
   );
 }
diff --git a/public/app/features/explore/__snapshots__/Explore.test.tsx.snap b/public/app/features/explore/__snapshots__/Explore.test.tsx.snap
index 193e62eaaa0..a27d07eef1e 100644
--- a/public/app/features/explore/__snapshots__/Explore.test.tsx.snap
+++ b/public/app/features/explore/__snapshots__/Explore.test.tsx.snap
@@ -3,6 +3,7 @@
 exports[`Explore should render component 1`] = `
 <CustomScrollbar
   autoHeightMin="100%"
+  scrollRefCallback={[Function]}
 >
   <Connect(UnConnectedExploreToolbar)
     exploreId="left"