diff --git a/Network/Gearman/Client.hs b/Network/Gearman/Client.hs new file mode 100644 --- /dev/null +++ b/Network/Gearman/Client.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.Gearman.Client (connectGearman, submitJob, submitJobBg) where + +import qualified Control.Monad.State as S +import qualified Data.ByteString.Char8 as B +import Network.Socket hiding (send, sendTo, recv, recvFrom) +import qualified Data.HashMap.Strict as H +import qualified Data.Pool as Pool + +import Network.Gearman.Protocol +import Network.Gearman.Internal + +connectGearman :: B.ByteString -> HostName -> Port -> IO (Either GearmanError GearmanClient) +connectGearman i h p = do + addrInfo <- getAddrInfo Nothing (Just h) (Just $ show p) + + sk <- socket (addrFamily (head addrInfo)) Stream defaultProtocol + connect sk (addrAddress (head addrInfo)) + + pool <- mkPool addrInfo + + _ <- return $ writePacket sk (mkRequest SET_CLIENT_ID i) >> response sk + return $ Right $ GearmanClient sk pool (Just i) H.empty + + where mkPool :: [AddrInfo] -> IO (Pool.Pool Socket) + mkPool a = Pool.createPool (mkConn a) sClose 1 10 10 + + mkConn :: [AddrInfo] -> IO Socket + mkConn a = do + sk <- socket (addrFamily (head a)) Stream defaultProtocol + connect sk (addrAddress (head a)) + return sk + +submitJob :: Function -> B.ByteString -> Gearman (Either GearmanError B.ByteString) +submitJob f d = do + packet <- submit $ B.concat [f, "\0\0", d] + case packet of + Left p -> return $ Left p + Right p -> let _:r:_ = B.split '\0' $ _msg p in return $ Right r + +submit :: B.ByteString -> Gearman (Either GearmanError Packet) +submit req = S.get >>= \env -> Pool.withResource (_pool env) $ \s -> + writePacket s (mkRequest SUBMIT_JOB req) >> response s + +response :: Socket -> Gearman (Either GearmanError Packet) +response s = do + p <- readPacket s + case (_type . _hdr) p of + WORK_COMPLETE -> return $ Right p + WORK_DATA -> return $ Left "WORK_DATA NOT IMPLEMENTED" + WORK_STATUS -> return $ Left "WORK_STATUS NOT IMPLEMENTED" + WORK_EXCEPTION -> return $ Left "WORK_EXCEPTION NOT IMPLEMENTED" + WORK_FAIL -> return $ Left "WORK_FAIL NOT IMPLEMENTED" + WORK_WARNING -> return $ Left "WORK_WARNING NOT IMPLEMENTED" + _ -> response s + +submitJobBg :: Function -> B.ByteString -> Gearman () +submitJobBg f d = S.get >>= \env -> Pool.withResource (_pool env) $ \s -> + writePacket s (mkRequest SUBMIT_JOB_BG req) >> readPacket s >> return () + + where req = B.concat [f, "\0\0", d] + diff --git a/Network/Gearman/Internal.hs b/Network/Gearman/Internal.hs new file mode 100644 --- /dev/null +++ b/Network/Gearman/Internal.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.Gearman.Internal ( Port + , Function + , GearmanClient(..) + , GearmanError + , Gearman() + , withGearman + ) where + +import qualified Control.Monad.State as S +import qualified Data.ByteString as B +import Network.Socket hiding (send, sendTo, recv, recvFrom) +import qualified Data.HashMap.Strict as H +import qualified Data.Pool as Pool + +type Port = Int +type Function = B.ByteString + +data GearmanClient = GearmanClient { + _sock :: Socket + , _pool :: Pool.Pool Socket + , _id :: Maybe B.ByteString + , _fns :: H.HashMap B.ByteString (B.ByteString -> B.ByteString) +} + +type GearmanError = B.ByteString +type Gearman = S.StateT GearmanClient IO + +withGearman :: GearmanClient -> Gearman a -> IO a +withGearman env action = S.runStateT action env >>= (\(v,_) -> return v) diff --git a/Network/Gearman/Worker.hs b/Network/Gearman/Worker.hs new file mode 100644 --- /dev/null +++ b/Network/Gearman/Worker.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.Gearman.Worker(registerWorker, unregisterWorker, runWorker) where + +import Control.Concurrent +import Control.Monad +import qualified Control.Monad.State as S +import qualified Data.ByteString.Char8 as B +import qualified Data.HashMap.Strict as H + +import Network.Gearman.Protocol +import Network.Gearman.Internal + +runWorker :: GearmanClient -> Gearman () -> IO ThreadId +runWorker c f = forkIO $ withGearman c (f >> gmLoop) + +registerWorker :: Function -> (B.ByteString -> B.ByteString) -> Gearman () +registerWorker n f = S.get >>= \env -> do + let ht = H.insert n f (_fns env) + S.put env{ _fns = ht } + void $ writePacket (_sock env) (mkRequest CAN_DO n) + +unregisterWorker :: Function -> Gearman () +unregisterWorker n = S.get >>= \env -> do + let ht = H.delete n (_fns env) + S.put env { _fns = ht } + + void $ writePacket (_sock env) (mkRequest CANT_DO n) + +gmWait :: Gearman () +gmWait = S.get >>= \env -> do + packet <- readPacket (_sock env) + case (_type . _hdr) packet of + NO_JOB -> writePacket (_sock env) (mkRequest PRE_SLEEP "") >> gmWait + JOB_ASSIGN -> doJob (_msg packet) >>= \p -> void $ writePacket (_sock env) p + NOOP -> return () + _ -> gmWait + +gmLoop :: Gearman () +gmLoop = S.get >>= \env -> forever $ writePacket (_sock env) (mkRequest GRAB_JOB "") >> gmWait + +doJob :: B.ByteString -> Gearman Packet +doJob m = S.get >>= \env -> do + let [h,f,d] = B.split '\0' m + return $ case H.lookup f (_fns env) of + Just fn -> let resp = B.concat [h, "\0", fn d] in mkRequest WORK_COMPLETE resp + Nothing -> mkRequest WORK_FAIL "" diff --git a/hgearman.cabal b/hgearman.cabal --- a/hgearman.cabal +++ b/hgearman.cabal @@ -2,7 +2,7 @@ -- documentation, see http://haskell.org/cabal/users-guide/ name: hgearman -version: 0.1.0.0 +version: 0.1.0.1 synopsis: A Gearman client for Haskell. description: Gearman client for Haskell (not reccomended for production use). license: MIT @@ -34,5 +34,8 @@ ,resource-pool >= 0.2.3.2 && < 0.2.4.0 exposed-modules: Network.Gearman - --other-modules: Network.Gearman.Protocol + , Network.Gearman.Client + , Network.Gearman.Worker + , Network.Gearman.Internal + default-language: Haskell2010