<commit_before>// This package provides all the input fields logic and customization methods.
package fields

import (
	"github.com/coscms/forms/common"
	"github.com/coscms/forms/widgets"
	"html/template"
	"strings"
)

// Field is a generic type containing all data associated to an input field.
type Field struct {
	fieldType      string
	Widget         widgets.WidgetInterface // Public Widget field for widget customization
	name           string
	placeholder    string
	class          []string
	id             string
	params         map[string]string
	css            map[string]string
	label          string
	labelClass     []string
	tag            map[string]struct{}
	value          string
	helptext       string
	errors         []string
	additionalData map[string]interface{}
}

// FieldInterface defines the interface an object must implement to be used in a form. Every method returns a FieldInterface object
// to allow methods chaining.
type FieldInterface interface {
	Name() string
	Render(appendData ...map[string]interface{}) template.HTML
	AddClass(class string) FieldInterface
	RemoveClass(class string) FieldInterface
	AddTag(class string) FieldInterface
	RemoveTag(class string) FieldInterface
	SetId(id string) FieldInterface
	SetParam(key, value string) FieldInterface
	DeleteParam(key string) FieldInterface
	AddCss(key, value string) FieldInterface
	RemoveCss(key string) FieldInterface
	SetStyle(style string) FieldInterface
	SetLabel(label string) FieldInterface
	AddLabelClass(class string) FieldInterface
	RemoveLabelClass(class string) FieldInterface
	SetValue(value string) FieldInterface
	Disabled() FieldInterface
	Enabled() FieldInterface
	SetHelptext(text string) FieldInterface
	AddError(err string) FieldInterface
	MultipleChoice() FieldInterface
	SingleChoice() FieldInterface
	AddSelected(opt ...string) FieldInterface
	RemoveSelected(opt string) FieldInterface
	SetSelectChoices(choices map[string][]InputChoice) FieldInterface
	SetRadioChoices(choices []InputChoice) FieldInterface
	SetText(text string) FieldInterface
}

// FieldWithType creates an empty field of the given type and identified by name.
func FieldWithType(name, t string) *Field {
	return &Field{
		fieldType:      t,
		Widget:         nil,
		name:           name,
		placeholder:    "",
		class:          []string{},
		id:             "",
		params:         map[string]string{},
		css:            map[string]string{},
		label:          "",
		labelClass:     []string{},
		tag:            map[string]struct{}{},
		value:          "",
		helptext:       "",
		errors:         []string{},
		additionalData: map[string]interface{}{},
	}
}

// SetStyle sets the style (e.g.: BASE, BOOTSTRAP) of the field, correctly populating the Widget field.
func (f *Field) SetStyle(style string) FieldInterface {
	f.Widget = widgets.BaseWidget(style, f.fieldType)
	return f
}

// Name returns the name of the field.
func (f *Field) Name() string {
	return strings.TrimSuffix(f.name, "[]")
}

func (f *Field) dataForRender(appendData ...map[string]interface{}) map[string]interface{} {
	safeParams := make(map[template.HTMLAttr]string)
	for k, v := range f.params {
		safeParams[template.HTMLAttr(k)] = v
	}
	data := map[string]interface{}{
		"classes":      f.class,
		"id":           f.id,
		"name":         f.name,
		"params":       safeParams,
		"css":          f.css,
		"type":         f.fieldType,
		"label":        f.label,
		"labelClasses": f.labelClass,
		"tags":         f.tag,
		"value":        f.value,
		"helptext":     f.helptext,
		"errors":       f.errors,
		"container":    "form",
	}
	for k, v := range f.additionalData {
		data[k] = v
	}
	for _, val := range appendData {
		for k, v := range val {
			data[k] = v
		}
	}
	return data
}

// Render packs all data and executes widget render method.
func (f *Field) Render(appendData ...map[string]interface{}) template.HTML {
	if f.Widget != nil {
		data := f.dataForRender(appendData...)
		return template.HTML(f.Widget.Render(data))
	}
	return template.HTML("")
}

// AddClass adds a class to the field.
func (f *Field) AddClass(class string) FieldInterface {
	f.class = append(f.class, class)
	return f
}

// RemoveClass removes a class from the field, if it was present.
func (f *Field) RemoveClass(class string) FieldInterface {
	ind := -1
	for i, v := range f.class {
		if v == class {
			ind = i
			break
		}
	}

	if ind != -1 {
		f.class = append(f.class[:ind], f.class[ind+1:]...)
	}
	return f
}

// SetId associates the given id to the field, overwriting any previous id.
func (f *Field) SetId(id string) FieldInterface {
	f.id = id
	return f
}

// SetLabel saves the label to be rendered along with the field.
func (f *Field) SetLabel(label string) FieldInterface {
	f.label = label
	return f
}

// SetLablClass allows to define custom classes for the label.
func (f *Field) AddLabelClass(class string) FieldInterface {
	f.labelClass = append(f.labelClass, class)
	return f
}

// RemoveLabelClass removes the given class from the field label.
func (f *Field) RemoveLabelClass(class string) FieldInterface {
	ind := -1
	for i, v := range f.labelClass {
		if v == class {
			ind = i
			break
		}
	}

	if ind != -1 {
		f.labelClass = append(f.labelClass[:ind], f.labelClass[ind+1:]...)
	}
	return f
}

// SetParam adds a parameter (defined as key-value pair) in the field.
func (f *Field) SetParam(key, value string) FieldInterface {
	f.params[key] = value
	return f
}

// DeleteParam removes a parameter identified by key from the field.
func (f *Field) DeleteParam(key string) FieldInterface {
	delete(f.params, key)
	return f
}

// AddCss adds a custom CSS style the field.
func (f *Field) AddCss(key, value string) FieldInterface {
	f.css[key] = value
	return f
}

// RemoveCss removes CSS options identified by key from the field.
func (f *Field) RemoveCss(key string) FieldInterface {
	delete(f.css, key)
	return f
}

// Disabled add the "disabled" tag to the field, making it unresponsive in some environments (e.g. Bootstrap).
func (f *Field) Disabled() FieldInterface {
	f.AddTag("disabled")
	return f
}

// Enabled removes the "disabled" tag from the field, making it responsive.
func (f *Field) Enabled() FieldInterface {
	f.RemoveTag("disabled")
	return f
}

// AddTag adds a no-value parameter (e.g.: checked, disabled) to the field.
func (f *Field) AddTag(tag string) FieldInterface {
	f.tag[tag] = struct{}{}
	return f
}

// RemoveTag removes a no-value parameter from the field.
func (f *Field) RemoveTag(tag string) FieldInterface {
	delete(f.tag, tag)
	return f
}

// SetValue sets the value parameter for the field.
func (f *Field) SetValue(value string) FieldInterface {
	f.value = value
	return f
}

// SetHelptext saves the field helptext.
func (f *Field) SetHelptext(text string) FieldInterface {
	f.helptext = text
	return f
}

// AddError adds an error string to the field. It's valid only for Bootstrap forms.
func (f *Field) AddError(err string) FieldInterface {
	f.errors = append(f.errors, err)
	return f
}

// MultipleChoice configures the SelectField to accept and display multiple choices.
// It has no effect if type is not SELECT.
func (f *Field) MultipleChoice() FieldInterface {
	if f.fieldType == formcommon.SELECT {
		f.AddTag("multiple")
		// fix name if necessary
		if !strings.HasSuffix(f.name, "[]") {
			f.name = f.name + "[]"
		}
	}
	return f
}

// SingleChoice configures the Field to accept and display only one choice (valid for SelectFields only).
// It has no effect if type is not SELECT.
func (f *Field) SingleChoice() FieldInterface {
	if f.fieldType == formcommon.SELECT {
		f.RemoveTag("multiple")
		if strings.HasSuffix(f.name, "[]") {
			f.name = strings.TrimSuffix(f.name, "[]")
		}
	}
	return f
}

// If the field is configured as "multiple", AddSelected adds a selected value to the field (valid for SelectFields only).
// It has no effect if type is not SELECT.
func (f *Field) AddSelected(opt ...string) FieldInterface {
	if f.fieldType == formcommon.SELECT {
		for _, v := range opt {
			f.additionalData["multValues"].(map[string]struct{})[v] = struct{}{}
		}
	}
	return f
}

// If the field is configured as "multiple", AddSelected removes the selected value from the field (valid for SelectFields only).
// It has no effect if type is not SELECT.
func (f *Field) RemoveSelected(opt string) FieldInterface {
	if f.fieldType == formcommon.SELECT {
		if _, ok := f.additionalData["multValues"]; ok {
			delete(f.additionalData["multValues"].(map[string]struct{}), opt)
		}
	}
	return f
}

// SetSelectChoices takes as input a dictionary whose key-value entries are defined as follows: key is the group name (the empty string
// is the default group that is not explicitly rendered) and value is the list of choices belonging to that group.
// Grouping is only useful for Select fields, while groups are ignored in Radio fields.
// It has no effect if type is not SELECT.
func (f *Field) SetSelectChoices(choices map[string][]InputChoice) FieldInterface {
	if f.fieldType == formcommon.SELECT {
		f.additionalData["choices"] = choices
	}
	return f
}

// SetRadioChoices sets an array of InputChoice objects as the possible choices for a radio field. It has no effect if type is not RADIO.
func (f *Field) SetRadioChoices(choices []InputChoice) FieldInterface {
	if f.fieldType == formcommon.RADIO {
		f.additionalData["choices"] = choices
	}
	return f
}

// SetText saves the provided text as content of the field, usually a TextAreaField.
func (f *Field) SetText(text string) FieldInterface {
	if f.fieldType == formcommon.BUTTON ||
		f.fieldType == formcommon.SUBMIT ||
		f.fieldType == formcommon.RESET ||
		f.fieldType == formcommon.STATIC ||
		f.fieldType == formcommon.TEXTAREA {
		f.additionalData["text"] = text
	}
	return f
}
<commit_msg>update<commit_after>// This package provides all the input fields logic and customization methods.
package fields

import (
	"github.com/coscms/forms/common"
	"github.com/coscms/forms/widgets"
	"html/template"
	"strings"
)

// Field is a generic type containing all data associated to an input field.
type Field struct {
	fieldType      string
	Widget         widgets.WidgetInterface // Public Widget field for widget customization
	name           string
	class          []string
	id             string
	params         map[string]string
	css            map[string]string
	label          string
	labelClass     []string
	tag            map[string]struct{}
	value          string
	helptext       string
	errors         []string
	additionalData map[string]interface{}
}

// FieldInterface defines the interface an object must implement to be used in a form. Every method returns a FieldInterface object
// to allow methods chaining.
type FieldInterface interface {
	Name() string
	Render(appendData ...map[string]interface{}) template.HTML
	AddClass(class string) FieldInterface
	RemoveClass(class string) FieldInterface
	AddTag(class string) FieldInterface
	RemoveTag(class string) FieldInterface
	SetId(id string) FieldInterface
	SetParam(key, value string) FieldInterface
	DeleteParam(key string) FieldInterface
	AddCss(key, value string) FieldInterface
	RemoveCss(key string) FieldInterface
	SetStyle(style string) FieldInterface
	SetLabel(label string) FieldInterface
	AddLabelClass(class string) FieldInterface
	RemoveLabelClass(class string) FieldInterface
	SetValue(value string) FieldInterface
	Disabled() FieldInterface
	Enabled() FieldInterface
	SetHelptext(text string) FieldInterface
	AddError(err string) FieldInterface
	MultipleChoice() FieldInterface
	SingleChoice() FieldInterface
	AddSelected(opt ...string) FieldInterface
	RemoveSelected(opt string) FieldInterface
	SetSelectChoices(choices map[string][]InputChoice) FieldInterface
	SetRadioChoices(choices []InputChoice) FieldInterface
	SetText(text string) FieldInterface
}

// FieldWithType creates an empty field of the given type and identified by name.
func FieldWithType(name, t string) *Field {
	return &Field{
		fieldType:      t,
		Widget:         nil,
		name:           name,
		class:          []string{},
		id:             "",
		params:         map[string]string{},
		css:            map[string]string{},
		label:          "",
		labelClass:     []string{},
		tag:            map[string]struct{}{},
		value:          "",
		helptext:       "",
		errors:         []string{},
		additionalData: map[string]interface{}{},
	}
}

