diff --git a/packages/react/src/FeatureFlags/FeatureFlags.tsx b/packages/react/src/FeatureFlags/FeatureFlags.tsx index 7ea32fcea3d..3cf0806cb92 100644 --- a/packages/react/src/FeatureFlags/FeatureFlags.tsx +++ b/packages/react/src/FeatureFlags/FeatureFlags.tsx @@ -1,5 +1,5 @@ import type React from 'react' -import {useContext, useMemo, useEffect, useRef} from 'react' +import {useContext, useMemo, useEffect} from 'react' import {FeatureFlagContext} from './FeatureFlagContext' import {FeatureFlagScope, type FeatureFlags} from './FeatureFlagScope' @@ -7,11 +7,28 @@ export type FeatureFlagsProps = React.PropsWithChildren<{ flags: FeatureFlags }> -// WeakMap-based ref counting for data-dialog-scroll-optimized attribute -// Keys are component instances, values track if the instance has contributed to the count -const dialogScrollOptimizedInstances = new WeakMap() +/** + * Ref count for data-dialog-scroll-optimized attribute management. + * + * NOTE: This is temporary infrastructure while we feature flag the CSS :has() + * performance optimization (primer_react_css_has_selector_perf). Once the flag + * is removed and the optimization is the default behavior, this ref counting + * can be removed - the attribute can simply always be present. + * + * @internal - Not part of the public API + */ let dialogScrollOptimizedCount = 0 +/** + * Reset the ref count for testing purposes only. + * + * @internal - Not part of the public API. Only exported for test isolation. + */ +export function __resetDialogScrollOptimizedCount(): void { + dialogScrollOptimizedCount = 0 + document.body.removeAttribute('data-dialog-scroll-optimized') +} + export function FeatureFlags({children, flags}: FeatureFlagsProps) { const parentFeatureFlags = useContext(FeatureFlagContext) const value = useMemo(() => { @@ -20,24 +37,17 @@ export function FeatureFlags({children, flags}: FeatureFlagsProps) { }, [parentFeatureFlags, flags]) const isOptimizationEnabled = value.enabled('primer_react_css_has_selector_perf') - const instanceRef = useRef({}) // Set body attribute for CSS :has() optimization when flag is enabled useEffect(() => { if (isOptimizationEnabled) { - const instance = instanceRef.current - if (!dialogScrollOptimizedInstances.get(instance)) { - dialogScrollOptimizedInstances.set(instance, true) - dialogScrollOptimizedCount++ - document.body.setAttribute('data-dialog-scroll-optimized', '') - } + dialogScrollOptimizedCount++ + document.body.setAttribute('data-dialog-scroll-optimized', '') + return () => { - if (dialogScrollOptimizedInstances.get(instance)) { - dialogScrollOptimizedInstances.delete(instance) - dialogScrollOptimizedCount-- - if (dialogScrollOptimizedCount === 0) { - document.body.removeAttribute('data-dialog-scroll-optimized') - } + dialogScrollOptimizedCount-- + if (dialogScrollOptimizedCount === 0) { + document.body.removeAttribute('data-dialog-scroll-optimized') } } } diff --git a/packages/react/src/FeatureFlags/__tests__/FeatureFlags.test.tsx b/packages/react/src/FeatureFlags/__tests__/FeatureFlags.test.tsx index de99183b739..f09f5b4dd4e 100644 --- a/packages/react/src/FeatureFlags/__tests__/FeatureFlags.test.tsx +++ b/packages/react/src/FeatureFlags/__tests__/FeatureFlags.test.tsx @@ -1,11 +1,12 @@ import {describe, expect, it, beforeEach} from 'vitest' import {render} from '@testing-library/react' import {FeatureFlags, useFeatureFlag} from '../../FeatureFlags' +import {__resetDialogScrollOptimizedCount} from '../FeatureFlags' describe('FeatureFlags', () => { beforeEach(() => { - // Clean up body attributes between tests - document.body.removeAttribute('data-dialog-scroll-optimized') + // Reset module state between tests for isolation + __resetDialogScrollOptimizedCount() }) it('should allow a component to check if a feature flag is enabled', () => {