Skip to content

Commit

Permalink
chore: add some basic ratelimit support
Browse files Browse the repository at this point in the history
  • Loading branch information
kashalls committed Aug 7, 2024
1 parent 205dc97 commit e80caa8
Show file tree
Hide file tree
Showing 5 changed files with 192 additions and 7 deletions.
16 changes: 12 additions & 4 deletions cmd/kromgo/init/configuration/configuration.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,20 @@ import (
)

type ServerConfig struct {
ServerHost string `env:"SERVER_HOST" envDefault:"localhost"`
ServerPort int `env:"SERVER_PORT" envDefault:"8080"`
ServerHost string `env:"SERVER_HOST" envDefault:"localhost"`
ServerPort int `env:"SERVER_PORT" envDefault:"8080"`

HealthHost string `env:"HEALTH_HOST" envDefault:"localhost"`
HealthPort int `env:"HEALTH_PORT" envDefault:"8888"`

ServerReadTimeout time.Duration `env:"SERVER_READ_TIMEOUT"`
ServerWriteTimeout time.Duration `env:"SERVER_WRITE_TIMEOUT"`
HealthHost string `env:"HEALTH_HOST" envDefault:"localhost"`
HealthPort int `env:"HEALTH_PORT" envDefault:"8888"`

RatelimitEnable bool `env:"RATELIMIT_ENABLE"`
RatelimitAll bool `env:"RATELIMIT_ALL"`
RatelimitByRealIP bool `env:"RATELIMIT_BY_REAL_IP"`
RatelimitRequestLimit int `env:"RATELIMIT_REQUEST_LIMIT" envDefault:"100"`
RatelimitWindowLength time.Duration `env:"RATELIMIT_WINDOW_LENGTH" envDefault:"60000"`
}

// KromgoConfig struct for configuration environmental variables
Expand Down
12 changes: 12 additions & 0 deletions cmd/kromgo/init/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"time"

