From d8a137bf943eb03ae03c0f7a8a40d4f514e2ef79 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Fri, 28 Jun 2019 13:58:37 +0300
Subject: [PATCH] Add `<summary>_sum` and `<summary>_count` values for Summary

This improves output compatibility with Prometheus Summary type
defined at https://prometheus.io/docs/concepts/metric_types/#summary
---
 summary.go      | 18 ++++++++++++++++--
 summary_test.go | 11 +++++++----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/summary.go b/summary.go
index 41454e4..3bb002d 100644
--- a/summary.go
+++ b/summary.go
@@ -24,6 +24,9 @@ type Summary struct {
 	quantiles      []float64
 	quantileValues []float64
 
+	sum   float64
+	count uint64
+
 	window time.Duration
 }
 
@@ -83,6 +86,8 @@ func (sm *Summary) Update(v float64) {
 	sm.mu.Lock()
 	sm.curr.Update(v)
 	sm.next.Update(v)
+	sm.sum += v
+	sm.count++
 	sm.mu.Unlock()
 }
 
@@ -93,9 +98,18 @@ func (sm *Summary) UpdateDuration(startTime time.Time) {
 }
 
 func (sm *Summary) marshalTo(prefix string, w io.Writer) {
-	// Do nothing. Quantile values should be already updated by the caller
-	// via sm.updateQuantiles() call.
+	// Marshal only *_sum and *_count values.
+	// Quantile values should be already updated by the caller via sm.updateQuantiles() call.
 	// sm.quantileValues will be marshaled later via quantileValue.marshalTo.
+	sm.mu.Lock()
+	sum := sm.sum
+	count := sm.count
+	sm.mu.Unlock()
+
+	if count > 0 {
+		fmt.Fprintf(w, "%s_sum %g\n", prefix, sum)
+		fmt.Fprintf(w, "%s_count %d\n", prefix, count)
+	}
 }
 
 func (sm *Summary) updateQuantiles() {
diff --git a/summary_test.go b/summary_test.go
index 007a417..d8f45ca 100644
--- a/summary_test.go
+++ b/summary_test.go
@@ -28,8 +28,8 @@ func TestSummarySerial(t *testing.T) {
 		s.UpdateDuration(t.Add(-time.Millisecond * time.Duration(i)))
 	}
 
-	// Make sure the summary doesn't print anything on marshalTo call
-	testMarshalTo(t, s, "prefix", "")
+	// Make sure the summary prints <prefix>_sum and <prefix>_count on marshalTo call
+	testMarshalTo(t, s, "prefix", fmt.Sprintf("prefix_sum %g\nprefix_count %d\n", s.sum, s.count))
 
 	// Verify s.quantileValues
 	s.updateQuantiles()
@@ -54,12 +54,12 @@ func TestSummaryConcurrent(t *testing.T) {
 		for i := 0; i < 10; i++ {
 			s.Update(float64(i))
 		}
-		testMarshalTo(t, s, "prefix", "")
 		return nil
 	})
 	if err != nil {
 		t.Fatal(err)
 	}
+	testMarshalTo(t, s, "prefix", "prefix_sum 225\nprefix_count 50\n")
 }
 
 func TestSummaryWithTags(t *testing.T) {
@@ -92,10 +92,13 @@ func TestSummarySmallWindow(t *testing.T) {
 		s.Update(123)
 	}
 	// Wait for window update and verify that the summary has been cleared.
-	time.Sleep(5 * window)
+	time.Sleep(2 * window)
 	var bb bytes.Buffer
 	WritePrometheus(&bb, false)
 	result := bb.String()
+	// <name>_sum and <name>_count are present in the output.
+	// Only <name>{quantile} shouldn't be present.
+	name += "{"
 	if strings.Contains(result, name) {
 		t.Fatalf("summary %s cannot be present in the WritePrometheus output; got\n%s", name, result)
 	}