git.iko.soy/distributed-process
on 17ad36a help sign in
distributed-process/distributed-process-azure/src/Control/Distributed/Process/Backend/Azure/GenericMain.hs
401 lines 13 kB history raw
1-- | Generic main
2module Control.Distributed.Process.Backend.Azure.GenericMain
3 ( genericMain
4 , ProcessPair(..)
5 ) where
7import Prelude hiding (catch)
8import System.Exit (exitSuccess, exitFailure)
9import System.IO
10 ( hFlush
11 , stdout
12 , stdin
13 , stderr
14 , hSetBinaryMode
15 , hClose
16 , Handle
17 )
18import Data.Foldable (forM_)
19import Data.Binary (decode)
20import qualified Data.ByteString.Lazy as BSL (ByteString, hGet, toChunks, length)
21import qualified Data.ByteString as BSS (hGet, length)
22import Control.Monad (unless, forM, void)
23import Control.Monad.Reader (ask)
24import Control.Exception (SomeException)
25import Control.Applicative ((<$>), (<*>), optional)
26import Control.Monad.IO.Class (liftIO)
27import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, readMVar)
29-- Posix
30import qualified System.Posix.Process as Posix (forkProcess, createSession)
32-- SSH
33import qualified Network.SSH.Client.LibSSH2.Foreign as SSH
34 ( initialize
35 , exit
36 )
38-- Command line options
39import Options.Applicative
40 ( Parser
41 , strOption
42 , long
43 , (&)
44 , metavar
45 , help
46 , subparser
47 , command
48 , info
49 , progDesc
50 , execParser
51 , helper
52 , fullDesc
53 , header
54 , switch
55 )
57-- CH
58import Control.Distributed.Process
59 ( RemoteTable
60 , Process
61 , unClosure
62 , catch
63 )
64import Control.Distributed.Process.Node
65 ( newLocalNode
66 , runProcess
67 , forkProcess
68 , initRemoteTable
69 , LocalNode
70 )
71import Control.Distributed.Process.Internal.Types
72 ( LocalProcess(processQueue)
73 , payloadToMessage
74 , Message
75 )
76import Control.Distributed.Process.Internal.CQueue (CQueue, enqueue)
78-- Azure
79import Control.Distributed.Process.Backend.Azure
80 ( AzureParameters(azureSshUserName, azureSetup)
81 , defaultAzureParameters
82 , initializeBackend
83 , cloudServices
84 , VirtualMachine(vmName)
85 , Backend(findVMs, copyToVM, checkMD5, callOnVM, spawnOnVM)
86 , ProcessPair(..)
87 , RemoteProcess
88 , remoteSend'
89 , remoteThrow
90 )
92-- Transport
93import Network.Transport.Internal (decodeInt32)
94import Network.Transport.TCP (createTransport, defaultTCPParameters)
96--------------------------------------------------------------------------------
97-- Main --
98--------------------------------------------------------------------------------
100genericMain :: (RemoteTable -> RemoteTable) -- ^ Standard CH remote table
101 -> (String -> IO (ProcessPair ())) -- ^ Closures to support in 'run'
102 -> (String -> IO (RemoteProcess ())) -- ^ Closures to support in @run --background@
103 -> IO ()
104genericMain remoteTable callable spawnable = do
105 _ <- SSH.initialize True
106 cmd <- execParser opts
107 case cmd of
108 List {} -> do
109 params <- azureParameters (azureOptions cmd) Nothing
110 css <- cloudServices (azureSetup params)
111 mapM_ print css
112 CopyTo {} -> do
113 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
114 backend <- initializeBackend params (targetService (target cmd))
115 vms <- findMatchingVMs backend (targetVM (target cmd))
116 forM_ vms $ \vm -> do
117 putStr (vmName vm ++ ": ") >> hFlush stdout
118 copyToVM backend vm
119 putStrLn "Done"
120 CheckMD5 {} -> do
121 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
122 backend <- initializeBackend params (targetService (target cmd))
123 vms <- findMatchingVMs backend (targetVM (target cmd))
124 matches <- forM vms $ \vm -> do
125 unless (status cmd) $ putStr (vmName vm ++ ": ") >> hFlush stdout
126 match <- checkMD5 backend vm
127 unless (status cmd) $ putStrLn $ if match then "OK" else "FAILED"
128 return match
129 if and matches
130 then exitSuccess
131 else exitFailure
132 RunOn {} | background cmd -> do
133 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
134 backend <- initializeBackend params (targetService (target cmd))
135 vms <- findMatchingVMs backend (targetVM (target cmd))
136 rProc <- spawnable (closureId cmd)
137 forM_ vms $ \vm -> do
138 putStr (vmName vm ++ ": ") >> hFlush stdout
139 spawnOnVM backend vm (remotePort cmd) rProc
140 putStrLn "OK"
141 RunOn {} {- not (background cmd) -} -> do
142 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
143 backend <- initializeBackend params (targetService (target cmd))
144 vms <- findMatchingVMs backend (targetVM (target cmd))
145 procPair <- callable (closureId cmd)
146 forM_ vms $ \vm -> do
147 putStr (vmName vm ++ ": ") >> hFlush stdout
148 callOnVM backend vm (remotePort cmd) procPair
149 OnVmCommand (vmCmd@OnVmRun {}) ->
150 onVmRun (remoteTable initRemoteTable)
151 (onVmIP vmCmd)
152 (onVmPort vmCmd)
153 (onVmBackground vmCmd)
154 SSH.exit
155 where
156 opts = info (helper <*> commandParser)
157 ( fullDesc
158 & header "Cloud Haskell backend for Azure"
159 )
161findMatchingVMs :: Backend -> Maybe String -> IO [VirtualMachine]
162findMatchingVMs backend Nothing = findVMs backend
163findMatchingVMs backend (Just vm) = filter ((== vm) . vmName) `fmap` findVMs backend
165azureParameters :: AzureOptions -> Maybe SshOptions -> IO AzureParameters
166azureParameters opts Nothing =
167 defaultAzureParameters (subscriptionId opts)
168 (pathToCert opts)
169 (pathToKey opts)
170azureParameters opts (Just sshOpts) = do
171 params <- defaultAzureParameters (subscriptionId opts)
172 (pathToCert opts)
173 (pathToKey opts)
174 return params {
175 azureSshUserName = remoteUser sshOpts
176 }
178--------------------------------------------------------------------------------
179-- Executing a closure on the VM --
180--------------------------------------------------------------------------------
182onVmRun :: RemoteTable -> String -> String -> Bool -> IO ()
183onVmRun rtable host port bg = do
184 hSetBinaryMode stdin True
185 hSetBinaryMode stdout True
186 mProcEnc <- getWithLength stdin
187 forM_ mProcEnc $ \procEnc -> do
188 let proc = decode procEnc
189 lprocMVar <- newEmptyMVar :: IO (MVar LocalProcess)
190 if bg
191 then detach $ startCH proc lprocMVar runProcess (\_ -> return ())
192 else do
193 startCH proc lprocMVar forkProcess remoteThrow
194 lproc <- readMVar lprocMVar
195 queueFromHandle stdin (processQueue lproc)
196 where
197 startCH :: RemoteProcess ()
198 -> MVar LocalProcess
199 -> (LocalNode -> Process () -> IO a)
200 -> (SomeException -> Process ())
201 -> IO ()
202 startCH rproc lprocMVar go exceptionHandler = do
203 mTransport <- createTransport host port defaultTCPParameters
204 case mTransport of
205 Left err -> remoteSend' 1 (show err)
206 Right transport -> do
207 node <- newLocalNode transport rtable
208 void . go node $ do
209 ask >>= liftIO . putMVar lprocMVar
210 let backend = error "TODO: backend not initialized in onVmRun"
211 proc <- unClosure rproc :: Process (Backend -> Process ())
212 catch (proc backend) exceptionHandler
214-- | Read a 4-byte length @l@ and then an @l@-byte payload
215--
216-- Returns Nothing on EOF
217getWithLength :: Handle -> IO (Maybe BSL.ByteString)
218getWithLength h = do
219 lenEnc <- BSS.hGet h 4
220 if BSS.length lenEnc < 4
221 then return Nothing
222 else do
223 let len = decodeInt32 lenEnc
224 bs <- BSL.hGet h len
225 if BSL.length bs < fromIntegral len
226 then return Nothing
227 else return (Just bs)
229queueFromHandle :: Handle -> CQueue Message -> IO ()
230queueFromHandle h q = do
231 mPayload <- getWithLength stdin
232 forM_ mPayload $ \payload -> do
233 enqueue q $ payloadToMessage (BSL.toChunks payload)
234 queueFromHandle h q
236detach :: IO () -> IO ()
237detach io = do
238 mapM_ hClose [stdin, stdout, stderr]
239 void . Posix.forkProcess $ void Posix.createSession >> io
241--------------------------------------------------------------------------------
242-- Command line options --
243--------------------------------------------------------------------------------
245data AzureOptions = AzureOptions {
246 subscriptionId :: String
247 , pathToCert :: FilePath
248 , pathToKey :: FilePath
249 }
250 deriving Show
252data SshOptions = SshOptions {
253 remoteUser :: String
254 }
255 deriving Show
257data Target = Target {
258 targetService :: String
259 , targetVM :: Maybe String
260 }
261 deriving Show
263data Command =
264 List {
265 azureOptions :: AzureOptions
266 }
267 | CopyTo {
268 azureOptions :: AzureOptions
269 , sshOptions :: SshOptions
270 , target :: Target
271 }
272 | CheckMD5 {
273 azureOptions :: AzureOptions
274 , sshOptions :: SshOptions
275 , target :: Target
276 , status :: Bool
277 }
278 | RunOn {
279 azureOptions :: AzureOptions
280 , sshOptions :: SshOptions
281 , target :: Target
282 , remotePort :: String
283 , closureId :: String
284 , background :: Bool
285 }
286 | OnVmCommand {
287 _onVmCommand :: OnVmCommand
288 }
289 deriving Show
291data OnVmCommand =
292 OnVmRun {
293 onVmIP :: String
294 , onVmPort :: String
295 , onVmBackground :: Bool
296 }
297 deriving Show
299azureOptionsParser :: Parser AzureOptions
300azureOptionsParser = AzureOptions
301 <$> strOption ( long "subscription-id"
302 & metavar "SID"
303 & help "Azure subscription ID"
304 )
305 <*> strOption ( long "certificate"
306 & metavar "CERT"
307 & help "X509 certificate"
308 )
309 <*> strOption ( long "private"
310 & metavar "PRI"
311 & help "Private key in PKCS#1 format"
312 )
314sshOptionsParser :: Parser SshOptions
315sshOptionsParser = SshOptions
316 <$> strOption ( long "user"
317 & metavar "USER"
318 & help "Remove SSH username"
319 )
321listParser :: Parser Command
322listParser = List <$> azureOptionsParser
324copyToParser :: Parser Command
325copyToParser = CopyTo
326 <$> azureOptionsParser
327 <*> sshOptionsParser
328 <*> targetParser
330targetParser :: Parser Target
331targetParser = Target
332 <$> strOption ( long "cloud-service"
333 & metavar "CS"
334 & help "Cloud service name"
335 )
336 <*> optional (strOption ( long "virtual-machine"
337 & metavar "VM"
338 & help "Virtual machine name (all VMs if unspecified)"
339 ))
341checkMD5Parser :: Parser Command
342checkMD5Parser = CheckMD5
343 <$> azureOptionsParser
344 <*> sshOptionsParser
345 <*> targetParser
346 <*> switch ( long "status"
347 & help "Don't output anything, status code shows success"
348 )
350commandParser :: Parser Command
351commandParser = subparser
352 ( command "list" (info listParser
353 (progDesc "List Azure cloud services"))
354 & command "install" (info copyToParser
355 (progDesc "Install the executable"))
356 & command "md5" (info checkMD5Parser
357 (progDesc "Check if the remote and local MD5 hash match"))
358 & command "run" (info runOnParser
359 (progDesc "Run the executable"))
360 & command "onvm" (info onVmCommandParser
361 (progDesc "Commands used when running ON the vm (usually used internally only)"))
362 )
364runOnParser :: Parser Command
365runOnParser = RunOn
366 <$> azureOptionsParser
367 <*> sshOptionsParser
368 <*> targetParser
369 <*> strOption ( long "port"
370 & metavar "PORT"
371 & help "Port number of the CH instance"
372 )
373 <*> strOption ( long "closure"
374 & metavar "PROC"
375 & help "Process to run on the CH instance"
376 )
377 <*> switch ( long "background"
378 & help "Run the process in the background"
379 )
381onVmRunParser :: Parser OnVmCommand
382onVmRunParser = OnVmRun
383 <$> strOption ( long "host"
384 & metavar "IP"
385 & help "IP address"
386 )
387 <*> strOption ( long "port"
388 & metavar "PORT"
389 & help "port number"
390 )
391 <*> switch ( long "background"
392 & help "Run the process in the background"
393 )
395onVmCommandParser :: Parser Command
396onVmCommandParser = OnVmCommand <$> subparser
397 ( command "run" (info onVmRunParser
398 (progDesc "Run the executable"))
399 )