Add reading the file names from env vars
[ganeti-local] / hbal.hs
diff --git a/hbal.hs b/hbal.hs
index 3d36700..bcd517a 100644 (file)
--- a/hbal.hs
+++ b/hbal.hs
@@ -6,6 +6,7 @@ module Main (main) where
 
 import Data.List
 import Data.Function
+import Data.Maybe (isJust, fromJust, fromMaybe)
 import Monad
 import System
 import System.IO
@@ -16,35 +17,44 @@ import Text.Printf (printf)
 
 import qualified Ganeti.HTools.Container as Container
 import qualified Ganeti.HTools.Cluster as Cluster
-import qualified Ganeti.HTools.Version as Version
+import qualified Ganeti.HTools.Node as Node
+import qualified Ganeti.HTools.CLI as CLI
 import Ganeti.HTools.Rapi
 import Ganeti.HTools.Utils
 
 -- | Command line options structure.
 data Options = Options
-    { optShowNodes :: Bool     -- ^ Whether to show node status
-    , optShowCmds  :: Bool     -- ^ Whether to show the command list
-    , optOneline   :: Bool     -- ^ Switch output to a single line
-    , optNodef     :: FilePath -- ^ Path to the nodes file
-    , optInstf     :: FilePath -- ^ Path to the instances file
-    , optMaxLength :: Int      -- ^ Stop after this many steps
-    , optMaster    :: String   -- ^ Collect data from RAPI
-    , optVerbose   :: Int      -- ^ Verbosity level
-    , optShowVer   :: Bool     -- ^ Just show the program version
+    { optShowNodes :: Bool           -- ^ Whether to show node status
+    , optShowCmds  :: Maybe FilePath -- ^ Whether to show the command list
+    , optOneline   :: Bool           -- ^ Switch output to a single line
+    , optNodef     :: FilePath       -- ^ Path to the nodes file
+    , optNodeSet   :: Bool           -- ^ The nodes have been set by options
+    , optInstf     :: FilePath       -- ^ Path to the instances file
+    , optInstSet   :: Bool           -- ^ The insts have been set by options
+    , optMaxLength :: Int            -- ^ Stop after this many steps
+    , optMaster    :: String         -- ^ Collect data from RAPI
+    , optVerbose   :: Int            -- ^ Verbosity level
+    , optOffline   :: [String]       -- ^ Names of offline nodes
+    , optShowVer   :: Bool           -- ^ Just show the program version
+    , optShowHelp  :: Bool           -- ^ Just show the help
     } deriving Show
 
 -- | Default values for the command line options.
 defaultOptions :: Options
 defaultOptions  = Options
  { optShowNodes = False
- , optShowCmds  = False
+ , optShowCmds  = Nothing
  , optOneline   = False
  , optNodef     = "nodes"
+ , optNodeSet   = False
  , optInstf     = "instances"
+ , optInstSet   = False
  , optMaxLength = -1
  , optMaster    = ""
  , optVerbose   = 0
+ , optOffline   = []
  , optShowVer   = False
+ , optShowHelp  = False
  }
 
 -- | Options list and functions
@@ -54,16 +64,19 @@ options =
       (NoArg (\ opts -> opts { optShowNodes = True }))
       "print the final node list"
     , Option ['C']     ["print-commands"]
-      (NoArg (\ opts -> opts { optShowCmds = True }))
-      "print the ganeti command list for reaching the solution"
+      (OptArg ((\ f opts -> opts { optShowCmds = Just f }) . fromMaybe "-")
+                  "FILE")
+      "print the ganeti command list for reaching the solution,\
+      \if an argument is passed then write the commands to a file named\
+      \ as such"
     , Option ['o']     ["oneline"]
       (NoArg (\ opts -> opts { optOneline = True }))
       "print the ganeti command list for reaching the solution"
     , Option ['n']     ["nodes"]
-      (ReqArg (\ f opts -> opts { optNodef = f }) "FILE")
+      (ReqArg (\ f opts -> opts { optNodef = f, optNodeSet = True }) "FILE")
       "the node list FILE"
     , Option ['i']     ["instances"]
-      (ReqArg (\ f opts -> opts { optInstf =  f }) "FILE")
+      (ReqArg (\ f opts -> opts { optInstf =  f, optInstSet = True }) "FILE")
       "the instance list FILE"
     , Option ['m']     ["master"]
       (ReqArg (\ m opts -> opts { optMaster = m }) "ADDRESS")
@@ -73,33 +86,27 @@ options =
       "cap the solution at this many moves (useful for very unbalanced \
       \clusters)"
     , Option ['v']     ["verbose"]
-      (NoArg (\ opts -> let nv = (optVerbose opts)
-                        in opts { optVerbose = nv + 1 }))
+      (NoArg (\ opts -> opts { optVerbose = (optVerbose opts) + 1 }))
       "increase the verbosity level"
+    , Option ['O']     ["offline"]
+      (ReqArg (\ n opts -> opts { optOffline = n:optOffline opts }) "NODE")
+       " set node as offline"
     , Option ['V']     ["version"]
       (NoArg (\ opts -> opts { optShowVer = True}))
       "show the version of the program"
+    , Option ['h']     ["help"]
+      (NoArg (\ opts -> opts { optShowHelp = True}))
+      "show help"
     ]
 
