<commit_before>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package optimize

import (
	"fmt"
	"math"
	"time"

	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

// Local finds a local minimum of a function using a sequential algorithm.
// In order to maximize a function, multiply the output by -1.
//
// The first argument is of Function type representing the function to be minimized.
// Type switching is used to see if the function implements Gradient, FunctionGradient
// and Statuser.
//
// The second argument is the initial location at which to start the minimization.
// The initial location must be supplied, and must have a length equal to the
// problem dimension.
//
// The third argument contains the settings for the minimization. It is here that
// gradient tolerance, etc. are specified. The DefaultSettings() function
// can be called for a Settings struct with the default values initialized.
// If settings == nil, the default settings are used. Please see the documentation
// for the Settings structure for more information. The optimization Method used
// may also contain settings, see documentation for the appropriate optimizer.
//
// The final argument is the optimization method to use. If method == nil, then
// an appropriate default is chosen based on the properties of the other arguments
// (dimension, gradient-free or gradient-based, etc.). The optimization
// methods in this package are designed such that reasonable defaults occur
// if options are not specified explicitly. For example, the code
//  method := &Bfgs{}
// creates a pointer to a new Bfgs struct. When minimize is called, the settings
// in the method will be populated with default values. The methods are also
// designed such that they can be reused in future calls to method.
//
// Local returns a Result struct and any error that occurred. Please see the
// documentation of Result for more information.
//
// Please be aware that the default behavior of Local is to find the minimum.
// For certain functions and optimization methods, this process can take many
// function evaluations. If you would like to put limits on this, for example
// maximum runtime or maximum function evaluations, please modify the Settings
// input struct.
func Local(f Function, initX []float64, settings *Settings, method Method) (*Result, error) {
	if len(initX) == 0 {
		panic("local: initial X has zero length")
	}

	startTime := time.Now()
	funcInfo := newFunctionInfo(f)
	if method == nil {
		method = getDefaultMethod(funcInfo)
	}
	if err := funcInfo.satisfies(method); err != nil {
		return nil, err
	}

	if funcInfo.IsStatuser {
		_, err := funcInfo.statuser.Status()
		if err != nil {
			return nil, err
		}
	}

	if settings == nil {
		settings = DefaultSettings()
	}

	if settings.Recorder != nil {
		// Initialize Recorder first. If it fails, we avoid the (possibly
		// time-consuming) evaluation of F and DF at the starting location.
		err := settings.Recorder.Init(&funcInfo.FunctionInfo)
		if err != nil {
			return nil, err
		}
	}

	stats := &Stats{}
	optLoc, evalType, err := getStartingLocation(funcInfo, method, initX, stats, settings)
	if err != nil {
		return nil, err
	}

	// Runtime is the only Stats field that needs to be updated here.
	stats.Runtime = time.Since(startTime)
	// Send optLoc to Recorder before checking it for convergence.
	if settings.Recorder != nil {
		err = settings.Recorder.Record(optLoc, evalType, InitIteration, stats)
	}

	// Check if the starting location satisfies the convergence criteria.
	status := checkConvergence(optLoc, InitIteration, stats, settings)
	if status == NotTerminated && err == nil {
		// The starting location is not good enough, we need to perform a
		// minimization. The optimal location will be stored in-place in
		// optLoc.
		status, err = minimize(settings, method, funcInfo, stats, optLoc, startTime)
	}

	if settings.Recorder != nil && err == nil {
		// Send the optimal location to Recorder.
		err = settings.Recorder.Record(optLoc, NoEvaluation, PostIteration, stats)
	}
	stats.Runtime = time.Since(startTime)
	return &Result{
		Location: *optLoc,
		Stats:    *stats,
		Status:   status,
	}, err
}

func minimize(settings *Settings, method Method, funcInfo *functionInfo, stats *Stats, optLoc *Location, startTime time.Time) (status Status, err error) {
	loc := &Location{}
	copyLocation(loc, optLoc)
	xNext := make([]float64, len(loc.X))

	methodStatus, methodIsStatuser := method.(Statuser)

	evalType, iterType, err := method.Init(loc, &funcInfo.FunctionInfo, xNext)
	if err != nil {
		return Failure, err
	}

	for {
		if funcInfo.IsStatuser {
			// Check the function status before evaluating.
			status, err = funcInfo.statuser.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		// Perform evalType evaluation of the function at xNext and store the
		// result in location.
		evaluate(funcInfo, evalType, xNext, loc, stats)
		// Update the stats and optLoc.
		update(loc, optLoc, stats, iterType, startTime)
		// Get the convergence status before recording the new location.
		status = checkConvergence(optLoc, iterType, stats, settings)

		if settings.Recorder != nil {
			err = settings.Recorder.Record(loc, evalType, iterType, stats)
			if err != nil {
				if status == NotTerminated {
					status = Failure
				}
				return
			}
		}

		if status != NotTerminated {
			return
		}

		if methodIsStatuser {
			status, err = methodStatus.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		// Find the next location (stored in-place into xNext).
		evalType, iterType, err = method.Iterate(loc, xNext)
		if err != nil {
			status = Failure
			return
		}
	}
	panic("unreachable")
}

func copyLocation(dst, src *Location) {
	dst.X = resize(dst.X, len(src.X))
	copy(dst.X, src.X)

	dst.F = src.F

	dst.Gradient = resize(dst.Gradient, len(src.Gradient))
	copy(dst.Gradient, src.Gradient)
}

func getDefaultMethod(funcInfo *functionInfo) Method {
	if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
		return &BFGS{}
	}
	// TODO: Implement a gradient-free method
	panic("optimize: gradient-free methods not yet coded")
}

// getStartingLocation allocates and initializes the starting location for the minimization.
func getStartingLocation(funcInfo *functionInfo, method Method, initX []float64, stats *Stats, settings *Settings) (*Location, EvaluationType, error) {
	dim := len(initX)
	loc := &Location{
		X: make([]float64, dim),
	}
	copy(loc.X, initX)
	if method.Needs().Gradient {
		loc.Gradient = make([]float64, dim)
	}
	if method.Needs().Hessian {
		loc.Hessian = mat64.NewSymDense(dim, nil)
	}

	evalType := NoEvaluation
	if settings.UseInitialData {
		loc.F = settings.InitialFunctionValue
		if loc.Gradient != nil {
			initG := settings.InitialGradient
			if len(initG) != dim {
				panic("local: initial location size mismatch")
			}
			copy(loc.Gradient, initG)
		}
	} else {
		evalType = FuncEvaluation
		if loc.Gradient != nil {
			evalType = FuncGradEvaluation
		}
		evaluate(funcInfo, evalType, loc.X, loc, stats)
	}

	if math.IsNaN(loc.F) {
		return loc, evalType, ErrNaN
	}
	if math.IsInf(loc.F, 1) {
		return loc, evalType, ErrInf
	}
	for _, v := range loc.Gradient {
		if math.IsInf(v, 0) {
			return loc, evalType, ErrGradInf
		}
		if math.IsNaN(v) {
			return loc, evalType, ErrGradNaN
		}
	}

	return loc, evalType, nil
}

func checkConvergence(loc *Location, iterType IterationType, stats *Stats, settings *Settings) Status {
	if iterType == MajorIteration || iterType == InitIteration {
		if loc.Gradient != nil {
			norm := floats.Norm(loc.Gradient, math.Inf(1))
			if norm < settings.GradientAbsTol {
				return GradientAbsoluteConvergence
			}
		}
		if loc.F < settings.FunctionAbsTol {
			return FunctionAbsoluteConvergence
		}
	}

	// Check every step for negative infinity because it could break the
	// linesearches and -inf is the best you can do anyway.
	if math.IsInf(loc.F, -1) {
		return FunctionNegativeInfinity
	}

	if settings.FuncEvaluations > 0 {
		totalFun := stats.FuncEvaluations + stats.FuncGradEvaluations + stats.FuncGradHessEvaluations
		if totalFun >= settings.FuncEvaluations {
			return FunctionEvaluationLimit
		}
	}

	if settings.GradEvaluations > 0 {
		totalGrad := stats.GradEvaluations + stats.FuncGradEvaluations + stats.FuncGradHessEvaluations
		if totalGrad >= settings.GradEvaluations {
			return GradientEvaluationLimit
		}
	}

	if settings.HessEvaluations > 0 {
		totalHess := stats.HessEvaluations + stats.FuncGradHessEvaluations
		if totalHess >= settings.HessEvaluations {
			return HessianEvaluationLimit
		}
	}

	if settings.Runtime > 0 {
		// TODO(vladimir-ch): It would be nice to update Runtime here.
		if stats.Runtime >= settings.Runtime {
			return RuntimeLimit
		}
	}

	if iterType == MajorIteration && settings.MajorIterations > 0 {
		if stats.MajorIterations >= settings.MajorIterations {
			return IterationLimit
		}
	}
	return NotTerminated
}

func invalidate(loc *Location, f, grad, hess bool) {
	if f {
		loc.F = math.NaN()
	}
	if grad {
		for i := range loc.Gradient {
			loc.Gradient[i] = math.NaN()
		}
	}
	if hess && loc.Hessian != nil {
		n := loc.Hessian.Symmetric()
		for i := 0; i < n; i++ {
			for j := i; j < n; j++ {
				loc.Hessian.SetSym(i, j, math.NaN())
			}
		}
	}
}

// evaluate evaluates the function given by funcInfo at xNext, stores the
// answer into loc and updates stats. If loc.X is not equal to xNext, then
// unused fields of loc are set to NaN.
// evaluate panics if the function does not support the requested evalType.
func evaluate(funcInfo *functionInfo, evalType EvaluationType, xNext []float64, loc *Location, stats *Stats) {
	different := !floats.Equal(loc.X, xNext)
	if different {
		copy(loc.X, xNext)
	}
	switch evalType {
	case FuncEvaluation:
		if different {
			invalidate(loc, false, true, true)
		}
		loc.F = funcInfo.function.Func(loc.X)
		stats.FuncEvaluations++
		return
	case GradEvaluation:
		if funcInfo.IsGradient {
			if different {
				invalidate(loc, true, false, true)
			}
			funcInfo.gradient.Grad(loc.X, loc.Gradient)
			stats.GradEvaluations++
			return
		}
		if funcInfo.IsFunctionGradient {
			if different {
				invalidate(loc, false, false, true)
			}
			loc.F = funcInfo.functionGradient.FuncGrad(loc.X, loc.Gradient)
			stats.FuncGradEvaluations++
			return
		}
		if funcInfo.IsFunctionGradientHessian {
			if loc.Hessian == nil {
				loc.Hessian = mat64.NewSymDense(len(loc.X), nil)
			}
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
	case HessEvaluation:
		if funcInfo.IsHessian {
			if different {
				invalidate(loc, true, true, false)
			}
			funcInfo.hessian.Hess(loc.X, loc.Hessian)
			stats.HessEvaluations++
			return
		}
		if funcInfo.IsFunctionGradientHessian {
			if loc.Gradient == nil {
				loc.Gradient = make([]float64, len(loc.X))
			}
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
	case FuncGradEvaluation:
		if funcInfo.IsFunctionGradient {
			if different {
				invalidate(loc, false, false, true)
			}
			loc.F = funcInfo.functionGradient.FuncGrad(loc.X, loc.Gradient)
			stats.FuncGradEvaluations++
			return
		}
		if funcInfo.IsGradient {
			if different {
				invalidate(loc, false, false, true)
			}
			loc.F = funcInfo.function.Func(loc.X)
			stats.FuncEvaluations++
			funcInfo.gradient.Grad(loc.X, loc.Gradient)
			stats.GradEvaluations++
			return
		}
		if funcInfo.IsFunctionGradientHessian {
			if loc.Hessian == nil {
				loc.Hessian = mat64.NewSymDense(len(loc.X), nil)
			}
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
	case FuncGradHessEvaluation:
		if funcInfo.IsFunctionGradientHessian {
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
		if funcInfo.IsGradient && funcInfo.IsHessian {
			loc.F = funcInfo.function.Func(loc.X)
			stats.FuncEvaluations++
			funcInfo.gradient.Grad(loc.X, loc.Gradient)
			stats.GradEvaluations++
			funcInfo.hessian.Hess(loc.X, loc.Hessian)
			stats.HessEvaluations++
			return
		}
	case NoEvaluation:
		if different {
			// Optimizers should not request NoEvaluation at a new location.
			// The intent and therefore an appropriate action are both unclear.
			panic("no evaluation requested at new location")
		}
		return
	default:
		panic(fmt.Sprintf("unknown evaluation type %v", evalType))
	}
	panic(fmt.Sprintf("objective function does not support %v", evalType))
}

// update updates the stats given the new evaluation
func update(loc *Location, optLoc *Location, stats *Stats, iterType IterationType, startTime time.Time) {
	if iterType == MajorIteration {
		stats.MajorIterations++
	}
	if loc.F <= optLoc.F {
		copyLocation(optLoc, loc)
	}
	stats.Runtime = time.Since(startTime)
}
<commit_msg>Copy Hessian in copyLocation()<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package optimize

import (
	"fmt"
	"math"
	"time"

	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

// Local finds a local minimum of a function using a sequential algorithm.
// In order to maximize a function, multiply the output by -1.
//
// The first argument is of Function type representing the function to be minimized.
// Type switching is used to see if the function implements Gradient, FunctionGradient
// and Statuser.
//
// The second argument is the initial location at which to start the minimization.
// The initial location must be supplied, and must have a length equal to the
// problem dimension.
//
// The third argument contains the settings for the minimization. It is here that
// gradient tolerance, etc. are specified. The DefaultSettings() function
// can be called for a Settings struct with the default values initialized.
// If settings == nil, the default settings are used. Please see the documentation
// for the Settings structure for more information. The optimization Method used
// may also contain settings, see documentation for the appropriate optimizer.
//
// The final argument is the optimization method to use. If method == nil, then
// an appropriate default is chosen based on the properties of the other arguments
// (dimension, gradient-free or gradient-based, etc.). The optimization
// methods in this package are designed such that reasonable defaults occur
// if options are not specified explicitly. For example, the code
//  method := &Bfgs{}
// creates a pointer to a new Bfgs struct. When minimize is called, the settings
// in the method will be populated with default values. The methods are also
// designed such that they can be reused in future calls to method.
//
// Local returns a Result struct and any error that occurred. Please see the
// documentation of Result for more information.
//
// Please be aware that the default behavior of Local is to find the minimum.
// For certain functions and optimization methods, this process can take many
// function evaluations. If you would like to put limits on this, for example
// maximum runtime or maximum function evaluations, please modify the Settings
// input struct.
func Local(f Function, initX []float64, settings *Settings, method Method) (*Result, error) {
	if len(initX) == 0 {
		panic("local: initial X has zero length")
	}

	startTime := time.Now()
	funcInfo := newFunctionInfo(f)
	if method == nil {
		method = getDefaultMethod(funcInfo)
	}
	if err := funcInfo.satisfies(method); err != nil {
		return nil, err
	}

	if funcInfo.IsStatuser {
		_, err := funcInfo.statuser.Status()
		if err != nil {
			return nil, err
		}
	}

	if settings == nil {
		settings = DefaultSettings()
	}

	if settings.Recorder != nil {
		// Initialize Recorder first. If it fails, we avoid the (possibly
		// time-consuming) evaluation of F and DF at the starting location.
		err := settings.Recorder.Init(&funcInfo.FunctionInfo)
		if err != nil {
			return nil, err
		}
	}

	stats := &Stats{}
	optLoc, evalType, err := getStartingLocation(funcInfo, method, initX, stats, settings)
	if err != nil {
		return nil, err
	}

	// Runtime is the only Stats field that needs to be updated here.
	stats.Runtime = time.Since(startTime)
	// Send optLoc to Recorder before checking it for convergence.
	if settings.Recorder != nil {
		err = settings.Recorder.Record(optLoc, evalType, InitIteration, stats)
	}

	// Check if the starting location satisfies the convergence criteria.
	status := checkConvergence(optLoc, InitIteration, stats, settings)
	if status == NotTerminated && err == nil {
		// The starting location is not good enough, we need to perform a
		// minimization. The optimal location will be stored in-place in
		// optLoc.
		status, err = minimize(settings, method, funcInfo, stats, optLoc, startTime)
	}

	if settings.Recorder != nil && err == nil {
		// Send the optimal location to Recorder.
		err = settings.Recorder.Record(optLoc, NoEvaluation, PostIteration, stats)
	}
	stats.Runtime = time.Since(startTime)
	return &Result{
		Location: *optLoc,
		Stats:    *stats,
		Status:   status,
	}, err
}

func minimize(settings *Settings, method Method, funcInfo *functionInfo, stats *Stats, optLoc *Location, startTime time.Time) (status Status, err error) {
	loc := &Location{}
	copyLocation(loc, optLoc)
	xNext := make([]float64, len(loc.X))

	methodStatus, methodIsStatuser := method.(Statuser)

	evalType, iterType, err := method.Init(loc, &funcInfo.FunctionInfo, xNext)
	if err != nil {
		return Failure, err
	}

	for {
		if funcInfo.IsStatuser {
			// Check the function status before evaluating.
			status, err = funcInfo.statuser.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		// Perform evalType evaluation of the function at xNext and store the
		// result in location.
		evaluate(funcInfo, evalType, xNext, loc, stats)
		// Update the stats and optLoc.
		update(loc, optLoc, stats, iterType, startTime)
		// Get the convergence status before recording the new location.
		status = checkConvergence(optLoc, iterType, stats, settings)

		if settings.Recorder != nil {
			err = settings.Recorder.Record(loc, evalType, iterType, stats)
			if err != nil {
				if status == NotTerminated {
					status = Failure
				}
				return
			}
		}

		if status != NotTerminated {
			return
		}

		if methodIsStatuser {
			status, err = methodStatus.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		// Find the next location (stored in-place into xNext).
		evalType, iterType, err = method.Iterate(loc, xNext)
		if err != nil {
			status = Failure
			return
		}
	}
	panic("unreachable")
}

func copyLocation(dst, src *Location) {
	dst.X = resize(dst.X, len(src.X))
	copy(dst.X, src.X)

	dst.F = src.F

	dst.Gradient = resize(dst.Gradient, len(src.Gradient))
	copy(dst.Gradient, src.Gradient)

	if src.Hessian != nil {
		if dst.Hessian == nil || dst.Hessian.Symmetric() != len(src.X) {
			dst.Hessian = mat64.NewSymDense(len(src.X), nil)
		}
		dst.Hessian.CopySym(src.Hessian)
	}
}

func getDefaultMethod(funcInfo *functionInfo) Method {
	if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
		return &BFGS{}
	}
	// TODO: Implement a gradient-free method
	panic("optimize: gradient-free methods not yet coded")
}

// getStartingLocation allocates and initializes the starting location for the minimization.
func getStartingLocation(funcInfo *functionInfo, method Method, initX []float64, stats *Stats, settings *Settings) (*Location, EvaluationType, error) {
	dim := len(initX)
	loc := &Location{
		X: make([]float64, dim),
	}
	copy(loc.X, initX)
	if method.Needs().Gradient {
		loc.Gradient = make([]float64, dim)
	}
	if method.Needs().Hessian {
		loc.Hessian = mat64.NewSymDense(dim, nil)
	}

	evalType := NoEvaluation
	if settings.UseInitialData {
		loc.F = settings.InitialFunctionValue
		if loc.Gradient != nil {
			initG := settings.InitialGradient
			if len(initG) != dim {
				panic("local: initial location size mismatch")
			}
			copy(loc.Gradient, initG)
		}
	} else {
		evalType = FuncEvaluation
		if loc.Gradient != nil {
			evalType = FuncGradEvaluation
		}
		evaluate(funcInfo, evalType, loc.X, loc, stats)
	}

	if math.IsNaN(loc.F) {
		return loc, evalType, ErrNaN
	}
	if math.IsInf(loc.F, 1) {
		return loc, evalType, ErrInf
	}
	for _, v := range loc.Gradient {
		if math.IsInf(v, 0) {
			return loc, evalType, ErrGradInf
		}
		if math.IsNaN(v) {
			return loc, evalType, ErrGradNaN
		}
	}

	return loc, evalType, nil
}

func checkConvergence(loc *Location, iterType IterationType, stats *Stats, settings *Settings) Status {
	if iterType == MajorIteration || iterType == InitIteration {
		if loc.Gradient != nil {
			norm := floats.Norm(loc.Gradient, math.Inf(1))
			if norm < settings.GradientAbsTol {
				return GradientAbsoluteConvergence
			}
		}
		if loc.F < settings.FunctionAbsTol {
			return FunctionAbsoluteConvergence
		}
	}

	// Check every step for negative infinity because it could break the
	// linesearches and -inf is the best you can do anyway.
	if math.IsInf(loc.F, -1) {
		return FunctionNegativeInfinity
	}

	if settings.FuncEvaluations > 0 {
		totalFun := stats.FuncEvaluations + stats.FuncGradEvaluations + stats.FuncGradHessEvaluations
		if totalFun >= settings.FuncEvaluations {
			return FunctionEvaluationLimit
		}
	}

	if settings.GradEvaluations > 0 {
		totalGrad := stats.GradEvaluations + stats.FuncGradEvaluations + stats.FuncGradHessEvaluations
		if totalGrad >= settings.GradEvaluations {
			return GradientEvaluationLimit
		}
	}

	if settings.HessEvaluations > 0 {
		totalHess := stats.HessEvaluations + stats.FuncGradHessEvaluations
		if totalHess >= settings.HessEvaluations {
			return HessianEvaluationLimit
		}
	}

	if settings.Runtime > 0 {
		// TODO(vladimir-ch): It would be nice to update Runtime here.
		if stats.Runtime >= settings.Runtime {
			return RuntimeLimit
		}
	}

	if iterType == MajorIteration && settings.MajorIterations > 0 {
		if stats.MajorIterations >= settings.MajorIterations {
			return IterationLimit
		}
	}
	return NotTerminated
}

func invalidate(loc *Location, f, grad, hess bool) {
	if f {
		loc.F = math.NaN()
	}
	if grad {
		for i := range loc.Gradient {
			loc.Gradient[i] = math.NaN()
		}
	}
	if hess && loc.Hessian != nil {
		n := loc.Hessian.Symmetric()
		for i := 0; i < n; i++ {
			for j := i; j < n; j++ {
				loc.Hessian.SetSym(i, j, math.NaN())
			}
		}
	}
}

// evaluate evaluates the function given by funcInfo at xNext, stores the
// answer into loc and updates stats. If loc.X is not equal to xNext, then
// unused fields of loc are set to NaN.
// evaluate panics if the function does not support the requested evalType.
func evaluate(funcInfo *functionInfo, evalType EvaluationType, xNext []float64, loc *Location, stats *Stats) {
	different := !floats.Equal(loc.X, xNext)
	if different {
		copy(loc.X, xNext)
	}
	switch evalType {
	case FuncEvaluation:
		if different {
			invalidate(loc, false, true, true)
		}
		loc.F = funcInfo.function.Func(loc.X)
		stats.FuncEvaluations++
		return
	case GradEvaluation:
		if funcInfo.IsGradient {
			if different {
				invalidate(loc, true, false, true)
			}
			funcInfo.gradient.Grad(loc.X, loc.Gradient)
			stats.GradEvaluations++
			return
		}
		if funcInfo.IsFunctionGradient {
			if different {
				invalidate(loc, false, false, true)
			}
			loc.F = funcInfo.functionGradient.FuncGrad(loc.X, loc.Gradient)
			stats.FuncGradEvaluations++
			return
		}
		if funcInfo.IsFunctionGradientHessian {
			if loc.Hessian == nil {
				loc.Hessian = mat64.NewSymDense(len(loc.X), nil)
			}
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
	case HessEvaluation:
		if funcInfo.IsHessian {
			if different {
				invalidate(loc, true, true, false)
			}
			funcInfo.hessian.Hess(loc.X, loc.Hessian)
			stats.HessEvaluations++
			return
		}
		if funcInfo.IsFunctionGradientHessian {
			if loc.Gradient == nil {
				loc.Gradient = make([]float64, len(loc.X))
			}
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
	case FuncGradEvaluation:
		if funcInfo.IsFunctionGradient {
			if different {
				invalidate(loc, false, false, true)
			}
			loc.F = funcInfo.functionGradient.FuncGrad(loc.X, loc.Gradient)
			stats.FuncGradEvaluations++
			return
		}
		if funcInfo.IsGradient {
			if different {
				invalidate(loc, false, false, true)
			}
			loc.F = funcInfo.function.Func(loc.X)
			stats.FuncEvaluations++
			funcInfo.gradient.Grad(loc.X, loc.Gradient)
			stats.GradEvaluations++
			return
		}
		if funcInfo.IsFunctionGradientHessian {
			if loc.Hessian == nil {
				loc.Hessian = mat64.NewSymDense(len(loc.X), nil)
			}
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
	case FuncGradHessEvaluation:
		if funcInfo.IsFunctionGradientHessian {
			loc.F = funcInfo.functionGradientHessian.FuncGradHess(loc.X, loc.Gradient, loc.Hessian)
			stats.FuncGradHessEvaluations++
			return
		}
		if funcInfo.IsGradient && funcInfo.IsHessian {
			loc.F = funcInfo.function.Func(loc.X)
			stats.FuncEvaluations++
			funcInfo.gradient.Grad(loc.X, loc.Gradient)
			stats.GradEvaluations++
			funcInfo.hessian.Hess(loc.X, loc.Hessian)
			stats.HessEvaluations++
			return
		}
	case NoEvaluation:
		if different {
			// Optimizers should not request NoEvaluation at a new location.
			// The intent and therefore an appropriate action are both unclear.
			panic("no evaluation requested at new location")
		}
		return
	default:
		panic(fmt.Sprintf("unknown evaluation type %v", evalType))
	}
	panic(fmt.Sprintf("objective function does not support %v", evalType))
}

// update updates the stats given the new evaluation
func update(loc *Location, optLoc *Location, stats *Stats, iterType IterationType, startTime time.Time) {
	if iterType == MajorIteration {
		stats.MajorIterations++
	}
	if loc.F <= optLoc.F {
		copyLocation(optLoc, loc)
	}
	stats.Runtime = time.Since(startTime)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/kouhin/envflag"
)

var (
	previousRead     int
	previousReadTime time.Time
)

type MeterRead struct {
	Time    *time.Time
	Offset  int
	Length  int
	Message *MeterMessage
}

type MeterMessage struct {
	ID          int
	Type        int
	TamperPhy   int
	TamperEnc   int
	Consumption int
	ChecksumVal int
}

func main() {
	var (
		meterId = flag.String("meter-id", "REQUIRED", "ID of the meter to read from")
	)
	if err := envflag.Parse(); err != nil {
		panic(err)
	}

	if *meterId == "REQUIRED" {
		fmt.Println("meter-id is a required field")
		os.Exit(1)
	}

	cmdName := "rtlamr"
	cmdArgs := []string{"-format=json", "-filterid=" + *meterId}

	cmd := exec.Command(cmdName, cmdArgs...)
	cmdReader, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}

	scanner := bufio.NewScanner(cmdReader)
	go func() {
		for scanner.Scan() {
			text := scanner.Text()
			line := &MeterRead{}
			if err := json.Unmarshal([]byte(text), &line); err != nil {
				fmt.Fprintln(os.Stderr, fmt.Sprintf("Error unmarshaling line (%s):| %s", err, text))
				continue
			}
			fmt.Printf("received | %#v\n", line)
			processLine(line)
		}
	}()

	err = cmd.Start()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Error starting Cmd", err)
		os.Exit(1)
	}

	err = cmd.Wait()
	if err != nil {
		time.Sleep(2000) // let the scanner finish. there's gotta be a better way
		fmt.Fprintln(os.Stderr, "Command exited", err)
		os.Exit(1)
	}
}

func processLine(read *MeterRead) {
	if previousRead == 0 {
		// we have no baseline, set and wait for more
		previousRead = read.Message.Consumption
		previousReadTime = *read.Time
		return
	}

	timeBetweenReads := read.Time.Sub(previousReadTime)
	if timeBetweenReads.Hours() < 1.0 {
		// not long enough, bail out
		return
	}

	// it's been more than an hour, so we can report data. Compare the two readings, then factor in
	// that the time might have been more than an hour to get a cf/hr measurement.
	consumed := float64(read.Message.Consumption - previousRead)
	minsBetweenReads := timeBetweenReads.Minutes()
	consumedPerMin := consumed / minsBetweenReads
	consumedPerHour := consumedPerMin * 60
	reportMetric(consumedPerHour)
}

func reportMetric(cfHr float64) {
	fmt.Printf("At %s, consumption rate is %f cf/hr", time.Now().Format("2006-01-02T15:04:05.999999-07:00"), cfHr)
}
<commit_msg>We should reset that. Also make output a little nicer.<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/kouhin/envflag"
)

var (
	previousRead     int
	previousReadTime time.Time
)

type MeterRead struct {
	Time    *time.Time
	Offset  int
	Length  int
	Message *MeterMessage
}

type MeterMessage struct {
	ID          int
	Type        int
	TamperPhy   int
	TamperEnc   int
	Consumption int
	ChecksumVal int
}

func main() {
	var (
		meterId = flag.String("meter-id", "REQUIRED", "ID of the meter to read from")
	)
	if err := envflag.Parse(); err != nil {
		panic(err)
	}

	if *meterId == "REQUIRED" {
		fmt.Println("meter-id is a required field")
		os.Exit(1)
	}

	cmdName := "rtlamr"
	cmdArgs := []string{"-format=json", "-filterid=" + *meterId}

	cmd := exec.Command(cmdName, cmdArgs...)
	cmdReader, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}

	scanner := bufio.NewScanner(cmdReader)
	go func() {
		for scanner.Scan() {
			text := scanner.Text()
			line := &MeterRead{}
			if err := json.Unmarshal([]byte(text), &line); err != nil {
				fmt.Fprintln(os.Stderr, fmt.Sprintf("Error unmarshaling line (%s):| %s", err, text))
				continue
			}
			fmt.Printf("received | %q\n", line)
			processLine(line)
		}
	}()

	err = cmd.Start()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Error starting Cmd", err)
		os.Exit(1)
	}

	err = cmd.Wait()
	if err != nil {
		time.Sleep(2000) // let the scanner finish. there's gotta be a better way
		fmt.Fprintln(os.Stderr, "Command exited", err)
		os.Exit(1)
	}
}

func processLine(read *MeterRead) {
	if previousRead == 0 {
		// we have no baseline, set and wait for more
		previousRead = read.Message.Consumption
		previousReadTime = *read.Time
		return
	}

	timeBetweenReads := read.Time.Sub(previousReadTime)
	if timeBetweenReads.Hours() < 1.0 {
		// not long enough, bail out
		return
	}

	// it's been more than an hour, so we can report data. Compare the two readings, then factor in
	// that the time might have been more than an hour to get a cf/hr measurement.
	consumed := float64(read.Message.Consumption - previousRead)
	minsBetweenReads := timeBetweenReads.Minutes()
	consumedPerMin := consumed / minsBetweenReads
	consumedPerHour := consumedPerMin * 60
	reportMetric(consumedPerHour)
	// and start again.
	previousRead = read.Message.Consumption
	previousReadTime = *read.Time
}

func reportMetric(cfHr float64) {
	fmt.Printf("At %s, consumption rate is %f cf/hr\n", time.Now().Format("2006-01-02T15:04:05.999999-07:00"), cfHr)
}
<|endoftext|>


<commit_before>/*
Package prometheus provides middleware to add Prometheus metrics.

Example:
```
package main
import (
    "github.com/labstack/echo/v4"
    "github.com/labstack/echo-contrib/prometheus"
)
func main() {
    e := echo.New()
    // Enable metrics middleware
    p := prometheus.NewPrometheus("echo", nil)
    p.Use(e)

    e.Logger.Fatal(e.Start(":1323"))
}
```
*/
package prometheus

import (
	"bytes"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/labstack/gommon/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/expfmt"
)

var defaultMetricPath = "/metrics"
var defaultSubsystem = "echo"

// Standard default metrics
//	counter, counter_vec, gauge, gauge_vec,
//	histogram, histogram_vec, summary, summary_vec
var reqCnt = &Metric{
	ID:          "reqCnt",
	Name:        "requests_total",
	Description: "How many HTTP requests processed, partitioned by status code and HTTP method.",
	Type:        "counter_vec",
	Args:        []string{"code", "method", "host", "url"}}

var reqDur = &Metric{
	ID:          "reqDur",
	Name:        "request_duration_seconds",
	Description: "The HTTP request latencies in seconds.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var resSz = &Metric{
	ID:          "resSz",
	Name:        "response_size_bytes",
	Description: "The HTTP response sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var reqSz = &Metric{
	ID:          "reqSz",
	Name:        "request_size_bytes",
	Description: "The HTTP request sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var standardMetrics = []*Metric{
	reqCnt,
	reqDur,
	resSz,
	reqSz,
}

/*
RequestCounterURLLabelMappingFunc is a function which can be supplied to the middleware to control
the cardinality of the request counter's "url" label, which might be required in some contexts.
For instance, if for a "/customer/:name" route you don't want to generate a time series for every
possible customer name, you could use this function:

func(c echo.Context) string {
	url := c.Request.URL.Path
	for _, p := range c.Params {
		if p.Key == "name" {
			url = strings.Replace(url, p.Value, ":name", 1)
			break
		}
	}
	return url
}

which would map "/customer/alice" and "/customer/bob" to their template "/customer/:name".
*/
type RequestCounterURLLabelMappingFunc func(c echo.Context) string

// Metric is a definition for the name, description, type, ID, and
// prometheus.Collector type (i.e. CounterVec, Summary, etc) of each metric
type Metric struct {
	MetricCollector prometheus.Collector
	ID              string
	Name            string
	Description     string
	Type            string
	Args            []string
}

// Prometheus contains the metrics gathered by the instance and its path
type Prometheus struct {
	reqCnt               *prometheus.CounterVec
	reqDur, reqSz, resSz *prometheus.HistogramVec
	router               *echo.Echo
	listenAddress        string
	Ppg                  PushGateway

	MetricsList []*Metric
	MetricsPath string
	Subsystem   string
	Skipper     middleware.Skipper

	RequestCounterURLLabelMappingFunc RequestCounterURLLabelMappingFunc

	// Context string to use as a prometheus URL label
	URLLabelFromContext string
}

// PushGateway contains the configuration for pushing to a Prometheus pushgateway (optional)
type PushGateway struct {
	// Push interval in seconds
	PushIntervalSeconds time.Duration

	// Push Gateway URL in format http://domain:port
	// where JOBNAME can be any string of your choice
	PushGatewayURL string

	// pushgateway job name, defaults to "echo"
	Job string
}

// NewPrometheus generates a new set of metrics with a certain subsystem name
func NewPrometheus(subsystem string, skipper middleware.Skipper, customMetricsList ...[]*Metric) *Prometheus {
	var metricsList []*Metric
	if skipper == nil {
		skipper = middleware.DefaultSkipper
	}

	if len(customMetricsList) > 1 {
		panic("Too many args. NewPrometheus( string, <optional []*Metric> ).")
	} else if len(customMetricsList) == 1 {
		metricsList = customMetricsList[0]
	}

	for _, metric := range standardMetrics {
		metricsList = append(metricsList, metric)
	}

	p := &Prometheus{
		MetricsList: metricsList,
		MetricsPath: defaultMetricPath,
		Subsystem:   defaultSubsystem,
		Skipper:     skipper,
		RequestCounterURLLabelMappingFunc: func(c echo.Context) string {
			return c.Path() // i.e. by default do nothing, i.e. return URL as is
		},
	}

	p.registerMetrics(subsystem)

	return p
}

// SetPushGateway sends metrics to a remote pushgateway exposed on pushGatewayURL
// every pushIntervalSeconds. Metrics are fetched from
func (p *Prometheus) SetPushGateway(pushGatewayURL string, pushIntervalSeconds time.Duration) {
	p.Ppg.PushGatewayURL = pushGatewayURL
	p.Ppg.PushIntervalSeconds = pushIntervalSeconds
	p.startPushTicker()
}

// SetPushGatewayJob job name, defaults to "echo"
func (p *Prometheus) SetPushGatewayJob(j string) {
	p.Ppg.Job = j
}

// SetListenAddress for exposing metrics on address. If not set, it will be exposed at the
// same address of the echo engine that is being used
// func (p *Prometheus) SetListenAddress(address string) {
// 	p.listenAddress = address
// 	if p.listenAddress != "" {
// 		p.router = echo.Echo().Router()
// 	}
// }

// SetListenAddressWithRouter for using a separate router to expose metrics. (this keeps things like GET /metrics out of
// your content's access log).
// func (p *Prometheus) SetListenAddressWithRouter(listenAddress string, r *echo.Echo) {
// 	p.listenAddress = listenAddress
// 	if len(p.listenAddress) > 0 {
// 		p.router = r
// 	}
// }

// SetMetricsPath set metrics paths
func (p *Prometheus) SetMetricsPath(e *echo.Echo) {
	if p.listenAddress != "" {
		p.router.GET(p.MetricsPath, prometheusHandler())
		p.runServer()
	} else {
		e.GET(p.MetricsPath, prometheusHandler())
	}
}

func (p *Prometheus) runServer() {
	if p.listenAddress != "" {
		go p.router.Start(p.listenAddress)
	}
}

func (p *Prometheus) getMetrics() []byte {
	out := &bytes.Buffer{}
	metricFamilies, _ := prometheus.DefaultGatherer.Gather()
	for i := range metricFamilies {
		expfmt.MetricFamilyToText(out, metricFamilies[i])

	}
	return out.Bytes()
}

func (p *Prometheus) getPushGatewayURL() string {
	h, _ := os.Hostname()
	if p.Ppg.Job == "" {
		p.Ppg.Job = "echo"
	}
	return p.Ppg.PushGatewayURL + "/metrics/job/" + p.Ppg.Job + "/instance/" + h
}

func (p *Prometheus) sendMetricsToPushGateway(metrics []byte) {
	req, err := http.NewRequest("POST", p.getPushGatewayURL(), bytes.NewBuffer(metrics))
	client := &http.Client{}
	if _, err = client.Do(req); err != nil {
		log.Errorf("Error sending to push gateway: %v", err)
	}
}

func (p *Prometheus) startPushTicker() {
	ticker := time.NewTicker(time.Second * p.Ppg.PushIntervalSeconds)
	go func() {
		for range ticker.C {
			p.sendMetricsToPushGateway(p.getMetrics())
		}
	}()
}

// NewMetric associates prometheus.Collector based on Metric.Type
func NewMetric(m *Metric, subsystem string) prometheus.Collector {
	var metric prometheus.Collector
	switch m.Type {
	case "counter_vec":
		metric = prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "counter":
		metric = prometheus.NewCounter(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "gauge_vec":
		metric = prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "gauge":
		metric = prometheus.NewGauge(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "histogram_vec":
		metric = prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "histogram":
		metric = prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "summary_vec":
		metric = prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "summary":
		metric = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	}
	return metric
}

func (p *Prometheus) registerMetrics(subsystem string) {

	for _, metricDef := range p.MetricsList {
		metric := NewMetric(metricDef, subsystem)
		if err := prometheus.Register(metric); err != nil {
			log.Errorf("%s could not be registered in Prometheus: %v", metricDef.Name, err)
		}
		switch metricDef {
		case reqCnt:
			p.reqCnt = metric.(*prometheus.CounterVec)
		case reqDur:
			p.reqDur = metric.(*prometheus.HistogramVec)
		case resSz:
			p.resSz = metric.(*prometheus.HistogramVec)
		case reqSz:
			p.reqSz = metric.(*prometheus.HistogramVec)
		}
		metricDef.MetricCollector = metric
	}
}

// Use adds the middleware to the Echo engine.
func (p *Prometheus) Use(e *echo.Echo) {
	e.Use(p.HandlerFunc)
	p.SetMetricsPath(e)
}

// HandlerFunc defines handler function for middleware
func (p *Prometheus) HandlerFunc(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) (err error) {
		if c.Path() == p.MetricsPath {
			return next(c)
		}
		if p.Skipper(c) {
			return next(c)
		}

		start := time.Now()
		reqSz := computeApproximateRequestSize(c.Request())

		if err = next(c); err != nil {
			c.Error(err)
		}

		status := strconv.Itoa(c.Response().Status)
		url := p.RequestCounterURLLabelMappingFunc(c)

		elapsed := float64(time.Since(start)) / float64(time.Second)
		resSz := float64(c.Response().Size)

		p.reqDur.WithLabelValues(status, c.Request().Method, url).Observe(elapsed)

		if len(p.URLLabelFromContext) > 0 {
			u := c.Get(p.URLLabelFromContext)
			if u == nil {
				u = "unknown"
			}
			url = u.(string)
		}

		p.reqCnt.WithLabelValues(status, c.Request().Method, c.Request().Host, url).Inc()
		p.reqSz.WithLabelValues(status, c.Request().Method, url).Observe(float64(reqSz))
		p.resSz.WithLabelValues(status, c.Request().Method, url).Observe(resSz)

		return
	}
}

func prometheusHandler() echo.HandlerFunc {
	h := promhttp.Handler()
	return func(c echo.Context) error {
		h.ServeHTTP(c.Response(), c.Request())
		return nil
	}
}

func computeApproximateRequestSize(r *http.Request) int {
	s := 0
	if r.URL != nil {
		s = len(r.URL.Path)
	}

	s += len(r.Method)
	s += len(r.Proto)
	for name, values := range r.Header {
		s += len(name)
		for _, value := range values {
			s += len(value)
		}
	}
	s += len(r.Host)

	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

	if r.ContentLength != -1 {
		s += int(r.ContentLength)
	}
	return s
}
<commit_msg>Prometheus add host relabel (#53)<commit_after>/*
Package prometheus provides middleware to add Prometheus metrics.

Example:
```
package main
import (
    "github.com/labstack/echo/v4"
    "github.com/labstack/echo-contrib/prometheus"
)
func main() {
    e := echo.New()
    // Enable metrics middleware
    p := prometheus.NewPrometheus("echo", nil)
    p.Use(e)

    e.Logger.Fatal(e.Start(":1323"))
}
```
*/
package prometheus

import (
	"bytes"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/labstack/gommon/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/expfmt"
)

var defaultMetricPath = "/metrics"
var defaultSubsystem = "echo"

// Standard default metrics
//	counter, counter_vec, gauge, gauge_vec,
//	histogram, histogram_vec, summary, summary_vec
var reqCnt = &Metric{
	ID:          "reqCnt",
	Name:        "requests_total",
	Description: "How many HTTP requests processed, partitioned by status code and HTTP method.",
	Type:        "counter_vec",
	Args:        []string{"code", "method", "host", "url"}}

var reqDur = &Metric{
	ID:          "reqDur",
	Name:        "request_duration_seconds",
	Description: "The HTTP request latencies in seconds.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var resSz = &Metric{
	ID:          "resSz",
	Name:        "response_size_bytes",
	Description: "The HTTP response sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var reqSz = &Metric{
	ID:          "reqSz",
	Name:        "request_size_bytes",
	Description: "The HTTP request sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var standardMetrics = []*Metric{
	reqCnt,
	reqDur,
	resSz,
	reqSz,
}

/*
RequestCounterLabelMappingFunc is a function which can be supplied to the middleware to control
the cardinality of the request counter's "url" label, which might be required in some contexts.
For instance, if for a "/customer/:name" route you don't want to generate a time series for every
possible customer name, you could use this function:

func(c echo.Context) string {
	url := c.Request.URL.Path
	for _, p := range c.Params {
		if p.Key == "name" {
			url = strings.Replace(url, p.Value, ":name", 1)
			break
		}
	}
	return url
}

which would map "/customer/alice" and "/customer/bob" to their template "/customer/:name".
It can also be applied for the "Host" label
*/
type RequestCounterLabelMappingFunc func(c echo.Context) string

// Metric is a definition for the name, description, type, ID, and
// prometheus.Collector type (i.e. CounterVec, Summary, etc) of each metric
type Metric struct {
	MetricCollector prometheus.Collector
	ID              string
	Name            string
	Description     string
	Type            string
	Args            []string
}

// Prometheus contains the metrics gathered by the instance and its path
type Prometheus struct {
	reqCnt               *prometheus.CounterVec
	reqDur, reqSz, resSz *prometheus.HistogramVec
	router               *echo.Echo
	listenAddress        string
	Ppg                  PushGateway

	MetricsList []*Metric
	MetricsPath string
	Subsystem   string
	Skipper     middleware.Skipper

	RequestCounterURLLabelMappingFunc  RequestCounterLabelMappingFunc
	RequestCounterHostLabelMappingFunc RequestCounterLabelMappingFunc

	// Context string to use as a prometheus URL label
	URLLabelFromContext string
}

// PushGateway contains the configuration for pushing to a Prometheus pushgateway (optional)
type PushGateway struct {
	// Push interval in seconds
	PushIntervalSeconds time.Duration

	// Push Gateway URL in format http://domain:port
	// where JOBNAME can be any string of your choice
	PushGatewayURL string

	// pushgateway job name, defaults to "echo"
	Job string
}

// NewPrometheus generates a new set of metrics with a certain subsystem name
func NewPrometheus(subsystem string, skipper middleware.Skipper, customMetricsList ...[]*Metric) *Prometheus {
	var metricsList []*Metric
	if skipper == nil {
		skipper = middleware.DefaultSkipper
	}

	if len(customMetricsList) > 1 {
		panic("Too many args. NewPrometheus( string, <optional []*Metric> ).")
	} else if len(customMetricsList) == 1 {
		metricsList = customMetricsList[0]
	}

	for _, metric := range standardMetrics {
		metricsList = append(metricsList, metric)
	}

	p := &Prometheus{
		MetricsList: metricsList,
		MetricsPath: defaultMetricPath,
		Subsystem:   defaultSubsystem,
		Skipper:     skipper,
		RequestCounterURLLabelMappingFunc: func(c echo.Context) string {
			return c.Path() // i.e. by default do nothing, i.e. return URL as is
		},
		RequestCounterHostLabelMappingFunc: func(c echo.Context) string {
			return c.Request().Host
		},
	}

	p.registerMetrics(subsystem)

	return p
}

// SetPushGateway sends metrics to a remote pushgateway exposed on pushGatewayURL
// every pushIntervalSeconds. Metrics are fetched from
func (p *Prometheus) SetPushGateway(pushGatewayURL string, pushIntervalSeconds time.Duration) {
	p.Ppg.PushGatewayURL = pushGatewayURL
	p.Ppg.PushIntervalSeconds = pushIntervalSeconds
	p.startPushTicker()
}

// SetPushGatewayJob job name, defaults to "echo"
func (p *Prometheus) SetPushGatewayJob(j string) {
	p.Ppg.Job = j
}

// SetListenAddress for exposing metrics on address. If not set, it will be exposed at the
// same address of the echo engine that is being used
// func (p *Prometheus) SetListenAddress(address string) {
// 	p.listenAddress = address
// 	if p.listenAddress != "" {
// 		p.router = echo.Echo().Router()
// 	}
// }

// SetListenAddressWithRouter for using a separate router to expose metrics. (this keeps things like GET /metrics out of
// your content's access log).
// func (p *Prometheus) SetListenAddressWithRouter(listenAddress string, r *echo.Echo) {
// 	p.listenAddress = listenAddress
// 	if len(p.listenAddress) > 0 {
// 		p.router = r
// 	}
// }

// SetMetricsPath set metrics paths
func (p *Prometheus) SetMetricsPath(e *echo.Echo) {
	if p.listenAddress != "" {
		p.router.GET(p.MetricsPath, prometheusHandler())
		p.runServer()
	} else {
		e.GET(p.MetricsPath, prometheusHandler())
	}
}

func (p *Prometheus) runServer() {
	if p.listenAddress != "" {
		go p.router.Start(p.listenAddress)
	}
}

func (p *Prometheus) getMetrics() []byte {
	out := &bytes.Buffer{}
	metricFamilies, _ := prometheus.DefaultGatherer.Gather()
	for i := range metricFamilies {
		expfmt.MetricFamilyToText(out, metricFamilies[i])

	}
	return out.Bytes()
}

func (p *Prometheus) getPushGatewayURL() string {
	h, _ := os.Hostname()
	if p.Ppg.Job == "" {
		p.Ppg.Job = "echo"
	}
	return p.Ppg.PushGatewayURL + "/metrics/job/" + p.Ppg.Job + "/instance/" + h
}

func (p *Prometheus) sendMetricsToPushGateway(metrics []byte) {
	req, err := http.NewRequest("POST", p.getPushGatewayURL(), bytes.NewBuffer(metrics))
	client := &http.Client{}
	if _, err = client.Do(req); err != nil {
		log.Errorf("Error sending to push gateway: %v", err)
	}
}

func (p *Prometheus) startPushTicker() {
	ticker := time.NewTicker(time.Second * p.Ppg.PushIntervalSeconds)
	go func() {
		for range ticker.C {
			p.sendMetricsToPushGateway(p.getMetrics())
		}
	}()
}

// NewMetric associates prometheus.Collector based on Metric.Type
func NewMetric(m *Metric, subsystem string) prometheus.Collector {
	var metric prometheus.Collector
	switch m.Type {
	case "counter_vec":
		metric = prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "counter":
		metric = prometheus.NewCounter(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "gauge_vec":
		metric = prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "gauge":
		metric = prometheus.NewGauge(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "histogram_vec":
		metric = prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "histogram":
		metric = prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "summary_vec":
		metric = prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "summary":
		metric = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	}
	return metric
}

func (p *Prometheus) registerMetrics(subsystem string) {

	for _, metricDef := range p.MetricsList {
		metric := NewMetric(metricDef, subsystem)
		if err := prometheus.Register(metric); err != nil {
			log.Errorf("%s could not be registered in Prometheus: %v", metricDef.Name, err)
		}
		switch metricDef {
		case reqCnt:
			p.reqCnt = metric.(*prometheus.CounterVec)
		case reqDur:
			p.reqDur = metric.(*prometheus.HistogramVec)
		case resSz:
			p.resSz = metric.(*prometheus.HistogramVec)
		case reqSz:
			p.reqSz = metric.(*prometheus.HistogramVec)
		}
		metricDef.MetricCollector = metric
	}
}

// Use adds the middleware to the Echo engine.
func (p *Prometheus) Use(e *echo.Echo) {
	e.Use(p.HandlerFunc)
	p.SetMetricsPath(e)
}

// HandlerFunc defines handler function for middleware
func (p *Prometheus) HandlerFunc(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) (err error) {
		if c.Path() == p.MetricsPath {
			return next(c)
		}
		if p.Skipper(c) {
			return next(c)
		}

		start := time.Now()
		reqSz := computeApproximateRequestSize(c.Request())

		if err = next(c); err != nil {
			c.Error(err)
		}

		status := strconv.Itoa(c.Response().Status)
		url := p.RequestCounterURLLabelMappingFunc(c)

		elapsed := float64(time.Since(start)) / float64(time.Second)
		resSz := float64(c.Response().Size)

		p.reqDur.WithLabelValues(status, c.Request().Method, url).Observe(elapsed)

		if len(p.URLLabelFromContext) > 0 {
			u := c.Get(p.URLLabelFromContext)
			if u == nil {
				u = "unknown"
			}
			url = u.(string)
		}

		p.reqCnt.WithLabelValues(status, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Inc()
		p.reqSz.WithLabelValues(status, c.Request().Method, url).Observe(float64(reqSz))
		p.resSz.WithLabelValues(status, c.Request().Method, url).Observe(resSz)

		return
	}
}

func prometheusHandler() echo.HandlerFunc {
	h := promhttp.Handler()
	return func(c echo.Context) error {
		h.ServeHTTP(c.Response(), c.Request())
		return nil
	}
}

func computeApproximateRequestSize(r *http.Request) int {
	s := 0
	if r.URL != nil {
		s = len(r.URL.Path)
	}

	s += len(r.Method)
	s += len(r.Proto)
	for name, values := range r.Header {
		s += len(name)
		for _, value := range values {
			s += len(value)
		}
	}
	s += len(r.Host)

	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

	if r.ContentLength != -1 {
		s += int(r.ContentLength)
	}
	return s
}
<|endoftext|>


<commit_before>// +build go1.9

package prometheus

import (
	"fmt"
	"log"
	"math"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
)

var (
	// DefaultPrometheusOpts is the default set of options used when creating a
	// PrometheusSink.
	DefaultPrometheusOpts = PrometheusOpts{
		Expiration: 60 * time.Second,
	}
)

// PrometheusOpts is used to configure the Prometheus Sink
type PrometheusOpts struct {
	// Expiration is the duration a metric is valid for, after which it will be
	// untracked. If the value is zero, a metric is never expired.
	Expiration time.Duration
	Registerer prometheus.Registerer

	// Gauges, Summaries, and Counters allow us to pre-declare metrics by giving their Name, Help, and ConstLabels to
	// the PrometheusSink when it is created. Metrics declared in this way will be initialized at zero and will not be
	// deleted when their expiry is reached.
	// - Gauges and Summaries will be set to NaN when they expire.
	// - Counters continue to Collect their last known value.
	// Ex:
	// PrometheusOpts{
	//     Expiration: 10 * time.Second,
	//     Gauges: []GaugeDefinition{
	//         {
	//	         Name: []string{ "application", "component", "measurement"},
	//           Help: "application_component_measurement provides an example of how to declare static metrics",
	//           ConstLabels: []metrics.Label{ { Name: "my_label", Value: "does_not_change" }, },
	//         },
	//     },
	// }
	GaugeDefinitions   []GaugeDefinition
	SummaryDefinitions []SummaryDefinition
	CounterDefinitions []CounterDefinition
}

type PrometheusSink struct {
	// If these will ever be copied, they should be converted to *sync.Map values and initialized appropriately
	gauges     sync.Map
	summaries  sync.Map
	counters   sync.Map
	expiration time.Duration
}

// GaugeDefinition can be provided to PrometheusOpts to declare a constant gauge that is not deleted on expiry.
type GaugeDefinition struct {
	Name []string
	ConstLabels []metrics.Label
	Help string
}

type gauge struct {
	prometheus.Gauge
	updatedAt time.Time
	// canDelete is set if the metric is created during runtime so we know it's ephemeral and can delete it on expiry.
	canDelete bool
}

// SummaryDefinition can be provided to PrometheusOpts to declare a constant summary that is not deleted on expiry.
type SummaryDefinition struct {
	Name []string
	ConstLabels []metrics.Label
	Help string
}

type summary struct {
	prometheus.Summary
	updatedAt time.Time
	canDelete bool
}

// CounterDefinition can be provided to PrometheusOpts to declare a constant counter that is not deleted on expiry.
type CounterDefinition struct {
	Name []string
	ConstLabels []metrics.Label
	Help string
}

type counter struct {
	prometheus.Counter
	updatedAt time.Time
	canDelete bool
}

// NewPrometheusSink creates a new PrometheusSink using the default options.
func NewPrometheusSink() (*PrometheusSink, error) {
	return NewPrometheusSinkFrom(DefaultPrometheusOpts)
}

// NewPrometheusSinkFrom creates a new PrometheusSink using the passed options.
func NewPrometheusSinkFrom(opts PrometheusOpts) (*PrometheusSink, error) {
	sink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: opts.Expiration,
	}

	initGauges(&sink.gauges, opts.GaugeDefinitions)
	initSummaries(&sink.summaries, opts.SummaryDefinitions)
	initCounters(&sink.counters, opts.CounterDefinitions)

	reg := opts.Registerer
	if reg == nil {
		reg = prometheus.DefaultRegisterer
	}

	return sink, reg.Register(sink)
}

// Describe is needed to meet the Collector interface.
func (p *PrometheusSink) Describe(c chan<- *prometheus.Desc) {
	// We must emit some description otherwise an error is returned. This
	// description isn't shown to the user!
	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(c)
}

// Collect meets the collection interface and allows us to enforce our expiration
// logic to clean up ephemeral metrics if their value haven't been set for a
// duration exceeding our allowed expiration time.
func (p *PrometheusSink) Collect(c chan<- prometheus.Metric) {
	expire := p.expiration != 0
	now := time.Now()
	p.gauges.Range(func(k, v interface{}) bool {
		if v == nil {
			return true
		}
		g := v.(*gauge)
		lastUpdate := g.updatedAt
		if expire && lastUpdate.Add(p.expiration).Before(now) {
			if g.canDelete {
				p.gauges.Delete(k)
				return true
			}
			// We have not observed the gauge this interval so we don't know its value.
			g.Set(math.NaN())
		}
		g.Collect(c)
		return true
	})
	p.summaries.Range(func(k, v interface{}) bool {
		if v == nil {
			return true
		}
		s := v.(*summary)
		lastUpdate := s.updatedAt
		if expire && lastUpdate.Add(p.expiration).Before(now) {
			if s.canDelete {
				p.summaries.Delete(k)
				return true
			}
			// We have observed nothing in this interval.
			s.Observe(math.NaN())
		}
		s.Collect(c)
		return true
	})
	p.counters.Range(func(k, v interface{}) bool {
		if v == nil {
			return true
		}
		count := v.(*counter)
		lastUpdate := count.updatedAt
		if expire && lastUpdate.Add(p.expiration).Before(now) {
			if count.canDelete {
				p.counters.Delete(k)
				return true
			}
			// Counters remain at their previous value when not observed, so we do not set it to NaN.
		}
		count.Collect(c)
		return true
	})
}

func initGauges(m *sync.Map, gauges []GaugeDefinition) {
	for _, g := range gauges {
		key, hash := flattenKey(g.Name, g.ConstLabels)
		pG := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        g.Help,
			ConstLabels: prometheusLabels(g.ConstLabels),
		})
		pG.Set(float64(0)) // Initialize at zero
		m.Store(hash, &gauge{ Gauge: pG })
	}
	return
}

func initSummaries(m *sync.Map, summaries []SummaryDefinition) {
	for _, s := range summaries {
		key, hash := flattenKey(s.Name, s.ConstLabels)
		pS := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        s.Help,
			ConstLabels: prometheusLabels(s.ConstLabels),
		})
		pS.Observe(float64(0)) // Initialize at zero
		m.Store(hash, &summary{ Summary: pS })
	}
	return
}

func initCounters(m *sync.Map, counters []CounterDefinition) {
	for _, c := range counters {
		key, hash := flattenKey(c.Name, c.ConstLabels)
		pC := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        c.Help,
			ConstLabels: prometheusLabels(c.ConstLabels),
		})
		pC.Add(float64(0)) // Initialize at zero
		m.Store(hash, &counter{ Counter: pC })
	}
	return
}

var forbiddenChars = regexp.MustCompile("[ .=\\-/]")

func flattenKey(parts []string, labels []metrics.Label) (string, string) {
	key := strings.Join(parts, "_")
	key = forbiddenChars.ReplaceAllString(key, "_")

	hash := key
	for _, label := range labels {
		hash += fmt.Sprintf(";%s=%s", label.Name, label.Value)
	}

	return key, hash
}

func prometheusLabels(labels []metrics.Label) prometheus.Labels {
	l := make(prometheus.Labels)
	for _, label := range labels {
		l[label.Name] = label.Value
	}
	return l
}

func (p *PrometheusSink) SetGauge(parts []string, val float32) {
	p.SetGaugeWithLabels(parts, val, nil)
}

func (p *PrometheusSink) SetGaugeWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	pg, ok := p.gauges.Load(hash)

	// The sync.Map underlying gauges stores pointers to our structs. If we need to make updates,
	// rather than modifying the underlying value directly, which would be racy, we make a local
	// copy by dereferencing the pointer we get back, making the appropriate changes, and then
	// storing a pointer to our local copy. The underlying Prometheus types are threadsafe,
	// so there's no issues there. It's possible for racy updates to occur to the updatedAt
	// value, but since we're always setting it to time.Now(), it doesn't really matter.
	if ok {
		localGauge := *pg.(*gauge)
		localGauge.Set(float64(val))
		localGauge.updatedAt = time.Now()
		p.gauges.Store(hash, &localGauge)

	// The gauge does not exist, create the gauge and allow it to be deleted
	} else {
		g := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		g.Set(float64(val))
		pg = &gauge{
			Gauge: g,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.gauges.Store(hash, pg)
	}
}

func (p *PrometheusSink) AddSample(parts []string, val float32) {
	p.AddSampleWithLabels(parts, val, nil)
}

func (p *PrometheusSink) AddSampleWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	ps, ok := p.summaries.Load(hash)

	// Does the summary already exist for this sample type?
	if ok {
		localSummary := *ps.(*summary)
		localSummary.Observe(float64(val))
		localSummary.updatedAt = time.Now()
		p.summaries.Store(hash, &localSummary)

	// The summary does not exist, create the Summary and allow it to be deleted
	} else {
		s := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        key,
			MaxAge:      10 * time.Second,
			ConstLabels: prometheusLabels(labels),
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		})
		s.Observe(float64(val))
		ps = &summary{
			Summary: s,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.summaries.Store(hash, ps)
	}
}

// EmitKey is not implemented. Prometheus doesn’t offer a type for which an
// arbitrary number of values is retained, as Prometheus works with a pull
// model, rather than a push model.
func (p *PrometheusSink) EmitKey(key []string, val float32) {
}

func (p *PrometheusSink) IncrCounter(parts []string, val float32) {
	p.IncrCounterWithLabels(parts, val, nil)
}

func (p *PrometheusSink) IncrCounterWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	pc, ok := p.counters.Load(hash)

	// Does the counter exist?
	if ok {
		localCounter := *pc.(*counter)
		localCounter.Add(float64(val))
		localCounter.updatedAt = time.Now()
		p.counters.Store(hash, &localCounter)

	// The counter does not exist yet, create it and allow it to be deleted
	} else {
		c := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		c.Add(float64(val))
		pc = &counter{
			Counter: c,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.counters.Store(hash, pc)
	}
}

// PrometheusPushSink wraps a normal prometheus sink and provides an address and facilities to export it to an address
// on an interval.
type PrometheusPushSink struct {
	*PrometheusSink
	pusher       *push.Pusher
	address      string
	pushInterval time.Duration
	stopChan     chan struct{}
}

// NewPrometheusPushSink creates a PrometheusPushSink by taking an address, interval, and destination name.
func NewPrometheusPushSink(address string, pushInterval time.Duration, name string) (*PrometheusPushSink, error) {
	promSink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: 60 * time.Second,
	}

	pusher := push.New(address, name).Collector(promSink)

	sink := &PrometheusPushSink{
		promSink,
		pusher,
		address,
		pushInterval,
		make(chan struct{}),
	}

	sink.flushMetrics()
	return sink, nil
}