// SetStyle sets the style (e.g.: BASE, BOOTSTRAP) of the field, correctly populating the Widget field.
func (f *Field) SetStyle(style string) FieldInterface {
	f.Widget = widgets.BaseWidget(style, f.fieldType)
	return f
}

// Name returns the name of the field.
func (f *Field) Name() string {
	return strings.TrimSuffix(f.name, "[]")
}

func (f *Field) dataForRender(appendData ...map[string]interface{}) map[string]interface{} {
	safeParams := make(map[template.HTMLAttr]string)
	for k, v := range f.params {
		safeParams[template.HTMLAttr(k)] = v
	}
	data := map[string]interface{}{
		"classes":      f.class,
		"id":           f.id,
		"name":         f.name,
		"params":       safeParams,
		"css":          f.css,
		"type":         f.fieldType,
		"label":        f.label,
		"labelClasses": f.labelClass,
		"tags":         f.tag,
		"value":        f.value,
		"helptext":     f.helptext,
		"errors":       f.errors,
		"container":    "form",
	}
	for k, v := range f.additionalData {
		data[k] = v
	}
	for _, val := range appendData {
		for k, v := range val {
			data[k] = v
		}
	}
	return data
}

// Render packs all data and executes widget render method.
func (f *Field) Render(appendData ...map[string]interface{}) template.HTML {
	if f.Widget != nil {
		data := f.dataForRender(appendData...)
		return template.HTML(f.Widget.Render(data))
	}
	return template.HTML("")
}

// AddClass adds a class to the field.
func (f *Field) AddClass(class string) FieldInterface {
	f.class = append(f.class, class)
	return f
}

// RemoveClass removes a class from the field, if it was present.
func (f *Field) RemoveClass(class string) FieldInterface {
	ind := -1
	for i, v := range f.class {
		if v == class {
			ind = i
			break
		}
	}

	if ind != -1 {
		f.class = append(f.class[:ind], f.class[ind+1:]...)
	}
	return f
}

// SetId associates the given id to the field, overwriting any previous id.
func (f *Field) SetId(id string) FieldInterface {
	f.id = id
	return f
}

// SetLabel saves the label to be rendered along with the field.
func (f *Field) SetLabel(label string) FieldInterface {
	f.label = label
	return f
}

// SetLablClass allows to define custom classes for the label.
func (f *Field) AddLabelClass(class string) FieldInterface {
	f.labelClass = append(f.labelClass, class)
	return f
}

// RemoveLabelClass removes the given class from the field label.
func (f *Field) RemoveLabelClass(class string) FieldInterface {
	ind := -1
	for i, v := range f.labelClass {
		if v == class {
			ind = i
			break
		}
	}

	if ind != -1 {
		f.labelClass = append(f.labelClass[:ind], f.labelClass[ind+1:]...)
	}
	return f
}

// SetParam adds a parameter (defined as key-value pair) in the field.
func (f *Field) SetParam(key, value string) FieldInterface {
	f.params[key] = value
	return f
}

// DeleteParam removes a parameter identified by key from the field.
func (f *Field) DeleteParam(key string) FieldInterface {
	delete(f.params, key)
	return f
}

// AddCss adds a custom CSS style the field.
func (f *Field) AddCss(key, value string) FieldInterface {
	f.css[key] = value
	return f
}

// RemoveCss removes CSS options identified by key from the field.
func (f *Field) RemoveCss(key string) FieldInterface {
	delete(f.css, key)
	return f
}

// Disabled add the "disabled" tag to the field, making it unresponsive in some environments (e.g. Bootstrap).
func (f *Field) Disabled() FieldInterface {
	f.AddTag("disabled")
	return f
}

// Enabled removes the "disabled" tag from the field, making it responsive.
func (f *Field) Enabled() FieldInterface {
	f.RemoveTag("disabled")
	return f
}

// AddTag adds a no-value parameter (e.g.: checked, disabled) to the field.
func (f *Field) AddTag(tag string) FieldInterface {
	f.tag[tag] = struct{}{}
	return f
}

// RemoveTag removes a no-value parameter from the field.
func (f *Field) RemoveTag(tag string) FieldInterface {
	delete(f.tag, tag)
	return f
}

// SetValue sets the value parameter for the field.
func (f *Field) SetValue(value string) FieldInterface {
	f.value = value
	return f
}

// SetHelptext saves the field helptext.
func (f *Field) SetHelptext(text string) FieldInterface {
	f.helptext = text
	return f
}

// AddError adds an error string to the field. It's valid only for Bootstrap forms.
func (f *Field) AddError(err string) FieldInterface {
	f.errors = append(f.errors, err)
	return f
}

// MultipleChoice configures the SelectField to accept and display multiple choices.
// It has no effect if type is not SELECT.
func (f *Field) MultipleChoice() FieldInterface {
	if f.fieldType == formcommon.SELECT {
		f.AddTag("multiple")
		// fix name if necessary
		if !strings.HasSuffix(f.name, "[]") {
			f.name = f.name + "[]"
		}
	}
	return f
}

// SingleChoice configures the Field to accept and display only one choice (valid for SelectFields only).
// It has no effect if type is not SELECT.
func (f *Field) SingleChoice() FieldInterface {
	if f.fieldType == formcommon.SELECT {
		f.RemoveTag("multiple")
		if strings.HasSuffix(f.name, "[]") {
			f.name = strings.TrimSuffix(f.name, "[]")
		}
	}
	return f
}

// If the field is configured as "multiple", AddSelected adds a selected value to the field (valid for SelectFields only).
// It has no effect if type is not SELECT.
func (f *Field) AddSelected(opt ...string) FieldInterface {
	if f.fieldType == formcommon.SELECT {
		for _, v := range opt {
			f.additionalData["multValues"].(map[string]struct{})[v] = struct{}{}
		}
	}
	return f
}

// If the field is configured as "multiple", AddSelected removes the selected value from the field (valid for SelectFields only).
// It has no effect if type is not SELECT.
func (f *Field) RemoveSelected(opt string) FieldInterface {
	if f.fieldType == formcommon.SELECT {
		if _, ok := f.additionalData["multValues"]; ok {
			delete(f.additionalData["multValues"].(map[string]struct{}), opt)
		}
	}
	return f
}

// SetSelectChoices takes as input a dictionary whose key-value entries are defined as follows: key is the group name (the empty string
// is the default group that is not explicitly rendered) and value is the list of choices belonging to that group.
// Grouping is only useful for Select fields, while groups are ignored in Radio fields.
// It has no effect if type is not SELECT.
func (f *Field) SetSelectChoices(choices map[string][]InputChoice) FieldInterface {
	if f.fieldType == formcommon.SELECT {
		f.additionalData["choices"] = choices
	}
	return f
}

// SetRadioChoices sets an array of InputChoice objects as the possible choices for a radio field. It has no effect if type is not RADIO.
func (f *Field) SetRadioChoices(choices []InputChoice) FieldInterface {
	if f.fieldType == formcommon.RADIO {
		f.additionalData["choices"] = choices
	}
	return f
}

// SetText saves the provided text as content of the field, usually a TextAreaField.
func (f *Field) SetText(text string) FieldInterface {
	if f.fieldType == formcommon.BUTTON ||
		f.fieldType == formcommon.SUBMIT ||
		f.fieldType == formcommon.RESET ||
		f.fieldType == formcommon.STATIC ||
		f.fieldType == formcommon.TEXTAREA {
		f.additionalData["text"] = text
	}
	return f
}
<|endoftext|>


<commit_before>package stats

import (
	"math"
	"reflect"
	"runtime"
	"testing"
)

var data1 = Float64Data{-10, -10.001, 5, 1.1, 2, 3, 4.20, 5}
var data2 = Float64Data{-9, -9.001, 4, .1, 1, 2, 3.20, 5}

func getFunctionName(i interface{}) string {
	return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
}

func checkResult(result float64, err error, name string, f float64, t *testing.T) {
	if err != nil {
		t.Errorf("%s returned an error", name)
	}
	if result != f {
		t.Errorf("%s() => %v != %v", name, result, f)
	}
}

func TestInterfaceMethods(t *testing.T) {
	// Test Get
	a := data1.Get(1)
	if a != -10.001 {
		t.Errorf("Get(2) => %.1f != %.1f", a, -10.001)
	}

	// Test Len
	l := data1.Len()
	if l != 8 {
		t.Errorf("Len() => %v != %v", l, 8)
	}

	// Test Less
	b := data1.Less(0, 5)
	if b != true {
		t.Errorf("Less() => %v != %v", b, true)
	}

	// Test Swap
	data1.Swap(0, 2)
	if data1.Get(0) != 5 {
		t.Errorf("Len() => %v != %v", l, 8)
	}
}

func TestHelperMethods(t *testing.T) {

	// Test Min
	m, _ := data1.Min()
	if m != -10.001 {
		t.Errorf("Min() => %v != %v", m, -10.001)
	}

	// Test Max
	m, _ = data1.Max()
	if m != 5 {
		t.Errorf("Max() => %v != %v", m, 5)
	}

	// Test Sum
	m, _ = data1.Sum()
	if m != 0.2990000000000004 {
		t.Errorf("Sum() => %v != %v", m, 0.2990000000000004)
	}

	// Test Mean
	m, _ = data1.Mean()
	if m != 0.03737500000000005 {
		t.Errorf("Mean() => %v != %v", m, 0.03737500000000005)
	}

	// Test GeometricMean
	m, _ = data1.GeometricMean()
	if m != 4.028070682618703 {
		t.Errorf("GeometricMean() => %v != %v", m, 4.028070682618703)
	}

	// Test HarmonicMean
	m, _ = data1.HarmonicMean()
	if !math.IsNaN(m) {
		t.Errorf("HarmonicMean() => %v != %v", m, math.NaN())
	}

	// Test Median
	m, _ = data1.Median()
	if m != 2.5 {
		t.Errorf("Median() => %v != %v", m, 2.5)
	}

	// Test Mode
	mo, _ := data1.Mode()
	if !reflect.DeepEqual(mo, []float64{5.0}) {
		t.Errorf("Mode() => %.1f != %.1f", mo, []float64{5.0})
	}
}

func assertFloat64(fn func() (float64, error), f float64, t *testing.T) {
	res, err := fn()
	checkResult(res, err, getFunctionName(fn), f, t)
}

func TestMedianAbsoluteDeviationMethods(t *testing.T) {
	assertFloat64(data1.MedianAbsoluteDeviation, 2.1, t)
	assertFloat64(data1.MedianAbsoluteDeviationPopulation, 2.1, t)
}

func TestStandardDeviationMethods(t *testing.T) {
	assertFloat64(data1.StandardDeviation, 5.935684731720091, t)
	assertFloat64(data1.StandardDeviationPopulation, 5.935684731720091, t)
	assertFloat64(data1.StandardDeviationSample, 6.345513892000508, t)
}

func TestVarianceMethods(t *testing.T) {
	assertFloat64(data1.Variance, 35.232353234375005, t)
	assertFloat64(data1.PopulationVariance, 35.232353234375005, t)
	assertFloat64(data1.SampleVariance, 40.26554655357143, t)

}

func assertPercentiles(fn func(i float64) (float64, error), i float64, f float64, t *testing.T) {
	res, err := fn(i)
	checkResult(res, err, getFunctionName(fn), f, t)
}

func TestPercentileMethods(t *testing.T) {
	assertPercentiles(data1.Percentile, 75, 4.6, t)
	assertPercentiles(data1.PercentileNearestRank, 75, 4.2, t)

}

func assertOtherDataMethods(fn func(d Float64Data) (float64, error), d Float64Data, f float64, t *testing.T) {
	res, err := fn(d)
	checkResult(res, err, getFunctionName(fn), f, t)
}

