/
main.go
88 lines (76 loc) · 2.04 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
package main
import (
"fmt"
"github.com/discoriver/massh"
"golang.org/x/crypto/ssh"
"sync"
"time"
)
func main() {
j := massh.Job{
Command: "echo \"Hello, World\"",
}
j2 := massh.Job{
Command: "echo \"Hello, World 2\"",
}
j3 := massh.Job{
Command: "echo \"Hello, World 3\"",
}
sshc := &ssh.ClientConfig{
// Fake credentials
User: "u01",
Auth: []ssh.AuthMethod{ssh.Password("password")},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
Timeout: time.Duration(2) * time.Second,
}
cfg := massh.NewConfig()
cfg.SSHConfig = sshc
cfg.JobStack = &[]massh.Job{j, j2, j3}
cfg.WorkerPool = 10
cfg.SetHosts([]string{"192.168.1.118"})
resChan := make(chan massh.Result)
// This should be the last responsibility from the massh package. Handling the Result channel is up to the user.
err := cfg.Stream(resChan)
if err != nil {
panic(err)
}
var wg sync.WaitGroup
numberOfExpectedCompletions := len(cfg.Hosts) * len(*cfg.JobStack)
// This can probably be cleaner. We're hindered somewhat, I think, by reading a channel from a channel.
for {
select {
case result := <-resChan:
wg.Add(1)
go func() {
if result.Error != nil {
fmt.Printf("%s: %s\n", result.Host, result.Error)
wg.Done()
} else {
readStream(result, &wg)
}
}()
default:
if massh.NumberOfStreamingHostsCompleted == numberOfExpectedCompletions {
// We want to wait for all goroutines to complete before we declare that the work is finished, as
// it's possible for us to execute this code before the gofunc above has completed if left unchecked.
wg.Wait()
// This should always be the last thing written. Waiting above ensures this.
fmt.Println("Everything returned.")
return
}
}
}
}
// Read Stdout stream
func readStream(res massh.Result, wg *sync.WaitGroup) error {
for {
select {
case d := <-res.StdOutStream:
fmt.Printf("%s: %s", res.Host, d)
case <-res.DoneChannel:
// Confirm that the host has exited.
fmt.Printf("%s: Finished\n", res.Host)
wg.Done()
}
}
}