使用通道捕获Go​​routine的输出和错误

Question

I have a for-loop that calls a function runCommand() which runs a remote command on a switch and prints the output. The function is called in a goroutine on each iteration and I am using a sync.Waitgroup to synchronize the goroutines. Now, I need a way to capture the output and any errors of my runCommand() function into a channel. I have read many articles and watched a lot of videos on using channels with goroutines, but this is the first time I have ever written a concurrent application and I can't seem to wrap my head around the idea.

Basically, my program takes in a list of hostnames from the command line then asynchronously connects to each host, runs a configuration command on it, and prints the output. It is ok for my program to continue configuring the remaining hosts if one has an error.

How would I idiomatically send the output or error(s) of each call to runCommand() to a channel then receive the output or error(s) for printing?

Here is my code:

package main

import (
    "fmt"
    "golang.org/x/crypto/ssh"
    "os"
    "time"
    "sync"
)

func main() {
    hosts := os.Args[1:]
    clientConf := configureClient("user", "password")

    var wg sync.WaitGroup
    for _, host := range hosts {
        wg.Add(1)
        go runCommand(host, &clientConf, &wg)
    }
    wg.Wait()

    fmt.Println("Configuration complete!")
}

// Run a remote command
func runCommand(host string, config *ssh.ClientConfig, wg *sync.WaitGroup) {
    defer wg.Done()
    // Connect to the client
    client, err := ssh.Dial("tcp", host+":22", config)
    if err != nil {
        fmt.Println(err)
        return
    }
    defer client.Close()
    // Create a session
    session, err := client.NewSession()
    if err != nil {
        fmt.Println(err)
        return
    }
    defer session.Close()
    // Get the session output
    output, err := session.Output("show lldp ne")
    if err != nil {
        fmt.Println(err)
        return
    }
    fmt.Print(string(output))
    fmt.Printf("Connection to %s closed.
", host)
}

// Set up client configuration
func configureClient(user, password string) ssh.ClientConfig {
    var sshConf ssh.Config
    sshConf.SetDefaults()
    // Append supported ciphers
    sshConf.Ciphers = append(sshConf.Ciphers, "aes128-cbc", "aes256-cbc", "3des-cbc", "des-cbc", "aes192-cbc")
    // Create client config
    clientConf := &ssh.ClientConfig{
        Config:          sshConf,
        User:            user,
        Auth:            []ssh.AuthMethod{ssh.Password(password)},
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         time.Second * 5,
    }
    return *clientConf
}

EDIT: I got rid of the Waitgroup, as suggested, and now I need to keep track of which output belongs to which host by printing the hostname before printing its output and printing a Connection to <host> closed. message when the gorouttine completes. For example:

$ go run main.go host1[,host2[,...]]
Connecting to <host1>
[Output]
...
[Error]
Connection to <host1> closed.

Connecting to <host2>
...
Connection to <host2> closed.

Configuration complete!

I know the above won't necessarily process host1 and host2 in order, But I need to print the correct host value for the connecting and closing messages before and after the output/error(s), respectively. I tried defering printing the closing message in the runCommand() function, but the message is printed out before the output/error(s). And printing the closing message in the for-loop after each goroutine call doesn't work as expected either.

Updated code:

package main

import (
    "fmt"
    "golang.org/x/crypto/ssh"
    "os"
    "time"
)

type CmdResult struct {
    Host string
    Output string
    Err error
}

func main() {
    start := time.Now()

    hosts := os.Args[1:]
    clientConf := configureClient("user", "password")
    results := make(chan CmdResult)

    for _, host := range hosts {
        go runCommand(host, &clientConf, results)
    }
    for i := 0; i < len(hosts); i++ {
        output := <- results
        fmt.Println(output.Host)
        if output.Output != "" {
            fmt.Printf("%s
", output.Output)
        }
        if output.Err != nil {
            fmt.Printf("Error: %v
", output.Err)
        }
    }
    fmt.Printf("Configuration complete! [%s]
", time.Since(start).String())
}

// Run a remote command
func runCommand(host string, config *ssh.ClientConfig, ch chan CmdResult) {
    // This is printing before the output/error(s).
    // Does the same when moved to the bottom of this function.
    defer fmt.Printf("Connection to %s closed.
", host)

    // Connect to the client
    client, err := ssh.Dial("tcp", host+":22", config)
    if err != nil {
        ch <- CmdResult{host, "", err}
        return
    }
    defer client.Close()
    // Create a session
    session, err := client.NewSession()
    if err != nil {
        ch <- CmdResult{host, "", err}
        return
    }
    defer session.Close()
    // Get the session output
    output, err := session.Output("show lldp ne")
    if err != nil {
        ch <- CmdResult{host, "", err}
        return
    }
    ch <- CmdResult{host, string(output), nil}
}

// Set up client configuration
func configureClient(user, password string) ssh.ClientConfig {
    var sshConf ssh.Config
    sshConf.SetDefaults()
    // Append supported ciphers
    sshConf.Ciphers = append(sshConf.Ciphers, "aes128-cbc", "aes256-cbc", "3des-cbc", "des-cbc", "aes192-cbc")
    // Create client config
    clientConf := &ssh.ClientConfig{
        Config:          sshConf,
        User:            user,
        Auth:            []ssh.AuthMethod{ssh.Password(password)},
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         time.Second * 5,
    }
    return *clientConf
}

Answer 1

If you use an unbuffered channel, you actually don't need the sync.WaitGroup, because you can call the receive operator on the channel once for every goroutine that will send on the channel. Each receive operation will block until a send statement is ready, resulting in the same behavior as a WaitGroup.

To make this happen, change runCommand to execute a send statement exactly once before the function exits, under all conditions.

First, create a type to send over the channel:

type CommandResult struct {
    Output      string
    Err         error
}

And edit your main() {...} to execute a receive operation on the channel the same number of times as the number of goroutines that will send to the channel:

func main() {
    ch := make(chan CommandResult) // initialize an unbuffered channel
    // rest of your setup
    for _, host := range hosts {
        go runCommand(host, &clientConf, ch) // pass in the channel
    }
    for x := 0; x < len(hosts); x++ {
        fmt.Println(<-ch) // this will block until one is ready to send
    }

And edit your runCommand function to accept the channel, remove references to WaitGroup, and execute the send exactly once under all conditions:

func runCommand(host string, config *ssh.ClientConfig, ch chan CommandResult) {
    // do stuff that generates output, err; then when ready to exit function:
    ch <- CommandResult{output, err}
}

---

EDIT: Question updated with stdout message order requirements

I'd like to get nicely formatted output that ignores the order of events

In this case, remove all print messages from runCommand, you're going to put all output into the element you're passing on the channel so it can be grouped together. Edit the CommandResult type to contain additional fields you want to organize, such as:

type CommandResult struct {
    Host        string
    Output      string
    Err         error
}

If you don't need to sort your results, you can just move on to printing the data received, e.g.

for x := 0; x < len(hosts); x++ {
    r := <-ch
    fmt.Printf("Host: %s----
Output: %s
", r.Host, r.Output)
    if r.Err != nil {
        fmt.Printf("Error: %s
", r.Err)
    }
}

If you do need to sort your results, then in your main goroutine, add the elements received on the channel to a slice:

    ...
    results := make([]CommandResult, 0, len(hosts))
    for x := 0; x < len(hosts); x++ {
        results = append(results, <-ch) // this will block until one is ready to send
    }

Then you can use the sort package in the Go standard library to sort your results for printing. For example, you could sort them alphabetically by host. Or you could put the results into a map with host string as the key instead of a slice to allow you to print in the order of the original host list.