func TestOtherDataMethods(t *testing.T) {
	assertOtherDataMethods(data1.Correlation, data2, 0.20875473597605448, t)
	assertOtherDataMethods(data1.Pearson, data2, 0.20875473597605448, t)
	assertOtherDataMethods(data1.InterQuartileRange, data2, 8.05, t)
	assertOtherDataMethods(data1.Midhinge, data2, -0.42500000000000004, t)
	assertOtherDataMethods(data1.Trimean, data2, 0.5375, t)
	assertOtherDataMethods(data1.Covariance, data2, 7.3814215535714265, t)
	assertOtherDataMethods(data1.CovariancePopulation, data2, 6.458743859374998, t)
}

func TestSampleMethod(t *testing.T) {
	// Test Sample method
	_, err := data1.Sample(5, true)
	if err != nil {
		t.Errorf("%s returned an error", getFunctionName(data1.Sample))
	}
}

func TestQuartileMethods(t *testing.T) {
	// Test QuartileOutliers method
	_, err := data1.QuartileOutliers()
	if err != nil {
		t.Errorf("%s returned an error", getFunctionName(data1.QuartileOutliers))
	}

	// Test Quartile method
	_, err = data1.Quartile(data2)
	if err != nil {
		t.Errorf("%s returned an error", getFunctionName(data1.Quartile))
	}
}

// Here we show the regular way of doing it
// with a plain old slice of float64s
func BenchmarkRegularAPI(b *testing.B) {
	for i := 0; i < b.N; i++ {
		data := []float64{-10, -7, -3.11, 5, 1.1, 2, 3, 4.20, 5, 18}
		Min(data)
		Max(data)
		Sum(data)
		Mean(data)
		Median(data)
		Mode(data)
	}
}

// Here's where things get interesting
// and we start to use the included
// Float64Data type and methods
func BenchmarkMethodsAPI(b *testing.B) {
	for i := 0; i < b.N; i++ {
		data := Float64Data{-10, -7, -3.11, 5, 1.1, 2, 3, 4.20, 5, 18}
		data.Min()
		data.Max()
		data.Sum()
		data.Mean()
		data.Median()
		data.Mode()
	}
}
<commit_msg>Fix InterQuartileRange method test<commit_after>package stats

import (
	"math"
	"reflect"
	"runtime"
	"testing"
)

var data1 = Float64Data{-10, -10.001, 5, 1.1, 2, 3, 4.20, 5}
var data2 = Float64Data{-9, -9.001, 4, .1, 1, 2, 3.20, 5}

func getFunctionName(i interface{}) string {
	return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
}

func checkResult(result float64, err error, name string, f float64, t *testing.T) {
	if err != nil {
		t.Errorf("%s returned an error", name)
	}
	if result != f {
		t.Errorf("%s() => %v != %v", name, result, f)
	}
}

func TestInterfaceMethods(t *testing.T) {
	// Test Get
	a := data1.Get(1)
	if a != -10.001 {
		t.Errorf("Get(2) => %.1f != %.1f", a, -10.001)
	}

	// Test Len
	l := data1.Len()
	if l != 8 {
		t.Errorf("Len() => %v != %v", l, 8)
	}

	// Test Less
	b := data1.Less(0, 5)
	if b != true {
		t.Errorf("Less() => %v != %v", b, true)
	}

	// Test Swap
	data1.Swap(0, 2)
	if data1.Get(0) != 5 {
		t.Errorf("Len() => %v != %v", l, 8)
	}
}

func TestHelperMethods(t *testing.T) {

	// Test Min
	m, _ := data1.Min()
	if m != -10.001 {
		t.Errorf("Min() => %v != %v", m, -10.001)
	}

	// Test Max
	m, _ = data1.Max()
	if m != 5 {
		t.Errorf("Max() => %v != %v", m, 5)
	}

	// Test Sum
	m, _ = data1.Sum()
	if m != 0.2990000000000004 {
		t.Errorf("Sum() => %v != %v", m, 0.2990000000000004)
	}

	// Test Mean
	m, _ = data1.Mean()
	if m != 0.03737500000000005 {
		t.Errorf("Mean() => %v != %v", m, 0.03737500000000005)
	}

	// Test GeometricMean
	m, _ = data1.GeometricMean()
	if m != 4.028070682618703 {
		t.Errorf("GeometricMean() => %v != %v", m, 4.028070682618703)
	}

	// Test HarmonicMean
	m, _ = data1.HarmonicMean()
	if !math.IsNaN(m) {
		t.Errorf("HarmonicMean() => %v != %v", m, math.NaN())
	}

	// Test Median
	m, _ = data1.Median()
	if m != 2.5 {
		t.Errorf("Median() => %v != %v", m, 2.5)
	}

	// Test Mode
	mo, _ := data1.Mode()
	if !reflect.DeepEqual(mo, []float64{5.0}) {
		t.Errorf("Mode() => %.1f != %.1f", mo, []float64{5.0})
	}

	// Test InterQuartileRange
	iqr, _ := data1.InterQuartileRange()
	if iqr != 9.05 {
		t.Errorf("InterQuartileRange() => %v != %v", iqr, 9.05)
	}
}

func assertFloat64(fn func() (float64, error), f float64, t *testing.T) {
	res, err := fn()
	checkResult(res, err, getFunctionName(fn), f, t)
}

func TestMedianAbsoluteDeviationMethods(t *testing.T) {
	assertFloat64(data1.MedianAbsoluteDeviation, 2.1, t)
	assertFloat64(data1.MedianAbsoluteDeviationPopulation, 2.1, t)
}

func TestStandardDeviationMethods(t *testing.T) {
	assertFloat64(data1.StandardDeviation, 5.935684731720091, t)
	assertFloat64(data1.StandardDeviationPopulation, 5.935684731720091, t)
	assertFloat64(data1.StandardDeviationSample, 6.345513892000508, t)
}

func TestVarianceMethods(t *testing.T) {
	assertFloat64(data1.Variance, 35.232353234375005, t)
	assertFloat64(data1.PopulationVariance, 35.232353234375005, t)
	assertFloat64(data1.SampleVariance, 40.26554655357143, t)

}

func assertPercentiles(fn func(i float64) (float64, error), i float64, f float64, t *testing.T) {
	res, err := fn(i)
	checkResult(res, err, getFunctionName(fn), f, t)
}

func TestPercentileMethods(t *testing.T) {
	assertPercentiles(data1.Percentile, 75, 4.6, t)
	assertPercentiles(data1.PercentileNearestRank, 75, 4.2, t)

}

func assertOtherDataMethods(fn func(d Float64Data) (float64, error), d Float64Data, f float64, t *testing.T) {
	res, err := fn(d)
	checkResult(res, err, getFunctionName(fn), f, t)
}

func TestOtherDataMethods(t *testing.T) {
	assertOtherDataMethods(data1.Correlation, data2, 0.20875473597605448, t)
	assertOtherDataMethods(data1.Pearson, data2, 0.20875473597605448, t)
	assertOtherDataMethods(data1.Midhinge, data2, -0.42500000000000004, t)
	assertOtherDataMethods(data1.Trimean, data2, 0.5375, t)
	assertOtherDataMethods(data1.Covariance, data2, 7.3814215535714265, t)
	assertOtherDataMethods(data1.CovariancePopulation, data2, 6.458743859374998, t)
}

func TestSampleMethod(t *testing.T) {
	// Test Sample method
	_, err := data1.Sample(5, true)
	if err != nil {
		t.Errorf("%s returned an error", getFunctionName(data1.Sample))
	}
}

func TestQuartileMethods(t *testing.T) {
	// Test QuartileOutliers method
	_, err := data1.QuartileOutliers()
	if err != nil {
		t.Errorf("%s returned an error", getFunctionName(data1.QuartileOutliers))
	}

	// Test Quartile method
	_, err = data1.Quartile(data2)
	if err != nil {
		t.Errorf("%s returned an error", getFunctionName(data1.Quartile))
	}
}

// Here we show the regular way of doing it
// with a plain old slice of float64s
func BenchmarkRegularAPI(b *testing.B) {
	for i := 0; i < b.N; i++ {
		data := []float64{-10, -7, -3.11, 5, 1.1, 2, 3, 4.20, 5, 18}
		Min(data)
		Max(data)
		Sum(data)
		Mean(data)
		Median(data)
		Mode(data)
	}
}

// Here's where things get interesting
// and we start to use the included
// Float64Data type and methods
func BenchmarkMethodsAPI(b *testing.B) {
	for i := 0; i < b.N; i++ {
		data := Float64Data{-10, -7, -3.11, 5, 1.1, 2, 3, 4.20, 5, 18}
		data.Min()
		data.Max()
		data.Sum()
		data.Mean()
		data.Median()
		data.Mode()
	}
}
<|endoftext|>


<commit_before>// +build !appengine

package goji

import (
	"flag"
	"log"
	"net/http"
	"time"

	"github.com/zenazn/goji/bind"
	"github.com/zenazn/goji/graceful"
)

func init() {
	bind.WithFlag()
	if fl := log.Flags(); fl&log.Ltime != 0 {
		log.SetFlags(fl | log.Lmicroseconds)
	}
	graceful.DoubleKickWindow(2 * time.Second)
}

// Serve starts Goji using reasonable defaults.
func Serve() {
	if !flag.Parsed() {
		flag.Parse()
	}

	DefaultMux.Compile()
	// Install our handler at the root of the standard net/http default mux.
	// This allows packages like expvar to continue working as expected.
	http.Handle("/", DefaultMux)

	listener := bind.Default()
	log.Println("Starting Goji on", listener.Addr())

	graceful.HandleSignals()
	bind.Ready()
	graceful.PreHook(func() { log.Printf("Goji received signal, gracefully stopping") })
	graceful.PostHook(func() { log.Printf("Goji stopped") })

	err := graceful.Serve(listener, http.DefaultServeMux)

	if err != nil {
		log.Fatal(err)
	}

	graceful.Wait()
}
<commit_msg>Add goji.ServeTLS and goji.ServeListener functions.<commit_after>// +build !appengine

package goji

import (
	"crypto/tls"
	"flag"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/zenazn/goji/bind"
	"github.com/zenazn/goji/graceful"
)

func init() {
	bind.WithFlag()
	if fl := log.Flags(); fl&log.Ltime != 0 {
		log.SetFlags(fl | log.Lmicroseconds)
	}
	graceful.DoubleKickWindow(2 * time.Second)
}

// Serve starts Goji using reasonable defaults.
func Serve() {
	ServeListener(bind.Default())
}

// Like Serve, but enables TLS using the given config.
func ServeTLS(config *tls.Config) {
	ServeListener(tls.NewListener(bind.Default(), config))
}

// Like Serve, but runs Goji on top of an arbitrary net.Listener.
func ServeListener(listener net.Listener) {
	if !flag.Parsed() {
		flag.Parse()
	}

	DefaultMux.Compile()
	// Install our handler at the root of the standard net/http default mux.
	// This allows packages like expvar to continue working as expected.
	http.Handle("/", DefaultMux)

	log.Println("Starting Goji on", listener.Addr())

	graceful.HandleSignals()
	bind.Ready()
	graceful.PreHook(func() { log.Printf("Goji received signal, gracefully stopping") })
	graceful.PostHook(func() { log.Printf("Goji stopped") })

	err := graceful.Serve(listener, http.DefaultServeMux)

	if err != nil {
		log.Fatal(err)
	}

	graceful.Wait()
}
<|endoftext|>


<commit_before>package vtls

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"sync"
	"time"
)

var (
	day  = 24 * time.Hour
	year = 365 * day

	certTime = time.Minute
)

type certCache struct {
	m   map[string]*tls.Certificate
	mut *sync.RWMutex
	crt Certifier
}

func newCertCache(crt Certifier) *certCache {
	return &certCache{
		m:   map[string]*tls.Certificate{},
		mut: &sync.RWMutex{},
		crt: crt,
	}
}

func (cc *certCache) add(name string) (*tls.Certificate, error) {
	crt, err := cc.crt.Certify(name, certTime)
	if err != nil {
		return nil, err
	}

	cc.mut.Lock()
	cc.m[name] = &crt
	cc.mut.Unlock()
	return &crt, nil
}

