hadoop-rpc-1.1.0.0: src/Network/Hadoop/Rpc.hs
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Network.Hadoop.Rpc
( Connection(..)
, Protocol(..)
, User
, Method
, RawRequest
, RawResponse
, initConnectionV9
, invokeAsync
, invoke
) where
import Control.Applicative
import Control.Concurrent (ThreadId, forkIO, newEmptyMVar, putMVar, takeMVar)
import Control.Concurrent.STM
import Control.Exception (SomeException(..), throwIO, handle)
import Control.Monad (forever, when)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as L
import qualified Data.HashMap.Strict as H
import Data.Hashable (Hashable)
import Data.Maybe (fromMaybe, isNothing)
import Data.Monoid
import Data.ProtocolBuffers
import Data.ProtocolBuffers.Orphans ()
import Data.Serialize.Get
import Data.Serialize.Put
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.UUID as UUID
import qualified Network.Protocol.SASL.GNU as GS
import Network.Socket (Socket)
import System.Random (randomIO)
import Prelude
import Data.Hadoop.Configuration
import qualified Data.Hadoop.Protobuf.Headers as P
import Data.Hadoop.Types
import Network.Hadoop.Sasl
import qualified Network.Hadoop.Stream as S
import Network.Hadoop.Types
type CallId = Int
newtype ClientId = ClientId { unClientId :: ByteString }
------------------------------------------------------------------------
-- 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
, csClientId :: !ClientId
, csCallId :: !(TVar CallId)
, csRecvCallbacks :: !(TVar (H.HashMap CallId (RawResponse -> IO ())))
, csSendQueue :: !(TQueue (Method, RawRequest, RawResponse -> IO ()))
, csFatalError :: !(TVar (Maybe SomeException))
}
initConnectionV9 :: HadoopConfig -> NameNode -> Protocol -> Socket -> IO Connection
initConnectionV9 config@HadoopConfig{..} nameNode protocol sock = do
csStream <- S.mkSocketStream sock
csClientId <- mkClientId
S.runPut csStream $ do
putByteString "hrpc"
putWord8 9 -- version
putWord8 0 -- rpc service class (0 = default/protobuf, 1 = built-in, 2 = writable, 3 = protobuf
-- auth protocol (0 = none, -33/0xDF = sasl)
putWord8 $ maybe 0 (const 0xDF) (nnPrincipal nameNode)
case nnPrincipal nameNode of
Nothing -> return ()
Just principal -> saslAuth principal csStream
S.runPut csStream $
putMessage $ delimitedBytesL (rpcRequestHeaderProto csClientId (-3))
<> delimitedBytesL (contextProto protocol hcUser)
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
invokeSasl :: (Encode a, Decode b) => S.Stream -> a -> IO b
invokeSasl stream msgr = do
let message = delimitedBytesL (rpcRequestHeaderProto (ClientId "") (-33))
<> delimitedBytesL msgr
S.runPut stream (putMessage message)
mget <- S.maybeGet stream $ do
n <- fromIntegral <$> getWord32be
-- TODO Would be nice if we didn't have to isolate here
-- TODO and could stream instead. We could stream if we
-- TODO were able to read the varint length prefix
-- TODO ourselves and keep track of how many bytes were
-- TODO remaining instead of calling `getRemaining`.
isolate n $ do
hdr <- decodeLengthPrefixedMessage
msg <- case getField (P.rspStatus hdr) of
P.Success -> Right <$> getRemaining
_ -> return . Left . SomeException $ rspError hdr
return (hdr, msg)
case mget of
Nothing -> throwIO ConnectionClosed
Just (_, msg) ->
case msg of
Left ex -> throwIO ex
Right m -> case fromDelimitedBytes m of
Left ex -> throwIO ex
Right x -> return x
saslAuth :: Principal -> S.Stream -> IO ()
saslAuth Principal{..} c = do
negResp <- liftIO $ invokeSasl c RpcSaslProto
{ spState = putField Negotiate
, spVersion = mempty
, spToken = mempty
, spAuths = mempty
}
case getField $ spState negResp of
Negotiate -> return ()
_ -> throwIO $ SomeException FailedNegotiation
let auths = getField $ spAuths negResp
authsMap = zip (map (getField . saMechanism) auths) auths
res <- GS.runSASL $ do
-- Select a common mechanism
mSelectedMech <- GS.clientSuggestMechanism $ map (GS.Mechanism . fst) authsMap
GS.Mechanism selectedMech <- liftIO $
maybe (throwIO $ SomeException NoSharedMechanism) return mSelectedMech
-- Safe because `selectedMech` was selected from one of the keys in `authsMap`
let Just selectedAuth = lookup selectedMech authsMap
GS.runClient (GS.Mechanism selectedMech) $ do
GS.setProperty GS.PropertyService $ T.encodeUtf8 pService
GS.setProperty GS.PropertyHostname $ T.encodeUtf8 pHost
GS.setProperty GS.PropertyAuthID $ T.encodeUtf8 (authUser hcUser)
(token,_) <- GS.step ""
r <- liftIO $ invokeSasl c RpcSaslProto
{ spState = putField Initiate
, spVersion = mempty
, spToken = putField $ Just token
, spAuths = putField $ [selectedAuth]
}
saslAuthLoop c r
case res of
Left err -> throwIO $ SomeException $ SaslException err
Right _ -> return ()
saslAuthLoop :: S.Stream -> RpcSaslProto -> GS.Session ()
saslAuthLoop c r = do
case getField $ spState r of
Challenge -> do
let mToken = getField $ spToken r
case mToken of
Nothing -> liftIO $ throwIO $ SomeException NoToken
Just token -> do
(resToken, _) <- GS.step token
resp <- liftIO $ invokeSasl c RpcSaslProto
{ spToken = putField $ Just resToken
, spState = putField $ Response
, spVersion = mempty
, spAuths = mempty
}
saslAuthLoop c resp
Success -> return ()
_ -> liftIO $ throwIO $ SomeException UnexptededState
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 ()
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 $ delimitedBytesL (rpcRequestHeaderProto csClientId callId)
<> delimitedBytesL (requestHeaderProto protocol method)
<> L.fromStrict requestBytes
S.runPut csStream (putMessage bs)
forkRecv :: ConnectionState -> IO ThreadId
forkRecv cs@ConnectionState{..} = forkIO $ handle (onSocketError cs) $ forever $ do
mget <- S.maybeGet csStream $ do
n <- fromIntegral <$> getWord32be
-- TODO Would be nice if we didn't have to isolate here
-- TODO and could stream instead. We could stream if we
-- TODO were able to read the varint length prefix
-- TODO ourselves and keep track of how many bytes were
-- TODO remaining instead of calling `getRemaining`.
isolate n $ do
hdr <- decodeLengthPrefixedMessage
msg <- case getField (P.rspStatus hdr) of
P.Success -> Right <$> getRemaining
_ -> return . Left . SomeException $ rspError hdr
return (hdr, msg)
case mget of
Nothing -> throwIO ConnectionClosed
Just (hdr, msg) -> do
onResponse <- fromMaybe (return $ return ())
<$> lookupDelete csRecvCallbacks (fromIntegral $ getField $ P.rspCallId hdr)
onResponse msg
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 ()
mkClientId :: IO ClientId
mkClientId = ClientId . L.toStrict . UUID.toByteString <$> randomIO
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)
------------------------------------------------------------------------
contextProto :: Protocol -> UserDetails -> P.IpcConnectionContext
contextProto protocol user = P.IpcConnectionContext
{ P.ctxProtocol = putField (Just (prName protocol))
, P.ctxUserInfo = putField (Just P.UserInformation
{ P.effectiveUser = putField (Just (authUser user))
, P.realUser = mempty
})
}
rpcRequestHeaderProto :: ClientId -> CallId -> P.RpcRequestHeader
rpcRequestHeaderProto clientId callId = P.RpcRequestHeader
{ P.reqKind = putField (Just P.ProtocolBuffer)
, P.reqOp = putField (Just P.FinalPacket)
, P.reqCallId = putField (fromIntegral callId)
, P.reqClientId = putField (unClientId clientId)
, P.reqRetryCount = putField (Just (-1))
}
requestHeaderProto :: Protocol -> Method -> P.RequestHeader
requestHeaderProto protocol method = P.RequestHeader
{ P.reqMethodName = putField method
, P.reqProtocolName = putField (prName protocol)
, P.reqProtocolVersion = putField (fromIntegral (prVersion protocol))
}
rspError :: P.RpcResponseHeader -> RemoteError
rspError rsp = RemoteError (fromMaybe "unknown error" $ getField $ P.rspExceptionClassName rsp)
(fromMaybe "unknown error" $ getField $ P.rspErrorMsg rsp)
putMessage :: L.ByteString -> Put
putMessage body = do
putWord32be (fromIntegral (L.length body))
putLazyByteString body
getRemaining :: Get ByteString
getRemaining = do
n <- remaining
getByteString n
------------------------------------------------------------------------
invoke :: (Decode b, Encode a) => Connection -> Text -> a -> IO b
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 (delimitedBytes arg) k'
where
k' (Left err) = k (Left err)
k' (Right bs) = k (fromDelimitedBytes bs)
delimitedBytes :: Encode a => a -> ByteString
delimitedBytes = runPut . encodeLengthPrefixedMessage
delimitedBytesL :: Encode a => a -> L.ByteString
delimitedBytesL = L.fromStrict . delimitedBytes
fromDelimitedBytes :: Decode a => ByteString -> Either SomeException a
fromDelimitedBytes bs = case runGetState decodeLengthPrefixedMessage 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