Using ContT to abort an "infinite" loop (based on "Adventures in Looping")

loop.hs

{-
  An alternative to the (almost) never ending looping problem
  using the @ContT@ monad transformer.

  Based on the post by Drew Olson:
  https://blog.drewolson.org/adventures-in-looping
-}

module Main where

import           Control.Concurrent (threadDelay)
import           Control.Monad      (forever, void)
import           Control.Monad.Cont (callCC, liftIO, runContT)
import           System.Random      (randomRIO)


main :: IO ()
main = forever $ do
  wsUrl <- fetchConnectionUrl
  conn <- connectWebSocket wsUrl

  void . flip runContT pure $ callCC $ \abort -> forever $ do
    message <- liftIO $ readMessage conn

    case message of
      MessageA   -> liftIO $ putStrLn "Message A"
      MessageB   -> liftIO $ putStrLn "Message B"
      Disconnect -> do
        liftIO $ putStrLn "Disconnect!"
        abort ()

main2 :: IO ()
main2 = forever $ do
  wsUrl <- fetchConnectionUrl
  conn <- connectWebSocket wsUrl

  AbortReason reason <- flip runContT pure $ callCC $ \abort -> forever $ do
    message <- liftIO $ readMessage conn

    case message of
      MessageA   -> liftIO $ putStrLn "Message A"
      MessageB   -> liftIO $ putStrLn "Message B"
      Disconnect -> do
        liftIO $ putStrLn "Disconnect!"
        abort $ AbortReason "something went wrong!"
  putStrLn $ "disconnected! reason: " <> reason

data Message = MessageA
             | MessageB
             | Disconnect

newtype AbortReason = AbortReason String

sleep :: IO ()
sleep = threadDelay 300000

fetchConnectionUrl :: IO ()
fetchConnectionUrl = sleep >> putStrLn "connecting..."

connectWebSocket :: a -> IO ()
connectWebSocket _ = sleep

readMessage :: a -> IO Message
readMessage _ = do
  sleep
  n <- randomRIO (1, 3) :: IO Int
  case n of
    1 -> pure MessageA
    2 -> pure MessageB
    3 -> pure Disconnect
    _ -> error "impossible!"

zzz_example_run.txt

> main
connecting...
Message B
Message A
Message B
Message B
Message A
Message B
Message A
Message A
Disconnect!
connecting...
Message A
Message B
Message A
Disconnect!
connecting...
Disconnect!
connecting...
Disconnect!
connecting...
Message A
Message A
^CInterrupted.
> main2
connecting...
Disconnect!
disconnected! reason: something went wrong!
connecting...
Message B
Disconnect!
disconnected! reason: something went wrong!
connecting...
Message A
Message B
Message B
Disconnect!
disconnected! reason: something went wrong!
connecting...
^CInterrupted.