func (cc *certCache) get(name string) (*tls.Certificate, error) {
	lkr := cc.mut.RLocker()
	lkr.Lock()

	if c, ok := cc.m[name]; ok {
		n := time.Now()
		if n.After(c.Leaf.NotBefore) && n.Before(c.Leaf.NotAfter) {
			lkr.Unlock()
			return c, nil
		}
	}
	lkr.Unlock()

	return cc.add(name)
}

// ListenAndServeTLS mostly mirrors the http.ListenAndServeTLS API, but
// generates the certificates for the server automatically via vault, with a
// short TTL. The function only needs an additional Certifier parameter which
// can generate signed certificates in order to work properly.
func ListenAndServeTLS(addr string, handler http.Handler, crt Certifier) error {
	certs := newCertCache(crt)

	tl, err := tls.Listen("tcp", addr, &tls.Config{
		GetCertificate: func(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if h.ServerName == "" {
				return nil, fmt.Errorf("Cannot generate certs without TLS SNI (no server name was indicated)")
			}
			return certs.get(h.ServerName)
		},
	})

	if err != nil {
		return err
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}

	return http.Serve(tl, handler)
}
<commit_msg>Add default ttl package variable<commit_after>package vtls

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"sync"
	"time"
)

var (
	day  = 24 * time.Hour
	year = 365 * day

	//DefaultTTL is the default TTL the library will request for certificates
	DefaultTTL = day
)

type certCache struct {
	m   map[string]*tls.Certificate
	mut *sync.RWMutex
	crt Certifier
	ttl time.Duration
}

func newCertCache(crt Certifier) *certCache {
	return &certCache{
		m:   map[string]*tls.Certificate{},
		mut: &sync.RWMutex{},
		crt: crt,
		ttl: DefaultTTL,
	}
}

func (cc *certCache) add(name string) (*tls.Certificate, error) {
	crt, err := cc.crt.Certify(name, cc.ttl)
	if err != nil {
		return nil, err
	}

	cc.mut.Lock()
	cc.m[name] = &crt
	cc.mut.Unlock()
	return &crt, nil
}

func (cc *certCache) get(name string) (*tls.Certificate, error) {
	lkr := cc.mut.RLocker()
	lkr.Lock()

	if c, ok := cc.m[name]; ok {
		n := time.Now()
		if n.After(c.Leaf.NotBefore) && n.Before(c.Leaf.NotAfter) {
			lkr.Unlock()
			return c, nil
		}
	}
	lkr.Unlock()

	return cc.add(name)
}

// ListenAndServeTLS mostly mirrors the http.ListenAndServeTLS API, but
// generates the certificates for the server automatically via vault, with a
// short TTL. The function only needs an additional Certifier parameter which
// can generate signed certificates in order to work properly.
func ListenAndServeTLS(addr string, handler http.Handler, crt Certifier) error {
	certs := newCertCache(crt)

	tl, err := tls.Listen("tcp", addr, &tls.Config{
		GetCertificate: func(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if h.ServerName == "" {
				return nil, fmt.Errorf("Cannot generate certs without TLS SNI (no server name was indicated)")
			}
			return certs.get(h.ServerName)
		},
	})

	if err != nil {
		return err
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}

	return http.Serve(tl, handler)
}
<|endoftext|>


<commit_before>package volumewatcher

import (
	"testing"
	"time"

	memdb "github.com/hashicorp/go-memdb"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/stretchr/testify/require"
)

// TestVolumeWatch_EnableDisable tests the watcher registration logic that needs
// to happen during leader step-up/step-down
func TestVolumeWatch_EnableDisable(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockRPCServer{}
	srv.state = state.TestStateStore(t)
	index := uint64(100)

	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")
	watcher.SetEnabled(true, srv.State())

	plugin := mock.CSIPlugin()
	node := testNode(plugin, srv.State())
	alloc := mock.Alloc()
	alloc.ClientStatus = structs.AllocClientStatusComplete
	vol := testVolume(plugin, alloc, node.ID)

	index++
	err := srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	claim := &structs.CSIVolumeClaim{
		Mode:  structs.CSIVolumeClaimGC,
		State: structs.CSIVolumeClaimStateNodeDetached,
	}
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)
	require.Eventually(func() bool {
		return 1 == len(watcher.watchers)
	}, time.Second, 10*time.Millisecond)

	watcher.SetEnabled(false, srv.State())
	require.Equal(0, len(watcher.watchers))
}

// TestVolumeWatch_Checkpoint tests the checkpointing of progress across
// leader leader step-up/step-down
func TestVolumeWatch_Checkpoint(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockRPCServer{}
	srv.state = state.TestStateStore(t)
	index := uint64(100)

	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")

	plugin := mock.CSIPlugin()
	node := testNode(plugin, srv.State())
	alloc := mock.Alloc()
	alloc.ClientStatus = structs.AllocClientStatusComplete
	vol := testVolume(plugin, alloc, node.ID)

	watcher.SetEnabled(true, srv.State())

	index++
	err := srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	// we should get or start up a watcher when we get an update for
	// the volume from the state store
	require.Eventually(func() bool {
		return 1 == len(watcher.watchers)
	}, time.Second, 10*time.Millisecond)

	// step-down (this is sync, but step-up is async)
	watcher.SetEnabled(false, srv.State())
	require.Equal(0, len(watcher.watchers))

	// step-up again
	watcher.SetEnabled(true, srv.State())
	require.Eventually(func() bool {
		return 1 == len(watcher.watchers) &&
			!watcher.watchers[vol.ID+vol.Namespace].isRunning()
	}, time.Second, 10*time.Millisecond)
}

// TestVolumeWatch_StartStop tests the start and stop of the watcher when
// it receives notifcations and has completed its work
func TestVolumeWatch_StartStop(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockStatefulRPCServer{}
	srv.state = state.TestStateStore(t)
	index := uint64(100)
	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")

	watcher.SetEnabled(true, srv.State())
	require.Equal(0, len(watcher.watchers))

	plugin := mock.CSIPlugin()
	node := testNode(plugin, srv.State())
	alloc1 := mock.Alloc()
	alloc1.ClientStatus = structs.AllocClientStatusRunning
	alloc2 := mock.Alloc()
	alloc2.Job = alloc1.Job
	alloc2.ClientStatus = structs.AllocClientStatusRunning
	index++
	err := srv.State().UpsertJob(structs.MsgTypeTestSetup, index, alloc1.Job)
	require.NoError(err)
	index++
	err = srv.State().UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1, alloc2})
	require.NoError(err)

	// register a volume
	vol := testVolume(plugin, alloc1, node.ID)
	index++
	err = srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	// assert we get a watcher; there are no claims so it should immediately stop
	require.Eventually(func() bool {
		return 1 == len(watcher.watchers) &&
			!watcher.watchers[vol.ID+vol.Namespace].isRunning()
	}, time.Second*2, 10*time.Millisecond)

	// claim the volume for both allocs
	claim := &structs.CSIVolumeClaim{
		AllocationID: alloc1.ID,
		NodeID:       node.ID,
		Mode:         structs.CSIVolumeClaimRead,
	}
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)
	claim.AllocationID = alloc2.ID
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)

	// reap the volume and assert nothing has happened
	claim = &structs.CSIVolumeClaim{
		AllocationID: alloc1.ID,
		NodeID:       node.ID,
	}
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)

	ws := memdb.NewWatchSet()
	vol, _ = srv.State().CSIVolumeByID(ws, vol.Namespace, vol.ID)
	require.Equal(2, len(vol.ReadAllocs))

	// alloc becomes terminal
	alloc1.ClientStatus = structs.AllocClientStatusComplete
	index++
	err = srv.State().UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1})
	require.NoError(err)
	index++
	claim.State = structs.CSIVolumeClaimStateReadyToFree
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)

	// 1 claim has been released and watcher stops
	require.Eventually(func() bool {
		ws := memdb.NewWatchSet()
		vol, _ := srv.State().CSIVolumeByID(ws, vol.Namespace, vol.ID)
		return len(vol.ReadAllocs) == 1 && len(vol.PastClaims) == 0
	}, time.Second*2, 10*time.Millisecond)

	require.Eventually(func() bool {
		return !watcher.watchers[vol.ID+vol.Namespace].isRunning()
	}, time.Second*5, 10*time.Millisecond)
}

// TestVolumeWatch_RegisterDeregister tests the start and stop of
// watchers around registration
func TestVolumeWatch_RegisterDeregister(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockStatefulRPCServer{}
	srv.state = state.TestStateStore(t)

	index := uint64(100)

	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")

	watcher.SetEnabled(true, srv.State())
	require.Equal(0, len(watcher.watchers))

	plugin := mock.CSIPlugin()
	alloc := mock.Alloc()
	alloc.ClientStatus = structs.AllocClientStatusComplete

	// register a volume without claims
	vol := mock.CSIVolume(plugin)
	index++
	err := srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	// watcher should be started but immediately stopped
	require.Eventually(func() bool {
		return 1 == len(watcher.watchers)
	}, time.Second, 10*time.Millisecond)

	require.False(watcher.watchers[vol.ID+vol.Namespace].isRunning())
}
<commit_msg>volumewatcher: fix test data race.<commit_after>package volumewatcher

import (
	"testing"
	"time"

	memdb "github.com/hashicorp/go-memdb"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/stretchr/testify/require"
)

// TestVolumeWatch_EnableDisable tests the watcher registration logic that needs
// to happen during leader step-up/step-down
func TestVolumeWatch_EnableDisable(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockRPCServer{}
	srv.state = state.TestStateStore(t)
	index := uint64(100)

	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")
	watcher.SetEnabled(true, srv.State())

	plugin := mock.CSIPlugin()
	node := testNode(plugin, srv.State())
	alloc := mock.Alloc()
	alloc.ClientStatus = structs.AllocClientStatusComplete
	vol := testVolume(plugin, alloc, node.ID)

	index++
	err := srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	claim := &structs.CSIVolumeClaim{
		Mode:  structs.CSIVolumeClaimGC,
		State: structs.CSIVolumeClaimStateNodeDetached,
	}
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)
	require.Eventually(func() bool {
		watcher.wlock.RLock()
		defer watcher.wlock.RUnlock()
		return 1 == len(watcher.watchers)
	}, time.Second, 10*time.Millisecond)

	watcher.SetEnabled(false, nil)
	require.Equal(0, len(watcher.watchers))
}

// TestVolumeWatch_Checkpoint tests the checkpointing of progress across
// leader leader step-up/step-down
func TestVolumeWatch_Checkpoint(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockRPCServer{}
	srv.state = state.TestStateStore(t)
	index := uint64(100)

	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")

	plugin := mock.CSIPlugin()
	node := testNode(plugin, srv.State())
	alloc := mock.Alloc()
	alloc.ClientStatus = structs.AllocClientStatusComplete
	vol := testVolume(plugin, alloc, node.ID)

	watcher.SetEnabled(true, srv.State())

	index++
	err := srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	// we should get or start up a watcher when we get an update for
	// the volume from the state store
	require.Eventually(func() bool {
		watcher.wlock.RLock()
		defer watcher.wlock.RUnlock()
		return 1 == len(watcher.watchers)
	}, time.Second, 10*time.Millisecond)

	// step-down (this is sync, but step-up is async)
	watcher.SetEnabled(false, nil)
	require.Equal(0, len(watcher.watchers))

	// step-up again
	watcher.SetEnabled(true, srv.State())
	require.Eventually(func() bool {
		watcher.wlock.RLock()
		defer watcher.wlock.RUnlock()
		return 1 == len(watcher.watchers) &&
			!watcher.watchers[vol.ID+vol.Namespace].isRunning()
	}, time.Second, 10*time.Millisecond)
}