func (s *PrometheusPushSink) flushMetrics() {
	ticker := time.NewTicker(s.pushInterval)

	go func() {
		for {
			select {
			case <-ticker.C:
				err := s.pusher.Push()
				if err != nil {
					log.Printf("[ERR] Error pushing to Prometheus! Err: %s", err)
				}
			case <-s.stopChan:
				ticker.Stop()
				return
			}
		}
	}()
}

func (s *PrometheusPushSink) Shutdown() {
	close(s.stopChan)
}
<commit_msg>fix the configuration for summaries created via SummaryDefinitions<commit_after>// +build go1.9

package prometheus

import (
	"fmt"
	"log"
	"math"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
)

var (
	// DefaultPrometheusOpts is the default set of options used when creating a
	// PrometheusSink.
	DefaultPrometheusOpts = PrometheusOpts{
		Expiration: 60 * time.Second,
	}
)

// PrometheusOpts is used to configure the Prometheus Sink
type PrometheusOpts struct {
	// Expiration is the duration a metric is valid for, after which it will be
	// untracked. If the value is zero, a metric is never expired.
	Expiration time.Duration
	Registerer prometheus.Registerer

	// Gauges, Summaries, and Counters allow us to pre-declare metrics by giving their Name, Help, and ConstLabels to
	// the PrometheusSink when it is created. Metrics declared in this way will be initialized at zero and will not be
	// deleted when their expiry is reached.
	// - Gauges and Summaries will be set to NaN when they expire.
	// - Counters continue to Collect their last known value.
	// Ex:
	// PrometheusOpts{
	//     Expiration: 10 * time.Second,
	//     Gauges: []GaugeDefinition{
	//         {
	//	         Name: []string{ "application", "component", "measurement"},
	//           Help: "application_component_measurement provides an example of how to declare static metrics",
	//           ConstLabels: []metrics.Label{ { Name: "my_label", Value: "does_not_change" }, },
	//         },
	//     },
	// }
	GaugeDefinitions   []GaugeDefinition
	SummaryDefinitions []SummaryDefinition
	CounterDefinitions []CounterDefinition
}

