Skip to content

Commit 6da3ad1

Browse files
Biki-dasBiki-das
committed
perf(hooks): Optimize useSyncExternalStoreWithSelector selector calls
Add intelligent selector execution tracking to useSyncExternalStoreWithSelector to avoid unnecessary re-renders and selector calls. The hook now tracks which properties are accessed during selection and only re-runs the selector when those specific properties change. - Add Proxy-based property access tracking - Track accessed properties to determine relevant state changes - Skip selector execution when only irrelevant state changes - Add tests verifying selector optimization behavior - Preserve existing memoization and isEqual behavior This improves performance by reducing unnecessary selector executions when unrelated parts of the state change.
1 parent c80b336 commit 6da3ad1

File tree

2 files changed

+300
-17
lines changed

2 files changed

+300
-17
lines changed

‎packages/use-sync-external-store/src/__tests__/useSyncExternalStoreWithSelector-test.js‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
/**
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
*
4+
* This source code is licensed under the MIT license found in the
5+
* LICENSE file in the root directory of this source tree.
6+
*
7+
* @emails react-core
8+
*/
9+
10+
'use strict';
11+
12+
let useSyncExternalStoreWithSelector;
13+
let React;
14+
let ReactDOM;
15+
let ReactDOMClient;
16+
let ReactFeatureFlags;
17+
let act;
18+
19+
describe('useSyncExternalStoreWithSelector', () => {
20+
beforeEach(() => {
21+
jest.resetModules();
22+
23+
if (gate(flags => flags.enableUseSyncExternalStoreShim)) {
24+
// Remove useSyncExternalStore from the React imports so that we use the
25+
// shim instead. Also removing startTransition, since we use that to
26+
// detect outdated 18 alphas that don't yet include useSyncExternalStore.
27+
//
28+
// Longer term, we'll probably test this branch using an actual build
29+
// of React 17.
30+
jest.mock('react', () => {
31+
const {
32+
// eslint-disable-next-line no-unused-vars
33+
startTransition: _,
34+
// eslint-disable-next-line no-unused-vars
35+
useSyncExternalStore: __,
36+
...otherExports
37+
} = jest.requireActual('react');
38+
return otherExports;
39+
});
40+
}
41+
42+
React = require('react');
43+
ReactDOM = require('react-dom');
44+
ReactDOMClient = require('react-dom/client');
45+
ReactFeatureFlags = require('shared/ReactFeatureFlags');
46+
47+
const internalAct = require('internal-test-utils').act;
48+
49+
// The internal act implementation doesn't batch updates by default, since
50+
// it's mostly used to test concurrent mode. But since these tests run
51+
// in both concurrent and legacy mode, I'm adding batching here.
52+
act = cb => internalAct(() => ReactDOM.unstable_batchedUpdates(cb));
53+
54+
if (gate(flags => flags.source)) {
55+
// The `shim/with-selector` module composes the main
56+
// `use-sync-external-store` entrypoint. In the compiled artifacts, this
57+
// is resolved to the `shim` implementation by our build config, but when
58+
// running the tests against the source files, we need to tell Jest how to
59+
// resolve it. Because this is a source module, this mock has no affect on
60+
// the build tests.
61+
jest.mock('use-sync-external-store/src/useSyncExternalStore', () =>
62+
jest.requireActual('use-sync-external-store/shim'),
63+
);
64+
}
65+
useSyncExternalStoreWithSelector =
66+
require('use-sync-external-store/shim/with-selector').useSyncExternalStoreWithSelector;
67+
});
68+
69+
function createRoot(container) {
70+
// This wrapper function exists so we can test both legacy roots and
71+
// concurrent roots.
72+
if (gate(flags => !flags.enableUseSyncExternalStoreShim)) {
73+
// The native implementation only exists in 18+, so we test using
74+
// concurrent mode.
75+
return ReactDOMClient.createRoot(container);
76+
} else {
77+
// For legacy mode, use ReactDOM.createRoot instead of ReactDOM.render
78+
const root = ReactDOMClient.createRoot(container);
79+
return {
80+
render(children) {
81+
root.render(children);
82+
},
83+
};
84+
}
85+
}
86+
87+
function createExternalStore(initialState) {
88+
const listeners = new Set();
89+
let currentState = initialState;
90+
return {
91+
set(text) {
92+
currentState = text;
93+
ReactDOM.unstable_batchedUpdates(() => {
94+
listeners.forEach(listener => listener());
95+
});
96+
},
97+
subscribe(listener) {
98+
listeners.add(listener);
99+
return () => listeners.delete(listener);
100+
},
101+
getState() {
102+
return currentState;
103+
},
104+
getSubscriberCount() {
105+
return listeners.size;
106+
},
107+
};
108+
}
109+
110+
test('should call selector on change accessible segment', async () => {
111+
const store = createExternalStore({a: '1', b: '2'});
112+
113+
const selectorFn = jest.fn();
114+
const selector = state => {
115+
selectorFn();
116+
return state.a;
117+
};
118+
119+
function App() {
120+
const data = useSyncExternalStoreWithSelector(
121+
store.subscribe,
122+
store.getState,
123+
null,
124+
selector,
125+
);
126+
return <>{data}</>;
127+
}
128+
129+
const container = document.createElement('div');
130+
const root = createRoot(container);
131+
await act(() => {
132+
root.render(<App />);
133+
});
134+
135+
expect(selectorFn).toHaveBeenCalledTimes(1);
136+
137+
await expect(async () => {
138+
await act(() => {
139+
store.set({a: '2', b: '2'});
140+
});
141+
}).toWarnDev(
142+
ReactFeatureFlags.enableUseRefAccessWarning
143+
? ['Warning: App: Unsafe read of a mutable value during render.']
144+
: [],
145+
);
146+
147+
expect(selectorFn).toHaveBeenCalledTimes(2);
148+
});
149+
150+
test('should not call selector if nothing changed', async () => {
151+
const store = createExternalStore({a: '1', b: '2'});
152+
153+
const selectorFn = jest.fn();
154+
const selector = state => {
155+
selectorFn();
156+
return state.a;
157+
};
158+
159+
function App() {
160+
const data = useSyncExternalStoreWithSelector(
161+
store.subscribe,
162+
store.getState,
163+
null,
164+
selector,
165+
);
166+
return <>{data}</>;
167+
}
168+
169+
const container = document.createElement('div');
170+
const root = createRoot(container);
171+
await act(() => {
172+
root.render(<App />);
173+
});
174+
175+
expect(selectorFn).toHaveBeenCalledTimes(1);
176+
177+
await act(() => {
178+
store.set({a: '1', b: '2'});
179+
});
180+
181+
expect(selectorFn).toHaveBeenCalledTimes(1);
182+
});
183+
184+
test('should not call selector on change not accessible segment', async () => {
185+
const store = createExternalStore({a: '1', b: '2'});
186+
187+
const selectorFn = jest.fn();
188+
const selector = state => {
189+
selectorFn();
190+
return state.a;
191+
};
192+
193+
function App() {
194+
const data = useSyncExternalStoreWithSelector(
195+
store.subscribe,
196+
store.getState,
197+
null,
198+
selector,
199+
);
200+
return <>{data}</>;
201+
}
202+
203+
const container = document.createElement('div');
204+
const root = createRoot(container);
205+
await act(() => {
206+
root.render(<App />);
207+
});
208+
209+
expect(selectorFn).toHaveBeenCalledTimes(1);
210+
211+
await act(() => {
212+
store.set({a: '1', b: '3'});
213+
});
214+
215+
expect(selectorFn).toHaveBeenCalledTimes(1);
216+
});
217+
});

