From db20f89ae581b88eeccb225cbcc3b1973b20b566 Mon Sep 17 00:00:00 2001
From: Michael Grosser <michael@grosser.it>
Date: Sun, 24 May 2020 12:40:35 -0700
Subject: [PATCH] add Time function to report time a passed function takes

---
 statsd/noop.go        |  6 ++++++
 statsd/noop_test.go   |  1 +
 statsd/statsd.go      | 11 +++++++++++
 statsd/statsd_test.go | 12 +++++++++++-
 4 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/statsd/noop.go b/statsd/noop.go
index 010783333..35b91563e 100644
--- a/statsd/noop.go
+++ b/statsd/noop.go
@@ -41,6 +41,12 @@ func (n *NoOpClient) Set(name string, value string, tags []string, rate float64)
 	return nil
 }
 
+// Time executes the passed function and returns nil
+func (n *NoOpClient) Time(name string, tags []string, rate float64, fn func()) error {
+	fn()
+	return nil
+}
+
 // Timing does nothing and returns nil
 func (n *NoOpClient) Timing(name string, value time.Duration, tags []string, rate float64) error {
 	return nil
diff --git a/statsd/noop_test.go b/statsd/noop_test.go
index afff3121d..f9e9059dc 100644
--- a/statsd/noop_test.go
+++ b/statsd/noop_test.go
@@ -19,6 +19,7 @@ func TestNoOpClient(t *testing.T) {
 	a.Nil(c.Decr("asd", tags, 56.0))
 	a.Nil(c.Incr("asd", tags, 56.0))
 	a.Nil(c.Set("asd", "asd", tags, 56.0))
+	a.Nil(c.Time("asd", tags, 56.0, func() {}))
 	a.Nil(c.Timing("asd", time.Second, tags, 56.0))
 	a.Nil(c.TimeInMilliseconds("asd", 1234.5, tags, 56.0))
 	a.Nil(c.Event(nil))
diff --git a/statsd/statsd.go b/statsd/statsd.go
index 6eadb283a..6e33af07b 100644
--- a/statsd/statsd.go
+++ b/statsd/statsd.go
@@ -166,6 +166,9 @@ type ClientInterface interface {
 	// Set counts the number of unique elements in a group.
 	Set(name string, value string, tags []string, rate float64) error
 
+	// Time measures and sends timing information
+	Time(name string, tags []string, rate float64, fn func()) error
+
 	// Timing sends timing information, it is an alias for TimeInMilliseconds
 	Timing(name string, value time.Duration, tags []string, rate float64) error
 
@@ -544,6 +547,14 @@ func (c *Client) Set(name string, value string, tags []string, rate float64) err
 	return c.send(metric{metricType: set, name: name, svalue: value, tags: tags, rate: rate})
 }
 
+// Time measures and sends timing information
+func (c *Client) Time(name string, tags []string, rate float64, fn func()) error {
+	start := time.Now()
+	fn()
+	duration := time.Since(start)
+	return c.Timing(name, duration, tags, rate)
+}
+
 // Timing sends timing information, it is an alias for TimeInMilliseconds
 func (c *Client) Timing(name string, value time.Duration, tags []string, rate float64) error {
 	return c.TimeInMilliseconds(name, value.Seconds()*1000, tags, rate)
diff --git a/statsd/statsd_test.go b/statsd/statsd_test.go
index 1a9eb2b92..a0bd03e53 100644
--- a/statsd/statsd_test.go
+++ b/statsd/statsd_test.go
@@ -46,6 +46,7 @@ func testTelemetry(t *testing.T, client *Client, expectedTelemetryTags []string)
 	client.Decr("Decr", nil, 1)
 	client.Incr("Incr", nil, 1)
 	client.Set("Set", "value", nil, 1)
+	client.Time("Time", nil, 1, func() {})
 	client.Timing("Timing", 21, nil, 1)
 	client.TimeInMilliseconds("TimeInMilliseconds", 21, nil, 1)
 	client.SimpleEvent("hello", "world")
@@ -54,7 +55,7 @@ func testTelemetry(t *testing.T, client *Client, expectedTelemetryTags []string)
 	metrics := client.flushTelemetry()
 
 	expectedMetricsName := map[string]int64{
-		"datadog.dogstatsd.client.metrics":                   9,
+		"datadog.dogstatsd.client.metrics":                   10,
 		"datadog.dogstatsd.client.events":                    1,
 		"datadog.dogstatsd.client.service_checks":            1,
 		"datadog.dogstatsd.client.metric_dropped_on_receive": 0,
@@ -179,3 +180,12 @@ func TestCloneWithExtraOptions(t *testing.T) {
 	assert.Equal(t, cloneClient.addrOption, addr)
 	assert.Len(t, cloneClient.options, 3)
 }
+
+func testTime(t *testing.T, client *Client, expectedTelemetryTags []string) {
+	executed := 0
+	fn := func() {
+		executed = executed + 1
+	}
+	client.Time("Time", nil, 1, fn)
+	assert.Equal(t, executed, 1)
+}