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

Last change on this file since 2615 was 2306, checked in by ezyang, 12 years ago
Don't starve non-find commands, handle concurrent writes to file.
File size: 4.7 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 handle_mvar base r =
76    withMVar handle_mvar $ \handle -> do
77        let line = base ++ ":" ++ r ++ "\n"
78        putStr line
79        hPutStr handle line
80
81exec :: TVar Int -> String -> [String] -> IO String
82exec pool bin args = do
83    (_, stdout, _) <- limit pool $ readProcessWithExitCode bin args ""
84    return stdout
85
86main = do
87    let lockfile = destdir ++ ".lock"
88    (_, host, _) <- readProcessWithExitCode "hostname" [] ""
89    pid <- getProcessID
90    whenM (doesFileExist lockfile) (error "Another parallel-find already in progress")
91    -- XXX if we lose the race the error message isn't as good
92    bracket_ (openFd lockfile WriteOnly (Just 0o644) (defaultFileFlags {exclusive = True})
93                >>= fdToHandle
94                >>= \h -> hPutStrLn h (host ++ " " ++ show pid) >> hClose h)
95             (removeFile lockfile)
96             (prepare >> parfind)
97
98prepare = do
99    whenM (doesDirectoryExist destdir) $ do
100        uniq <- show <$> epochTime
101        -- XXX does the wrong thing if you lose the race
102        renameDirectory destdir (destdir ++ uniq)
103    createDirectory destdir
104
105parfind = do
106    findpool <- newTVarIO 50
107    pool <- newTVarIO 10 -- git/fs gets its own pool so they don't starve
108    children <- newMVar []
109    userlines <- lines <$> readFile "/mit/scripts/admin/backup/userlist"
110    let userdirs = filter ((/= "dn:") . fst) -- XXX should be done by generator of userlist
111                 . catMaybes
112                 . map (\s -> second tail    -- proof obligation discharged by elemIndex
113                           .  (`splitAt` s)
114                          <$> elemIndex ' ' s)
115                 $  userlines
116    forM_ userdirs $ \(cn, homedir) -> forkChild children $ do
117        subchildren <- newMVar []
118        let scriptsdir = homedir </> "web_scripts"
119        matches <- lines <$> exec findpool "find" [scriptsdir, "-xdev", "-name", ".scripts-version", "-o", "-name", ".scripts"]
120        withFile (destdir </> cn) WriteMode $ \h -> do
121            mh <- newMVar h
122            forM_ matches $ \dir -> forkChild subchildren . handle (\(SomeException e) -> putStrLn (dir ++ ": " ++ show e)) $ do
123                let base = takeDirectory dir
124                whenM (checkPerm pool base) $ do
125                if ".scripts" `isSuffixOf` dir
126                    then newVersion pool cn base >>= writeOut mh base
127                    else whenM (not <$> doesDirectoryExist (base </> ".scripts")) $ oldVersion base >>= writeOut mh base
128            waitForChildren subchildren
129    waitForChildren children
Note: See TracBrowser for help on using the repository browser.