diff --git a/Network/Beanstalk.hs b/Network/Beanstalk.hs --- a/Network/Beanstalk.hs +++ b/Network/Beanstalk.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveDataTypeable, OverloadedStrings #-} ----------------------------------------------------------------------------- -- | -- Module : Network.Beanstalk @@ -24,24 +24,30 @@ -- * Exception Predicates isNotFoundException, isBadFormatException, isTimedOutException, isOutOfMemoryException, isInternalErrorException, isJobTooBigException, - isDeadlineSoonException, isNotIgnoredException, + isDeadlineSoonException, isNotIgnoredException, isDrainingException, -- * Data Types Job(..), BeanstalkServer, JobState(..), BeanstalkException(..) ) where import Data.Bits -import Network.Socket -import Network.BSD +import Network.Socket hiding (send, sendTo, recv, recvFrom) +import Network.Socket.ByteString import Data.List import System.IO -import Text.ParserCombinators.Parsec -import Data.Yaml.Syck import Data.Typeable import qualified Data.Map as M import qualified Control.Exception as E import Data.Maybe import Control.Monad import Control.Concurrent.MVar +import Control.Applicative hiding (many) +import Data.Attoparsec as P +import qualified Data.Attoparsec.Char8 as P8 +import qualified Data.ByteString.Char8 as B +import Blaze.ByteString.Builder +import Blaze.ByteString.Builder.Char8 +import Blaze.ByteString.Builder.ByteString +import Data.Monoid (mappend) -- | Beanstalk Server, wrapped in an 'MVar' for synchronizing access -- to the server socket. As many of these can be created as are @@ -55,7 +61,7 @@ Job { -- | Job numeric identifier job_id :: Int, -- | Job body - job_body :: String} + job_body :: B.ByteString} deriving (Show, Read, Eq) -- | States describing the lifecycle of a job. @@ -184,17 +190,19 @@ -- 'buryJob', or 'touchJob' command being run, the job -- will be placed back on the ready queue. The minimum -- value is 1. - -> String -- ^ Job body. + -> B.ByteString -- ^ Job body. -> IO (JobState, Int) -- ^ State of the newly created job and its ID putJob bs priority delay ttr job_body = withMVar bs task where task s = - do let job_size = length job_body - send s ("put " ++ - (show priority) ++ " " ++ - (show delay) ++ " " ++ - (show ttr) ++ " " ++ - (show job_size) ++ "\r\n") - send s (job_body ++ "\r\n") + do let job_size = B.length job_body + sendAll s $ toByteString (fromByteString "put " `mappend` + fromShow priority `mappend` fromChar ' ' `mappend` + fromShow delay `mappend` fromChar ' ' `mappend` + fromShow ttr `mappend` fromChar ' ' `mappend` + fromShow job_size `mappend` + fromByteString "\r\n" `mappend` + fromByteString job_body `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response let (state, jobid) = parsePut response @@ -207,18 +215,18 @@ -> IO Job -- ^ Job reserved by this client reserveJob bs = withMVar bs task where task s = - do send s "reserve\r\n" + do sendAll s "reserve\r\n" response <- readLine s checkForBeanstalkErrors response let (jobid, bytes) = parseReserve response - (jobContent, bytesRead) <- recvLen s (bytes) - recv s 2 -- Ending CRLF - return (Job (read jobid) jobContent) + jobContent <- recvBytes s bytes + recvBytes s 2 -- Ending CRLF + return (Job jobid jobContent) -- | Reserve a job from the watched tube list, blocking for the specified number -- of seconds or until a job is returned. If no jobs are found before the --- timeout value, a TimedOutException will be thrown. If another reserved job --- is about to exceed its time-to-run, a DeadlineSoonException will be thrown. +-- timeout value, a 'TimedOutException' will be thrown. If another reserved job +-- is about to exceed its time-to-run, a 'DeadlineSoonException' will be thrown. reserveJobWithTimeout :: BeanstalkServer -- ^ Beanstalk server -> Int -- ^ Time in seconds to wait for a job -- to become available. Once this time @@ -227,13 +235,15 @@ -> IO Job -- ^ Job reserved by this client reserveJobWithTimeout bs seconds = withMVar bs task where task s = - do send s ("reserve-with-timeout "++(show seconds)++"\r\n") + do sendAll s $ toByteString (fromByteString "reserve-with-timeout " `mappend` + fromShow seconds `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response let (jobid, bytes) = parseReserve response - (jobContent, bytesRead) <- recvLen s (bytes) - recv s 2 -- Ending CRLF - return (Job (read jobid) jobContent) + jobContent <- recvBytes s bytes + recvBytes s 2 -- Ending CRLF + return (Job jobid jobContent) -- | Delete a job to indicate that it has been completed. If the job -- does not exist, was not reserved by this client, or is not in the @@ -243,7 +253,9 @@ -> IO () deleteJob bs jobid = withMVar bs task where task s = - do send s ("delete "++(show jobid)++"\r\n") + do sendAll s $ toByteString (fromByteString "delete " `mappend` + fromShow jobid `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response @@ -255,10 +267,10 @@ -> IO () releaseJob bs jobid priority delay = withMVar bs task where task s = - do send s ("release " ++ - (show jobid) ++ " " ++ - (show priority) ++ " " ++ - (show delay) ++ "\r\n") + do sendAll s $ toByteString (fromByteString "release " `mappend` + fromShow jobid `mappend` fromChar ' ' `mappend` + fromShow priority `mappend` fromChar ' ' `mappend` + fromShow delay `mappend` fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response @@ -269,11 +281,13 @@ -> IO () buryJob bs jobid pri = withMVar bs task where task s = - do send s ("bury "++(show jobid)++" "++(show pri)++"\r\n") + do sendAll s $ toByteString (fromByteString "bury " `mappend` + fromShow jobid `mappend` fromChar ' ' `mappend` + fromShow pri `mappend` fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response -checkForBeanstalkErrors :: String -- ^ beanstalkd server response +checkForBeanstalkErrors :: B.ByteString -- ^ beanstalkd server response -> IO () checkForBeanstalkErrors input = do eop OutOfMemoryException "OUT_OF_MEMORY\r\n" @@ -287,34 +301,32 @@ eop DeadlineSoonException "DEADLINE_SOON\r\n" eop TimedOutException "TIMED_OUT\r\n" eop NotIgnoredException "NOT_IGNORED\r\n" - where eop e s = exceptionOnParse e (parse (string s) "errorParser" input) - --- | When an error is successfully parsed, throw the given exception. -exceptionOnParse :: BeanstalkException -> Either a b -> IO () -exceptionOnParse e x = case x of - Right _ -> E.throwIO e - Left _ -> return () + where eop e s = if B.take (B.length s) input == s then E.throwIO e else return () -- | Assign a tube for new jobs created with put command. If the tube -- does not already exist, it will be created. Initially, all -- sessions will use the tube named \"default\". useTube :: BeanstalkServer -- ^ Beanstalk server - -> String -- ^ Name of tube to watch + -> B.ByteString -- ^ Name of tube to watch -> IO () useTube bs name = withMVar bs task where task s = - do send s ("use "++name++"\r\n"); + do sendAll s $ toByteString (fromByteString "use " `mappend` + fromByteString name `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response -- | Add a named tube to the watch list, those tubes which -- 'reserveJob' will request jobs from. watchTube :: BeanstalkServer -- ^ Beanstalk server - -> String -- ^ Name of tube to watch + -> B.ByteString -- ^ Name of tube to watch -> IO Int -- ^ Number of tubes currently being watched watchTube bs name = withMVar bs task where task s = - do send s ("watch "++name++"\r\n"); + do sendAll s $ toByteString (fromByteString "watch " `mappend` + fromByteString name `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response return $ parseWatching response @@ -323,11 +335,13 @@ -- is the only one currently being watched, a 'NotIgnoredException' -- is thrown. ignoreTube :: BeanstalkServer -- ^ Beanstalk server - -> String -- ^ Name of tube to ignore + -> B.ByteString -- ^ Name of tube to ignore -> IO Int -- ^ Number of tubes currently being watched ignoreTube bs name = withMVar bs task where task s = - do send s ("ignore "++name++"\r\n"); + do sendAll s $ toByteString (fromByteString "ignore " `mappend` + fromByteString name `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response return $ parseWatching response @@ -336,35 +350,35 @@ peekJob :: BeanstalkServer -- ^ Beanstalk server -> Int -- ^ ID of job to get information about -> IO Job -- ^ Job definition -peekJob bs jobid = genericPeek bs ("peek "++(show jobid)) +peekJob bs jobid = genericPeek bs (fromByteString "peek " `mappend` fromShow jobid) -- | Inspect the next ready job on the currently used tube. peekReadyJob :: BeanstalkServer -- ^ Beanstalk server -> IO Job -- ^ Job definition -peekReadyJob bs = genericPeek bs "peek-ready" +peekReadyJob bs = genericPeek bs (fromByteString "peek-ready") -- | Inspect the delayed job with shortest delay remaining on the currently used tube. peekDelayedJob :: BeanstalkServer -- ^ Beanstalk server -> IO Job -- ^ Job definition -peekDelayedJob bs = genericPeek bs "peek-delayed" +peekDelayedJob bs = genericPeek bs (fromByteString "peek-delayed") -- | Inspect the next buried job on the currently used tube. peekBuriedJob :: BeanstalkServer -- ^ Beanstalk server -> IO Job -- ^ Job definition -peekBuriedJob bs = genericPeek bs "peek-buried" +peekBuriedJob bs = genericPeek bs (fromByteString "peek-buried") -- Essence of the peek command. Variations (peek, peek-ready, -- peek-delayed, peek-buried) just provide the command string, while -- this function actually executes it and parses the results. -genericPeek :: BeanstalkServer -> String -> IO Job +genericPeek :: BeanstalkServer -> Builder -> IO Job genericPeek bs cmd = withMVar bs task where task s = - do send s (cmd++"\r\n") + do sendAll s $ toByteString (cmd `mappend` fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response let (jobid, bytes) = parseFoundIdLen response - (content,bytesRead) <- recvLen s (bytes) - recv s 2 -- Ending CRLF + content <- recvBytes s bytes + recvBytes s 2 -- Ending CRLF return (Job jobid content) -- | Update the Time-To-Run (TTR) value for a job, giving a worker more time before job expiry. @@ -373,7 +387,9 @@ -> IO () touchJob bs jobid = withMVar bs task where task s = - do send s ("touch "++(show jobid)++"\r\n") + do sendAll s $ toByteString (fromByteString "touch " `mappend` + fromShow jobid `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response return () @@ -386,31 +402,30 @@ -> IO Int -- ^ Number of jobs actually kicked kickJobs bs maxcount = withMVar bs task where task s = - do send s ("kick "++(show maxcount)++"\r\n") + do sendAll s $ toByteString (fromByteString "kick " `mappend` + fromShow maxcount `mappend` + fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response return (parseKicked response) -parseKicked :: String -> Int +parseKicked :: B.ByteString -> Int parseKicked input = case kparse of - Right a -> read a - Left _ -> 0 -- Error - where kparse = parse (string "KICKED " >> many1 digit) "KickedParser" input + Done _ x -> x + _ -> 0 -- Error + where kparse = parse (P.string "KICKED " *> P8.decimal) input -- Essence of the various stats commands. Variations provide the -- command, while this function actually executes it and parses the -- results. -genericStats :: BeanstalkServer -> String -> IO (M.Map String String) +genericStats :: BeanstalkServer -> Builder -> IO (M.Map B.ByteString B.ByteString) genericStats bs cmd = withMVar bs task where task s = - do send s (cmd++"\r\n") + do sendAll s $ toByteString (cmd `mappend` fromByteString "\r\n") statHeader <- readLine s checkForBeanstalkErrors statHeader let bytes = parseOkLen statHeader - (statContent, bytesRead) <- recvLen s (bytes) - recv s 2 -- Ending CRLF - yamlN <- parseYaml statContent - return $ yamlMapToHMap yamlN + readYamlDict s -- | Return statistical information about a job. Keys that can be -- expected to be returned are the following: @@ -441,8 +456,8 @@ -- See the Beanstalk protocol docs for the definitive list and definitions. statsJob :: BeanstalkServer -- ^ Beanstalk server -> Int -- ^ ID of job - -> IO (M.Map String String) -- ^ Key-value map of job statistics -statsJob bs jobid = genericStats bs ("stats-job "++(show jobid)) + -> IO (M.Map B.ByteString B.ByteString) -- ^ Key-value map of job statistics +statsJob bs jobid = genericStats bs (fromByteString "stats-job " `mappend` fromShow jobid) -- | Return statistical information about a tube. Keys that can be -- expected to be returned are the following: @@ -480,22 +495,23 @@ -- -- See the Beanstalk protocol docs for the definitive list and definitions. statsTube :: BeanstalkServer -- ^ Beanstalk server - -> String -- ^ Name of tube - -> IO (M.Map String String) -- ^ Key-value map of tube statistics -statsTube bs tube = genericStats bs ("stats-tube "++tube) + -> B.ByteString -- ^ Name of tube + -> IO (M.Map B.ByteString B.ByteString) -- ^ Key-value map of tube statistics +statsTube bs tube = + genericStats bs (fromByteString "stats-tube " `mappend` fromByteString tube) -- | Print stats to screen in a readable format. -printStats :: M.Map String String -- ^ Key-value map of statistic names and values +printStats :: M.Map B.ByteString B.ByteString -- ^ Key-value map of statistic names and values -> IO () -- ^ Screen output showing all \"key => value\" pairs printStats stats = do let kv = M.assocs stats - mapM_ (\(k,v) -> putStrLn (k ++ " => " ++ v)) kv + mapM_ (\(k,v) -> B.putStr k >> B.putStr " => " >> B.putStrLn v) kv -- | Pretty print a list. -printList :: [String] -- ^ List of names +printList :: [B.ByteString] -- ^ List of names -> IO () -- ^ Screen output showing results with prefixed counter printList list = - do mapM_ (\(n,x) -> putStrLn (" "++(show n)++". "++x)) (zip [1..] (list)) + do mapM_ (\(n,x) -> putStr (" "++(show n)++". ") >> B.putStrLn x) (zip [(1::Int)..] (list)) -- | Return statistical information about the server, across all -- clients. Keys that can be expected to be returned are the @@ -604,70 +620,66 @@ -- -- See the Beanstalk protocol docs for the definitive list and definitions. statsServer :: BeanstalkServer -- ^ Beanstalk server - -> IO (M.Map String String) -- ^ Key-value map of server statistics -statsServer bs = genericStats bs "stats" + -> IO (M.Map B.ByteString B.ByteString) -- ^ Key-value map of server statistics +statsServer bs = genericStats bs (fromByteString "stats") -- | Pause a tube for a specified time, so that reservations are no longer accepted. pauseTube :: BeanstalkServer -- ^ Beanstalk server - -> String -- ^ Name of tube to pause + -> B.ByteString -- ^ Name of tube to pause -> Int -- ^ Number of seconds before reservations are accepted again -> IO () pauseTube bs tube delay = withMVar bs task where task s = - do send s ("pause-tube "++tube++" "++(show delay)++"\r\n") + do sendAll s $ toByteString (fromByteString "pause-tube " `mappend` + fromByteString tube `mappend` + fromChar ' ' `mappend` + fromShow delay `mappend` fromByteString "\r\n") response <- readLine s checkForBeanstalkErrors response return () -- | List all existing tubes. listTubes :: BeanstalkServer -- ^ Beanstalk server - -> IO [String] -- ^ Names of all tubes on the server + -> IO [B.ByteString] -- ^ Names of all tubes on the server listTubes bs = genericList bs "list-tubes" -- | List all watched tubes. listTubesWatched :: BeanstalkServer -- ^ Beanstalk server - -> IO [String] -- ^ Names of all currently watched tubes + -> IO [B.ByteString] -- ^ Names of all currently watched tubes listTubesWatched bs = genericList bs "list-tubes-watched" -- | List used tube. listTubeUsed :: BeanstalkServer -- ^ Beanstalk server - -> IO String -- ^ Name of current used tube + -> IO B.ByteString -- ^ Name of current used tube listTubeUsed bs = withMVar bs task where task s = - do send s ("list-tube-used\r\n") + do sendAll s ("list-tube-used\r\n") response <- readLine s checkForBeanstalkErrors response let tubeName = parseUsedTube response return tubeName -parseUsedTube :: String -> String +parseUsedTube :: B.ByteString -> B.ByteString parseUsedTube input = - case (parse usedTubeParser "UsedTubeParser" input) of - Right x -> x - Left _ -> "" + case (parse usedTubeParser input) of + Done _ x -> x + _ -> "" -nameParser = do initial <- leadingNameParser - rest <- many1 (leadingNameParser <|> char '-') - return (initial : rest) - where leadingNameParser = alphaNum <|> oneOf "+/;.$_()" +nameParser = B.cons <$> leadingNameParser <*> followingNameParser + where leadingNameParser = P8.satisfy $ P8.inClass "a-zA-Z0-9+/;.$_()" + followingNameParser = P.takeWhile $ inClass "-a-zA-Z0-9+/;.$_()" -usedTubeParser = do string "USING " - tube <- nameParser - string "\r\n" - return tube +usedTubeParser = P.string "USING " *> nameParser <* P.string "\r\n" -- Essence of list commands that return YAML lists. -genericList :: BeanstalkServer -> String -> IO [String] +genericList :: BeanstalkServer -> B.ByteString -> IO [B.ByteString] genericList bs cmd = withMVar bs task where task s = - do send s (cmd++"\r\n") + do sendAll s (cmd `B.append` "\r\n") lHeader <- readLine s checkForBeanstalkErrors lHeader let bytes = parseOkLen lHeader - (content, bytesRead) <- recvLen s (bytes) - recv s 2 -- Ending CRLF - yamlN <- parseYaml content - return $ yamlListToHList yamlN + readYamlList s -- | Count number of jobs in a tube with a state in a given list. -- This is not part of the beanstalk protocol spec, so multiple @@ -675,116 +687,123 @@ -- may not be consistent (it does not represent one snapshot in -- time). jobCountWithState :: BeanstalkServer -- ^ Beanstalk server - -> String -- ^ Name of tube to inspect + -> B.ByteString -- ^ Name of tube to inspect -> [JobState] -- ^ List of valid states for count -> IO Int -- ^ Number of jobs with a state in the valid list jobCountWithState bs tube validStatuses = do ts <- statsTube bs tube let readyCount = case (elem READY validStatuses) of - True -> read (fromJust (M.lookup "current-jobs-ready" ts)) + True -> parseIntBS (fromJust (M.lookup "current-jobs-ready" ts)) False -> 0 let reservedCount = case (elem RESERVED validStatuses) of - True -> read (fromJust (M.lookup "current-jobs-reserved" ts)) + True -> parseIntBS (fromJust (M.lookup "current-jobs-reserved" ts)) False -> 0 let delayedCount = case (elem DELAYED validStatuses) of - True -> read (fromJust (M.lookup "current-jobs-delayed" ts)) + True -> parseIntBS (fromJust (M.lookup "current-jobs-delayed" ts)) False -> 0 let buriedCount = case (elem BURIED validStatuses) of - True -> read (fromJust (M.lookup "current-jobs-buried" ts)) + True -> parseIntBS (fromJust (M.lookup "current-jobs-buried" ts)) False -> 0 return (readyCount+reservedCount+delayedCount+buriedCount) -yamlListToHList :: YamlNode -> [String] -yamlListToHList y = elems where - elist = (n_elem y) - ESeq list = elist - yelems = map n_elem list - elems = map (\(EStr x) -> unpackBuf x) yelems - -yamlMapToHMap :: YamlNode -> M.Map String String -yamlMapToHMap y = M.fromList elems where - emap = (n_elem y) - EMap maplist = emap -- [(YamlNode,YamlNode)] - yelems = map (\(x,y) -> (n_elem x, n_elem y)) maplist - elems = map (\(EStr x, EStr y) -> (unpackBuf x, unpackBuf y)) yelems - --- Read a single character from socket without handling errors. -readChar :: Socket -> IO Char -readChar s = recv s 1 >>= return . head - -- Read up to and including a newline. Any errors result in a string -- starting with "Error: " -readLine :: Socket -> IO String +readLine :: Socket -> IO B.ByteString readLine s = - catch readLine' (\err -> return ("Error: " ++ show err)) + catch readLine' (\err -> return (B.pack $ "Error: " ++ show err)) where - readLine' = do c <- readChar s - if c == '\n' - then return (c:[]) - else do l <- readLine s - return (c:l) + readLine' = readline'' (fromByteString B.empty) >>= return . toByteString + where readline'' b = do c <- recvBytes s 1 + if B.head c == '\n' + then return (b `mappend` fromByteString c) + else readline'' (b `mappend` fromByteString c) -- Parse response from watch/ignore command to determine how many -- tubes are currently being watched. -parseWatching :: String -> Int +-- Parse response from watch/ignore command to determine how many +-- tubes are currently being watched. +parseWatching :: B.ByteString -> Int parseWatching input = - case (parse (string "WATCHING " >> many1 digit) "WatchParser" input) of - Right x -> read x - Left _ -> 0 + case (parse (P.string "WATCHING " *> P8.decimal) input) of + Done _ x -> x + _ -> 0 -- Parse response from put command. -parsePut :: String -> (JobState, Int) +parsePut :: B.ByteString -> (JobState, Int) parsePut input = - case (parse putParser "PutParser" input) of - Right x -> x - Left _ -> (READY, 0) -- Error + case (parse putParser input) of + Done _ x -> x + _ -> (READY, 0) -- Error -putParser = do stateStr <- many1 letter - char ' ' - jobid <- many1 digit - let state = case stateStr of - "BURIED" -> BURIED - _ -> READY - return (state, read jobid) +putParser = (,) <$> (parseStateStr <$> takeTill (==32) <* P8.space) <*> P8.decimal + where parseStateStr "BURIED" = BURIED + parseStateStr _ = READY -- Get Job ID and size. -parseReserve :: String -> (String,Int) +parseReserve :: B.ByteString -> (Int,Int) parseReserve input = - case (parse reservedParser "ReservedParser" input) of - Right (x,y) -> (x, read y) - Left _ -> ("",0) + case (parse reservedParser input) of + Done _ x -> x + _ -> (-1, 0) -- Error --- Parse response from reserve command, including job id and bytes of body --- to come next. -reservedParser :: GenParser Char st (String,String) -reservedParser = do string "RESERVED" - char ' ' - x <- many1 digit - char ' ' - y <- many1 digit - return (x,y) +reservedParser = P.string "RESERVED " *> ((,) <$> P8.decimal <* P8.space <*> P8.decimal) -- Get number of bytes from an OK response string. -parseOkLen :: String -> Int +parseOkLen :: B.ByteString -> Int parseOkLen input = - case (parse okLenParser "okLenParser" input) of - Right len -> read len - Left err -> 0 - --- Parser for first line of stats for data length indicator. -okLenParser :: GenParser Char st String -okLenParser = string "OK " >> many1 digit + case (parse (P.string "OK " *> P8.decimal) input) of + Done _ x -> x + _ -> 0 -- Get job id and number of bytes from FOUND response string. -parseFoundIdLen :: String -> (Int,Int) +parseFoundIdLen :: B.ByteString -> (Int,Int) parseFoundIdLen input = - case (parse foundIdLenParser "FoundIdLenParser" input) of - Right x -> x - Left _ -> (0,0) + case (parse foundIdLenParser input) of + Done _ x -> x + _ -> (0,0) -foundIdLenParser :: GenParser Char st (Int,Int) -foundIdLenParser = do string "FOUND " - jobid <- many1 digit - string " " - bytes <- many1 digit - return (read jobid, read bytes) +foundIdLenParser = P.string "FOUND " *> ((,) <$> P8.decimal <* P8.space <*> P8.decimal) + +-- Read a YAML list, and parse it. +readYamlList s = do result <- parseWith (recv s 1024) yamlListParser "" + case result of + Done _ x -> return x + _ -> return [] + +yamlListParser = start_line *> many list_item <* P.string "\r\n" + where start_line = P.string "---" *> P8.endOfLine + list_item = P.string "- " *> takeTill P8.isEndOfLine <* P8.endOfLine + +-- Read a YAML dict, parse it, and return a Map +readYamlDict s = do result <- parseWith (recv s 1024) yamlDictParser "" + case result of + Done _ x -> return x + _ -> return M.empty + +yamlDictParser = start_line *> (M.fromList <$> many dict_item) <* P.string "\r\n" + where start_line = P.string "---" *> P8.endOfLine + dict_item = (,) <$> dict_key <*> dict_value + dict_key = takeTill colon_or_newline <* P.string ": " + colon_or_newline byte = byte == 58 || P8.isEndOfLine byte + dict_value = takeTill P8.isEndOfLine <* P8.endOfLine + +-- Parse an integer ByteString +parseIntBS :: B.ByteString -> Int +parseIntBS input = + case feed (parse P8.decimal input) "" of + Done _ x -> x + _ -> error "Beanstalk.parseIntBS: no parse" + +-- Read bytes from a socket, and return the bytestring. Blocks until the given +-- number of bytes has been read. +recvBytes :: Socket -> Int -> IO B.ByteString +recvBytes s bytes = recv' (fromByteString "") bytes + where recv' b 0 = return $ toByteString b + recv' b n = do + chunk <- recv s (min n 1024) + let n' = n - B.length chunk + if n' == n + then ioError $ userError $ "Could not read " ++ (show bytes) ++ + " bytes from beanstalkd; server disconnect." + else + recv' (b `mappend` fromByteString chunk) n' diff --git a/Tests.hs b/Tests.hs --- a/Tests.hs +++ b/Tests.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE OverloadedStrings #-} ----------------------------------------------------------------------------- -- | -- Module : Beanstalk Tests @@ -25,6 +26,7 @@ import qualified Control.Exception as E import Control.Monad import Control.Concurrent(threadDelay) +import qualified Data.ByteString.Char8 as B bs_host = "localhost" bs_port = "11300" @@ -137,7 +139,7 @@ TestCase ( do (bs, tt) <- connectAndSelectRandomTube randString <- randomName - let body = "My test job body, " ++ randString + let body = "My test job body, " `B.append` randString (_,put_job_id) <- putJob bs 1 0 60 body rsv_job <- reserveJob bs assertEqual "Reserved job ID should match what was put" put_job_id (job_id rsv_job) @@ -151,7 +153,7 @@ TestCase ( do (bs, tt) <- connectAndSelectRandomTube randString <- randomName - let body = "My test job body, " ++ randString + let body = "My test job body, " `B.append` randString (_,put_job_id) <- putJob bs 1 0 60 body rsv_job <- reserveJobWithTimeout bs 2 assertEqual "Reserved job should match job that was just put" @@ -163,9 +165,9 @@ TestCase ( do (bs, tt) <- connectAndSelectRandomTube randString <- randomName - let body = "My test job body, " ++ randString + let body = "My test job body, " `B.append` randString (_,put_job_id) <- putJob bs 1 0 60 body - let next_body = "My test job body, " ++ randString + let next_body = "My test job body, " `B.append` randString (_,put_next_job_id) <- putJob bs 1 0 60 next_body peeked_job <- peekJob bs put_job_id assertEqual "Peeked job id should match job id that was just put" @@ -183,7 +185,7 @@ TestCase ( do (bs, tt) <- connectAndSelectRandomTube randString <- randomName - let body = "My test job body, " ++ randString + let body = "My test job body, " `B.append` randString (_,put_job_id) <- putJob bs 1 5 60 body kicked <- kickJobs bs 1 assertEqual "Kick should indicate one job kicked" 1 kicked @@ -196,7 +198,7 @@ assertJobsCount bs tt [READY] 0 "New tube has no jobs" -- Put a job on the tube randString <- randomName - let body = "My test job body, " ++ randString + let body = "My test job body, " `B.append` randString (_,put_job_id) <- putJob bs 1 0 60 body assertJobsCount bs tt [READY] 1 "Put adds job to tube" -- Reserve the job @@ -252,7 +254,7 @@ do (bs, tt) <- connectAndSelectRandomTube assertJobsCount bs tt [READY] 0 "New tube has no jobs" randString <- randomName - let jobcontent = "new job "++randString + let jobcontent = "new job " `B.append` randString (_,put_job_id) <- putJob bs 1 0 60 jobcontent assertJobsCount bs tt [READY] 1 "Put creates new ready job" job <- peekJob bs put_job_id @@ -296,8 +298,8 @@ let priority = 99 (job_state ,put_job_id) <- putJob bs priority 0 60 "new job" job_stats <- statsJob bs put_job_id - assertEqual "Job ID matches" put_job_id (read (fromJust (M.lookup "id" job_stats))) - assertEqual "Job priority matches" priority (read (fromJust (M.lookup "pri" job_stats))) + assertEqual "Job ID matches" put_job_id (read $ B.unpack (fromJust (M.lookup "id" job_stats))) + assertEqual "Job priority matches" priority (read $ B.unpack (fromJust (M.lookup "pri" job_stats))) ) -- Test finding server statistics. @@ -305,7 +307,7 @@ TestCase ( do (bs, tt) <- connectAndSelectRandomTube stats <- statsServer bs - assertBool "More than 1 job has been created" (1 < (read (fromJust (M.lookup "total-jobs" stats)))) + assertBool "More than 1 job has been created" (1 < (read $ B.unpack (fromJust (M.lookup "total-jobs" stats)))) ) -- Test listing all tubes for the server. @@ -347,8 +349,8 @@ jobstat_before <- statsJob bs jobid touchJob bs jobid jobstat_after <- statsJob bs jobid - let ttr_before = ((read (fromJust (M.lookup "time-left" jobstat_before)))::Int) - let ttr_after = ((read (fromJust (M.lookup "time-left" jobstat_after)))::Int) + let ttr_before = ((read $ B.unpack (fromJust (M.lookup "time-left" jobstat_before)))::Int) + let ttr_after = ((read $ B.unpack (fromJust (M.lookup "time-left" jobstat_after)))::Int) assertBool "TTR extended by touch" (ttr_after >= ttr_before) ) @@ -359,7 +361,7 @@ tubestat_before <- statsTube bs tt pauseTube bs tt 1000 tubestat_after <- statsTube bs tt - let paused_rem = ((read (fromJust (M.lookup "pause-time-left" tubestat_after)))::Int) + let paused_rem = ((read $ B.unpack (fromJust (M.lookup "pause-time-left" tubestat_after)))::Int) -- Check that at least 990 seconds still remains of the -- original 1000 seconds we paused the tube for. assertBool "Tube has at least 990 seconds before un-pausing" (paused_rem > 990) @@ -380,7 +382,7 @@ TestCase ( do (bs, tt) <- connectAndSelectRandomTube rname <- randomName - e <- E.tryJust (guard . isBadFormatException) (statsTube bs ("-"++rname)) + e <- E.tryJust (guard . isBadFormatException) (statsTube bs ("-" `B.append` rname)) case e of Right _ -> assertFailure "Using tube name starting with hyphen should fail" Left _ -> return () @@ -448,7 +450,7 @@ -- Assert a number of jobs on a given tube with one of the states -- listed. -assertJobsCount :: BeanstalkServer -> String -> [JobState] -> Int -> String -> IO () +assertJobsCount :: BeanstalkServer -> B.ByteString -> [JobState] -> Int -> String -> IO () assertJobsCount bs tube states jobs msg = do ts <- statsTube bs tube jobsReady <- jobCountWithState bs tube states @@ -457,7 +459,7 @@ -- Configure a new beanstalkd connection to use&watch a single tube -- with a random name. -connectAndSelectRandomTube :: IO (BeanstalkServer, String) +connectAndSelectRandomTube :: IO (BeanstalkServer, B.ByteString) connectAndSelectRandomTube = do bs <- connectBeanstalk bs_host bs_port tt <- randomName @@ -467,7 +469,7 @@ return (bs, tt) -- Generate random tube names for test separation. -randomName :: IO String +randomName :: IO B.ByteString randomName = do rdata <- randomIO :: IO Integer - return (show (abs rdata)) + return $ B.pack (show (abs rdata)) diff --git a/hbeanstalk.cabal b/hbeanstalk.cabal --- a/hbeanstalk.cabal +++ b/hbeanstalk.cabal @@ -1,5 +1,5 @@ Name: hbeanstalk -Version: 0.1 +Version: 0.2 License: BSD3 License-file: LICENSE Cabal-Version: >= 1.6 @@ -24,7 +24,9 @@ Library - Build-depends: base >= 4 && < 5, network, containers >= 0.3.0.0, HsSyck >= 0.45, parsec >= 2.1.0.1 + Build-depends: base >= 4 && < 5, network, containers >= 0.3.0.0, + blaze-builder >= 0.2.1.0, network-bytestring >= 0.1.2.0, + bytestring >= 0.9.1.7, attoparsec >= 0.8.3.0 Exposed-modules: Network.Beanstalk