type PrometheusSink struct {
	// If these will ever be copied, they should be converted to *sync.Map values and initialized appropriately
	gauges     sync.Map
	summaries  sync.Map
	counters   sync.Map
	expiration time.Duration
}

// GaugeDefinition can be provided to PrometheusOpts to declare a constant gauge that is not deleted on expiry.
type GaugeDefinition struct {
	Name []string
	ConstLabels []metrics.Label
	Help string
}

type gauge struct {
	prometheus.Gauge
	updatedAt time.Time
	// canDelete is set if the metric is created during runtime so we know it's ephemeral and can delete it on expiry.
	canDelete bool
}

// SummaryDefinition can be provided to PrometheusOpts to declare a constant summary that is not deleted on expiry.
type SummaryDefinition struct {
	Name []string
	ConstLabels []metrics.Label
	Help string
}

type summary struct {
	prometheus.Summary
	updatedAt time.Time
	canDelete bool
}

// CounterDefinition can be provided to PrometheusOpts to declare a constant counter that is not deleted on expiry.
type CounterDefinition struct {
	Name []string
	ConstLabels []metrics.Label
	Help string
}

type counter struct {
	prometheus.Counter
	updatedAt time.Time
	canDelete bool
}

// NewPrometheusSink creates a new PrometheusSink using the default options.
func NewPrometheusSink() (*PrometheusSink, error) {
	return NewPrometheusSinkFrom(DefaultPrometheusOpts)
}

