Fix RollbarContext without a payload (#102), make onRender work and document it (#88)

README.md

@@ -48,6 +48,7 @@ next month.
     - [Pass a Fallback UI](#pass-a-fallback-ui)
   - [`RollbarContext` Component](#rollbarcontext-component)
     - [Basic Usage](#basic-usage)
+    - [Using with `ErrorBoundary`](#using-with-errorboundary)
     - [Using with React Router](#using-with-react-router)
 - [Functions](#functions)
   - [`historyContext` to create `history.listener`](#historycontext-to-create-historylistener)
@@ -357,8 +358,8 @@ export function App(props) {
 Use the `RollbarContext` component to declaratively set the `context` value used by [Rollbar.js] when it's sending any
 messages to [Rollbar].
 
-This works for your `ErrorBoundary` from above or any other log or message sent to [Rollbar] while the `RollbarContext`
-is mounted on the tree.
+The context applies to any log or message sent to [Rollbar] while the `RollbarContext` is mounted on the tree. For
+errors that your `ErrorBoundary` from above catches, see [Using with `ErrorBoundary`](#using-with-errorboundary).
 
 Like `ErrorBoundary` above, `RollbarContext` relies on a [`Provider`] for an instance of a [Rollbar.js] client.
 
@@ -376,6 +377,41 @@ function HomePage() {
 }
 ```
 
+A change to the `context` prop is applied, and the previous context is restored on unmount.
+
+`RollbarContext` components can be nested, including with the `useRollbarContext` hook. The innermost one that's
+mounted sets the context. When it unmounts, the next one out applies again, with its current `context`.
+
+#### Using with `ErrorBoundary`
+
+Put the `RollbarContext` outside the `ErrorBoundary`, and add the `onRender` prop:
+
+```javascript
+<RollbarContext context="home" onRender>
+  <ErrorBoundary>
+    <HomePage />
+  </ErrorBoundary>
+</RollbarContext>
+```
+
+Outside, because when the `ErrorBoundary` catches an error, React removes everything inside it before the error is
+reported. A `RollbarContext` inside it has already been removed by then, so an error thrown after the first render
+is reported with the previous context.
+
+`onRender`, because by default `RollbarContext` sets the context when it mounts, which React does after the children
+have rendered and mounted. An error thrown while they're first rendering would be reported with the previous
+context. With `onRender`, `RollbarContext` sets the context during its first render, before the children render.
+
+That means the context is set before React has committed anything. React can throw the render away, for example
+when an `ErrorBoundary` around the `RollbarContext` catches an error, and nothing mounts when rendering on the
+server. So once React has finished, and an `ErrorBoundary` has reported the error, `RollbarContext` puts back the
+context of whatever is mounted. Until then, anything else sent to [Rollbar] also gets this context. That's why
+`onRender` isn't the default.
+
+On the server, give the `Provider` a `config` rather than a shared `instance`. If the instance's config doesn't
+set `payload.context`, `RollbarContext` can only put back an empty context, and on the server that replaces the
+context rollbar.js takes from the request's route for later errors.
+
 #### Using with React Router
 
 It's useful to set the `context` in [Rollbar] associated with areas of your application. On the server it's usually
@@ -554,6 +590,9 @@ function ContactDetails({ contactId }) {
 As an alternative to the [`RollbarContext`] component, you can use the `useRollbarContext` hook in your [Functional Component]
 to set the `context` in the [Rollbar.js] client provided by the [`Provider`] above in the React Tree.
 
+The hook sets the context in an effect, so like `RollbarContext` without `onRender`, it doesn't apply to errors
+thrown while the component and its children are first rendering.
+
 Here's an example of using it in several components:
 
 ```javascript

src/context-stack.js

@@ -0,0 +1,80 @@
+// The RollbarContext components and useRollbarContext hooks that are mounted
+// and setting a context, per Rollbar client. The innermost one decides the
+// client's context. When the last one goes away, the context from before the
+// first one is restored.
+const stacks = new WeakMap();
+
+let lastOrder = 0;
+
+// Parents render before their children, so an order number taken during the
+// first render ranks nested contexts from outermost to innermost, even though
+// React mounts (and runs effects for) children first. A context that mounts
+// later under an existing one also ranks after it.
+export function nextContextOrder() {
+  lastOrder += 1;
+  return lastOrder;
+}
+
+function getStack(rollbar) {
+  let stack = stacks.get(rollbar);
+  if (!stack) {
+    // rollbar.js has no default payload, so options.payload is undefined
+    // unless the config sets it.
+    stack = { base: rollbar.options.payload?.context, contexts: new Map() };
+    stacks.set(rollbar, stack);
+  }
+  return stack;
+}
+
+function applyStack(rollbar, stack) {
+  if (stack.contexts.size) {
+    const innermost = Math.max(...stack.contexts.keys());
+    rollbar.configure({ payload: { context: stack.contexts.get(innermost) } });
+    return;
+  }
+  stacks.delete(rollbar);
+  // configure() ignores undefined values and there's no way to remove the key,
+  // so restoring an unset context needs ''. In the browser that's sent the same
+  // as an unset context. On the server it isn't: rollbar.js takes the context
+  // from the request's route, and payload.context, even '', replaces it. Only
+  // onRender restores on the server, since nothing mounts there; the README
+  // covers this.
+  rollbar.configure({ payload: { context: stack.base ?? '' } });
+}
+
+// Adds the context for `order`, or applies its new value if it's already
+// there.
+export function setContext(rollbar, order, context) {
+  const stack = getStack(rollbar);
+  stack.contexts.set(order, context);
+  applyStack(rollbar, stack);
+}
+
+export function removeContext(rollbar, order) {
+  const stack = stacks.get(rollbar);
+  if (stack?.contexts.delete(order)) {
+    applyStack(rollbar, stack);
+  }
+}
+
+// For RollbarContext's onRender: sets `context` while the component renders,
+// before it has mounted and been added with setContext. React can throw that
+// render away without mounting anything, for example when an ErrorBoundary
+// around the component catches an error from its children, and nothing mounts
+// on the server. So a microtask applies the context of whatever is mounted
+// again. React commits in the same task that it finishes rendering in, and an
+// ErrorBoundary reports during the commit, so the microtask runs after both.
+// If a transition yields partway through rendering, the microtask runs then,
+// but when a child throws, React renders again from the start, synchronously,
+// before it commits.
+export function setRenderContext(rollbar, context) {
+  // Taken before the context changes, so that it's the one restored if
+  // nothing is mounted.
+  const stack = getStack(rollbar);
+  rollbar.configure({ payload: { context } });
+  Promise.resolve().then(() => {
+    if (stacks.get(rollbar) === stack) {
+      applyStack(rollbar, stack);
+    }
+  });
+}

src/hooks/use-rollbar-context.js

@@ -1,8 +1,9 @@
 'use client';
 
 import invariant from 'tiny-invariant';
-import { useEffect, useLayoutEffect } from 'react';
+import { useEffect, useLayoutEffect, useState } from 'react';
 import { useRollbar } from './use-rollbar';
+import { nextContextOrder, removeContext, setContext } from '../context-stack';
 
 // Simple version does its job
 // export function useRollbarContext(context) {
@@ -13,11 +14,11 @@ import { useRollbar } from './use-rollbar';
 export function useRollbarContext(ctx = '', isLayout = false) {
   invariant(typeof ctx === 'string', '`ctx` must be a string');
   const rollbar = useRollbar();
-  (isLayout ? useLayoutEffect : useEffect)(() => {
-    const origCtx = rollbar.options.payload.context;
-    rollbar.configure({ payload: { context: ctx } });
-    return () => {
-      rollbar.configure({ payload: { context: origCtx } });
-    };
+  // Where this component sits among nested contexts; see context-stack.js.
+  const [order] = useState(nextContextOrder);
+  const useEffectOfType = isLayout ? useLayoutEffect : useEffect;
+  useEffectOfType(() => {
+    setContext(rollbar, order, ctx);
   }, [ctx]);
+  useEffectOfType(() => () => removeContext(rollbar, order), []);
 }

src/rollbar-context.js

@@ -3,6 +3,12 @@
 import { Component } from 'react';
 import PropTypes from 'prop-types';
 import { Context, getRollbarFromContext } from './provider';
+import {
+  nextContextOrder,
+  removeContext,
+  setContext,
+  setRenderContext,
+} from './context-stack';
 
 export class RollbarContext extends Component {
   static propTypes = {
@@ -17,48 +23,42 @@ export class RollbarContext extends Component {
 
   static contextType = Context;
 
-  firstRender = true;
+  // Where this component sits among nested contexts; see context-stack.js.
+  order = nextContextOrder();
+  active = false;
 
-  constructor(props) {
-    super(props);
-    this.state = { previousContext: null };
-  }
-
-  changeContext = (storePrevious = true) => {
-    const rollbar = getRollbarFromContext(this.context);
-    const { context } = this.props;
-    if (storePrevious) {
-      this.setState({ previousContext: rollbar.options.payload.context });
-    }
-    rollbar.configure({ payload: { context } });
+  changeContext = () => {
+    this.active = true;
+    setContext(
+      getRollbarFromContext(this.context),
+      this.order,
+      this.props.context,
+    );
   };
 
   componentDidMount() {
-    const { onRender } = this.props;
-    if (!onRender) {
-      this.changeContext(true);
-    }
+    this.changeContext();
   }
 
-  componentDidUpdate() {
-    const { onRender } = this.props;
-    if (!onRender) {
-      this.changeContext(false);
+  componentDidUpdate(prevProps) {
+    const { onRender, context } = this.props;
+    if (!onRender || context !== prevProps.context) {
+      this.changeContext();
     }
   }
 
   componentWillUnmount() {
-    const rollbar = getRollbarFromContext(this.context);
-    const { previousContext } = this.state;
-    rollbar.configure({ payload: { context: previousContext } });
+    removeContext(getRollbarFromContext(this.context), this.order);
+    this.active = false;
   }
 
   render() {
-    const { onRender } = this.props;
-    if (onRender && this.firstRender) {
-      this.changeContext(true);
+    const { onRender, context } = this.props;
+    if (onRender && !this.active) {
+      // Before the children render, so that errors they throw are reported
+      // with this context. The component is added to the stack on mount.
+      setRenderContext(getRollbarFromContext(this.context), context);
     }
-    this.firstRender = false;
     return this.props.children;
   }
 }