aws-kinesis-client 0.1.0.2 → 0.2.0.0
raw patch · 5 files changed
+168/−29 lines, 5 filesdep +aesondep +containersdep +unordered-containers
Dependencies added: aeson, containers, unordered-containers
Files
- CHANGELOG.md +7/−0
- aws-kinesis-client.cabal +6/−2
- cli/CLI.hs +16/−2
- cli/CLI/Options.hs +22/−3
- src/Aws/Kinesis/Client/Consumer.hs +117/−22
CHANGELOG.md view
@@ -1,3 +1,10 @@+### v0.2.0.0++- [Consumer] Add the ability to save & restore stream state (i.e. last consumed+ sequence number per shard).++- [Consumer CLI] Add --save-state, --restore-state options+ ### v0.1.0.2 - Support specifying AWS keys as options to the CLI (either directly or from a
aws-kinesis-client.cabal view
@@ -1,5 +1,5 @@ name: aws-kinesis-client-version: 0.1.0.2+version: 0.2.0.0 synopsis: A producer & consumer client library for AWS Kinesis -- description: license: Apache-2.0@@ -24,9 +24,11 @@ -- other-extensions: build-depends: base >=4.7 && <4.8, base-unicode-symbols,+ aeson >=0.8, aws >=0.10.5, aws-kinesis >=0.1.2, conduit >=1.2.3.1,+ containers >=0.5.5.1, data-carousel >=0.1.0.0, either >=4.3.2.1, errors >=1.4.7,@@ -45,7 +47,8 @@ stm-conduit >=2.5.3, stm-queue-extras >=0.2.0.0, text >=1.2.0.3,- transformers >=0.3.0.0+ transformers >=0.3.0.0,+ unordered-containers >=0.2.5.1 hs-source-dirs: src default-language: Haskell2010 ghc-options: -Wall@@ -57,6 +60,7 @@ other-modules: CLI.Options build-depends: base >=4.7 && <4.8, base-unicode-symbols,+ aeson >=0.8, aws >=0.10.5, aws-general >=0.1.1, aws-kinesis >=0.1.2,
cli/CLI.hs view
@@ -36,7 +36,6 @@ -- Maintainer: Jon Sterling <jsterling@alephcloud.com> -- Stability: experimental --- module Main ( main ) where@@ -51,6 +50,7 @@ import Control.Exception import Control.Lens+import Control.Monad import Control.Monad.Trans import Control.Monad.Trans.Control import Control.Monad.Trans.Either@@ -60,9 +60,12 @@ import Control.Monad.Error.Class import Control.Monad.Error.Hoist +import qualified Data.Aeson as A import qualified Data.ByteString.Char8 as B8+import qualified Data.ByteString.Lazy.Char8 as BL8 import Data.Conduit import qualified Data.Conduit.List as CL+import Data.Traversable import Data.Typeable import Options.Applicative@@ -110,6 +113,13 @@ CLIOptions{..} ← ask manager ← managedHttpManager credentials ← lift fetchCredentials+ savedStreamState ←+ for _clioStateIn $+ either (fail ∘ ("Invalid saved state: " ++)) return+ <=< fmap A.eitherDecode+ ∘ liftIO+ ∘ BL8.readFile+ consumer ← managedKinesisConsumer $ ConsumerKit { _ckKinesisKit = KinesisKit { _kkManager = manager@@ -123,11 +133,15 @@ , _ckStreamName = _clioStreamName , _ckBatchSize = 100 , _ckIteratorType = _clioIteratorType+ , _ckSavedStreamState = savedStreamState } lift $ consumerSource consumer $$ limitConduit =$ CL.mapM_ (liftIO ∘ B8.putStrLn ∘ recordData)- return ()++ void ∘ for _clioStateOut $ \outPath → do+ state ← lift $ consumerStreamState consumer+ liftIO ∘ BL8.writeFile outPath $ A.encode state main ∷ IO () main =
cli/CLI/Options.hs view
@@ -28,13 +28,14 @@ -- Maintainer: Jon Sterling <jsterling@alephcloud.com> -- Stability: experimental --- module CLI.Options ( CLIOptions(..) , clioStreamName , clioLimit , clioIteratorType , clioAccessKeys+, clioStateIn+, clioStateOut , AccessKeys(..) , akAccessKeyId , akSecretAccessKey@@ -63,6 +64,8 @@ , _clioLimit ∷ !Int , _clioIteratorType ∷ !ShardIteratorType , _clioAccessKeys ∷ !(Either AccessKeys FilePath)+ , _clioStateIn ∷ !(Maybe FilePath)+ , _clioStateOut ∷ !(Maybe FilePath) } deriving Show makeLenses ''AccessKeys@@ -103,7 +106,7 @@ strOption $ long "access-keys-path" ⊕ metavar "PATH"- ⊕ help "The path to a file containing your access keys. To be formatted \"default ID SECRET\""+ ⊕ help "The path to a file containing your access keys. To be formatted \"default ID SECRET\"; you may provide this instead of the command line options for credentials" streamNameParser ∷ Parser StreamName streamNameParser =@@ -131,6 +134,20 @@ ⊕ value TrimHorizon ⊕ showDefault +stateOutParser ∷ Parser FilePath+stateOutParser =+ strOption $+ long "save-state"+ ⊕ help "Write the last known state of each shard to a file"+ ⊕ metavar "FILE"++stateInParser ∷ Parser FilePath+stateInParser =+ strOption $+ long "restore-state"+ ⊕ help "Read a saved stream state from a file. For any shards whose state is restored, the 'AFTER_SEQUENCE_NUMBER' iterator type will be used; other shards will use the iterator type you have specified. Some shards may have been merged or closed between when the state was saved and restored; at this point, no effort has been made to do anything here beyond the obvious (shards are identified by their shard-id)."+ ⊕ metavar "FILE"+ optionsParser ∷ Parser CLIOptions optionsParser = pure CLIOptions@@ -138,11 +155,13 @@ ⊛ limitParser ⊛ iteratorTypeParser ⊛ (Left <$> accessKeysParser <|> Right <$> accessKeysPathParser)+ ⊛ optional stateInParser+ ⊛ optional stateOutParser parserInfo ∷ ParserInfo CLIOptions parserInfo = info (helper ⊛ optionsParser) $ fullDesc- ⊕ progDesc "Fetch `L` records from a Kinesis stream `SN`."+ ⊕ progDesc "Fetch a given number of records from a Kinesis stream; unlike the standard command line utilities, this interface is suitable for use with a sharded stream. If you both specify a saved stream state to be restored and an iterator type, the latter will be used on any shards which are not contained in the saved state. Minimally, you must specify your AWS credentials, a stream name, and a limit." ⊕ header "The Kinesis Consumer CLI"
src/Aws/Kinesis/Client/Consumer.hs view
@@ -23,6 +23,7 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UnicodeSyntax #-} @@ -43,6 +44,7 @@ , consumerSource , readConsumer , tryReadConsumer+, consumerStreamState -- * Consumer Environment , ConsumerKit(..)@@ -51,11 +53,14 @@ , ckBatchSize , ConsumerError(..) , MonadConsumer+, SavedStreamState ) where import qualified Aws.Kinesis as Kin import Aws.Kinesis.Client.Common +import Control.Applicative+import Control.Applicative.Unicode import Control.Concurrent.Async.Lifted import Control.Concurrent.Lifted hiding (yield) import Control.Concurrent.STM@@ -68,7 +73,11 @@ import Control.Monad.Reader import Control.Monad.Trans.Control import Control.Monad.Unicode+import qualified Data.Aeson as Æ import qualified Data.Carousel as CR+import qualified Data.Map as M+import qualified Data.HashMap.Strict as HM+import Data.Traversable (for) import Data.Conduit import qualified Data.Conduit.List as CL import Prelude.Unicode@@ -79,9 +88,10 @@ = ShardState { _ssIterator ∷ !(TVar (Maybe Kin.ShardIterator)) , _ssShardId ∷ !Kin.ShardId+ , _ssLastSequenceNumber ∷ !(TVar (Maybe Kin.SequenceNumber)) } --- | A lens for '_ssIterator'.+-- | A getter for '_ssIterator'. -- ssIterator ∷ Getter ShardState (TVar (Maybe Kin.ShardIterator)) ssIterator = to _ssIterator@@ -91,6 +101,11 @@ ssShardId ∷ Lens' ShardState Kin.ShardId ssShardId = lens _ssShardId $ \ss sid → ss { _ssShardId = sid } +-- | A getter for '_ssLastSequenceNumber'.+--+ssLastSequenceNumber ∷ Getter ShardState (TVar (Maybe Kin.SequenceNumber))+ssLastSequenceNumber = to _ssLastSequenceNumber+ -- | 'ShardState' is quotiented by shard ID. -- instance Eq ShardState where@@ -119,6 +134,11 @@ , _ckIteratorType ∷ !Kin.ShardIteratorType -- ^ The type of iterator to consume.++ , _ckSavedStreamState ∷ !(Maybe SavedStreamState)+ -- ^ Optionally, an initial stream state. The iterator type in+ -- '_ckIteratorType' will be used for any shards not present in the saved+ -- stream state; otherwise, 'Kin.AfterSequenceNumber' will be used. } -- | A lens for '_ckKinesisKit'.@@ -141,15 +161,58 @@ ckIteratorType ∷ Lens' ConsumerKit Kin.ShardIteratorType ckIteratorType = lens _ckIteratorType $ \ck it → ck { _ckIteratorType = it } +-- | A lens for '_ckSavedStreamState'.+--+ckSavedStreamState ∷ Lens' ConsumerKit (Maybe SavedStreamState)+ckSavedStreamState = lens _ckSavedStreamState $ \ck ss → ck { _ckSavedStreamState = ss }+++type MessageQueueItem = (ShardState, Kin.Record)+type StreamState = CR.Carousel ShardState++newtype SavedStreamState+ = SavedStreamState+ { _savedStreamState ∷ M.Map Kin.ShardId Kin.SequenceNumber+ }++-- | An iso for 'SavedStreamState'.+--+_SavedStreamState ∷ Iso' SavedStreamState (M.Map Kin.ShardId Kin.SequenceNumber)+_SavedStreamState = iso _savedStreamState SavedStreamState++instance Æ.ToJSON SavedStreamState where+ toJSON (SavedStreamState m) =+ Æ.Object ∘ HM.fromList ∘ flip fmap (M.toList m) $ \(sid, sn) →+ let Æ.String sid' = Æ.toJSON sid+ in sid' Æ..= sn++instance Æ.FromJSON SavedStreamState where+ parseJSON =+ Æ.withObject "SavedStreamState" $ \xs → do+ fmap (SavedStreamState ∘ M.fromList) ∘ for (HM.toList xs) $ \(sid, sn) → do+ pure (,)+ ⊛ Æ.parseJSON (Æ.String sid)+ ⊛ Æ.parseJSON sn++ -- | The 'KinesisConsumer' maintains state about which shards to pull from. ---newtype KinesisConsumer = KinesisConsumer { _kcMessageQueue ∷ TBQueue Kin.Record }+data KinesisConsumer+ = KinesisConsumer+ { _kcMessageQueue ∷ !(TBQueue MessageQueueItem)+ , _kcStreamState ∷ !(TVar StreamState)+ } -- | A getter for '_kcMessageQueue'. ---kcMessageQueue ∷ Getter KinesisConsumer (TBQueue Kin.Record)+kcMessageQueue ∷ Getter KinesisConsumer (TBQueue MessageQueueItem) kcMessageQueue = to _kcMessageQueue +-- | A getter for '_kcStreamState'.+--+kcStreamState ∷ Getter KinesisConsumer (TVar StreamState)+kcStreamState = to _kcStreamState+ -- | The basic effect modality required for operating the consumer. -- type MonadConsumer m@@ -198,7 +261,7 @@ producerLoop = forever $ handleError (\_ → liftIO $ threadDelay 2000000) $ do recordsCount ← replenishMessages messageQueue state- liftIO . threadDelay $+ liftIO ∘ threadDelay $ case recordsCount of 0 → 5000000 _ → 70000@@ -208,7 +271,7 @@ link reshardingHandle withAsync producerLoop $ \producerHandle → do link producerHandle- res ← lift ∘ inner $ KinesisConsumer messageQueue+ res ← lift ∘ inner $ KinesisConsumer messageQueue state return res -- | This requests new information from Kinesis and reconciles that with an@@ -216,11 +279,12 @@ -- updateStreamState ∷ MonadConsumerInternal m- ⇒ CR.Carousel ShardState- → m (CR.Carousel ShardState)+ ⇒ StreamState+ → m StreamState updateStreamState state = do streamName ← view ckStreamName- iteratorType ← view ckIteratorType+ defaultIteratorType ← view ckIteratorType+ savedState ← view ckSavedStreamState mapError KinesisError ∘ mapEnvironment ckKinesisKit $ do let existingShardIds = state ^. CR.clList <&> view ssShardId@@ -231,25 +295,36 @@ newShards ← shardSource $$ CL.consume shardStates ← forM newShards $ \Kin.Shard{..} → do- Kin.GetShardIteratorResponse it ← runKinesis Kin.GetShardIterator+ let startingSequenceNumber =+ savedState ^? _Just ∘ _SavedStreamState ∘ ix shardShardId+ iteratorType =+ maybe+ defaultIteratorType+ (const Kin.AfterSequenceNumber)+ startingSequenceNumber++ Kin.GetShardIteratorResponse it ← runKinesis Kin.GetShardIterator { Kin.getShardIteratorShardId = shardShardId , Kin.getShardIteratorShardIteratorType = iteratorType- , Kin.getShardIteratorStartingSequenceNumber = Nothing+ , Kin.getShardIteratorStartingSequenceNumber = startingSequenceNumber , Kin.getShardIteratorStreamName = streamName }- iteratorVar ← liftIO ∘ newTVarIO $ Just it- return ShardState- { _ssIterator = iteratorVar- , _ssShardId = shardShardId- }+ liftIO $ do+ iteratorVar ← newTVarIO $ Just it+ sequenceNumberVar ← newTVarIO Nothing+ return ShardState+ { _ssIterator = iteratorVar+ , _ssShardId = shardShardId+ , _ssLastSequenceNumber = sequenceNumberVar+ } 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)+ ⇒ TBQueue MessageQueueItem+ → TVar StreamState → m Int replenishMessages messageQueue shardsVar = do bufferSize ← view ckBatchSize@@ -268,7 +343,7 @@ liftIO ∘ atomically $ do writeTVar (shard ^. ssIterator) getRecordsResNextShardIterator- forM_ getRecordsResRecords $ writeTBQueue messageQueue+ forM_ getRecordsResRecords $ writeTBQueue messageQueue ∘ (shard ,) modifyTVar shardsVar CR.moveRight return $ length getRecordsResRecords@@ -280,8 +355,10 @@ ⇒ KinesisConsumer → m Kin.Record readConsumer consumer =- liftIO ∘ atomically $- consumer ^! kcMessageQueue ∘ act readTBQueue+ liftIO ∘ atomically $ do+ (ss, rec) ← consumer ^! kcMessageQueue ∘ act readTBQueue+ writeTVar (ss ^. ssLastSequenceNumber) ∘ Just $ Kin.recordSequenceNumber rec+ return rec -- | Try to read a single record from the consumer; if there is non queued up, -- then 'Nothing' will be returned.@@ -291,8 +368,11 @@ ⇒ KinesisConsumer → m (Maybe Kin.Record) tryReadConsumer consumer =- liftIO ∘ atomically $- consumer ^! kcMessageQueue ∘ act tryReadTBQueue+ liftIO ∘ atomically $ do+ mitem ← consumer ^! kcMessageQueue ∘ act tryReadTBQueue+ for mitem $ \(ss, rec) → do+ writeTVar (ss ^. ssLastSequenceNumber) ∘ Just $ Kin.recordSequenceNumber rec+ return rec -- | A conduit for getting records. --@@ -304,3 +384,18 @@ forever $ lift (readConsumer consumer) ≫= yield++-- | Get the last read sequence number at each shard.+--+consumerStreamState+ ∷ MonadConsumer m+ ⇒ KinesisConsumer+ → m SavedStreamState+consumerStreamState consumer =+ liftIO ∘ atomically $ do+ shards ← consumer ^! kcStreamState ∘ act readTVar ∘ CR.clList+ pairs ← for shards $ \ss →+ (ss ^. ssShardId,) <$>+ ss ^! ssLastSequenceNumber ∘ act readTVar+ return ∘ SavedStreamState ∘ M.fromList $ pairs ≫= \(sid, msn) →+ msn ^.. _Just ∘ to (sid,)