mirror of
https://github.com/appleboy/drone-ssh.git
synced 2025-05-09 18:23:21 +08:00
feat: refactor code for parallel execution on multiple hosts
- Add `trimValues` function for cleaning up slice values - Remove unused `wg.Done()` call - Modify `Exec` function to launch goroutines for each host in `Config.Host` - Add test for `ScriptStop` with multiple hosts and sync mode Signed-off-by: Bo-Yi.Wu <appleboy.tw@gmail.com>
This commit is contained in:
parent
f733bb5ebf
commit
e650f4c9c2
84
plugin.go
84
plugin.go
@ -65,6 +65,7 @@ func (p Plugin) hostPort(host string) (string, string) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (p Plugin) exec(host string, wg *sync.WaitGroup, errChannel chan error) {
|
func (p Plugin) exec(host string, wg *sync.WaitGroup, errChannel chan error) {
|
||||||
|
defer wg.Done()
|
||||||
host, port := p.hostPort(host)
|
host, port := p.hostPort(host)
|
||||||
// Create MakeConfig instance with remote username, server address and path to private key.
|
// Create MakeConfig instance with remote username, server address and path to private key.
|
||||||
ssh := &easyssh.MakeConfig{
|
ssh := &easyssh.MakeConfig{
|
||||||
@ -117,38 +118,36 @@ func (p Plugin) exec(host string, wg *sync.WaitGroup, errChannel chan error) {
|
|||||||
stdoutChan, stderrChan, doneChan, errChan, err := ssh.Stream(strings.Join(p.Config.Script, "\n"), p.Config.CommandTimeout)
|
stdoutChan, stderrChan, doneChan, errChan, err := ssh.Stream(strings.Join(p.Config.Script, "\n"), p.Config.CommandTimeout)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
errChannel <- err
|
errChannel <- err
|
||||||
} else {
|
return
|
||||||
// read from the output channel until the done signal is passed
|
}
|
||||||
isTimeout := true
|
// read from the output channel until the done signal is passed
|
||||||
loop:
|
isTimeout := true
|
||||||
for {
|
loop:
|
||||||
select {
|
for {
|
||||||
case isTimeout = <-doneChan:
|
select {
|
||||||
break loop
|
case isTimeout = <-doneChan:
|
||||||
case outline := <-stdoutChan:
|
break loop
|
||||||
if outline != "" {
|
case outline := <-stdoutChan:
|
||||||
p.log(host, "out:", outline)
|
if outline != "" {
|
||||||
}
|
p.log(host, "out:", outline)
|
||||||
case errline := <-stderrChan:
|
|
||||||
if errline != "" {
|
|
||||||
p.log(host, "err:", errline)
|
|
||||||
}
|
|
||||||
case err = <-errChan:
|
|
||||||
}
|
}
|
||||||
}
|
case errline := <-stderrChan:
|
||||||
|
if errline != "" {
|
||||||
// get exit code or command error.
|
p.log(host, "err:", errline)
|
||||||
if err != nil {
|
}
|
||||||
errChannel <- err
|
case err = <-errChan:
|
||||||
}
|
|
||||||
|
|
||||||
// command time out
|
|
||||||
if !isTimeout {
|
|
||||||
errChannel <- errCommandTimeOut
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
wg.Done()
|
// get exit code or command error.
|
||||||
|
if err != nil {
|
||||||
|
errChannel <- err
|
||||||
|
}
|
||||||
|
|
||||||
|
// command time out
|
||||||
|
if !isTimeout {
|
||||||
|
errChannel <- errCommandTimeOut
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p Plugin) log(host string, message ...interface{}) {
|
func (p Plugin) log(host string, message ...interface{}) {
|
||||||
@ -164,6 +163,8 @@ func (p Plugin) log(host string, message ...interface{}) {
|
|||||||
|
|
||||||
// Exec executes the plugin.
|
// Exec executes the plugin.
|
||||||
func (p Plugin) Exec() error {
|
func (p Plugin) Exec() error {
|
||||||
|
p.Config.Host = trimValues(p.Config.Host)
|
||||||
|
|
||||||
if len(p.Config.Host) == 0 {
|
if len(p.Config.Host) == 0 {
|
||||||
return errMissingHost
|
return errMissingHost
|
||||||
}
|
}
|
||||||
@ -176,10 +177,14 @@ func (p Plugin) Exec() error {
|
|||||||
wg.Add(len(p.Config.Host))
|
wg.Add(len(p.Config.Host))
|
||||||
errChannel := make(chan error)
|
errChannel := make(chan error)
|
||||||
finished := make(chan struct{})
|
finished := make(chan struct{})
|
||||||
for _, host := range p.Config.Host {
|
if p.Config.Sync {
|
||||||
if p.Config.Sync {
|
go func() {
|
||||||
p.exec(host, &wg, errChannel)
|
for _, host := range p.Config.Host {
|
||||||
} else {
|
p.exec(host, &wg, errChannel)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
} else {
|
||||||
|
for _, host := range p.Config.Host {
|
||||||
go p.exec(host, &wg, errChannel)
|
go p.exec(host, &wg, errChannel)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -230,3 +235,18 @@ func (p Plugin) scriptCommands() []string {
|
|||||||
|
|
||||||
return commands
|
return commands
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func trimValues(keys []string) []string {
|
||||||
|
var newKeys []string
|
||||||
|
|
||||||
|
for _, value := range keys {
|
||||||
|
value = strings.TrimSpace(value)
|
||||||
|
if len(value) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
newKeys = append(newKeys, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
return newKeys
|
||||||
|
}
|
||||||
|
@ -440,6 +440,41 @@ func TestFingerprint(t *testing.T) {
|
|||||||
assert.Equal(t, unindent(expected), unindent(buffer.String()))
|
assert.Equal(t, unindent(expected), unindent(buffer.String()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestScriptStopWithMultipleHostAndSyncMode(t *testing.T) {
|
||||||
|
var (
|
||||||
|
buffer bytes.Buffer
|
||||||
|
expected = `
|
||||||
|
======CMD======
|
||||||
|
mkdir a/b/c
|
||||||
|
mkdir d/e/f
|
||||||
|
======END======
|
||||||
|
err: mkdir: can't create directory 'a/b/c': No such file or directory
|
||||||
|
`
|
||||||
|
)
|
||||||
|
|
||||||
|
plugin := Plugin{
|
||||||
|
Config: Config{
|
||||||
|
Host: []string{"", "localhost"},
|
||||||
|
Username: "drone-scp",
|
||||||
|
Port: 22,
|
||||||
|
KeyPath: "./tests/.ssh/id_rsa",
|
||||||
|
Script: []string{
|
||||||
|
"mkdir a/b/c",
|
||||||
|
"mkdir d/e/f",
|
||||||
|
},
|
||||||
|
CommandTimeout: 10 * time.Second,
|
||||||
|
ScriptStop: true,
|
||||||
|
Sync: true,
|
||||||
|
},
|
||||||
|
Writer: &buffer,
|
||||||
|
}
|
||||||
|
|
||||||
|
err := plugin.Exec()
|
||||||
|
assert.NotNil(t, err)
|
||||||
|
|
||||||
|
assert.Equal(t, unindent(expected), unindent(buffer.String()))
|
||||||
|
}
|
||||||
|
|
||||||
func TestScriptStop(t *testing.T) {
|
func TestScriptStop(t *testing.T) {
|
||||||
var (
|
var (
|
||||||
buffer bytes.Buffer
|
buffer bytes.Buffer
|
||||||
|
Loading…
Reference in New Issue
Block a user