// TestVolumeWatch_StartStop tests the start and stop of the watcher when
// it receives notifcations and has completed its work
func TestVolumeWatch_StartStop(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockStatefulRPCServer{}
	srv.state = state.TestStateStore(t)
	index := uint64(100)
	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")

	watcher.SetEnabled(true, srv.State())
	require.Equal(0, len(watcher.watchers))

	plugin := mock.CSIPlugin()
	node := testNode(plugin, srv.State())
	alloc1 := mock.Alloc()
	alloc1.ClientStatus = structs.AllocClientStatusRunning
	alloc2 := mock.Alloc()
	alloc2.Job = alloc1.Job
	alloc2.ClientStatus = structs.AllocClientStatusRunning
	index++
	err := srv.State().UpsertJob(structs.MsgTypeTestSetup, index, alloc1.Job)
	require.NoError(err)
	index++
	err = srv.State().UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1, alloc2})
	require.NoError(err)

	// register a volume
	vol := testVolume(plugin, alloc1, node.ID)
	index++
	err = srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	// assert we get a watcher; there are no claims so it should immediately stop
	require.Eventually(func() bool {
		watcher.wlock.RLock()
		defer watcher.wlock.RUnlock()
		return 1 == len(watcher.watchers) &&
			!watcher.watchers[vol.ID+vol.Namespace].isRunning()
	}, time.Second*2, 10*time.Millisecond)

	// claim the volume for both allocs
	claim := &structs.CSIVolumeClaim{
		AllocationID: alloc1.ID,
		NodeID:       node.ID,
		Mode:         structs.CSIVolumeClaimRead,
	}
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)
	claim.AllocationID = alloc2.ID
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)

	// reap the volume and assert nothing has happened
	claim = &structs.CSIVolumeClaim{
		AllocationID: alloc1.ID,
		NodeID:       node.ID,
	}
	index++
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)

	ws := memdb.NewWatchSet()
	vol, _ = srv.State().CSIVolumeByID(ws, vol.Namespace, vol.ID)
	require.Equal(2, len(vol.ReadAllocs))

	// alloc becomes terminal
	alloc1.ClientStatus = structs.AllocClientStatusComplete
	index++
	err = srv.State().UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1})
	require.NoError(err)
	index++
	claim.State = structs.CSIVolumeClaimStateReadyToFree
	err = srv.State().CSIVolumeClaim(index, vol.Namespace, vol.ID, claim)
	require.NoError(err)

	// 1 claim has been released and watcher stops
	require.Eventually(func() bool {
		ws := memdb.NewWatchSet()
		vol, _ := srv.State().CSIVolumeByID(ws, vol.Namespace, vol.ID)
		return len(vol.ReadAllocs) == 1 && len(vol.PastClaims) == 0
	}, time.Second*2, 10*time.Millisecond)

	require.Eventually(func() bool {
		watcher.wlock.RLock()
		defer watcher.wlock.RUnlock()
		return !watcher.watchers[vol.ID+vol.Namespace].isRunning()
	}, time.Second*5, 10*time.Millisecond)
}

// TestVolumeWatch_RegisterDeregister tests the start and stop of
// watchers around registration
func TestVolumeWatch_RegisterDeregister(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	srv := &MockStatefulRPCServer{}
	srv.state = state.TestStateStore(t)

	index := uint64(100)

	watcher := NewVolumesWatcher(testlog.HCLogger(t), srv, "")

	watcher.SetEnabled(true, srv.State())
	require.Equal(0, len(watcher.watchers))

	plugin := mock.CSIPlugin()
	alloc := mock.Alloc()
	alloc.ClientStatus = structs.AllocClientStatusComplete

	// register a volume without claims
	vol := mock.CSIVolume(plugin)
	index++
	err := srv.State().CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(err)

	// watcher should be started but immediately stopped
	require.Eventually(func() bool {
		watcher.wlock.RLock()
		defer watcher.wlock.RUnlock()
		return 1 == len(watcher.watchers)
	}, time.Second, 10*time.Millisecond)

	require.False(watcher.watchers[vol.ID+vol.Namespace].isRunning())
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"gopkg.in/resty.v0"

	"github.com/boltdb/bolt"
	"github.com/codegangsta/cli"
)

var Client *resty.Client

// initCommon provides common startup initialization
func initCommon(c *cli.Context) {
	var (
		err  error
		resp *resty.Response
	)
	if err = configSetup(c); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read the configuration: "+
			"%s\n", err.Error())
		os.Exit(1)
	}

	// open database
	if err = store.Open(
		Cfg.Run.PathBoltDB,
		os.FileMode(uint32(Cfg.Run.ModeBoltDB)),
		&bolt.Options{Timeout: Cfg.Run.TimeoutBoltDB},
	); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to open database: %s\n", err)
		os.Exit(1)
	}

	// ensure database content structure is in place
	if err = store.EnsureBuckets(); err != nil {
		fmt.Fprintf(os.Stderr, "Database bucket error: %s\n", err)
		os.Exit(1)
	}

	// setup our REST client
	Client = resty.New().SetRESTMode().
		SetTimeout(Cfg.Run.TimeoutResty).
		SetDisableWarn(true).
		SetHeader(`User-Agent`, `somaadm 0.4.8`).
		SetHostURL(Cfg.Run.SomaAPI.String()).
		SetContentLength(true)

	if Cfg.Run.SomaAPI.Scheme == `https` {
		Client = Client.SetRootCertificate(Cfg.Run.CertPath)
	}

	// check configured API
	if resp, err = Client.R().Head(`/`); err != nil {
		fmt.Fprintf(os.Stderr, "Error tasting the API endpoint: %s\n",
			err.Error())
	} else if resp.StatusCode() != 204 {
		fmt.Fprintf(os.Stderr, "Error, API Url returned %d instead of 204."+
			" Sure this is SOMA?\n", resp.StatusCode())
		os.Exit(1)
	}

	// check who we talked to
	if resp.Header().Get(`X-Powered-By`) != `SOMA Configuration System` {
		fmt.Fprintf(os.Stderr, `Just FYI, at the end of that API URL`+
			` is not SOMA`)
		os.Exit(1)
	}
}

// boottime is the pre-run target for bootstrapping SOMA
func boottime(action cli.ActionFunc) cli.ActionFunc {
	return func(c *cli.Context) error {
		initCommon(c)

		return action(c)
	}
}

// runtime is the regular pre-run target
func runtime(action cli.ActionFunc) cli.ActionFunc {
	return func(c *cli.Context) error {
		initCommon(c)

		return action(c)
	}
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Update somaadm TLS configuration<commit_after>package main

import (
	"crypto/tls"
	"fmt"
	"os"

	"gopkg.in/resty.v0"

	"github.com/boltdb/bolt"
	"github.com/codegangsta/cli"
)

var Client *resty.Client

// initCommon provides common startup initialization
func initCommon(c *cli.Context) {
	var (
		err     error
		resp    *resty.Response
		session tls.ClientSessionCache
	)
	if err = configSetup(c); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read the configuration: "+
			"%s\n", err.Error())
		os.Exit(1)
	}

	// open database
	if err = store.Open(
		Cfg.Run.PathBoltDB,
		os.FileMode(uint32(Cfg.Run.ModeBoltDB)),
		&bolt.Options{Timeout: Cfg.Run.TimeoutBoltDB},
	); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to open database: %s\n", err)
		os.Exit(1)
	}

	// ensure database content structure is in place
	if err = store.EnsureBuckets(); err != nil {
		fmt.Fprintf(os.Stderr, "Database bucket error: %s\n", err)
		os.Exit(1)
	}

	// setup our REST client
	Client = resty.New().SetRESTMode().
		SetTimeout(Cfg.Run.TimeoutResty).
		SetDisableWarn(true).
		SetHeader(`User-Agent`, `somaadm 0.4.8`).
		SetHostURL(Cfg.Run.SomaAPI.String()).
		SetContentLength(true)

	if Cfg.Run.SomaAPI.Scheme == `https` {
		session = tls.NewLRUClientSessionCache(64)

		// SetTLSClientConfig replaces, SetRootCertificate updates the
		// tls configuration - option ordering is important
		Client = Client.SetTLSClientConfig(&tls.Config{
			ClientSessionCache: session,
			MinVersion:         tls.VersionTLS12,
			MaxVersion:         tls.VersionTLS12,
			CipherSuites:       []uint16{tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256},
		}).SetRootCertificate(Cfg.Run.CertPath)
	}

	// check configured API
	if resp, err = Client.R().Head(`/`); err != nil {
		fmt.Fprintf(os.Stderr, "Error tasting the API endpoint: %s\n",
			err.Error())
	} else if resp.StatusCode() != 204 {
		fmt.Fprintf(os.Stderr, "Error, API Url returned %d instead of 204."+
			" Sure this is SOMA?\n", resp.StatusCode())
		os.Exit(1)
	}

	// check who we talked to
	if resp.Header().Get(`X-Powered-By`) != `SOMA Configuration System` {
		fmt.Fprintf(os.Stderr, `Just FYI, at the end of that API URL`+
			` is not SOMA`)
		os.Exit(1)
	}
}

// boottime is the pre-run target for bootstrapping SOMA
func boottime(action cli.ActionFunc) cli.ActionFunc {
	return func(c *cli.Context) error {
		initCommon(c)

		return action(c)
	}
}

// runtime is the regular pre-run target
func runtime(action cli.ActionFunc) cli.ActionFunc {
	return func(c *cli.Context) error {
		initCommon(c)

		return action(c)
	}
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"testing"

	"database/sql"

	"code.google.com/p/go.net/context"
	"github.com/piotrkowalczuk/charon"
	"github.com/piotrkowalczuk/mnemosyne"
	"github.com/piotrkowalczuk/mnemosyne/mnemosynetest"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/stretchr/testify/mock"
)

func TestHandler(t *testing.T) {
	var (
		id  int64
		ctx context.Context
		err error
		act *actor
		tkn mnemosyne.Token
	)

	Convey("retrieveActor", t, func() {
		userRepositoryMock := &MockUserRepository{}
		permissionRepositoryMock := &MockPermissionRepository{}
		sessionMock := &mnemosynetest.Mnemosyne{}
		h := &handler{
			session: sessionMock,
		}
		h.repository.user = userRepositoryMock
		h.repository.permission = permissionRepositoryMock

		Convey("As unauthenticated user", func() {
			ctx = context.Background()
			sessionMock.On("FromContext", mock.Anything).
				Once().
				Return(nil, errors.New("mnemosyned: test error"))

			Convey("should return an error", func() {
				act, err = h.retrieveActor(ctx)

				So(err, ShouldNotBeNil)
				So(act, ShouldBeNil)
			})
		})
		Convey("As authenticated user", func() {
			id = 7856282
			tkn = mnemosyne.NewToken([]byte("0000000001"), []byte("hash"))
			ctx = mnemosyne.NewTokenContext(context.Background(), tkn)
			sessionMock.On("FromContext", ctx).
				Once().
				Return(&mnemosyne.Session{
				Token:     &tkn,
				SubjectId: charon.SubjectIDFromInt64(id).String(),
			}, nil)

			Convey("When user exists", func() {
				userRepositoryMock.On("FindOneByID", id).
					Once().
					Return(&userEntity{ID: id}, nil)

				Convey("And it has some permissions", func() {
					permissionRepositoryMock.On("FindByUserID", id).
						Once().
						Return([]*permissionEntity{
						{
							Subsystem: charon.PermissionCanRetrieve.Subsystem(),
							Module:    charon.PermissionCanRetrieve.Module(),
							Action:    charon.PermissionCanRetrieve.Action(),
						},
						{
							Subsystem: charon.UserCanRetrieveAsOwner.Subsystem(),
							Module:    charon.UserCanRetrieveAsOwner.Module(),
							Action:    charon.UserCanRetrieveAsOwner.Action(),
						},
					}, nil)

					Convey("Then it should be retrieved without any error", func() {
						act, err = h.retrieveActor(ctx)

						So(err, ShouldBeNil)
						So(act, ShouldNotBeNil)
						So(act.user.ID, ShouldEqual, id)
						So(act.permissions, ShouldHaveLength, 2)
					})
				})
				Convey("And it has no permissions", func() {
					permissionRepositoryMock.On("FindByUserID", id).
						Once().
						Return(nil, sql.ErrNoRows)

					Convey("Then it should be retrieved without any error", func() {
						act, err = h.retrieveActor(ctx)

						So(err, ShouldBeNil)
						So(act, ShouldNotBeNil)
						So(act.user.ID, ShouldEqual, id)
						So(act.permissions, ShouldHaveLength, 0)
					})
				})
			})
		})
	})
}
<commit_msg>handler_test import fix<commit_after>package main

