func main() { // Prefixer accepts anything that implements io.Reader interface prefixReader := prefixer.New(os.Stdin, "> ") // Read all prefixed lines from STDIN into a buffer buffer, _ := ioutil.ReadAll(prefixReader) // Write buffer to STDOUT os.Stdout.Write(buffer) }
// Run runs set of commands on multiple hosts defined by network sequentially. // TODO: This megamoth method needs a big refactor and should be split // to multiple smaller methods. func (sup *Stackup) Run(network *Network, commands ...*Command) error { if len(commands) == 0 { return errors.New("no commands to be run") } // Process all ENVs into a string of form // `export FOO="bar"; export BAR="baz";`. env := `` for _, v := range append(sup.conf.Env, network.Env...) { env += v.AsExport() + " " } // Create clients for every host (either SSH or Localhost). var ( clients []Client bastion *SSHClient ) if network.Bastion != "" { bastion = &SSHClient{} if err := bastion.Connect(network.Bastion); err != nil { return err } } for i, host := range network.Hosts { // Localhost client. if host == "localhost" { local := &LocalhostClient{ env: env + `export SUP_HOST="` + host + `";`, } if err := local.Connect(host); err != nil { return err } clients = append(clients, local) continue } // SSH client. remote := &SSHClient{ env: env + `export SUP_HOST="` + host + `";`, color: Colors[i%len(Colors)], } if bastion != nil { if err := remote.ConnectWith(host, bastion.DialThrough); err != nil { return err } } else { if err := remote.Connect(host); err != nil { return err } } defer remote.Close() clients = append(clients, remote) } maxLen := 0 for _, c := range clients { _, prefixLen := c.Prefix() if prefixLen > maxLen { maxLen = prefixLen } } // Run command or run multiple commands defined by target sequentially. for _, cmd := range commands { // Translate command into task(s). tasks, err := CreateTasks(cmd, clients, env) if err != nil { return fmt.Errorf("CreateTasks(): %s", err) } // Run tasks sequentially. for _, task := range tasks { var writers []io.Writer var wg sync.WaitGroup // Run tasks on the provided clients. for _, c := range task.Clients { prefix, prefixLen := c.Prefix() if len(prefix) < maxLen { // Left padding. prefix = strings.Repeat(" ", maxLen-prefixLen) + prefix } err := c.Run(task) if err != nil { return fmt.Errorf("%s%v", prefix, err) } // Copy over tasks's STDOUT. wg.Add(1) go func(c Client) { defer wg.Done() _, err := io.Copy(os.Stdout, prefixer.New(c.Stdout(), prefix)) if err != nil && err != io.EOF { // TODO: io.Copy() should not return io.EOF at all. // Upstream bug? Or prefixer.WriteTo() bug? fmt.Fprintf(os.Stderr, "%sSTDOUT: %v", prefix, err) } }(c) // Copy over tasks's STDERR. wg.Add(1) go func(c Client) { defer wg.Done() _, err := io.Copy(os.Stderr, prefixer.New(c.Stderr(), prefix)) if err != nil && err != io.EOF { fmt.Fprintf(os.Stderr, "%sSTDERR: %v", prefix, err) } }(c) writers = append(writers, c.Stdin()) } // Copy over task's STDIN. if task.Input != nil { go func() { writer := io.MultiWriter(writers...) _, err := io.Copy(writer, task.Input) if err != nil && err != io.EOF { fmt.Fprintf(os.Stderr, "STDIN: %v", err) } // TODO: Use MultiWriteCloser (not in Stdlib), so we can writer.Close() instead? for _, c := range clients { c.WriteClose() } }() } // Wait for all I/O operations first. wg.Wait() // Make sure each client finishes the task, return on failure. for _, c := range task.Clients { wg.Add(1) go func(c Client) { defer wg.Done() if err := c.Wait(); err != nil { prefix, prefixLen := c.Prefix() if len(prefix) < maxLen { // Left padding. prefix = strings.Repeat(" ", maxLen-prefixLen) + prefix } if e, ok := err.(*ssh.ExitError); ok && e.ExitStatus() != 15 { // TODO: Store all the errors, and print them after Wait(). fmt.Fprintf(os.Stderr, "%sexit %v\n", prefix, e.ExitStatus()) os.Exit(e.ExitStatus()) } fmt.Fprintf(os.Stderr, "%s%v\n", prefix, err) os.Exit(1) } }(c) } // Wait for all commands to finish. wg.Wait() } } return nil }