// NewPrometheusSinkFrom creates a new PrometheusSink using the passed options.
func NewPrometheusSinkFrom(opts PrometheusOpts) (*PrometheusSink, error) {
	sink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: opts.Expiration,
	}

	initGauges(&sink.gauges, opts.GaugeDefinitions)
	initSummaries(&sink.summaries, opts.SummaryDefinitions)
	initCounters(&sink.counters, opts.CounterDefinitions)

	reg := opts.Registerer
	if reg == nil {
		reg = prometheus.DefaultRegisterer
	}

	return sink, reg.Register(sink)
}

// Describe is needed to meet the Collector interface.
func (p *PrometheusSink) Describe(c chan<- *prometheus.Desc) {
	// We must emit some description otherwise an error is returned. This
	// description isn't shown to the user!
	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(c)
}

// Collect meets the collection interface and allows us to enforce our expiration
// logic to clean up ephemeral metrics if their value haven't been set for a
// duration exceeding our allowed expiration time.
func (p *PrometheusSink) Collect(c chan<- prometheus.Metric) {
	expire := p.expiration != 0
	now := time.Now()
	p.gauges.Range(func(k, v interface{}) bool {
		if v == nil {
			return true
		}
		g := v.(*gauge)
		lastUpdate := g.updatedAt
		if expire && lastUpdate.Add(p.expiration).Before(now) {
			if g.canDelete {
				p.gauges.Delete(k)
				return true
			}
			// We have not observed the gauge this interval so we don't know its value.
			g.Set(math.NaN())
		}
		g.Collect(c)
		return true
	})
	p.summaries.Range(func(k, v interface{}) bool {
		if v == nil {
			return true
		}
		s := v.(*summary)
		lastUpdate := s.updatedAt
		if expire && lastUpdate.Add(p.expiration).Before(now) {
			if s.canDelete {
				p.summaries.Delete(k)
				return true
			}
			// We have observed nothing in this interval.
			s.Observe(math.NaN())
		}
		s.Collect(c)
		return true
	})
	p.counters.Range(func(k, v interface{}) bool {
		if v == nil {
			return true
		}
		count := v.(*counter)
		lastUpdate := count.updatedAt
		if expire && lastUpdate.Add(p.expiration).Before(now) {
			if count.canDelete {
				p.counters.Delete(k)
				return true
			}
			// Counters remain at their previous value when not observed, so we do not set it to NaN.
		}
		count.Collect(c)
		return true
	})
}