import (
	"database/sql"
	"errors"
	"testing"

	"github.com/piotrkowalczuk/charon"
	"github.com/piotrkowalczuk/mnemosyne"
	"github.com/piotrkowalczuk/mnemosyne/mnemosynetest"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/stretchr/testify/mock"
	"golang.org/x/net/context"
)

func TestHandler(t *testing.T) {
	var (
		id  int64
		ctx context.Context
		err error
		act *actor
		tkn mnemosyne.Token
	)

	Convey("retrieveActor", t, func() {
		userRepositoryMock := &MockUserRepository{}
		permissionRepositoryMock := &MockPermissionRepository{}
		sessionMock := &mnemosynetest.Mnemosyne{}
		h := &handler{
			session: sessionMock,
		}
		h.repository.user = userRepositoryMock
		h.repository.permission = permissionRepositoryMock

		Convey("As unauthenticated user", func() {
			ctx = context.Background()
			sessionMock.On("FromContext", mock.Anything).
				Once().
				Return(nil, errors.New("mnemosyned: test error"))

			Convey("should return an error", func() {
				act, err = h.retrieveActor(ctx)

				So(err, ShouldNotBeNil)
				So(act, ShouldBeNil)
			})
		})
		Convey("As authenticated user", func() {
			id = 7856282
			tkn = mnemosyne.NewToken([]byte("0000000001"), []byte("hash"))
			ctx = mnemosyne.NewTokenContext(context.Background(), tkn)
			sessionMock.On("FromContext", ctx).
				Once().
				Return(&mnemosyne.Session{
				Token:     &tkn,
				SubjectId: charon.SubjectIDFromInt64(id).String(),
			}, nil)

			Convey("When user exists", func() {
				userRepositoryMock.On("FindOneByID", id).
					Once().
					Return(&userEntity{ID: id}, nil)

				Convey("And it has some permissions", func() {
					permissionRepositoryMock.On("FindByUserID", id).
						Once().
						Return([]*permissionEntity{
						{
							Subsystem: charon.PermissionCanRetrieve.Subsystem(),
							Module:    charon.PermissionCanRetrieve.Module(),
							Action:    charon.PermissionCanRetrieve.Action(),
						},
						{
							Subsystem: charon.UserCanRetrieveAsOwner.Subsystem(),
							Module:    charon.UserCanRetrieveAsOwner.Module(),
							Action:    charon.UserCanRetrieveAsOwner.Action(),
						},
					}, nil)

					Convey("Then it should be retrieved without any error", func() {
						act, err = h.retrieveActor(ctx)

						So(err, ShouldBeNil)
						So(act, ShouldNotBeNil)
						So(act.user.ID, ShouldEqual, id)
						So(act.permissions, ShouldHaveLength, 2)
					})
				})
				Convey("And it has no permissions", func() {
					permissionRepositoryMock.On("FindByUserID", id).
						Once().
						Return(nil, sql.ErrNoRows)

					Convey("Then it should be retrieved without any error", func() {
						act, err = h.retrieveActor(ctx)

						So(err, ShouldBeNil)
						So(act, ShouldNotBeNil)
						So(act.user.ID, ShouldEqual, id)
						So(act.permissions, ShouldHaveLength, 0)
					})
				})
			})
		})
	})
}
<|endoftext|>


<commit_before>package orgtracker

import (
	sl "github.com/djosephsen/slacker/slackerlib"
	"encoding/json"
	"regexp"
	"fmt"
	"strings"
)

const ORGS string = `OT:ORGS`

type Org struct{
	Name string
	Members	map[string]sl.User
}

func getOrgs(bot *sl.Sbot) (map[string]Org, error){
// load in the orgs struct from brain
 	orgs := make(map[string]Org)
	brain := *bot.Brain
	if orgson, err := brain.Get(ORGS); err != nil{
		//if notFound,_ := regexp.MatchString(
		if err.Error() != `key OT:ORGS was not found`{
			return orgs, err
		}
	}else if err := json.Unmarshal(orgson, &orgs); err != nil{
		return orgs, err
	}
	return orgs, nil
}

func setOrgs(bot *sl.Sbot, orgs map[string]Org) error{
// write the orgs struct to brain
	brain := *bot.Brain
	if orgson, err := json.Marshal(&orgs); err !=nil{
		return err
	}else if err = brain.Set(ORGS, orgson); err != nil{
		return err
	}
	return nil
}

var QueryPeeps = sl.MessageHandler{
	Name: `OrgTracker, Query Peeps`,
	Usage:`"<botname> who is from <orgname>" returns the users from <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)who is from (\w+)`,
	Run:	func(e *sl.Event, match []string){
		orgName := match[1]
		orgID := strings.ToLower(orgName)
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}

		if org, exists := orgs[orgID]; exists{
			users := fmt.Sprintf("Members from %s:\n",orgName)
			for _,peep := range org.Members{
				users = fmt.Sprintf("%s\n%s", users, peep.Name)
			}
			e.Respond(users)
			return
		}else{
			e.Respond(fmt.Sprintf("sorry, no org called: %s", orgName))
			return
		}
	},
}

var ManageOrg = sl.MessageHandler{
	Name: `OrgTracker: Manage Org`,
	Usage:`"<botname> (add|delete) org <orgname>" adds or deletes an org called <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)(add|delete) org (\w+)`,
	Run:	func(e *sl.Event, match []string){
		cmd:=match[1]
		orgName:=match[2]
		orgID := strings.ToLower(orgName)
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}

		if isAdd,_ := regexp.MatchString( `(?i)add`, cmd); isAdd{
			newOrg:=Org{
				Name:	strings.ToLower(orgName),
				Members: make(map[string]sl.User),
			}
			orgs[orgID] = newOrg
			if err := setOrgs(e.Sbot, orgs); err != nil{
				e.Reply(fmt.Sprintf("sorry I couldn't add %s because my brain says: %s", orgName, err))
				return
			}
			e.Reply(fmt.Sprintf("sure thing. Org: %s added", orgName))
		}

		if isDelete,_ := regexp.MatchString( `(?i)delete`, cmd); isDelete{
			delete(orgs,orgID)
			if err := setOrgs(e.Sbot, orgs); err != nil{
				e.Reply(fmt.Sprintf("blerg, couldn't delete %s because my brain says: %s", orgName, err))
				return
			}

			e.Reply(fmt.Sprintf("sure thing. Org: %s deleted", orgName))
		}
	},
}

var JoinOrg = sl.MessageHandler{
	Name: `OrgTracker: Join Org`,
	Usage:`"<botname> (join|leave) org <orgname>" adds or removes you to/from <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)(join|leave) org (\w+)`,
	Run:	func(e *sl.Event, match []string){
		cmd:=match[1]
		orgName:=match[2]
		orgID := strings.ToLower(orgName)

		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}

		if isJoin,_ := regexp.MatchString( `(?i)join`, cmd); isJoin{
			if org, exists := orgs[orgID]; exists{
				user:=e.Sbot.Meta.GetUser(e.User)
				if _,exists := org.Members[user.ID]; exists{
					e.Reply(fmt.Sprintf("user: %s already belongs to %s (sorry)",user.Name, orgName))
				}else{
					org.Members[user.ID] = *user
					if err := setOrgs(e.Sbot, orgs); err != nil{
						e.Reply(fmt.Sprintf("derp.. I couldn't add %s. Brain trouble: %s", user.Name, err))
						return
					}
					e.Reply(fmt.Sprintf("OK! user: %s now belongs to %s",user.Name, orgName))
					return
				}
			}else{
				e.Reply(fmt.Sprintf("No such org: %s (sorry)",orgName))
			}
		}
		if isLeave,_ := regexp.MatchString( `(?i)leave`, cmd); isLeave{
			if org,exists := orgs[orgID]; exists{
				user:=e.Sbot.Meta.GetUser(e.User)
				delete (org.Members,user.ID)
				e.Reply(fmt.Sprintf("OK! %s is no longer in %s",user.Name, orgName))
			}else{
				e.Reply(fmt.Sprintf("No such org: %s (sorry)",orgName))
			}
		}
	},
}
<commit_msg>added listing<commit_after>package orgtracker

import (
	sl "github.com/djosephsen/slacker/slackerlib"
	"encoding/json"
	"regexp"
	"fmt"
	"strings"
)

const ORGS string = `OT:ORGS`

type Org struct{
	Name string
	Members	map[string]sl.User
}

func getOrgs(bot *sl.Sbot) (map[string]Org, error){
// load in the orgs struct from brain
 	orgs := make(map[string]Org)
	brain := *bot.Brain
	if orgson, err := brain.Get(ORGS); err != nil{
		//if notFound,_ := regexp.MatchString(
		if err.Error() != `key OT:ORGS was not found`{
			return orgs, err
		}
	}else if err := json.Unmarshal(orgson, &orgs); err != nil{
		return orgs, err
	}
	return orgs, nil
}

func setOrgs(bot *sl.Sbot, orgs map[string]Org) error{
// write the orgs struct to brain
	brain := *bot.Brain
	if orgson, err := json.Marshal(&orgs); err !=nil{
		return err
	}else if err = brain.Set(ORGS, orgson); err != nil{
		return err
	}
	return nil
}

var QueryPeeps = sl.MessageHandler{
	Name: `OrgTracker, Query Peeps`,
	Usage:`"<botname> who is from <orgname>" returns the users from <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)who is from (\w+)`,
	Run:	func(e *sl.Event, match []string){
		orgName := match[1]
		orgID := strings.ToLower(orgName)
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}

		if org, exists := orgs[orgID]; exists{
			users := fmt.Sprintf("Members from %s:\n",orgName)
			for _,peep := range org.Members{
				users = fmt.Sprintf("%s\n%s", users, peep.Name)
			}
			e.Respond(users)
			return
		}else{
			e.Respond(fmt.Sprintf("sorry, no org called: %s", orgName))
			return
		}
	},
}

var ManageOrg = sl.MessageHandler{
	Name: `OrgTracker: Manage Org`,
	Usage:`"<botname> (add|delete) org <orgname>" adds or deletes an org called <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)(add|delete) org (\w+)`,
	Run:	func(e *sl.Event, match []string){
		cmd:=match[1]
		orgName:=match[2]
		orgID := strings.ToLower(orgName)
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}

		if isAdd,_ := regexp.MatchString( `(?i)add`, cmd); isAdd{
			newOrg:=Org{
				Name:	strings.ToLower(orgName),
				Members: make(map[string]sl.User),
			}
			orgs[orgID] = newOrg
			if err := setOrgs(e.Sbot, orgs); err != nil{
				e.Reply(fmt.Sprintf("sorry I couldn't add %s because my brain says: %s", orgName, err))
				return
			}
			e.Reply(fmt.Sprintf("sure thing. Org: %s added", orgName))
		}

		if isDelete,_ := regexp.MatchString( `(?i)delete`, cmd); isDelete{
			delete(orgs,orgID)
			if err := setOrgs(e.Sbot, orgs); err != nil{
				e.Reply(fmt.Sprintf("blerg, couldn't delete %s because my brain says: %s", orgName, err))
				return
			}

			e.Reply(fmt.Sprintf("sure thing. Org: %s deleted", orgName))
		}
	},
}

