Anonymous avatar Anonymous committed a333393

Initial version of mercurial command server.

Tests don't pass yet, but MercurialCommandServer.hs works.

Comments (0)

Files changed (3)

Data/FileStore/Mercurial.hs

 import Data.FileStore.Types
 import Data.Maybe (fromJust)
 import System.Exit
-import Data.FileStore.Utils (withSanityCheck, hashsMatch, runShellCommand, withVerifyDir, grepSearchRepo) 
+import Data.FileStore.Utils (withSanityCheck, hashsMatch, withVerifyDir, grepSearchRepo) 
+import Data.FileStore.MercurialCommandServer
 import Data.ByteString.Lazy.UTF8 (toString)
 import qualified Data.ByteString.Lazy as B
 import qualified Text.ParserCombinators.Parsec as P
   , idsMatch          = const hashsMatch repo
   }
 
--- | 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, B.ByteString)
-runMercurialCommand repo command args = do
-  (status, err, out) <- runShellCommand repo Nothing "hg" (command : args)
-  return (status, toString err, out)
 
 -- | 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
-  (status, err, _) <- runMercurialCommand repo "init" []
+  (status, err, _) <- rawRunMercurialCommand repo "init" []
   if status == ExitSuccess
      then
        -- Add a hook so that changes made remotely via hg will be reflected in

Data/FileStore/MercurialCommandServer.hs

+{-# LANGUAGE DeriveDataTypeable #-}
+{- |
+   Module      : Data.FileStore.MercurialCommandServer
+   Copyright   : Copyright (C) 2011 John Lenz (lenz@math.uic.edu)
+   License     : BSD 3
+
+   Maintainer  : John MacFarlane <jgm@berkeley.edu>
+   Stability   : alpha
+   Portability : GHC 6.10 required
+
+   In version 1.9, mercurial introduced a command server which allows
+   a single instance of mercurial to be launched and multiple commands
+   can be executed without requiring mercurial to start and stop.  See
+   http://mercurial.selenic.com/wiki/CommandServer
+-}
+
+module Data.FileStore.MercurialCommandServer
+    ( runMercurialCommand
+    , rawRunMercurialCommand
+    , checkMercurialVersion
+    )
+where
+
+import Control.Exception (Exception, onException, throwIO)
+import Control.Monad (when)
+import Data.Typeable (Typeable)
+import Data.IORef
+import Data.List (intercalate)
+import Data.FileStore.Utils (runShellCommand)
+import System.Exit (ExitCode(..))
+import System.Process (runInteractiveProcess)
+import System.IO (Handle, hClose, hReady, 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
+
+-- | 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
+
+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
+    hello <- readMessage hout
+    case hello of
+       MessageO _ -> return (hin,hout,herr)
+       MessageE x -> throwIO $ MercurialServerException (UTF8.toString x)
+       _          -> throwIO $ MercurialServerException "unknown hello message"
+
+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
+
+-- | 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
+    hPutStr hin "runcommand\n"
+    let fcmd = formatCommand cmd args
+    BL.hPut hin $ P.runPut (P.putWord32be $ 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
+  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])
+
+data MercurialMessage = MessageO B.ByteString
+                      | MessageR B.ByteString
+                      | MessageE B.ByteString
+
+data MercurialServerException = MercurialServerException String
+  deriving (Show,Typeable)
+instance Exception MercurialServerException
+
+-- | Read a single message
+readMessage :: Handle -> IO MercurialMessage
+readMessage hout = do
+    buf <- B.hGet hout 1
+    when (buf == B.empty) $
+       throwIO $ MercurialServerException "Unknown channel"
+    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)) $
+       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
 
 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,
     Exposed-modules:     Data.FileStore, Data.FileStore.Types, Data.FileStore.Git, Data.FileStore.Darcs, Data.FileStore.Mercurial,
                          -- Data.FileStore.Sqlite3,
                          Data.FileStore.Utils, Data.FileStore.Generic
-    Other-modules:       Paths_filestore, Data.FileStore.DarcsXml
+    Other-modules:       Paths_filestore,
+                         Data.FileStore.DarcsXml,
+                         Data.FileStore.MercurialCommandServer
 
     if flag(maxcount) 
         cpp-options: -DUSE_MAXCOUNT
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.