17840: Deduplicate flag-parsing code.
[arvados.git] / lib / costanalyzer / costanalyzer.go
index 3fd8097f82234646a8ed284d5e55fcaaa50973aa..a3673c9794d42d9db2da9b3cded67faaf3be67c4 100644 (file)
@@ -17,6 +17,7 @@ import (
        "strings"
        "time"
 
+       "git.arvados.org/arvados.git/lib/cmd"
        "git.arvados.org/arvados.git/sdk/go/arvados"
        "git.arvados.org/arvados.git/sdk/go/arvadosclient"
        "git.arvados.org/arvados.git/sdk/go/keepclient"
@@ -39,11 +40,16 @@ type nodeInfo struct {
        Preemptible  bool
 }
 
-type containerInfo struct {
+type consumption struct {
        cost     float64
        duration float64
 }
 
+func (c *consumption) Add(n consumption) {
+       c.cost += n.cost
+       c.duration += n.duration
+}
+
 type arrayFlags []string
 
 func (i *arrayFlags) String() string {
@@ -57,10 +63,9 @@ func (i *arrayFlags) Set(value string) error {
        return nil
 }
 
-func (c *command) parseFlags(prog string, args []string, logger *logrus.Logger, stderr io.Writer) (exitCode int, err error) {
+func (c *command) parseFlags(prog string, args []string, logger *logrus.Logger, stderr io.Writer) (ok bool, exitCode int) {
        var beginStr, endStr string
        flags := flag.NewFlagSet("", flag.ContinueOnError)
-       flags.SetOutput(stderr)
        flags.Usage = func() {
                fmt.Fprintf(flags.Output(), `
 Usage:
@@ -130,22 +135,14 @@ Options:
        flags.StringVar(&beginStr, "begin", "", fmt.Sprintf("timestamp `begin` for date range operation (format: %s)", timestampFormat))
        flags.StringVar(&endStr, "end", "", fmt.Sprintf("timestamp `end` for date range operation (format: %s)", timestampFormat))
        flags.BoolVar(&c.cache, "cache", true, "create and use a local disk cache of Arvados objects")
-       err = flags.Parse(args)
-       if err == flag.ErrHelp {
-               err = nil
-               exitCode = 1
-               return
-       } else if err != nil {
-               exitCode = 2
-               return
+       if ok, code := cmd.ParseFlags(flags, prog, args, "[uuid ...]", stderr); !ok {
+               return false, code
        }
        c.uuids = flags.Args()
 
        if (len(beginStr) != 0 && len(endStr) == 0) || (len(beginStr) == 0 && len(endStr) != 0) {
-               flags.Usage()
-               err = fmt.Errorf("When specifying a date range, both begin and end must be specified")
-               exitCode = 2
-               return
+               fmt.Fprintf(stderr, "When specifying a date range, both begin and end must be specified (try -help)\n")
+               return false, 2
        }
 
        if len(beginStr) != 0 {
@@ -153,30 +150,26 @@ Options:
                c.begin, errB = time.Parse(timestampFormat, beginStr)
                c.end, errE = time.Parse(timestampFormat, endStr)
                if (errB != nil) || (errE != nil) {
-                       flags.Usage()
-                       err = fmt.Errorf("When specifying a date range, both begin and end must be of the format %s %+v, %+v", timestampFormat, errB, errE)
-                       exitCode = 2
-                       return
+                       fmt.Fprintf(stderr, "When specifying a date range, both begin and end must be of the format %s %+v, %+v\n", timestampFormat, errB, errE)
+                       return false, 2
                }
        }
 
        if (len(c.uuids) < 1) && (len(beginStr) == 0) {
-               flags.Usage()
-               err = fmt.Errorf("error: no uuid(s) provided")
-               exitCode = 2
-               return
+               fmt.Fprintf(stderr, "error: no uuid(s) provided (try -help)\n")
+               return false, 2
        }
 
        lvl, err := logrus.ParseLevel(*loglevel)
        if err != nil {
-               exitCode = 2
-               return
+               fmt.Fprintf(stderr, "invalid argument to -log-level: %s\n", err)
+               return false, 2
        }
        logger.SetLevel(lvl)
        if !c.cache {
                logger.Debug("Caching disabled")
        }
-       return
+       return true, 0
 }
 
 func ensureDirectory(logger *logrus.Logger, dir string) (err error) {
@@ -194,7 +187,9 @@ func ensureDirectory(logger *logrus.Logger, dir string) (err error) {
        return
 }
 
-func addContainerLine(logger *logrus.Logger, node nodeInfo, cr arvados.ContainerRequest, container arvados.Container) (csv string, cost float64, duration float64) {
+func addContainerLine(logger *logrus.Logger, node nodeInfo, cr arvados.ContainerRequest, container arvados.Container) (string, consumption) {
+       var csv string
+       var containerConsumption consumption
        csv = cr.UUID + ","
        csv += cr.Name + ","
        csv += container.UUID + ","
@@ -222,10 +217,10 @@ func addContainerLine(logger *logrus.Logger, node nodeInfo, cr arvados.Container
                price = node.Price
                size = node.ProviderType
        }
-       cost = delta.Seconds() / 3600 * price
-       duration = delta.Seconds()
-       csv += size + "," + fmt.Sprintf("%+v", node.Preemptible) + "," + strconv.FormatFloat(price, 'f', 8, 64) + "," + strconv.FormatFloat(cost, 'f', 8, 64) + "\n"
-       return
+       containerConsumption.cost = delta.Seconds() / 3600 * price
+       containerConsumption.duration = delta.Seconds()
+       csv += size + "," + fmt.Sprintf("%+v", node.Preemptible) + "," + strconv.FormatFloat(price, 'f', 8, 64) + "," + strconv.FormatFloat(containerConsumption.cost, 'f', 8, 64) + "\n"
+       return csv, containerConsumption
 }
 
 func loadCachedObject(logger *logrus.Logger, file string, uuid string, object interface{}) (reload bool) {
@@ -360,8 +355,8 @@ func getNode(arv *arvadosclient.ArvadosClient, ac *arvados.Client, kc *keepclien
        return
 }
 
-func handleProject(logger *logrus.Logger, uuid string, arv *arvadosclient.ArvadosClient, ac *arvados.Client, kc *keepclient.KeepClient, resultsDir string, cache bool) (cost map[string]containerInfo, err error) {
-       cost = make(map[string]containerInfo)
+func handleProject(logger *logrus.Logger, uuid string, arv *arvadosclient.ArvadosClient, ac *arvados.Client, kc *keepclient.KeepClient, resultsDir string, cache bool) (cost map[string]consumption, err error) {
+       cost = make(map[string]consumption)
 
        var project arvados.Group
        err = loadObject(logger, ac, uuid, uuid, cache, &project)
@@ -394,11 +389,11 @@ func handleProject(logger *logrus.Logger, uuid string, arv *arvadosclient.Arvado
                items := value.([]interface{})
                for _, item := range items {
                        itemMap := item.(map[string]interface{})
-                       crCsv, err := generateCrCsv(logger, itemMap["uuid"].(string), arv, ac, kc, resultsDir, cache)
+                       crInfo, err := generateCrInfo(logger, itemMap["uuid"].(string), arv, ac, kc, resultsDir, cache)
                        if err != nil {
                                return nil, fmt.Errorf("error generating container_request CSV: %s", err.Error())
                        }
-                       for k, v := range crCsv {
+                       for k, v := range crInfo {
                                cost[k] = v
                        }
                }
@@ -408,14 +403,13 @@ func handleProject(logger *logrus.Logger, uuid string, arv *arvadosclient.Arvado
        return
 }
 
-func generateCrCsv(logger *logrus.Logger, uuid string, arv *arvadosclient.ArvadosClient, ac *arvados.Client, kc *keepclient.KeepClient, resultsDir string, cache bool) (cost map[string]containerInfo, err error) {
+func generateCrInfo(logger *logrus.Logger, uuid string, arv *arvadosclient.ArvadosClient, ac *arvados.Client, kc *keepclient.KeepClient, resultsDir string, cache bool) (cost map[string]consumption, err error) {
 
-       cost = make(map[string]containerInfo)
+       cost = make(map[string]consumption)
 
        csv := "CR UUID,CR name,Container UUID,State,Started At,Finished At,Duration in seconds,Compute node type,Preemptible,Hourly node cost,Total cost\n"
        var tmpCsv string
-       var tmpTotalCost, tmpTotalDuration float64
-       var totalCost, totalDuration float64
+       var total, tmpTotal consumption
        logger.Debugf("Processing %s", uuid)
 
        var crUUID = uuid
@@ -458,10 +452,9 @@ func generateCrCsv(logger *logrus.Logger, uuid string, arv *arvadosclient.Arvado
                logger.Errorf("Skipping container request %s: error getting node %s: %s", cr.UUID, cr.UUID, err)
                return nil, nil
        }
-       tmpCsv, totalCost, totalDuration = addContainerLine(logger, topNode, cr, container)
+       tmpCsv, total = addContainerLine(logger, topNode, cr, container)
        csv += tmpCsv
-       totalCost += tmpTotalCost
-       cost[container.UUID] = containerInfo{cost: totalCost, duration: totalDuration}
+       cost[container.UUID] = total
 
        // Find all container requests that have the container we found above as requesting_container_uuid
        var childCrs arvados.ContainerRequestList
@@ -498,14 +491,14 @@ func generateCrCsv(logger *logrus.Logger, uuid string, arv *arvadosclient.Arvado
                if err != nil {
                        return nil, fmt.Errorf("error loading object %s: %s", cr2.ContainerUUID, err)
                }
-               tmpCsv, tmpTotalCost, tmpTotalDuration = addContainerLine(logger, node, cr2, c2)
-               cost[cr2.ContainerUUID] = containerInfo{cost: tmpTotalCost, duration: tmpTotalDuration}
+               tmpCsv, tmpTotal = addContainerLine(logger, node, cr2, c2)
+               cost[cr2.ContainerUUID] = tmpTotal
                csv += tmpCsv
-               totalCost += tmpTotalCost
+               total.Add(tmpTotal)
        }
        logger.Debug("Done collecting child containers")
 
-       csv += "TOTAL,,,,," + strconv.FormatFloat(totalDuration, 'f', 3, 64) + ",,,," + strconv.FormatFloat(totalCost, 'f', 2, 64) + "\n"
+       csv += "TOTAL,,,,,," + strconv.FormatFloat(total.duration, 'f', 3, 64) + ",,,," + strconv.FormatFloat(total.cost, 'f', 2, 64) + "\n"
 
        if resultsDir != "" {
                // Write the resulting CSV file
@@ -521,9 +514,9 @@ func generateCrCsv(logger *logrus.Logger, uuid string, arv *arvadosclient.Arvado
 }
 
 func (c *command) costAnalyzer(prog string, args []string, logger *logrus.Logger, stdout, stderr io.Writer) (exitcode int, err error) {
-       exitcode, err = c.parseFlags(prog, args, logger, stderr)
-
-       if exitcode != 0 {
+       var ok bool
+       ok, exitcode = c.parseFlags(prog, args, logger, stderr)
+       if !ok {
                return
        }
        if c.resultsDir != "" {
@@ -590,7 +583,7 @@ func (c *command) costAnalyzer(prog string, args []string, logger *logrus.Logger
                }
        }()
 
-       cost := make(map[string]containerInfo)
+       cost := make(map[string]consumption)
 
        for uuid := range uuidChannel {
                logger.Debugf("Considering %s", uuid)
@@ -605,15 +598,15 @@ func (c *command) costAnalyzer(prog string, args []string, logger *logrus.Logger
                                cost[k] = v
                        }
                } else if strings.Contains(uuid, "-xvhdp-") || strings.Contains(uuid, "-4zz18-") {
-                       // This is a container request
-                       var crCsv map[string]containerInfo
-                       crCsv, err = generateCrCsv(logger, uuid, arv, ac, kc, c.resultsDir, c.cache)
+                       // This is a container request or collection
+                       var crInfo map[string]consumption
+                       crInfo, err = generateCrInfo(logger, uuid, arv, ac, kc, c.resultsDir, c.cache)
                        if err != nil {
                                err = fmt.Errorf("error generating CSV for uuid %s: %s", uuid, err.Error())
                                exitcode = 2
                                return
                        }
-                       for k, v := range crCsv {
+                       for k, v := range crInfo {
                                cost[k] = v
                        }
                } else if strings.Contains(uuid, "-tpzed-") {
@@ -641,14 +634,13 @@ func (c *command) costAnalyzer(prog string, args []string, logger *logrus.Logger
                csv += "# " + uuid + "\n"
        }
 
-       var totalCost, totalDuration float64
+       var total consumption
        for k, v := range cost {
                csv += k + "," + strconv.FormatFloat(v.duration, 'f', 3, 64) + "," + strconv.FormatFloat(v.cost, 'f', 8, 64) + "\n"
-               totalCost += v.cost
-               totalDuration += v.duration
+               total.Add(v)
        }
 
-       csv += "TOTAL," + strconv.FormatFloat(totalDuration, 'f', 3, 64) + "," + strconv.FormatFloat(totalCost, 'f', 2, 64) + "\n"
+       csv += "TOTAL," + strconv.FormatFloat(total.duration, 'f', 3, 64) + "," + strconv.FormatFloat(total.cost, 'f', 2, 64) + "\n"
 
        if c.resultsDir != "" {
                // Write the resulting CSV file
@@ -663,7 +655,7 @@ func (c *command) costAnalyzer(prog string, args []string, logger *logrus.Logger
        }
 
        // Output the total dollar amount on stdout
-       fmt.Fprintf(stdout, "%s\n", strconv.FormatFloat(totalCost, 'f', 2, 64))
+       fmt.Fprintf(stdout, "%s\n", strconv.FormatFloat(total.cost, 'f', 2, 64))
 
        return
 }