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 +2/−5
- src/Data/Map/Strict/Decaying.hs +0/−66
- src/Data/Map/Strict/Expiring.hs +0/−107
- src/Network/MQTT/Client.hs +157/−73
- test/AuditSpec.hs +199/−0
- test/ClientSpec.hs +18/−0
- test/DecayingSpec.hs +0/−50
- test/ExpiringSpec.hs +0/−122
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