[9.0] [embeddable] fix race condition in useBatchedPublishingSubjects (elastic#216399) (elastic#216527)

kibanamachine · nreese · web-flow · commit eff6a4e3531b · 2025-03-31T20:06:24.000+01:00
# Backport This will backport the following commits from `main` to `9.0`: - [[embeddable] fix race condition in useBatchedPublishingSubjects (elastic#216399)](elastic#216399)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport)  Co-authored-by: Nathan Reese <reese.nathan@elastic.co>
diff --git a/src/platform/packages/shared/presentation/presentation_publishing/publishing_subject/publishing_batcher.ts b/src/platform/packages/shared/presentation/presentation_publishing/publishing_subject/publishing_batcher.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { useEffect, useRef, useState } from 'react';
+import { useEffect, useMemo, useRef, useState } from 'react';
 import { combineLatest, debounceTime, skip } from 'rxjs';
 import { AnyPublishingSubject, PublishingSubject, UnwrapPublishingSubjectTuple } from './types';
 
@@ -120,22 +120,37 @@ export const useBatchedPublishingSubjects = <
 
   /**
    * Subscribe to all subjects and update the latest values when any of them change.
+   *
+   * Can not set up subscription in useEffect.
+   * useEffect introduces a race condition where subscriptions may emit after values are set with useState
+   * but before subscription is setup in useEffect.
+   *
+   * Can not set up subscription in useRef.
+   * useRef executes initialization logic every render.
    */
-  useEffect(() => {
-    const subscription = combineLatest(subjects)
-      .pipe(
-        // When a new observer subscribes to a BehaviorSubject, it immediately receives the current value. Skip this emit.
-        skip(1),
-        debounceTime(0)
-      )
-      .subscribe((values) => {
-        setLatestPublishedValues(values as UnwrapPublishingSubjectTuple<SubjectsType>);
-      });
-    return () => subscription.unsubscribe();
+  const subscription = useMemo(
+    () =>
+      combineLatest(subjects)
+        .pipe(
+          // When a new observer subscribes to a BehaviorSubject, it immediately receives the current value. Skip this emit.
+          skip(1),
+          debounceTime(0)
+        )
+        .subscribe((values) => {
+          setLatestPublishedValues(values as UnwrapPublishingSubjectTuple<SubjectsType>);
+        }),
     // 'subjects' gets a new reference on each call because of spread
     // Use 'useBatchedOptionalPublishingSubjects' when 'subjects' are expected to change.
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
+    []
+  );
+
+  /**
+   * Clean up subscription on unmount.
+   */
+  useEffect(() => {
+    return () => subscription.unsubscribe();
+  }, [subscription]);
 
   return latestPublishedValues;
 };
diff --git a/src/platform/packages/shared/presentation/presentation_publishing/publishing_subject/publishing_subject.test.tsx b/src/platform/packages/shared/presentation/presentation_publishing/publishing_subject/publishing_subject.test.tsx
@@ -118,6 +118,30 @@ describe('publishing subject', () => {
       expect(renderCount).toBe(2);
     });
 
+    test('useBatchedPublishingSubjects should synchronously subscribe to observables to avoid race conditions', async () => {
+      function Component() {
+        const [value1] = useBatchedPublishingSubjects(subject1);
+
+        // synchronously emit new values for observables
+        // this will cause test to fail if subscriptions are not setup synchronously
+        incrementAll();
+
+        return (
+          <>
+            <span>{`value1: ${value1}`}</span>
+          </>
+        );
+      }
+      render(<Component />);
+      await waitFor(() => {
+        expect(
+          // If there is a race condition, then 'value1: 0' will be rendered
+          // because value1 will have the original value '0' instead of latest value
+          screen.getByText('value1: 1')
+        ).toBeInTheDocument();
+      });
+    });
+
     test('should batch state updates when using useBatchedOptionalPublishingSubjects', async () => {
       let renderCount = 0;
       function Component() {
