fix: correct lock behavior for nested invocations, fixes #57

.size-limit

@@ -2,6 +2,6 @@
   {
     "path": "dist/es2019/index.js",
     "ignore": ["react-dom", "tslib"],
-    "limit": "1250 B"
+    "limit": "1350 B"
   }
 ]

.size.json

@@ -2,7 +2,7 @@
   {
     "name": "dist/es2019/index.js",
     "passed": true,
-    "size": 1227,
-    "sizeLimit": 1250
+    "size": 1285,
+    "sizeLimit": 1350
   }
 ]

__tests__/index.tsx

@@ -5,13 +5,13 @@ import React, { useState } from 'react';
 import { RemoveScrollBar } from '../src';
 import { lockAttribute } from '../src/component';
 
-const renderTest = () => {
+const renderTest = (actionName = 'toggle') => {
   const Test = () => {
     const [lock, setLock] = useState(false);
 
     return (
       <>
-        <button onClick={() => setLock(!lock)}>Toggle</button>
+        <button onClick={() => setLock(!lock)}>{actionName}</button>
         {lock ? <RemoveScrollBar /> : null}
       </>
     );
@@ -21,16 +21,39 @@ const renderTest = () => {
 };
 
 describe('RemoveScrollBar', () => {
-  it('should toggle the lock attribute on mount/unmount', () => {
+  it('should toggle the lock attribute on mount/unmount', async () => {
     const { getByRole } = renderTest();
     const button = getByRole('button');
 
     expect(document.body.getAttribute(lockAttribute)).toBeNull();
+    expect(window.getComputedStyle(document.body).overflow).toBe('');
 
     button.click();
     expect(document.body.getAttribute(lockAttribute)).toBeDefined();
+    await new Promise((resolve) => setTimeout(resolve, 1));
+    expect(window.getComputedStyle(document.body).overflow).toBe('hidden');
 
     button.click();
     expect(document.body.getAttribute(lockAttribute)).toBeNull();
+    expect(window.getComputedStyle(document.body).overflow).toBe('');
+  });
+
+  it('should handle nested cases', () => {
+    const t1 = renderTest('toggle1');
+    const t2 = renderTest('toggle2');
+    const button1 = t1.getByRole('button', { name: 'toggle1' });
+    const button2 = t2.getByRole('button', { name: 'toggle2' });
+
+    expect(document.body.getAttribute(lockAttribute)).toBeNull();
+
+    button1.click();
+    button2.click();
+    expect(document.body.getAttribute(lockAttribute)).toBeDefined();
+
+    button1.click();
+    expect(document.body.getAttribute(lockAttribute)).toBeDefined();
+
+    button2.click();
+    expect(document.body.getAttribute(lockAttribute)).toBeNull();
   });
 });

src/component.tsx

@@ -68,25 +68,40 @@ const getStyles = (
   }
 `;
 
+const getCurrentUseCounter = () => {
+  const counter = parseInt(document.body.getAttribute(lockAttribute) || '0', 10);
+
+  return isFinite(counter) ? counter : 0;
+};
+
+export const useLockAttribute = () => {
+  React.useEffect(() => {
+    document.body.setAttribute(lockAttribute, (getCurrentUseCounter() + 1).toString());
+
+    return () => {
+      const newCounter = getCurrentUseCounter() - 1;
+
+      if (newCounter <= 0) {
+        document.body.removeAttribute(lockAttribute);
+      } else {
+        document.body.setAttribute(lockAttribute, newCounter.toString());
+      }
+    };
+  }, []);
+};
+
 /**
  * Removes page scrollbar and blocks page scroll when mounted
  */
-export const RemoveScrollBar: React.FC<BodyScroll> = (props) => {
-  const { noRelative, noImportant, gapMode = 'margin' } = props;
+export const RemoveScrollBar: React.FC<BodyScroll> = ({ noRelative, noImportant, gapMode = 'margin' }) => {
+  useLockAttribute();
+
   /*
    gap will be measured on every component mount
    however it will be used only by the "first" invocation
    due to singleton nature of <Style
    */
   const gap = React.useMemo(() => getGapWidth(gapMode), [gapMode]);
 
-  React.useEffect(() => {
-    document.body.setAttribute(lockAttribute, '');
-
-    return () => {
-      document.body.removeAttribute(lockAttribute);
-    };
-  }, []);
-
   return <Style styles={getStyles(gap, !noRelative, gapMode, !noImportant ? '!important' : '')} />;
 };