diff --git a/hadoop-rpc.cabal b/hadoop-rpc.cabal --- a/hadoop-rpc.cabal +++ b/hadoop-rpc.cabal @@ -1,5 +1,5 @@ name: hadoop-rpc -version: 0.1.0.0 +version: 0.1.1.0 synopsis: Use the Hadoop RPC interface from Haskell. @@ -30,6 +30,7 @@ exposed-modules: Data.Hadoop.Configuration + Data.Hadoop.HdfsPath Data.Hadoop.Protobuf.ClientNameNode Data.Hadoop.Protobuf.DataTransfer Data.Hadoop.Protobuf.Hdfs @@ -49,9 +50,11 @@ , bytestring >= 0.10 , cereal >= 0.4 , exceptions >= 0.6 + , hashable >= 1.2.1 , network >= 2.5 , protobuf >= 0.2.0.4 , socks >= 0.5 + , stm >= 2.4 , text >= 1.1 , transformers >= 0.4 , unix >= 2.7 diff --git a/src/Data/Hadoop/HdfsPath.hs b/src/Data/Hadoop/HdfsPath.hs new file mode 100644 --- /dev/null +++ b/src/Data/Hadoop/HdfsPath.hs @@ -0,0 +1,26 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Data.Hadoop.HdfsPath + ( HdfsPath + , () + , combine + ) where + +import qualified Data.ByteString.Char8 as B +import Data.Monoid ((<>)) + +import Data.Hadoop.Types + +------------------------------------------------------------------------ + +infixr 5 + +() :: HdfsPath -> HdfsPath -> HdfsPath +() = combine + +combine :: HdfsPath -> HdfsPath -> HdfsPath +combine xs ys | B.null xs = ys + | B.null ys = xs + | B.head ys == '/' = ys + | B.last xs == '/' = xs <> ys + | otherwise = xs <> "/" <> ys diff --git a/src/Data/Hadoop/Protobuf/ClientNameNode.hs b/src/Data/Hadoop/Protobuf/ClientNameNode.hs --- a/src/Data/Hadoop/Protobuf/ClientNameNode.hs +++ b/src/Data/Hadoop/Protobuf/ClientNameNode.hs @@ -8,7 +8,6 @@ import Data.ProtocolBuffers import Data.Text (Text) import Data.Int (Int64) -import Data.Word (Word64) import GHC.Generics (Generic) import Data.Hadoop.Protobuf.Hdfs diff --git a/src/Data/Hadoop/Types.hs b/src/Data/Hadoop/Types.hs --- a/src/Data/Hadoop/Types.hs +++ b/src/Data/Hadoop/Types.hs @@ -7,6 +7,7 @@ import Data.Data (Data) import Data.Text (Text) import Data.Typeable (Typeable) +import qualified Data.Vector as V import Data.Word (Word16, Word64) ------------------------------------------------------------------------ @@ -46,6 +47,20 @@ ------------------------------------------------------------------------ +data DecodeError = DecodeError !Text + deriving (Show, Eq, Data, Typeable) + +instance Exception DecodeError + +------------------------------------------------------------------------ + +data ConnectionClosed = ConnectionClosed + deriving (Show, Eq, Data, Typeable) + +instance Exception ConnectionClosed + +------------------------------------------------------------------------ + type User = Text type Group = Text type Permission = Word16 @@ -68,6 +83,12 @@ , csQuota :: !Word64 , csSpaceConsumed :: !Word64 , csSpaceQuota :: !Word64 + } deriving (Eq, Ord, Show) + +-- | Partial directory listing. +data PartialListing = PartialListing + { lsRemaining :: !Int -- ^ number of files left to fetch + , lsFiles :: !(V.Vector FileStatus) } deriving (Eq, Ord, Show) -- | Status of a file, directory or symbolic link. diff --git a/src/Network/Hadoop/Hdfs.hs b/src/Network/Hadoop/Hdfs.hs --- a/src/Network/Hadoop/Hdfs.hs +++ b/src/Network/Hadoop/Hdfs.hs @@ -16,6 +16,7 @@ , getListing , getListing' + , getListingRecursive , getFileInfo , getContentSummary , mkdirs @@ -25,13 +26,17 @@ ) where import Control.Applicative (Applicative(..), (<$>)) -import Control.Exception (throw) -import Control.Monad (ap) +import Control.Concurrent.STM +import Control.Exception (SomeException(..), throw) +import Control.Monad (ap, when) import Control.Monad.Catch (MonadMask(..), MonadThrow(..), MonadCatch(..)) import Control.Monad.IO.Class (MonadIO(..)) import Data.ByteString (ByteString) +import qualified Data.ByteString.Char8 as B import Data.Maybe (fromMaybe) +import Data.Monoid ((<>)) import Data.Text (Text) +import qualified Data.Text as T import qualified Data.Text.Encoding as T import qualified Data.Vector as V import Data.Word (Word32) @@ -42,6 +47,7 @@ import Data.ProtocolBuffers.Orphans () import Data.Hadoop.Configuration +import Data.Hadoop.HdfsPath import Data.Hadoop.Types import Network.Hadoop.Rpc import qualified Network.Hadoop.Socket as S @@ -58,8 +64,8 @@ (<*>) = ap instance Monad Hdfs where - return x = Hdfs $ \_ -> return x - m >>= k = Hdfs $ \c -> unHdfs m c >>= \x -> unHdfs (k x) c + return = pure + m >>= k = Hdfs $ \c -> unHdfs m c >>= \x -> unHdfs (k x) c instance MonadIO Hdfs where liftIO io = Hdfs $ const io @@ -112,50 +118,104 @@ ------------------------------------------------------------------------ +getListingRecursive :: HdfsPath + -> Hdfs (TBQueue (Maybe (HdfsPath, Either SomeException (V.Vector FileStatus)))) +getListingRecursive initialPath = do + conn <- getConnection + queue <- liftIO (newTBQueueIO 10) + outstanding <- liftIO (newTVarIO 0) + liftIO (getListingRecursive' conn outstanding queue initialPath) + return queue + +getListingRecursive' :: Connection + -> TVar Int + -> TBQueue (Maybe (HdfsPath, Either SomeException (V.Vector FileStatus))) + -> HdfsPath + -> IO () +getListingRecursive' conn outstanding queue rootPath = do + getInitial rootPath + where + getInitial path = getPartial path B.empty + + getPartial path startAfter = do + atomically $ modifyTVar' outstanding succ + getPartialAsync conn path startAfter (onPartial path) + + onPartial path (Left err) = enqueueResult path (Left err) + onPartial path (Right PartialListing{..}) = do + when (lsRemaining /= 0) $ + getPartial path (lastFileName lsFiles) + + V.mapM_ getInitial $ V.map (\x -> path x) $ dirs lsFiles + + enqueueResult path (Right lsFiles) + + enqueueResult path result = atomically $ do + writeTBQueue queue $ Just (path, result) + modifyTVar' outstanding pred + n <- readTVar outstanding + when (n == 0) (writeTBQueue queue Nothing) + + dirs :: V.Vector FileStatus -> V.Vector HdfsPath + dirs = V.map fsPath . V.filter ((Dir ==) . fsFileType) + +getPartialAsync :: Connection + -> HdfsPath + -> HdfsPath + -> (Either SomeException PartialListing -> IO ()) + -> IO () +getPartialAsync c path startAfter k = invokeAsync c "getListing" request k' + where + k' :: Either SomeException P.GetListingResponse -> IO () + k' (Left err) = k (Left err) + k' (Right proto) = k (fromProto proto) + + fromProto :: P.GetListingResponse -> Either SomeException PartialListing + fromProto dl = case getField (P.lsDirList dl) of + Nothing -> Left notExist + Just x -> Right (fromProtoDirectoryListing x) + + notExist :: SomeException + notExist = SomeException $ RemoteError ("Directory does not exist: " <> T.decodeUtf8 path) T.empty + + request = P.GetListingRequest + { P.lsSrc = putField (T.decodeUtf8 path) + , P.lsStartAfter = putField startAfter + , P.lsNeedLocation = putField False + } + +------------------------------------------------------------------------ + getListing :: HdfsPath -> Hdfs (Maybe (V.Vector FileStatus)) getListing path = do mDirList <- getPartialListing path "" case mDirList of - Nothing -> return Nothing - Just dirList -> do - let p = partialListing dirList - if hasRemainingEntries dirList - then Just <$> loop [p] (lastFileName p) - else return (Just p) + Nothing -> return Nothing + Just (PartialListing 0 fs) -> return (Just fs) + Just (PartialListing _ fs) -> Just <$> loop [fs] (lastFileName fs) where - partialListing :: P.DirectoryListing -> V.Vector FileStatus - partialListing = V.map fromProtoFileStatus - . V.fromList - . getField - . P.dlPartialListing - - hasRemainingEntries :: P.DirectoryListing -> Bool - hasRemainingEntries = (/= 0) . getField . P.dlRemaingEntries - - lastFileName :: V.Vector FileStatus -> ByteString - lastFileName v | V.null v = "" - | otherwise = fsPath . V.last $ v - loop :: [V.Vector FileStatus] -> ByteString -> Hdfs (V.Vector FileStatus) loop ps startAfter = do - dirList <- fromMaybe emptyListing <$> getPartialListing path startAfter - - let p = partialListing dirList - ps' = ps ++ [p] + PartialListing{..} <- fromMaybe (PartialListing 0 V.empty) + <$> getPartialListing path startAfter - if hasRemainingEntries dirList - then loop ps' (lastFileName p) - else return (V.concat ps') + let ps' = ps ++ [lsFiles] - emptyListing = P.DirectoryListing (putField []) (putField 0) + if lsRemaining == 0 + then return (V.concat ps') + else loop ps' (lastFileName lsFiles) getListing' :: HdfsPath -> Hdfs (V.Vector FileStatus) getListing' path = fromMaybe V.empty <$> getListing path +lastFileName :: V.Vector FileStatus -> ByteString +lastFileName v | V.null v = "" + | otherwise = fsPath (V.last v) + ------------------------------------------------------------------------ -getPartialListing :: HdfsPath -> ByteString -> Hdfs (Maybe P.DirectoryListing) -getPartialListing path startAfter = getField . P.lsDirList <$> +getPartialListing :: HdfsPath -> HdfsPath -> Hdfs (Maybe PartialListing) +getPartialListing path startAfter = fmap fromProtoDirectoryListing . getField . P.lsDirList <$> hdfsInvoke "getListing" P.GetListingRequest { P.lsSrc = putField (T.decodeUtf8 path) , P.lsStartAfter = putField startAfter @@ -223,19 +283,25 @@ , csSpaceQuota = getField $ P.csSpaceQuota p } +fromProtoDirectoryListing :: P.DirectoryListing -> PartialListing +fromProtoDirectoryListing p = PartialListing + { lsRemaining = fromIntegral . getField $ P.dlRemaingEntries p + , lsFiles = V.map fromProtoFileStatus . V.fromList . getField $ P.dlPartialListing p + } + fromProtoFileStatus :: P.FileStatus -> FileStatus fromProtoFileStatus p = FileStatus - { fsFileType = fromProtoFileType $ getField $ P.fsFileType p + { fsFileType = fromProtoFileType . getField $ P.fsFileType p , fsPath = getField $ P.fsPath p , fsLength = getField $ P.fsLength p - , fsPermission = fromIntegral $ getField $ P.fpPerm $ getField $ P.fsPermission p + , fsPermission = fromIntegral . getField . P.fpPerm . getField $ P.fsPermission p , fsOwner = getField $ P.fsOwner p , fsGroup = getField $ P.fsGroup p , fsModificationTime = getField $ P.fsModificationTime p , fsAccessTime = getField $ P.fsAccessTime p , fsSymLink = getField $ P.fsSymLink p - , fsBlockReplication = fromIntegral $ fromMaybe 0 $ getField $ P.fsBlockReplication p - , fsBlockSize = fromMaybe 0 $ getField $ P.fsBlockSize p + , fsBlockReplication = fromIntegral . fromMaybe 0 . getField $ P.fsBlockReplication p + , fsBlockSize = fromMaybe 0 . getField $ P.fsBlockSize p } fromProtoFileType :: P.FileType -> FileType diff --git a/src/Network/Hadoop/Rpc.hs b/src/Network/Hadoop/Rpc.hs --- a/src/Network/Hadoop/Rpc.hs +++ b/src/Network/Hadoop/Rpc.hs @@ -11,16 +11,20 @@ , RawResponse , initConnectionV7 + , invokeAsync , invoke ) where import Control.Applicative ((<$>), (<*>)) -import Control.Exception (throwIO) -import Control.Monad.IO.Class (liftIO) - +import Control.Concurrent (ThreadId, forkIO, newEmptyMVar, putMVar, takeMVar) +import Control.Concurrent.STM +import Control.Exception (SomeException(..), throwIO, handle) +import Control.Monad (forever, when) import Data.ByteString (ByteString) import qualified Data.ByteString.Char8 as B -import Data.IORef +import qualified Data.HashMap.Strict as H +import Data.Hashable (Hashable) +import Data.Maybe (fromMaybe, isNothing) import Data.Monoid (mempty) import Data.Text (Text) import qualified Data.Text as T @@ -39,10 +43,10 @@ ------------------------------------------------------------------------ data Connection = Connection - { cnVersion :: !Int - , cnConfig :: !HadoopConfig - , cnProtocol :: !Protocol - , invokeRaw :: !(Method -> RawRequest -> IO RawResponse) + { cnVersion :: !Int + , cnConfig :: !HadoopConfig + , cnProtocol :: !Protocol + , invokeRaw :: !(Method -> RawRequest -> (RawResponse -> IO ()) -> IO ()) } data Protocol = Protocol @@ -52,17 +56,28 @@ type Method = Text type RawRequest = ByteString -type RawResponse = ByteString +type RawResponse = Either SomeException ByteString +type CallId = Int + ------------------------------------------------------------------------ -- hadoop-2.1.0-beta is on version 9 -- see https://issues.apache.org/jira/browse/HADOOP-8990 for differences +data ConnectionState = ConnectionState + { csStream :: !S.Stream + , csCallId :: !(TVar CallId) + , csRecvCallbacks :: !(TVar (H.HashMap CallId (RawResponse -> IO ()))) + , csSendQueue :: !(TQueue (Method, RawRequest, RawResponse -> IO ())) + , csFatalError :: !(TVar (Maybe SomeException)) + } + initConnectionV7 :: HadoopConfig -> Protocol -> Socket -> IO Connection initConnectionV7 config@HadoopConfig{..} protocol sock = do - stream <- S.mkSocketStream sock - S.runPut stream $ do + csStream <- S.mkSocketStream sock + + S.runPut csStream $ do putByteString "hrpc" putWord8 7 -- version putWord8 80 -- auth method (80 = simple, 81 = kerberos/gssapi, 82 = token/digest-md5) @@ -72,26 +87,73 @@ putWord32be (fromIntegral (B.length bs)) putByteString bs - ref <- newIORef 0 - return (Connection 7 config protocol (sendAndWait stream ref)) + csCallId <- newTVarIO 0 + csRecvCallbacks <- newTVarIO H.empty + csSendQueue <- newTQueueIO + csFatalError <- newTVarIO Nothing + + let cs = ConnectionState{..} + + _ <- forkSend cs + _ <- forkRecv cs + + return (Connection 7 config protocol (enqueue cs)) where - sendAndWait :: S.Stream -> IORef Int -> Method -> ByteString -> IO ByteString - sendAndWait stream ref method requestBytes = do - callId <- atomicModifyIORef' ref (\x -> (succ x, x)) + enqueue :: ConnectionState + -> Method + -> RawRequest + -> (RawResponse -> IO ()) + -> IO () + enqueue ConnectionState{..} method bs k = do + merr <- atomically $ do + merr <- readTVar csFatalError + when (isNothing merr) $ writeTQueue csSendQueue (method, bs, k) + return merr + case merr of + Just err -> throwIO err + Nothing -> return () - S.runPut stream $ do - let bs = runPut $ encodeLengthPrefixedMessage (requestHeader callId) - >> encodeLengthPrefixedMessage (request method requestBytes) + forkSend :: ConnectionState -> IO ThreadId + forkSend cs@ConnectionState{..} = forkIO $ handle (onSocketError cs) $ forever $ do + bs <- atomically $ do + (method, requestBytes, k) <- readTQueue csSendQueue + callId <- readTVar csCallId + modifyTVar' csCallId succ + modifyTVar' csRecvCallbacks (H.insert callId k) + + return $ runPut $ encodeLengthPrefixedMessage (requestHeaderProto callId) + >> encodeLengthPrefixedMessage (requestProto method requestBytes) + + S.runPut csStream $ do putWord32be (fromIntegral (B.length bs)) putByteString bs - responseHdr <- S.maybeGet stream decodeLengthPrefixedMessage - case getField . rspStatus <$> responseHdr of - Just Success -> S.runGet stream getResponse - Just _ -> S.runGet stream getError >>= liftIO . throwIO - Nothing -> throwClosed + forkRecv :: ConnectionState -> IO ThreadId + forkRecv cs@ConnectionState{..} = forkIO $ handle (onSocketError cs) $ forever $ do + hdr <- S.maybeGet csStream decodeLengthPrefixedMessage + case hdr of + Nothing -> throwIO ConnectionClosed + Just rspHdr -> do + onResponse <- fromMaybe (return $ return ()) + <$> lookupDelete csRecvCallbacks (fromIntegral $ getField $ rspCallId rspHdr) + case getField (rspStatus rspHdr) of + Success -> S.runGet csStream getResponse >>= onResponse . Right + _ -> S.runGet csStream getError >>= onResponse . Left . SomeException + onSocketError :: ConnectionState -> SomeException -> IO () + onSocketError ConnectionState{..} ex = do + ks <- atomically $ do + writeTVar csFatalError (Just ex) + sks <- map (\(_,_,k) -> k) <$> unfoldM (tryReadTQueue csSendQueue) + rks <- H.elems <$> readTVar csRecvCallbacks + return (sks ++ rks) + + mapM_ (\k -> handle ignore $ k $ Left ex) ks + + ignore :: SomeException -> IO () + ignore _ = return () + context = IpcConnectionContext { ctxProtocol = putField (Just (prName protocol)) , ctxUserInfo = putField (Just UserInformation @@ -100,21 +162,34 @@ }) } - requestHeader callId = RpcRequestHeader + requestHeaderProto callId = RpcRequestHeader { reqKind = putField (Just ProtocolBuffer) , reqOp = putField (Just FinalPacket) , reqCallId = putField (fromIntegral callId) } - request method bytes = RpcRequest + requestProto method bytes = RpcRequest { reqMethodName = putField method , reqBytes = putField (Just bytes) , reqProtocolName = putField (prName protocol) , reqProtocolVersion = putField (fromIntegral (prVersion protocol)) } - throwClosed = throwIO (RemoteError "ConnectionClosed" "The socket connection was closed") +unfoldM :: Monad m => m (Maybe a) -> m [a] +unfoldM f = go [] + where + go xs = do + m <- f + case m of + Nothing -> return xs + Just x -> go (xs ++ [x]) +lookupDelete :: (Eq k, Hashable k) => TVar (H.HashMap k v) -> k -> IO (Maybe v) +lookupDelete var k = atomically $ do + hm <- readTVar var + writeTVar var (H.delete k hm) + return (H.lookup k hm) + ------------------------------------------------------------------------ getResponse :: Get ByteString @@ -132,10 +207,27 @@ ------------------------------------------------------------------------ invoke :: (Decode b, Encode a) => Connection -> Text -> a -> IO b -invoke Connection{..} method arg = decodeBytes =<< invokeRaw method (encodeBytes arg) +invoke connection method arg = do + mv <- newEmptyMVar + invokeAsync connection method arg (putMVar mv) + e <- takeMVar mv + case e of + Left ex -> throwIO ex + Right x -> return x + +invokeAsync :: (Decode b, Encode a) => Connection -> Text -> a -> (Either SomeException b -> IO ()) -> IO () +invokeAsync Connection{..} method arg k = invokeRaw method (encodeBytes arg) k' where - encodeBytes = runPut . encodeMessage - decodeBytes bs = case runGetState decodeMessage bs 0 of - Left err -> throwIO (RemoteError "DecodeError" (T.pack err)) - Right (x, "") -> return x - Right (_, _) -> throwIO (RemoteError "DecodeError" "decoded response but did not consume enough bytes") + k' (Left err) = k (Left err) + k' (Right bs) = k (decodeBytes bs) + +encodeBytes :: Encode a => a -> ByteString +encodeBytes = runPut . encodeMessage + +decodeBytes :: Decode a => ByteString -> Either SomeException a +decodeBytes bs = case runGetState decodeMessage bs 0 of + Left err -> decodeError (T.pack err) + Right (x, "") -> Right x + Right (_, _) -> decodeError "decoded response but did not consume enough bytes" + where + decodeError = Left . SomeException . DecodeError