ssh connection caching

Ssh connection caching is now enabled automatically by git-annex. Only one
ssh connection is made to each host per git-annex run, which can speed some
things up a lot, as well as avoiding repeated password prompts. Concurrent
git-annex processes also share ssh connections. Cached ssh connections are
shut down when git-annex exits.

Note: The rsync special remote does not yet participate in the ssh
connection caching.

2 files changed, 150 insertions, 0 deletions

diff --git a/Annex/LockPool.hs b/Annex/LockPool.hs
new file mode 100644
index 000000000..3fede5739
--- /dev/null
+++ b/Annex/LockPool.hs
@@ -0,0 +1,43 @@
+{- git-annex lock pool
+ -
+ - Copyright 2012 Joey Hess <joey@kitenet.net>
+ -
+ - Licensed under the GNU GPL version 3 or higher.
+ -}
+
+module Annex.LockPool where
+
+import qualified Data.Map as M
+import System.Posix.Types (Fd)
+
+import Common.Annex
+import Annex
+
+{- Create a specified lock file, and takes a shared lock. -}
+lockFile :: FilePath -> Annex ()
+lockFile file = go =<< fromPool file
+	where
+		go (Just _) = return () -- already locked
+		go Nothing = do
+			fd <- liftIO $ openFd file ReadOnly (Just stdFileMode) defaultFileFlags
+			liftIO $ waitToSetLock fd (ReadLock, AbsoluteSeek, 0, 0)
+			changePool $ M.insert file fd
+
+unlockFile :: FilePath -> Annex ()
+unlockFile file = go =<< fromPool file
+	where
+		go Nothing = return ()
+		go (Just fd) = do
+			liftIO $ closeFd fd
+			changePool $ M.delete file
+
+getPool :: Annex (M.Map FilePath Fd)
+getPool = getState lockpool
+
+fromPool :: FilePath -> Annex (Maybe Fd)
+fromPool file = M.lookup file <$> getPool
+
+changePool :: (M.Map FilePath Fd -> M.Map FilePath Fd) -> Annex ()
+changePool a = do
+	m <- getPool
+	changeState $ \s -> s { lockpool = a m }
diff --git a/Annex/Ssh.hs b/Annex/Ssh.hs
new file mode 100644
index 000000000..cd832a373
--- /dev/null
+++ b/Annex/Ssh.hs
@@ -0,0 +1,107 @@
+{- git-annex ssh interface, with connection caching
+ -
+ - Copyright 2012 Joey Hess <joey@kitenet.net>
+ -
+ - Licensed under the GNU GPL version 3 or higher.
+ -}
+
+module Annex.Ssh (
+	sshParams,
+	sshCleanup,
+) where
+
+import qualified Data.Map as M
+import System.IO.Error (try)
+
+import Common.Annex
+import Annex.LockPool
+
+{- Generates parameters to ssh to a given host (or user@host) on a given
+ - port, with connection caching. -}
+sshParams :: (String, Maybe Integer) -> Annex [CommandParam]
+sshParams (host, port) = do
+	cleanstale
+	(socketfile, params) <- sshInfo (host, port)
+	liftIO $ createDirectoryIfMissing True $ parentDir socketfile
+	lockFile $ socket2lock socketfile
+	return params
+	where
+		-- If the lock pool is empty, this is the first ssh of this
+		-- run. There could be stale ssh connections hanging around
+		-- from a previous git-annex run that was interrupted.
+		cleanstale = whenM (null . filter isLock . M.keys <$> getPool) $
+			sshCleanup
+
+sshInfo :: (String, Maybe Integer) -> Annex (FilePath, [CommandParam])
+sshInfo (host, port) = do
+	dir <- fromRepo $ gitAnnexSshDir
+	let socketfile = dir </> hostport2socket host port
+	return $ (socketfile, cacheParams socketfile ++ portParams port ++ [Param host])
+
+cacheParams :: FilePath -> [CommandParam]
+cacheParams socketfile =
+	[ Param "-S", Param socketfile
+	, Params "-o ControlMaster=auto -o ControlPersist=yes"
+	]
+
+portParams :: Maybe Integer -> [CommandParam]
+portParams Nothing = []
+portParams (Just port) = [Param "-p", Param $ show port]
+
+{- Stop any unused ssh processes. -}
+sshCleanup :: Annex ()
+sshCleanup = do
+	dir <- fromRepo $ gitAnnexSshDir
+	liftIO $ createDirectoryIfMissing True dir
+	sockets <- filter (not . isLock) <$> liftIO (dirContents dir)
+	forM_ sockets cleanup
+	where
+		cleanup socketfile = do
+			-- Drop any shared lock we have, and take an
+			-- exclusive lock, without blocking. If the lock
+			-- succeeds, nothing is using this ssh, and it can
+			-- be stopped.
+			let lockfile = socket2lock socketfile
+			unlockFile lockfile
+			fd <- liftIO $ openFd lockfile ReadWrite (Just stdFileMode) defaultFileFlags
+			v <- liftIO $ try $ setLock fd (WriteLock, AbsoluteSeek, 0, 0)
+			case v of
+				Left _ -> return ()
+				Right _ -> stopssh socketfile
+			liftIO $ closeFd fd
+		stopssh socketfile = do
+			(_, params) <- sshInfo $ socket2hostport socketfile
+			_ <- liftIO $ do
+				-- "ssh -O stop" is noisy on stderr even with -q
+				let cmd = unwords $ toCommand $
+					[ Params "-O stop"
+					] ++ params
+				_ <- boolSystem "sh"
+					[ Param "-c"
+					, Param $ "ssh " ++ cmd ++ " >/dev/null 2>/dev/null"
+					]
+				--try $ removeFile socketfile
+				return ()
+			-- Cannot remove the lock file; other processes may
+			-- be waiting on our exclusive lock to use it.
+			return ()
+
+hostport2socket :: String -> Maybe Integer -> FilePath
+hostport2socket host Nothing = host
+hostport2socket host (Just port) = host ++ "!" ++ show port
+
+socket2hostport :: FilePath -> (String, Maybe Integer)
+socket2hostport socket
+	| null p = (h, Nothing)
+	| otherwise = (h, readMaybe p)
+	where
+		(h, p) = separate (== '!') $ takeFileName socket
+
+socket2lock :: FilePath -> FilePath
+socket2lock socket = socket ++ lockExt
+
+isLock :: FilePath -> Bool
+isLock f = lockExt `isSuffixOf` f
+
+lockExt :: String
+lockExt = ".lock"