Merge pull request #26 from qa-dev/optional-statsd

Optional statsd.

Author: Andrei Skomorkhov

config/config.go

@@ -9,10 +9,10 @@ import (
 
 // Config - settings of application.
 type Config struct {
-	Logger Logger `json:"logger"`
-	DB     DB     `json:"db"`
-	Grid   Grid   `json:"grid"`
-	Statsd Statsd `json:"statsd"`
+	Logger Logger  `json:"logger"`
+	DB     DB      `json:"db"`
+	Grid   Grid    `json:"grid"`
+	Statsd *Statsd `json:"statsd,omitempty"`
 }
 
 // Grid general settings

logger/logger.go

@@ -8,10 +8,10 @@ import (
 )
 
 // Init - initialisation of logger.
-func Init(logger config.Logger) error {
-	level, err := logrus.ParseLevel(logger.Level)
+func Init(cfgLogger config.Logger) error {
+	level, err := logrus.ParseLevel(cfgLogger.Level)
 	if err != nil {
-		return fmt.Errorf("Parse log level, %v", err)
+		return fmt.Errorf("parse log level, %v", err)
 	}
 	logrus.Infof("Set log level to: %v", level)
 	logrus.SetLevel(level)

main.go

@@ -32,16 +32,7 @@ func main() {
 		log.Fatalf("Problem in init logger, %s", err)
 	}
 
-	statsdClient, err := metrics.NewStatsd(
-		cfg.Statsd.Host,
-		cfg.Statsd.Port,
-		cfg.Statsd.Protocol,
-		cfg.Statsd.Prefix,
-		cfg.Statsd.Enable)
-
-	if nil != err {
-		log.Errorf("Statsd create socked error: %s", err)
-	}
+	middlewareWrap := middleware.NewWrap(log.StandardLogger())
 
 	busyNodeDuration, err := time.ParseDuration(cfg.Grid.BusyNodeDuration)
 	if err != nil {
@@ -84,8 +75,7 @@ func main() {
 	poolInstance.SetBusyNodeDuration(busyNodeDuration)
 	poolInstance.SetReservedNodeDuration(reservedNodeDuration)
 
-	poolMetricsSender := poolMetrics.NewSender(statsdClient, poolInstance, time.Second*1) // todo: move to config
-	go poolMetricsSender.SendAll()
+
 
 	go func() {
 		for {
@@ -103,13 +93,29 @@ func main() {
 		}
 	}()
 
-	m := middleware.NewLogMiddleware(statsdClient)
-	http.Handle("/wd/hub/session", m.Log(&handlers.CreateSession{Pool: poolInstance, ClientFactory: clientFactory})) //selenium
-	http.Handle("/session", m.Log(&handlers.CreateSession{Pool: poolInstance, ClientFactory: clientFactory}))        //wda
-	http.Handle("/grid/register", m.Log(&handlers.RegisterNode{Pool: poolInstance}))
+
+
+	if cfg.Statsd != nil {
+		statsdClient, err := metrics.NewStatsd(
+			cfg.Statsd.Host,
+			cfg.Statsd.Port,
+			cfg.Statsd.Protocol,
+			cfg.Statsd.Prefix,
+			cfg.Statsd.Enable)
+		poolMetricsSender := poolMetrics.NewSender(statsdClient, poolInstance, time.Second*1) // todo: move to config
+		go poolMetricsSender.SendAll()
+		if err != nil {
+			log.Errorf("Statsd create socked error: %s", err)
+		}
+		middlewareWrap.Add(middleware.NewStatsd(log.StandardLogger(), statsdClient, true).RegisterMetrics)
+	}
+	
+	http.Handle("/wd/hub/session", middlewareWrap.Do(&handlers.CreateSession{Pool: poolInstance, ClientFactory: clientFactory})) //selenium
+	http.Handle("/session", middlewareWrap.Do(&handlers.CreateSession{Pool: poolInstance, ClientFactory: clientFactory}))        //wda
+	http.Handle("/grid/register", middlewareWrap.Do(&handlers.RegisterNode{Pool: poolInstance}))
 	http.Handle("/grid/api/proxy", &handlers.APIProxy{Pool: poolInstance})
 	http.HandleFunc("/_info", heartbeat)
-	http.Handle("/", m.Log(&handlers.UseSession{Pool: poolInstance, Cache: cache}))
+	http.Handle("/", middlewareWrap.Do(&handlers.UseSession{Pool: poolInstance, Cache: cache}))
 
 	server := &http.Server{Addr: fmt.Sprintf(":%v", cfg.Grid.Port)}
 	serverError := make(chan error)

middleware/log.go renamed to middleware/statsd.go

@@ -2,41 +2,38 @@ package middleware
 
 import (
 	"net/http"
-	"runtime/debug"
 	"strings"
 
 	"fmt"
-	log "github.com/Sirupsen/logrus"
+	"github.com/Sirupsen/logrus"
 	"gopkg.in/alexcesaro/statsd.v2"
 	"regexp"
 )
 
-// LogMiddleware - wraps adds logging to the query handlers.
-type LogMiddleware struct {
-	statsd *statsd.Client
+// Statsd is statsd metrics middleware.
+type Statsd struct {
+	logger  *logrus.Logger
+	client  *statsd.Client
+	withLog bool
 }
 
-// NewLogMiddleware - constructor of LogMiddleware.
-func NewLogMiddleware(statsd *statsd.Client) *LogMiddleware {
-	return &LogMiddleware{
-		statsd: statsd,
+// NewStatsd construct Statsd.
+func NewStatsd(logger *logrus.Logger, client *statsd.Client, withLog bool) *Statsd {
+	return &Statsd{
+		logger:  logger,
+		client:  client,
+		withLog: withLog,
 	}
 }
 
-// Log - wraps http.Handler for runtime logging.
-func (m *LogMiddleware) Log(handler http.Handler) http.Handler {
+// RegisterMetrics send metrics to statsd
+func (s *Statsd) RegisterMetrics(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
 
-		defer func() {
-			if err := recover(); err != nil {
-				log.Fatalf("Panic: %+v\n%s", err, debug.Stack())
-			}
-		}()
-
 		path := req.URL.Path
 		handlerName := prepareHandlerName(path)
 
-		requestTimer := m.statsd.NewTiming()
+		requestTimer := s.client.NewTiming()
 
 		lrw := &LoggedResponseWriter{responseWriter: resp}
 
@@ -49,10 +46,11 @@ func (m *LogMiddleware) Log(handler http.Handler) http.Handler {
 			handlerName,
 		))
 
-		//Example: 200 POST /rec/ (127.0.0.1) 1.460s
-		log.Infof("%v %v %v (%v) %.3fs",
-			lrw.Status(), req.Method, path, req.RemoteAddr, requestTimer.Duration().Seconds())
-
+		if s.withLog {
+			//Example: 200 POST /rec/ (127.0.0.1) 1.460s
+			s.logger.Infof("%v %v %v (%v) %.3fs",
+				lrw.Status(), req.Method, path, req.RemoteAddr, requestTimer.Duration().Seconds())
+		}
 	})
 }
 

middleware/wrap.go

@@ -0,0 +1,38 @@
+package middleware
+
+import (
+	"net/http"
+	"runtime/debug"
+	"github.com/Sirupsen/logrus"
+)
+
+type Wrap struct {
+	logger *logrus.Entry
+	list   []func(http.Handler) http.Handler
+}
+
+func NewWrap(logger *logrus.Logger) *Wrap {
+	return &Wrap{
+		logger: logger.WithField("component", "middlewareWrap"),
+		list:   make([]func(http.Handler) http.Handler, 0),
+	}
+}
+
+func (s *Wrap) Add(middleware func(http.Handler) http.Handler) {
+	s.list = append(s.list, middleware)
+}
+
+func (s *Wrap) Do(handler http.Handler) http.Handler {
+	for i := 0; i < len(s.list); i++ {
+		handler = s.list[i](handler)
+	}
+	return http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
+		defer func() {
+			if err := recover(); err != nil {
+				s.logger.Errorf("Panic: %+v\n%s", err, debug.Stack())
+			}
+		}()
+
+		handler.ServeHTTP(resp, req)
+	})
+}

middleware/wrap_test.go

@@ -0,0 +1,52 @@
+package middleware
+
+import (
+	"testing"
+	"github.com/Sirupsen/logrus"
+	"net/http"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWrap_Add(t *testing.T) {
+	a := assert.New(t)
+	wrap := NewWrap(logrus.StandardLogger())
+	eLen := 5
+	for i := 0; i < eLen; i++ {
+		wrap.Add(func(handler http.Handler) http.Handler { return handler })
+	}
+	a.Len(wrap.list, eLen)
+}
+
+func TestWrap_Do(t *testing.T) {
+	a := assert.New(t)
+	wrap := NewWrap(logrus.StandardLogger())
+	checkArray := make([]int, 0)
+
+	wrap.Add(
+		func(handler http.Handler) http.Handler {
+			checkArray = append(checkArray, 1)
+			return handler
+		},
+	)
+
+	wrap.Add(
+		func(handler http.Handler) http.Handler {
+			checkArray = append(checkArray, 2)
+			return handler
+		},
+	)
+
+	wrap.Add(
+		func(handler http.Handler) http.Handler {
+			checkArray = append(checkArray, 3)
+			return handler
+		},
+	)
+
+	wrap.Do(http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {}))
+
+	a.Len(checkArray, 3)
+	a.Equal(1, checkArray[0])
+	a.Equal(2, checkArray[1])
+	a.Equal(3, checkArray[2])
+}