haskell: nonblocking enumerator

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

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