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

# Backport This will backport the following commits from `main` to `9.0`: - [[embeddable] fix race condition in useBatchedPublishingSubjects (#216399)](https://github.com/elastic/kibana/pull/216399)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport)  Co-authored-by: Nathan Reese <reese.nathan@elastic.co>
2025-04-23 17:28:26 -04:00 · 2025-03-31 21:06:24 +02:00 · 2025-03-31 21:06:24 +02:00 · eff6a4e353
commit eff6a4e353
parent bb3ac40a5c
2 changed files with 52 additions and 13 deletions
--- 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;
 };
--- 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() {