var JoinOrg = sl.MessageHandler{
	Name: `OrgTracker: Join Org`,
	Usage:`"<botname> (join|leave) org <orgname>" adds or removes you to/from <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)(join|leave) org (\w+)`,
	Run:	func(e *sl.Event, match []string){
		cmd:=match[1]
		orgName:=match[2]
		orgID := strings.ToLower(orgName)
		var org Org
		var exists bool

		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}

		if isJoin,_ := regexp.MatchString( `(?i)join`, cmd); isJoin{
			if org, exists = orgs[orgID]; !exists{
				e.Reply(fmt.Sprintf("(Creating new org: %s first)",orgName))
				org=Org{
					Name:	strings.ToLower(orgName),
					Members: make(map[string]sl.User),
				}
				orgs[orgID]=org
			}
			user:=e.Sbot.Meta.GetUser(e.User)
			if _,exists := org.Members[user.ID]; exists{
				e.Reply(fmt.Sprintf("user: %s already belongs to %s (sorry)",user.Name, orgName))
			}else{
				org.Members[user.ID] = *user
				if err := setOrgs(e.Sbot, orgs); err != nil{
					e.Reply(fmt.Sprintf("derp.. I couldn't add %s. Brain trouble: %s", user.Name, err))
					return
				}
				e.Reply(fmt.Sprintf("OK! user: %s now belongs to %s",user.Name, orgName))
				return
			}
		}
		if isLeave,_ := regexp.MatchString( `(?i)leave`, cmd); isLeave{
			if org,exists := orgs[orgID]; exists{
				user:=e.Sbot.Meta.GetUser(e.User)
				delete (org.Members,user.ID)
				e.Reply(fmt.Sprintf("OK! %s is no longer in %s",user.Name, orgName))
			}else{
				e.Reply(fmt.Sprintf("No such org: %s (sorry)",orgName))
			}
		}
	},
}

var ListOrgs = sl.MessageHandler{
	Name: `OrgTracker: List Orgs`,
	Usage:`"<botname> list orgs" lists the available orgs`,
	Method: `RESPOND`,
	Pattern: `(?i)list orgs`,
	Run:	func(e *sl.Event, match []string){
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}
		reply:=`Existing Orgs:`
		for orgid,org := range orgs{
			reply=fmt.Sprintf("%s\n%s, (%d members)",reply, orgid, len(org.Members))
		}
		e.Reply(reply)
	},
}
<|endoftext|>


<commit_before>package orgtracker

import (
	sl "github.com/djosephsen/slacker/slackerlib"
	"encoding/json"
	"regexp"
	"fmt"
	"strings"
)

const ORGS string = `OT:ORGS`

type Org struct{
	Name string
	Members	map[string]sl.User
}

var OrgTracker = sl.MessageHandler{
	Name: `OrgTracker`,
	Usage:`<botname> [add|delete|join|leave|list] org <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)(add|delete|join|leave|list) org *(\w*)`,
	Run:	func(e *sl.Event, match []string){
		var orgName, orgID string
		cmd := match[1]
		if len(match) >= 2{
			orgName = match[2]
			orgID = strings.ToLower(orgName)
		}
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}
		if isAdd,_ := regexp.MatchString( `(?i)add`, cmd); isAdd{
			if err := addOrg(e.Sbot,orgs,orgID); err!=nil{
				e.Reply(fmt.Sprintf("sorry I couldn't add %s because: %s", orgName, err))
				return
			}
			e.Reply(fmt.Sprintf("sure thing. Org: %s added", orgName))
			return
		}else if isDel,_ := regexp.MatchString( `(?i)delete`, cmd); isDel{
			if err := deleteOrg(e.Sbot,orgs,orgID); err!=nil{
				e.Reply(fmt.Sprintf("gar. %s", err))
				return
			}
			e.Reply(fmt.Sprintf("Ok. Org: %s deleted", orgName))
			return
		}else if isJoin,_ := regexp.MatchString( `(?i)join`, cmd); isJoin{
			if _, exists := orgs[orgID]; !exists{
				e.Reply(fmt.Sprintf("(Creating new org: %s first)",orgName))
				if err := addOrg(e.Sbot,orgs,orgID); err!=nil{
					e.Reply(fmt.Sprintf("sorry I couldn't add %s because: %s", orgName, err))
					return
				}
			}
			if err := joinOrg(e,orgs,orgID); err!=nil{
					e.Reply(fmt.Sprintf("derp. %s", err))
					return
			}
			user:=e.Sbot.Meta.GetUser(e.User)
			e.Reply(fmt.Sprintf("OK! user: %s now belongs to %s",user.Name, orgName))
			return
		}else if isLeave,_ := regexp.MatchString( `(?i)leave`, cmd); isLeave{
			if err := leaveOrg(e,orgs,orgID); err!=nil{
					e.Reply(fmt.Sprintf("bleh. %s", err))
					return
			}
			user:=e.Sbot.Meta.GetUser(e.User)
			e.Reply(fmt.Sprintf("OK! removed user: %s from org: %s",user.Name, orgName))
			return
		}else if isList,_ := regexp.MatchString( `(?i)list`, cmd); isList{
			if reply, err := listOrg(orgs,orgID); err!=nil{
				e.Reply(fmt.Sprintf("sorry. %s", err))
				return
			}else{
				e.Reply(reply)
				return
			}
		}
	},
}

func listOrg(orgs map[string]Org, orgID string) (string, error){
	if orgID != ``{
		if org, exists := orgs[orgID]; exists{
			users := fmt.Sprintf("Members from %s:\n",orgID)
			for _,peep := range org.Members{
				users = fmt.Sprintf("%s\n%s", users, peep.Name)
			}
			return users, nil
		}else{
			return ``,fmt.Errorf("no org called: %s", orgID)
		}
	}else{
		reply:=`Existing Orgs:`
		for orgid,org := range orgs{
			reply=fmt.Sprintf("%s\n%s, (%d members)",reply, orgid, len(org.Members))
		}
		return reply,nil
	}
}

func addOrg(bot *sl.Sbot, orgs map[string]Org, orgID string) error{
	newOrg:=Org{
		Name:	orgID,
		Members: make(map[string]sl.User),
	}
	orgs[orgID] = newOrg
	if err := setOrgs(bot, orgs); err != nil{
		return fmt.Errorf("I couldn't add %s because my brain says: %s", orgID, err)
	}
	return nil
}

func deleteOrg(bot *sl.Sbot, orgs map[string]Org, orgID string) error{
	delete(orgs,orgID)
	if err := setOrgs(bot, orgs); err != nil{
		return fmt.Errorf("I couldn't delete %s because my brain says: %s", orgID, err)
	}
	return nil
}

func joinOrg(e *sl.Event, orgs map[string]Org, orgID string) error{
	user := e.Sbot.Meta.GetUser(e.User)
	org := orgs[orgID]
	if _,exists := org.Members[user.ID]; exists{
		return fmt.Errorf("user: %s already belongs to %s (sorry)",user.Name, orgID)
	}else{
		org.Members[user.ID] = *user
		if err := setOrgs(e.Sbot, orgs); err != nil{
			return fmt.Errorf("I couldn't add %s. Brain trouble: %s", user.Name, err)
		}
		return nil
	}
}

func leaveOrg(e *sl.Event, orgs map[string]Org, orgID string) error{
	if org,exists := orgs[orgID]; exists{
		user:=e.Sbot.Meta.GetUser(e.User)
		delete (org.Members,user.ID)
		if err := setOrgs(e.Sbot, orgs); err != nil{
			return fmt.Errorf("I couldn't delete %s. Brain trouble: %s", user.Name, err)
		}
		return nil
	}else{
		return fmt.Errorf("No such org: %s (sorry)",orgID)
	}
}

func getOrgs(bot *sl.Sbot) (map[string]Org, error){
// load in the orgs struct from brain
 	orgs := make(map[string]Org)
	brain := *bot.Brain
	if orgson, err := brain.Get(ORGS); err != nil{
		//if notFound,_ := regexp.MatchString(
		if err.Error() != `key OT:ORGS was not found`{
			return orgs, err
		}
	}else if err := json.Unmarshal(orgson, &orgs); err != nil{
		return orgs, err
	}
	return orgs, nil
}

func setOrgs(bot *sl.Sbot, orgs map[string]Org) error{
// write the orgs struct to brain
	brain := *bot.Brain
	if orgson, err := json.Marshal(&orgs); err !=nil{
		return err
	}else if err = brain.Set(ORGS, orgson); err != nil{
		return err
	}
	return nil
}

var WhoIsFrom = sl.MessageHandler{
	Name: `OrgTracker: Who-IS-FROM`,
	Usage:`"<botname> who is from <org>" :: lists members who are from org`,
	Method: `RESPOND`,
	Pattern: `(?i)who is from (\w*)`,
	Run:	func(e *sl.Event, match []string){
		orgName := match[1]
		orgID := strings.ToLower(orgName)
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}
		if reply, err := listOrg(orgs,orgID); err!=nil{
			e.Reply(fmt.Sprintf("sorry. %s", err))
			return
		}else{
			e.Reply(reply)
			return
		}
	},
}

<commit_msg>plural in regex<commit_after>package orgtracker

import (
	sl "github.com/djosephsen/slacker/slackerlib"
	"encoding/json"
	"regexp"
	"fmt"
	"strings"
)

const ORGS string = `OT:ORGS`

type Org struct{
	Name string
	Members	map[string]sl.User
}

var OrgTracker = sl.MessageHandler{
	Name: `OrgTracker`,
	Usage:`<botname> [add|delete|join|leave|list] org <orgname>`,
	Method: `RESPOND`,
	Pattern: `(?i)(add|delete|join|leave|list) orgs* *(\w*)`,
	Run:	func(e *sl.Event, match []string){
		var orgName, orgID string
		cmd := match[1]
		if len(match) >= 2{
			orgName = match[2]
			orgID = strings.ToLower(orgName)
		}
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}
		if isAdd,_ := regexp.MatchString( `(?i)add`, cmd); isAdd{
			if err := addOrg(e.Sbot,orgs,orgID); err!=nil{
				e.Reply(fmt.Sprintf("sorry I couldn't add %s because: %s", orgName, err))
				return
			}
			e.Reply(fmt.Sprintf("sure thing. Org: %s added", orgName))
			return
		}else if isDel,_ := regexp.MatchString( `(?i)delete`, cmd); isDel{
			if err := deleteOrg(e.Sbot,orgs,orgID); err!=nil{
				e.Reply(fmt.Sprintf("gar. %s", err))
				return
			}
			e.Reply(fmt.Sprintf("Ok. Org: %s deleted", orgName))
			return
		}else if isJoin,_ := regexp.MatchString( `(?i)join`, cmd); isJoin{
			if _, exists := orgs[orgID]; !exists{
				e.Reply(fmt.Sprintf("(Creating new org: %s first)",orgName))
				if err := addOrg(e.Sbot,orgs,orgID); err!=nil{
					e.Reply(fmt.Sprintf("sorry I couldn't add %s because: %s", orgName, err))
					return
				}
			}
			if err := joinOrg(e,orgs,orgID); err!=nil{
					e.Reply(fmt.Sprintf("derp. %s", err))
					return
			}
			user:=e.Sbot.Meta.GetUser(e.User)
			e.Reply(fmt.Sprintf("OK! user: %s now belongs to %s",user.Name, orgName))
			return
		}else if isLeave,_ := regexp.MatchString( `(?i)leave`, cmd); isLeave{
			if err := leaveOrg(e,orgs,orgID); err!=nil{
					e.Reply(fmt.Sprintf("bleh. %s", err))
					return
			}
			user:=e.Sbot.Meta.GetUser(e.User)
			e.Reply(fmt.Sprintf("OK! removed user: %s from org: %s",user.Name, orgName))
			return
		}else if isList,_ := regexp.MatchString( `(?i)list`, cmd); isList{
			if reply, err := listOrg(orgs,orgID); err!=nil{
				e.Reply(fmt.Sprintf("sorry. %s", err))
				return
			}else{
				e.Reply(reply)
				return
			}
		}
	},
}

func listOrg(orgs map[string]Org, orgID string) (string, error){
	if orgID != ``{
		if org, exists := orgs[orgID]; exists{
			users := fmt.Sprintf("Members from %s:\n",orgID)
			for _,peep := range org.Members{
				users = fmt.Sprintf("%s\n%s", users, peep.Name)
			}
			return users, nil
		}else{
			return ``,fmt.Errorf("no org called: %s", orgID)
		}
	}else{
		reply:=`Existing Orgs:`
		for orgid,org := range orgs{
			reply=fmt.Sprintf("%s\n%s, (%d members)",reply, orgid, len(org.Members))
		}
		return reply,nil
	}
}