"github.com/go-chi/chi/v5"
"github.com/go-chi/httprate"
"github.com/kashalls/kromgo/cmd/kromgo/init/configuration"
"github.com/kashalls/kromgo/cmd/kromgo/init/log"
"github.com/kashalls/kromgo/pkg/kromgo"
Expand All @@ -34,6 +35,17 @@ func ReadinessHandler(w http.ResponseWriter, r *http.Request) {
func Init(config configuration.KromgoConfig, serverConfig configuration.ServerConfig) (*http.Server, *http.Server) {

mainRouter := chi.NewRouter()

if serverConfig.RatelimitEnable {
if serverConfig.RatelimitAll {
mainRouter.Use(httprate.LimitAll(serverConfig.RatelimitRequestLimit, serverConfig.RatelimitWindowLength))
} else if serverConfig.RatelimitByRealIP {
mainRouter.Use(httprate.LimitByRealIP(serverConfig.RatelimitRequestLimit, serverConfig.RatelimitWindowLength))
} else {
mainRouter.Use(httprate.LimitByIP(serverConfig.RatelimitRequestLimit, serverConfig.RatelimitWindowLength))
}
}

mainRouter.Get("/{metric}", func(w http.ResponseWriter, r *http.Request) {
kromgo.KromgoRequestHandler(w, r, config)
})
Expand Down
6 changes: 5 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,11 @@ go 1.22.1
require (
github.com/caarlos0/env/v11 v11.2.0
github.com/go-chi/chi/v5 v5.1.0
github.com/go-chi/httprate v0.12.0
github.com/invopop/jsonschema v0.12.0
github.com/prometheus/client_golang v1.19.1
github.com/prometheus/common v0.55.0
github.com/stretchr/testify v1.9.0
go.uber.org/zap v1.27.0
gopkg.in/yaml.v3 v3.0.1
)
Expand All @@ -16,12 +18,14 @@ require (
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/buger/jsonparser v1.1.1 // indirect
github.com/cespare/xxhash/v2 v2.2.0 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/mailru/easyjson v0.7.7 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/prometheus/client_model v0.6.1 // indirect
github.com/prometheus/procfs v0.15.1 // indirect
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
Expand Down
8 changes: 6 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,15 @@ github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMU
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/caarlos0/env/v11 v11.2.0 h1:kvB1ZmwdWgI3JsuuVUE7z4cY/6Ujr03D0w2WkOOH4Xs=
github.com/caarlos0/env/v11 v11.2.0/go.mod h1:LwgkYk1kDvfGpHthrWWLof3Ny7PezzFwS4QrsJdHTMo=
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/go-chi/chi/v5 v5.1.0 h1:acVI1TYaD+hhedDJ3r54HyA6sExp3HfXq7QWEEY/xMw=
github.com/go-chi/chi/v5 v5.1.0/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8=
github.com/go-chi/httprate v0.12.0 h1:08D/te3pOTJe5+VAZTQrHxwdsH2NyliiUoRD1naKaMg=
github.com/go-chi/httprate v0.12.0/go.mod h1:TUepLXaz/pCjmCtf/obgOQJ2Sz6rC8fSf5cAt5cnTt0=
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
Expand Down Expand Up @@ -66,6 +68,8 @@ golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
golang.org/x/oauth2 v0.21.0 h1:tsimM75w1tF/uws5rbeHzIWxEqElMehnc+iW793zsZs=
golang.org/x/oauth2 v0.21.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
Expand Down
157 changes: 157 additions & 0 deletions pkg/kromgo/utils_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,157 @@
package kromgo

import (
"testing"

"github.com/kashalls/kromgo/cmd/kromgo/init/configuration"
"github.com/prometheus/common/model"
"github.com/stretchr/testify/assert"
)

func TestGetColorConfig_MatchingRange(t *testing.T) {
colors := []configuration.MetricColor{
{Min: 0, Max: 10, Color: "blue", ValueOverride: "low"},
{Min: 11, Max: 20, Color: "green", ValueOverride: "medium"},
{Min: 21, Max: 30, Color: "red", ValueOverride: "high"},
}

value := 15.0

result := GetColorConfig(colors, value)

expected := configuration.MetricColor{Min: 11, Max: 20, Color: "green", ValueOverride: "medium"}
assert.Equal(t, expected, result)
}

func TestGetColorConfig_ExactMatch(t *testing.T) {
colors := []configuration.MetricColor{
{Min: 0, Max: 10, Color: "blue"},
{Min: 10, Max: 20, Color: "green"},
}

value := 10.0

result := GetColorConfig(colors, value)

expected := configuration.MetricColor{Min: 10, Max: 20, Color: "green"}
assert.Equal(t, expected, result)
}

func TestGetColorConfig_NoMatch(t *testing.T) {
colors := []configuration.MetricColor{
{Min: 0, Max: 10, Color: "blue"},
{Min: 11, Max: 20, Color: "green"},
}

value := 25.0

result := GetColorConfig(colors, value)

expected := configuration.MetricColor{Min: 25, Max: 25}
assert.Equal(t, expected, result)
}

func TestGetColorConfig_EmptyColors(t *testing.T) {
colors := []configuration.MetricColor{}

value := 10.0

result := GetColorConfig(colors, value)

expected := configuration.MetricColor{Min: 10, Max: 10}
assert.Equal(t, expected, result)
}

func TestGetColorConfig_ValueBelowMin(t *testing.T) {
colors := []configuration.MetricColor{
{Min: 10, Max: 20, Color: "green"},
{Min: 21, Max: 30, Color: "red"},
}

value := 5.0

result := GetColorConfig(colors, value)

expected := configuration.MetricColor{Min: 5, Max: 5}
assert.Equal(t, expected, result)
}

func TestGetColorConfig_ValueAboveMax(t *testing.T) {
colors := []configuration.MetricColor{
{Min: 0, Max: 10, Color: "blue"},
{Min: 11, Max: 20, Color: "green"},
}

value := 25.0

result := GetColorConfig(colors, value)

expected := configuration.MetricColor{Min: 25, Max: 25}
assert.Equal(t, expected, result)
}

func TestExtractLabelValue_LabelExists(t *testing.T) {
vector := model.Vector{
&model.Sample{
Metric: model.Metric{
"label1": "value1",
"label2": "value2",
},
},
}

labelName := "label1"
expectedValue := "value1"

value, err := ExtractLabelValue(vector, labelName)

assert.NoError(t, err)
assert.Equal(t, expectedValue, value)
}

func TestExtractLabelValue_LabelDoesNotExist(t *testing.T) {
vector := model.Vector{
&model.Sample{
Metric: model.Metric{
"label1": "value1",
},
},
}

labelName := "label2"

value, err := ExtractLabelValue(vector, labelName)

assert.Error(t, err)
assert.Equal(t, "", value)
assert.Equal(t, "label 'label2' not found in the query result", err.Error())
}

func TestExtractLabelValue_EmptyVector(t *testing.T) {
vector := model.Vector{}
labelName := "label1"

value, err := ExtractLabelValue(vector, labelName)

assert.Error(t, err)
assert.Equal(t, "", value)
assert.Equal(t, "label 'label1' not found in the query result", err.Error())
}

func TestExtractLabelValue_LabelEmptyValue(t *testing.T) {
vector := model.Vector{
&model.Sample{
Metric: model.Metric{
"label1": "", // Empty string value for the label
},
},
}

labelName := "label1"
expectedValue := ""

value, err := ExtractLabelValue(vector, labelName)

assert.NoError(t, err)
assert.Equal(t, expectedValue, value)
}

0 comments on commit e80caa8

Please sign in to comment.