feat: record cmd duration in telemetry (#1122)

* feat: record cmd duration in telemetry Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix lint Signed-off-by: Keming <kemingyang@tensorchord.ai> * functional options Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix lint Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Oct 31, 2022 · 07768b7 · 07768b7
1 parent 0bdca3c
commit 07768b7
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 17 deletions.
diff --git a/pkg/app/build.go b/pkg/app/build.go
@@ -17,6 +17,7 @@ package app
 import (
 	"fmt"
 	"path/filepath"
+	"time"
 
 	"github.com/cockroachdb/errors"
 	"github.com/sirupsen/logrus"
@@ -100,11 +101,14 @@ To build and push the image to a registry:
 }
 
 func build(clicontext *cli.Context) error {
-	telemetry.GetReporter().Telemetry("build", nil)
 	opt, err := ParseBuildOpt(clicontext)
 	if err != nil {
 		return err
 	}
+	defer func(start time.Time) {
+		telemetry.GetReporter().Telemetry(
+			"build", telemetry.AddField("duration", time.Since(start).Seconds()))
+	}(time.Now())
 
 	logger := logrus.WithFields(logrus.Fields{
 		"build-context": opt.BuildContextDir,

diff --git a/pkg/app/destroy.go b/pkg/app/destroy.go
@@ -72,9 +72,7 @@ func destroy(clicontext *cli.Context) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to get the current context")
 	}
-
-	r := string(context.Runner)
-	telemetry.GetReporter().Telemetry("destroy", &r)
+	telemetry.GetReporter().Telemetry("destroy", telemetry.AddField("runner", context.Runner))
 
 	opt := envd.Options{Context: context}
 	envdEngine, err := envd.New(clicontext.Context, opt)

diff --git a/pkg/app/env_ls.go b/pkg/app/env_ls.go
@@ -44,8 +44,7 @@ func getEnvironment(clicontext *cli.Context) error {
 		return errors.Wrap(err, "failed to get the current context")
 	}
 
-	runner := string(context.Runner)
-	telemetry.GetReporter().Telemetry("environment_list", &runner)
+	telemetry.GetReporter().Telemetry("environment_list", telemetry.AddField("runner", context.Runner))
 	opt := envd.Options{
 		Context: context,
 	}

diff --git a/pkg/app/image_prune.go b/pkg/app/image_prune.go
@@ -35,8 +35,7 @@ var CommandPruneImages = &cli.Command{
 }
 
 func pruneImages(clicontext *cli.Context) error {
-	r := "docker"
-	telemetry.GetReporter().Telemetry("image_prune", &r)
+	telemetry.GetReporter().Telemetry("image_prune", telemetry.AddField("runner", "docker"))
 
 	cli, err := docker.NewClient(clicontext.Context)
 	if err != nil {

diff --git a/pkg/app/telemetry/reporter.go b/pkg/app/telemetry/reporter.go
@@ -30,8 +30,10 @@ import (
 	"github.com/tensorchord/envd/pkg/version"
 )
 
+type TelemetryField func(*segmentio.Properties)
+
 type Reporter interface {
-	Telemetry(command string, runner *string)
+	Telemetry(command string, fields ...TelemetryField)
 }
 
 type defaultReporter struct {
@@ -153,7 +155,13 @@ func (r *defaultReporter) Identify() {
 	}
 }
 
-func (r *defaultReporter) Telemetry(command string, runner *string) {
+func AddField(name string, value interface{}) TelemetryField {
+	return func(p *segmentio.Properties) {
+		p.Set(name, value)
+	}
+}
+
+func (r *defaultReporter) Telemetry(command string, fields ...TelemetryField) {
 	if r.enabled {
 		logrus.WithFields(logrus.Fields{
 			"UID":     r.UID,
@@ -164,11 +172,13 @@ func (r *defaultReporter) Telemetry(command string, runner *string) {
 			Event:      command,
 			Properties: segmentio.NewProperties(),
 		}
-		if runner != nil {
-			t.Properties = t.Properties.Set("runner", runner)
+		for _, field := range fields {
+			field(&t.Properties)
 		}
 		if err := r.client.Enqueue(t); err != nil {
 			logrus.Warn(err)
 		}
+		// make sure the msg can be sent out
+		r.client.Close()
 	}
 }
diff --git a/pkg/app/top.go b/pkg/app/top.go
@@ -51,9 +51,7 @@ func top(clicontext *cli.Context) error {
 	opt := envd.Options{
 		Context: context,
 	}
-
-	r := string(context.Runner)
-	telemetry.GetReporter().Telemetry("top", &r)
+	telemetry.GetReporter().Telemetry("top", telemetry.AddField("runner", context.Runner))
 
 	envdEngine, err := envd.New(clicontext.Context, opt)
 	if err != nil {

diff --git a/pkg/app/up.go b/pkg/app/up.go
@@ -131,8 +131,7 @@ func up(clicontext *cli.Context) error {
 	if c.Runner == types.RunnerTypeEnvdServer {
 		buildOpt.OutputOpts = fmt.Sprintf("type=image,name=%s,push=true", buildOpt.Tag)
 	}
-	r := string(c.Runner)
-	telemetry.GetReporter().Telemetry("up", &r)
+	start := time.Now()
 
 	ctr := filepath.Base(buildOpt.BuildContextDir)
 	detach := clicontext.Bool("detach")
@@ -219,6 +218,10 @@ func up(clicontext *cli.Context) error {
 		logrus.Infof("failed to add entry %s to your SSH config file: %s", ctr, err)
 		return errors.Wrap(err, "failed to add entry to your SSH config file")
 	}
+	telemetry.GetReporter().Telemetry(
+		"up",
+		telemetry.AddField("runner", c.Runner),
+		telemetry.AddField("duration", time.Since(start).Seconds()))
 
 	if !detach {
 		if err := engine.Attach(ctr, hostname,