func initGauges(m *sync.Map, gauges []GaugeDefinition) {
	for _, g := range gauges {
		key, hash := flattenKey(g.Name, g.ConstLabels)
		pG := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        g.Help,
			ConstLabels: prometheusLabels(g.ConstLabels),
		})
		pG.Set(float64(0)) // Initialize at zero
		m.Store(hash, &gauge{ Gauge: pG })
	}
	return
}

func initSummaries(m *sync.Map, summaries []SummaryDefinition) {
	for _, s := range summaries {
		key, hash := flattenKey(s.Name, s.ConstLabels)
		pS := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        s.Help,
			MaxAge:      10 * time.Second,
			ConstLabels: prometheusLabels(s.ConstLabels),
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		})
		pS.Observe(float64(0)) // Initialize at zero
		m.Store(hash, &summary{ Summary: pS })
	}
	return
}

func initCounters(m *sync.Map, counters []CounterDefinition) {
	for _, c := range counters {
		key, hash := flattenKey(c.Name, c.ConstLabels)
		pC := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        c.Help,
			ConstLabels: prometheusLabels(c.ConstLabels),
		})
		pC.Add(float64(0)) // Initialize at zero
		m.Store(hash, &counter{ Counter: pC })
	}
	return
}

var forbiddenChars = regexp.MustCompile("[ .=\\-/]")

