git.iko.soy/distributed-process
on ryan help sign in
distributed-process/distributed-process-azure/src/Control/Distributed/Process/Backend/Azure/GenericMain.hs
364 lines 11 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 )
17import Data.Binary (decode)
18import qualified Data.ByteString.Lazy as BSL (getContents, length)
19import Control.Monad (unless, forM, forM_, join, void)
20import Control.Exception (throwIO, SomeException, evaluate)
21import Control.Applicative ((<$>), (<*>), (<|>))
22import Control.Monad.IO.Class (liftIO)
24-- Posix
25import System.Posix.Process (forkProcess, createSession)
27-- SSH
28import qualified Network.SSH.Client.LibSSH2.Foreign as SSH
29 ( initialize
30 , exit
31 )
33-- Command line options
34import Options.Applicative
35 ( Parser
36 , strOption
37 , long
38 , (&)
39 , metavar
40 , help
41 , subparser
42 , command
43 , info
44 , progDesc
45 , execParser
46 , helper
47 , fullDesc
48 , header
49 , switch
50 )
52-- CH
53import Control.Distributed.Process
54 ( RemoteTable
55 , Closure
56 , Process
57 , unClosure
58 , Static
59 , catch
60 )
61import Control.Distributed.Process.Node (newLocalNode, runProcess, initRemoteTable)
62import Control.Distributed.Process.Serializable (Serializable)
63import Control.Distributed.Process.Closure (SerializableDict)
64import Network.Transport.TCP (createTransport, defaultTCPParameters)
65import Control.Distributed.Process.Backend.Azure
66 ( AzureParameters(azureSshUserName)
67 , defaultAzureParameters
68 , initializeBackend
69 , cloudServices
70 , CloudService(cloudServiceName, cloudServiceVMs)
71 , VirtualMachine(vmName)
72 , Backend(copyToVM, checkMD5, callOnVM, spawnOnVM)
73 )
74import qualified Control.Distributed.Process.Backend.Azure as Azure (remoteTable)
76--------------------------------------------------------------------------------
77-- Main --
78--------------------------------------------------------------------------------
80data ProcessPair b = forall a. Serializable a => ProcessPair {
81 ppairRemote :: Closure (Process a)
82 , ppairLocal :: a -> IO b
83 , ppairDict :: Static (SerializableDict a)
84 }
86genericMain :: (RemoteTable -> RemoteTable) -- ^ Standard CH remote table
87 -> (String -> IO (ProcessPair ())) -- ^ Closures to support in 'run'
88 -> (String -> IO (Closure (Process ()))) -- ^ Closures to support in @run --background@
89 -> IO ()
90genericMain remoteTable callable spawnable = do
91 _ <- SSH.initialize True
92 cmd <- execParser opts
93 case cmd of
94 List {} -> do
95 params <- azureParameters (azureOptions cmd) Nothing
96 backend <- initializeBackend params
97 css <- cloudServices backend
98 mapM_ print css
99 CopyTo {} -> do
100 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
101 backend <- initializeBackend params
102 css <- cloudServices backend
103 forM_ (findTarget (target cmd) css) $ \vm -> do
104 putStr (vmName vm ++ ": ") >> hFlush stdout
105 copyToVM backend vm
106 putStrLn "Done"
107 CheckMD5 {} -> do
108 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
109 backend <- initializeBackend params
110 css <- cloudServices backend
111 matches <- forM (findTarget (target cmd) css) $ \vm -> do
112 unless (status cmd) $ putStr (vmName vm ++ ": ") >> hFlush stdout
113 match <- checkMD5 backend vm
114 unless (status cmd) $ putStrLn $ if match then "OK" else "FAILED"
115 return match
116 if and matches
117 then exitSuccess
118 else exitFailure
119 RunOn {} | background cmd -> do
120 rProc <- spawnable (closureId cmd)
121 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
122 backend <- initializeBackend params
123 css <- cloudServices backend
124 forM_ (findTarget (target cmd) css) $ \vm -> do
125 putStr (vmName vm ++ ": ") >> hFlush stdout
126 spawnOnVM backend vm (remotePort cmd) rProc
127 putStrLn "OK"
128 RunOn {} {- not (background cmd) -} -> do
129 procPair <- callable (closureId cmd)
130 params <- azureParameters (azureOptions cmd) (Just (sshOptions cmd))
131 backend <- initializeBackend params
132 css <- cloudServices backend
133 forM_ (findTarget (target cmd) css) $ \vm -> do
134 putStr (vmName vm ++ ": ") >> hFlush stdout
135 case procPair of
136 ProcessPair rProc lProc dict -> do
137 result <- callOnVM backend dict vm (remotePort cmd) rProc
138 lProc result
139 OnVmCommand (vmCmd@OnVmRun {}) ->
140 onVmRun (remoteTable . Azure.remoteTable $ initRemoteTable)
141 (onVmIP vmCmd)
142 (onVmPort vmCmd)
143 (onVmBackground vmCmd)
144 SSH.exit
145 where
146 opts = info (helper <*> commandParser)
147 ( fullDesc
148 & header "Cloud Haskell backend for Azure"
149 )
151findTarget :: Target -> [CloudService] -> [VirtualMachine]
152findTarget (CloudService cs) css =
153 concatMap cloudServiceVMs . filter ((== cs) . cloudServiceName) $ css
154findTarget (VirtualMachine virtualMachine) css =
155 [ vm | vm <- concatMap cloudServiceVMs css
156 , vmName vm == virtualMachine
157 ]
159azureParameters :: AzureOptions -> Maybe SshOptions -> IO AzureParameters
160azureParameters opts Nothing =
161 defaultAzureParameters (subscriptionId opts)
162 (pathToCert opts)
163 (pathToKey opts)
164azureParameters opts (Just sshOpts) = do
165 params <- defaultAzureParameters (subscriptionId opts)
166 (pathToCert opts)
167 (pathToKey opts)
168 return params {
169 azureSshUserName = remoteUser sshOpts
170 }
172onVmRun :: RemoteTable -> String -> String -> Bool -> IO ()
173onVmRun rtable host port bg = do
174 hSetBinaryMode stdin True
175 hSetBinaryMode stdout True
176 procEnc <- BSL.getContents
177 -- Force evaluation (so that we can safely close stdin)
178 _length <- evaluate (BSL.length procEnc)
179 let proc = decode procEnc
180 if bg
181 then do
182 hClose stdin
183 hClose stdout
184 hClose stderr
185 void . forkProcess $ do
186 void createSession
187 startCH proc
188 else
189 startCH proc
190 where
191 startCH :: Closure (Process ()) -> IO ()
192 startCH proc = do
193 mTransport <- createTransport host port defaultTCPParameters
194 case mTransport of
195 Left err -> throwIO err
196 Right transport -> do
197 node <- newLocalNode transport rtable
198 runProcess node $
199 catch (join . unClosure $ proc)
200 (\e -> liftIO (print (e :: SomeException) >> throwIO e))
202--------------------------------------------------------------------------------
203-- Command line options --
204--------------------------------------------------------------------------------
206data AzureOptions = AzureOptions {
207 subscriptionId :: String
208 , pathToCert :: FilePath
209 , pathToKey :: FilePath
210 }
211 deriving Show
213data SshOptions = SshOptions {
214 remoteUser :: String
215 }
216 deriving Show
218data Target =
219 VirtualMachine String
220 | CloudService String
221 deriving Show
223data Command =
224 List {
225 azureOptions :: AzureOptions
226 }
227 | CopyTo {
228 azureOptions :: AzureOptions
229 , sshOptions :: SshOptions
230 , target :: Target
231 }
232 | CheckMD5 {
233 azureOptions :: AzureOptions
234 , sshOptions :: SshOptions
235 , target :: Target
236 , status :: Bool
237 }
238 | RunOn {
239 azureOptions :: AzureOptions
240 , sshOptions :: SshOptions
241 , target :: Target
242 , remotePort :: String
243 , closureId :: String
244 , background :: Bool
245 }
246 | OnVmCommand {
247 onVmCommand :: OnVmCommand
248 }
249 deriving Show
251data OnVmCommand =
252 OnVmRun {
253 onVmIP :: String
254 , onVmPort :: String
255 , onVmBackground :: Bool
256 }
257 deriving Show
259azureOptionsParser :: Parser AzureOptions
260azureOptionsParser = AzureOptions
261 <$> strOption ( long "subscription-id"
262 & metavar "SID"
263 & help "Azure subscription ID"
264 )
265 <*> strOption ( long "certificate"
266 & metavar "CERT"
267 & help "X509 certificate"
268 )
269 <*> strOption ( long "private"
270 & metavar "PRI"
271 & help "Private key in PKCS#1 format"
272 )
274sshOptionsParser :: Parser SshOptions
275sshOptionsParser = SshOptions
276 <$> strOption ( long "user"
277 & metavar "USER"
278 & help "Remove SSH username"
279 )
281listParser :: Parser Command
282listParser = List <$> azureOptionsParser
284copyToParser :: Parser Command
285copyToParser = CopyTo
286 <$> azureOptionsParser
287 <*> sshOptionsParser
288 <*> targetParser
290targetParser :: Parser Target
291targetParser =
292 ( VirtualMachine <$> strOption ( long "virtual-machine"
293 & metavar "VM"
294 & help "Virtual machine name"
295 )
296 )
297 <|>
298 ( CloudService <$> strOption ( long "cloud-service"
299 & metavar "CS"
300 & help "Cloud service name"
301 )
302 )
304checkMD5Parser :: Parser Command
305checkMD5Parser = CheckMD5
306 <$> azureOptionsParser
307 <*> sshOptionsParser
308 <*> targetParser
309 <*> switch ( long "status"
310 & help "Don't output anything, status code shows success"
311 )
313commandParser :: Parser Command
314commandParser = subparser
315 ( command "list" (info listParser
316 (progDesc "List Azure cloud services"))
317 & command "install" (info copyToParser
318 (progDesc "Install the executable"))
319 & command "md5" (info checkMD5Parser
320 (progDesc "Check if the remote and local MD5 hash match"))
321 & command "run" (info runOnParser
322 (progDesc "Run the executable"))
323 & command "onvm" (info onVmCommandParser
324 (progDesc "Commands used when running ON the vm (usually used internally only)"))
325 )
327runOnParser :: Parser Command
328runOnParser = RunOn
329 <$> azureOptionsParser
330 <*> sshOptionsParser
331 <*> targetParser
332 <*> strOption ( long "port"
333 & metavar "PORT"
334 & help "Port number of the CH instance"
335 )
336 <*> strOption ( long "closure"
337 & metavar "PROC"
338 & help "Process to run on the CH instance"
339 )
340 <*> switch ( long "background"
341 & help "Run the process in the background"
342 )
344onVmRunParser :: Parser OnVmCommand
345onVmRunParser = OnVmRun
346 <$> strOption ( long "host"
347 & metavar "IP"
348 & help "IP address"
349 )
350 <*> strOption ( long "port"
351 & metavar "PORT"
352 & help "port number"
353 )
354 <*> switch ( long "background"
355 & help "Run the process in the background"
356 )
358onVmCommandParser :: Parser Command
359onVmCommandParser = OnVmCommand <$> subparser
360 ( command "run" (info onVmRunParser
361 (progDesc "Run the executable"))
362 )