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

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"crypto/tls"
     6  	"io"
     7  	"log"
     8  	"net/url"
     9  	"os"
    10  	"os/signal"
    11  	"strconv"
    12  	"strings"
    13  	"syscall"
    14  
    15  	"github.com/heroku/hk/Godeps/_workspace/src/github.com/bgentry/heroku-go"
    16  	"github.com/heroku/hk/term"
    17  )
    18  
    19  var (
    20  	detachedRun bool
    21  	dynoSize    string
    22  )
    23  
    24  var cmdRun = &Command{
    25  	Run:      runRun,
    26  	Usage:    "run [-s <size>] [-d] <command> [<argument>...]",
    27  	NeedsApp: true,
    28  	Category: "dyno",
    29  	Short:    "run a process in a dyno",
    30  	Long: `
    31  Run a process on Heroku. Flags such as` + " `-a` " + `may be parsed out of
    32  the command unless the command is quoted or provided after a
    33  double-dash (--).
    34  
    35  Options:
    36  
    37      -s <size>  set the size for this dyno (e.g. 2X)
    38      -d         run in detached mode instead of attached to terminal
    39  
    40  Examples:
    41  
    42      $ hk run echo "hello"
    43      Running ` + "`echo \"hello\"`" + ` on myapp as run.1234:
    44      "hello"
    45  
    46      $ hk run console
    47      Running ` + "`console`" + ` on myapp as run.5678:
    48      Loading production environment (Rails 3.2.14)
    49      irb(main):001:0> ...
    50  
    51      $ hk run -d -s 2X bin/my_worker
    52      Ran ` + "`bin/my_worker`" + ` on myapp as run.4321, detached.
    53  
    54      $ hk run -a myapp -- ls -a /
    55      Running ` + "`ls -a bin /`" + ` on myapp as run.8650:
    56      /:
    57      .  ..  app  bin  dev  etc  home  lib  lib64  lost+found  proc  sbin  tmp  usr  var
    58  `,
    59  }
    60  
    61  func init() {
    62  	cmdRun.Flag.BoolVarP(&detachedRun, "detached", "d", false, "detached")
    63  	cmdRun.Flag.StringVarP(&dynoSize, "size", "s", "", "dyno size")
    64  }
    65  
    66  func runRun(cmd *Command, args []string) {
    67  	if len(args) == 0 {
    68  		cmd.PrintUsage()
    69  		os.Exit(2)
    70  	}
    71  	appname := mustApp()
    72  
    73  	cols, err := term.Cols()
    74  	if err != nil {
    75  		printFatal(err.Error())
    76  	}
    77  	lines, err := term.Lines()
    78  	if err != nil {
    79  		printFatal(err.Error())
    80  	}
    81  
    82  	attached := !detachedRun
    83  	opts := heroku.DynoCreateOpts{Attach: &attached}
    84  	if attached {
    85  		env := map[string]string{
    86  			"COLUMNS": strconv.Itoa(cols),
    87  			"LINES":   strconv.Itoa(lines),
    88  			"TERM":    os.Getenv("TERM"),
    89  		}
    90  		opts.Env = &env
    91  	}
    92  	if dynoSize != "" {
    93  		if !strings.HasSuffix(dynoSize, "X") {
    94  			cmd.PrintUsage()
    95  			os.Exit(2)
    96  		}
    97  		opts.Size = &dynoSize
    98  	}
    99  
   100  	command := strings.Join(args, " ")
   101  	dyno, err := client.DynoCreate(appname, command, &opts)
   102  	must(err)
   103  
   104  	if detachedRun {
   105  		log.Printf("Ran `%s` on %s as %s, detached.", dyno.Command, appname, dyno.Name)
   106  		return
   107  	}
   108  	log.Printf("Running `%s` on %s as %s:", dyno.Command, appname, dyno.Name)
   109  
   110  	u, err := url.Parse(*dyno.AttachURL)
   111  	if err != nil {
   112  		printFatal(err.Error())
   113  	}
   114  
   115  	cn, err := tls.Dial("tcp", u.Host, nil)
   116  	if err != nil {
   117  		printFatal(err.Error())
   118  	}
   119  	defer cn.Close()
   120  
   121  	br := bufio.NewReader(cn)
   122  
   123  	_, err = io.WriteString(cn, u.Path[1:]+"\r\n")
   124  	if err != nil {
   125  		printFatal(err.Error())
   126  	}
   127  
   128  	for {
   129  		_, pre, err := br.ReadLine()
   130  		if err != nil {
   131  			printFatal(err.Error())
   132  		}
   133  		if !pre {
   134  			break
   135  		}
   136  	}
   137  
   138  	if term.IsTerminal(os.Stdin) && term.IsTerminal(os.Stdout) {
   139  		err = term.MakeRaw(os.Stdin)
   140  		if err != nil {
   141  			printFatal(err.Error())
   142  		}
   143  		defer term.Restore(os.Stdin)
   144  
   145  		sig := make(chan os.Signal)
   146  		signal.Notify(sig, os.Signal(syscall.SIGQUIT), os.Interrupt)
   147  		go func() {
   148  			defer term.Restore(os.Stdin)
   149  			for sg := range sig {
   150  				switch sg {
   151  				case os.Interrupt:
   152  					cn.Write([]byte{3})
   153  				case os.Signal(syscall.SIGQUIT):
   154  					cn.Write([]byte{28})
   155  				default:
   156  					panic("not reached")
   157  				}
   158  			}
   159  		}()
   160  	}
   161  
   162  	errc := make(chan error)
   163  	cp := func(a io.Writer, b io.Reader) {
   164  		_, err := io.Copy(a, b)
   165  		errc <- err
   166  	}
   167  
   168  	go cp(os.Stdout, br)
   169  	go cp(cn, os.Stdin)
   170  	if err = <-errc; err != nil {
   171  		printFatal(err.Error())
   172  	}
   173  }