Implementing database metrics, configuring DB connection settings and integrating fullPath rate limiter (#43)

* Register and sync db metrics

* fix mock

* Add RateLimiterByFullPath

* TestRateLimitByFullPath

Author: lucaslopezf

pkg/zdb/methods.go

@@ -132,6 +132,11 @@ func (z *zDatabase) Scopes(funcs ...func(ZDatabase) ZDatabase) ZDatabase {
 	return wrap(z.db.Scopes(gormFuncs...))
 }
 
+func (z *zDatabase) GetDBStats() (sql.DBStats, error) {
+	sqlDB, err := z.db.DB()
+	return sqlDB.Stats(), err
+}
+
 func (z *zDatabase) Error() error {
 	return z.db.Error
 }

pkg/zdb/metrics.go

@@ -0,0 +1,64 @@
+package zdb
+
+import (
+	"fmt"
+	"github.com/zondax/golem/pkg/metrics"
+	"github.com/zondax/golem/pkg/metrics/collectors"
+	"go.uber.org/zap"
+	"time"
+)
+
+const (
+	defaultInterval                = time.Minute
+	dbOpenConnectionsMetricName    = "db_open_connections"
+	dbIdleConnectionsMetricName    = "db_idle_connections"
+	dbMaxOpenConnectionsMetricName = "db_max_open_connections"
+	dbWaitDurationMetricName       = "db_wait_duration"
+	dbInUseConnectionsMetricName   = "db_in_use_connections"
+)
+
+func SetupAndMonitorDBMetrics(appName string, metricsServer metrics.TaskMetrics, db ZDatabase, updateInterval time.Duration) []error {
+	if updateInterval <= 0 {
+		updateInterval = defaultInterval
+	}
+
+	var errs []error
+	register := func(name, help string, labels []string, handler metrics.MetricHandler) {
+		if err := metricsServer.RegisterMetric(name, help, labels, handler); err != nil {
+			errs = append(errs, err)
+		}
+	}
+
+	register(getMetricName(appName, dbOpenConnectionsMetricName), "Number of open database connections.", nil, &collectors.Gauge{})
+	register(getMetricName(appName, dbIdleConnectionsMetricName), "Number of idle database connections in the pool.", nil, &collectors.Gauge{})
+	register(getMetricName(appName, dbMaxOpenConnectionsMetricName), "Maximum number of open database connections.", nil, &collectors.Gauge{})
+	register(getMetricName(appName, dbWaitDurationMetricName), "Total time waited for new database connections.", nil, &collectors.Histogram{})
+	register(getMetricName(appName, dbInUseConnectionsMetricName), "Number of database connections currently in use.", nil, &collectors.Gauge{})
+
+	if len(errs) > 0 {
+		return errs
+	}
+
+	go func() {
+		ticker := time.NewTicker(updateInterval)
+		for range ticker.C {
+			stats, err := db.GetDBStats()
+			if err != nil {
+				zap.S().Errorf("Error while getting db stats: %v", err)
+				continue
+			}
+
+			_ = metricsServer.UpdateMetric(getMetricName(appName, dbOpenConnectionsMetricName), float64(stats.OpenConnections))
+			_ = metricsServer.UpdateMetric(getMetricName(appName, dbIdleConnectionsMetricName), float64(stats.Idle))
+			_ = metricsServer.UpdateMetric(getMetricName(appName, dbMaxOpenConnectionsMetricName), float64(stats.MaxOpenConnections))
+			_ = metricsServer.UpdateMetric(getMetricName(appName, dbWaitDurationMetricName), float64(stats.WaitDuration.Milliseconds()))
+			_ = metricsServer.UpdateMetric(getMetricName(appName, dbInUseConnectionsMetricName), float64(stats.InUse))
+		}
+	}()
+
+	return nil
+}
+
+func getMetricName(appName, metricType string) string {
+	return fmt.Sprintf("zdatabase_%s_%s", appName, metricType)
+}

pkg/zdb/mock_zdb.go

@@ -142,6 +142,11 @@ func (m *MockZDatabase) Group(name string) ZDatabase {
 	return args.Get(0).(ZDatabase)
 }
 
+func (m *MockZDatabase) GetDBStats() (sql.DBStats, error) {
+	args := m.Called()
+	return args.Get(0).(sql.DBStats), args.Error(1)
+}
+
 // MockDBConnector
 
 type MockDBConnector struct {

pkg/zdb/zdb.go

@@ -44,6 +44,7 @@ type ZDatabase interface {
 	Scopes(funcs ...func(ZDatabase) ZDatabase) ZDatabase
 	RowsAffected() int64
 	GetDbConnection() *gorm.DB
+	GetDBStats() (sql.DBStats, error)
 }
 
 type zDatabase struct {

pkg/zdb/zdbconfig/zdbconfig.go

@@ -1,6 +1,9 @@
 package zdbconfig
 
-import "gorm.io/gorm"
+import (
+	"gorm.io/gorm"
+	"time"
+)
 
 type ConnectionParams struct {
 	User     string
@@ -17,6 +20,9 @@ type Config struct {
 	MaxAttempts      int
 	ConnectionParams ConnectionParams
 	LogConfig        LogConfig
+	MaxIdleConns     int
+	MaxOpenConns     int
+	ConnMaxLifetime  time.Duration
 }
 
 type LogConfig struct {

pkg/zdb/zdbconnector/clickhouse.go

@@ -24,6 +24,23 @@ func (c *ClickHouseConnector) Connect(config *zdbconfig.Config) (*gorm.DB, error
 	if err != nil {
 		return nil, err
 	}
+	db, err := dbConn.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	if config.MaxIdleConns != 0 {
+		db.SetMaxIdleConns(config.MaxIdleConns)
+	}
+
+	if config.MaxOpenConns != 0 {
+		db.SetMaxOpenConns(config.MaxOpenConns)
+	}
+
+	if config.ConnMaxLifetime != 0 {
+		db.SetConnMaxLifetime(config.ConnMaxLifetime)
+	}
+
 	return dbConn, nil
 }
 

pkg/zrouter/zmiddlewares/rate_limit.go

@@ -3,9 +3,44 @@ package zmiddlewares
 import (
 	"golang.org/x/time/rate"
 	"net/http"
+	"sync"
 	"time"
 )
 
+var (
+	limiters = make(map[string]*rate.Limiter)
+	mu       sync.Mutex
+)
+
+func getLimiter(key string, maxRPM int) *rate.Limiter {
+	mu.Lock()
+	defer mu.Unlock()
+
+	limiter, exists := limiters[key]
+	if !exists {
+		limiter = rate.NewLimiter(rate.Every(time.Minute/time.Duration(maxRPM)), maxRPM)
+		limiters[key] = limiter
+	}
+
+	return limiter
+}
+
+func RateLimitByFullPath(maxRPM int) Middleware {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			key := r.URL.Path
+			limiter := getLimiter(key, maxRPM)
+
+			if !limiter.Allow() {
+				http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
 func RateLimit(maxRPM int) Middleware {
 	limiter := rate.NewLimiter(rate.Every(time.Minute/time.Duration(maxRPM)), maxRPM)
 

pkg/zrouter/zmiddlewares/rate_limit_test.go

@@ -28,3 +28,24 @@ func TestRateLimit(t *testing.T) {
 	r.ServeHTTP(rec, req)
 	assert.Equal(t, http.StatusTooManyRequests, rec.Code)
 }
+
+func TestRateLimitByFullPath(t *testing.T) {
+	r := chi.NewRouter()
+
+	r.Use(RateLimitByFullPath(1))
+
+	r.Get("/tests/{id}", func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("OK"))
+	})
+
+	req := httptest.NewRequest("GET", "/tests/1", nil)
+
+	rec := httptest.NewRecorder()
+	r.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "OK", rec.Body.String())
+
+	rec = httptest.NewRecorder()
+	r.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusTooManyRequests, rec.Code)
+}