diff --git a/net-mqtt.cabal b/net-mqtt.cabal --- a/net-mqtt.cabal +++ b/net-mqtt.cabal @@ -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 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: diff --git a/src/Data/Map/Strict/Decaying.hs b/src/Data/Map/Strict/Decaying.hs deleted file mode 100644 --- a/src/Data/Map/Strict/Decaying.hs +++ /dev/null @@ -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) diff --git a/src/Data/Map/Strict/Expiring.hs b/src/Data/Map/Strict/Expiring.hs deleted file mode 100644 --- a/src/Data/Map/Strict/Expiring.hs +++ /dev/null @@ -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) diff --git a/src/Network/MQTT/Client.hs b/src/Network/MQTT/Client.hs --- a/src/Network/MQTT/Client.hs +++ b/src/Network/MQTT/Client.hs @@ -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) diff --git a/test/AuditSpec.hs b/test/AuditSpec.hs new file mode 100644 --- /dev/null +++ b/test/AuditSpec.hs @@ -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 diff --git a/test/ClientSpec.hs b/test/ClientSpec.hs --- a/test/ClientSpec.hs +++ b/test/ClientSpec.hs @@ -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 diff --git a/test/DecayingSpec.hs b/test/DecayingSpec.hs deleted file mode 100644 --- a/test/DecayingSpec.hs +++ /dev/null @@ -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 diff --git a/test/ExpiringSpec.hs b/test/ExpiringSpec.hs deleted file mode 100644 --- a/test/ExpiringSpec.hs +++ /dev/null @@ -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