interactive shell improvements

- support case when script exec time > sampling rate
- extended example config file

data/item.go

@@ -14,16 +14,22 @@ import (
 	"time"
 )
 
-const interactiveShellStartupTimeout = 100 * time.Millisecond
+const (
+	interactiveShellStartupTimeout  = 100 * time.Millisecond
+	interactiveShellMinAwaitTimeout = 100 * time.Millisecond
+	interactiveShellMaxAwaitTimeout = 1 * time.Second
+	interactiveShellErrorThreshold  = 10
+)
 
 type Item struct {
-	Label            string
+	label            string
-	SampleScript     string
+	sampleScript     string
-	InitScript       *string
+	initScript       *string
-	TransformScript  *string
+	transformScript  *string
-	Color            *ui.Color
+	color            *ui.Color
-	RateMs           int
+	rateMs           int
-	InteractiveShell *InteractiveShell
+	errorsCount      int
+	interactiveShell *InteractiveShell
 }
 
 type InteractiveShell struct {
@@ -38,12 +44,12 @@ func NewItems(cfgs []config.Item, rateMs int) []*Item {
 
 	for _, i := range cfgs {
 		item := &Item{
-			Label:           *i.Label,
+			label:           *i.Label,
-			SampleScript:    *i.SampleScript,
+			sampleScript:    *i.SampleScript,
-			InitScript:      i.InitScript,
+			initScript:      i.InitScript,
-			TransformScript: i.TransformScript,
+			transformScript: i.TransformScript,
-			Color:           i.Color,
+			color:           i.Color,
-			RateMs:          rateMs,
+			rateMs:          rateMs,
 		}
 		items = append(items, item)
 	}
@@ -53,17 +59,17 @@ func NewItems(cfgs []config.Item, rateMs int) []*Item {
 
 func (i *Item) nextValue(variables []string) (string, error) {
 
-	if i.InitScript != nil && i.InteractiveShell == nil {
+	if i.initScript != nil && i.interactiveShell == nil {
 		err := i.initInteractiveShell(variables)
 		if err != nil {
 			return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err))
 		}
 	}
 
-	if i.InitScript != nil {
+	if i.initScript != nil {
 		return i.executeInteractiveShellCmd(variables)
 	} else {
-		return i.executeCmd(variables, i.SampleScript)
+		return i.executeCmd(variables, i.sampleScript)
 	}
 }
 
@@ -83,7 +89,7 @@ func (i *Item) executeCmd(variables []string, script string) (string, error) {
 
 func (i *Item) initInteractiveShell(variables []string) error {
 
-	cmd := exec.Command("sh", "-c", *i.InitScript)
+	cmd := exec.Command("sh", "-c", *i.initScript)
 	enrichEnvVariables(cmd, variables)
 
 	file, err := pty.Start(cmd)
@@ -100,7 +106,7 @@ func (i *Item) initInteractiveShell(variables []string) error {
 		}
 	}()
 
-	i.InteractiveShell = &InteractiveShell{
+	i.interactiveShell = &InteractiveShell{
 		Channel: channel,
 		File:    file,
 		Cmd:     cmd,
@@ -118,43 +124,79 @@ func (i *Item) initInteractiveShell(variables []string) error {
 
 func (i *Item) executeInteractiveShellCmd(variables []string) (string, error) {
 
-	_, err := io.WriteString(i.InteractiveShell.File, fmt.Sprintf(" %s\n", i.SampleScript))
+	_, err := io.WriteString(i.interactiveShell.File, fmt.Sprintf(" %s\n", i.sampleScript))
 	if err != nil {
+		i.errorsCount++
+		if i.errorsCount > interactiveShellErrorThreshold {
+			i.interactiveShell = nil // restart session
+			i.errorsCount = 0
+		}
 		return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err))
 	}
 
-	timeout := make(chan bool, 1)
+	softTimeout := make(chan bool, 1)
+	hardTimeout := make(chan bool, 1)
 
 	go func() {
-		time.Sleep(time.Duration(i.RateMs))
+		time.Sleep(i.getAwaitTimeout() / 4)
-		timeout <- true
+		softTimeout <- true
+		time.Sleep(i.getAwaitTimeout() * 100)
+		hardTimeout <- true
 	}()
 
-	var outputText strings.Builder
+	var builder strings.Builder
+	softTimeoutElapsed := false
 
+await:
 	for {
 		select {
-		case output := <-i.InteractiveShell.Channel:
+		case output := <-i.interactiveShell.Channel:
-			if !strings.Contains(output, i.SampleScript) && len(output) > 0 {
+			o := cleanupOutput(output)
-				outputText.WriteString(output)
+			if len(o) > 0 && !strings.Contains(o, i.sampleScript) {
-				outputText.WriteString("\n")
+				builder.WriteString(o)
+				builder.WriteString("\n")
+				if softTimeoutElapsed {
+					break await
 				}
-		case <-timeout:
+			}
-			sample := cleanupOutput(outputText.String())
+		case <-softTimeout:
+			if builder.Len() > 0 {
+				break await
+			} else {
+				softTimeoutElapsed = true
+			}
+		case <-hardTimeout:
+			break await
+		}
+	}
+
+	sample := strings.TrimSpace(builder.String())
+
 	return i.transformInteractiveShellCmd(sample)
-		}
-	}
 }
 
 func (i *Item) transformInteractiveShellCmd(sample string) (string, error) {
 
-	if i.TransformScript != nil && len(sample) > 0 {
+	if i.transformScript != nil && len(sample) > 0 {
-		return i.executeCmd([]string{"sample=" + sample}, *i.TransformScript)
+		return i.executeCmd([]string{"sample=" + sample}, *i.transformScript)
 	}
 
 	return sample, nil
 }
 
+func (i *Item) getAwaitTimeout() time.Duration {
+
+	timeout := time.Duration(i.rateMs) * time.Millisecond
+
+	if timeout > interactiveShellMaxAwaitTimeout {
+		return interactiveShellMaxAwaitTimeout
+	} else if timeout < interactiveShellMinAwaitTimeout {
+		return interactiveShellMinAwaitTimeout
+	}
+
+	return timeout
+}
+
 func enrichEnvVariables(cmd *exec.Cmd, variables []string) {
 	cmd.Env = os.Environ()
 	for _, variable := range variables {

data/sampler.go

@@ -53,14 +53,14 @@ func (s *Sampler) sample(item *Item, options config.Options) {
 	val, err := item.nextValue(s.variables)
 
 	if len(val) > 0 {
-		sample := &Sample{Label: item.Label, Value: val, Color: item.Color}
+		sample := &Sample{Label: item.label, Value: val, Color: item.color}
 		s.consumer.SampleChannel <- sample
 		s.triggersChannel <- sample
 	} else if err != nil {
 		s.consumer.AlertChannel <- &Alert{
 			Title: "SAMPLING FAILURE",
 			Text:  getErrorMessage(err),
-			Color: item.Color,
+			Color: item.color,
 		}
 	}
 }

example-interactive-shell.yml

@@ -7,23 +7,24 @@ variables:
   sshconnection: ssh -i ~/sqshq.pem ec2-user@3.215.108.82
 textboxes:
 - title: Neo4j
-  position: [[0, 0], [13, 40]]
+  position: [[0, 0], [10, 40]]
   init: $neo4jconnection
-  sample: match (n) return count(n);
+  sample: RETURN rand();
   transform: echo "$sample" | tail -n 1
 - title: Postgres
-  position: [[13, 0], [14, 40]]
+  position: [[10, 0], [9, 40]]
   init: $postgresconnection
   sample: select random();
 - title: MySQL
-  position: [[27, 0], [14, 40]]
+  position: [[19, 0], [10, 40]]
   init: $mysqlconnection
   sample: select rand();
 - title: MongoDB
-  position: [[41, 0], [13, 40]]
+  position: [[29, 0], [10, 40]]
+  rate-ms: 500
   init: $mongoconnection
-  sample: db.getCollection('posts').find({status:'ACTIVE'}).itcount()
+  sample: sleep(3000);Date.now();
 - title: SSH
-  position: [[54, 0], [13, 40]]
+  position: [[39, 0], [41, 40]]
   init: $sshconnection
-  sample: ps -A -o %cpu | awk '{s+=$1} END {print s}'
+  sample: top