Simplify DOCS_RENDERED and just use STORY_RENDERED for hooks.

This means inline docs stories need to emit `STORY_RENDERED` when they render, which isn't *too* complicated.

addons/docs/src/blocks/Canvas.tsx

@@ -70,14 +70,10 @@ export const Canvas: FC<CanvasProps> = (props) => {
   const sourceContext = useContext(SourceContext);
   const { isLoading, previewProps } = getPreviewProps(props, docsContext, sourceContext);
   const { children } = props;
-  return (
+
+  return isLoading ? null : (
     <MDXProvider components={resetComponents}>
-      {/* We use `isLoading` as a key here to make a new instance of the PurePreview when we have
-          the proper set of props for the Preview. Otherwise, the preview will store an incorrect
-          value for the sourceState into its internal state. */}
-      <PurePreview key={isLoading.toString()} {...previewProps}>
-        {children}
-      </PurePreview>
+      <PurePreview {...previewProps}>{children}</PurePreview>
     </MDXProvider>
   );
 };

addons/docs/src/blocks/Story.tsx

@@ -6,13 +6,14 @@ import React, {
   useContext,
   useRef,
   useEffect,
-  useMemo,
 } from 'react';
 import { MDXProvider } from '@mdx-js/react';
+import global from 'global';
 import { resetComponents, Story as PureStory } from '@storybook/components';
 import { StoryId, toId, storyNameFromExport, StoryAnnotations, AnyFramework } from '@storybook/csf';
 import { Story as StoryType } from '@storybook/store';
-import global from 'global';
+import { addons } from '@storybook/addons';
+import Events from '@storybook/core-events';
 
 import { CURRENT_SELECTION } from './types';
 import { DocsContext, DocsContextProps } from './DocsContext';
@@ -62,7 +63,8 @@ export const getStoryId = (props: StoryProps, context: DocsContextProps): StoryI
 export const getStoryProps = <TFramework extends AnyFramework>(
   { height, inline }: StoryProps,
   story: StoryType<TFramework>,
-  context: DocsContextProps<TFramework>
+  context: DocsContextProps<TFramework>,
+  onStoryFnCalled: () => void
 ): PureStoryProps => {
   const { name: storyName, parameters } = story;
   const { docs = {} } = parameters;
@@ -80,12 +82,20 @@ export const getStoryProps = <TFramework extends AnyFramework>(
     );
   }
 
-  const boundStoryFn = () =>
-    story.unboundStoryFn({
+  const boundStoryFn = () => {
+    const storyResult = story.unboundStoryFn({
       ...context.getStoryContext(story),
       loaded: {},
     });
 
+    // We need to wait until the bound story function has actually been called before we
+    // consider the story rendered. Certain frameworks (i.e. angular) don't actually render
+    // the component in the very first react render cycle, and so we can't just wait until the
+    // `PureStory` component has been rendered to consider the underlying story "rendered".
+    onStoryFnCalled();
+    return storyResult;
+  };
+
   return {
     inline: storyIsInline,
     id: story.id,
@@ -100,17 +110,10 @@ export const getStoryProps = <TFramework extends AnyFramework>(
 
 const Story: FunctionComponent<StoryProps> = (props) => {
   const context = useContext(DocsContext);
+  const channel = addons.getChannel();
   const ref = useRef();
-  const story = useStory(getStoryId(props, context), context);
-
-  // Ensure we wait until this story is properly rendered in the docs context.
-  // The purpose of this is to ensure that that the `DOCS_RENDERED` event isn't emitted
-  // until all stories on the page have rendered.
-  const { id: storyId, registerRenderingStory } = context;
-  const storyRendered = useMemo(registerRenderingStory, [storyId]);
-  useEffect(() => {
-    if (story) storyRendered();
-  }, [story]);
+  const storyId = getStoryId(props, context);
+  const story = useStory(storyId, context);
 
   useEffect(() => {
     let cleanup: () => void;
@@ -140,7 +143,13 @@ const Story: FunctionComponent<StoryProps> = (props) => {
   if (!story) {
     return <div>Loading...</div>;
   }
-  const storyProps = getStoryProps(props, story, context);
+
+  // If we are rendering a old-style inline Story via `PureStory` below, we want to emit
+  // the `STORY_RENDERED` event when it renders. The modern mode below calls out to
+  // `Preview.renderStoryToDom()` which itself emits the event.
+  const storyProps = getStoryProps(props, story, context, () =>
+    channel.emit(Events.STORY_RENDERED, storyId)
+  );
   if (!storyProps) {
     return null;
   }

lib/addons/src/hooks.ts

@@ -5,13 +5,13 @@ import {
   DecoratorFunction,
   DecoratorApplicator,
   StoryContext,
+  StoryId,
   Args,
   LegacyStoryFn,
 } from '@storybook/csf';
 import {
   FORCE_RE_RENDER,
   STORY_RENDERED,
-  DOCS_RENDERED,
   UPDATE_STORY_ARGS,
   RESET_STORY_ARGS,
   UPDATE_GLOBALS,
@@ -33,8 +33,6 @@ interface Effect {
 
 type AbstractFunction = (...args: any[]) => any;
 
-const RenderEvents = [STORY_RENDERED, DOCS_RENDERED];
-
 export class HooksContext<TFramework extends AnyFramework> {
   hookListsMap: WeakMap<AbstractFunction, Hook[]>;
 
@@ -58,7 +56,8 @@ export class HooksContext<TFramework extends AnyFramework> {
 
   currentContext: StoryContext<TFramework> | null;
 
-  renderListener = () => {
+  renderListener = (storyId: StoryId) => {
+    if (storyId !== this.currentContext.id) return;
     this.triggerEffects();
     this.currentContext = null;
     this.removeRenderListeners();
@@ -119,12 +118,12 @@ export class HooksContext<TFramework extends AnyFramework> {
   addRenderListeners() {
     this.removeRenderListeners();
     const channel = addons.getChannel();
-    RenderEvents.forEach((e) => channel.on(e, this.renderListener));
+    channel.on(STORY_RENDERED, this.renderListener);
   }
 
   removeRenderListeners() {
     const channel = addons.getChannel();
-    RenderEvents.forEach((e) => channel.removeListener(e, this.renderListener));
+    channel.removeListener(STORY_RENDERED, this.renderListener);
   }
 }
 

lib/preview-web/src/PreviewWeb.test.ts

@@ -507,28 +507,6 @@ describe('PreviewWeb', () => {
 
         expect(mockChannel.emit).toHaveBeenCalledWith(Events.DOCS_RENDERED, 'component-one--a');
       });
-
-      it('emits DOCS_RENDERED after all stories are rendered', async () => {
-        document.location.search = '?id=component-one--a&viewMode=docs';
-        const [reactDomGate, openReactDomGate] = createGate();
-
-        let rendered;
-        (ReactDOM.render as jest.Mock).mockImplementationOnce((docsElement, element, cb) => {
-          rendered = docsElement.props.context.registerRenderingStory();
-          openReactDomGate();
-          cb();
-        });
-
-        await new PreviewWeb().initialize({ importFn, getProjectAnnotations });
-
-        // Wait for `ReactDOM.render()` to be called. We should still be waiting for the story
-        await reactDomGate;
-        expect(mockChannel.emit).not.toHaveBeenCalledWith(Events.DOCS_RENDERED, 'component-one--a');
-
-        rendered();
-        await waitForRender();
-        expect(mockChannel.emit).toHaveBeenCalledWith(Events.DOCS_RENDERED, 'component-one--a');
-      });
     });
   });
 

lib/preview-web/src/PreviewWeb.tsx

@@ -369,17 +369,6 @@ export class PreviewWeb<TFramework extends AnyFramework> {
       componentStories: () => this.storyStore.componentStoriesFromCSFFile({ csfFile }),
       loadStory: (storyId: StoryId) => this.storyStore.loadStory({ storyId }),
       renderStoryToElement: this.renderStoryToElement.bind(this),
-      // Keep track of the stories that are rendered by the <Story/> component and don't emit
-      // the DOCS_RENDERED event(below) until they have all marked themselves as rendered.
-      registerRenderingStory: () => {
-        let rendered: (v: void) => void;
-        renderingStoryPromises.push(
-          new Promise((resolve) => {
-            rendered = resolve;
-          })
-        );
-        return rendered;
-      },
       getStoryContext: (renderedStory: Story<TFramework>) =>
         ({
           ...this.storyStore.getStoryContext(renderedStory),
@@ -405,10 +394,7 @@ export class PreviewWeb<TFramework extends AnyFramework> {
         </DocsContainer>
       );
 
-      ReactDOM.render(docsElement, element, async () => {
-        await Promise.all(renderingStoryPromises);
-        this.channel.emit(Events.DOCS_RENDERED, id);
-      });
+      ReactDOM.render(docsElement, element, () => this.channel.emit(Events.DOCS_RENDERED, id));
     };
 
     // Initially render right away

lib/preview-web/src/types.ts

@@ -17,7 +17,6 @@ export interface DocsContextProps<TFramework extends AnyFramework = AnyFramework
   loadStory: (id: StoryId) => Promise<Story<TFramework>>;
   renderStoryToElement: PreviewWeb<TFramework>['renderStoryToElement'];
   getStoryContext: (story: Story<TFramework>) => StoryContextForLoaders<TFramework>;
-  registerRenderingStory: () => (v: void) => void;
 
   /**
    * mdxStoryNameToKey is an MDX-compiler-generated mapping of an MDX story's