func flattenKey(parts []string, labels []metrics.Label) (string, string) {
	key := strings.Join(parts, "_")
	key = forbiddenChars.ReplaceAllString(key, "_")

	hash := key
	for _, label := range labels {
		hash += fmt.Sprintf(";%s=%s", label.Name, label.Value)
	}

	return key, hash
}

func prometheusLabels(labels []metrics.Label) prometheus.Labels {
	l := make(prometheus.Labels)
	for _, label := range labels {
		l[label.Name] = label.Value
	}
	return l
}

func (p *PrometheusSink) SetGauge(parts []string, val float32) {
	p.SetGaugeWithLabels(parts, val, nil)
}

func (p *PrometheusSink) SetGaugeWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	pg, ok := p.gauges.Load(hash)

	// The sync.Map underlying gauges stores pointers to our structs. If we need to make updates,
	// rather than modifying the underlying value directly, which would be racy, we make a local
	// copy by dereferencing the pointer we get back, making the appropriate changes, and then
	// storing a pointer to our local copy. The underlying Prometheus types are threadsafe,
	// so there's no issues there. It's possible for racy updates to occur to the updatedAt
	// value, but since we're always setting it to time.Now(), it doesn't really matter.
	if ok {
		localGauge := *pg.(*gauge)
		localGauge.Set(float64(val))
		localGauge.updatedAt = time.Now()
		p.gauges.Store(hash, &localGauge)

	// The gauge does not exist, create the gauge and allow it to be deleted
	} else {
		g := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		g.Set(float64(val))
		pg = &gauge{
			Gauge: g,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.gauges.Store(hash, pg)
	}
}

