diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/aws-kinesis-client.cabal b/aws-kinesis-client.cabal new file mode 100644 --- /dev/null +++ b/aws-kinesis-client.cabal @@ -0,0 +1,81 @@ +name: aws-kinesis-client +version: 0.1.0.0 +synopsis: A producer & consumer client library for AWS Kinesis +-- description: +license: Apache-2.0 +license-file: LICENSE +author: Jon Sterling +maintainer: jon@jonmsterling.com +copyright: Copyright (c) 2013-2014 PivotCloud, Inc. +category: Cloud +build-type: Simple +cabal-version: >=1.10 + +library + exposed-modules: Aws.Kinesis.Client.Common, + Aws.Kinesis.Client.Consumer + Aws.Kinesis.Client.Producer + -- other-modules: + -- other-extensions: + build-depends: base >=4.7 && <4.8, + base-unicode-symbols, + aws >=0.10.5, + aws-kinesis >=0.1.2, + conduit >=1.2.3.1, + data-carousel >=0.1.0.0, + either >=4.3.2.1, + errors >=1.4.7, + http-conduit >=2.1.5, + kan-extensions >=4.2, + lens >=4.6 && < 4.7, + lifted-async >=0.3.0, + lifted-base >=0.2.3.3, + monad-control >=1.0, + mtl >=2.2.1, + random >=1.1, + resourcet >=1.1.3.3, + stm >=2.4.4, + stm-chans >=3.0.0.2, + stm-conduit >=2.5.3, + stm-queue-extras >=0.2.0.0, + text >=1.2.0.3, + transformers >=0.3.0.0 + hs-source-dirs: src + default-language: Haskell2010 + +executable kinesis-cli + default-language: Haskell2010 + hs-source-dirs: cli + main-is: CLI.hs + other-modules: CLI.Options, + CLI.Record + build-depends: base >=4.7 && <4.8, + base-unicode-symbols, + aws >=0.10.5, + aws-kinesis >=0.1.2, + aws-kinesis-client, + conduit >=1.2.3.1, + data-carousel >=0.1.0.0, + either >=4.3.2.1, + errors >=1.4.7, + http-conduit >=2.1.5, + kan-extensions >=4.2, + hourglass >=0.2.6, + lens >=4.6 && < 4.7, + lifted-async >=0.3.0, + lifted-base >=0.2.3.3, + monad-control >=1.0, + mtl >=2.2.1, + random >=1.1, + resourcet >=1.1.3.3, + stm >=2.4.4, + stm-chans >=3.0.0.2, + stm-conduit >=2.5.3, + stm-queue-extras >=0.2.0.0, + text >=1.2.0.3, + transformers >=0.3.0.0, + aeson >=0.8, + aws-general >=0.1.1, + bytestring >=0.10.4, + optparse-applicative >=0.11, + time >=1.4.2 diff --git a/cli/CLI.hs b/cli/CLI.hs new file mode 100644 --- /dev/null +++ b/cli/CLI.hs @@ -0,0 +1,148 @@ +-- Copyright (c) 2013-2014 PivotCloud, Inc. +-- +-- Aws.Kinesis.Client.Consumer +-- +-- Please feel free to contact us at licensing@pivotmail.com with any +-- contributions, additions, or other feedback; we would love to hear from +-- you. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); you may +-- not use this file except in compliance with the License. You may obtain a +-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +-- License for the specific language governing permissions and limitations +-- under the License. + +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE UnicodeSyntax #-} + +-- | +-- Module: Main +-- Copyright: Copyright © 2013-2014 PivotCloud, Inc. +-- License: Apache-2.0 +-- Maintainer: Jon Sterling +-- Stability: experimental +-- + +module Main where + +import Aws.Aws +import Aws.General +import Aws.Kinesis hiding (Record) +import Aws.Kinesis.Client.Common +import Aws.Kinesis.Client.Consumer + +import CLI.Options +import CLI.Record + +import Control.Applicative +import Control.Monad +import Control.Monad.Trans +import Control.Monad.Trans.Control +import Control.Monad.Trans.Either +import Control.Monad.Trans.Resource +import Control.Monad.Codensity +import Control.Monad.Reader.Class +import Control.Monad.Trans.Reader (ReaderT(..)) +import Control.Monad.Error.Class + +import Data.Aeson +import qualified Data.ByteString.Lazy as BL +import Data.Conduit +import Data.Monoid +import qualified Data.Conduit.List as CL +import qualified Data.Text.Lens as T + +import Options.Applicative +import qualified Network.HTTP.Conduit as HC +import Prelude.Unicode +import Control.Monad.Unicode + +type MonadCLI m + = ( MonadReader CLIOptions m + , MonadIO m + , MonadBaseControl IO m + , MonadError ConsumerError m + ) + +identityConduit + ∷ Monad m + ⇒ Conduit a m a +identityConduit = CL.map id + +limitConduit + ∷ MonadCLI m + ⇒ Conduit a m a +limitConduit = + lift (asks clioLimit) ≫= + CL.isolate + +filterConduit + ∷ MonadCLI m + ⇒ Conduit Record m Record +filterConduit = do + CLIOptions{..} ← lift ask + CL.filter (\Record{..} → maybe True (rTimestamp ≥) clioStartDate) + =$ takeTill (\Record{..} → maybe False (rTimestamp >) clioEndDate) + +takeTill + ∷ Monad m + ⇒ (i → Bool) + → Conduit i m i +takeTill f = loop + where + loop = await ≫= maybe (return ()) (\x → unless (f x) $ yield x ≫ loop) + +app + ∷ MonadCLI m + ⇒ Codensity m () +app = do + CLIOptions{..} ← ask + manager ← managedHttpManager + awsConfiguration ← liftIO baseConfiguration + consumer ← managedKinesisConsumer $ ConsumerKit + { _ckKinesisKit = KinesisKit + { _kkManager = manager + , _kkConfiguration = awsConfiguration + , _kkKinesisConfiguration = KinesisConfiguration UsWest2 + } + , _ckStreamName = clioStreamName + , _ckBatchSize = 100 + , _ckIteratorType = clioIteratorType + } + + lift $ consumerSource consumer $$ + case clioRaw of + False → + CL.mapMaybe (decode ∘ BL.fromChunks ∘ (:[]) ∘ recordData) + =$ filterConduit + =$ limitConduit + =$ CL.mapM_ (liftIO ∘ print) + True → + limitConduit + =$ CL.mapM_ (liftIO ∘ print ∘ recordData) + return () + +main ∷ IO () +main = + eitherT (fail ∘ show) return $ + liftIO (execParser parserInfo) + ≫= runReaderT (lowerCodensity app) + +managedHttpManager + ∷ ( MonadIO m + , MonadBaseControl IO m + ) + ⇒ Codensity m HC.Manager +managedHttpManager = + Codensity $ HC.withManager ∘ (lift ∘) diff --git a/cli/CLI/Options.hs b/cli/CLI/Options.hs new file mode 100644 --- /dev/null +++ b/cli/CLI/Options.hs @@ -0,0 +1,143 @@ +-- Copyright (c) 2013-2014 PivotCloud, Inc. +-- +-- Aws.Kinesis.Client.Consumer +-- +-- Please feel free to contact us at licensing@pivotmail.com with any +-- contributions, additions, or other feedback; we would love to hear from +-- you. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); you may +-- not use this file except in compliance with the License. You may obtain a +-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +-- License for the specific language governing permissions and limitations +-- under the License. + +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE UnicodeSyntax #-} + +-- | +-- Module: CLI.Options +-- Copyright: Copyright © 2013-2014 PivotCloud, Inc. +-- License: Apache-2.0 +-- Maintainer: Jon Sterling +-- Stability: experimental +-- + +module CLI.Options +( CLIOptions(..) +, optionsParser +, parserInfo +) where + +import Aws.Kinesis +import Control.Applicative.Unicode +import Control.Lens +import Control.Monad.Unicode +import Data.Monoid.Unicode +import qualified Data.Text as T +import qualified Data.Text.Lens as T +import Data.Hourglass +import Options.Applicative +import Prelude.Unicode + +data CLIOptions + = CLIOptions + { clioStreamName ∷ StreamName + , clioLimit ∷ Int + , clioIteratorType ∷ ShardIteratorType + , clioStartDate ∷ Maybe DateTime + , clioEndDate ∷ Maybe DateTime + -- , clioTimeDuration ∷ Maybe Seconds + , clioRaw ∷ Bool + } deriving Show + +eitherTextReader + ∷ ( T.IsText i + , T.IsText e + ) + ⇒ (i → Either e a) + → ReadM a +eitherTextReader f = + eitherReader $ + (_Left %~ view T.unpacked) ∘ f ∘ view T.packed + +streamNameParser ∷ Parser StreamName +streamNameParser = + option (eitherTextReader streamName) $ + long "stream-name" + ⊕ short 's' + ⊕ metavar "SN" + ⊕ help "Fetch from the Kinesis stream named `SN`" + +limitParser ∷ Parser Int +limitParser = + option auto $ + long "limit" + ⊕ short 'l' + ⊕ metavar "L" + ⊕ help "Fetch `L` records" + +iteratorTypeParser ∷ Parser ShardIteratorType +iteratorTypeParser = + option auto $ + long "iterator-type" + ⊕ short 'i' + ⊕ metavar "IT" + ⊕ help "Iterator type (Latest|TrimHorizon)" + ⊕ value TrimHorizon + ⊕ showDefault + +readDateTime + ∷ String + → ReadM DateTime +readDateTime = + maybe (readerError "Invalid DateTime") return + ∘ timeParse ISO8601_DateAndTime + +startDateParser ∷ Parser DateTime +startDateParser = + option (str ≫= readDateTime) $ + long "start-date" + ⊕ short 's' + ⊕ metavar "SD" + ⊕ help "Start Date (ISO 8601)" + +endDateParser ∷ Parser DateTime +endDateParser = + option (str ≫= readDateTime) $ + long "end-date" + ⊕ short 'e' + ⊕ metavar "ED" + ⊕ help "End Date (ISO 8601)" + +timeDurationParser ∷ Parser Seconds +timeDurationParser = + option auto $ + long "end-date" + ⊕ short 'e' + ⊕ metavar "ED" + ⊕ help "Time window from start (in seconds)" + +optionsParser ∷ Parser CLIOptions +optionsParser = + CLIOptions + <$> streamNameParser + ⊛ limitParser + ⊛ iteratorTypeParser + ⊛ optional startDateParser + ⊛ optional endDateParser + -- ⊛ optional timeDurationParser + ⊛ switch (long "raw" ⊕ help "Treat records as raw text") + +parserInfo ∷ ParserInfo CLIOptions +parserInfo = + info (helper ⊛ optionsParser) $ + fullDesc + ⊕ progDesc "Fetch `L` records from a Kinesis stream `SN`. Put your AWS keys in ~/.aws-keys" + ⊕ header "The Kinesis Consumer CLI" + diff --git a/cli/CLI/Record.hs b/cli/CLI/Record.hs new file mode 100644 --- /dev/null +++ b/cli/CLI/Record.hs @@ -0,0 +1,79 @@ +-- Copyright (c) 2013-2014 PivotCloud, Inc. +-- +-- Aws.Kinesis.Client.Consumer +-- +-- Please feel free to contact us at licensing@pivotmail.com with any +-- contributions, additions, or other feedback; we would love to hear from +-- you. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); you may +-- not use this file except in compliance with the License. You may obtain a +-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +-- License for the specific language governing permissions and limitations +-- under the License. + +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE UnicodeSyntax #-} + +-- | +-- Module: CLI.Record +-- Copyright: Copyright © 2013-2014 PivotCloud, Inc. +-- License: Apache-2.0 +-- Maintainer: Jon Sterling +-- Stability: experimental +-- + +module CLI.Record where + +import Data.Aeson +import Data.Aeson.Types +import Data.Time +import qualified Data.Text as T +import qualified Data.ByteString as BS +import Control.Applicative +import Control.Applicative.Unicode +import Data.Hourglass +import Prelude.Unicode + +data Record + = Record + { rTimestamp ∷ DateTime + , rMessage ∷ T.Text + } deriving (Show, Eq) + +instance FromJSON DateTime where + parseJSON = + withText "DateTime" $ + maybe (fail "Invalid DateTime") return + ∘ timeParse fmt + ∘ T.unpack + where + fmt = + [ Format_Year4 + , hyphen + , Format_Month2 + , hyphen + , Format_Day2 + , Format_Text 'T' + , Format_Hour + , colon + , Format_Minute + , colon + , Format_Second + ] + hyphen = Format_Text '-' + colon = Format_Text ':' + +instance FromJSON Record where + parseJSON = + withObject "Record" $ \xs → + Record <$> xs .: "time" ⊛ xs .: "message" + +instance Ord Record where + compare r1 r2 = compare (rTimestamp r1) (rTimestamp r2) + diff --git a/src/Aws/Kinesis/Client/Common.hs b/src/Aws/Kinesis/Client/Common.hs new file mode 100644 --- /dev/null +++ b/src/Aws/Kinesis/Client/Common.hs @@ -0,0 +1,196 @@ +-- Copyright (c) 2013-2014 PivotCloud, Inc. +-- +-- Aws.Kinesis.Client.Common +-- +-- Please feel free to contact us at licensing@pivotmail.com with any +-- contributions, additions, or other feedback; we would love to hear from +-- you. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); you may +-- not use this file except in compliance with the License. You may obtain a +-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +-- License for the specific language governing permissions and limitations +-- under the License. + +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UnicodeSyntax #-} + +-- | +-- Module: Aws.Kinesis.Client.Common +-- Copyright: Copyright © 2013-2014 PivotCloud, Inc. +-- License: Apache-2.0 +-- Maintainer: Jon Sterling +-- Stability: experimental +-- +module Aws.Kinesis.Client.Common +( KinesisKit(..) +, kkConfiguration +, kkKinesisConfiguration +, kkManager +, MonadKinesis +, runKinesis + + -- * Fetching Shards +, streamShardSource +, streamOpenShardSource +, shardIsOpen + + -- * Miscellaneous monad stuff +, mapError +, handleError +, mapEnvironment +) where + +import qualified Aws +import qualified Aws.Core as Aws +import qualified Aws.Kinesis as Kin +import Control.Exception +import Control.Lens +import Control.Error +import Control.Monad +import Control.Monad.Reader.Class +import Control.Monad.Error.Class +import Control.Monad.Trans +import Control.Monad.Trans.Reader +import Control.Monad.Trans.Resource +import Control.Monad.Unicode +import Data.Conduit +import qualified Data.Conduit.List as CL +import qualified Network.HTTP.Conduit as HC +import Prelude.Unicode + +-- | The 'KinesisKit' contains what is necessary to make a request to Kinesis. +-- +data KinesisKit + = KinesisKit + { _kkConfiguration ∷ !Aws.Configuration + , _kkKinesisConfiguration ∷ !(Kin.KinesisConfiguration Aws.NormalQuery) + , _kkManager ∷ !HC.Manager + } + +-- | A lens for '_kkConfiguration'. +-- +kkConfiguration ∷ Lens' KinesisKit Aws.Configuration +kkConfiguration = lens _kkConfiguration $ \kk cfg → kk { _kkConfiguration = cfg } + +-- | A lens for '_kkKinesisConfiguration'. +-- +kkKinesisConfiguration ∷ Lens' KinesisKit (Kin.KinesisConfiguration Aws.NormalQuery) +kkKinesisConfiguration = lens _kkKinesisConfiguration $ \kk cfg → kk { _kkKinesisConfiguration = cfg } + +-- | A lens for '_kkManager'. +-- +kkManager ∷ Lens' KinesisKit HC.Manager +kkManager = lens _kkManager $ \kk mgr → kk { _kkManager = mgr } + +-- | The minimal effect modality for running Kinesis commands. +-- +type MonadKinesis m + = ( MonadIO m + , MonadReader KinesisKit m + , MonadError SomeException m + ) + +-- | Run a Kinesis request inside 'MonadKinesis'. +-- +runKinesis + ∷ ( MonadKinesis m + , Aws.ServiceConfiguration req ~ Kin.KinesisConfiguration + , Aws.Transaction req resp + ) + ⇒ req + → m resp +runKinesis req = do + KinesisKit{..} ← view id + eitherT throwError return ∘ syncIO ∘ runResourceT $ + Aws.pureAws + _kkConfiguration + _kkKinesisConfiguration + _kkManager + req + +shardIsOpen + ∷ Kin.Shard + → Bool +shardIsOpen Kin.Shard{..} = + isNothing $ shardSequenceNumberRange ^. _2 + +fetchShardsConduit + ∷ MonadKinesis m + ⇒ Kin.StreamName + → Conduit (Maybe Kin.ShardId) m Kin.Shard +fetchShardsConduit streamName = + awaitForever $ \mshardId → do + let req = Kin.DescribeStream + { Kin.describeStreamExclusiveStartShardId = mshardId + , Kin.describeStreamLimit = Nothing + , Kin.describeStreamStreamName = streamName + } + resp@(Kin.DescribeStreamResponse Kin.StreamDescription{..}) ← + lift $ runKinesis req + yield `mapM_` streamDescriptionShards + void ∘ traverse (leftover ∘ Just) $ + Kin.describeStreamExclusiveStartShardId =<< + Aws.nextIteratedRequest req resp + return () + + +-- | A 'Source' of shards for a stream. +-- +streamShardSource + ∷ MonadKinesis m + ⇒ Kin.StreamName + → Source m Kin.Shard +streamShardSource streamName = + CL.sourceList [Nothing] $= fetchShardsConduit streamName + +-- | A 'Source' of open shards for a stream. +-- +streamOpenShardSource + ∷ MonadKinesis m + ⇒ Kin.StreamName + → Source m Kin.Shard +streamOpenShardSource streamName = + flip mapOutputMaybe (streamShardSource streamName) $ \shard → + if shardIsOpen shard + then Just shard + else Nothing + +-- | This function may be used to transform a computation with errors in one +-- type to one with errors in another type. Whilst the argument is fixed at +-- 'EitherT', this may be used where the argument is in an /arbitrary/ monad, +-- since it will simply get instantiated at 'EitherT'. +-- +mapError + ∷ MonadError e' m + ⇒ (e → e') + → EitherT e m a + → m a +mapError e = eitherT (throwError ∘ e) return + +-- | This is a verion of 'catchError' with its arguments flipped. +-- +handleError + ∷ MonadError e m + ⇒ (e → m α) + → m α + → m α +handleError = flip catchError + +-- | Analogous to 'withReader', but supports a result in 'MonadReader'. +-- +mapEnvironment + ∷ MonadReader r' m + ⇒ Getter r' r + → ReaderT r m a + → m a +mapEnvironment l m = view l ≫= runReaderT m + diff --git a/src/Aws/Kinesis/Client/Consumer.hs b/src/Aws/Kinesis/Client/Consumer.hs new file mode 100644 --- /dev/null +++ b/src/Aws/Kinesis/Client/Consumer.hs @@ -0,0 +1,303 @@ +-- Copyright (c) 2013-2014 PivotCloud, Inc. +-- +-- Aws.Kinesis.Client.Consumer +-- +-- Please feel free to contact us at licensing@pivotmail.com with any +-- contributions, additions, or other feedback; we would love to hear from +-- you. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); you may +-- not use this file except in compliance with the License. You may obtain a +-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +-- License for the specific language governing permissions and limitations +-- under the License. + +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UnicodeSyntax #-} + +-- | +-- Module: Aws.Kinesis.Client.Consumer +-- Copyright: Copyright © 2013-2014 PivotCloud, Inc. +-- License: Apache-2.0 +-- Maintainer: Jon Sterling +-- Stability: experimental +-- +module Aws.Kinesis.Client.Consumer +( -- * The Consumer + KinesisConsumer +, managedKinesisConsumer +, withKinesisConsumer + + -- * Commands +, consumerSource +, readConsumer +, tryReadConsumer + + -- * Consumer Environment +, ConsumerKit(..) +, ckKinesisKit +, ckStreamName +, ckBatchSize +, ConsumerError(..) +, MonadConsumer +) where + +import qualified Aws.Kinesis as Kin +import Aws.Kinesis.Client.Common + +import Control.Concurrent.Async.Lifted +import Control.Concurrent.Lifted hiding (yield) +import Control.Concurrent.STM +import Control.Concurrent.STM.Queue +import Control.Exception +import Control.Lens +import Control.Monad.Codensity +import Control.Monad.Error.Class +import Control.Monad.Reader +import Control.Monad.Trans.Control +import Control.Monad.Unicode +import qualified Data.Carousel as CR +import Data.Conduit +import qualified Data.Conduit.List as CL +import Prelude.Unicode + +-- | The internal representation for shards used by the consumer. +-- +data ShardState + = ShardState + { _ssIterator ∷ !(TVar (Maybe Kin.ShardIterator)) + , _ssShardId ∷ !Kin.ShardId + } + +-- | A lens for '_ssIterator'. +-- +ssIterator ∷ Getter ShardState (TVar (Maybe Kin.ShardIterator)) +ssIterator = to _ssIterator + +-- | A lens for '_ssShardId'. +-- +ssShardId ∷ Lens' ShardState Kin.ShardId +ssShardId = lens _ssShardId $ \ss sid → ss { _ssShardId = sid } + +-- | 'ShardState' is quotiented by shard ID. +-- +instance Eq ShardState where + ss == ss' = ss ^. ssShardId ≡ ss' ^. ssShardId + +data ConsumerError + = NoShards + -- ^ A stream must always have at least one open shard. + + | KinesisError !SomeException + -- ^ An error which derives from a request made to Kinesis. + + deriving Show + +-- | The 'ConsumerKit' contains what is needed to initialize a 'KinesisConsumer'. +data ConsumerKit + = ConsumerKit + { _ckKinesisKit ∷ !KinesisKit + -- ^ The credentials and configuration for making requests to AWS Kinesis. + + , _ckStreamName ∷ !Kin.StreamName + -- ^ The name of the stream to consume from. + + , _ckBatchSize ∷ {-# UNPACK #-} !Int + -- ^ The number of records to fetch at once from the stream. + + , _ckIteratorType ∷ !Kin.ShardIteratorType + -- ^ The type of iterator to consume. + } + +-- | A lens for '_ckKinesisKit'. +-- +ckKinesisKit ∷ Lens' ConsumerKit KinesisKit +ckKinesisKit = lens _ckKinesisKit $ \ck kk → ck { _ckKinesisKit = kk } + +-- | A lens for '_ckStreamName'. +-- +ckStreamName ∷ Lens' ConsumerKit Kin.StreamName +ckStreamName = lens _ckStreamName $ \ck sn → ck { _ckStreamName = sn } + +-- | A lens for '_ckBatchSize'. +-- +ckBatchSize ∷ Lens' ConsumerKit Int +ckBatchSize = lens _ckBatchSize $ \ck bs → ck { _ckBatchSize = bs } + +-- | A lens for '_ckIteratorType'. +-- +ckIteratorType ∷ Lens' ConsumerKit Kin.ShardIteratorType +ckIteratorType = lens _ckIteratorType $ \ck it → ck { _ckIteratorType = it } + +-- | The 'KinesisConsumer' maintains state about which shards to pull from. +-- +newtype KinesisConsumer = KinesisConsumer { _kcMessageQueue ∷ TBQueue Kin.Record } + +-- | A getter for '_kcMessageQueue'. +-- +kcMessageQueue ∷ Getter KinesisConsumer (TBQueue Kin.Record) +kcMessageQueue = to _kcMessageQueue + +-- | The basic effect modality required for operating the consumer. +-- +type MonadConsumer m + = ( MonadIO m + , MonadBaseControl IO m + , MonadError ConsumerError m + ) + +type MonadConsumerInternal m + = ( MonadConsumer m + , MonadReader ConsumerKit m + ) + +-- | This constructs a 'KinesisConsumer' and closes it when you have done with +-- it; this is equivalent to 'withKinesisConsumer', except that the +-- continuation is replaced with returning the consumer in 'Codensity'. +-- +managedKinesisConsumer + ∷ MonadConsumer m + ⇒ ConsumerKit + → Codensity m KinesisConsumer +managedKinesisConsumer kit = + Codensity $ withKinesisConsumer kit + +-- | This constructs a 'KinesisConsumer' and closes it when you have done with +-- it. +-- +withKinesisConsumer + ∷ MonadConsumer m + ⇒ ConsumerKit + → (KinesisConsumer → m α) + → m α +withKinesisConsumer kit inner = + flip runReaderT kit $ do + batchSize ← view ckBatchSize + messageQueue ← liftIO ∘ newTBQueueIO $ batchSize * 10 + + state ← updateStreamState CR.empty ≫= liftIO ∘ newTVarIO + let reshardingLoop = forever $ + handleError (\_ → liftIO $ threadDelay 3000000) $ do + liftIO (readTVarIO state) + ≫= updateStreamState + ≫= liftIO ∘ atomically ∘ writeTVar state + liftIO $ threadDelay 10000000 + + producerLoop = forever $ + handleError (\_ → liftIO $ threadDelay 2000000) $ do + recordsCount ← replenishMessages messageQueue state + when (recordsCount ≡ 0) $ + liftIO $ threadDelay 5000000 + + + withAsync reshardingLoop $ \reshardingHandle → do + link reshardingHandle + withAsync producerLoop $ \producerHandle → do + link producerHandle + res ← lift ∘ inner $ KinesisConsumer messageQueue + return res + +-- | This requests new information from Kinesis and reconciles that with an +-- existing carousel of shard states. +-- +updateStreamState + ∷ MonadConsumerInternal m + ⇒ CR.Carousel ShardState + → m (CR.Carousel ShardState) +updateStreamState state = do + streamName ← view ckStreamName + iteratorType ← view ckIteratorType + + mapError KinesisError ∘ mapEnvironment ckKinesisKit $ do + let existingShardIds = state ^. CR.clList <&> view ssShardId + shardSource = flip mapOutputMaybe (streamOpenShardSource streamName) $ \sh@Kin.Shard{..} → + if shardShardId `elem` existingShardIds + then Nothing + else Just sh + + newShards ← shardSource $$ CL.consume + shardStates ← forM newShards $ \Kin.Shard{..} → do + Kin.GetShardIteratorResponse it ← runKinesis Kin.GetShardIterator + { Kin.getShardIteratorShardId = shardShardId + , Kin.getShardIteratorShardIteratorType = iteratorType + , Kin.getShardIteratorStartingSequenceNumber = Nothing + , Kin.getShardIteratorStreamName = streamName + } + iteratorVar ← liftIO ∘ newTVarIO $ Just it + return ShardState + { _ssIterator = iteratorVar + , _ssShardId = shardShardId + } + return ∘ CR.nub $ CR.append shardStates state + +-- | Waits for a message queue to be emptied and fills it up again. +-- +replenishMessages + ∷ MonadConsumerInternal m + ⇒ TBQueue Kin.Record + → TVar (CR.Carousel ShardState) + → m Int +replenishMessages messageQueue shardsVar = do + bufferSize ← view ckBatchSize + liftIO ∘ atomically ∘ awaitQueueEmpty $ messageQueue + (shard, iterator) ← liftIO ∘ atomically $ do + mshard ← shardsVar ^!? act readTVar ∘ CR.cursor + shard ← maybe retry return mshard + miterator ← shard ^! ssIterator ∘ act readTVar + iterator ← maybe retry return miterator + return (shard, iterator) + + Kin.GetRecordsResponse{..} ← mapError KinesisError ∘ mapEnvironment ckKinesisKit $ runKinesis Kin.GetRecords + { getRecordsLimit = Just bufferSize + , getRecordsShardIterator = iterator + } + + liftIO ∘ atomically $ do + writeTVar (shard ^. ssIterator) getRecordsResNextShardIterator + forM_ getRecordsResRecords $ writeTBQueue messageQueue + modifyTVar shardsVar CR.moveRight + + return $ length getRecordsResRecords + +-- | Await and read a single record from the consumer. +-- +readConsumer + ∷ MonadConsumer m + ⇒ KinesisConsumer + → m Kin.Record +readConsumer consumer = + liftIO ∘ atomically $ + consumer ^! kcMessageQueue ∘ act readTBQueue + +-- | Try to read a single record from the consumer; if there is non queued up, +-- then 'Nothing' will be returned. +-- +tryReadConsumer + ∷ MonadConsumer m + ⇒ KinesisConsumer + → m (Maybe Kin.Record) +tryReadConsumer consumer = + liftIO ∘ atomically $ + consumer ^! kcMessageQueue ∘ act tryReadTBQueue + +-- | A conduit for getting records. +-- +consumerSource + ∷ MonadConsumer m + ⇒ KinesisConsumer + → Source m Kin.Record +consumerSource consumer = + forever $ + lift (readConsumer consumer) + ≫= yield diff --git a/src/Aws/Kinesis/Client/Producer.hs b/src/Aws/Kinesis/Client/Producer.hs new file mode 100644 --- /dev/null +++ b/src/Aws/Kinesis/Client/Producer.hs @@ -0,0 +1,567 @@ +-- Copyright (c) 2013-2014 PivotCloud, Inc. +-- +-- Aws.Kinesis.Client.Producer +-- +-- Please feel free to contact us at licensing@pivotmail.com with any +-- contributions, additions, or other feedback; we would love to hear from +-- you. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); you may +-- not use this file except in compliance with the License. You may obtain a +-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +-- License for the specific language governing permissions and limitations +-- under the License. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE UnicodeSyntax #-} + +-- | +-- Module: Aws.Kinesis.Client.Producer +-- Copyright: Copyright © 2013-2014 PivotCloud, Inc. +-- License: Apache-2.0 +-- Maintainer: Jon Sterling +-- Stability: experimental +-- +module Aws.Kinesis.Client.Producer +( -- * The Producer + KinesisProducer +, withKinesisProducer +, managedKinesisProducer + + -- * Commands +, writeProducer +, Message + + -- * Producer Environment +, MonadProducer +, ProducerKit(..) +, pkKinesisKit +, pkStreamName +, pkBatchPolicy +, pkMessageQueueBounds +, pkMaxConcurrency + +, ProducerError(..) +, _KinesisError +, _MessageNotEnqueued +, _InvalidConcurrentConsumerCount + +, BatchPolicy +, defaultBatchPolicy +, bpBatchSize +, bpEndpoint +, RecordEndpoint(..) +) where + +import qualified Aws.Kinesis as Kin +import qualified Aws.Kinesis.Commands.PutRecords as Kin +import Aws.Kinesis.Client.Common + +import Control.Applicative +import Control.Concurrent.Async.Lifted +import Control.Concurrent.Lifted hiding (yield) +import Control.Concurrent.STM +import Control.Concurrent.STM.TBMQueue +import Control.Exception.Lifted +import Control.Lens +import Control.Monad +import Control.Monad.Codensity +import Control.Monad.Error.Class +import Control.Monad.Reader +import Control.Monad.Trans.Control +import Control.Monad.Trans.Either +import Data.Conduit +import Data.Conduit.TQueue +import qualified Data.Conduit.List as CL +import Data.Maybe +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import Data.Traversable +import Data.Typeable +import Prelude.Unicode +import qualified System.Random as R + +-- | There are two endpoints which may be used to send records to Kinesis. +-- +data RecordEndpoint + = PutRecordEndpoint + -- ^ Use the @PutRecord@ endpoint, which sends records one at a time. + + | PutRecordsEndpoint + -- ^ Use the @PutRecords@ endpoint, which sends records in batches. + + deriving (Eq, Show) + +-- | The producer batches records according to a user-specified policy. +-- +data BatchPolicy + = BatchPolicy + { _bpBatchSize ∷ {-# UNPACK #-} !Int + , _bpEndpoint ∷ !RecordEndpoint + } deriving (Eq, Show) + +-- | The number of records to send in a single request. This is only used +-- when the endpoint is set to 'PutRecordsEndpoint'. +-- +bpBatchSize ∷ Lens' BatchPolicy Int +bpBatchSize = lens _bpBatchSize $ \bp bs → bp { _bpBatchSize = bs } + +-- | The endpoint to use when sending records to Kinesis. +-- +bpEndpoint ∷ Lens' BatchPolicy RecordEndpoint +bpEndpoint = lens _bpEndpoint $ \bp ep → bp { _bpEndpoint = ep } + +-- | The default batching policy sends '200' records per 'PutRecordsEndpoint' +-- request. +-- +defaultBatchPolicy ∷ BatchPolicy +defaultBatchPolicy = BatchPolicy + { _bpBatchSize = 200 + , _bpEndpoint = PutRecordsEndpoint + } + + +type Message = T.Text + +data MessageQueueItem + = MessageQueueItem + { _mqiMessage ∷ !Message + , _mqiPartitionKey ∷ !Kin.PartitionKey + } deriving (Eq, Show) + +mqiMessage ∷ Lens' MessageQueueItem Message +mqiMessage = lens _mqiMessage $ \i m → i { _mqiMessage = m } + +mqiPartitionKey ∷ Lens' MessageQueueItem Kin.PartitionKey +mqiPartitionKey = lens _mqiPartitionKey $ \i s → i { _mqiPartitionKey = s } + +-- | The basic input required to construct a Kinesis producer. +-- +data ProducerKit + = ProducerKit + { _pkKinesisKit ∷ !KinesisKit + -- ^ The basic information required to send requests to AWS Kinesis. + + , _pkStreamName ∷ !Kin.StreamName + -- ^ The name of the stream to send records to. + + , _pkBatchPolicy ∷ !BatchPolicy + -- ^ The record batching policy for the producer. + + , _pkMessageQueueBounds ∷ {-# UNPACK #-} !Int + -- ^ The maximum number of records that may be enqueued at one time. + + , _pkMaxConcurrency ∷ {-# UNPACK #-} !Int + -- ^ The number of requests to run concurrently (minimum: 1). + } + +-- | A lens for '_pkKinesisKit'. +-- +pkKinesisKit ∷ Lens' ProducerKit KinesisKit +pkKinesisKit = lens _pkKinesisKit $ \pk kk → pk { _pkKinesisKit = kk } + +-- | A lens for '_pkStreamName'. +-- +pkStreamName ∷ Lens' ProducerKit Kin.StreamName +pkStreamName = lens _pkStreamName $ \pk sn → pk { _pkStreamName = sn } + +-- | A lens for '_pkBatchPolicy'. +-- +pkBatchPolicy ∷ Lens' ProducerKit BatchPolicy +pkBatchPolicy = lens _pkBatchPolicy $ \pk bp → pk { _pkBatchPolicy = bp } + +-- | A lens for '_pkMessageQueueBounds'. +-- +pkMessageQueueBounds ∷ Lens' ProducerKit Int +pkMessageQueueBounds = lens _pkMessageQueueBounds $ \pk qb → pk { _pkMessageQueueBounds = qb } + +-- | A lens for '_pkMaxConcurrency'. +-- +pkMaxConcurrency ∷ Lens' ProducerKit Int +pkMaxConcurrency = lens _pkMaxConcurrency $ \pk n → pk { _pkMaxConcurrency = n } + +-- | The (abstract) Kinesis producer client. +-- +newtype KinesisProducer + = KinesisProducer + { _kpMessageQueue ∷ TBMQueue MessageQueueItem + } + +kpMessageQueue ∷ Getter KinesisProducer (TBMQueue MessageQueueItem) +kpMessageQueue = to _kpMessageQueue + +data ProducerError + = KinesisError !SomeException + -- ^ Represents an error which occured as a result of a request to Kinesis. + + | MessageNotEnqueued Message + -- ^ Thrown when a message could not be enqueued since the queue was full. + -- This error must be handled at the call-site. + + | InvalidConcurrentConsumerCount + -- ^ Thrown when 'pkMaxConcurrency' is set with an invalid value. + + deriving (Typeable, Show) + +instance Exception ProducerError + +-- | A prism for 'KinesisError'. +-- +_KinesisError ∷ Prism' ProducerError SomeException +_KinesisError = + prism KinesisError $ \case + KinesisError e → Right e + e → Left e + +-- | A prism for 'MessageNotEnqueued'. +-- +_MessageNotEnqueued ∷ Prism' ProducerError Message +_MessageNotEnqueued = + prism MessageNotEnqueued $ \case + MessageNotEnqueued m → Right m + e → Left e + +-- | A prism for 'InvalidConcurrentConsumerCount'. +-- +_InvalidConcurrentConsumerCount ∷ Prism' ProducerError () +_InvalidConcurrentConsumerCount = + prism (const InvalidConcurrentConsumerCount) $ \case + InvalidConcurrentConsumerCount → Right () + e → Left e + +-- | The basic effect modality required to use the Kinesis producer. +-- +type MonadProducer m + = ( MonadIO m + , MonadBaseControl IO m + , MonadError ProducerError m + ) + +type MonadProducerInternal m + = ( MonadProducer m + , MonadReader ProducerKit m + ) + +-- | Lifts something in 'MonadKinesis' to 'MonadProducer'. +-- +liftKinesis + ∷ MonadProducerInternal m + ⇒ EitherT SomeException (ReaderT KinesisKit m) α + → m α +liftKinesis = + mapEnvironment pkKinesisKit + ∘ mapError KinesisError + +-- | Generates a valid 'Kin.PartitionKey'. +-- +generatePartitionKey + ∷ R.RandomGen g + ⇒ g + → Kin.PartitionKey +generatePartitionKey gen = + let name = take 25 $ R.randomRs ('a','z') gen in + Kin.partitionKey (T.pack name) + & either (error ∘ T.unpack) id + +-- | This will take up to @n@ items from a 'TBMQueue'. +-- +takeTBMQueue + ∷ Int + → TBMQueue α + → STM [α] +takeTBMQueue n q + | n <= 0 = return [] + | otherwise = do + res ← tryReadTBMQueue q + case res of + Just (Just x) → (x:) <$> takeTBMQueue (n - 1) q + _ → return [] + +-- | A policy for chunking the contents of the message queue. +-- +data ChunkingPolicy + = ChunkingPolicy + { _cpMaxChunkSize ∷ Int + -- ^ The largest chunk size that is permitted. + + , _cpThrottlingDelay ∷ Int + -- ^ The time after which a chunk should be committed, even if the maximum + -- chunk size has not yet been reached. + } + +-- | A 'Source' that reads chunks off a bounded STM queue according some +-- 'ChunkingPolicy'. +-- +chunkedSourceTBMQueue + ∷ ( MonadIO m + , MonadBaseControl IO m + ) + ⇒ ChunkingPolicy + → TBMQueue α + → Source m [α] +chunkedSourceTBMQueue bp@ChunkingPolicy{..} q = do + terminateNow ← liftIO ∘ atomically $ isClosedTBMQueue q + unless terminateNow $ do + items ← liftIO ∘ atomically $ takeTBMQueue _cpMaxChunkSize q + unless (null items) $ do + yield items + + when (length items < _cpMaxChunkSize) $ + threadDelay _cpThrottlingDelay + + chunkedSourceTBMQueue bp q + +-- | Transform a 'Source' into a chunked 'Source' according to some +-- 'ChunkingPolicy'. +-- +chunkSource + ∷ ( MonadIO m + , MonadBaseControl IO m + ) + ⇒ ChunkingPolicy + → Source m α + → Source m [α] +chunkSource cp src = do + queue ← liftIO $ newTBMQueueIO $ _cpMaxChunkSize cp + worker ← lift ∘ async $ src $$+ sinkTBMQueue queue True + addCleanup (\_ → cancel worker) $ do + lift $ link worker + chunkedSourceTBMQueue cp queue + return () + +-- | A conduit for concurently sending multiple records to Kinesis using the +-- @PutRecord@ endpoint. +-- +concurrentPutRecordSink + ∷ MonadProducerInternal m + ⇒ Sink [MessageQueueItem] m () +concurrentPutRecordSink = do + maxWorkerCount ← view pkMaxConcurrency + awaitForever $ \messages → do + lift ∘ flip (mapConcurrentlyN maxWorkerCount 100) messages $ \m → do + CL.sourceList [m] $$ putRecordSink + + +-- | A conduit for sending a record to Kinesis using the @PutRecord@ endpoint; +-- this is a conduit in order to restore failed messages as leftovers. +-- +putRecordSink + ∷ MonadProducerInternal m + ⇒ Sink MessageQueueItem m () +putRecordSink = do + streamName ← view pkStreamName + + awaitForever $ \item → do + let handler e = do + liftIO $ do + putStrLn $ "Error: " ++ show e + putStrLn "Will wait 5s" + threadDelay 5000000 + leftover item + + handleError handler $ do + let partitionKey = item ^. mqiPartitionKey + void ∘ lift ∘ liftKinesis $ runKinesis Kin.PutRecord + { Kin.putRecordData = item ^. mqiMessage ∘ to T.encodeUtf8 + , Kin.putRecordExplicitHashKey = Nothing + , Kin.putRecordPartitionKey = partitionKey + , Kin.putRecordSequenceNumberForOrdering = Nothing + , Kin.putRecordStreamName = streamName + } + +splitEvery + ∷ Int + → [α] + → [[α]] +splitEvery _ [] = [] +splitEvery n list = first : splitEvery n rest + where + (first,rest) = splitAt n list + +-- | A conduit for sending records to Kinesis using the @PutRecords@ endpoint. +-- This is a conduit in order to restore failed messages as leftovers. +-- +putRecordsSink + ∷ MonadProducerInternal m + ⇒ Sink [MessageQueueItem] m () +putRecordsSink = do + streamName ← view pkStreamName + batchSize ← view $ pkBatchPolicy ∘ bpBatchSize + maxWorkerCount ← view pkMaxConcurrency + awaitForever $ \messages → do + let batches = splitEvery batchSize messages + leftovers ← lift ∘ flip (mapConcurrentlyN maxWorkerCount 100) batches $ \ms → do + let handler e = do + liftIO $ print e + return ms + + handleError handler $ do + items ← for ms $ \m → do + let partitionKey = m ^. mqiPartitionKey + return Kin.PutRecordsRequestEntry + { Kin.putRecordsRequestEntryData = m ^. mqiMessage ∘ to T.encodeUtf8 + , Kin.putRecordsRequestEntryExplicitHashKey = Nothing + , Kin.putRecordsRequestEntryPartitionKey = partitionKey + } + + Kin.PutRecordsResponse{..} ← liftKinesis $ runKinesis Kin.PutRecords + { Kin.putRecordsRecords = items + , Kin.putRecordsStreamName = streamName + } + let processResult m m' + | isJust (Kin.putRecordsResponseRecordErrorCode m') = Just m + | otherwise = Nothing + return ∘ catMaybes $ zipWith processResult ms putRecordsResponseRecords + + forM_ leftovers $ \mss → + unless (null mss) $ leftover mss + +sendMessagesSink + ∷ MonadProducerInternal m + ⇒ Sink [MessageQueueItem] m () +sendMessagesSink = do + batchPolicy ← view pkBatchPolicy + case batchPolicy ^. bpEndpoint of + PutRecordsEndpoint → putRecordsSink + PutRecordEndpoint → concurrentPutRecordSink + +-- | Enqueues a message to Kinesis on the next shard. If a message cannot be +-- enqueued (because the client has exceeded its queue size), the +-- 'MessageNotEnqueued' exception will be thrown. +-- +writeProducer + ∷ MonadProducer m + ⇒ KinesisProducer + → Message + → m () +writeProducer producer !msg = do + gen ← liftIO R.newStdGen + result ← liftIO ∘ atomically $ do + tryWriteTBMQueue (producer ^. kpMessageQueue) MessageQueueItem + { _mqiMessage = msg + , _mqiPartitionKey = generatePartitionKey gen + } + case result of + Just True → return () + _ → throwError $ MessageNotEnqueued msg + +-- | This is a 'Source' that returns all the items presently in a queue: it +-- terminates when the queue is empty. +-- +exhaustTBMQueue + ∷ TBMQueue α + → Source STM α +exhaustTBMQueue q = do + mx ← lift $ tryReadTBMQueue q + case mx of + Just (Just x) → do + yield x + exhaustTBMQueue q + _ → return () + +-- | This constructs a 'KinesisProducer' and closes it when you have done with +-- it. This is equivalent to 'withKinesisProducer', but replaces the +-- continuation with a return in 'Codensity'. +-- +managedKinesisProducer + ∷ ∀ m + . ( MonadIO m + , MonadBaseControl IO m + , MonadError ProducerError m + ) + ⇒ ProducerKit + → Codensity m KinesisProducer +managedKinesisProducer kit = do + when (kit ^. pkMaxConcurrency < 1) ∘ lift $ + throwError InvalidConcurrentConsumerCount + + messageQueue ← liftIO ∘ newTBMQueueIO $ kit ^. pkMessageQueueBounds + + let chunkingPolicy = ChunkingPolicy ((kit ^. pkBatchPolicy ∘ bpBatchSize) * (kit ^. pkMaxConcurrency)) 5000000 + -- TODO: this 'forever' is only here to restart if we get killed. + -- Replace with proper error handling. + consumerLoop ∷ m () = flip runReaderT kit ∘ forever $ + chunkSource chunkingPolicy (sourceTBMQueue messageQueue) + $$ sendMessagesSink + + let cleanupConsumer consumerHandle = do + liftIO ∘ atomically $ closeTBMQueue messageQueue + flip runReaderT kit $ do + leftovers ← liftIO ∘ atomically $ + exhaustTBMQueue messageQueue + $$ CL.consume + chunkSource chunkingPolicy (CL.sourceList leftovers) + $$ sendMessagesSink + cancel consumerHandle + + consumerHandle ← managedBracket (async consumerLoop) cleanupConsumer + + Codensity $ \inner → do + link consumerHandle + res ← inner $ KinesisProducer messageQueue + () ← wait consumerHandle + return res + + +-- | This constructs a 'KinesisProducer' and closes it when you have done with +-- it. +-- +withKinesisProducer + ∷ ( MonadIO m + , MonadBaseControl IO m + , MonadError ProducerError m + ) + ⇒ ProducerKit + → (KinesisProducer → m α) + → m α +withKinesisProducer = + runCodensity ∘ managedKinesisProducer + +managedBracket + ∷ MonadBaseControl IO m + ⇒ m α + → (α → m β) + → Codensity m α +managedBracket action cleanup = + Codensity $ bracket action cleanup + +-- | map at most n actions concurrently +-- +mapConcurrentlyN + ∷ ( MonadIO m + , MonadBaseControl IO m + , Traversable t + ) + ⇒ Int -- ^ number of concurrent actions + → Int -- ^ startup delay between actions in milliseconds + → (a → m b) + → t a + → m (t b) +mapConcurrentlyN n delay f t = do + sem ← liftIO $ newQSem n + mapConcurrently (run sem) t_ + where + (_, t_) = mapAccumL (\i v → (succ i, (i,v))) 0 t + run sem (i,a) = + liftBaseOp_ (bracket_ (waitQSem sem) (signalQSem sem)) $ do + liftIO ∘ threadDelay $ 1000 * delay * i + f a +