Added host rules + minor refactoring

config.go

@@ -1,12 +1,18 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/pkg/errors"
+	"github.com/ryanuber/go-glob"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/time/rate"
 	"io/ioutil"
 	"os"
+	"reflect"
+	"runtime"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -16,7 +22,47 @@ type HostConfig struct {
 	EveryStr string            `json:"every"`
 	Burst    int               `json:"burst"`
 	Headers  map[string]string `json:"headers"`
+	RawRules []*RawHostRule    `json:"rules"`
 	Every    time.Duration
+	Rules    []*HostRule
+}
+
+type RawHostRule struct {
+	Condition string `json:"condition"`
+	Action    string `json:"action"`
+	Arg       string `json:"arg"`
+}
+
+type HostRuleAction int
+
+const (
+	DontRetry     HostRuleAction = 0
+	MultiplyEvery HostRuleAction = 1
+	SetEvery      HostRuleAction = 2
+	ForceRetry    HostRuleAction = 3
+	ShouldRetry   HostRuleAction = 4
+)
+
+func (a HostRuleAction) String() string {
+	switch a {
+	case DontRetry:
+		return "dont_retry"
+	case MultiplyEvery:
+		return "multiply_every"
+	case SetEvery:
+		return "set_every"
+	case ForceRetry:
+		return "force_retry"
+	case ShouldRetry:
+		return "should_retry"
+	}
+	return "???"
+}
+
+type HostRule struct {
+	Matches func(r *RequestCtx) bool
+	Action  HostRuleAction
+	Arg     float64
 }
 
 type ProxyConfig struct {
@@ -30,6 +76,7 @@ var config struct {
 	WaitStr       string        `json:"wait"`
 	Multiplier    float64       `json:"multiplier"`
 	Retries       int           `json:"retries"`
+	RetriesHard   int           `json:"retries_hard"`
 	Hosts         []*HostConfig `json:"hosts"`
 	Proxies       []ProxyConfig `json:"proxies"`
 	Wait          int64
@@ -37,16 +84,202 @@ var config struct {
 	DefaultConfig *HostConfig
 }
 
-func loadConfig() {
+func parseRule(raw *RawHostRule) (*HostRule, error) {
+	//TODO: for the love of god someone please refactor this func
+
+	rule := &HostRule{}
+	var err error
+
+	switch raw.Action {
+	case "should_retry":
+		rule.Action = ShouldRetry
+	case "dont_retry":
+		rule.Action = DontRetry
+	case "multiply_every":
+		rule.Action = MultiplyEvery
+		rule.Arg, err = strconv.ParseFloat(raw.Arg, 64)
+	case "set_every":
+		rule.Action = SetEvery
+		var duration time.Duration
+		duration, err = time.ParseDuration(raw.Arg)
+		if err != nil {
+			return nil, err
+		}
+		rule.Arg = 1 / duration.Seconds()
+	case "force_retry":
+		rule.Action = ForceRetry
+	default:
+		return nil, errors.Errorf("Invalid argument for action: %s", raw.Action)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	switch {
+	case strings.Contains(raw.Condition, "!="):
+		op1Str, op2Str := split(raw.Condition, "!=")
+		op1Func := parseOperand1(op1Str)
+		if op1Func == nil {
+			return nil, errors.Errorf("Invalid rule: %s", raw.Condition)
+		}
+
+		if isGlob(op2Str) {
+			rule.Matches = func(ctx *RequestCtx) bool {
+				return !glob.Glob(op2Str, op1Func(ctx))
+			}
+		} else {
+			op2Str = strings.Replace(op2Str, "\\*", "*", -1)
+			rule.Matches = func(ctx *RequestCtx) bool {
+				return op1Func(ctx) != op2Str
+			}
+		}
+	case strings.Contains(raw.Condition, "="):
+		op1Str, op2Str := split(raw.Condition, "=")
+		op1Func := parseOperand1(op1Str)
+		if op1Func == nil {
+			return nil, errors.Errorf("Invalid rule: %s", raw.Condition)
+		}
+
+		if isGlob(op2Str) {
+			rule.Matches = func(ctx *RequestCtx) bool {
+				return glob.Glob(op2Str, op1Func(ctx))
+			}
+		} else {
+			op2Str = strings.Replace(op2Str, "\\*", "*", -1)
+			rule.Matches = func(ctx *RequestCtx) bool {
+				return op1Func(ctx) == op2Str
+			}
+		}
+	case strings.Contains(raw.Condition, ">"):
+		op1Str, op2Str := split(raw.Condition, ">")
+		op1Func := parseOperand1(op1Str)
+		if op1Func == nil {
+			return nil, errors.Errorf("Invalid rule: %s", raw.Condition)
+		}
+		op2Num, err := parseOperand2(op1Str, op2Str)
+		if err != nil {
+			return nil, err
+		}
+
+		rule.Matches = func(ctx *RequestCtx) bool {
+			op1Num, err := strconv.ParseFloat(op1Func(ctx), 64)
+			handleRuleErr(err)
+			return op1Num > op2Num
+		}
+	case strings.Contains(raw.Condition, "<"):
+		op1Str, op2Str := split(raw.Condition, "<")
+		op1Func := parseOperand1(op1Str)
+		if op1Func == nil {
+			return nil, errors.Errorf("Invalid rule: %s", raw.Condition)
+		}
+		op2Num, err := parseOperand2(op1Str, op2Str)
+		if err != nil {
+			return nil, err
+		}
+
+		rule.Matches = func(ctx *RequestCtx) bool {
+			op1Num, err := strconv.ParseFloat(op1Func(ctx), 64)
+			handleRuleErr(err)
+			return op1Num < op2Num
+		}
+	}
+
+	return rule, nil
+}
+
+func handleRuleErr(err error) {
+	if err != nil {
+		logrus.WithError(err).Warn("Error computing rule")
+	}
+}
+
+func split(str, subStr string) (string, string) {
+
+	str1 := str[:strings.Index(str, subStr)]
+	str2 := str[strings.Index(str, subStr)+len(subStr):]
+
+	return str1, str2
+}
+
+func parseOperand2(op1, op2 string) (float64, error) {
+	if op1 == "response_time" {
+		res, err := time.ParseDuration(op2)
+		if err != nil {
+			return -1, err
+		}
+		return res.Seconds(), nil
+	}
+
+	return strconv.ParseFloat(op2, 64)
+}
+
+func parseOperand1(op string) func(ctx *RequestCtx) string {
+	switch {
+	case op == "body":
+		return func(ctx *RequestCtx) string {
+
+			if ctx.Response == nil {
+				return ""
+			}
+			bodyBytes, err := ioutil.ReadAll(ctx.Response.Body)
+			if err != nil {
+				return ""
+			}
+			err = ctx.Response.Body.Close()
+			if err != nil {
+				return ""
+			}
+			ctx.Response.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))
+
+			return string(bodyBytes)
+		}
+	case op == "status":
+		return func(ctx *RequestCtx) string {
+			if ctx.Response == nil {
+				return ""
+			}
+			return strconv.Itoa(ctx.Response.StatusCode)
+		}
+	case op == "response_time":
+		return func(ctx *RequestCtx) string {
+			return strconv.FormatFloat(time.Now().Sub(ctx.RequestTime).Seconds(), 'f', 6, 64)
+		}
+	case strings.HasPrefix(op, "header:"):
+		header := op[strings.Index(op, ":")+1:]
+		return func(ctx *RequestCtx) string {
+			if ctx.Response == nil {
+				return ""
+			}
+			return ctx.Response.Header.Get(header)
+		}
+	default:
+		return nil
+	}
+}
+
+func isGlob(op string) bool {
+	tmpStr := strings.Replace(op, "\\*", "_", -1)
+
+	return strings.Contains(tmpStr, "*")
+}
+
+func loadConfig() error {
 
 	configFile, err := os.Open("config.json")
-	handleErr(err)
+	if err != nil {
+		return err
+	}
 
 	configBytes, err := ioutil.ReadAll(configFile)
-	handleErr(err)
+	if err != nil {
+		return err
+	}
 
 	err = json.Unmarshal(configBytes, &config)
-	handleErr(err)
+	if err != nil {
+		return err
+	}
 
 	validateConfig()
 
@@ -54,18 +287,52 @@ func loadConfig() {
 	wait, err := time.ParseDuration(config.WaitStr)
 	config.Wait = int64(wait)
 
-	for _, conf := range config.Hosts {
+	for i, conf := range config.Hosts {
 		if conf.EveryStr == "" {
-			conf.Every = config.DefaultConfig.Every
+			// Look 'upwards' for every
+			for _, prevConf := range config.Hosts[:i] {
+				if glob.Glob(prevConf.Host, conf.Host) {
+					conf.Every = prevConf.Every
+				}
+			}
 		} else {
 			conf.Every, err = time.ParseDuration(conf.EveryStr)
 			handleErr(err)
 		}
 
-		if config.DefaultConfig != nil && conf.Burst == 0 {
-			conf.Burst = config.DefaultConfig.Burst
+		if conf.Burst == 0 {
+			// Look 'upwards' for burst
+			for _, prevConf := range config.Hosts[:i] {
+				if glob.Glob(prevConf.Host, conf.Host) {
+					conf.Burst = prevConf.Burst
+				}
+			}
 		}
+		if conf.Burst == 0 {
+			return errors.Errorf("Burst must be > 0 (Host: %s)", conf.Host)
+		}
+
+		for _, rawRule := range conf.RawRules {
+			r, err := parseRule(rawRule)
+			handleErr(err)
+			conf.Rules = append(conf.Rules, r)
+
+			logrus.WithFields(logrus.Fields{
+				"arg":       r.Arg,
+				"action":    r.Action,
+				"matchFunc": runtime.FuncForPC(reflect.ValueOf(r.Matches).Pointer()).Name(),
+			}).Info("Rule")
+		}
+
+		logrus.WithFields(logrus.Fields{
+			"every":   conf.Every,
+			"burst":   conf.Burst,
+			"headers": conf.Headers,
+			"host":    conf.Host,
+		}).Info("Host")
 	}
+
+	return nil
 }
 
 func validateConfig() {
@@ -91,18 +358,28 @@ func validateConfig() {
 
 func applyConfig(proxy *Proxy) {
 
-	for _, conf := range config.Hosts {
-		proxy.Limiters[conf.Host] = &ExpiringLimiter{
-			rate.NewLimiter(rate.Every(conf.Every), conf.Burst),
-			time.Now(),
-		}
+	//Reverse order
+	for i := len(config.Hosts) - 1; i >= 0; i-- {
+
+		conf := config.Hosts[i]
+
+		proxy.Limiters = append(proxy.Limiters, &ExpiringLimiter{
+			HostGlob:  conf.Host,
+			IsGlob:    isGlob(conf.Host),
+			Limiter:   rate.NewLimiter(rate.Every(conf.Every), conf.Burst),
+			LastRead:  time.Now(),
+			CanDelete: false,
+		})
 	}
 }
 
 func (b *Balancer) reloadConfig() {
 
 	b.proxyMutex.Lock()
-	loadConfig()
+	err := loadConfig()
+	if err != nil {
+		panic(err)
+	}
 
 	if b.proxies != nil {
 		b.proxies = b.proxies[:0]