func addOrg(bot *sl.Sbot, orgs map[string]Org, orgID string) error{
	newOrg:=Org{
		Name:	orgID,
		Members: make(map[string]sl.User),
	}
	orgs[orgID] = newOrg
	if err := setOrgs(bot, orgs); err != nil{
		return fmt.Errorf("I couldn't add %s because my brain says: %s", orgID, err)
	}
	return nil
}

func deleteOrg(bot *sl.Sbot, orgs map[string]Org, orgID string) error{
	delete(orgs,orgID)
	if err := setOrgs(bot, orgs); err != nil{
		return fmt.Errorf("I couldn't delete %s because my brain says: %s", orgID, err)
	}
	return nil
}

func joinOrg(e *sl.Event, orgs map[string]Org, orgID string) error{
	user := e.Sbot.Meta.GetUser(e.User)
	org := orgs[orgID]
	if _,exists := org.Members[user.ID]; exists{
		return fmt.Errorf("user: %s already belongs to %s (sorry)",user.Name, orgID)
	}else{
		org.Members[user.ID] = *user
		if err := setOrgs(e.Sbot, orgs); err != nil{
			return fmt.Errorf("I couldn't add %s. Brain trouble: %s", user.Name, err)
		}
		return nil
	}
}

func leaveOrg(e *sl.Event, orgs map[string]Org, orgID string) error{
	if org,exists := orgs[orgID]; exists{
		user:=e.Sbot.Meta.GetUser(e.User)
		delete (org.Members,user.ID)
		if err := setOrgs(e.Sbot, orgs); err != nil{
			return fmt.Errorf("I couldn't delete %s. Brain trouble: %s", user.Name, err)
		}
		return nil
	}else{
		return fmt.Errorf("No such org: %s (sorry)",orgID)
	}
}

func getOrgs(bot *sl.Sbot) (map[string]Org, error){
// load in the orgs struct from brain
 	orgs := make(map[string]Org)
	brain := *bot.Brain
	if orgson, err := brain.Get(ORGS); err != nil{
		//if notFound,_ := regexp.MatchString(
		if err.Error() != `key OT:ORGS was not found`{
			return orgs, err
		}
	}else if err := json.Unmarshal(orgson, &orgs); err != nil{
		return orgs, err
	}
	return orgs, nil
}

func setOrgs(bot *sl.Sbot, orgs map[string]Org) error{
// write the orgs struct to brain
	brain := *bot.Brain
	if orgson, err := json.Marshal(&orgs); err !=nil{
		return err
	}else if err = brain.Set(ORGS, orgson); err != nil{
		return err
	}
	return nil
}

var WhoIsFrom = sl.MessageHandler{
	Name: `OrgTracker: Who-IS-FROM`,
	Usage:`"<botname> who is from <org>" :: lists members who are from org`,
	Method: `RESPOND`,
	Pattern: `(?i)who is from (\w*)`,
	Run:	func(e *sl.Event, match []string){
		orgName := match[1]
		orgID := strings.ToLower(orgName)
		orgs, err := getOrgs(e.Sbot)
		if err != nil{
			e.Respond(fmt.Sprintf("ack! I couldn't load my orgs struct! %s", err))
			sl.Logger.Debug(err)
			return
		}
		if reply, err := listOrg(orgs,orgID); err!=nil{
			e.Reply(fmt.Sprintf("sorry. %s", err))
			return
		}else{
			e.Reply(reply)
			return
		}
	},
}

<|endoftext|>


<commit_before>package gocql

import (
	"sync"
	"time"
)

type ExecutableQuery interface {
	execute(conn *Conn) *Iter
	attempt(keyspace string, end, start time.Time, iter *Iter, host *HostInfo)
	retryPolicy() RetryPolicy
	speculativeExecutionPolicy() SpeculativeExecutionPolicy
	GetRoutingKey() ([]byte, error)
	Keyspace() string
	Cancel()
	IsIdempotent() bool
	RetryableQuery
}

type queryExecutor struct {
	pool   *policyConnPool
	policy HostSelectionPolicy
}

type queryResponse struct {
	iter *Iter
	err  error
}

func (q *queryExecutor) attemptQuery(qry ExecutableQuery, conn *Conn) *Iter {
	start := time.Now()
	iter := qry.execute(conn)
	end := time.Now()

	qry.attempt(q.pool.keyspace, end, start, iter, conn.host)

	return iter
}

func (q *queryExecutor) executeQuery(qry ExecutableQuery) (*Iter, error) {

	// check if the query is not marked as idempotent, if
	// it is, we force the policy to NonSpeculative
	sp := qry.speculativeExecutionPolicy()
	if !qry.IsIdempotent() {
		sp = NonSpeculativeExecution{}
	}

	results := make(chan queryResponse, 1)
	stop := make(chan struct{})
	defer close(stop)
	var specWG sync.WaitGroup

	// Launch the main execution
	specWG.Add(1)
	go q.run(qry, &specWG, results, stop)

	// The speculative executions are launched _in addition_ to the main
	// execution, on a timer. So Speculation{2} would make 3 executions running
	// in total.
	go func() {
		// Handle the closing of the resources. We do it here because it's
		// right after we finish launching executions. Otherwise clearing the
		// wait group is complicated.
		defer func() {
			specWG.Wait()
			close(results)
		}()

		// setup a ticker
		ticker := time.NewTicker(sp.Delay())
		defer ticker.Stop()

		for i := 0; i < sp.Attempts(); i++ {
			select {
			case <-ticker.C:
				// Launch the additional execution
				specWG.Add(1)
				go q.run(qry, &specWG, results, stop)
			case <-qry.GetContext().Done():
				// not starting additional executions
				return
			case <-stop:
				// not starting additional executions
				return
			}
		}
	}()

	res := <-results
	if res.iter == nil && res.err == nil {
		// if we're here, the results channel was closed, so no more hosts
		return nil, ErrNoConnections
	}
	return res.iter, res.err
}

func (q *queryExecutor) run(qry ExecutableQuery, specWG *sync.WaitGroup, results chan queryResponse, stop chan struct{}) {
	// Handle the wait group
	defer specWG.Done()

	hostIter := q.policy.Pick(qry)
	selectedHost := hostIter()
	rt := qry.retryPolicy()

	var iter *Iter
	for selectedHost != nil {
		host := selectedHost.Info()
		if host == nil || !host.IsUp() {
			continue
		}

		pool, ok := q.pool.getPool(host)
		if !ok {
			continue
		}

		conn := pool.Pick()
		if conn == nil {
			continue
		}

		select {
		case <-stop:
			// stop this execution and return
			return
		default:
			// Run the query
			iter = q.attemptQuery(qry, conn)
			iter.host = selectedHost.Info()
			// Update host
			selectedHost.Mark(iter.err)

			// Exit if the query was successful
			// or no retry policy defined or retry attempts were reached
			if iter.err == nil || rt == nil || !rt.Attempt(qry) {
				results <- queryResponse{iter: iter}
				return
			}

			// If query is unsuccessful, check the error with RetryPolicy to retry
			switch rt.GetRetryType(iter.err) {
			case Retry:
				// retry on the same host
				continue
			case Rethrow:
				results <- queryResponse{err: iter.err}
				return
			case Ignore:
				results <- queryResponse{iter: iter}
				return
			case RetryNextHost:
				// retry on the next host
				selectedHost = hostIter()
				if selectedHost == nil {
					results <- queryResponse{iter: iter}
					return
				}
				continue
			default:
				// Undefined? Return nil and error, this will panic in the requester
				results <- queryResponse{iter: nil, err: ErrUnknownRetryType}
				return
			}
		}

	}
	// All hosts are exhausted, return nothing
}
<commit_msg>Prevent infinite loop (#1219)<commit_after>package gocql

import (
	"sync"
	"time"
)

type ExecutableQuery interface {
	execute(conn *Conn) *Iter
	attempt(keyspace string, end, start time.Time, iter *Iter, host *HostInfo)
	retryPolicy() RetryPolicy
	speculativeExecutionPolicy() SpeculativeExecutionPolicy
	GetRoutingKey() ([]byte, error)
	Keyspace() string
	Cancel()
	IsIdempotent() bool
	RetryableQuery
}

type queryExecutor struct {
	pool   *policyConnPool
	policy HostSelectionPolicy
}

type queryResponse struct {
	iter *Iter
	err  error
}

func (q *queryExecutor) attemptQuery(qry ExecutableQuery, conn *Conn) *Iter {
	start := time.Now()
	iter := qry.execute(conn)
	end := time.Now()

	qry.attempt(q.pool.keyspace, end, start, iter, conn.host)

	return iter
}

func (q *queryExecutor) executeQuery(qry ExecutableQuery) (*Iter, error) {

	// check if the query is not marked as idempotent, if
	// it is, we force the policy to NonSpeculative
	sp := qry.speculativeExecutionPolicy()
	if !qry.IsIdempotent() {
		sp = NonSpeculativeExecution{}
	}

	results := make(chan queryResponse, 1)
	stop := make(chan struct{})
	defer close(stop)
	var specWG sync.WaitGroup

	// Launch the main execution
	specWG.Add(1)
	go q.run(qry, &specWG, results, stop)

	// The speculative executions are launched _in addition_ to the main
	// execution, on a timer. So Speculation{2} would make 3 executions running
	// in total.
	go func() {
		// Handle the closing of the resources. We do it here because it's
		// right after we finish launching executions. Otherwise clearing the
		// wait group is complicated.
		defer func() {
			specWG.Wait()
			close(results)
		}()

		// setup a ticker
		ticker := time.NewTicker(sp.Delay())
		defer ticker.Stop()

		for i := 0; i < sp.Attempts(); i++ {
			select {
			case <-ticker.C:
				// Launch the additional execution
				specWG.Add(1)
				go q.run(qry, &specWG, results, stop)
			case <-qry.GetContext().Done():
				// not starting additional executions
				return
			case <-stop:
				// not starting additional executions
				return
			}
		}
	}()

	res := <-results
	if res.iter == nil && res.err == nil {
		// if we're here, the results channel was closed, so no more hosts
		return nil, ErrNoConnections
	}
	return res.iter, res.err
}

func (q *queryExecutor) run(qry ExecutableQuery, specWG *sync.WaitGroup, results chan queryResponse, stop chan struct{}) {
	// Handle the wait group
	defer specWG.Done()

	hostIter := q.policy.Pick(qry)
	selectedHost := hostIter()
	rt := qry.retryPolicy()

	var iter *Iter
	for selectedHost != nil {
		host := selectedHost.Info()
		if host == nil || !host.IsUp() {
			selectedHost = hostIter()
			continue
		}

		pool, ok := q.pool.getPool(host)
		if !ok {
			selectedHost = hostIter()
			continue
		}

		conn := pool.Pick()
		if conn == nil {
			selectedHost = hostIter()
			continue
		}

		select {
		case <-stop:
			// stop this execution and return
			return
		default:
			// Run the query
			iter = q.attemptQuery(qry, conn)
			iter.host = selectedHost.Info()
			// Update host
			selectedHost.Mark(iter.err)

			// Exit if the query was successful
			// or no retry policy defined or retry attempts were reached
			if iter.err == nil || rt == nil || !rt.Attempt(qry) {
				results <- queryResponse{iter: iter}
				return
			}

			// If query is unsuccessful, check the error with RetryPolicy to retry
			switch rt.GetRetryType(iter.err) {
			case Retry:
				// retry on the same host
				continue
			case Rethrow:
				results <- queryResponse{err: iter.err}
				return
			case Ignore:
				results <- queryResponse{iter: iter}
				return
			case RetryNextHost:
				// retry on the next host
				selectedHost = hostIter()
				if selectedHost == nil {
					results <- queryResponse{iter: iter}
					return
				}
				continue
			default:
				// Undefined? Return nil and error, this will panic in the requester
				results <- queryResponse{iter: nil, err: ErrUnknownRetryType}
				return
			}
		}

	}
	// All hosts are exhausted, return nothing
}
<|endoftext|>