kapacitor-scale/handler/lib.go

package handler

import (
	"errors"
	"fmt"
	"log"
	"strconv"
	"time"

	"git.higgsboson.tk/Mic92/kapacitor-scale/scaling"
	"github.com/influxdata/kapacitor/udf"
	"github.com/influxdata/kapacitor/udf/agent"
	"github.com/pk-rawat/gostr/src"
)

type Handler struct {
	Id, When, By    string
	MinInstances    int64
	MaxInstances    int64
	Cooldown        time.Duration
	Simulate, Debug bool
	kapacitorAgent  *agent.Agent
	scaleAgent      *scaling.Agent
}

func New(kapacitorAgent *agent.Agent, scaleAgent *scaling.Agent) *Handler {
	h := Handler{}
	h.kapacitorAgent = kapacitorAgent
	h.scaleAgent = scaleAgent
	return &h
}

// Return the InfoResponse. Describing the properties of this UDF agent.
func (*Handler) Info() (*udf.InfoResponse, error) {
	info := &udf.InfoResponse{
		Wants:    udf.EdgeType_STREAM,
		Provides: udf.EdgeType_STREAM,
		Options: map[string]*udf.OptionInfo{
			"id":            {ValueTypes: []udf.ValueType{udf.ValueType_STRING}},
			"when":          {ValueTypes: []udf.ValueType{udf.ValueType_STRING}},
			"by":            {ValueTypes: []udf.ValueType{udf.ValueType_STRING}},
			"min_instances": {ValueTypes: []udf.ValueType{udf.ValueType_INT}},
			"max_instances": {ValueTypes: []udf.ValueType{udf.ValueType_INT}},
			"cooldown":      {ValueTypes: []udf.ValueType{udf.ValueType_STRING}},
			"simulate":      {ValueTypes: []udf.ValueType{udf.ValueType_BOOL}},
			"debug":         {ValueTypes: []udf.ValueType{udf.ValueType_BOOL}},
		},
	}
	return info, nil
}

func (h *Handler) debug(format string, args ...interface{}) {
	if h.Debug {
		log.Printf(format, args...)
	}
}

// Initialze the handler based of the provided options.
func (h *Handler) Init(r *udf.InitRequest) (*udf.InitResponse, error) {
	init := &udf.InitResponse{Success: true, Error: ""}
	h.Debug = false
	h.Simulate = false
	h.Cooldown = time.Minute
	h.By = "current + 1"
	h.MinInstances = 1
	h.MaxInstances = 3

	var cooldown string
	for _, opt := range r.Options {
		switch opt.Name {
		case "when":
			h.When = opt.Values[0].Value.(*udf.OptionValue_StringValue).StringValue
		case "by":
			h.By = opt.Values[0].Value.(*udf.OptionValue_StringValue).StringValue
		case "min_instances":
			h.MinInstances = opt.Values[0].Value.(*udf.OptionValue_IntValue).IntValue
		case "max_instances":
			h.MaxInstances = opt.Values[0].Value.(*udf.OptionValue_IntValue).IntValue
		case "cooldown":
			cooldown = opt.Values[0].Value.(*udf.OptionValue_StringValue).StringValue
		case "id":
			h.Id = opt.Values[0].Value.(*udf.OptionValue_StringValue).StringValue
		case "simulate":
			h.Simulate = opt.Values[0].Value.(*udf.OptionValue_BoolValue).BoolValue
		case "debug":
			h.Debug = opt.Values[0].Value.(*udf.OptionValue_BoolValue).BoolValue
		}
	}

	if h.When == "" {
		init.Success = false
		init.Error += " must supply `when` expression;"
	}
	if h.By == "" {
		init.Success = false
		init.Error += " must supply `by` expression;"
	}
	if h.MinInstances < 0 {
		init.Success = false
		init.Error += " `MinInstances` must be greater equal 0;"
	}
	if h.MaxInstances < 0 {
		init.Success = false
		init.Error += " `MaxInstances` must be greater equal 0;"
	}

	if h.MaxInstances < h.MinInstances {
		init.Success = false
		init.Error += " `MaxInstances` must be greater equal minimum instances;"
	}

	var err error
	h.Cooldown, err = time.ParseDuration(cooldown)
	if err != nil {
		init.Success = false
		init.Error += fmt.Sprintf(" `cooldown` '%s' has an invalid format: %v", cooldown, err)
	}
	if h.Cooldown < 0 {
		init.Success = false
		init.Error += " `cooldown` must be greater equal 0s"
	}

	return init, nil
}

