git.iko.soy/distributed-process
on d54d67e help sign in
distributed-process/distributed-process-azure/src/Control/Distributed/Process/Backend/Azure/GenericMain.hs
407 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 , remoteThrow
89 )
91-- Transport
92import Network.Transport.Internal (decodeInt32)
93import Network.Transport.TCP (createTransport, defaultTCPParameters)
95--------------------------------------------------------------------------------
96-- Main --
97--------------------------------------------------------------------------------
99genericMain :: (RemoteTable -> RemoteTable) -- ^ Standard CH remote table
100 -> (String -> IO (ProcessPair ())) -- ^ Closures to support in 'run'
101 -> (String -> IO (RemoteProcess ())) -- ^ Closures to support in @run --background@
102 -> IO ()
103genericMain remoteTable callable spawnable = do
104 _ <- SSH.initialize True
105 cmd <- execParser opts
106 case cmd of
107 List {} -> do
108 params <- azureParameters (azureOptions cmd) Nothing
109 css <- cloudServices (azureSetup params)
110 mapM_ print css
111 CopyTo {} -> do
112 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
113 backend <- initializeBackend params (targetService (target cmd))
114 vms <- findMatchingVMs backend (targetVM (target cmd))
115 forM_ vms $ \vm -> do
116 putStr (vmName vm ++ ": ") >> hFlush stdout
117 copyToVM backend vm
118 putStrLn "Done"
119 CheckMD5 {} -> do
120 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
121 backend <- initializeBackend params (targetService (target cmd))
122 vms <- findMatchingVMs backend (targetVM (target cmd))
123 matches <- forM vms $ \vm -> do
124 unless (status cmd) $ putStr (vmName vm ++ ": ") >> hFlush stdout
125 match <- checkMD5 backend vm
126 unless (status cmd) $ putStrLn $ if match then "OK" else "FAILED"
127 return match
128 if and matches
129 then exitSuccess
130 else exitFailure
131 RunOn {} | background cmd -> do
132 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
133 backend <- initializeBackend params (targetService (target cmd))
134 vms <- findMatchingVMs backend (targetVM (target cmd))
135 rProc <- spawnable (closureId cmd)
136 forM_ vms $ \vm -> do
137 putStr (vmName vm ++ ": ") >> hFlush stdout
138 spawnOnVM backend vm (remotePort cmd) rProc
139 putStrLn "OK"
140 RunOn {} {- not (background cmd) -} -> do
141 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
142 backend <- initializeBackend params (targetService (target cmd))
143 vms <- findMatchingVMs backend (targetVM (target cmd))
144 procPair <- callable (closureId cmd)
145 forM_ vms $ \vm -> do
146 putStr (vmName vm ++ ": ") >> hFlush stdout
147 callOnVM backend vm (remotePort cmd) procPair
148 OnVmCommand (vmCmd@OnVmRun {}) ->
149 onVmRun (remoteTable initRemoteTable)
150 (onVmIP vmCmd)
151 (onVmPort vmCmd)
152 (onVmService 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 -> String -> Bool -> IO ()
183onVmRun rtable host port cloudService bg = do
184 hSetBinaryMode stdin True
185 hSetBinaryMode stdout True
186 Just procEnc <- getWithLength stdin
187 Just paramsEnc <- getWithLength stdin
188 backend <- initializeBackend (decode paramsEnc) cloudService
189 let proc = decode procEnc
190 lprocMVar <- newEmptyMVar :: IO (MVar LocalProcess)
191 if bg
192 then detach $ startCH proc lprocMVar backend runProcess (\_ -> return ())
193 else do
194 startCH proc lprocMVar backend forkProcess (liftIO . remoteThrow)
195 lproc <- readMVar lprocMVar
196 queueFromHandle stdin (processQueue lproc)
197 where
198 startCH :: RemoteProcess ()
199 -> MVar LocalProcess
200 -> Backend
201 -> (LocalNode -> Process () -> IO a)
202 -> (SomeException -> Process ())
203 -> IO ()
204 startCH rproc lprocMVar backend go exceptionHandler = do
205 mTransport <- createTransport host port defaultTCPParameters
206 case mTransport of
207 Left err -> remoteThrow err
208 Right transport -> do
209 node <- newLocalNode transport rtable
210 void . go node $ do
211 ask >>= liftIO . putMVar lprocMVar
212 proc <- unClosure rproc :: Process (Backend -> Process ())
213 catch (proc backend) exceptionHandler
215-- | Read a 4-byte length @l@ and then an @l@-byte payload
216--
217-- Returns Nothing on EOF
218getWithLength :: Handle -> IO (Maybe BSL.ByteString)
219getWithLength h = do
220 lenEnc <- BSS.hGet h 4
221 if BSS.length lenEnc < 4
222 then return Nothing
223 else do
224 let len = decodeInt32 lenEnc
225 bs <- BSL.hGet h len
226 if BSL.length bs < fromIntegral len
227 then return Nothing
228 else return (Just bs)
230queueFromHandle :: Handle -> CQueue Message -> IO ()
231queueFromHandle h q = do
232 mPayload <- getWithLength stdin
233 forM_ mPayload $ \payload -> do
234 enqueue q $ payloadToMessage (BSL.toChunks payload)
235 queueFromHandle h q
237detach :: IO () -> IO ()
238detach io = do
239 mapM_ hClose [stdin, stdout, stderr]
240 void . Posix.forkProcess $ void Posix.createSession >> io
242--------------------------------------------------------------------------------
243-- Command line options --
244--------------------------------------------------------------------------------
246data AzureOptions = AzureOptions {
247 subscriptionId :: String
248 , pathToCert :: FilePath
249 , pathToKey :: FilePath
250 }
251 deriving Show
253data SshOptions = SshOptions {
254 remoteUser :: String
255 }
256 deriving Show
258data Target = Target {
259 targetService :: String
260 , targetVM :: Maybe String
261 }
262 deriving Show
264data Command =
265 List {
266 azureOptions :: AzureOptions
267 }
268 | CopyTo {
269 azureOptions :: AzureOptions
270 , sshOptions :: SshOptions
271 , target :: Target
272 }
273 | CheckMD5 {
274 azureOptions :: AzureOptions
275 , sshOptions :: SshOptions
276 , target :: Target
277 , status :: Bool
278 }
279 | RunOn {
280 azureOptions :: AzureOptions
281 , sshOptions :: SshOptions
282 , target :: Target
283 , remotePort :: String
284 , closureId :: String
285 , background :: Bool
286 }
287 | OnVmCommand {
288 _onVmCommand :: OnVmCommand
289 }
290 deriving Show
292data OnVmCommand =
293 OnVmRun {
294 onVmIP :: String
295 , onVmPort :: String
296 , onVmService :: String
297 , onVmBackground :: Bool
298 }
299 deriving Show
301azureOptionsParser :: Parser AzureOptions
302azureOptionsParser = AzureOptions
303 <$> strOption ( long "subscription-id"
304 & metavar "SID"
305 & help "Azure subscription ID"
306 )
307 <*> strOption ( long "certificate"
308 & metavar "CERT"
309 & help "X509 certificate"
310 )
311 <*> strOption ( long "private"
312 & metavar "PRI"
313 & help "Private key in PKCS#1 format"
314 )
316sshOptionsParser :: Parser SshOptions
317sshOptionsParser = SshOptions
318 <$> strOption ( long "user"
319 & metavar "USER"
320 & help "Remove SSH username"
321 )
323listParser :: Parser Command
324listParser = List <$> azureOptionsParser
326copyToParser :: Parser Command
327copyToParser = CopyTo
328 <$> azureOptionsParser
329 <*> sshOptionsParser
330 <*> targetParser
332targetParser :: Parser Target
333targetParser = Target
334 <$> strOption ( long "cloud-service"
335 & metavar "CS"
336 & help "Cloud service name"
337 )
338 <*> optional (strOption ( long "virtual-machine"
339 & metavar "VM"
340 & help "Virtual machine name (all VMs if unspecified)"
341 ))
343checkMD5Parser :: Parser Command
344checkMD5Parser = CheckMD5
345 <$> azureOptionsParser
346 <*> sshOptionsParser
347 <*> targetParser
348 <*> switch ( long "status"
349 & help "Don't output anything, status code shows success"
350 )
352commandParser :: Parser Command
353commandParser = subparser
354 ( command "list" (info listParser
355 (progDesc "List Azure cloud services"))
356 & command "install" (info copyToParser
357 (progDesc "Install the executable"))
358 & command "md5" (info checkMD5Parser
359 (progDesc "Check if the remote and local MD5 hash match"))
360 & command "run" (info runOnParser
361 (progDesc "Run the executable"))
362 & command "onvm" (info onVmCommandParser
363 (progDesc "Commands used when running ON the vm (usually used internally only)"))
364 )
366runOnParser :: Parser Command
367runOnParser = RunOn
368 <$> azureOptionsParser
369 <*> sshOptionsParser
370 <*> targetParser
371 <*> strOption ( long "port"
372 & metavar "PORT"
373 & help "Port number of the CH instance"
374 )
375 <*> strOption ( long "closure"
376 & metavar "PROC"
377 & help "Process to run on the CH instance"
378 )
379 <*> switch ( long "background"
380 & help "Run the process in the background"
381 )
383onVmRunParser :: Parser OnVmCommand
384onVmRunParser = OnVmRun
385 <$> strOption ( long "host"
386 & metavar "IP"
387 & help "IP address"
388 )
389 <*> strOption ( long "port"
390 & metavar "PORT"
391 & help "port number"
392 )
393 <*> strOption ( long "cloud-service"
394 & metavar "CS"
395 & help "Cloud service name"
396 )
397 <*> switch ( long "background"
398 & help "Run the process in the background"
399 )
401onVmCommandParser :: Parser Command
402onVmCommandParser = OnVmCommand <$> subparser
403 ( command "run" (info onVmRunParser
404 (progDesc "Run the executable"))
405 )