packages feed

net-mqtt 0.8.7.3 → 0.8.8.0

raw patch · 8 files changed

+376/−423 lines, 8 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

- Data.Map.Strict.Decaying: data Map k a
- Data.Map.Strict.Decaying: delete :: Ord k => k -> Map k v -> STM ()
- Data.Map.Strict.Decaying: elems :: Map k a -> STM [a]
- Data.Map.Strict.Decaying: findWithDefault :: Ord k => v -> k -> Map k v -> STM v
- Data.Map.Strict.Decaying: insert :: Ord k => k -> v -> Map k v -> STM ()
- Data.Map.Strict.Decaying: new :: Ord k => NominalDiffTime -> IO (Map k a)
- Data.Map.Strict.Decaying: tick :: Ord k => Map k v -> IO ()
- Data.Map.Strict.Decaying: updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> STM (Maybe a)
- Data.Map.Strict.Expiring: assocs :: Ord g => Map g k a -> [(k, a)]
- Data.Map.Strict.Expiring: data Map g k a
- Data.Map.Strict.Expiring: delete :: (Ord k, Ord g) => k -> Map g k a -> Map g k a
- Data.Map.Strict.Expiring: generation :: Map g k a -> g
- Data.Map.Strict.Expiring: insert :: (Ord k, Ord g) => g -> k -> a -> Map g k a -> Map g k a
- Data.Map.Strict.Expiring: inspect :: Ord k => Map g k a -> (Int, g, Int)
- Data.Map.Strict.Expiring: instance (GHC.Show.Show a, GHC.Show.Show g) => GHC.Show.Show (Data.Map.Strict.Expiring.Entry g a)
- Data.Map.Strict.Expiring: instance (GHC.Show.Show k, GHC.Show.Show a, GHC.Show.Show g) => GHC.Show.Show (Data.Map.Strict.Expiring.Map g k a)
- Data.Map.Strict.Expiring: instance GHC.Base.Functor (Data.Map.Strict.Expiring.Entry g)
- Data.Map.Strict.Expiring: instance GHC.Base.Functor (Data.Map.Strict.Expiring.Map g k)
- Data.Map.Strict.Expiring: instance GHC.Classes.Ord g => Data.Foldable.Foldable (Data.Map.Strict.Expiring.Map g k)
- Data.Map.Strict.Expiring: lookup :: (Ord k, Ord g) => k -> Map g k a -> Maybe a
- Data.Map.Strict.Expiring: new :: g -> Map g k a
- Data.Map.Strict.Expiring: newGen :: (Ord k, Ord g) => g -> Map g k a -> Map g k a
- Data.Map.Strict.Expiring: updateLookupWithKey :: (Ord g, Ord k) => g -> (k -> a -> Maybe a) -> k -> Map g k a -> (Maybe a, Map g k a)
+ Network.MQTT.Client: instance GHC.Classes.Eq Network.MQTT.Client.InboundQoS2
+ Network.MQTT.Client: instance GHC.Show.Show Network.MQTT.Client.InboundQoS2

Files

