@@ -30,6 +30,8 @@ module Control.RefCount (
3030 -- * Test API
3131 , checkForgottenRefs
3232 , ignoreForgottenRefs
33+ , enableForgottenRefChecks
34+ , disableForgottenRefChecks
3335 ) where
3436
3537import Control.DeepSeq
@@ -451,9 +453,11 @@ data RefTracker = RefTracker !RefId
451453globalRefIdSupply :: PrimVar RealWorld Int
452454globalRefIdSupply = unsafePerformIO $ newPrimVar 0
453455
456+ data Enabled a = Enabled ! a | Disabled
457+
454458{-# NOINLINE globalForgottenRef #-}
455- globalForgottenRef :: IORef (Maybe (RefId , CallStack ))
456- globalForgottenRef = unsafePerformIO $ newIORef Nothing
459+ globalForgottenRef :: IORef (Enabled ( Maybe (RefId , CallStack ) ))
460+ globalForgottenRef = unsafePerformIO $ newIORef ( Enabled Nothing )
457461
458462-- | This version of 'unsafeIOToPrim' is strict in the result of the arument
459463-- action.
@@ -492,27 +496,29 @@ finaliserRefTracker inner refid allocSite = do
492496 -- Add it to a global var which we can poll elsewhere.
493497 mref <- readIORef globalForgottenRef
494498 case mref of
499+ Disabled -> pure ()
495500 -- Just keep one, but keep the last allocated one.
496501 -- The reason for last is that when there are nested structures with
497502 -- refs then the last allocated is likely to be the outermost, which
498503 -- is the best place to start hunting for ref leaks. Otherwise one can
499504 -- go on a wild goose chase tracking down inner refs that were only
500505 -- forgotten due to an outer ref being forgotten.
501- Just (refid', _) | refid < refid' -> return ()
502- _ -> writeIORef globalForgottenRef (Just (refid, allocSite))
506+ Enabled ( Just (refid', _) ) | refid < refid' -> return ()
507+ Enabled _ -> writeIORef globalForgottenRef (Enabled ( Just (refid, allocSite) ))
503508
504509assertNoForgottenRefs :: (PrimMonad m , MonadThrow m ) => m ()
505510assertNoForgottenRefs = do
506511 mrefs <- unsafeIOToPrimStrict $ readIORef globalForgottenRef
507512 case mrefs of
508- Nothing -> return ()
509- Just (refid, allocSite) -> do
513+ Disabled -> return ()
514+ Enabled Nothing -> return ()
515+ Enabled (Just (refid, allocSite)) -> do
510516 -- Clear the var so we don't assert again.
511517 --
512518 -- Using the strict version is important here: if @m ~ IOSim s@, then
513519 -- using the non-strict version will lead to @RefNeverReleased@
514520 -- exceptions.
515- unsafeIOToPrimStrict $ writeIORef globalForgottenRef Nothing
521+ unsafeIOToPrimStrict $ writeIORef globalForgottenRef ( Enabled Nothing )
516522 throwIO (RefNeverReleased refid allocSite)
517523
518524
@@ -592,3 +598,26 @@ performMajorGCWithBlockingIfAvailable = do
592598 performMajorGC
593599#endif
594600#endif
601+
602+ -- | Enable forgotten reference checks.
603+ enableForgottenRefChecks :: IO ()
604+
605+ -- | Disable forgotten reference checks. This will error if there are already
606+ -- forgotten references while we are trying to disable the checks.
607+ disableForgottenRefChecks :: IO ()
608+
609+ #ifdef NO_IGNORE_ASSERTS
610+ enableForgottenRefChecks =
611+ modifyIORef globalForgottenRef $ \ case
612+ Disabled -> Enabled Nothing
613+ Enabled _ -> error " enableForgottenRefChecks: already enabled"
614+
615+ disableForgottenRefChecks =
616+ modifyIORef globalForgottenRef $ \ case
617+ Disabled -> error " disableForgottenRefChecks: already disabled"
618+ Enabled Nothing -> Disabled
619+ Enabled _ -> error " disableForgottenRefChecks: can not disable when there are forgotten references"
620+ #else
621+ enableForgottenRefChecks = pure ()
622+ disableForgottenRefChecks = pure ()
623+ #endif
0 commit comments