func (p *PrometheusSink) AddSample(parts []string, val float32) {
	p.AddSampleWithLabels(parts, val, nil)
}

func (p *PrometheusSink) AddSampleWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	ps, ok := p.summaries.Load(hash)

	// Does the summary already exist for this sample type?
	if ok {
		localSummary := *ps.(*summary)
		localSummary.Observe(float64(val))
		localSummary.updatedAt = time.Now()
		p.summaries.Store(hash, &localSummary)

	// The summary does not exist, create the Summary and allow it to be deleted
	} else {
		s := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        key,
			MaxAge:      10 * time.Second,
			ConstLabels: prometheusLabels(labels),
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		})
		s.Observe(float64(val))
		ps = &summary{
			Summary: s,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.summaries.Store(hash, ps)
	}
}

// EmitKey is not implemented. Prometheus doesn’t offer a type for which an
// arbitrary number of values is retained, as Prometheus works with a pull
// model, rather than a push model.
func (p *PrometheusSink) EmitKey(key []string, val float32) {
}

func (p *PrometheusSink) IncrCounter(parts []string, val float32) {
	p.IncrCounterWithLabels(parts, val, nil)
}

func (p *PrometheusSink) IncrCounterWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	pc, ok := p.counters.Load(hash)

	// Does the counter exist?
	if ok {
		localCounter := *pc.(*counter)
		localCounter.Add(float64(val))
		localCounter.updatedAt = time.Now()
		p.counters.Store(hash, &localCounter)

	// The counter does not exist yet, create it and allow it to be deleted
	} else {
		c := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		c.Add(float64(val))
		pc = &counter{
			Counter: c,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.counters.Store(hash, pc)
	}
}