net-mqtt.cabal view
@@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack  name:           net-mqtt-version:        0.8.7.3+version:        0.8.8.0 synopsis:       An MQTT Protocol Implementation. description:    Please see the README on GitHub at <https://github.com/dustin/mqtt-hs#readme> category:       Network@@ -27,8 +27,6 @@  library   exposed-modules:-      Data.Map.Strict.Decaying-      Data.Map.Strict.Expiring       Network.MQTT.Arbitrary       Network.MQTT.Client       Network.MQTT.Topic@@ -127,11 +125,10 @@   type: exitcode-stdio-1.0   main-is: Main.hs   other-modules:+      AuditSpec       ClientSpec-      DecayingSpec       Example1       Example2-      ExpiringSpec       Spec       Paths_net_mqtt   hs-source-dirs:
− src/Data/Map/Strict/Decaying.hs
@@ -1,66 +0,0 @@-{-# LANGUAGE RecordWildCards #-}-module Data.Map.Strict.Decaying-  ( Map,-    new,-    insert,-    delete,-    elems,-    findWithDefault,-    updateLookupWithKey,--    -- * Testing Hooks-    tick,-  )-where--import           Control.Concurrent.Async    (async, cancel, link)-import           Control.Concurrent.STM      (STM, atomically)-import           Control.Concurrent.STM.TVar (TVar, mkWeakTVar, modifyTVar', newTVarIO, readTVar, stateTVar, writeTVar)-import           Control.Exception           (mask)-import           Control.Monad               (forever, join, void, (<=<))-import           Control.Monad.Loops         (whileJust_)-import           Data.Foldable               (toList)-import qualified Data.Map.Strict.Expiring    as Map-import           Data.Maybe                  (fromMaybe, mapMaybe)-import           Data.Time                   (NominalDiffTime)-import           Data.Time.Clock.POSIX       (POSIXTime, getPOSIXTime)-import           GHC.Conc                    (threadDelay)-import           System.Mem.Weak             (deRefWeak)--data Map k a = Map {-  mapVar :: TVar (Map.Map POSIXTime k a),-  maxAge :: NominalDiffTime-}--expiry :: NominalDiffTime -> Map.Map POSIXTime k a -> POSIXTime-expiry ma = (+ ma) . Map.generation--insert :: Ord k => k -> v -> Map k v -> STM ()-insert k v (Map m ma) = modifyTVar' m (\m -> Map.insert (expiry ma m) k v m)--delete :: Ord k => k -> Map k v -> STM ()-delete k (Map m _) = modifyTVar' m (Map.delete k)--findWithDefault :: Ord k => v -> k -> Map k v -> STM v-findWithDefault d k Map{..} = fromMaybe d . Map.lookup k <$> readTVar mapVar---- | All visible records.-elems :: Map k a -> STM [a]-elems Map{..} = toList <$> readTVar mapVar--tick :: Ord k => Map k v -> IO ()-tick Map{..} = getPOSIXTime >>= \t -> atomically $ modifyTVar' mapVar (Map.newGen t)--updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> STM (Maybe a)-updateLookupWithKey f k (Map mv ma) = stateTVar mv (\m -> Map.updateLookupWithKey (expiry ma m) f k m)--new :: Ord k => NominalDiffTime -> IO (Map k a)-new maxAge = mask $ \restore -> do-  now <- getPOSIXTime-  var <- newTVarIO (Map.new now)-  wVar <- mkWeakTVar var $ pure ()-  let m = Map var maxAge-  link <=< restore $ async $ whileJust_ (deRefWeak wVar) (\m' -> tick' m' *> threadDelay 1000000)-  pure m-  where-    tick' m = getPOSIXTime >>= \t -> atomically $ modifyTVar' m (Map.newGen t)
− src/Data/Map/Strict/Expiring.hs
@@ -1,107 +0,0 @@-{-# LANGUAGE DeriveFunctor   #-}-{-# LANGUAGE RecordWildCards #-}-{-# LANGUAGE TupleSections   #-}-module Data.Map.Strict.Expiring (-    Map,-    new,-    generation,-    newGen,--    insert,-    delete,-    lookup,-    updateLookupWithKey,-    assocs,--    -- * for testing-    inspect-) where--import           Data.Foldable   (fold)-import qualified Data.Map.Strict as Map-import           Data.Maybe      (fromMaybe, mapMaybe)-import           Data.Set        (Set)-import qualified Data.Set        as Set-import           Prelude         hiding (lookup, map)--data Entry g a = Entry {-  value :: !a,-  gen   :: !g-} deriving (Functor, Show)---- | A map of values that expire after a given generation.-data Map g k a = Map {-  -- | Primary store of values.-  map        :: !(Map.Map k (Entry g a)),-  -- | The current generation-  generation :: !g,-  -- | A map of generations to keys that are expiring at that generation.-  aging      :: !(Map.Map g (Set k))-} deriving (Functor, Show)--instance Ord g => Foldable (Map g k) where-    foldMap f = foldMap (f . value) . Map.elems . map---- | Make a new empty Map at the starting generation.-new :: g -> Map g k a-new g = Map Map.empty g Map.empty---- | 𝑂(log𝑛). Assign the next generation and expire any data this new generation invalidates.--- The generation may never decrease.  Attempts to decrease it are ignored.-newGen :: (Ord k, Ord g) => g -> Map g k a -> Map g k a-newGen g m-    | g > generation m = expire m { generation = g }-    | otherwise = m---- | 𝑂(log𝑛). Insert a new value into the map to expire after the given generation.--- alterF :: (Functor f, Ord k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)-insert :: (Ord k, Ord g) => g -> k -> a -> Map g k a -> Map g k a-insert g _ _ m | g < generation m = m-insert g k v m@Map{..} = case Map.alterF (, Just (Entry v g)) k map of-    (Just old, m') -> m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen old) k aging)}-    (Nothing, m')  -> m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) aging}---- | 𝑂(log𝑛). Lookup and update.--- The function returns changed value, if it is updated. Returns the original key value if the map entry is deleted.-updateLookupWithKey :: (Ord g, Ord k) => g -> (k -> a -> Maybe a) -> k -> Map g k a -> (Maybe a, Map g k a)-updateLookupWithKey g _ _ m | g < generation m = (Nothing, m)-updateLookupWithKey g f k m@Map{..} = case Map.alterF f' k map of-    ((Nothing, _), m')   -> (Nothing, m)-    ((Just old, Nothing), m')  -> (Just (value old), m{map=m', aging = removeAging (gen old) k aging})-    ((Just old, Just new), m') -> (Just (value new), m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen old) k aging)})-    where-        f' Nothing = ((Nothing, Nothing), Nothing)-        f' (Just e) = case f k (value e) of-            Nothing -> ((Just e, Nothing), Nothing)-            Just v  -> ((Just e, Just (Entry v g)), Just (Entry v g))--removeAging :: (Ord g, Ord k) => g -> k -> Map.Map g (Set k) -> Map.Map g (Set k)-removeAging g k = Map.update (nonNull . Set.delete k) g-    where nonNull s = if Set.null s then Nothing else Just s---- | 𝑂(log𝑛). Lookup a value in the map.--- This will not return any items that have expired.-lookup :: (Ord k, Ord g) => k -> Map g k a -> Maybe a-lookup k = fmap value . Map.lookup k . map---- | 𝑂(log𝑛). Delete an item.-delete :: (Ord k, Ord g) => k -> Map g k a -> Map g k a-delete k m@Map{..} = case Map.lookup k map of-  Nothing        -> m-  Just Entry{..} -> m { map = Map.delete k map, aging = removeAging gen k aging }---- | 𝑂(𝑛). Return all current key/value associations.-assocs :: Ord g => Map g k a -> [(k,a)]-assocs Map{..} = fmap value <$> Map.assocs map---- | 𝑂(log𝑛).  Expire older generation items.-expire :: (Ord g, Ord k) => Map g k a -> Map g k a-expire m@Map{..} = m{ map = map', aging = aging'}-    where-        (todo, exact, later) = Map.splitLookup generation aging-        aging' = later <> maybe mempty (Map.singleton generation) exact-        map' = foldr Map.delete map (fold todo)---- | Inspect stored size for testing.-inspect :: Ord k => Map g k a -> (Int, g, Int)-inspect Map{..} = (Map.size map, generation, length $ fold aging)
src/Network/MQTT/Client.hs view
@@ -38,7 +38,7 @@   ) where  import           Control.Concurrent         (myThreadId, threadDelay)-import           Control.Concurrent.Async   (Async, async, asyncThreadId, cancel, cancelWith, link, race_, wait,+import           Control.Concurrent.Async   (Async, AsyncCancelled, async, asyncThreadId, cancel, cancelWith, link, race_, wait, waitCatch,                                              waitAnyCancel) import           Control.Concurrent.STM     (STM, TChan, TQueue, TVar, atomically, check, modifyTVar', newTChan,                                              newTChanIO, newTQueueIO, newTVarIO, orElse, readTChan, readTQueue, readTVar,@@ -57,11 +57,9 @@ import           Data.Conduit.Network       (AppData, appSink, appSource, clientSettings, runTCPClient) import           Data.Conduit.Network.TLS   (runTLSClient, tlsClientConfig, tlsClientTLSSettings) import           Data.Default.Class         (def)-import           Data.Foldable              (traverse_) import           Data.Map.Strict            (Map) import qualified Data.Map.Strict            as Map-import qualified Data.Map.Strict.Decaying   as Decaying-import           Data.Maybe                 (fromJust, fromMaybe, isJust)+import           Data.Maybe                 (fromJust, fromMaybe, isJust, isNothing) import           Data.Text                  (Text) import qualified Data.Text.Lazy             as TL import qualified Data.Text.Encoding         as TE@@ -82,11 +80,19 @@  data ConnState = Starting                | Connected+               | Stopping -- ^ Shutting down on request (DISCONNECT queued, or stopClient)                | Stopped                | Disconnected                | DiscoErr DisconnectRequest                | ConnErr ConnACKFlags deriving (Eq, Show) +-- | Where an inbound QoS 2 message is.  It's delivered on PUBLISH, and+-- the PUBCOMP waits for both the PUBREL and the callback to finish.+data InboundQoS2 = Delivering -- ^ callback running, no PUBREL yet+                 | Released   -- ^ PUBREL received, callback still running+                 | Delivered  -- ^ callback done, waiting for PUBREL+  deriving (Eq, Show)+ data DispatchType = DSubACK | DUnsubACK | DPubACK | DPubREC | DPubREL | DPubCOMP   deriving (Eq, Show, Ord, Enum, Bounded) @@ -104,7 +110,8 @@   -- | Callbacks are guaranteed to be invoked in the same order messages are received.   -- Slow callbacks do not block the underlying TCP connection. Pending QoS 0 messages   -- can accumulate while a callback is running. With MQTT 5, pending QoS 1 and 2-  -- messages are bounded by the Receive Maximum.+  -- messages are bounded by the Receive Maximum.  Messages received before the+  -- connection ends are still delivered after it ends, unless 'stopClient' is called.   -- Typically slower than `SimpleCallback`.   | OrderedCallback (MQTTClient -> Topic -> BL.ByteString -> [Property] -> IO ())   -- | A LowLevelCallback receives the client and the entire publish request, providing@@ -123,13 +130,13 @@   , _pktID        :: TVar Word16   , _cb           :: MessageCallback   , _acks         :: TVar (Map (DispatchType,Word16) (TChan MQTTPkt))-  , _inflight     :: Decaying.Map Word16 PublishRequest+  , _inflight     :: TVar (Map Word16 InboundQoS2) -- ^ Inbound QoS 2 packet IDs awaiting PUBREL   , _st           :: TVar ConnState   , _ct           :: TVar (Maybe (Async ()))   , _outA         :: TVar (Map Topic Word16)   , _inA          :: TVar (Map Word16 Topic)   , _connACKFlags :: TVar ConnACKFlags-  , _corr         :: Decaying.Map BL.ByteString MessageCallback+  , _corr         :: TVar (Map BL.ByteString MessageCallback)   , _cbQ          :: TQueue (IO ())   , _cbHandle     :: TVar (Maybe (Async ()))   , _outFlight    :: TVar Word16 -- ^ QoS 1 and 2 PUBLISHes in flight to the server@@ -246,16 +253,20 @@     adapt mk f = mk (f . adaptor)     adaptor s = (wsSource s, wsSink s) -    endpoint = uriPath <> uriQuery+    -- An empty path would make a request line of "GET  HTTP/1.1".+    endpoint = (if null uriPath then "/" else uriPath) <> uriQuery      cf :: Bool -> String -> Int -> String -> WS.ConnectionOptions -> WS.Headers -> WS.ClientApp () -> IO ()     cf False = WS.runClientWith     cf True  = runWSS +    -- The connection closing ends the stream, as EOF does over TCP.     wsSource :: WS.Connection -> ConduitT () BCS.ByteString IO ()-    wsSource ws = forever $ do-      bs <- liftIO $ WS.receiveData ws-      unless (BCS.null bs) $ yield bs+    wsSource ws = liftIO (E.try (WS.receiveData ws)) >>= \case+      Left (WS.CloseRequest _ _) -> pure ()+      Left WS.ConnectionClosed   -> pure ()+      Left e                     -> liftIO (E.throwIO e)+      Right bs                   -> unless (BCS.null bs) (yield bs) *> wsSource ws      wsSink :: WS.Connection -> ConduitT BCS.ByteString Void IO ()     wsSink ws = maybe (pure ()) (\bs -> liftIO (WS.sendBinaryData ws bs) >> wsSink ws) =<< await@@ -273,7 +284,8 @@       E.bracket (connectTo context connectionParams) connectionClose         (\conn -> do             stream <- makeStream (reader conn) (writer conn)-            WS.runClientWithStream stream host path options hdrs' app)+            -- Host carries the port, as it does for ws://.+            WS.runClientWithStream stream (host <> ":" <> show port) path options hdrs' app)          where           reader conn =@@ -308,13 +320,13 @@   _ch <- newTChanIO   _pktID <- newTVarIO 1   _acks <- newTVarIO mempty-  _inflight <- Decaying.new 60+  _inflight <- newTVarIO mempty   _st <- newTVarIO Starting   _ct <- newTVarIO Nothing   _outA <- newTVarIO mempty   _inA <- newTVarIO mempty   _connACKFlags <- newTVarIO (ConnACKFlags NewSession ConnUnspecifiedError mempty)-  _corr <- Decaying.new 600+  _corr <- newTVarIO mempty   _cbQ <- newTQueueIO   _cbHandle <- newTVarIO Nothing   _outFlight <- newTVarIO 0@@ -322,11 +334,14 @@       cli = MQTTClient{..}    t <- namedAsync "MQTT clientThread" $ clientThread cli-  s <- atomically (waitForLaunch cli t) -  when (s == Disconnected) $ wait t--  atomically $ checkConnected cli+  -- If we don't get connected (including a timeout or other async+  -- exception, e.g. from connectURI), the client thread must not carry on+  -- and connect anyway with nobody holding the client.+  flip E.onException (cancel t) $ do+    s <- atomically (waitForLaunch cli t)+    when (s == Disconnected) $ wait t+    atomically $ checkConnected cli    -- Ordered callbacks get messages from the start, not only after a   -- subscribe (a resumed session can deliver without one).@@ -342,6 +357,7 @@           modifyTVar' (_st cli) $ \case             Starting -> Disconnected             Connected -> Disconnected+            Stopping -> Stopped             x -> x      start c (_,sink) = do@@ -370,19 +386,31 @@                     .| conduitParser (parsePacket _protocol)                     .| C.mapM_ (\(_,x) -> liftIO (dispatch c pch $! x)) -        onceConnected = atomically $ check . (== Connected) =<< readTVar _st+        -- Past Starting, not exactly Connected: a disconnect can move us on+        -- to Stopping before these threads first look.+        onceConnected = atomically $ check . (/= Starting) =<< readTVar _st -        processOut = runConduit $-          C.repeatM (liftIO (atomically $ checkConnected c >> readTChan _ch))-          .| C.map (BL.toStrict . toByteString _protocol)-          .| sink+        processOut = runConduit $ outgoing .| C.map (BL.toStrict . toByteString _protocol) .| sink +        -- Packets go out while connected, and while stopping (to get the+        -- DISCONNECT out).  After a DISCONNECT, the client closes the+        -- connection itself [MQTT-3.14.4-1].+        outgoing = do+          p <- liftIO . atomically $ do+            st <- readTVar _st+            unless (st == Stopping) $ checkConnected c+            readTChan _ch+          yield p+          case p of+            DisconnectPkt _ -> pure ()+            _               -> outgoing+         -- Ping at least twice per Server Keep Alive, if the server set         -- one shorter than our period [MQTT-3.2.2-21].         doPing = do           ConnACKFlags _ _ props <- readTVarIO _connACKFlags           let period = minimum (_pingPeriod : [fromIntegral k * 500000 | PropServerKeepAlive k <- props, k > 0])-          forever $ threadDelay period >> sendPacketIO c PingPkt+          forever $ threadDelay period *> atomically (sendQuietly c PingPkt)          watchdog ch = forever $ do           toch <- registerDelay _pingPatience@@ -397,15 +425,26 @@ -- | Wait for a client to terminate its connection. -- An exception is thrown if the client didn't terminate expectedly. waitForClient :: MQTTClient -> IO ()-waitForClient c@MQTTClient{..} = flip E.finally (stopCallbackThread c) $ do-  void . traverse wait =<< readTVarIO _ct-  maybe (pure ()) E.throwIO =<< atomically (stateX c Stopped)+waitForClient c@MQTTClient{..} = do+  r <- traverse waitCatch =<< readTVarIO _ct+  st <- readTVarIO _st+  case (st, r) of+    -- A stopped client is fine however its thread ended (stopClient+    -- cancels it).+    (Stopped, _)                        -> pure ()+    (_, Just (Left e)) | notCanceled e  -> E.throwIO e+    _                                   -> maybe (pure ()) E.throwIO =<< atomically (stateX c Stopped) +  where notCanceled e = isNothing (E.fromException e :: Maybe AsyncCancelled)+ -- | Stops the client and closes the connection without sending a DISCONNECT -- message to the broker. This will cause the last-will message to be delivered -- by the broker if it has been defined. stopClient :: MQTTClient -> IO () stopClient MQTTClient{..} = do+  -- Stopping first, so the client thread ends as Stopped rather than+  -- Disconnected for anyone waiting on it.+  atomically (writeTVar _st Stopping)   void . traverse cancel =<< readTVarIO _ct   void . traverse cancel =<< readTVarIO _cbHandle   atomically (writeTVar _st Stopped)@@ -419,6 +458,7 @@     f :: ConnState -> Maybe E.SomeException     f Connected    = if want == Connected then Nothing else je "unexpectedly connected"     f Stopped      = if want == Stopped then Nothing else je "unexpectedly stopped"+    f Stopping     = je "stopping"     f Disconnected = je "disconnected"     f Starting     = je "died while starting"     f (DiscoErr x) = Just . E.toException . Discod $ x@@ -470,21 +510,39 @@               sendPacketIO c (DisconnectPkt (DisconnectRequest DiscoProtocolError mempty))               E.throwIO (MQTTException ("invalid topic in PUBLISH: " <> show _pubTopic)) -        pub p@PublishRequest{_pubQoS=QoS0} = atomically (resolve p) >>= notify Nothing+        pub p@PublishRequest{_pubQoS=QoS0} = atomically (resolve p) >>= notify (pure ())         pub p@PublishRequest{_pubQoS=QoS1, _pubPktID} =-          notify (Just (PubACKPkt (PubACK _pubPktID 0 mempty))) =<< atomically (resolve p)-        pub p@PublishRequest{_pubQoS=QoS2} = atomically $ do-          p'@PublishRequest{..} <- resolve p-          Decaying.insert _pubPktID p' _inflight-          sendPacket c (PubRECPkt (PubREC _pubPktID 0 mempty))+          notify (sendQuietly c (PubACKPkt (PubACK _pubPktID 0 mempty))) =<< atomically (resolve p)+        -- QoS 2 is delivered on PUBLISH, since the PUBREC takes ownership of+        -- the message [MQTT-4.3.3-9]: delivering on PUBREL would lose it if+        -- the connection dropped in between.  Only the packet ID is kept,+        -- to drop duplicates until the PUBREL.+        pub p@PublishRequest{_pubQoS=QoS2, _pubPktID=i} = do+          (p', fresh) <- atomically $ do+            p' <- resolve p+            seen <- Map.member i <$> readTVar _inflight+            unless seen $ modifyTVar' _inflight (Map.insert i Delivering)+            sendQuietly c (PubRECPkt (PubREC i 0 mempty))+            pure (p', not seen)+          when fresh $ notify (qos2Done i) p' -        pubd i = atomically (Decaying.updateLookupWithKey (\_ _ -> Nothing) i _inflight) >>= \case-            Nothing -> sendPacketIO c (PubCOMPPkt (PubCOMP i 0x92 mempty))-            Just p  -> notify (Just (PubCOMPPkt (PubCOMP i 0 mempty))) p+        qos2Done i = Map.lookup i <$> readTVar _inflight >>= \case+          Just Delivering -> modifyTVar' _inflight (Map.insert i Delivered)+          Just Released   -> modifyTVar' _inflight (Map.delete i) *> pubcomp i+          _               -> pure () -        notify rpkt p@PublishRequest{..} = do-          atomically $ Decaying.delete _pubPktID _inflight-          cb <- maybe (pure _cb) (\cd -> atomically (Decaying.findWithDefault _cb cd _corr)) cdata+        -- A PUBREL for an ID we don't know was delivered already (e.g. on+        -- an earlier connection of this session), so it's completed too.+        pubd i = atomically $ Map.lookup i <$> readTVar _inflight >>= \case+          Just Delivering -> modifyTVar' _inflight (Map.insert i Released)+          Just Released   -> pure ()+          _               -> modifyTVar' _inflight (Map.delete i) *> pubcomp i++        pubcomp i = sendQuietly c (PubCOMPPkt (PubCOMP i 0 mempty))++        -- done runs once the callback has finished (sending the ack).+        notify done p@PublishRequest{..} = do+          cb <- maybe (pure _cb) (\cd -> Map.findWithDefault _cb cd <$> readTVarIO _corr) cdata           E.evaluate . force =<< case cb of                                    NoCallback                -> pure ()                                    SimpleCallback f          -> call (f c (blToTopic _pubTopic) _pubBody _pubProps)@@ -493,9 +551,8 @@                                    OrderedLowLevelCallback f -> callOrd (f c p)              where-              call a = link =<< namedAsync "notifier" (a >> respond)-              callOrd a = atomically $ writeTQueue _cbQ (a *> respond)-              respond = traverse_ (sendPacketIO c) rpkt+              call a = link =<< namedAsync "notifier" (a *> atomically done)+              callOrd a = atomically $ writeTQueue _cbQ (a *> atomically done)               cdata = foldr f Nothing _pubProps                 where f (PropCorrelationData x) _ = Just x                       f _ o                       = o@@ -517,7 +574,7 @@           maybe (nak dt) (`writeTChan` pkt) . Map.lookup (dt, pid) =<< readTVar _acks              where-              nak DPubREC = sendPacket c (PubRELPkt  (PubREL  pid 0x92 mempty))+              nak DPubREC = sendQuietly c (PubRELPkt  (PubREL  pid 0x92 mempty))               nak _       = pure ()  @@ -528,7 +585,7 @@ -- Run OrderedCallbacks in a background thread. Does nothing for other callback types. -- We keep the async handle in a TVar and make sure only one of these threads is running. runCallbackThread :: MQTTClient -> IO ()-runCallbackThread c@MQTTClient{_cb, _cbQ, _cbHandle}+runCallbackThread c@MQTTClient{_cb, _cbQ, _cbHandle, _st}   | isOrdered _cb = do       -- We always spawn a thread, but may kill it if we already have       -- one.  The new thread won't start until we atomically confirm@@ -549,17 +606,17 @@           isOrdered (OrderedLowLevelCallback _) = True           isOrdered _                           = False           waitFor latch = atomically (check =<< readTVar latch)-          -- Keep running callbacks from the queue.  One that fails ends the-          -- connection, as with unordered callbacks, rather than leaving it-          -- up with nothing to run the rest.-          runOrderedCallbacks = forever (join (atomically (readTQueue _cbQ))) `E.catch` \(e :: E.SomeException) ->+          -- Keep running callbacks from the queue.  Once the connection+          -- has ended, finish what's queued (messages already received),+          -- then stop.  A callback that fails ends the connection, as with+          -- unordered callbacks, rather than leaving it up with nothing to+          -- run the rest.+          runOrderedCallbacks = loop `E.catch` \(e :: E.SomeException) ->             case E.fromException e of               Just (_ :: E.SomeAsyncException) -> E.throwIO e               Nothing                          -> killConn c e >> E.throwIO e---- Stop the background thread for OrderedCallbacks. Does nothing for other callback types.-stopCallbackThread :: MQTTClient -> IO ()-stopCallbackThread MQTTClient{_cbHandle} = maybe (pure ()) cancel =<< readTVarIO _cbHandle+          loop = join . atomically $ ((*> loop) <$> readTQueue _cbQ) `orElse` (readTVar _st >>= check . ended >> pure (pure ()))+          ended st = st `notElem` [Starting, Connected, Stopping]  maybeCancelWith :: E.Exception e => e -> Maybe (Async ()) -> IO () maybeCancelWith e = void . traverse (`cancelWith` e)@@ -584,6 +641,12 @@ sendPacketIO :: MQTTClient -> MQTTPkt -> IO () sendPacketIO c = atomically . sendPacket c +-- | Queue a packet if the connection is still up, otherwise drop it.  For+-- acks and pings, which shouldn't end (or race) a connection that is+-- already going away.+sendQuietly :: MQTTClient -> MQTTPkt -> STM ()+sendQuietly MQTTClient{..} p = readTVar _st >>= \st -> when (st == Connected || st == Stopping) (writeTChan _ch p)+ textToBL :: Text -> BL.ByteString textToBL = BL.fromStrict . TE.encodeUtf8 @@ -599,20 +662,31 @@ blToTopic :: BL.ByteString -> Topic blToTopic = fromJust . mkTopic . blToText +-- | Reserve the next packet ID that isn't still in use by an exchange+-- waiting for its response.  If all of them are, wait for one to free up. reservePktID :: MQTTClient -> [DispatchType] -> STM (TChan MQTTPkt, Word16) reservePktID c@MQTTClient{..} dts = do   checkConnected c   ch <- newTChan-  pid <- readTVar _pktID-  modifyTVar' _pktID $ if pid == maxBound then const 1 else succ-  modifyTVar' _acks (Map.union (Map.fromList [((t, pid), ch) | t <- dts]))-  pure (ch,pid)+  acks <- readTVar _acks+  start <- readTVar _pktID+  let next pid = if pid == maxBound then 1 else succ pid+      inUse pid = any (\t -> Map.member (t, pid) acks) [minBound ..]+      candidates = take (fromIntegral (maxBound :: Word16)) (iterate next start)+  case filter (not . inUse) candidates of+    []      -> retry+    (pid:_) -> do+      writeTVar _pktID (next pid)+      writeTVar _acks (Map.union (Map.fromList [((t, pid), ch) | t <- dts]) acks)+      pure (ch,pid) +-- Releasing an ID doesn't need a connection, and requiring one would+-- replace whatever exception is in flight with "disconnected". releasePktID :: MQTTClient -> (DispatchType,Word16) -> STM ()-releasePktID c@MQTTClient{..} k = checkConnected c >> modifyTVar' _acks (Map.delete k)+releasePktID MQTTClient{..} k = modifyTVar' _acks (Map.delete k)  releasePktIDs :: MQTTClient -> [(DispatchType,Word16)] -> STM ()-releasePktIDs c@MQTTClient{..} ks = checkConnected c >> modifyTVar' _acks deleteMany+releasePktIDs MQTTClient{..} ks = modifyTVar' _acks deleteMany   where deleteMany m = Map.withoutKeys m (Set.fromList ks)  sendAndWait :: MQTTClient -> DispatchType -> (Word16 -> MQTTPkt) -> IO MQTTPkt@@ -674,16 +748,21 @@ -- runs in the same transaction that queues the PUBLISH, so whatever it -- decides (e.g. a topic alias) goes out in the order it was decided. publishWith :: MQTTClient -> STM (Topic, [Property]) -> BL.ByteString -> Bool -> QoS -> IO ()-publishWith c@MQTTClient{_outFlight, _connACKFlags} mk m r q = do-  (ch,pid) <- atomically $ do-    (ch,pid) <- reservePktID c types-    when (q > QoS0) takeSlot-    (t, props) <- mk-    sendPacket c (pkt t props pid)-    pure (ch,pid)-  E.finally (when (q > QoS0) $ satisfyQoS ch pid) $ do-    when (q > QoS0) . atomically $ modifyTVar' _outFlight (subtract 1)-    atomically $ releasePktIDs c [(t',pid) | t' <- types]+publishWith c@MQTTClient{_outFlight, _connACKFlags} mk m r q+  -- QoS 0 has no packet ID and nothing to wait for.+  | q == QoS0 = atomically $ do+      (t, props) <- mk+      sendPacket c (pkt t props 0)+  | otherwise = do+      (ch,pid) <- atomically $ do+        (ch,pid) <- reservePktID c types+        takeSlot+        (t, props) <- mk+        sendPacket c (pkt t props pid)+        pure (ch,pid)+      E.finally (satisfyQoS ch pid) . atomically $ do+        modifyTVar' _outFlight (subtract 1)+        releasePktIDs c [(t',pid) | t' <- types]      where       types = [DPubACK, DPubREC, DPubCOMP]@@ -732,11 +811,16 @@ disconnect :: MQTTClient -> DiscoReason -> [Property] -> IO () disconnect c@MQTTClient{..} reason props = race_ getDisconnected orDieTrying   where+    -- Stopping in the same transaction that queues the DISCONNECT, so the+    -- connection ending (from either side) after it counts as a stop.     getDisconnected = do-      sendPacketIO c (DisconnectPkt $ DisconnectRequest reason props)-      void . traverse wait =<< readTVarIO _ct+      atomically $ do+        checkConnected c+        writeTVar _st Stopping+        writeTChan _ch (DisconnectPkt $ DisconnectRequest reason props)+      void . traverse waitCatch =<< readTVarIO _ct       atomically $ writeTVar _st Stopped-    orDieTrying = threadDelay 10000000 >> killConn c Timeout+    orDieTrying = threadDelay 10000000 *> atomically (writeTVar _st Disconnected) *> killConn c Timeout  -- | Disconnect with 'DiscoNormalDisconnection' and no properties. normalDisconnect :: MQTTClient -> IO ()@@ -807,8 +891,8 @@ -- -- This registration will remain in place until unregisterCorrelated is called to remove it. registerCorrelated :: MQTTClient -> BL.ByteString -> MessageCallback -> STM ()-registerCorrelated MQTTClient{_corr} bs cb = Decaying.insert bs cb _corr+registerCorrelated MQTTClient{_corr} bs cb = modifyTVar' _corr (Map.insert bs cb)  -- | Unregister a callback handler for the given correlated data identifier. unregisterCorrelated :: MQTTClient -> BL.ByteString -> STM ()-unregisterCorrelated MQTTClient{_corr} bs = Decaying.delete bs _corr+unregisterCorrelated MQTTClient{_corr} bs = modifyTVar' _corr (Map.delete bs)
+ test/AuditSpec.hs view
@@ -0,0 +1,199 @@+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE LambdaCase          #-}+{-# LANGUAGE OverloadedStrings   #-}+{-# LANGUAGE ScopedTypeVariables #-}++-- Reproductions for the Nov 2026 net-mqtt audit.  Each test asserts the+-- correct behaviour, so a failure is a reproduced bug.+module AuditSpec where++import           Control.Concurrent              (threadDelay)+import           Control.Concurrent.Async        (async, wait, waitCatch, cancel)+import           Control.Concurrent.STM+import qualified Control.Exception               as E+import           Control.Monad                   (forM, replicateM_, when, unless)+import           Control.Monad.IO.Class          (liftIO)+import qualified Data.Attoparsec.ByteString.Lazy as A+import qualified Data.ByteString                 as BS+import qualified Data.ByteString.Lazy            as BL+import           Data.Conduit                    (awaitForever, yield)+import           Data.Either                     (isLeft)+import           Data.Time.Clock                 (diffUTCTime, getCurrentTime)+import           System.Timeout                  (timeout)+import           Test.Tasty.HUnit++import           Network.MQTT.Client+import           Network.MQTT.Types              as MT++data Broker = Broker {+  written   :: TVar [MQTTPkt] -- newest first, after the CONNECT+  , inject  :: MQTTPkt -> IO ()+  , close   :: IO ()          -- end the connection from the broker side+  , connack :: IO ()          -- send the CONNACK (when held)+  , conn    :: (MQTTConduit -> IO ()) -> IO ()+  }++-- | A fake broker.  It records what the client writes (parsed), can inject+-- packets, can close the connection, and optionally closes as soon as+-- the client sends DISCONNECT (as real brokers do).+broker :: Bool -> Bool -> IO Broker+broker holdConnack closeOnDisco = do+  w <- newTVarIO []+  inj <- newTChanIO+  closed <- newTVarIO False+  ack <- newTVarIO (not holdConnack)+  first <- newTVarIO True+  let enc = BL.toStrict . toByteString Protocol50+      next = (Just <$> readTChan inj) `orElse` (readTVar closed >>= check >> pure Nothing)+      loop = liftIO (atomically next) >>= maybe (pure ()) (\p -> yield (enc p) >> loop)+      src = do+        liftIO . atomically $ readTVar ack >>= check+        yield (enc $ ConnACKPkt (ConnACKFlags NewSession ConnAccepted []))+        loop+      sink = awaitForever $ \bs -> liftIO $ do+        isFirst <- atomically $ readTVar first <* writeTVar first False+        unless isFirst $ case A.eitherResult (A.parse (parsePacket Protocol50) (BL.fromStrict bs)) of+          Right p -> do+            atomically $ modifyTVar' w (p:)+            case p of+              DisconnectPkt _ | closeOnDisco -> atomically $ writeTVar closed True+              _                              -> pure ()+          Left e -> error e+  pure Broker{written=w, inject=atomically . writeTChan inj, close=atomically (writeTVar closed True),+              connack=atomically (writeTVar ack True), conn=(\f -> f (src, sink))}++v5 :: MQTTConfig+v5 = mqttConfig{_protocol=Protocol50}++waitWritten :: Broker -> (MQTTPkt -> Bool) -> IO (Maybe MQTTPkt)+waitWritten b f = timeout 3000000 . atomically $ do+  ws <- readTVar (written b)+  case filter f ws of+    (x:_) -> pure x+    _     -> retry++-- QoS 0 publishes still reserve a packet ID.  After 65535 of them the+-- counter wraps onto the ID of a QoS 1 publish that is still waiting for+-- its PUBACK; the QoS 0 publish replaces, then deletes, the waiter's ack+-- registration, so the real PUBACK is dropped and the publish hangs.+unit_audit_qos0WrapLosesQoS1Ack :: Assertion+unit_audit_qos0WrapLosesQoS1Ack = wrapCase 65535++-- Control: the same with fewer QoS 0 publishes passes.+unit_audit_qos0NoWrapControl :: Assertion+unit_audit_qos0NoWrapControl = wrapCase 1000++wrapCase :: Int -> Assertion+wrapCase n = do+  b <- broker False False+  c <- runMQTTConduit (conn b) v5+  p <- async $ publishq c "a/b" "x" False QoS1 []+  Just (PublishPkt PublishRequest{_pubPktID=pid}) <- waitWritten b (\case PublishPkt _ -> True; _ -> False)+  replicateM_ n $ publish c "a/b" "y" False+  inject b (PubACKPkt (PubACK pid 0 []))+  r <- timeout 3000000 (wait p)+  stopClient c+  assertEqual ("QoS 1 publish (pid " <> show pid <> ") completes after its PUBACK") (Just ()) r++-- waitForClient after stopClient should return normally (the state is+-- Stopped), but it can see the transient Disconnected written by markDisco+-- before stopClient writes Stopped.+unit_audit_waitForClientStopRace :: Assertion+unit_audit_waitForClientStopRace = do+  rs <- forM [1 .. 300 :: Int] $ \_ -> do+    b <- broker False False+    c <- runMQTTConduit (conn b) v5+    w <- async (waitForClient c)+    threadDelay 1000+    stopClient c+    either (Just . show) (const Nothing) <$> waitCatch w+  let bad = [e | Just e <- rs]+  assertEqual ("waitForClient threw after stopClient (" <> show (take 1 bad) <> ")") 0 (length bad)++-- Same race with normalDisconnect against a broker that closes on DISCONNECT.+unit_audit_waitForClientDisconnectRace :: Assertion+unit_audit_waitForClientDisconnectRace = do+  rs <- forM [1 .. 300 :: Int] $ \_ -> do+    b <- broker False True+    c <- runMQTTConduit (conn b) v5+    w <- async (waitForClient c)+    threadDelay 1000+    normalDisconnect c+    either (Just . show) (const Nothing) <$> waitCatch w+  let bad = [e | Just e <- rs]+  assertEqual ("waitForClient threw after normalDisconnect (" <> show (take 1 bad) <> ")") 0 (length bad)++-- After sending DISCONNECT the client must close the connection itself+-- [MQTT-3.14.4-1].  net-mqtt waits for the server to close it, so against a+-- broker that waits for the client, disconnect takes 10s and waitForClient+-- reports a Timeout.+unit_audit_disconnectDoesNotClose :: Assertion+unit_audit_disconnectDoesNotClose = do+  b <- broker False False+  c <- runMQTTConduit (conn b) v5+  w <- async (waitForClient c)+  t0 <- getCurrentTime+  normalDisconnect c+  t1 <- getCurrentTime+  r <- waitCatch w+  assertBool ("normalDisconnect took " <> show (diffUTCTime t1 t0) <> ", waitForClient: " <> show r)+    (diffUTCTime t1 t0 < 2 && not (isLeft r))++-- connectURI wraps the connect in a timeout, but the client thread is a+-- separate async that the timeout doesn't cancel.  If the CONNACK shows up+-- after the timeout, the "failed" client connects and keeps pinging.+unit_audit_connectTimeoutLeaksClient :: Assertion+unit_audit_connectTimeoutLeaksClient = do+  b <- broker True False+  r <- timeout 300000 (runMQTTConduit (conn b) v5{_pingPeriod=100000})+  assertBool "connect timed out" (null' r)+  connack b+  threadDelay 1000000+  pings <- length . filter (== PingPkt) <$> readTVarIO (written b)+  assertEqual "PINGREQs from a client whose connect timed out" 0 pings+  where null' = maybe True (const False)++-- QoS 2 inbound messages are only delivered on PUBREL.  The PUBREC already+-- told the broker the client owns the message [MQTT-4.3.3-9], so if the+-- connection drops in between, the broker resends only PUBREL on the+-- resumed session, the new client answers PUBCOMP 0x92, and the message is+-- never delivered.+unit_audit_qos2LostAcrossReconnect :: Assertion+unit_audit_qos2LostAcrossReconnect = do+  got <- newTVarIO (0 :: Int)+  let cfg = v5{Network.MQTT.Client._cleanSession=False, _msgCB=SimpleCallback (\_ _ _ _ -> atomically $ modifyTVar' got succ)}+  b1 <- broker False False+  c1 <- runMQTTConduit (conn b1) cfg+  inject b1 (PublishPkt PublishRequest{_pubDup=False, _pubQoS=QoS2, _pubRetain=False, _pubTopic="a/b",+                                       _pubPktID=7, _pubBody="x", _pubProps=[]})+  Just _ <- waitWritten b1 (\case PubRECPkt _ -> True; _ -> False)+  close b1+  _ <- E.try @E.SomeException (waitForClient c1)+  b2 <- broker False False+  c2 <- runMQTTConduit (conn b2) cfg+  inject b2 (PubRELPkt (PubREL 7 0 []))+  Just comp <- waitWritten b2 (\case PubCOMPPkt _ -> True; _ -> False)+  threadDelay 200000+  n <- readTVarIO got+  stopClient c2+  assertEqual ("deliveries of the QoS 2 message (client sent " <> show comp <> ")") 1 n++-- Messages already received and queued for an ordered callback are thrown+-- away when the connection ends (waitForClient cancels the callback thread).+unit_audit_orderedQueueDroppedOnDisconnect :: Assertion+unit_audit_orderedQueueDroppedOnDisconnect = do+  release <- newTVarIO False+  got <- newTVarIO (0 :: Int)+  b <- broker False False+  c <- runMQTTConduit (conn b) v5{_msgCB=OrderedCallback (\_ _ _ _ -> do+                                     atomically $ readTVar release >>= check+                                     atomically $ modifyTVar' got succ)}+  replicateM_ 3 $ inject b (PublishPkt PublishRequest{_pubDup=False, _pubQoS=QoS0, _pubRetain=False, _pubTopic="a/b",+                                                      _pubPktID=0, _pubBody="x", _pubProps=[]})+  threadDelay 100000+  close b+  _ <- E.try @E.SomeException (waitForClient c)+  atomically $ writeTVar release True+  threadDelay 300000+  n <- readTVarIO got+  assertEqual "QoS 0 messages received before the close that reached the callback" 3 n
test/ClientSpec.hs view
@@ -242,3 +242,21 @@     assertEqual "PUBCOMP before the callback completes" [] beforeRelease     atomically $ writeTVar release True     nthWritten 2 written >>= assertEqual "PUBCOMP" (PubCOMPPkt (PubCOMP 1 0 []))++-- An inbound QoS 2 message is delivered once on PUBLISH; a duplicate+-- PUBLISH before the PUBREL is acked again but not delivered again.+unit_qos2DuplicateDeliveredOnce :: Assertion+unit_qos2DuplicateDeliveredOnce = do+  (written, inject, conn) <- scriptedBroker []+  got <- newTVarIO (0 :: Int)+  c <- runMQTTConduit conn mqttConfig{_protocol=Protocol50,+                                      _msgCB=SimpleCallback (\_ _ _ _ -> atomically $ modifyTVar' got succ)}+  flip E.finally (stopClient c) $ do+    inject (incomingPublish QoS2 5)+    nthWritten 1 written >>= assertEqual "PUBREC" (PubRECPkt (PubREC 5 0 []))+    inject ((\(PublishPkt p) -> PublishPkt p{_pubDup=True}) (incomingPublish QoS2 5))+    nthWritten 2 written >>= assertEqual "second PUBREC" (PubRECPkt (PubREC 5 0 []))+    inject (PubRELPkt (PubREL 5 0 []))+    nthWritten 3 written >>= assertEqual "PUBCOMP" (PubCOMPPkt (PubCOMP 5 0 []))+    threadDelay 100000+    atomically (readTVar got) >>= assertEqual "deliveries" 1
− test/DecayingSpec.hs
@@ -1,50 +0,0 @@-{-# LANGUAGE BlockArguments    #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ViewPatterns      #-}--module DecayingSpec where--import           Control.Concurrent              (threadDelay)-import           Control.Concurrent.STM          (STM, atomically)-import           Control.Monad                   (foldM, mapM_)-import qualified Data.Attoparsec.ByteString.Lazy as A-import qualified Data.ByteString.Lazy            as L-import           Data.Foldable                   (traverse_)-import qualified Data.Map.Strict                 as Map-import qualified Data.Map.Strict.Decaying        as DecayingMap-import           Data.Set                        (Set)-import qualified Data.Set                        as Set--import           Test.QuickCheck--prop_decayingMapWorks :: [Int] -> Property-prop_decayingMapWorks keys = idempotentIOProperty $ do-  m <- DecayingMap.new 60-  atomically $ traverse_ (\x -> DecayingMap.insert x x m) keys-  found <- atomically $ traverse (\x -> DecayingMap.findWithDefault maxBound x m) keys-  pure $ found === keys--prop_decayingMapDecays :: [Int] -> Property-prop_decayingMapDecays keys = idempotentIOProperty $ do-  m <- DecayingMap.new 0.001-  atomically $ traverse_ (\x -> DecayingMap.insert x x m) keys-  threadDelay 5000-  DecayingMap.tick m-  found <- atomically $ DecayingMap.elems m-  pure $ found === []--prop_decayingMapUpdates :: Set Int -> Property-prop_decayingMapUpdates (Set.toList -> keys) = idempotentIOProperty $ do-  m <- DecayingMap.new 60-  atomically $ traverse_ (\x -> DecayingMap.insert x x m) keys-  updated <- atomically $ traverse (\x -> DecayingMap.updateLookupWithKey (\_ v -> Just (v + 1)) x m) keys-  found <- atomically $ traverse (\x -> DecayingMap.findWithDefault maxBound x m) keys-  pure $ (found === fmap (+ 1) keys .&&. Just found === sequenceA updated)--prop_decayingMapDeletes :: Set Int -> Property-prop_decayingMapDeletes (Set.toList -> keys) = (not . null) keys ==> idempotentIOProperty $ do-  m <- DecayingMap.new 60-  atomically $ traverse_ (\x -> DecayingMap.insert x x m) keys-  atomically $ traverse (`DecayingMap.delete` m) (tail keys)-  found <- atomically $ DecayingMap.elems m-  pure $ found === take 1 keys
− test/ExpiringSpec.hs
@@ -1,122 +0,0 @@-{-# LANGUAGE BlockArguments             #-}-{-# LANGUAGE DeriveGeneric              #-}-{-# LANGUAGE DerivingStrategies         #-}-{-# LANGUAGE GeneralizedNewtypeDeriving #-}-{-# LANGUAGE LambdaCase                 #-}-{-# LANGUAGE OverloadedStrings          #-}-{-# LANGUAGE TemplateHaskell            #-}-{-# LANGUAGE TupleSections              #-}--module ExpiringSpec where--import           Control.Lens-import           Data.Bool                (bool)-import           Data.Foldable            (foldl', toList, traverse_)-import           Data.Function            ((&))-import qualified Data.Map.Strict          as Map-import qualified Data.Map.Strict.Expiring as ExpiringMap-import           Data.Set                 (Set)-import qualified Data.Set                 as Set-import           GHC.Generics             (Generic)--import           Test.QuickCheck--data SomeKey = Key1 | Key2 | Key3 | Key4 | Key5-  deriving (Bounded, Enum, Eq, Ord, Show)--instance Arbitrary SomeKey where-  arbitrary = arbitraryBoundedEnum-  shrink = shrinkBoundedEnum--newtype GenOffset = GenOffset { getOffset :: Int }-  deriving (Eq, Ord)-  deriving newtype (Show, Num, Bounded)--instance Arbitrary GenOffset where-  arbitrary = GenOffset <$> choose (0, 5)-  shrink = fmap GenOffset . shrink . getOffset--data Mutation = Insert GenOffset SomeKey Int-              | Delete SomeKey-              | Update GenOffset SomeKey Int-              | UpdateNothing GenOffset SomeKey-              | NewGeneration GenOffset-  deriving (Show, Generic)--makePrisms ''Mutation--instance Arbitrary Mutation where-  arbitrary = oneof [Insert <$> arbitrary <*> arbitrary <*> arbitrary,-                     Delete <$> arbitrary,-                     Update <$> arbitrary <*> arbitrary <*> arbitrary,-                     UpdateNothing <$> arbitrary <*> arbitrary,-                     NewGeneration <$> arbitrary-                     ]-  shrink = genericShrink--allOpTypes :: [String]-allOpTypes = ["Insert", "Delete", "Update", "UpdateNothing", "NewGeneration"]---- Verify that after a series of operations, the map and expiring map return the same values for the given keys.-prop_doesMapStuff :: [Mutation] -> [SomeKey] -> Property-prop_doesMapStuff ops lookups =-  coverTable "mutation types" ((,5) <$> allOpTypes) $ -- The test paths should hit every mutation type (5% min)-  tabulate "mutation types" (takeWhile (/= ' ') . show <$> ops) $ -- We can identify one by the first word in its constructor-  checkCoverage $-  ((`Map.lookup` massocs) <$> lookups) === ((`ExpiringMap.lookup` eassocs) <$> lookups)-  where-    massocs = degenerate $ foldl' applyOpM (0, mempty) ops-    eassocs = foldl' applyOpE (ExpiringMap.new 0) ops--    -- The emulation stores the generation along with the value, so when we're done, we fmap away the generation.-    degenerate :: (GenOffset, Map.Map SomeKey (GenOffset, Int)) -> Map.Map SomeKey Int-    degenerate = fmap snd . snd--    applyOpM (gen, m) = \case-      Insert g k v      -> (gen, Map.insert k (gen+g, v) m)-      Delete k          -> (gen, Map.delete k m)-      Update g k v      -> (gen, snd $ Map.updateLookupWithKey (\_ _ -> Just (gen+g, v)) k m)-      UpdateNothing _ k -> (gen, snd $ Map.updateLookupWithKey (\_ _ -> Nothing) k m)-      NewGeneration n   -> (gen + n, Map.filter ((>= gen + n) . fst) m)--    applyOpE m = \case-      Insert g k v      -> ExpiringMap.insert (gen + g) k v m-      Delete k          -> ExpiringMap.delete k m-      Update g k v      -> snd $ ExpiringMap.updateLookupWithKey (gen + g) (\k' _ -> bool Nothing (Just v) (k == k')) k m-      UpdateNothing g k -> snd $ ExpiringMap.updateLookupWithKey (gen + g) (\_ _ -> Nothing) k m-      NewGeneration g   -> ExpiringMap.newGen (gen + g) m-      where gen = ExpiringMap.generation m--prop_updateReturn :: Int -> Property-prop_updateReturn x = (Just plus2, Just plus2, Nothing, Just plus2) === (up1, ExpiringMap.lookup x m', up2, up3)-    where-        m = ExpiringMap.insert 0 x 0 $ ExpiringMap.new 0-        plus2 = x + 2-        (up1, m') = ExpiringMap.updateLookupWithKey 0 (\_ v -> Just (x + 2)) x m -- New value returns new value-        (up2, m'') = ExpiringMap.updateLookupWithKey 0 (\_ v -> Just (x + 3)) (x + 1) m' -- Missing returns nothing-        up3 = fst $ ExpiringMap.updateLookupWithKey 0 (\_ v -> Nothing) x m'' -- Nothing returns previous value--prop_cannotAcceptExpired :: Positive Int -> Positive Int -> Int -> Property-prop_cannotAcceptExpired (Positive lowGen) (Positive offset) k = ExpiringMap.inspect m === ExpiringMap.inspect m'-  where-    highGen = lowGen + offset-    m = ExpiringMap.new highGen :: ExpiringMap.Map Int Int Int-    m' = ExpiringMap.insert lowGen k k m--prop_cannotUpdateExpired :: Positive Int -> Positive Int -> Int -> Property-prop_cannotUpdateExpired (Positive lowGen) (Positive offset) k = mv === Nothing .&&. ExpiringMap.lookup k m' === Just True-  where-    highGen = lowGen + offset-    m = ExpiringMap.insert highGen k True $ ExpiringMap.new highGen-    (mv, m') = ExpiringMap.updateLookupWithKey lowGen (\_ _ -> Just False) k m--prop_assocs :: [SomeKey] -> Property-prop_assocs keys = ExpiringMap.assocs m === Map.assocs (Map.fromList $ zip keys keys)-  where-    m = foldr (\k -> ExpiringMap.insert 0 k k) (ExpiringMap.new 0) keys--prop_generation :: Int -> Int -> Property-prop_generation g1 g2 = ExpiringMap.inspect m === (0, max g1 g2, 0)-  where-    m :: ExpiringMap.Map Int Int Int-    m = ExpiringMap.newGen g2 $ ExpiringMap.new g1