Commits

Anonymous committed 218113c

more work on the mercurial command server

Comments (0)

Files changed (3)

Data/FileStore/Mercurial.hs

 -- | Initialize a repository, creating the directory if needed.
 mercurialInit :: FilePath -> IO ()
 mercurialInit repo = do
-  checkMercurialVersion repo
   exists <- doesDirectoryExist repo
   when exists $ withVerifyDir repo $ throwIO RepositoryExists
   createDirectoryIfMissing True repo

Data/FileStore/MercurialCommandServer.hs

 module Data.FileStore.MercurialCommandServer
     ( runMercurialCommand
     , rawRunMercurialCommand
-    , checkMercurialVersion
     )
 where
 
+import Control.Applicative ((<$>))
 import Control.Exception (Exception, onException, throwIO)
 import Control.Monad (when)
+import Data.Bits (shiftL, shiftR, (.|.))
+import Data.IORef
+import Data.FileStore.Utils (runShellCommand)
+import Data.List (intercalate)
 import Data.Typeable (Typeable)
-import Data.IORef
-import Data.List (intercalate)
-import Data.FileStore.Utils (runShellCommand)
+import Data.Word (Word32)
 import System.Exit (ExitCode(..))
 import System.Process (runInteractiveProcess)
-import System.IO (Handle, hClose, hReady, hPutStr, hFlush)
+import System.IO (Handle, hClose, hPutStr, hFlush)
 import System.IO.Unsafe (unsafePerformIO)
 
-import qualified Data.Binary.Get as G
-import qualified Data.Binary.Put as P
 import qualified Data.ByteString as B
 import qualified Data.ByteString.UTF8 as UTF8
 import qualified Data.ByteString.Char8 as B8
 import qualified Data.ByteString.Lazy as BL
 import qualified Data.ByteString.Lazy.UTF8 as LUTF8
-import qualified Data.Map as M
+
+-- | Maximum number of servers to keep around
+maxPoolSize :: Int
+maxPoolSize = 2
 
 -- | Run a mercurial command and return error status, error output, standard output.  The repository
 -- is used as working directory.
 runMercurialCommand :: FilePath -> String -> [String] -> IO (ExitCode, String, BL.ByteString)
 runMercurialCommand repo command args = do
-  use <- useMercurialServer
-  if use
-    then runMercurialServer repo command args
-    else rawRunMercurialCommand repo command args
+  server <- getServer
+  case server of
+     Nothing -> rawRunMercurialCommand repo command args
+     Just h  -> do ret <- runMercurialServer repo command args h `onException` cleanupServer h
+                   putServer h
+                   return ret
 
+-- | Run a mercurial command directly without using the server.
 rawRunMercurialCommand :: FilePath -> String -> [String] -> IO (ExitCode, String, BL.ByteString)
 rawRunMercurialCommand repo command args = do
    (status, err, out) <- runShellCommand repo Nothing "hg" (command : args)
    return (status, LUTF8.toString err, out)
 