‎packages/use-sync-external-store/src/useSyncExternalStoreWithSelector.js‎

Lines changed: 83 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ const {useRef, useEffect, useMemo, useDebugValue} = React;
1717

1818
// Same as useSyncExternalStore, but supports selector and isEqual arguments.
1919
export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
20-
subscribe: (() => void) => () => void,
20+
subscribe: (onStoreChange: () => void) => () => void,
2121
getSnapshot: () => Snapshot,
2222
getServerSnapshot: void | null | (() => Snapshot),
2323
selector: (snapshot: Snapshot) => Selection,
@@ -54,12 +54,44 @@ export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
5454
let hasMemo = false;
5555
let memoizedSnapshot;
5656
let memoizedSelection: Selection;
57+
let lastUsedProps: string[] = [];
58+
let hasAccessed = false;
59+
const accessedProps: string[] = [];
60+
5761
const memoizedSelector = (nextSnapshot: Snapshot) => {
62+
const getProxy = (): Snapshot => {
63+
if (
64+
!(typeof nextSnapshot === 'object') ||
65+
typeof Proxy === 'undefined'
66+
) {
67+
return nextSnapshot;
68+
}
69+
70+
const handler = {
71+
get: (target: Snapshot, prop: string, receiver: any) => {
72+
const propertyName = prop.toString();
73+
74+
if (accessedProps.indexOf(propertyName) === -1) {
75+
accessedProps.push(propertyName);
76+
}
77+
78+
const value = Reflect.get(target, prop, receiver);
79+
80+
return value;
81+
},
82+
};
83+
84+
return (new Proxy(nextSnapshot, handler): any);
85+
};
86+
5887
if (!hasMemo) {
5988
// The first time the hook is called, there is no memoized result.
6089
hasMemo = true;
6190
memoizedSnapshot = nextSnapshot;
62-
const nextSelection = selector(nextSnapshot);
91+
const nextSelection = selector(getProxy());
92+
lastUsedProps = accessedProps;
93+
hasAccessed = true;
94+
6395
if (isEqual !== undefined) {
6496
// Even if the selector has changed, the currently rendered selection
6597
// may be equal to the new selection. We should attempt to reuse the
@@ -77,31 +109,65 @@ export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
77109
}
78110

79111
// We may be able to reuse the previous invocation's result.
80-
const prevSnapshot: Snapshot = (memoizedSnapshot: any);
81-
const prevSelection: Selection = (memoizedSelection: any);
112+
const prevSnapshot = memoizedSnapshot;
113+
const prevSelection = memoizedSelection;
114+
115+
const getChangedSegments = (): string[] | void => {
116+
if (
117+
prevSnapshot === undefined ||
118+
!hasAccessed ||
119+
lastUsedProps.length === 0
120+
) {
121+
return undefined;
122+
}
123+
124+
const result: string[] = [];
125+
126+
if (
127+
nextSnapshot !== null &&
128+
typeof nextSnapshot === 'object' &&
129+
prevSnapshot !== null &&
130+
typeof prevSnapshot === 'object'
131+
) {
132+
for (let i = 0; i < lastUsedProps.length; i++) {
133+
const segmentName = lastUsedProps[i];
134+
135+
if (nextSnapshot[segmentName] !== prevSnapshot[segmentName]) {
136+
result.push(segmentName);
137+
}
138+
}
139+
}
140+
141+
return result;
142+
};
82143

83144
if (is(prevSnapshot, nextSnapshot)) {
84145
// The snapshot is the same as last time. Reuse the previous selection.
85146
return prevSelection;
86147
}
87148

88149
// The snapshot has changed, so we need to compute a new selection.
89-
const nextSelection = selector(nextSnapshot);
90-
91-
// If a custom isEqual function is provided, use that to check if the data
92-
// has changed. If it hasn't, return the previous selection. That signals
93-
// to React that the selections are conceptually equal, and we can bail
94-
// out of rendering.
95-
if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
96-
// The snapshot still has changed, so make sure to update to not keep
97-
// old references alive
150+
151+
const changedSegments = getChangedSegments();
152+
if (changedSegments === undefined || changedSegments.length > 0) {
153+
const nextSelection = selector(getProxy());
154+
lastUsedProps = accessedProps;
155+
hasAccessed = true;
156+
157+
// If a custom isEqual function is provided, use that to check if the data
158+
// has changed. If it hasn't, return the previous selection. That signals
159+
// to React that the selections are conceptually equal, and we can bail
160+
// out of rendering.
161+
if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
162+
return prevSelection;
163+
}
164+
98165
memoizedSnapshot = nextSnapshot;
166+
memoizedSelection = nextSelection;
167+
return nextSelection;
168+
} else {
99169
return prevSelection;
100170
}
101-
102-
memoizedSnapshot = nextSnapshot;
103-
memoizedSelection = nextSelection;
104-
return nextSelection;
105171
};
106172
// Assigning this to a constant so that Flow knows it can't change.
107173
const maybeGetServerSnapshot =

0 commit comments

Comments
 (0)