Subtle bug in Async code

Async.hs

module Main where
import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception
import Control.Monad

-- To start with, an intermezzo of the book's definition of Async and some key
-- operations on it.
data Async a = Async { pid :: ! ThreadId
                     , getAsync :: STM (Either SomeException a)
                     }

async :: IO a -> IO (Async a)
async m = do
  var <- newEmptyTMVarIO
  pid <- forkFinally m (atomically . putTMVar var)
  return Async {pid = pid, getAsync = readTMVar var}

cancel :: Async a -> IO ()
cancel a = killThread (pid a)

withAsync :: IO a -> (Async a -> IO b) -> IO b
withAsync m c =
  -- The async documentation is misleading here. [bracket] is not a
  -- valid alternative to inlining the code because it will call
  -- [async m] in a masked context. This will mean that later in the
  -- call to [forkFinally] we'll say something to the effect of [mask
  -- $ \restore -> ... try (restore m) ...] intending it to be the
  -- case that [m] is run *unmasked* it simply won't be.
  --
  -- This is essential for doing cancellations too. Suppose instead that
  -- [m] is running masked and loops forever in a non-blocking way,
  -- then it is impossible to kill the task! This is because bracket
  -- will cause the forkIO it spawns with to inherit as
  -- "MaskedInterruptible" state instead of
  -- [Unmasked]. forkFinally's call to [restore] around [m] will
  -- restore *this* masking state instead of the desired one of
  -- [Unmasked]. This is not observable in the book examples as far as
  -- I can tell because non of them contain purely non-blocking
  -- infinite loops, they all do things like download URLs which will
  -- block and shift temporarily into [Unmasked] states while they do
  -- so.
  --
  -- Furthermore the [async] package inlines [bracket] for performance
  -- anyways and just skips this bug entirely.
  bracket (async m) cancel c

-- A fixed version, I think.
withAsync' :: IO a -> (Async a -> IO b) -> IO b
withAsync' m body = mask $ \restore -> do
  a <- restore (async m)
  res <- restore (body a) `catch` \e -> cancel a >> throwIO (e :: SomeException)
  cancel a
  return res

-- Minimal demonstration of the bug: this hangs forever.
main :: IO ()
main = withAsync (forever $ return ()) (\_ -> return ())