milena 0.1.0.0 → 0.2.0.0

raw patch · 5 files changed

+200/−119 lines, 5 filesdep +resource-pooldep ~bytestringdep ~digestdep ~lensPVP ok

version bump matches the API change (PVP)

Dependencies added: resource-pool

Dependency ranges changed: bytestring, digest, lens, mtl, transformers

API changes (from Hackage documentation)

Files

• Network/Kafka.hs +72/−102
• Network/Kafka/Producer.hs +108/−0
• Network/Kafka/Protocol.hs +3/−3
• milena.cabal +15/−13
• test/tests.hs +2/−1

@@ -12,9 +12,9 @@ import Control.Monad.Trans.State import Data.ByteString.Char8 (ByteString) import Data.Monoid ((<>))+import qualified Data.Pool as Pool import Data.Serialize.Get import System.IO-import System.Random (getStdRandom, randomR) import qualified Data.ByteString.Char8 as B import qualified Data.Map as M import qualified Network@@ -23,8 +23,6 @@  data KafkaState = KafkaState { -- | Name to use as a client ID.                                _stateName :: KafkaString-                               -- | An incrementing counter of requests.-                             , _stateCorrelationId :: CorrelationId                                -- | How many acknowledgements are required for producing.                              , _stateRequiredAcks :: RequiredAcks                                -- | Time in milliseconds to wait for messages to be produced by broker.@@ -35,20 +33,26 @@                              , _stateBufferSize :: MaxBytes                                -- | Maximum time in milliseconds to wait for response.                              , _stateWaitTime :: MaxWaitTime+                               -- | An incrementing counter of requests.+                             , _stateCorrelationId :: CorrelationId                                -- | Broker cache                              , _stateBrokers :: M.Map Leader Broker+                               -- | Connection cache+                             , _stateConnections :: M.Map Broker (Pool.Pool Handle)+                               -- | Topic metadata cache+                             , _stateTopicMetadata :: M.Map TopicName TopicMetadata                              }  makeLenses ''KafkaState -data KafkaConsumer = KafkaConsumer { _consumerState :: KafkaState-                                   , _consumerHandle :: Handle-                                   }+data KafkaClient = KafkaClient { _kafkaClientState :: KafkaState+                               , _kafkaClientHandle :: Handle+                               } -makeLenses ''KafkaConsumer+makeLenses ''KafkaClient  -- | The core Kafka monad.-type Kafka = StateT KafkaConsumer (EitherT KafkaClientError IO)+type Kafka = StateT KafkaClient (EitherT KafkaClientError IO)  type KafkaAddress = (Host, Port) type KafkaClientId = KafkaString@@ -62,6 +66,7 @@                       | KafkaDeserializationError String -- TODO: cereal is Stringly typed, should use tickle                         -- | Could not find a cached broker for the found leader.                       | KafkaInvalidBroker Leader+                      | KafkaFailedToFetchMetadata                         deriving (Eq, Show)  -- | Type of response to expect, used for 'KafkaExpected' error.@@ -82,6 +87,7 @@                                              , _palPartition :: Partition                                              , _palLeader :: Leader                                              }+                                             deriving (Show)  makeLenses ''PartitionAndLeader @@ -132,31 +138,37 @@ defaultState :: KafkaClientId -> KafkaState defaultState cid =     KafkaState cid-               defaultCorrelationId                defaultRequiredAcks                defaultRequestTimeout                defaultMinBytes                defaultMaxBytes                defaultMaxWaitTime+               defaultCorrelationId                M.empty+               M.empty+               M.empty  -- | Run the underlying Kafka monad at the given leader address and initial state. runKafka :: KafkaAddress -> KafkaState -> Kafka a -> IO (Either KafkaClientError a) runKafka (h, p) s k =-    bracket (Network.connectTo (h ^. hostString) (p ^. portId)) hClose $ runEitherT . evalStateT k . KafkaConsumer s+  bracket (Network.connectTo (h ^. hostString) (p ^. portId)) hClose $ runEitherT . evalStateT k . KafkaClient s  -- | Make a request, incrementing the `_stateCorrelationId`. makeRequest :: RequestMessage -> Kafka Request makeRequest m = do-  corid <- use (consumerState . stateCorrelationId)-  consumerState . stateCorrelationId += 1-  conid <- use (consumerState . stateName)+  corid <- use (kafkaClientState . stateCorrelationId)+  kafkaClientState . stateCorrelationId += 1+  conid <- use (kafkaClientState . stateName)   return $ Request (corid, ClientId conid, m)  -- | Perform a request and deserialize the response. doRequest :: Request -> Kafka Response-doRequest r = mapStateT (bimapEitherT KafkaDeserializationError id) $ do-  h <- use consumerHandle+doRequest r = do+  h <- use kafkaClientHandle+  doRequest' h r++doRequest' :: Handle -> Request -> Kafka Response+doRequest' h r = mapStateT (bimapEitherT KafkaDeserializationError id) $ do   dataLength <- lift . EitherT $ do     B.hPut h $ requestBytes r     hFlush h@@ -167,9 +179,15 @@  -- | Send a metadata request metadata :: MetadataRequest -> Kafka MetadataResponse-metadata request =-    makeRequest (MetadataRequest request) >>= doRequest >>= expectResponse ExpectedMetadata _MetadataResponse+metadata request = do+  h <- use kafkaClientHandle+  metadata' h request +-- | Send a metadata request+metadata' :: Handle -> MetadataRequest -> Kafka MetadataResponse+metadata' handle request =+    makeRequest (MetadataRequest request) >>= doRequest' handle >>= expectResponse ExpectedMetadata _MetadataResponse+ -- | Function to give an error when the response seems wrong. expectResponse :: KafkaExpectedResponse -> Getting (Leftmost b) ResponseMessage b -> Response -> Kafka b expectResponse e p = lift . maybe (left $ KafkaExpected e) return . firstOf (responseMessage . p)@@ -180,61 +198,6 @@ protocolTime EarliestTime = Time (-2) protocolTime (OtherTime o) = o --- * Messages---- | Group messages together with the leader they should be sent to.-partitionAndCollate :: [TopicAndMessage] -> Kafka (M.Map Leader (M.Map TopicAndPartition [TopicAndMessage]))-partitionAndCollate ks = recurse ks M.empty-      where recurse [] accum = return accum-            recurse (x:xs) accum = do-              topicPartitionsList <- brokerPartitionInfo $ _tamTopic x-              pal <- getPartition topicPartitionsList-              let leader = maybe (Leader Nothing) _palLeader pal-                  tp = TopicAndPartition <$> pal ^? folded . palTopic <*> pal ^? folded . palPartition-                  b = M.singleton leader $ maybe M.empty (`M.singleton` [x]) tp-                  accum' = M.unionWith (M.unionWith (<>)) accum b-              recurse xs accum'--getPartition :: [PartitionAndLeader] -> Kafka (Maybe PartitionAndLeader)-getPartition ps =-    liftIO $ (ps' ^?) . element <$> getStdRandom (randomR (0, length ps' - 1))-        where ps' = ps ^.. folded . filtered (has $ palLeader . leaderId . _Just)---- | Create a protocol message set from a list of messages.-groupMessagesToSet :: [TopicAndMessage] -> MessageSet-groupMessagesToSet xs = MessageSet $ uncurry msm <$> zip [0..] xs-    where msm n = MessageSetMember (Offset n) . _tamMessage---- | Find a leader and partition for the topic.-brokerPartitionInfo :: TopicName -> Kafka [PartitionAndLeader]-brokerPartitionInfo t = do-  md <- metadata $ MetadataReq [t]-  let brokers = md ^.. metadataResponseFields . _1 . folded-  consumerState . stateBrokers .= foldr addBroker M.empty brokers-  return $ pal <$> md ^.. topicsMetadata . folded . partitionsMetadata . folded-      where pal d = PartitionAndLeader t (d ^. partitionId) (d ^. partitionMetadataLeader)-            addBroker b = M.insert (Leader . Just $ b ^. brokerFields . _1 . nodeId) b---- | Default: @1@-defaultMessageCrc :: Crc-defaultMessageCrc = 1---- | Default: @0@-defaultMessageMagicByte :: MagicByte-defaultMessageMagicByte = 0---- | Default: @Nothing@-defaultMessageKey :: Key-defaultMessageKey = Key Nothing---- | Default: @0@-defaultMessageAttributes :: Attributes-defaultMessageAttributes = 0---- | Construct a message from a string of bytes using default attributes.-makeMessage :: ByteString -> Message-makeMessage m = Message (defaultMessageCrc, defaultMessageMagicByte, defaultMessageAttributes, defaultMessageKey, Value (Just (KBytes m)))- -- * Fetching  -- | Default: @-1@@@ -244,9 +207,9 @@ -- | Construct a fetch request from the values in the state. fetchRequest :: Offset -> Partition -> TopicName -> Kafka FetchRequest fetchRequest o p topic = do-  wt <- use (consumerState . stateWaitTime)-  ws <- use (consumerState . stateWaitSize)-  bs <- use (consumerState . stateBufferSize)+  wt <- use (kafkaClientState . stateWaitTime)+  ws <- use (kafkaClientState . stateWaitSize)+  bs <- use (kafkaClientState . stateBufferSize)   return $ FetchReq (ordinaryConsumerId, wt, ws, [(topic, [(p, o, bs)])])  -- | Execute a fetch request and get the raw fetch response.@@ -259,37 +222,44 @@ fetchMessages fr = (fr ^.. fetchResponseFields . folded) >>= tam     where tam a = TopicAndMessage (a ^. _1) <$> a ^.. _2 . folded . _4 . messageSetMembers . folded . setMessage --- * Producing+updateMetadatas :: [TopicName] -> Kafka ()+updateMetadatas ts = do+  md <- metadata $ MetadataReq ts+  let (brokers, tmds) = (md ^.. metadataResponseBrokers . folded, md ^.. topicsMetadata . folded)+  kafkaClientState . stateBrokers %= \m -> foldr addBroker m brokers+  kafkaClientState . stateTopicMetadata %= \m -> foldr addTopicMetadata m tmds+  return ()+    where addBroker :: Broker -> M.Map Leader Broker -> M.Map Leader Broker+          addBroker b = M.insert (Leader . Just $ b ^. brokerNode . nodeId) b+          addTopicMetadata :: TopicMetadata -> M.Map TopicName TopicMetadata -> M.Map TopicName TopicMetadata+          addTopicMetadata tm = M.insert (tm ^. topicMetadataName) tm --- | Execute a produce request and get the raw preduce response.-produce :: ProduceRequest -> Kafka ProduceResponse-produce request =-    makeRequest (ProduceRequest request) >>= doRequest >>= expectResponse ExpectedProduce _ProduceResponse+updateMetadata :: TopicName -> Kafka ()+updateMetadata t = updateMetadatas [t] --- | Construct a produce request with explicit arguments.-produceRequest :: RequiredAcks -> Timeout -> [(TopicAndPartition, MessageSet)] -> ProduceRequest-produceRequest ra ti ts =-    ProduceReq (ra, ti, M.toList . M.unionsWith (<>) $ fmap f ts)-        where f (TopicAndPartition t p, i) = M.singleton t [(p, i)]+updateAllMetadata :: Kafka ()+updateAllMetadata = updateMetadatas [] --- | Send messages to partition calculated by 'partitionAndCollate'.-produceMessages :: [TopicAndMessage] -> Kafka [ProduceResponse]-produceMessages tams = do-  m <- fmap (fmap groupMessagesToSet) <$> partitionAndCollate tams-  mapM (uncurry send) $ fmap M.toList <$> M.toList m+-- | Execute a handler action, creating a new Pool and updating the connections Map if needed.+withBrokerHandle :: Broker -> (Handle -> Kafka a) -> Kafka a+withBrokerHandle broker f = do+  conns <- use (kafkaClientState . stateConnections)+  let foundPool = conns ^. at broker+  pool <- case foundPool of+    Nothing -> do+      newPool <- liftIO $ mkPool broker+      kafkaClientState . stateConnections .= (at broker ?~ newPool $ conns)+      return newPool+    Just p -> return p+  Pool.withResource pool f+    where mkPool :: Broker -> IO (Pool.Pool Handle)+          mkPool b = Pool.createPool (createHandle b) hClose 1 10 1+          createHandle b = do+            let h = b ^. brokerHost ^. hostString+                p = b ^. brokerPort ^. portId+            Network.connectTo h p --- | Execute a produce request using the values in the state.-send :: Leader -> [(TopicAndPartition, MessageSet)] -> Kafka ProduceResponse-send l ts = do-  foundBroker <- use (consumerState . stateBrokers . at l)-  broker <- lift $ maybe (left $ KafkaInvalidBroker l) right foundBroker-  requiredAcks <- use (consumerState . stateRequiredAcks)-  requestTimeout <- use (consumerState . stateRequestTimeout)-  let h' = broker ^. brokerFields . _2-      p' = broker ^. brokerFields . _3-  cstate <- use consumerState-  r <- liftIO . runKafka (h', p') cstate . produce $ produceRequest requiredAcks requestTimeout ts-  lift $ either left right r+  -- * Offsets 

@@ -0,0 +1,108 @@+module Network.Kafka.Producer where++import Control.Applicative+import Control.Lens+import Control.Monad.Trans (liftIO, lift)+import Control.Monad.Trans.Either+import Data.ByteString.Char8 (ByteString)+import Data.Monoid ((<>))+import System.IO+import qualified Data.Map as M+import System.Random (getStdRandom, randomR)++import Network.Kafka+import Network.Kafka.Protocol++-- * Producing++-- | Execute a produce request and get the raw preduce response.+produce :: Handle -> ProduceRequest -> Kafka ProduceResponse+produce handle request =+    makeRequest (ProduceRequest request) >>= doRequest' handle >>= expectResponse ExpectedProduce _ProduceResponse++-- | Construct a produce request with explicit arguments.+produceRequest :: RequiredAcks -> Timeout -> [(TopicAndPartition, MessageSet)] -> ProduceRequest+produceRequest ra ti ts =+    ProduceReq (ra, ti, M.toList . M.unionsWith (<>) $ fmap f ts)+        where f (TopicAndPartition t p, i) = M.singleton t [(p, i)]++-- | Send messages to partition calculated by 'partitionAndCollate'.+produceMessages :: [TopicAndMessage] -> Kafka [ProduceResponse]+produceMessages tams = do+  m <- fmap (fmap groupMessagesToSet) <$> partitionAndCollate tams+  mapM (uncurry send) $ fmap M.toList <$> M.toList m++-- | Create a protocol message set from a list of messages.+groupMessagesToSet :: [TopicAndMessage] -> MessageSet+groupMessagesToSet xs = MessageSet $ msm <$> xs+    where msm = MessageSetMember (Offset (-1)) . _tamMessage++-- | Group messages together with the leader they should be sent to.+partitionAndCollate :: [TopicAndMessage] -> Kafka (M.Map Leader (M.Map TopicAndPartition [TopicAndMessage]))+partitionAndCollate ks = recurse ks M.empty+      where recurse [] accum = return accum+            recurse (x:xs) accum = do+              topicPartitionsList <- brokerPartitionInfo $ _tamTopic x+              pal <- getPartition topicPartitionsList+              let leader = maybe (Leader Nothing) _palLeader pal+                  tp = TopicAndPartition <$> pal ^? folded . palTopic <*> pal ^? folded . palPartition+                  b = M.singleton leader $ maybe M.empty (`M.singleton` [x]) tp+                  accum' = M.unionWith (M.unionWith (<>)) accum b+              recurse xs accum'++-- | Execute a produce request using the values in the state.+send :: Leader -> [(TopicAndPartition, MessageSet)] -> Kafka ProduceResponse+send l ts = do+  let s = kafkaClientState . stateBrokers . at l+      topicNames = map (_tapTopic . fst) ts+  broker <- findMetadataOrElse topicNames s (KafkaInvalidBroker l)+  requiredAcks <- use (kafkaClientState . stateRequiredAcks)+  requestTimeout <- use (kafkaClientState . stateRequestTimeout)+  withBrokerHandle broker $ \handle -> produce handle $ produceRequest requiredAcks requestTimeout ts++-- | Find a leader and partition for the topic.+brokerPartitionInfo :: TopicName -> Kafka [PartitionAndLeader]+brokerPartitionInfo t = do+  let s = kafkaClientState . stateTopicMetadata . at t+  tmd <- findMetadataOrElse [t] s KafkaFailedToFetchMetadata+  return $ pal <$> tmd ^. partitionsMetadata+    where pal d = PartitionAndLeader t (d ^. partitionId) (d ^. partitionMetadataLeader)++findMetadataOrElse :: [TopicName] -> Getting (Maybe a) KafkaClient (Maybe a) -> KafkaClientError -> Kafka a+findMetadataOrElse ts s err = do+  maybeFound <- use s+  case maybeFound of+    Just x -> return x+    Nothing -> do+      updateMetadatas ts+      maybeFound' <- use s+      case maybeFound' of+        Just x -> return x+        Nothing -> lift $ left $ err++getPartition :: [PartitionAndLeader] -> Kafka (Maybe PartitionAndLeader)+getPartition ps =+    liftIO $ (ps' ^?) . element <$> getStdRandom (randomR (0, length ps' - 1))+        where ps' = ps ^.. folded . filtered (has $ palLeader . leaderId . _Just)++-- * Messages++-- | Default: @1@+defaultMessageCrc :: Crc+defaultMessageCrc = 1++-- | Default: @0@+defaultMessageMagicByte :: MagicByte+defaultMessageMagicByte = 0++-- | Default: @Nothing@+defaultMessageKey :: Key+defaultMessageKey = Key Nothing++-- | Default: @0@+defaultMessageAttributes :: Attributes+defaultMessageAttributes = 0++-- | Construct a message from a string of bytes using default attributes.+makeMessage :: ByteString -> Message+makeMessage m = Message (defaultMessageCrc, defaultMessageMagicByte, defaultMessageAttributes, defaultMessageKey, Value (Just (KBytes m)))

@@ -95,10 +95,10 @@   deriving (Show, Eq, Serializable, Deserializable)  newtype MetadataResponse = MetadataResp { _metadataResponseFields :: ([Broker], [TopicMetadata]) } deriving (Show, Eq, Deserializable)-newtype Broker = Broker { _brokerFields :: (NodeId, Host, Port) } deriving (Show, Eq, Deserializable)+newtype Broker = Broker { _brokerFields :: (NodeId, Host, Port) } deriving (Show, Eq, Ord, Deserializable) newtype NodeId = NodeId { _nodeId :: Int32 } deriving (Show, Eq, Ord, Deserializable, Num)-newtype Host = Host { _hostKString :: KafkaString } deriving (Show, Eq, Deserializable, IsString)-newtype Port = Port { _portInt :: Int32 } deriving (Show, Eq, Deserializable, Num)+newtype Host = Host { _hostKString :: KafkaString } deriving (Show, Eq, Ord, Deserializable, IsString)+newtype Port = Port { _portInt :: Int32 } deriving (Show, Eq, Ord, Deserializable, Num) newtype TopicMetadata = TopicMetadata { _topicMetadataFields :: (KafkaError, TopicName, [PartitionMetadata]) } deriving (Show, Eq, Deserializable) newtype PartitionMetadata = PartitionMetadata { _partitionMetadataFields :: (KafkaError, Partition, Leader, Replicas, Isr) } deriving (Show, Eq, Deserializable) newtype Leader = Leader { _leaderId :: Maybe Int32 } deriving (Show, Eq, Ord)

@@ -4,7 +4,7 @@ -- PVP summary:      +-+------- breaking API changes --                   | | +----- non-breaking API additions --                   | | | +--- code changes with no API change-version:             0.1.0.0+version:             0.2.0.0 synopsis:            A Kafka client for Haskell. description:         Kafka client for Haskell (not recommended for production use). license:             BSD3@@ -24,24 +24,26 @@ source-repository this   type: git   location: https://github.com/tylerholien/milena.git-  tag: 0.1.0.0+  tag: 0.2.0.0  library   default-language: Haskell2010   ghc-options:         -Wall   exposed-modules:     Network.Kafka,                        Network.Kafka.Protocol-  build-depends:       base       >=4.7      && <5,-                       mtl        >=2.1      && <2.2,-                       cereal     >=0.4      && <0.5,-                       bytestring >=0.10     && <0.12,-                       network    >=2.4      && <2.7,-                       digest     >=0.0.1.0  && <0.0.2,-                       containers >=0.5      && <0.6,-                       either     >=4.3      && <4.4,-                       random     >=1.0      && <1.2,-                       transformers >=0.3    && <0.4,-                       lens       >=4.4      && <4.5+                       Network.Kafka.Producer+  build-depends:       base          >=4.7      && <5,+                       mtl           >=2.1      && <2.3,+                       cereal        >=0.4      && <0.5,+                       bytestring    >=0.10     && <0.11,+                       network       >=2.4      && <2.7,+                       digest        >=0.0.1.0  && <0.1,+                       containers    >=0.5      && <0.6,+                       either        >=4.3      && <4.4,+                       random        >=1.0      && <1.2,+                       transformers  >=0.3      && <0.5,+                       lens          >=4.4      && <4.8,+                       resource-pool >=0.2.3.2  && <0.3  test-suite test   default-language: Haskell2010

@@ -5,6 +5,7 @@ import Data.Functor import Data.Either (isRight) import Network.Kafka+import Network.Kafka.Producer import Network.Kafka.Protocol (Leader(..)) import Test.Hspec import Test.Hspec.QuickCheck@@ -13,7 +14,7 @@ main :: IO () main = hspec $ do   let topic = "milena-test"-      run = runKafka ("localhost", 9092) $ defaultState "milena-test-client"+      run = runKafka ("kafka1.dev", 9093) $ defaultState "milena-test-client"       byteMessages = fmap (TopicAndMessage topic . makeMessage . B.pack)    describe "can talk to local Kafka server" $ do