Merge pull request #132 from alephcloud/connection-reuse

Connection reuse

Aws/DynamoDb/Core.hs

@@ -114,6 +114,7 @@ import           Data.Aeson
 import qualified Data.Aeson                   as A
 import           Data.Aeson.Types             (Pair, parseEither)
 import qualified Data.Aeson.Types             as A
+import qualified Data.Attoparsec.ByteString   as AttoB (endOfInput)
 import qualified Data.Attoparsec.Text         as Atto
 import           Data.Byteable
 import qualified Data.ByteString.Base16       as Base16
@@ -826,7 +827,7 @@ instance FromJSON AmazonError where
 -------------------------------------------------------------------------------
 ddbResponseConsumer :: A.FromJSON a => IORef DdbResponse -> HTTPResponseConsumer a
 ddbResponseConsumer ref resp = do
-    val <- HTTP.responseBody resp $$+- sinkParser A.json'
+    val <- HTTP.responseBody resp $$+- sinkParser (A.json' <* AttoB.endOfInput)
     case statusCode of
       200 -> rSuccess val
       _   -> rError val

aws.cabal

@@ -269,10 +269,12 @@ test-suite sqs-tests
         base == 4.*,
         bytestring >= 0.10,
         errors >= 1.4.7,
+        http-client >= 0.3,
         lifted-base >= 0.2,
         monad-control >= 0.3,
         mtl >= 2.1,
         quickcheck-instances >= 0.3,
+        resourcet >= 1.1,
         tagged >= 0.7,
         tasty >= 0.8,
         tasty-quickcheck >= 0.8,

tests/DynamoDb/Main.hs

@@ -22,15 +22,20 @@ module Main
 ( main
 ) where
 
+import Aws
 import qualified Aws.DynamoDb as DY
 
 import Control.Arrow (second)
 import Control.Error
 import Control.Monad
+import Control.Monad.IO.Class
 
+import Data.IORef
 import qualified Data.List as L
 import qualified Data.Text as T
 
+import qualified Network.HTTP.Client as HTTP
+
 import Test.Tasty
 import Test.QuickCheck.Instances ()
 
@@ -90,6 +95,7 @@ tests :: TestTree
 tests = testGroup "DynamoDb Tests"
     [ test_table
     -- , test_message
+    , test_core
     ]
 
 -- -------------------------------------------------------------------------- --
@@ -114,3 +120,38 @@ prop_createDescribeDeleteTable readCapacity writeCapacity tableName = do
     handleT (\e -> deleteTable >> left e) $ do
         retryT 6 . void . simpleDyT $ DY.DescribeTable tTableName
         deleteTable
+
+-- -------------------------------------------------------------------------- --
+-- Test core functionality
+
+test_core :: TestTree
+test_core = testGroup "Core Tests"
+        [ eitherTOnceTest0 "connectionReuse" prop_connectionReuse
+        ]
+
+prop_connectionReuse
+    :: EitherT T.Text IO ()
+prop_connectionReuse = do
+    c <- liftIO $ do
+        cfg <- baseConfiguration
+
+        -- counts the number of TCP connections
+        ref <- newIORef (0 :: Int)
+
+        void . HTTP.withManager (managerSettings ref) $ \manager -> runEitherT $
+            handleT (error . T.unpack) . replicateM_ 3 $ do
+                void $ dyT cfg manager DY.ListTables
+                mustFail . dyT cfg manager $ DY.DescribeTable "____"
+
+        readIORef ref
+    unless (c == 1) $
+        left "The TCP connection has not been reused"
+  where
+    managerSettings ref = HTTP.defaultManagerSettings
+        { HTTP.managerRawConnection = do
+            mkConn <- HTTP.managerRawConnection HTTP.defaultManagerSettings
+            return $ \a b c -> do
+                atomicModifyIORef ref $ \i -> (succ i, ())
+                mkConn a b c
+        }
+

tests/DynamoDb/Utils.hs

@@ -31,6 +31,7 @@ module DynamoDb.Utils
 -- * DynamoDb Utils
 , simpleDy
 , simpleDyT
+, dyT
 , withTable
 , withTable_
 , createTestTable
@@ -45,11 +46,14 @@ import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.Trans.Control
+import Control.Monad.Trans.Resource
 
 import Data.Monoid
 import qualified Data.Text as T
 import qualified Data.Text.IO as T
 
+import qualified Network.HTTP.Client as HTTP
+
 import Test.Tasty
 import Test.QuickCheck.Instances ()
 
@@ -95,6 +99,16 @@ simpleDyT
     -> EitherT T.Text m (MemoryResponse a)
 simpleDyT = tryT . simpleDy
 
+dyT
+    :: (Transaction r a, ServiceConfiguration r ~ DY.DdbConfiguration)
+    => Configuration
+    -> HTTP.Manager
+    -> r
+    -> EitherT T.Text IO a
+dyT cfg manager req = do
+    Response _ r <- liftIO . runResourceT $ aws cfg dyConfiguration manager req
+    hoistEither $ fmapL sshow r
+
 withTable
     :: T.Text -- ^ table Name
     -> Int -- ^ read capacity (#(non-consistent) reads * itemsize/4KB)

tests/Sqs/Main.hs

@@ -31,11 +31,15 @@ import Control.Error
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.Trans.Control
+import Control.Monad.Trans.Resource
 
+import Data.IORef
 import qualified Data.List as L
 import Data.Monoid
 import qualified Data.Text as T
 
+import qualified Network.HTTP.Client as HTTP
+
 import Test.Tasty
 import Test.QuickCheck.Instances ()
 
@@ -91,9 +95,10 @@ help = L.intercalate "\n"
     ]
 
 tests :: TestTree
-tests = testGroup "SQS Tests"
+tests = withQueueTest defaultQueueName $ \getQueueParams -> testGroup "SQS Tests"
     [ test_queue
-    , test_message
+    , test_message getQueueParams
+    , test_core getQueueParams
     ]
 
 -- -------------------------------------------------------------------------- --
@@ -131,6 +136,16 @@ sqsConfiguration = SQS.SqsConfiguration
     , SQS.sqsDefaultExpiry = 180
     }
 