--- | Command line parser, using the 'options' structure.
-parseOpts :: [String] -> IO (Options, [String])
-parseOpts argv =
-    case getOpt Permute options argv of
-      (o,n,[]  ) ->
-          return (foldl (flip id) defaultOptions o, n)
-      (_,_,errs) ->
-          ioError (userError (concat errs ++ usageInfo header options))
-      where header = printf "hbal %s\nUsage: hbal [OPTION...]"
-                     Version.version
-
 {- | Start computing the solution at the given depth and recurse until
 we find a valid solution or we exceed the maximum depth.
 
 -}
 iterateDepth :: Cluster.Table    -- ^ The starting table
              -> Int              -- ^ Remaining length
-             -> [(Int, String)]  -- ^ Node idx to name list
-             -> [(Int, String)]  -- ^ Inst idx to name list
+             -> Cluster.NameList -- ^ Node idx to name list
+             -> Cluster.NameList -- ^ Inst idx to name list
              -> Int              -- ^ Max node name len
              -> Int              -- ^ Max instance name len
              -> [[String]]       -- ^ Current command list
@@ -109,7 +116,8 @@ iterateDepth :: Cluster.Table    -- ^ The starting table
 iterateDepth ini_tbl max_rounds ktn kti nmlen imlen cmd_strs oneline =
     let Cluster.Table ini_nl ini_il ini_cv ini_plc = ini_tbl
         all_inst = Container.elems ini_il
-        node_idx = Container.keys ini_nl
+        node_idx = map Node.idx . filter (not . Node.offline) $
+                   Container.elems ini_nl
         fin_tbl = Cluster.checkMove node_idx ini_tbl all_inst
         (Cluster.Table _ _ fin_cv fin_plc) = fin_tbl
         ini_plc_len = length ini_plc
@@ -137,22 +145,52 @@ iterateDepth ini_tbl max_rounds ktn kti nmlen imlen cmd_strs oneline =
 main :: IO ()
 main = do
   cmd_args <- System.getArgs
-  (opts, _) <- parseOpts cmd_args
+  (opts, args) <- CLI.parseOpts cmd_args "hbal" options
+                  defaultOptions optShowHelp
+
+  unless (null args) $ do
+         hPutStrLn stderr "Error: this program doesn't take any arguments."
+         exitWith $ ExitFailure 1
 
   when (optShowVer opts) $ do
-         putStr $ showVersion "hbal"
+         putStr $ CLI.showVersion "hbal"
          exitWith ExitSuccess
 
-  let oneline = optOneline opts
+  (env_node, env_inst) <- CLI.parseEnv ()
+  let nodef = if optNodeSet opts then optNodef opts
+              else env_node
+      instf = if optInstSet opts then optInstf opts
+              else env_inst
+      oneline = optOneline opts
       verbose = optVerbose opts
       (node_data, inst_data) =
           case optMaster opts of
-            "" -> (readFile $ optNodef opts,
-                   readFile $ optInstf opts)
+            "" -> (readFile nodef,
+                   readFile instf)
             host -> (readData getNodes host,
                      readData getInstances host)
 
-  (nl, il, csf, ktn, kti) <- liftM2 Cluster.loadData node_data inst_data
+  (loaded_nl, il, csf, ktn, kti) <- liftM2 Cluster.loadData node_data inst_data
+  let (fix_msgs, fixed_nl) = Cluster.checkData loaded_nl il ktn kti
+
+  unless (null fix_msgs) $ do
+         putStrLn "Warning: cluster has inconsistent data:"
+         putStrLn . unlines . map (\s -> printf "  - %s" s) $ fix_msgs
+
+  let offline_names = optOffline opts
+      all_names = snd . unzip $ ktn
+      offline_wrong = filter (\n -> not $ elem n all_names) offline_names
+      offline_indices = fst . unzip .
+                        filter (\(_, n) -> elem n offline_names) $ ktn
+
+  when (length offline_wrong > 0) $ do
+         printf "Wrong node name(s) set as offline: %s\n"
+                (commaJoin offline_wrong)
+         exitWith $ ExitFailure 1
+
+  let nl = Container.map (\n -> if elem (Node.idx n) offline_indices
+                                then Node.setOffline n True
+                                else n) fixed_nl
 
   unless oneline $ printf "Loaded %d nodes, %d instances\n"
              (Container.size nl)
@@ -206,11 +244,21 @@ main = do
   unless (oneline || verbose == 0) $
          printf "Solution length=%d\n" (length ord_plc)
 
-  when (optShowCmds opts) $
+  let cmd_data = Cluster.formatCmds . reverse $ cmd_strs
+
+  when (isJust $ optShowCmds opts) $
        do
+         let out_path = fromJust $ optShowCmds opts
          putStrLn ""
-         putStrLn "Commands to run to reach the above solution:"
-         putStr . Cluster.formatCmds . reverse $ cmd_strs
+         (if out_path == "-" then
+              printf "Commands to run to reach the above solution:\n%s"
+                     (unlines . map ("  " ++) .
+                      filter (/= "check") .
+                      lines $ cmd_data)
+          else do
+            writeFile out_path (CLI.shTemplate ++ cmd_data)
+            printf "The commands have been written to file '%s'\n" out_path)
+
   when (optShowNodes opts) $
        do
          let (orig_mem, orig_disk) = Cluster.totalResources nl