From cea133f27ca4c02d063fbbcee9a7530b6b231d89 Mon Sep 17 00:00:00 2001
From: Sasha Klizhentas <klizhentas@gmail.com>
Date: Sun, 14 May 2017 12:47:04 -0700
Subject: [PATCH] fix bench goroutine leak and add more stats

---
 assets/monitoring/telegraf.conf |  6 ++++--
 lib/client/api.go               |  1 +
 lib/client/bench.go             |  2 ++
 lib/client/session.go           |  7 +++++++
 tool/tsh/common/tsh.go          | 16 ++++++++++++++++
 5 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/assets/monitoring/telegraf.conf b/assets/monitoring/telegraf.conf
index 33d1c8e463c..3c96a687312 100644
--- a/assets/monitoring/telegraf.conf
+++ b/assets/monitoring/telegraf.conf
@@ -66,8 +66,10 @@
   ## If not provided, will default to 5s. 0s means no timeout (not recommended).
   timeout = "5s"
 
-
-# Get all metrics from Kube-apiserver
+[[inputs.procstat]]
+  exe = "teleport"
+  prefix = "teleport"
+  
 [[inputs.prometheus]]
   # An array of urls to scrape metrics from.
   urls = ["http://127.0.0.1:3434/metrics"]
diff --git a/lib/client/api.go b/lib/client/api.go
index 926e4864505..9bdabade1b2 100644
--- a/lib/client/api.go
+++ b/lib/client/api.go
@@ -829,6 +829,7 @@ func (tc *TeleportClient) runCommand(
 				log.Error(err)
 				return
 			}
+			defer nodeSession.Close()
 			if err = nodeSession.runCommand(command, tc.OnShellCreated, tc.Config.Interactive); err != nil {
 				originErr := trace.Unwrap(err)
 				exitErr, ok := originErr.(*ssh.ExitError)
diff --git a/lib/client/bench.go b/lib/client/bench.go
index 49d75a1cc0f..feb3a673f5e 100644
--- a/lib/client/bench.go
+++ b/lib/client/bench.go
@@ -17,6 +17,7 @@ limitations under the License.
 package client
 
 import (
+	"bytes"
 	"context"
 	"io/ioutil"
 	"time"
@@ -55,6 +56,7 @@ type BenchmarkResult struct {
 func (tc *TeleportClient) Benchmark(ctx context.Context, bench Benchmark) (*BenchmarkResult, error) {
 	tc.Stdout = ioutil.Discard
 	tc.Stderr = ioutil.Discard
+	tc.Stdin = &bytes.Buffer{}
 
 	ctx, cancel := context.WithTimeout(ctx, bench.Duration)
 	defer cancel()
diff --git a/lib/client/session.go b/lib/client/session.go
index 33fc45132d2..d7556009825 100644
--- a/lib/client/session.go
+++ b/lib/client/session.go
@@ -471,3 +471,10 @@ func (ns *NodeSession) pipeInOut(shell io.ReadWriteCloser) {
 		}
 	}()
 }
+
+func (ns *NodeSession) Close() error {
+	if ns.closer != nil {
+		ns.closer.Close()
+	}
+	return nil
+}
diff --git a/tool/tsh/common/tsh.go b/tool/tsh/common/tsh.go
index edeb809fa35..a20c8322294 100644
--- a/tool/tsh/common/tsh.go
+++ b/tool/tsh/common/tsh.go
@@ -36,6 +36,7 @@ import (
 
 	"github.com/Sirupsen/logrus"
 	"github.com/buger/goterm"
+	gops "github.com/google/gops/agent"
 )
 
 // CLIConf stores command line arguments and flags:
@@ -92,6 +93,11 @@ type CLIConf struct {
 	BenchRate int
 	// Context is a context to control execution
 	Context context.Context
+	// Gops starts gops agent on a specified address
+	// if not specified, gops won't start
+	Gops bool
+	// GopsAddr specifies to gops addr to listen on
+	GopsAddr string
 }
 
 // Run executes TSH client. same as main() but easier to test
@@ -113,6 +119,8 @@ func Run(args []string, underTest bool) {
 	app.Flag("ttl", "Minutes to live for a SSH session").Int32Var(&cf.MinsToLive)
 	app.Flag("insecure", "Do not verify server's certificate and host name. Use only in test environments").Default("false").BoolVar(&cf.InsecureSkipVerify)
 	app.Flag("namespace", "Namespace of the cluster").Default(defaults.Namespace).StringVar(&cf.Namespace)
+	app.Flag("gops", "Start gops endpoint on a given address").Hidden().BoolVar(&cf.Gops)
+	app.Flag("gops-addr", "Specify gops addr to listen on").Hidden().StringVar(&cf.GopsAddr)
 	debugMode := app.Flag("debug", "Verbose logging to stdout").Short('d').Bool()
 	app.HelpFlag.Short('h')
 	ver := app.Command("version", "Print the version")
@@ -187,6 +195,14 @@ func Run(args []string, underTest bool) {
 	}()
 	cf.Context = ctx
 
+	if cf.Gops {
+		logrus.Debugf("starting gops agent")
+		err = gops.Listen(&gops.Options{Addr: cf.GopsAddr})
+		if err != nil {
+			logrus.Warningf("failed to start gops agent %v", err)
+		}
+	}
+
 	switch command {
 	case ver.FullCommand():
 		onVersion()