// PrometheusPushSink wraps a normal prometheus sink and provides an address and facilities to export it to an address
// on an interval.
type PrometheusPushSink struct {
	*PrometheusSink
	pusher       *push.Pusher
	address      string
	pushInterval time.Duration
	stopChan     chan struct{}
}

// NewPrometheusPushSink creates a PrometheusPushSink by taking an address, interval, and destination name.
func NewPrometheusPushSink(address string, pushInterval time.Duration, name string) (*PrometheusPushSink, error) {
	promSink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: 60 * time.Second,
	}

	pusher := push.New(address, name).Collector(promSink)

	sink := &PrometheusPushSink{
		promSink,
		pusher,
		address,
		pushInterval,
		make(chan struct{}),
	}

	sink.flushMetrics()
	return sink, nil
}

func (s *PrometheusPushSink) flushMetrics() {
	ticker := time.NewTicker(s.pushInterval)

	go func() {
		for {
			select {
			case <-ticker.C:
				err := s.pusher.Push()
				if err != nil {
					log.Printf("[ERR] Error pushing to Prometheus! Err: %s", err)
				}
			case <-s.stopChan:
				ticker.Stop()
				return
			}
		}
	}()
}

func (s *PrometheusPushSink) Shutdown() {
	close(s.stopChan)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"io"
	"log"
	"os"
	"strconv"
	"time"

	"github.com/flynn/go-discover/discover"
	"github.com/flynn/lorne/types"
	"github.com/flynn/rpcplus"
	"github.com/flynn/sampi/types"
	"github.com/titanous/go-dockerclient"
)

var state = NewState()
var Docker *docker.Client

func main() {
	go attachServer()
	go rpcServer()
	go allocatePorts()

	id := randomID()
	disc, err := discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}
	if err := disc.Register("flynn-lorne-rpc."+id, "1113", nil); err != nil {
		log.Fatal(err)
	}
	if err := disc.Register("flynn-lorne-attach."+id, "1114", nil); err != nil {
		log.Fatal(err)
	}

	services := disc.Services("flynn-sampi")
	time.Sleep(100 * time.Millisecond) // HAX: remove this when Online is blocking
	schedulers := services.Online()
	if len(schedulers) == 0 {
		log.Fatal("No sampi instances found")
	}

	scheduler, err := rpcplus.DialHTTP("tcp", schedulers[0].Addr)
	if err != nil {
		log.Fatal(err)
	}
	log.Print("Connected to scheduler")

	Docker, err = docker.NewClient("http://localhost:4243")
	if err != nil {
		log.Fatal(err)
	}

	go streamEvents(Docker)
	go syncScheduler(scheduler)

	host := sampi.Host{
		ID:        id,
		Resources: map[string]sampi.ResourceValue{"memory": sampi.ResourceValue{Value: 1024}},
	}

	jobs := make(chan *sampi.Job)
	scheduler.StreamGo("Scheduler.RegisterHost", host, jobs)
	log.Print("Host registered")
	for job := range jobs {
		log.Printf("%#v", job.Config)
		if job.TCPPorts > 0 {
			port := strconv.Itoa(<-portAllocator)
			job.Config.Env = append(job.Config.Env, "PORT="+port)
			job.Config.PortSpecs = []string{port + ":" + port}
		}
		state.AddJob(job)
		container, err := Docker.CreateContainer(job.Config)
		if err == docker.ErrNoSuchImage {
			err = Docker.PullImage(docker.PullImageOptions{Repository: job.Config.Image}, os.Stdout)
			if err != nil {
				log.Fatal(err)
			}
			container, err = Docker.CreateContainer(job.Config)
		}
		if err != nil {
			log.Fatal(err)
		}
		state.SetContainerID(job.ID, container.ID)
		state.WaitAttach(job.ID)
		if err := Docker.StartContainer(container.ID, nil); err != nil {
			log.Fatal(err)
		}
		state.SetStatusRunning(job.ID)
	}
}

func syncScheduler(client *rpcplus.Client) {
	events := make(chan lorne.Event)
	state.AddListener("all", events)
	for event := range events {
		if event.Event != "stop" {
			continue
		}
		log.Println("remove job", event.JobID)
		if err := client.Call("Scheduler.RemoveJobs", []string{event.JobID}, &struct{}{}); err != nil {
			log.Println("remove job", event.JobID, "error:", err)
			// TODO: try to reconnect?
		}
	}
}

func streamEvents(client *docker.Client) {
	stream, err := client.Events()
	if err != nil {
		log.Fatal(err)
	}
	for event := range stream.Events {
		log.Printf("%#v", event)
		if event.Status != "die" {
			continue
		}
		container, err := client.InspectContainer(event.ID)
		if err != nil {
			log.Println("inspect container", event.ID, "error:", err)
			// TODO: set job status anyway?
			continue
		}
		state.SetStatusDone(event.ID, container.State.ExitCode)
	}
	log.Println("events done", stream.Error)
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}

var portAllocator = make(chan int)

// TODO: fix this, horribly broken
const startPort = 55000
const endPort = 65535

func allocatePorts() {
	for i := startPort; i < endPort; i++ {
		portAllocator <- i
	}
	// TODO: handle wrap-around
}
<commit_msg>host: Update discover API usage<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"io"
	"log"
	"os"
	"strconv"

	"github.com/flynn/go-discover/discover"
	"github.com/flynn/lorne/types"
	"github.com/flynn/rpcplus"
	"github.com/flynn/sampi/types"
	"github.com/titanous/go-dockerclient"
)

var state = NewState()
var Docker *docker.Client

func main() {
	go attachServer()
	go rpcServer()
	go allocatePorts()

	id := randomID()
	disc, err := discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}
	if err := disc.Register("flynn-lorne-rpc."+id, "1113", nil); err != nil {
		log.Fatal(err)
	}
	if err := disc.Register("flynn-lorne-attach."+id, "1114", nil); err != nil {
		log.Fatal(err)
	}

	services, err := disc.Services("flynn-sampi")
	if err != nil {
		log.Fatal(err)
	}
	schedulers := services.Online()
	if len(schedulers) == 0 {
		log.Fatal("No sampi instances found")
	}

	scheduler, err := rpcplus.DialHTTP("tcp", schedulers[0].Addr)
	if err != nil {
		log.Fatal(err)
	}
	log.Print("Connected to scheduler")

	Docker, err = docker.NewClient("http://localhost:4243")
	if err != nil {
		log.Fatal(err)
	}

	go streamEvents(Docker)
	go syncScheduler(scheduler)

	host := sampi.Host{
		ID:        id,
		Resources: map[string]sampi.ResourceValue{"memory": sampi.ResourceValue{Value: 1024}},
	}

	jobs := make(chan *sampi.Job)
	scheduler.StreamGo("Scheduler.RegisterHost", host, jobs)
	log.Print("Host registered")
	for job := range jobs {
		log.Printf("%#v", job.Config)
		if job.TCPPorts > 0 {
			port := strconv.Itoa(<-portAllocator)
			job.Config.Env = append(job.Config.Env, "PORT="+port)
			job.Config.PortSpecs = []string{port + ":" + port}
		}
		state.AddJob(job)
		container, err := Docker.CreateContainer(job.Config)
		if err == docker.ErrNoSuchImage {
			err = Docker.PullImage(docker.PullImageOptions{Repository: job.Config.Image}, os.Stdout)
			if err != nil {
				log.Fatal(err)
			}
			container, err = Docker.CreateContainer(job.Config)
		}
		if err != nil {
			log.Fatal(err)
		}
		state.SetContainerID(job.ID, container.ID)
		state.WaitAttach(job.ID)
		if err := Docker.StartContainer(container.ID, nil); err != nil {
			log.Fatal(err)
		}
		state.SetStatusRunning(job.ID)
	}
}

func syncScheduler(client *rpcplus.Client) {
	events := make(chan lorne.Event)
	state.AddListener("all", events)
	for event := range events {
		if event.Event != "stop" {
			continue
		}
		log.Println("remove job", event.JobID)
		if err := client.Call("Scheduler.RemoveJobs", []string{event.JobID}, &struct{}{}); err != nil {
			log.Println("remove job", event.JobID, "error:", err)
			// TODO: try to reconnect?
		}
	}
}

func streamEvents(client *docker.Client) {
	stream, err := client.Events()
	if err != nil {
		log.Fatal(err)
	}
	for event := range stream.Events {
		log.Printf("%#v", event)
		if event.Status != "die" {
			continue
		}
		container, err := client.InspectContainer(event.ID)
		if err != nil {
			log.Println("inspect container", event.ID, "error:", err)
			// TODO: set job status anyway?
			continue
		}
		state.SetStatusDone(event.ID, container.State.ExitCode)
	}
	log.Println("events done", stream.Error)
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}

var portAllocator = make(chan int)

// TODO: fix this, horribly broken
const startPort = 55000
const endPort = 65535

func allocatePorts() {
	for i := startPort; i < endPort; i++ {
		portAllocator <- i
	}
	// TODO: handle wrap-around
}
<|endoftext|>