Merge pull request #110 from AlexeyRaga/sqs-lazy-messages

SQS: Read messages lazily
arbor · Oct 22, 2020 · 86ad3df · 86ad3df
2 parents aaa9728 + 4593c03
commit 86ad3df
Show file tree

Hide file tree

Showing 12 changed files with 49 additions and 13 deletions.
diff --git a/antiope-athena/antiope-athena.cabal b/antiope-athena/antiope-athena.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-athena
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-contract/antiope-contract.cabal b/antiope-contract/antiope-contract.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-contract
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-core/antiope-core.cabal b/antiope-core/antiope-core.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-core
-version:                7.5.2
+version:                7.5.3
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 category:               Services

diff --git a/antiope-dynamodb/antiope-dynamodb.cabal b/antiope-dynamodb/antiope-dynamodb.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-dynamodb
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-es/antiope-es.cabal b/antiope-es/antiope-es.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-es
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-messages/antiope-messages.cabal b/antiope-messages/antiope-messages.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-messages
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-optparse-applicative/antiope-optparse-applicative.cabal b/antiope-optparse-applicative/antiope-optparse-applicative.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-optparse-applicative
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-s3/antiope-s3.cabal b/antiope-s3/antiope-s3.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-s3
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-shell/antiope-shell.cabal b/antiope-shell/antiope-shell.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-shell
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-sns/antiope-sns.cabal b/antiope-sns/antiope-sns.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-sns
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-sqs/antiope-sqs.cabal b/antiope-sqs/antiope-sqs.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:                   antiope-sqs
-version:                7.5.2
+version:                7.5.3
 synopsis:               Please see the README on Github at <https://github.com/arbor/antiope#readme>
 description:            Please see the README on Github at <https://github.com/arbor/antiope#readme>.
 category:               Services

diff --git a/antiope-sqs/src/Antiope/SQS.hs b/antiope-sqs/src/Antiope/SQS.hs
@@ -14,6 +14,8 @@ module Antiope.SQS
   , keepAliveMessage
   , forAllMessages
   , forAllMessages'
+  , lazyReadAllMessages
+  , lazyReadAllMessages'
   , defaultReceiveMessage
 
   -- * Re-exports
@@ -23,7 +25,7 @@ module Antiope.SQS
 
 import Control.Lens
 import Control.Monad            (forM, forM_, join, void)
-import Control.Monad.IO.Unlift  (MonadUnliftIO)
+import Control.Monad.IO.Unlift  (MonadUnliftIO, askUnliftIO, liftIO, unliftIO)
 import Control.Monad.Loops      (unfoldWhileM)
 import Control.Monad.Trans      (lift)
 import Data.Coerce              (coerce)
@@ -37,7 +39,8 @@ import Network.AWS.SQS
 
 import Antiope.SQS.Types
 
-import qualified Network.AWS as AWS
+import qualified Network.AWS      as AWS
+import qualified System.IO.Unsafe as IO
 
 defaultReceiveMessage :: QueueUrl -> ReceiveMessage
 defaultReceiveMessage (QueueUrl url)
@@ -139,6 +142,39 @@ forAllMessages' :: (MonadUnliftIO m, HasEnv env)
 forAllMessages' env recMsg mode process =
   forAllMessages'' env recMsg mode (const process)
 
+-- | Reads all messages from the queue into a LAZY list.
+-- This function will read the provided SQS queue one batch at a time
+-- as the list is consumed.
+--
+-- Use 'lazyReadAllMessages'' to control the batch size.
+lazyReadAllMessages :: (MonadUnliftIO m, HasEnv env)
+  => env
+  -> QueueUrl
+  -> ConsumerMode
+  -> m [Message]
+lazyReadAllMessages env queue mode =
+  lazyReadAllMessages' env (defaultReceiveMessage queue) mode
+
+-- | Reads all messages from the queue into a LAZY list.
+-- This function will read the provided SQS queue one batch at a time as
+-- as the list is consumed.
+lazyReadAllMessages' :: (MonadUnliftIO m, HasEnv env)
+  => env
+  -> ReceiveMessage
+  -> ConsumerMode
+  -> m [Message]
+lazyReadAllMessages' env recMsg mode = do
+  u <- askUnliftIO
+  liftIO $ IO.unsafeInterleaveIO (go u)
+  where
+    go u = do
+      msgs <- unliftIO u (runResourceT $ runAWS env (readQueue' recMsg))
+      case (mode, msgs) of
+        (Drain, []) -> pure []
+        _           -> do
+          rest <- IO.unsafeInterleaveIO (go u)
+          pure (msgs ++ rest)
+
 forAllMessages'' :: (MonadUnliftIO m, HasEnv env)
   => env
   -> ReceiveMessage