Keep track of number of references to global context and cleanup (#1543)

* keep track of number of references to global context and cleanup

* refactor and add unit test

* fix: run prettier on dndprovider, add some additional checks in test

packages/core/react-dnd/src/__tests__/DndProvider.spec.tsx

@@ -0,0 +1,59 @@
+import { createDragDropManager } from 'dnd-core'
+import * as React from 'react'
+import { create } from 'react-test-renderer'
+import { DndContext, DndProvider } from '..'
+import TestBackend from 'react-dnd-test-backend'
+
+describe('DndProvider', () => {
+	it('reuses DragDropManager provided to it', () => {
+		let capturedManager
+		const manager = createDragDropManager(TestBackend, {}, {})
+
+		create(
+			<DndContext.Provider value={{ dragDropManager: manager }}>
+				<DndContext.Consumer>
+					{({ dragDropManager }) => {
+						capturedManager = dragDropManager
+						return null
+					}}
+				</DndContext.Consumer>
+			</DndContext.Provider>,
+		)
+
+		expect(capturedManager).toBe(manager)
+	})
+
+	it('stores DragDropManager in global context and cleans up on unmount', () => {
+		let capturedManager
+
+		const mountProvider = () =>
+			create(
+				<DndProvider backend={TestBackend}>
+					<DndContext.Consumer>
+						{({ dragDropManager }) => {
+							capturedManager = dragDropManager
+							return null
+						}}
+					</DndContext.Consumer>
+				</DndProvider>,
+			)
+
+		const globalInstance = () =>
+			global[Symbol.for('__REACT_DND_CONTEXT_INSTANCE__')]
+
+		// Single mount & unmount works
+		const root = mountProvider()
+		expect(globalInstance().dragDropManager).toEqual(capturedManager)
+		root.unmount()
+		expect(globalInstance()).toEqual(null)
+
+		// Two mounted components do a refcount
+		const rootA = mountProvider()
+		const rootB = mountProvider()
+		expect(globalInstance().dragDropManager).toEqual(capturedManager)
+		rootA.unmount()
+		expect(globalInstance().dragDropManager).toEqual(capturedManager)
+		rootB.unmount()
+		expect(globalInstance()).toEqual(null)
+	})
+})

packages/core/react-dnd/src/common/DndProvider.tsx

@@ -14,23 +14,58 @@ export type DndProviderProps<BackendContext, BackendOptions> =
 			debugMode?: boolean
 	  }
 
+let refCount = 0
+
 /**
  * A React component that provides the React-DnD context
  */
 export const DndProvider: React.FC<DndProviderProps<any, any>> = memo(
 	({ children, ...props }) => {
-		const context =
-			'manager' in props
-				? { dragDropManager: props.manager }
-				: createSingletonDndContext(
-						props.backend,
-						props.context,
-						props.options,
-						props.debugMode,
-				  )
-		return <DndContext.Provider value={context}>{children}</DndContext.Provider>
+		const [manager, isGlobalInstance] = getDndContextValue(props) // memoized from props
+
+		/**
+		 * If the global context was used to store the DND context
+		 * then where theres no more references to it we should
+		 * clean it up to avoid memory leaks
+		 */
+		React.useEffect(() => {
+			if (isGlobalInstance) {
+				refCount++
+			}
+
+			return () => {
+				if (isGlobalInstance) {
+					refCount--
+
+					if (refCount === 0) {
+						const context = getGlobalContext()
+						context[instanceSymbol] = null
+					}
+				}
+			}
+		}, [])
+
+		return <DndContext.Provider value={manager}>{children}</DndContext.Provider>
 	},
 )
+DndProvider.displayName = 'DndProvider'
+
+function getDndContextValue(props: DndProviderProps<any, any>) {
+	if ('manager' in props) {
+		const manager = { dragDropManager: props.manager }
+		return [manager, false]
+	}
+
+	const manager = createSingletonDndContext(
+		props.backend,
+		props.context,
+		props.options,
+		props.debugMode,
+	)
+	const isGlobalInstance = !props.context
+
+	return [manager, isGlobalInstance]
+}
 
 const instanceSymbol = Symbol.for('__REACT_DND_CONTEXT_INSTANCE__')
 function createSingletonDndContext<BackendContext, BackendOptions>(