Merge pull request #35 from InVisionApp/nostats

Add ability to turn off specific stats
InVisionApp · Aug 23, 2018 · c54f0c8 · c54f0c8
2 parents 1631744 + 7d1a2fb
commit c54f0c8
Show file tree

Hide file tree

Showing 2 changed files with 102 additions and 11 deletions.
diff --git a/rye.go b/rye.go
@@ -28,6 +28,11 @@ type MWHandler struct {
 type Config struct {
 	Statter  statsd.Statter
 	StatRate float32
+
+	// toggle types of stats sent
+	NoErrStats        bool
+	NoDurationStats   bool
+	NoStatusCodeStats bool
 }
 
 // JSONStatus is a simple container used for conveying status messages.
@@ -122,7 +127,7 @@ func (m *MWHandler) do(w http.ResponseWriter, r *http.Request, handler Handler)
 
 				// Now assume we have an error.
 				if m.Config.Statter != nil && resp.StatusCode >= 500 {
-					go m.Config.Statter.Inc("errors", 1, m.Config.StatRate)
+					go m.reportError()
 				}
 
 				// Write the error out
@@ -138,18 +143,10 @@ func (m *MWHandler) do(w http.ResponseWriter, r *http.Request, handler Handler)
 
 		if m.Config.Statter != nil {
 			// Record runtime metric
-			go m.Config.Statter.TimingDuration(
-				"handlers."+handlerName+".runtime",
-				time.Since(startTime), // delta
-				m.Config.StatRate,
-			)
+			go m.reportDuration(handlerName, startTime)
 
 			// Record status code metric (default 2xx)
-			go m.Config.Statter.Inc(
-				"handlers."+handlerName+"."+statusCode,
-				1,
-				m.Config.StatRate,
-			)
+			go m.reportStatusCode(handlerName, statusCode)
 		}
 	}()
 
@@ -162,6 +159,38 @@ func (m *MWHandler) do(w http.ResponseWriter, r *http.Request, handler Handler)
 	return false, r
 }
 
+func (m *MWHandler) reportError() {
+	if m.Config.NoErrStats {
+		return
+	}
+
+	m.Config.Statter.Inc("errors", 1, m.Config.StatRate)
+}
+
+func (m *MWHandler) reportDuration(handlerName string, startTime time.Time) {
+	if m.Config.NoDurationStats {
+		return
+	}
+
+	m.Config.Statter.TimingDuration(
+		"handlers."+handlerName+".runtime",
+		time.Since(startTime), // delta
+		m.Config.StatRate,
+	)
+}
+
+func (m *MWHandler) reportStatusCode(handlerName string, statusCode string) {
+	if m.Config.NoStatusCodeStats {
+		return
+	}
+
+	m.Config.Statter.Inc(
+		"handlers."+handlerName+"."+statusCode,
+		1,
+		m.Config.StatRate,
+	)
+}
+
 // WriteJSONStatus is a wrapper for WriteJSONResponse that returns a marshalled JSONStatus blob
 func WriteJSONStatus(rw http.ResponseWriter, status, message string, statusCode int) {
 	jsonData, _ := json.Marshal(&JSONStatus{

diff --git a/rye_test.go b/rye_test.go
@@ -267,6 +267,68 @@ var _ = Describe("Rye", func() {
 			})
 		})
 
+		Context("when error stats are turned off", func() {
+			It("should not call Inc or TimingDuration", func() {
+				ryeConfig := Config{
+					Statter:    fakeStatter,
+					NoErrStats: true,
+				}
+
+				handler := NewMWHandler(ryeConfig)
+				//use the failureHandler so an error would be reported
+				h := handler.Handle([]Handler{failureHandler})
+				h.ServeHTTP(response, request)
+
+				time.Sleep(time.Millisecond * 10)
+
+				Expect(fakeStatter.IncCallCount()).To(Equal(1))
+				metric, _, _ := fakeStatter.IncArgsForCall(0)
+				Expect(metric).ToNot(Equal("errors"))
+
+				Expect(fakeStatter.TimingDurationCallCount()).To(Equal(1))
+			})
+		})
+
+		Context("when statusCode stats are turned off", func() {
+			It("should not call Inc or TimingDuration", func() {
+				ryeConfig := Config{
+					Statter:           fakeStatter,
+					NoStatusCodeStats: true,
+				}
+
+				handler := NewMWHandler(ryeConfig)
+				//use the failureHandler so an error would be reported
+				h := handler.Handle([]Handler{failureHandler})
+				h.ServeHTTP(response, request)
+
+				time.Sleep(time.Millisecond * 10)
+
+				Expect(fakeStatter.IncCallCount()).To(Equal(1))
+				metric, _, _ := fakeStatter.IncArgsForCall(0)
+				Expect(metric).ToNot(ContainSubstring("handlers."))
+
+				Expect(fakeStatter.TimingDurationCallCount()).To(Equal(1))
+			})
+		})
+
+		Context("when timming stats are turned off", func() {
+			It("should not call Inc or TimingDuration", func() {
+				ryeConfig := Config{
+					Statter:         fakeStatter,
+					NoDurationStats: true,
+				}
+
+				handler := NewMWHandler(ryeConfig)
+				//use the failureHandler so an error would be reported
+				h := handler.Handle([]Handler{failureHandler})
+				h.ServeHTTP(response, request)
+
+				time.Sleep(time.Millisecond * 10)
+
+				Expect(fakeStatter.TimingDurationCallCount()).To(Equal(0))
+				Expect(fakeStatter.IncCallCount()).To(Equal(2))
+			})
+		})
 	})
 
 	Describe("getFuncName", func() {