github.com/koron/hk@v0.0.0-20150303213137-b8aeaa3ab34c/scale.go (about)

     1  package main
     2  
     3  import (
     4  	"errors"
     5  	"log"
     6  	"os"
     7  	"sort"
     8  	"strconv"
     9  	"strings"
    10  
    11  	"github.com/heroku/hk/Godeps/_workspace/src/github.com/bgentry/heroku-go"
    12  )
    13  
    14  var cmdScale = &Command{
    15  	Run:      runScale,
    16  	Usage:    "scale <type>=[<qty>]:[<size>]...",
    17  	NeedsApp: true,
    18  	Category: "dyno",
    19  	Short:    "change dyno quantities and sizes",
    20  	Long: `
    21  Scale changes the quantity of dynos (horizontal scale) and/or the
    22  dyno size (vertical scale) for each process type. Note that
    23  changing dyno size will restart all dynos of that type.
    24  
    25  Examples:
    26  
    27      $ hk scale web=2
    28      Scaled myapp to web=2:1X.
    29  
    30      $ hk scale web=2:1X worker=5:2X
    31      Scaled myapp to web=2:1X, worker=5:2X.
    32  
    33      $ hk scale web=PX worker=1X
    34      Scaled myapp to web=2:PX, worker=5:1X.
    35  `,
    36  }
    37  
    38  // takes args of the form "web=1", "worker=3X", web=4:2X etc
    39  func runScale(cmd *Command, args []string) {
    40  	appname := mustApp()
    41  	if len(args) == 0 {
    42  		cmd.PrintUsage()
    43  		os.Exit(2)
    44  	}
    45  	todo := make([]heroku.FormationBatchUpdateOpts, len(args))
    46  	types := make(map[string]bool)
    47  	for i, arg := range args {
    48  		pstype, qty, size, err := parseScaleArg(arg)
    49  		if err != nil {
    50  			cmd.PrintUsage()
    51  			os.Exit(2)
    52  		}
    53  		if _, exists := types[pstype]; exists {
    54  			// can only specify each process type once
    55  			printError("process type '%s' specified more than once", pstype)
    56  			cmd.PrintUsage()
    57  			os.Exit(2)
    58  		}
    59  		types[pstype] = true
    60  
    61  		opt := heroku.FormationBatchUpdateOpts{Process: pstype}
    62  		if qty != -1 {
    63  			opt.Quantity = &qty
    64  		}
    65  		if size != "" {
    66  			opt.Size = &size
    67  		}
    68  		todo[i] = opt
    69  	}
    70  
    71  	formations, err := client.FormationBatchUpdate(appname, todo)
    72  	must(err)
    73  
    74  	sortedFormations := formationsByType(formations)
    75  	sort.Sort(sortedFormations)
    76  	results := make([]string, len(types))
    77  	rindex := 0
    78  	for _, f := range sortedFormations {
    79  		if _, exists := types[f.Type]; exists {
    80  			results[rindex] = f.Type + "=" + strconv.Itoa(f.Quantity) + ":" + f.Size
    81  			rindex += 1
    82  		}
    83  	}
    84  	log.Printf("Scaled %s to %s.", appname, strings.Join(results, ", "))
    85  }
    86  
    87  var errInvalidScaleArg = errors.New("invalid argument")
    88  
    89  func parseScaleArg(arg string) (pstype string, qty int, size string, err error) {
    90  	qty = -1
    91  	iEquals := strings.IndexRune(arg, '=')
    92  	if fields := strings.Fields(arg); len(fields) > 1 || iEquals == -1 {
    93  		err = errInvalidScaleArg
    94  		return
    95  	}
    96  	pstype = arg[:iEquals]
    97  
    98  	rem := strings.ToUpper(arg[iEquals+1:])
    99  	if len(rem) == 0 {
   100  		err = errInvalidScaleArg
   101  		return
   102  	}
   103  
   104  	if iColon := strings.IndexRune(rem, ':'); iColon == -1 {
   105  		if iX := strings.IndexRune(rem, 'X'); iX == -1 {
   106  			qty, err = strconv.Atoi(rem)
   107  			if err != nil {
   108  				return pstype, -1, "", errInvalidScaleArg
   109  			}
   110  		} else {
   111  			size = rem
   112  		}
   113  	} else {
   114  		if iColon > 0 {
   115  			qty, err = strconv.Atoi(rem[:iColon])
   116  			if err != nil {
   117  				return pstype, -1, "", errInvalidScaleArg
   118  			}
   119  		}
   120  		if len(rem) > iColon+1 {
   121  			size = rem[iColon+1:]
   122  		}
   123  	}
   124  	if err != nil || qty == -1 && size == "" {
   125  		err = errInvalidScaleArg
   126  	}
   127  	return
   128  }
   129  
   130  type formationsByType []heroku.Formation
   131  
   132  func (f formationsByType) Len() int           { return len(f) }
   133  func (f formationsByType) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }
   134  func (f formationsByType) Less(i, j int) bool { return f[i].Type < f[j].Type }