+sqsT
+    :: (Transaction r a, ServiceConfiguration r ~ SQS.SqsConfiguration)
+    => Configuration
+    -> HTTP.Manager
+    -> r
+    -> EitherT T.Text IO a
+sqsT cfg manager req = do
+    Response _ r <- liftIO . runResourceT $ aws cfg sqsConfiguration manager req
+    hoistEither $ fmapL sshow r
+
 simpleSqs
     :: (AsMemoryResponse a, Transaction r a, ServiceConfiguration r ~ SQS.SqsConfiguration, MonadIO m)
     => r
@@ -188,19 +203,18 @@ prop_createListDeleteQueue queueName = do
 -- -------------------------------------------------------------------------- --
 -- Message Tests
 
-test_message :: TestTree
-test_message =
-    withQueueTest defaultQueueName $ \getQueueParams -> testGroup "Queue Tests"
-        [ eitherTOnceTest0 "SendReceiveDeleteMessage" $ do
-            (_, queue) <- liftIO getQueueParams
-            prop_sendReceiveDeleteMessage queue
-        , eitherTOnceTest0 "SendReceiveDeleteMessageLongPolling" $ do
-            (_, queue) <- liftIO getQueueParams
-            prop_sendReceiveDeleteMessageLongPolling queue
-        , eitherTOnceTest0 "SendReceiveDeleteMessageLongPolling1" $ do
-            (_, queue) <- liftIO getQueueParams
-            prop_sendReceiveDeleteMessageLongPolling1 queue
-        ]
+test_message :: IO (T.Text, SQS.QueueName) -> TestTree
+test_message getQueueParams = testGroup "Queue Tests"
+    [ eitherTOnceTest0 "SendReceiveDeleteMessage" $ do
+        (_, queue) <- liftIO getQueueParams
+        prop_sendReceiveDeleteMessage queue
+    , eitherTOnceTest0 "SendReceiveDeleteMessageLongPolling" $ do
+        (_, queue) <- liftIO getQueueParams
+        prop_sendReceiveDeleteMessageLongPolling queue
+    , eitherTOnceTest0 "SendReceiveDeleteMessageLongPolling1" $ do
+        (_, queue) <- liftIO getQueueParams
+        prop_sendReceiveDeleteMessageLongPolling1 queue
+    ]
 
 -- | Simple send and short-polling receive. First sends all messages
 -- and receives messages thereafter one by one.
@@ -316,3 +330,50 @@ prop_sendReceiveDeleteMessageLongPolling1 queue = do
     unless (sent == recv)
         $ left $ "received messages don't match send messages; sent: "
             <> sshow sent <> "; got: " <> sshow recv
+
+
+-- -------------------------------------------------------------------------- --
+-- Test core functionality
+
+test_core :: IO (T.Text, SQS.QueueName) -> TestTree
+test_core getQueueParams = testGroup "Core Tests"
+    [ eitherTOnceTest0 "connectionReuse" $ do
+        (_, queue) <- liftIO getQueueParams
+        prop_connectionReuse queue
+    ]
+
+prop_connectionReuse
+    :: SQS.QueueName
+    -> EitherT T.Text IO ()
+prop_connectionReuse queue = do
+    c <- liftIO $ do
+        cfg <- baseConfiguration
+
+        -- used for counting the number of TCP connections
+        ref <- newIORef (0 :: Int)
+
+        -- Use a single manager for all HTTP requests
+        void . HTTP.withManager (managerSettings ref) $ \manager -> runEitherT $
+
+            handleT (error . T.unpack) . replicateM_ 3 $ do
+                void . sqsT cfg manager $ SQS.ListQueues Nothing
+                mustFail . sqsT cfg manager $
+                    SQS.SendMessage "" (SQS.QueueName "" "") [] Nothing
+                void . sqsT cfg manager $
+                    SQS.SendMessage "test-message" queue [] Nothing
+                void . sqsT cfg manager $
+                    SQS.ReceiveMessage Nothing [] Nothing [] queue (Just 20)
+
+        readIORef ref
+    unless (c == 1) $
+        left "The TCP connection has not been reused"
+  where
+
+    managerSettings ref = HTTP.defaultManagerSettings
+        { HTTP.managerRawConnection = do
+            mkConn <- HTTP.managerRawConnection HTTP.defaultManagerSettings
+            return $ \a b c -> do
+                atomicModifyIORef ref $ \i -> (succ i, ())
+                mkConn a b c
+        }
+

tests/Utils.hs

@@ -21,6 +21,7 @@ module Utils
 
 -- * General Utils
 , sshow
+, mustFail
 , tryT
 , retryT
 , retryT_
@@ -131,6 +132,11 @@ retryT_ n f = go 1
 sshow :: (Show a, IsString b) => a -> b
 sshow = fromString . show
 
+mustFail :: Monad m => EitherT e m a -> EitherT T.Text m ()
+mustFail = EitherT . eitherT
+    (const . return $ Right ())
+    (const . return $ Left "operation succeeded when a failure was expected")
+
 evalTestTM
     :: Functor f
     => String -- ^ test name