diff --git a/haskell/src/Data/MessagePack.hs b/haskell/src/Data/MessagePack.hs index 92353c53..b53066b1 100644 --- a/haskell/src/Data/MessagePack.hs +++ b/haskell/src/Data/MessagePack.hs @@ -45,7 +45,6 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Lazy as L import Data.Functor.Identity import qualified Data.Iteratee as I -import qualified Data.Iteratee.IO as I import System.IO import Data.MessagePack.Object @@ -111,7 +110,7 @@ unpackFromStringI bs = -- | Iteratee interface to unpack from Handle unpackFromHandleI :: CIO.MonadCatchIO m => Handle -> I.Iteratee B.ByteString m a -> m a unpackFromHandleI h = - I.run . I.joinIM . I.enumHandle bufferSize h + I.run . I.joinIM . enumHandleNonBlocking bufferSize h -- | Iteratee interface to unpack from File unpackFromFileI :: CIO.MonadCatchIO m => FilePath -> I.Iteratee B.ByteString m a -> m a diff --git a/haskell/src/Data/MessagePack/Iteratee.hs b/haskell/src/Data/MessagePack/Iteratee.hs index 789b714a..4258cf68 100644 --- a/haskell/src/Data/MessagePack/Iteratee.hs +++ b/haskell/src/Data/MessagePack/Iteratee.hs @@ -15,13 +15,18 @@ module Data.MessagePack.Iteratee( -- * Iteratee version of deserializer getI, + -- * Non Blocking Enumerator + enumHandleNonBlocking, -- * Convert Parser to Iteratee parserToIteratee, ) where +import Control.Exception +import Control.Monad.IO.Class import qualified Data.Attoparsec as A import qualified Data.ByteString as B import qualified Data.Iteratee as I +import System.IO import Data.MessagePack.Parser @@ -29,6 +34,37 @@ import Data.MessagePack.Parser getI :: (Monad m, ObjectGet a) => I.Iteratee B.ByteString m a getI = parserToIteratee get +-- | Enumerator +enumHandleNonBlocking :: MonadIO m => Int -> Handle -> I.Enumerator B.ByteString m a +enumHandleNonBlocking bufSize h = + I.enumFromCallback $ readSome bufSize h + +readSome :: MonadIO m => Int -> Handle -> m (Either SomeException (Bool, B.ByteString)) +readSome bufSize h = liftIO $ do + ebs <- try $ hGetSome bufSize h + case ebs of + Left exc -> + return $ Left (exc :: SomeException) + Right bs | B.null bs -> + return $ Right (False, B.empty) + Right bs -> + return $ Right (True, bs) + +hGetSome :: Int -> Handle -> IO B.ByteString +hGetSome bufSize h = do + bs <- B.hGetNonBlocking h bufSize + if B.null bs + then do + hd <- B.hGet h 1 + if B.null hd + then do + return B.empty + else do + rest <- B.hGetNonBlocking h (bufSize - 1) + return $ B.cons (B.head hd) rest + else do + return bs + -- | Convert Parser to Iteratee parserToIteratee :: Monad m => A.Parser a -> I.Iteratee B.ByteString m a parserToIteratee p = I.icont (itr (A.parse p)) Nothing