interactive shell improvements

- support case when script exec time > sampling rate
- extended example config file
This commit is contained in:
sqshq 2019-05-30 23:01:43 -04:00
parent f5fdf635f0
commit 4495c8f8e8
3 changed files with 86 additions and 43 deletions

View File

@ -14,16 +14,22 @@ import (
"time" "time"
) )
const interactiveShellStartupTimeout = 100 * time.Millisecond const (
interactiveShellStartupTimeout = 100 * time.Millisecond
interactiveShellMinAwaitTimeout = 100 * time.Millisecond
interactiveShellMaxAwaitTimeout = 1 * time.Second
interactiveShellErrorThreshold = 10
)
type Item struct { 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 { type InteractiveShell struct {
@ -38,12 +44,12 @@ func NewItems(cfgs []config.Item, rateMs int) []*Item {
for _, i := range cfgs { for _, i := range cfgs {
item := &Item{ 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) items = append(items, item)
} }
@ -53,17 +59,17 @@ func NewItems(cfgs []config.Item, rateMs int) []*Item {
func (i *Item) nextValue(variables []string) (string, error) { 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) err := i.initInteractiveShell(variables)
if err != nil { if err != nil {
return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err)) return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err))
} }
} }
if i.InitScript != nil { if i.initScript != nil {
return i.executeInteractiveShellCmd(variables) return i.executeInteractiveShellCmd(variables)
} else { } 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 { func (i *Item) initInteractiveShell(variables []string) error {
cmd := exec.Command("sh", "-c", *i.InitScript) cmd := exec.Command("sh", "-c", *i.initScript)
enrichEnvVariables(cmd, variables) enrichEnvVariables(cmd, variables)
file, err := pty.Start(cmd) file, err := pty.Start(cmd)
@ -100,7 +106,7 @@ func (i *Item) initInteractiveShell(variables []string) error {
} }
}() }()
i.InteractiveShell = &InteractiveShell{ i.interactiveShell = &InteractiveShell{
Channel: channel, Channel: channel,
File: file, File: file,
Cmd: cmd, Cmd: cmd,
@ -118,43 +124,79 @@ func (i *Item) initInteractiveShell(variables []string) error {
func (i *Item) executeInteractiveShellCmd(variables []string) (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 { 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)) 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() { 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 { for {
select { 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) return i.transformInteractiveShellCmd(sample)
}
}
} }
func (i *Item) transformInteractiveShellCmd(sample string) (string, error) { 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 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) { func enrichEnvVariables(cmd *exec.Cmd, variables []string) {
cmd.Env = os.Environ() cmd.Env = os.Environ()
for _, variable := range variables { for _, variable := range variables {

View File

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

View File

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