Merge pull request #16 from DiscoRiver/stream-bug

Fix Stream bug

Author: DiscoRiver

.github/workflows/go-test.yml

@@ -1,10 +1,6 @@
 name: Go
 
-on:
-  push:
-    branches: [ master ]
-  pull_request:
-    branches: [ master ]
+on: [push, pull_request]
 
 jobs:
   build-and-test:

_examples/example_jobstack_streaming/main.go

@@ -57,10 +57,7 @@ func main() {
 					fmt.Printf("%s: %s\n", result.Host, result.Error)
 					wg.Done()
 				} else {
-					err := readStream(result, &wg)
-					if err != nil {
-						panic(err)
-					}
+					readStream(result, &wg)
 				}
 			}()
 		default:

_examples/example_streaming/main.go

@@ -54,10 +54,7 @@ func main() {
 					fmt.Printf("%s: %s\n", result.Host, result.Error)
 					wg.Done()
 				} else {
-					err := readStream(result, &wg)
-					if err != nil {
-						fmt.Println(err)
-					}
+					readStream(result, &wg)
 				}
 			}()
 		default:

session.go

@@ -1,9 +1,11 @@
 package massh
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"io"
+	"sync"
 )
 
 var (
@@ -133,8 +135,12 @@ func sshCommandStream(host string, config *Config, resultChannel chan Result) {
 	// Reading from our pipes as they're populated, and redirecting bytes to our stdout and stderr channels in Result.
 	//
 	// We're doing this before we start the ssh task so we can start churning through output as soon as it starts.
-	go readToBytesChannel(StdOutPipe, r.StdOutStream, r)
-	go readToBytesChannel(StdErrPipe, r.StdErrStream, r)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		readToBytesChannel(StdOutPipe, r.StdOutStream, r, &wg)
+		readToBytesChannel(StdErrPipe, r.StdErrStream, r, &wg)
+	}()
 
 	resultChannel <- r
 
@@ -145,21 +151,29 @@ func sshCommandStream(host string, config *Config, resultChannel chan Result) {
 	}
 
 	// Wait for the command to exit only after we've initiated all the output channels
+	wg.Wait()
 	session.Wait()
 
 	NumberOfStreamingHostsCompleted++
 }
 
 // readToBytesChannel reads from io.Reader and directs the data to a byte slice channel for streaming.
-func readToBytesChannel(reader io.Reader, stream chan []byte, r Result) {
-	var data = make([]byte, 1024)
+func readToBytesChannel(reader io.Reader, stream chan []byte, r Result, wg *sync.WaitGroup) {
+	defer func(){ wg.Done() }()
+
+	rdr := bufio.NewReader(reader)
+
 	for {
-		n, err := reader.Read(data)
+		line, err := rdr.ReadBytes('\n')
 		if err != nil {
-			r.Error = fmt.Errorf("couldn't read content to stream channel: %s", err)
-			return
+			if err == io.EOF {
+				return
+			} else {
+				r.Error = fmt.Errorf("couldn't read content to stream channel: %s", err)
+				return
+			}
 		}
-		stream <- data[:n]
+		stream <- line
 	}
 }
 

session_test.go

@@ -1,6 +1,7 @@
 package massh
 
 import (
+	"fmt"
 	"golang.org/x/crypto/ssh"
 	"strings"
 	"sync"
@@ -71,11 +72,7 @@ func TestSshCommandStream(t *testing.T) {
 
 					wg.Done()
 				} else {
-					err := readStream(result, &wg, t)
-					if err != nil {
-						t.Log(err)
-						t.FailNow()
-					}
+					readStream(result, &wg, t)
 				}
 			}()
 		default:
@@ -90,14 +87,62 @@ func TestSshCommandStream(t *testing.T) {
 	}
 }
 
-func readStream(res Result, wg *sync.WaitGroup, t *testing.T) error {
+// Test for bugs in lots of lines.
+func TestSshCommandStreamBigData(t *testing.T) {
+	defer func() {testConfig.Job = testJob}()
+	NumberOfStreamingHostsCompleted = 0
+
+	if err := testConfig.SetPrivateKeyAuth("~/.ssh/id_rsa", ""); err != nil {
+		t.Log(err)
+		t.FailNow()
+	}
+
+	testConfig.Job = &Job{
+		Command: "cat /var/log/auth.log",
+	}
+
+	resChan := make(chan Result)
+
+	// This should be the last responsibility from the massh package. Handling the Result channel is up to the user.
+	err := testConfig.Stream(resChan)
+	if err != nil {
+		t.Log(err)
+		t.FailNow()
+	}
+
+	var wg sync.WaitGroup
+	// This can probably be cleaner. We're hindered somewhat, I think, by reading a channel from a channel.
 	for {
 		select {
-		case d := <-res.StdOutStream:
-			if !strings.Contains(string(d), "Hello, World") {
-				t.Logf("Expected output from stream test not received from host %s: %s", res.Host, d)
-				t.Fail()
+		case result := <-resChan:
+			wg.Add(1)
+			go func() {
+				if result.Error != nil {
+					t.Logf("Unexpected error in stream test for host %s: %s", result.Host, result.Error)
+					t.Fail()
+
+					wg.Done()
+				} else {
+					readStream(result, &wg, t)
+				}
+			}()
+		default:
+			if NumberOfStreamingHostsCompleted == len(testConfig.Hosts) {
+				// 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()
+
+				return
 			}
+		}
+	}
+}
+
+func readStream(res Result, wg *sync.WaitGroup, t *testing.T) {
+	for {
+		select {
+		case d := <-res.StdOutStream:
+			fmt.Print(string(d))
 		case <-res.DoneChannel:
 			wg.Done()
 		}
@@ -232,11 +277,7 @@ func TestSshCommandStreamWithJobStack(t *testing.T) {
 
 					wg.Done()
 				} else {
-					err := readStream(result, &wg, t)
-					if err != nil {
-						t.Log(err)
-						t.FailNow()
-					}
+					readStream(result, &wg, t)
 				}
 			}()
 		default: