source: trunk/server/common/oursrc/hsparfind/hsparfind.hs @ 2305

Last change on this file since 2305 was 2305, checked in by ezyang, 9 years ago
Haskell version of parallel-find.
File size: 4.4 KB
Line 
1{-# LANGUAGE ViewPatterns #-}
2
3-- POSIX only
4
5import Prelude hiding (catch)
6
7import Data.Char
8import Data.List
9import Data.Maybe
10
11import Control.Arrow
12import Control.Monad
13import Control.Applicative
14import Control.Concurrent
15import Control.Concurrent.MVar
16import Control.Concurrent.STM
17import Control.Exception
18
19import System.FilePath
20import System.Process
21import System.IO
22import System.Directory
23import System.Exit
24import System.Posix hiding (createDirectory)
25
26destdir = "/mit/scripts/sec-tools/store/versions"
27
28whenM :: Monad m => m Bool -> m () -> m ()
29whenM p x = p >>= \b -> if b then x else return ()
30
31-- A simple semaphore implementation on a TVar Int.  Don't recursively
32-- call this while in a limit; you will be sad.
33limit :: TVar Int -> IO a -> IO a
34limit pool m = do
35    atomically $ do
36        i <- readTVar pool
37        check (i > 0)
38        writeTVar pool (i - 1)
39    m `finally` atomically (readTVar pool >>= writeTVar pool . (+1))
40
41-- These are cribbed off http://www.haskell.org/ghc/docs/5.00/set/sec-ghc-concurrency.html
42-- but with less unsafePerformIO
43
44-- Fork and register a child, so that it can be waited on
45forkChild :: MVar [MVar ()] -> IO () -> IO ()
46forkChild children m = do
47    c <- newEmptyMVar
48    forkIO (m `finally` putMVar c ())
49    cs <- takeMVar children
50    putMVar children (c:cs)
51
52-- Wait on all children
53waitForChildren :: MVar [MVar ()] -> IO ()
54waitForChildren children = do
55    cs' <- takeMVar children
56    case cs' of
57        [] -> return ()
58        (c:cs) -> do
59            putMVar children cs
60            takeMVar c
61            waitForChildren children
62
63-- Check if we have permissions
64checkPerm :: TVar Int -> FilePath -> IO Bool
65checkPerm pool base = ("system:scripts-security-upd rlidwk" `isInfixOf`) <$> exec pool "fs" ["listacl", base]
66
67newVersion pool cn base = do
68    stdout <- exec pool "sudo" ["-u", cn, "git", "--git-dir", base </> ".git", "describe", "--tags", "--always"]
69    -- XXX null stdout is an error condition, should say something
70    return (if null stdout then stdout else init stdout) -- munge off trailing newline
71oldVersion base =
72    -- XXX empty file is an error condition, should say something
73    last . lines <$> readFile (base </> ".scripts-version")
74
75writeOut cn base r = do
76    let line = base ++ ":" ++ r ++ "\n"
77    putStr line
78    appendFile (destdir </> cn) line
79
80exec :: TVar Int -> String -> [String] -> IO String
81exec pool bin args = do
82    (_, stdout, _) <- limit pool $ readProcessWithExitCode bin args ""
83    return stdout
84
85main = do
86    let lockfile = destdir ++ ".lock"
87    (_, host, _) <- readProcessWithExitCode "hostname" [] ""
88    pid <- getProcessID
89    whenM (doesFileExist lockfile) (error "Another parallel-find already in progress")
90    -- XXX if we lose the race the error message isn't as good
91    bracket_ (openFd lockfile WriteOnly (Just 0o644) (defaultFileFlags {exclusive = True})
92                >>= fdToHandle
93                >>= \h -> hPutStrLn h (host ++ " " ++ show pid) >> hClose h)
94             (removeFile lockfile)
95             (prepare >> parfind)
96
97prepare = do
98    whenM (doesDirectoryExist destdir) $ do
99        uniq <- show <$> epochTime
100        -- XXX does the wrong thing if you lose the race
101        renameDirectory destdir (destdir ++ uniq)
102    createDirectory destdir
103
104parfind = do
105    pool <- newTVarIO 40 -- number of child subprocesses to spawn simultaneously
106    children <- newMVar []
107    userlines <- lines <$> readFile "/mit/scripts/admin/backup/userlist"
108    let userdirs = filter ((/= "dn:") . fst) -- XXX should be done by generator of userlist
109                 . catMaybes
110                 . map (\s -> second tail    -- proof obligation discharged by elemIndex
111                           .  (`splitAt` s)
112                          <$> elemIndex ' ' s)
113                 $  userlines
114    forM_ userdirs $ \(cn, homedir) -> forkChild children $ do
115        let scriptsdir = homedir </> "web_scripts"
116        matches <- lines <$> exec pool "find" [scriptsdir, "-xdev", "-name", ".scripts-version", "-o", "-name", ".scripts"]
117        forM_ matches $ \dir -> forkIO . handle (\(SomeException e) -> putStrLn (dir ++ ": " ++ show e)) $ do
118            let base = takeDirectory dir
119            whenM (checkPerm pool base) $ do
120            if ".scripts" `isSuffixOf` dir
121                then newVersion pool cn base >>= writeOut cn base
122                else whenM (not <$> doesDirectoryExist (base </> ".scripts")) $ oldVersion base >>= writeOut cn base
123    waitForChildren children
Note: See TracBrowser for help on using the repository browser.