--- | needs to be called once.  TODO: write me
-checkMercurialVersion :: FilePath -> IO ()
-checkMercurialVersion _ = return ()
-
-data MercurialGlobalState = MercurialGlobalState {
-    useCommandServer :: Bool
-  , serverHandles    :: M.Map FilePath (Handle,Handle,Handle)
-} deriving (Show)
-
--- GAH, wish there was a better way.
--- See http://www.haskell.org/haskellwiki/Top_level_mutable_state
-mercurialGlobalVar :: IORef MercurialGlobalState
-{-# NOINLINE mercurialGlobalVar #-}
-mercurialGlobalVar = unsafePerformIO (newIORef (MercurialGlobalState True M.empty))
-
--- | returns if the mercurial version supports the command server.
-useMercurialServer :: IO Bool
-useMercurialServer = do
-    (MercurialGlobalState { useCommandServer = x}) <- readIORef mercurialGlobalVar
-    return x
-
-createServer :: FilePath -> IO (Handle,Handle,Handle)
-createServer repo = do
-    (hin,hout,herr,_) <- runInteractiveProcess "hg" ["serve", "--cmdserver", "pipe"] (Just repo) Nothing
+-- | Create a new command server for the given repository
+createServer :: IO (Handle,Handle,Handle)
+createServer = do
+    (hin,hout,herr,_) <- runInteractiveProcess "hg" ["serve", "--cmdserver", "pipe"] Nothing Nothing
     hello <- readMessage hout
     case hello of
        MessageO _ -> return (hin,hout,herr)
        MessageE x -> throwIO $ MercurialServerException (UTF8.toString x)
        _          -> throwIO $ MercurialServerException "unknown hello message"
 
+-- | Cleanup a command sever.  Mercurial will automatically exit itself
+--   when the handles are closed.
 cleanupServer :: (Handle,Handle,Handle) -> IO ()
 cleanupServer (hin,hout,herr) = hClose hin >> hClose hout >> hClose herr
 
--- | Run an action with some handles.  This function will create a server
---   if one does not exist or another thread is using it.  Only one server
---   is kept around.
-withServerHandles :: FilePath -> ((Handle,Handle,Handle) -> IO a) -> IO a
-withServerHandles repo f = do
-     h <- get
-     onException (run h) (cleanupServer h)
-
-  where get = do res <- atomicModifyIORef mercurialGlobalVar $ \state ->
-                           case M.lookup repo (serverHandles state) of
-                              Just h  -> (state { serverHandles = M.delete repo (serverHandles state)}, Right h)
-                              Nothing -> (state, Left ())
-                 case res of
-                     Left ()  -> createServer repo
-                     Right h' -> return h'
-
-        run h = do ret <- f h
-                   put h
-                   return ret
-
-        put h = do res <- atomicModifyIORef mercurialGlobalVar $ \state ->
-                             case M.lookup repo (serverHandles state) of
-                                 Just _  -> (state, Right ())
-                                 Nothing -> (state { serverHandles = M.insert repo h (serverHandles state)}, Left ())
-                   case res of
-                       Right () -> cleanupServer h
-                       Left  () -> return ()
-
 -- | format a command for sending to the server
-formatCommand :: String -> [String] -> B.ByteString
-formatCommand cmd args = UTF8.fromString $ intercalate "\0" $ cmd : args
+formatCommand :: FilePath -> String -> [String] -> B.ByteString
+formatCommand repo cmd args = UTF8.fromString $ intercalate "\0" allargs
+   where allargs = cmd : "-R" : repo : args
 
 -- | run a command using the mercurial server
-runMercurialServer :: FilePath -> String -> [String] -> IO (ExitCode, String, BL.ByteString)
-runMercurialServer repo cmd args = withServerHandles repo $ \(hin,hout,herr) -> do
+runMercurialServer :: FilePath -> String -> [String] -> (Handle,Handle,Handle) -> IO (ExitCode, String, BL.ByteString)
+runMercurialServer repo cmd args (hin,hout,herr) = do
     hPutStr hin "runcommand\n"
-    let fcmd = formatCommand cmd args
-    BL.hPut hin $ P.runPut (P.putWord32be $ fromIntegral $ B.length fcmd)
+    let fcmd = formatCommand repo cmd args
+    hWriteWord32be hin $ fromIntegral $ B.length fcmd
     B.hPut hin fcmd
     hFlush hin
     processUntilR hout herr
 
 -- | Read messages from the server until the command finishes or an error message appears
 processUntilR :: Handle -> Handle -> IO (ExitCode, String, BL.ByteString)
-processUntilR hout herr = loop BL.empty BL.empty
+processUntilR hout _ = loop BL.empty BL.empty
   where loop out err =
-          do iserr <- hReady herr
-             if iserr
-               then do errout <- B.hGetContents herr
-                       return (ExitFailure 1, UTF8.toString errout, BL.empty)
-               else do
-                 m <- readMessage hout
-                 case m of
-                    MessageO x -> loop (BL.append out $ BL.fromChunks [x]) err
-                    MessageE x -> loop out (BL.append err $ BL.fromChunks [x])
-                    MessageR x -> if (BL.null err)
-                                     then return (ExitSuccess, "", BL.append out $ BL.fromChunks [x])
-                                     else return (ExitFailure 1, LUTF8.toString err, BL.append out $ BL.fromChunks [x])
+          do m <- readMessage hout
+             case m of
+                MessageO x -> loop (BL.append out $ BL.fromChunks [x]) err
+                MessageE x -> loop out (BL.append err $ BL.fromChunks [x])
+                MessageR _ -> if (BL.null err)
+                                then return (ExitSuccess, "", out)
+                                else return (ExitFailure 1, LUTF8.toString err, out)
 
 data MercurialMessage = MessageO B.ByteString
                       | MessageR B.ByteString
     let c = B8.head buf
     when (c /= 'r' && c /= 'e' && c /= 'o' && c /= 'd') $
        throwIO $ MercurialServerException $ "Unknown channel " ++ show c
-    blen <- BL.hGet hout 4
-    when (buf == B.empty) $
-       throwIO $ MercurialServerException "Unknown length"
-    let len = G.runGet G.getWord32be blen
-    bdata <- B.hGet hout (fromIntegral len)
-    when (B.length bdata /= (fromIntegral len)) $
+    len <- hReadWord32be hout
+    bdata <- B.hGet hout len
+    when (B.length bdata /= len) $
        throwIO $ MercurialServerException "Mercurial did not produce enough output"
     case c of
       'r' -> return $ MessageR bdata
       'o' -> return $ MessageO bdata
       'e' -> return $ MessageE bdata
       'd' -> readMessage hout -- skip this message
-      _   -> throwIO $ MercurialServerException $ "Unknown channel2 " ++ show c
+      _   -> throwIO $ MercurialServerException $ "Unknown channel " ++ show c
+
+-- | Read a 32-bit big-endian into an Int
+hReadWord32be :: Handle -> IO Int
+hReadWord32be h = do
+    s <- B.hGet h 4
+    when (B.length s /= 4) $
+    	throwIO $ MercurialServerException "unable to read int"
+    return $! (fromIntegral (s `B.index` 0) `shiftL` 24) .|.
+              (fromIntegral (s `B.index` 1) `shiftL` 16) .|.
+              (fromIntegral (s `B.index` 2) `shiftL`  8) .|.
+              (fromIntegral (s `B.index` 3) )
+
+-- | Write a Word32 in big-endian to the handle
+hWriteWord32be :: Handle -> Word32 -> IO ()
+hWriteWord32be h w = B.hPut h buf
+  where buf = B.pack [  -- fromIntegeral to convert to Word8
+                fromIntegral (w `shiftR` 24),
+                fromIntegral (w `shiftR` 16),
+                fromIntegral (w `shiftR`  8),
+                fromIntegral w
+              ]
+
+-------------------------------------------------------------------
+-- Maintain a pool of mercurial servers.  Currently stored in a
+-- global IORef.  The code must provide two functions, to get
+-- and put a server from the pool.  The code above takes care of
+-- cleaning up if an exception occurs.
+-------------------------------------------------------------------
+
+data MercurialGlobalState = MercurialGlobalState {
+    useCommandServer :: Maybe Bool
+  , serverHandles    :: [(Handle,Handle,Handle)]
+} deriving (Show)
+
+-- | See http://www.haskell.org/haskellwiki/Top_level_mutable_state
+mercurialGlobalVar :: IORef MercurialGlobalState
+{-# NOINLINE mercurialGlobalVar #-}
+mercurialGlobalVar = unsafePerformIO (newIORef (MercurialGlobalState Nothing []))
+
+-- | Pull a server out of the pool.  Returns nothing if the mercurial version
+--   does not support servers.
+getServer :: IO (Maybe (Handle, Handle, Handle))
+getServer = do
+    use <- useCommandServer <$> readIORef mercurialGlobalVar
+    case use of
+      Just False -> return Nothing
+      Nothing    -> do isok <- checkVersion
+                       atomicModifyIORef mercurialGlobalVar $ \state ->
+                          (state { useCommandServer = Just isok }, ())
+                       getServer
+      Just True  -> allocateServer
+
+-- | Helper function called once we know that mercurial supports servers
+allocateServer :: IO (Maybe (Handle, Handle, Handle))
+allocateServer = do
+    ret <- atomicModifyIORef mercurialGlobalVar $ \state ->
+             case serverHandles state of
+                (x:xs) -> (state { serverHandles = xs}, Right x)
+                []     -> (state, Left ())
+    putStrLn $ "allocate: " ++ show ret
+    case ret of
+      Right x -> return $ Just x
+      Left () -> Just <$> createServer
+
+-- | Puts a server back in the pool if the pool is not full,
+--   otherwise closes the server.
+putServer :: (Handle,Handle,Handle) -> IO ()
+putServer h = do
+    ret <- atomicModifyIORef mercurialGlobalVar $ \state -> do
+              case serverHandles state of
+                  xs | length xs >= maxPoolSize -> (state, Right ())
+                  xs -> (state { serverHandles = h:xs}, Left ())
+    case ret of
+      Right () -> cleanupServer h
+      Left  () -> return ()
+
+-- | Check if the mercurial version supports servers
+checkVersion :: IO Bool
+checkVersion = return True
 
 Library
     Build-depends:       base >= 4 && < 5,
-                         binary >= 0.5 && < 0.6,
                          bytestring >= 0.9 && < 1.0,
-                         containers >= 0.4 && < 0.5,
                          utf8-string >= 0.3 && < 0.4,
                          filepath >= 1.1 && < 1.3,
                          directory >= 1.0 && < 1.2,
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.