// Create a snapshot of the running state of the process.
func (o *Handler) Snaphost() (*udf.SnapshotResponse, error) {
	return &udf.SnapshotResponse{}, nil
}

// Restore a previous snapshot.
func (o *Handler) Restore(req *udf.RestoreRequest) (*udf.RestoreResponse, error) {
	return &udf.RestoreResponse{
		Success: true,
	}, nil
}

// Start working with the next batch
func (h *Handler) BeginBatch(begin *udf.BeginBatch) error {
	return errors.New("batching not supported")
}

func (h *Handler) evaluateWhen(p *udf.Point) (bool, error) {
	fields := make(map[string]interface{})
	for k, v := range p.GetFieldsInt() {
		fields[k] = v
	}
	for k, v := range p.GetFieldsDouble() {
		fields[k] = v
	}
	res := gostr.Evaluate(h.When, fields).(string)
	doScale, err := strconv.ParseBool(res)
	if err != nil {
		return false, fmt.Errorf("the expression `when` should evaluate to true or false, got %s", res)
	}
	h.debug("evaluate '%s' for '%v' -> should scale: %v", h.When, fields, doScale)
	return doScale, nil
}

func (h *Handler) evaluateBy(s *scaling.Service) (int64, error) {
	scaleContext := make(map[string]interface{})
	scaleContext["current"] = s.CurrentInstances
	res := gostr.Evaluate(h.By, scaleContext).(string)
	amount, err := strconv.ParseInt(res, 10, 64)
	if err != nil {
		return -1, fmt.Errorf("the expression `scale` should evaluate to an integer value, got %s (tipp: there is an ROUND() method)", res)
	}
	if amount < h.MinInstances {
		return h.MinInstances, nil
	}
	if amount > h.MaxInstances {
		return h.MaxInstances, nil
	}
	return amount, nil
}

func (h *Handler) Point(p *udf.Point) error {
	doScale, err := h.evaluateWhen(p)
	if !doScale {
		return err
	}
	service, err := h.scaleAgent.RequestScaling(h.Id, time.Unix(0, p.Time))
	if err != nil {
		return fmt.Errorf("Failed start scaling: %v", err)
	}
	if service == nil {
		h.debug("skip scaling because of cooldown")
		return nil
	}
	defer service.Unlock()
	to, err := h.evaluateBy(service)
	if err != nil {
		return err
	}
	if to == service.CurrentInstances {
		h.debug("skip scaling service '%s' still %d", h.Id, to)
		return nil
	}
	h.debug("attempt to scale service '%s' from %d to %d", h.Id, service.CurrentInstances, to)
	if !h.Simulate {
		err = h.scaleAgent.Scale(h.Id, to)
		if err != nil {
			return err
		}
	}
	service.CurrentInstances = to
	service.CooldownUntil = time.Now().Add(h.Cooldown)
	p.FieldsDouble = make(map[string]float64)
	p.FieldsInt = map[string]int64{"scale": to}
	p.FieldsString = make(map[string]string)
	h.kapacitorAgent.Responses <- &udf.Response{
		Message: &udf.Response_Point{
			Point: p,
		},
	}
	return nil
}

func (o *Handler) EndBatch(end *udf.EndBatch) error {
	return nil
}

// Stop the handler gracefully.
func (o *Handler) Stop() {
	close(o.kapacitorAgent.Responses)
}