<commit_before>// Copyright 2015-2016, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cfgmodel

import (
	"fmt"
	"sync"
	"time"

	"github.com/corestoreio/csfw/config"
	"github.com/corestoreio/csfw/config/cfgpath"
	"github.com/corestoreio/csfw/config/element"
	"github.com/corestoreio/csfw/config/source"
	"github.com/corestoreio/csfw/store/scope"
	"github.com/corestoreio/csfw/util/errors"
)

// PkgBackend used for embedding in the PkgBackend type in each package. The
// mutex protects the init process.
type PkgBackend struct {
	sync.Mutex
}

// optionBox groups different types into one struct to allow multiple option
// functions applied to many different types within this package.
type optionBox struct {
	*baseValue
	*Obscure
	*StringCSV
	*CSV
	*IntCSV
}

// BaseOption as an optional argument for the New*() functions. These options
// will be applied to the underlying unexported baseValue type.
type Option func(*optionBox) error

// WithFieldFromSectionSlice extracts the element.Field from the global
// PackageConfiguration for retrieving the default value of a underlying type
// and for scope permission checking.
func WithFieldFromSectionSlice(cfgStruct element.SectionSlice) Option {
	return func(b *optionBox) error {
		f, _, err := cfgStruct.FindField(b.route)
		if err != nil {
			return errors.Wrap(err, "[cfgmodel] cfgStruct.FindField")
		}
		b.Field = &f
		return nil
	}
}

// WithField adds a Field to the model.
func WithField(f *element.Field) Option {
	return func(b *optionBox) error {
		b.Field = f
		return nil
	}
}

// WithSource sets a source slice for Options() and validation.
func WithSource(vl source.Slice) Option {
	return func(b *optionBox) error {
		b.Source = vl
		return nil
	}
}

// WithSourceByString sets a source slice for Options() and validation.
// Wrapper for source.NewByString
func WithSourceByString(pairs ...string) Option {
	return func(b *optionBox) (err error) {
		b.Source, err = source.NewByString(pairs...)
		return errors.Wrap(err, "[cfgmodel] WithSourceByString")
	}
}

// WithSourceByInt sets a source slice for Options() and validation.
// Wrapper for source.NewByInt
func WithSourceByInt(vli source.Ints) Option {
	return func(b *optionBox) error {
		b.Source = source.NewByInt(vli)
		return nil
	}
}

// baseValue defines the path in the "core_config_data" table like a/b/c. All
// other types in this package inherits from this path type.
type baseValue struct {
	// contains the path like web/cors/exposed_headers but has no scope
	route cfgpath.Route

	// Field is used for scope permission checks and retrieving the default
	// value. A nil field gets ignored. Field will be set through the option
	// functions at creation time of the struct.
	Field *element.Field

	// Source are all available options aka SourceModel in Mage slang. This
	// slice is also used for validation to get and write the correct values.
	// Validation gets triggered only when the slice has been set. The Options()
	// function will be used to access this slice.
	Source source.Slice
	// OptionError might contain an error when an applied function option returns an
	// error. Only used in the function MustNewValue()
	OptionError error
}

// NewValue creates a new baseValue type and the error gets packed into the field
// OptionError which you can check.
func NewValue(path string, opts ...Option) baseValue {
	b := baseValue{
		route: cfgpath.NewRoute(path),
	}
	b.OptionError = (&b).Option(opts...)
	return b
}

// Option sets the options and returns the last set previous option
func (bv *baseValue) Option(opts ...Option) error {
	ob := &optionBox{
		baseValue: bv,
	}
	for _, o := range opts {
		if err := o(ob); err != nil {
			return errors.Wrap(err, "[cfgmodel] baseValue.Option")
		}
	}
	bv = ob.baseValue
	return nil
}

func (bv baseValue) hasField() bool {
	return bv.Field != nil && bv.Field.ID.IsEmpty() == false
}

// Write writes a value v to the config.Writer without checking if the value has
// changed. Checks if the Scope matches as defined in the non-nil
// ConfigStructure.
// Error behaviour: Unauthorized
func (bv baseValue) Write(w config.Writer, v interface{}, s scope.Scope, scopeID int64) error {
	if bv.hasField() {
		if false == bv.Field.Scopes.Has(s) {
			return errors.NewUnauthorizedf("[cfgmodel] Scope permission insufficient: Have '%s'; Want '%s'", s, bv.Field.Scopes)
		}
	}
	pp, err := bv.ToPath(s, scopeID)
	if err != nil {
		return errors.Wrap(err, "[cfgmodel] ToPath")
	}
	return w.Write(pp, v)
}

// String returns the stringyfied route
func (bv baseValue) String() string {
	return bv.route.String()
}

// ToPath creates a new cfgpath.Path bound to a scope. If the argument scope
// does not match the defined scope in the Field, the error behaviour
// Unauthorized gets returned.
func (bv baseValue) ToPath(s scope.Scope, scopeID int64) (cfgpath.Path, error) {

	if err := bv.inScope(s, scopeID); err != nil {
		return cfgpath.Path{}, errors.Wrap(err, "[cfgmodel] ToPath")
	}

	p, err := cfgpath.New(bv.route)
	if err != nil {
		return cfgpath.Path{}, errors.Wrapf(err, "[cfgmodel] cfgpath.New: %q", bv.route)
	}
	return p.Bind(s, scopeID), nil
}

// Route returns a copy of the underlying route.
func (bv baseValue) Route() cfgpath.Route {
	return bv.route.Clone()
}

// InScope checks if a field from a path is allowed for current scope.
// Returns nil on success.
// Error behaviour: Unauthorized
func (bv baseValue) InScope(sg scope.Scoper) (err error) {
	return bv.inScope(sg.Scope())
}

func (bv baseValue) inScope(s scope.Scope, _ int64) (err error) {
	if bv.hasField() && !bv.Field.Scopes.Has(s) {
		err = errors.NewUnauthorizedf("[cfgmodel] Scope permission insufficient: Have '%s'; Want '%s'", s, bv.Field.Scopes)
	}
	return
}

// Options returns a source model for all available options for a configuration
// value.
//
// Usually this function gets customized in a sub-type. Customization can have
// different arguments, etc but must always call this function to set source
// slice.
func (bv baseValue) Options() source.Slice {
	return bv.Source
}

// FQ generates a fully qualified configuration path. Example:
// general/country/allow would transform with StrScope scope.StrStores and
// storeID e.g. 4 into: stores/4/general/country/allow
func (bv baseValue) FQ(s scope.Scope, scopeID int64) (string, error) {
	p, err := bv.ToPath(s, scopeID)
	return p.String(), errors.Wrap(err, "[cfgmodel] ToPath")
}

// MustFQ same as FQ but panics on error. Please use only for testing.
func (bv baseValue) MustFQ(s scope.Scope, scopeID int64) string {
	p, err := bv.ToPath(s, scopeID)
	if err != nil {
		panic(err)
	}
	return p.String()
}

// ValidateString checks if string v is contained in Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateString(v string) (err error) {
	if bv.Source != nil && false == bv.Source.ContainsValString(v) {
		jv, jErr := bv.Source.ToJSON()
		if jErr != nil {
			return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
		}
		err = errors.NewNotValidf("[cfgmodel] The value '%s' cannot be found within the allowed Options():\n%s", v, jv)
	}
	return
}

// ValidateInt checks if int v is contained in non-nil Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateInt(v int) (err error) {
	if bv.Source != nil && false == bv.Source.ContainsValInt(v) {
		jv, jErr := bv.Source.ToJSON()
		if jErr != nil {
			return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
		}
		err = errors.NewNotValidf("[cfgmodel] The value '%d' cannot be found within the allowed Options():\n%s", v, jv)
	}
	return
}

// ValidateFloat64 checks if float64 v is contained in non-nil Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateFloat64(v float64) (err error) {
	if bv.Source != nil && false == bv.Source.ContainsValFloat64(v) {
		jv, jErr := bv.Source.ToJSON()
		if jErr != nil {
			return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
		}
		err = errors.NewNotValidf("[cfgmodel] The value '%.14f' cannot be found within the allowed Options():\n%s", v, jv)
	}
	return
}

// ValidateTime checks if time.Time v is contained in non-nil Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateTime(v time.Time) (err error) {
	// todo:
	//if bv.Source != nil && false == bv.Source.ContainsValFloat64(v) {
	//jv, jErr := bv.Source.ToJSON()
	//if jErr != nil {
	//	return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
	//}
	//err = errors.NewNotValidf("[cfgmodel] The value '%s' cannot be found within the allowed Options():\n%s", v, jv)
	//}
	return errors.NewNotValidf("[cfgmodel] @todo once someone requires this feature")
}
<commit_msg>config/cfgmodel: Add route to error message for better understanding<commit_after>// Copyright 2015-2016, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cfgmodel

import (
	"fmt"
	"sync"
	"time"

	"github.com/corestoreio/csfw/config"
	"github.com/corestoreio/csfw/config/cfgpath"
	"github.com/corestoreio/csfw/config/element"
	"github.com/corestoreio/csfw/config/source"
	"github.com/corestoreio/csfw/store/scope"
	"github.com/corestoreio/csfw/util/errors"
)

// PkgBackend used for embedding in the PkgBackend type in each package. The
// mutex protects the init process.
type PkgBackend struct {
	sync.Mutex
}

// optionBox groups different types into one struct to allow multiple option
// functions applied to many different types within this package.
type optionBox struct {
	*baseValue
	*Obscure
	*StringCSV
	*CSV
	*IntCSV
}

// BaseOption as an optional argument for the New*() functions. These options
// will be applied to the underlying unexported baseValue type.
type Option func(*optionBox) error

// WithFieldFromSectionSlice extracts the element.Field from the global
// PackageConfiguration for retrieving the default value of a underlying type
// and for scope permission checking.
func WithFieldFromSectionSlice(cfgStruct element.SectionSlice) Option {
	return func(b *optionBox) error {
		f, _, err := cfgStruct.FindField(b.route)
		if err != nil {
			return errors.Wrap(err, "[cfgmodel] cfgStruct.FindField")
		}
		b.Field = &f
		return nil
	}
}

// WithField adds a Field to the model.
func WithField(f *element.Field) Option {
	return func(b *optionBox) error {
		b.Field = f
		return nil
	}
}

// WithSource sets a source slice for Options() and validation.
func WithSource(vl source.Slice) Option {
	return func(b *optionBox) error {
		b.Source = vl
		return nil
	}
}

// WithSourceByString sets a source slice for Options() and validation.
// Wrapper for source.NewByString
func WithSourceByString(pairs ...string) Option {
	return func(b *optionBox) (err error) {
		b.Source, err = source.NewByString(pairs...)
		return errors.Wrap(err, "[cfgmodel] WithSourceByString")
	}
}

// WithSourceByInt sets a source slice for Options() and validation.
// Wrapper for source.NewByInt
func WithSourceByInt(vli source.Ints) Option {
	return func(b *optionBox) error {
		b.Source = source.NewByInt(vli)
		return nil
	}
}

// baseValue defines the path in the "core_config_data" table like a/b/c. All
// other types in this package inherits from this path type.
type baseValue struct {
	// contains the path like web/cors/exposed_headers but has no scope
	route cfgpath.Route

	// Field is used for scope permission checks and retrieving the default
	// value. A nil field gets ignored. Field will be set through the option
	// functions at creation time of the struct.
	Field *element.Field

	// Source are all available options aka SourceModel in Mage slang. This
	// slice is also used for validation to get and write the correct values.
	// Validation gets triggered only when the slice has been set. The Options()
	// function will be used to access this slice.
	Source source.Slice
	// OptionError might contain an error when an applied function option returns an
	// error. Only used in the function MustNewValue()
	OptionError error
}

// NewValue creates a new baseValue type and the error gets packed into the field
// OptionError which you can check.
func NewValue(path string, opts ...Option) baseValue {
	b := baseValue{
		route: cfgpath.NewRoute(path),
	}
	b.OptionError = (&b).Option(opts...)
	return b
}

// Option sets the options and returns the last set previous option
func (bv *baseValue) Option(opts ...Option) error {
	ob := &optionBox{
		baseValue: bv,
	}
	for _, o := range opts {
		if err := o(ob); err != nil {
			return errors.Wrap(err, "[cfgmodel] baseValue.Option")
		}
	}
	bv = ob.baseValue
	return nil
}

func (bv baseValue) hasField() bool {
	return bv.Field != nil && bv.Field.ID.IsEmpty() == false
}

// Write writes a value v to the config.Writer without checking if the value has
// changed. Checks if the Scope matches as defined in the non-nil
// ConfigStructure.
// Error behaviour: Unauthorized
func (bv baseValue) Write(w config.Writer, v interface{}, s scope.Scope, scopeID int64) error {
	pp, err := bv.ToPath(s, scopeID)
	if err != nil {
		return errors.Wrap(err, "[cfgmodel] baseValue.ToPath")
	}
	return w.Write(pp, v)
}

// String returns the stringyfied route
func (bv baseValue) String() string {
	return bv.route.String()
}

// ToPath creates a new cfgpath.Path bound to a scope. If the argument scope
// does not match the defined scope in the Field, the error behaviour
// Unauthorized gets returned.
func (bv baseValue) ToPath(s scope.Scope, scopeID int64) (cfgpath.Path, error) {

	if err := bv.inScope(s, scopeID); err != nil {
		return cfgpath.Path{}, errors.Wrap(err, "[cfgmodel] ToPath")
	}

	p, err := cfgpath.New(bv.route)
	if err != nil {
		return cfgpath.Path{}, errors.Wrapf(err, "[cfgmodel] cfgpath.New: %q", bv.route)
	}
	return p.Bind(s, scopeID), nil
}

// Route returns a copy of the underlying route.
func (bv baseValue) Route() cfgpath.Route {
	return bv.route.Clone()
}

// InScope checks if a field from a path is allowed for current scope.
// Returns nil on success.
// Error behaviour: Unauthorized
func (bv baseValue) InScope(sg scope.Scoper) (err error) {
	return bv.inScope(sg.Scope())
}

func (bv baseValue) inScope(s scope.Scope, _ int64) (err error) {
	if bv.hasField() && !bv.Field.Scopes.Has(s) {
		err = errors.NewUnauthorizedf("[cfgmodel] Scope permission insufficient: Have %q; Want %q; Route: %q", s, bv.Field.Scopes, bv)
	}
	return
}

// Options returns a source model for all available options for a configuration
// value.
//
// Usually this function gets customized in a sub-type. Customization can have
// different arguments, etc but must always call this function to set source
// slice.
func (bv baseValue) Options() source.Slice {
	return bv.Source
}

// FQ generates a fully qualified configuration path. Example:
// general/country/allow would transform with StrScope scope.StrStores and
// storeID e.g. 4 into: stores/4/general/country/allow
func (bv baseValue) FQ(s scope.Scope, scopeID int64) (string, error) {
	p, err := bv.ToPath(s, scopeID)
	return p.String(), errors.Wrap(err, "[cfgmodel] ToPath")
}

// MustFQ same as FQ but panics on error. Please use only for testing.
func (bv baseValue) MustFQ(s scope.Scope, scopeID int64) string {
	p, err := bv.ToPath(s, scopeID)
	if err != nil {
		panic(err)
	}
	return p.String()
}

// ValidateString checks if string v is contained in Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateString(v string) (err error) {
	if bv.Source != nil && false == bv.Source.ContainsValString(v) {
		jv, jErr := bv.Source.ToJSON()
		if jErr != nil {
			return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
		}
		err = errors.NewNotValidf("[cfgmodel] The value '%s' cannot be found within the allowed Options():\n%s", v, jv)
	}
	return
}

// ValidateInt checks if int v is contained in non-nil Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateInt(v int) (err error) {
	if bv.Source != nil && false == bv.Source.ContainsValInt(v) {
		jv, jErr := bv.Source.ToJSON()
		if jErr != nil {
			return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
		}
		err = errors.NewNotValidf("[cfgmodel] The value '%d' cannot be found within the allowed Options():\n%s", v, jv)
	}
	return
}

// ValidateFloat64 checks if float64 v is contained in non-nil Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateFloat64(v float64) (err error) {
	if bv.Source != nil && false == bv.Source.ContainsValFloat64(v) {
		jv, jErr := bv.Source.ToJSON()
		if jErr != nil {
			return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
		}
		err = errors.NewNotValidf("[cfgmodel] The value '%.14f' cannot be found within the allowed Options():\n%s", v, jv)
	}
	return
}

// ValidateTime checks if time.Time v is contained in non-nil Source source.Slice.
// Error behaviour: NotValid
func (bv baseValue) ValidateTime(v time.Time) (err error) {
	// todo:
	//if bv.Source != nil && false == bv.Source.ContainsValFloat64(v) {
	//jv, jErr := bv.Source.ToJSON()
	//if jErr != nil {
	//	return errors.NewFatal(err, fmt.Sprintf("[cfgmodel] Source: %#v", bv.Source))
	//}
	//err = errors.NewNotValidf("[cfgmodel] The value '%s' cannot be found within the allowed Options():\n%s", v, jv)
	//}
	return errors.NewNotValidf("[cfgmodel] @todo once someone requires this feature")
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"

	"github.com/zclconf/go-cty/cty"
	"github.com/zclconf/go-cty/cty/gocty"
	ctyjson "github.com/zclconf/go-cty/cty/json"

	"github.com/gruntwork-io/terragrunt/errors"
	"github.com/gruntwork-io/terragrunt/remote"
)

// Serialize TerragruntConfig struct to a cty Value that can be used to reference the attributes in other config. Note
// that we can't straight up convert the struct using cty tags due to differences in the desired representation.
// Specifically, we want to reference blocks by named attributes, but blocks are rendered to lists in the
// TerragruntConfig struct, so we need to do some massaging of the data to convert the list of blocks in to a map going
// from the block name label to the block value.
func terragruntConfigAsCty(config *TerragruntConfig) (cty.Value, error) {
	output := map[string]cty.Value{}

	// Convert attributes that are primitive types
	output["terraform_binary"] = gostringToCty(config.TerraformBinary)
	output["terraform_version_constraint"] = gostringToCty(config.TerraformVersionConstraint)
	output["terragrunt_version_constraint"] = gostringToCty(config.TerragruntVersionConstraint)
	output["download_dir"] = gostringToCty(config.DownloadDir)
	output["iam_role"] = gostringToCty(config.IamRole)
	output["skip"] = goboolToCty(config.Skip)

	terraformConfigCty, err := terraformConfigAsCty(config.Terraform)
	if err != nil {
		return cty.NilVal, err
	}
	if terraformConfigCty != cty.NilVal {
		output["terraform"] = terraformConfigCty
	}

	remoteStateCty, err := remoteStateAsCty(config.RemoteState)
	if err != nil {
		return cty.NilVal, err
	}
	if remoteStateCty != cty.NilVal {
		output["remote_state"] = remoteStateCty
	}

	dependenciesCty, err := gostructToCty(config.Dependencies)
	if err != nil {
		return cty.NilVal, err
	}
	if dependenciesCty != cty.NilVal {
		output["dependencies"] = dependenciesCty
	}

	if config.PreventDestroy != nil {
		output["prevent_destroy"] = goboolToCty(*config.PreventDestroy)
	}

	dependencyCty, err := dependencyBlocksAsCty(config.TerragruntDependencies)
	if err != nil {
		return cty.NilVal, err
	}
	if dependencyCty != cty.NilVal {
		output["dependency"] = dependencyCty
	}

	generateCty, err := gostructToCty(config.GenerateConfigs)
	if err != nil {
		return cty.NilVal, err
	}
	if generateCty != cty.NilVal {
		output["generate"] = generateCty
	}

	retryableCty, err := gostructToCty(config.RetryableErrors)
	if err != nil {
		return cty.NilVal, err
	}
	if retryableCty != cty.NilVal {
		output["retryable_errors"] = retryableCty
	}

	inputsCty, err := convertToCtyWithJson(config.Inputs)
	if err != nil {
		return cty.NilVal, err
	}
	if inputsCty != cty.NilVal {
		output["inputs"] = inputsCty
	}

	localsCty, err := convertToCtyWithJson(config.Locals)
	if err != nil {
		return cty.NilVal, err
	}
	if localsCty != cty.NilVal {
		output["locals"] = localsCty
	}

	return convertValuesMapToCtyVal(output)
}

// ctyTerraformConfig is an alternate representation of TerraformConfig that converts internal blocks into a map that
// maps the name to the underlying struct, as opposed to a list representation.
type ctyTerraformConfig struct {
	ExtraArgs   map[string]TerraformExtraArguments `cty:"extra_arguments"`
	Source      *string                            `cty:"source"`
	BeforeHooks map[string]Hook                    `cty:"before_hook"`
	AfterHooks  map[string]Hook                    `cty:"after_hook"`
}

// Serialize TerraformConfig to a cty Value, but with maps instead of lists for the blocks.
func terraformConfigAsCty(config *TerraformConfig) (cty.Value, error) {
	if config == nil {
		return cty.NilVal, nil
	}

	configCty := ctyTerraformConfig{
		Source:      config.Source,
		ExtraArgs:   map[string]TerraformExtraArguments{},
		BeforeHooks: map[string]Hook{},
		AfterHooks:  map[string]Hook{},
	}

	for _, arg := range config.ExtraArgs {
		configCty.ExtraArgs[arg.Name] = arg
	}
	for _, hook := range config.BeforeHooks {
		configCty.BeforeHooks[hook.Name] = hook
	}
	for _, hook := range config.AfterHooks {
		configCty.AfterHooks[hook.Name] = hook
	}

	return gostructToCty(configCty)
}

// Serialize RemoteState to a cty Value. We can't directly serialize the struct because `config` is an arbitrary
// interface whose type we do not know, so we have to do a hack to go through json.
func remoteStateAsCty(remoteState *remote.RemoteState) (cty.Value, error) {
	if remoteState == nil {
		return cty.NilVal, nil
	}

	output := map[string]cty.Value{}
	output["backend"] = gostringToCty(remoteState.Backend)
	output["disable_init"] = goboolToCty(remoteState.DisableInit)
	output["disable_dependency_optimization"] = goboolToCty(remoteState.DisableDependencyOptimization)

	generateCty, err := gostructToCty(remoteState.Generate)
	if err != nil {
		return cty.NilVal, err
	}
	output["generate"] = generateCty

	ctyJsonVal, err := convertToCtyWithJson(remoteState.Config)
	if err != nil {
		return cty.NilVal, err
	}
	output["config"] = ctyJsonVal

	return convertValuesMapToCtyVal(output)
}

// Serialize the list of dependency blocks to a cty Value as a map that maps the block names to the cty representation.
func dependencyBlocksAsCty(dependencyBlocks []Dependency) (cty.Value, error) {
	out := map[string]cty.Value{}
	for _, block := range dependencyBlocks {
		blockCty, err := gostructToCty(block)
		if err != nil {
			return cty.NilVal, err
		}
		out[block.Name] = blockCty
	}
	return convertValuesMapToCtyVal(out)
}

// Converts arbitrary go types that are json serializable to a cty Value by using json as an intermediary
// representation. This avoids the strict type nature of cty, where you need to know the output type beforehand to
// serialize to cty.
func convertToCtyWithJson(val interface{}) (cty.Value, error) {
	jsonBytes, err := json.Marshal(val)
	if err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	var ctyJsonVal ctyjson.SimpleJSONValue
	if err := ctyJsonVal.UnmarshalJSON(jsonBytes); err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	return ctyJsonVal.Value, nil
}

// Converts an arbitrary go struct that has cty tags to a cty Value.
func gostructToCty(val interface{}) (cty.Value, error) {
	ctyType, err := gocty.ImpliedType(val)
	if err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	ctyOut, err := gocty.ToCtyValue(val, ctyType)
	if err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	return ctyOut, nil
}

// Converts primitive go strings to a cty Value.
func gostringToCty(val string) cty.Value {
	ctyOut, err := gocty.ToCtyValue(val, cty.String)
	if err != nil {
		// Since we are converting primitive strings, we should never get an error in this conversion.
		panic(err)
	}
	return ctyOut
}

// Converts primitive go bools to a cty Value.
func goboolToCty(val bool) cty.Value {
	ctyOut, err := gocty.ToCtyValue(val, cty.Bool)
	if err != nil {
		// Since we are converting primitive bools, we should never get an error in this conversion.
		panic(err)
	}
	return ctyOut
}
<commit_msg>Rename the gostructToCty function in config_as_cty.go to goTypeToCty<commit_after>package config

import (
	"encoding/json"

	"github.com/zclconf/go-cty/cty"
	"github.com/zclconf/go-cty/cty/gocty"
	ctyjson "github.com/zclconf/go-cty/cty/json"

	"github.com/gruntwork-io/terragrunt/errors"
	"github.com/gruntwork-io/terragrunt/remote"
)

// Serialize TerragruntConfig struct to a cty Value that can be used to reference the attributes in other config. Note
// that we can't straight up convert the struct using cty tags due to differences in the desired representation.
// Specifically, we want to reference blocks by named attributes, but blocks are rendered to lists in the
// TerragruntConfig struct, so we need to do some massaging of the data to convert the list of blocks in to a map going
// from the block name label to the block value.
func terragruntConfigAsCty(config *TerragruntConfig) (cty.Value, error) {
	output := map[string]cty.Value{}

	// Convert attributes that are primitive types
	output["terraform_binary"] = gostringToCty(config.TerraformBinary)
	output["terraform_version_constraint"] = gostringToCty(config.TerraformVersionConstraint)
	output["terragrunt_version_constraint"] = gostringToCty(config.TerragruntVersionConstraint)
	output["download_dir"] = gostringToCty(config.DownloadDir)
	output["iam_role"] = gostringToCty(config.IamRole)
	output["skip"] = goboolToCty(config.Skip)

	terraformConfigCty, err := terraformConfigAsCty(config.Terraform)
	if err != nil {
		return cty.NilVal, err
	}
	if terraformConfigCty != cty.NilVal {
		output["terraform"] = terraformConfigCty
	}

	remoteStateCty, err := remoteStateAsCty(config.RemoteState)
	if err != nil {
		return cty.NilVal, err
	}
	if remoteStateCty != cty.NilVal {
		output["remote_state"] = remoteStateCty
	}

	dependenciesCty, err := goTypeToCty(config.Dependencies)
	if err != nil {
		return cty.NilVal, err
	}
	if dependenciesCty != cty.NilVal {
		output["dependencies"] = dependenciesCty
	}

	if config.PreventDestroy != nil {
		output["prevent_destroy"] = goboolToCty(*config.PreventDestroy)
	}

	dependencyCty, err := dependencyBlocksAsCty(config.TerragruntDependencies)
	if err != nil {
		return cty.NilVal, err
	}
	if dependencyCty != cty.NilVal {
		output["dependency"] = dependencyCty
	}

	generateCty, err := goTypeToCty(config.GenerateConfigs)
	if err != nil {
		return cty.NilVal, err
	}
	if generateCty != cty.NilVal {
		output["generate"] = generateCty
	}

	retryableCty, err := goTypeToCty(config.RetryableErrors)
	if err != nil {
		return cty.NilVal, err
	}
	if retryableCty != cty.NilVal {
		output["retryable_errors"] = retryableCty
	}

	inputsCty, err := convertToCtyWithJson(config.Inputs)
	if err != nil {
		return cty.NilVal, err
	}
	if inputsCty != cty.NilVal {
		output["inputs"] = inputsCty
	}

	localsCty, err := convertToCtyWithJson(config.Locals)
	if err != nil {
		return cty.NilVal, err
	}
	if localsCty != cty.NilVal {
		output["locals"] = localsCty
	}

	return convertValuesMapToCtyVal(output)
}

// ctyTerraformConfig is an alternate representation of TerraformConfig that converts internal blocks into a map that
// maps the name to the underlying struct, as opposed to a list representation.
type ctyTerraformConfig struct {
	ExtraArgs   map[string]TerraformExtraArguments `cty:"extra_arguments"`
	Source      *string                            `cty:"source"`
	BeforeHooks map[string]Hook                    `cty:"before_hook"`
	AfterHooks  map[string]Hook                    `cty:"after_hook"`
}

// Serialize TerraformConfig to a cty Value, but with maps instead of lists for the blocks.
func terraformConfigAsCty(config *TerraformConfig) (cty.Value, error) {
	if config == nil {
		return cty.NilVal, nil
	}

	configCty := ctyTerraformConfig{
		Source:      config.Source,
		ExtraArgs:   map[string]TerraformExtraArguments{},
		BeforeHooks: map[string]Hook{},
		AfterHooks:  map[string]Hook{},
	}

	for _, arg := range config.ExtraArgs {
		configCty.ExtraArgs[arg.Name] = arg
	}
	for _, hook := range config.BeforeHooks {
		configCty.BeforeHooks[hook.Name] = hook
	}
	for _, hook := range config.AfterHooks {
		configCty.AfterHooks[hook.Name] = hook
	}

	return goTypeToCty(configCty)
}

// Serialize RemoteState to a cty Value. We can't directly serialize the struct because `config` is an arbitrary
// interface whose type we do not know, so we have to do a hack to go through json.
func remoteStateAsCty(remoteState *remote.RemoteState) (cty.Value, error) {
	if remoteState == nil {
		return cty.NilVal, nil
	}

	output := map[string]cty.Value{}
	output["backend"] = gostringToCty(remoteState.Backend)
	output["disable_init"] = goboolToCty(remoteState.DisableInit)
	output["disable_dependency_optimization"] = goboolToCty(remoteState.DisableDependencyOptimization)

	generateCty, err := goTypeToCty(remoteState.Generate)
	if err != nil {
		return cty.NilVal, err
	}
	output["generate"] = generateCty

	ctyJsonVal, err := convertToCtyWithJson(remoteState.Config)
	if err != nil {
		return cty.NilVal, err
	}
	output["config"] = ctyJsonVal

	return convertValuesMapToCtyVal(output)
}

// Serialize the list of dependency blocks to a cty Value as a map that maps the block names to the cty representation.
func dependencyBlocksAsCty(dependencyBlocks []Dependency) (cty.Value, error) {
	out := map[string]cty.Value{}
	for _, block := range dependencyBlocks {
		blockCty, err := goTypeToCty(block)
		if err != nil {
			return cty.NilVal, err
		}
		out[block.Name] = blockCty
	}
	return convertValuesMapToCtyVal(out)
}

// Converts arbitrary go types that are json serializable to a cty Value by using json as an intermediary
// representation. This avoids the strict type nature of cty, where you need to know the output type beforehand to
// serialize to cty.
func convertToCtyWithJson(val interface{}) (cty.Value, error) {
	jsonBytes, err := json.Marshal(val)
	if err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	var ctyJsonVal ctyjson.SimpleJSONValue
	if err := ctyJsonVal.UnmarshalJSON(jsonBytes); err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	return ctyJsonVal.Value, nil
}

// Converts arbitrary go type (struct that has cty tags, slice, map with string keys, string, bool, int
// uint, float, cty.Value) to a cty Value
func goTypeToCty(val interface{}) (cty.Value, error) {
	ctyType, err := gocty.ImpliedType(val)
	if err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	ctyOut, err := gocty.ToCtyValue(val, ctyType)
	if err != nil {
		return cty.NilVal, errors.WithStackTrace(err)
	}
	return ctyOut, nil
}

// Converts primitive go strings to a cty Value.
func gostringToCty(val string) cty.Value {
	ctyOut, err := gocty.ToCtyValue(val, cty.String)
	if err != nil {
		// Since we are converting primitive strings, we should never get an error in this conversion.
		panic(err)
	}
	return ctyOut
}

// Converts primitive go bools to a cty Value.
func goboolToCty(val bool) cty.Value {
	ctyOut, err := gocty.ToCtyValue(val, cty.Bool)
	if err != nil {
		// Since we are converting primitive bools, we should never get an error in this conversion.
		panic(err)
	}
	return ctyOut
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/sha256"
	"encoding/hex"
	"github.com/mholt/archiver/v3"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/server/filesystem"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

// Archiver represents a Server Archiver.
type Archiver struct {
	Server *Server
}

// Path returns the path to the server's archive.
func (a *Archiver) Path() string {
	return filepath.Join(config.Get().System.ArchiveDirectory, a.Name())
}

// Name returns the name of the server's archive.
func (a *Archiver) Name() string {
	return a.Server.Id() + ".tar.gz"
}

// Exists returns a boolean based off if the archive exists.
func (a *Archiver) Exists() bool {
	if _, err := os.Stat(a.Path()); os.IsNotExist(err) {
		return false
	}

	return true
}

// Stat stats the archive file.
func (a *Archiver) Stat() (*filesystem.Stat, error) {
	s, err := os.Stat(a.Path())
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &filesystem.Stat{
		Info:     s,
		Mimetype: "application/tar+gzip",
	}, nil
}

// Archive creates an archive of the server and deletes the previous one.
func (a *Archiver) Archive() error {
	path := a.Server.Filesystem().Path()

	// Get the list of root files and directories to archive.
	var files []string
	fileInfo, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	for _, file := range fileInfo {
		f := filepath.Join(path, file.Name())
		// If the file is a symlink we cannot safely assume that the result of a filepath.Join() will be
		// a safe destination. We need to check if the file is a symlink, and if so pass off to the SafePath
		// function to resolve it to the final destination.
		//
		// ioutil.ReadDir() calls Lstat, so this will work correctly. If it did not call Lstat, but rather
		// just did a normal Stat call, this would fail since that would be looking at the symlink destination
		// and not the actual file in this listing.
		if file.Mode()&os.ModeSymlink != 0 {
			f, err = a.Server.Filesystem().SafePath(filepath.Join(path, file.Name()))

			if err != nil {
				return err
			}
		}

		files = append(files, f)
	}

	stat, err := a.Stat()
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	// Check if the file exists.
	if stat != nil {
		if err := os.Remove(a.Path()); err != nil {
			return err
		}
	}

	return archiver.NewTarGz().Archive(files, a.Path())
}

// DeleteIfExists deletes the archive if it exists.
func (a *Archiver) DeleteIfExists() error {
	stat, err := a.Stat()
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	// Check if the file exists.
	if stat != nil {
		if err := os.Remove(a.Path()); err != nil {
			return err
		}
	}

	return nil
}

// Checksum computes a SHA256 checksum of the server's archive.
func (a *Archiver) Checksum() (string, error) {
	file, err := os.Open(a.Path())
	if err != nil {
		return "", err
	}
	defer file.Close()

	hash := sha256.New()

	buf := make([]byte, 1024*4)
	if _, err := io.CopyBuffer(hash, file, buf); err != nil {
		return "", err
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}
<commit_msg>Correctly detect os.IsNotExist error; ref pterodactyl/panel#2438<commit_after>package server

import (
	"crypto/sha256"
	"encoding/hex"
	"github.com/mholt/archiver/v3"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/server/filesystem"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

// Archiver represents a Server Archiver.
type Archiver struct {
	Server *Server
}

// Path returns the path to the server's archive.
func (a *Archiver) Path() string {
	return filepath.Join(config.Get().System.ArchiveDirectory, a.Name())
}

// Name returns the name of the server's archive.
func (a *Archiver) Name() string {
	return a.Server.Id() + ".tar.gz"
}

// Exists returns a boolean based off if the archive exists.
func (a *Archiver) Exists() bool {
	if _, err := os.Stat(a.Path()); os.IsNotExist(err) {
		return false
	}

	return true
}

// Stat stats the archive file.
func (a *Archiver) Stat() (*filesystem.Stat, error) {
	s, err := os.Stat(a.Path())
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &filesystem.Stat{
		Info:     s,
		Mimetype: "application/tar+gzip",
	}, nil
}

// Archive creates an archive of the server and deletes the previous one.
func (a *Archiver) Archive() error {
	path := a.Server.Filesystem().Path()

	// Get the list of root files and directories to archive.
	var files []string
	fileInfo, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	for _, file := range fileInfo {
		f := filepath.Join(path, file.Name())
		// If the file is a symlink we cannot safely assume that the result of a filepath.Join() will be
		// a safe destination. We need to check if the file is a symlink, and if so pass off to the SafePath
		// function to resolve it to the final destination.
		//
		// ioutil.ReadDir() calls Lstat, so this will work correctly. If it did not call Lstat, but rather
		// just did a normal Stat call, this would fail since that would be looking at the symlink destination
		// and not the actual file in this listing.
		if file.Mode()&os.ModeSymlink != 0 {
			f, err = a.Server.Filesystem().SafePath(filepath.Join(path, file.Name()))

			if err != nil {
				return err
			}
		}

		files = append(files, f)
	}

	stat, err := a.Stat()
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	// Check if the file exists.
	if stat != nil {
		if err := os.Remove(a.Path()); err != nil {
			return err
		}
	}

	return archiver.NewTarGz().Archive(files, a.Path())
}

// DeleteIfExists deletes the archive if it exists.
func (a *Archiver) DeleteIfExists() error {
	stat, err := a.Stat()
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}

	// Check if the file exists.
	if stat != nil {
		if err := os.Remove(a.Path()); err != nil {
			return err
		}
	}

	return nil
}

// Checksum computes a SHA256 checksum of the server's archive.
func (a *Archiver) Checksum() (string, error) {
	file, err := os.Open(a.Path())
	if err != nil {
		return "", err
	}
	defer file.Close()

	hash := sha256.New()

	buf := make([]byte, 1024*4)
	if _, err := io.CopyBuffer(hash, file, buf); err != nil {
		return "", err
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}
<|endoftext|>


<commit_before>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package glog

import (
	"context"
	"fmt"
	"github.com/fatih/color"
	"github.com/gogf/gf/container/gtype"
	"github.com/gogf/gf/internal/intlog"
	"github.com/gogf/gf/os/gfpool"
	"github.com/gogf/gf/os/gmlock"
	"github.com/gogf/gf/os/gtimer"
	"go.opentelemetry.io/otel/trace"
	"io"
	"os"
	"strings"
	"time"

	"github.com/gogf/gf/debug/gdebug"

	"github.com/gogf/gf/os/gfile"
	"github.com/gogf/gf/os/gtime"
	"github.com/gogf/gf/text/gregex"
	"github.com/gogf/gf/util/gconv"
)

// Logger is the struct for logging management.
type Logger struct {
	ctx    context.Context // Context for logging.
	init   *gtype.Bool     // Initialized.
	parent *Logger         // Parent logger, if it is not empty, it means the logger is used in chaining function.
	config Config          // Logger configuration.
}

const (
	defaultFileFormat                 = `{Y-m-d}.log`
	defaultFileFlags                  = os.O_CREATE | os.O_WRONLY | os.O_APPEND
	defaultFilePerm                   = os.FileMode(0666)
	defaultFileExpire                 = time.Minute
	pathFilterKey                     = "/os/glog/glog"
	memoryLockPrefixForPrintingToFile = "glog.printToFile:"
)

const (
	F_ASYNC      = 1 << iota // Print logging content asynchronously。
	F_FILE_LONG              // Print full file name and line number: /a/b/c/d.go:23.
	F_FILE_SHORT             // Print final file name element and line number: d.go:23. overrides F_FILE_LONG.
	F_TIME_DATE              // Print the date in the local time zone: 2009-01-23.
	F_TIME_TIME              // Print the time in the local time zone: 01:23:23.
	F_TIME_MILLI             // Print the time with milliseconds in the local time zone: 01:23:23.675.
	F_CALLER_FN              // Print Caller function name and package: main.main
	F_TIME_STD   = F_TIME_DATE | F_TIME_MILLI
)

// New creates and returns a custom logger.
func New() *Logger {
	logger := &Logger{
		init:   gtype.NewBool(),
		config: DefaultConfig(),
	}
	return logger
}

// NewWithWriter creates and returns a custom logger with io.Writer.
func NewWithWriter(writer io.Writer) *Logger {
	l := New()
	l.SetWriter(writer)
	return l
}

// Clone returns a new logger, which is the clone the current logger.
// It's commonly used for chaining operations.
func (l *Logger) Clone() *Logger {
	logger := New()
	logger.ctx = l.ctx
	logger.config = l.config
	logger.parent = l
	return logger
}

// getFilePath returns the logging file path.
// The logging file name must have extension name of "log".
func (l *Logger) getFilePath(now time.Time) string {
	// Content containing "{}" in the file name is formatted using gtime.
	file, _ := gregex.ReplaceStringFunc(`{.+?}`, l.config.File, func(s string) string {
		return gtime.New(now).Format(strings.Trim(s, "{}"))
	})
	file = gfile.Join(l.config.Path, file)
	return file
}

// print prints <s> to defined writer, logging file or passed <std>.
func (l *Logger) print(ctx context.Context, level int, values ...interface{}) {
	// Lazy initialize for rotation feature.
	// It uses atomic reading operation to enhance the performance checking.
	// It here uses CAP for performance and concurrent safety.
	p := l
	if p.parent != nil {
		p = p.parent
	}
	// It just initializes once for each logger.
	if p.config.RotateSize > 0 || p.config.RotateExpire > 0 {
		if !p.init.Val() && p.init.Cas(false, true) {
			gtimer.AddOnce(p.config.RotateCheckInterval, p.rotateChecksTimely)
			intlog.Printf(ctx, "logger rotation initialized: every %s", p.config.RotateCheckInterval.String())
		}
	}

	var (
		now   = time.Now()
		input = &HandlerInput{
			logger: l,
			index:  -1,
			Ctx:    ctx,
			Time:   now,
			Color:  defaultLevelColor[level],
			Level:  level,
		}
	)
	if l.config.HeaderPrint {
		// Time.
		timeFormat := ""
		if l.config.Flags&F_TIME_DATE > 0 {
			timeFormat += "2006-01-02"
		}
		if l.config.Flags&F_TIME_TIME > 0 {
			if timeFormat != "" {
				timeFormat += " "
			}
			timeFormat += "15:04:05"
		}
		if l.config.Flags&F_TIME_MILLI > 0 {
			if timeFormat != "" {
				timeFormat += " "
			}
			timeFormat += "15:04:05.000"
		}
		if len(timeFormat) > 0 {
			input.TimeFormat = now.Format(timeFormat)
		}

		// Level string.
		input.LevelFormat = l.getLevelPrefixWithBrackets(level)

		// Caller path and Fn name.
		if l.config.Flags&(F_FILE_LONG|F_FILE_SHORT|F_CALLER_FN) > 0 {
			callerFnName, path, line := gdebug.CallerWithFilter(pathFilterKey, l.config.StSkip)
			if l.config.Flags&F_CALLER_FN > 0 {
				input.CallerFunc = fmt.Sprintf(`[%s]`, callerFnName)
			}
			if l.config.Flags&F_FILE_LONG > 0 {
				input.CallerPath = fmt.Sprintf(`%s:%d:`, path, line)
			}
			if l.config.Flags&F_FILE_SHORT > 0 {
				input.CallerPath = fmt.Sprintf(`%s:%d:`, gfile.Basename(path), line)
			}
		}
		// Prefix.
		if len(l.config.Prefix) > 0 {
			input.Prefix = l.config.Prefix
		}
	}
	// Convert value to string.
	if ctx != nil {
		// Tracing values.
		spanCtx := trace.SpanContextFromContext(ctx)
		if traceId := spanCtx.TraceID(); traceId.IsValid() {
			input.CtxStr = "{" + traceId.String() + "}"
		}
		// Context values.
		if len(l.config.CtxKeys) > 0 {
			ctxStr := ""
			for _, key := range l.config.CtxKeys {
				if v := ctx.Value(key); v != nil {
					if ctxStr != "" {
						ctxStr += ", "
					}
					ctxStr += fmt.Sprintf("%s: %+v", key, v)
				}
			}
			if ctxStr != "" {
				input.CtxStr += "{" + ctxStr + "}"
			}
		}
	}
	var tempStr string
	for _, v := range values {
		tempStr = gconv.String(v)
		if len(input.Content) > 0 {
			if input.Content[len(input.Content)-1] == '\n' {
				// Remove one blank line(\n\n).
				if tempStr[0] == '\n' {
					input.Content += tempStr[1:]
				} else {
					input.Content += tempStr
				}
			} else {
				input.Content += " " + tempStr
			}
		} else {
			input.Content = tempStr
		}
	}
	if l.config.Flags&F_ASYNC > 0 {
		input.IsAsync = true
		err := asyncPool.Add(func() {
			input.Next()
		})
		if err != nil {
			intlog.Error(ctx, err)
		}
	} else {
		input.Next()
	}
}

// doPrint outputs the logging content according configuration.
func (l *Logger) doPrint(ctx context.Context, input *HandlerInput) {
	if l.config.Writer == nil {
		// Output content to disk file.
		if l.config.Path != "" {
			l.printToFile(ctx, input.Time, input)
		}
		// Allow output to stdout?
		if l.config.StdoutPrint {
			l.printToStdout(ctx, input)
		}
	} else {
		// Output to custom writer.
		l.printToWriter(ctx, input)
	}
}

// printToWriter writes buffer to writer.
func (l *Logger) printToWriter(ctx context.Context, input *HandlerInput) {
	if l.config.Writer != nil {
		var (
			buffer = input.getBuffer(l.config.WriterColorEnable)
		)
		if _, err := l.config.Writer.Write(buffer.Bytes()); err != nil {
			intlog.Error(ctx, err)
		}
	}
}

// printToStdout outputs logging content to stdout.
func (l *Logger) printToStdout(ctx context.Context, input *HandlerInput) {
	if l.config.StdoutPrint {
		// if _, err := os.Stdout.Write(input.getBuffer(true).Bytes()); err != nil {
		// For color in windows.
		if _, err := fmt.Fprintf(color.Output, input.getBuffer(true).String()); err != nil {
			intlog.Error(ctx, err)
		}
	}
}

// printToFile outputs logging content to disk file.
func (l *Logger) printToFile(ctx context.Context, t time.Time, input *HandlerInput) {
	var (
		buffer        = input.getBuffer(l.config.WriterColorEnable)
		logFilePath   = l.getFilePath(t)
		memoryLockKey = memoryLockPrefixForPrintingToFile + logFilePath
	)
	gmlock.Lock(memoryLockKey)
	defer gmlock.Unlock(memoryLockKey)

	// Rotation file size checks.
	if l.config.RotateSize > 0 {
		if gfile.Size(logFilePath) > l.config.RotateSize {
			l.rotateFileBySize(t)
		}
	}
	// Logging content outputting to disk file.
	if file := l.getFilePointer(ctx, logFilePath); file == nil {
		intlog.Errorf(ctx, `got nil file pointer for: %s`, logFilePath)
	} else {
		if _, err := file.Write(buffer.Bytes()); err != nil {
			intlog.Error(ctx, err)
		}
		if err := file.Close(); err != nil {
			intlog.Error(ctx, err)
		}
	}
}

// getFilePointer retrieves and returns a file pointer from file pool.
func (l *Logger) getFilePointer(ctx context.Context, path string) *gfpool.File {
	file, err := gfpool.Open(
		path,
		defaultFileFlags,
		defaultFilePerm,
		defaultFileExpire,
	)
	if err != nil {
		// panic(err)
		intlog.Error(ctx, err)
	}
	return file
}

// getCtx returns the context which is set through chaining operations.
// It returns an empty context if no context set previously.
func (l *Logger) getCtx() context.Context {
	if l.ctx != nil {
		return l.ctx
	}
	return context.TODO()
}

// printStd prints content <s> without stack.
func (l *Logger) printStd(level int, value ...interface{}) {
	l.print(l.getCtx(), level, value...)
}

// printStd prints content <s> with stack check.
func (l *Logger) printErr(level int, value ...interface{}) {
	if l.config.StStatus == 1 {
		if s := l.GetStack(); s != "" {
			value = append(value, "\nStack:\n"+s)
		}
	}
	// In matter of sequence, do not use stderr here, but use the same stdout.
	l.print(l.getCtx(), level, value...)
}

// format formats <values> using fmt.Sprintf.
func (l *Logger) format(format string, value ...interface{}) string {
	return fmt.Sprintf(format, value...)
}

// PrintStack prints the caller stack,
// the optional parameter <skip> specify the skipped stack offset from the end point.
func (l *Logger) PrintStack(skip ...int) {
	if s := l.GetStack(skip...); s != "" {
		l.Println("Stack:\n" + s)
	} else {
		l.Println()
	}
}

// GetStack returns the caller stack content,
// the optional parameter <skip> specify the skipped stack offset from the end point.
func (l *Logger) GetStack(skip ...int) string {
	stackSkip := l.config.StSkip
	if len(skip) > 0 {
		stackSkip += skip[0]
	}
	filters := []string{pathFilterKey}
	if l.config.StFilter != "" {
		filters = append(filters, l.config.StFilter)
	}
	return gdebug.StackWithFilters(filters, stackSkip)
}
<commit_msg>comment update for package glog<commit_after>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package glog

import (
	"context"
	"fmt"
	"github.com/fatih/color"
	"github.com/gogf/gf/container/gtype"
	"github.com/gogf/gf/internal/intlog"
	"github.com/gogf/gf/os/gfpool"
	"github.com/gogf/gf/os/gmlock"
	"github.com/gogf/gf/os/gtimer"
	"go.opentelemetry.io/otel/trace"
	"io"
	"os"
	"strings"
	"time"

	"github.com/gogf/gf/debug/gdebug"

	"github.com/gogf/gf/os/gfile"
	"github.com/gogf/gf/os/gtime"
	"github.com/gogf/gf/text/gregex"
	"github.com/gogf/gf/util/gconv"
)

// Logger is the struct for logging management.
type Logger struct {
	ctx    context.Context // Context for logging.
	init   *gtype.Bool     // Initialized.
	parent *Logger         // Parent logger, if it is not empty, it means the logger is used in chaining function.
	config Config          // Logger configuration.
}

const (
	defaultFileFormat                 = `{Y-m-d}.log`
	defaultFileFlags                  = os.O_CREATE | os.O_WRONLY | os.O_APPEND
	defaultFilePerm                   = os.FileMode(0666)
	defaultFileExpire                 = time.Minute
	pathFilterKey                     = "/os/glog/glog"
	memoryLockPrefixForPrintingToFile = "glog.printToFile:"
)

const (
	F_ASYNC      = 1 << iota // Print logging content asynchronously。
	F_FILE_LONG              // Print full file name and line number: /a/b/c/d.go:23.
	F_FILE_SHORT             // Print final file name element and line number: d.go:23. overrides F_FILE_LONG.
	F_TIME_DATE              // Print the date in the local time zone: 2009-01-23.
	F_TIME_TIME              // Print the time in the local time zone: 01:23:23.
	F_TIME_MILLI             // Print the time with milliseconds in the local time zone: 01:23:23.675.
	F_CALLER_FN              // Print Caller function name and package: main.main
	F_TIME_STD   = F_TIME_DATE | F_TIME_MILLI
)

// New creates and returns a custom logger.
func New() *Logger {
	logger := &Logger{
		init:   gtype.NewBool(),
		config: DefaultConfig(),
	}
	return logger
}

// NewWithWriter creates and returns a custom logger with io.Writer.
func NewWithWriter(writer io.Writer) *Logger {
	l := New()
	l.SetWriter(writer)
	return l
}

// Clone returns a new logger, which is the clone the current logger.
// It's commonly used for chaining operations.
func (l *Logger) Clone() *Logger {
	logger := New()
	logger.ctx = l.ctx
	logger.config = l.config
	logger.parent = l
	return logger
}

// getFilePath returns the logging file path.
// The logging file name must have extension name of "log".
func (l *Logger) getFilePath(now time.Time) string {
	// Content containing "{}" in the file name is formatted using gtime.
	file, _ := gregex.ReplaceStringFunc(`{.+?}`, l.config.File, func(s string) string {
		return gtime.New(now).Format(strings.Trim(s, "{}"))
	})
	file = gfile.Join(l.config.Path, file)
	return file
}

// print prints <s> to defined writer, logging file or passed <std>.
func (l *Logger) print(ctx context.Context, level int, values ...interface{}) {
	// Lazy initialize for rotation feature.
	// It uses atomic reading operation to enhance the performance checking.
	// It here uses CAP for performance and concurrent safety.
	p := l
	if p.parent != nil {
		p = p.parent
	}
	// It just initializes once for each logger.
	if p.config.RotateSize > 0 || p.config.RotateExpire > 0 {
		if !p.init.Val() && p.init.Cas(false, true) {
			gtimer.AddOnce(p.config.RotateCheckInterval, p.rotateChecksTimely)
			intlog.Printf(ctx, "logger rotation initialized: every %s", p.config.RotateCheckInterval.String())
		}
	}

	var (
		now   = time.Now()
		input = &HandlerInput{
			logger: l,
			index:  -1,
			Ctx:    ctx,
			Time:   now,
			Color:  defaultLevelColor[level],
			Level:  level,
		}
	)
	if l.config.HeaderPrint {
		// Time.
		timeFormat := ""
		if l.config.Flags&F_TIME_DATE > 0 {
			timeFormat += "2006-01-02"
		}
		if l.config.Flags&F_TIME_TIME > 0 {
			if timeFormat != "" {
				timeFormat += " "
			}
			timeFormat += "15:04:05"
		}
		if l.config.Flags&F_TIME_MILLI > 0 {
			if timeFormat != "" {
				timeFormat += " "
			}
			timeFormat += "15:04:05.000"
		}
		if len(timeFormat) > 0 {
			input.TimeFormat = now.Format(timeFormat)
		}

		// Level string.
		input.LevelFormat = l.getLevelPrefixWithBrackets(level)

		// Caller path and Fn name.
		if l.config.Flags&(F_FILE_LONG|F_FILE_SHORT|F_CALLER_FN) > 0 {
			callerFnName, path, line := gdebug.CallerWithFilter(pathFilterKey, l.config.StSkip)
			if l.config.Flags&F_CALLER_FN > 0 {
				input.CallerFunc = fmt.Sprintf(`[%s]`, callerFnName)
			}
			if l.config.Flags&F_FILE_LONG > 0 {
				input.CallerPath = fmt.Sprintf(`%s:%d:`, path, line)
			}
			if l.config.Flags&F_FILE_SHORT > 0 {
				input.CallerPath = fmt.Sprintf(`%s:%d:`, gfile.Basename(path), line)
			}
		}
		// Prefix.
		if len(l.config.Prefix) > 0 {
			input.Prefix = l.config.Prefix
		}
	}
	// Convert value to string.
	if ctx != nil {
		// Tracing values.
		spanCtx := trace.SpanContextFromContext(ctx)
		if traceId := spanCtx.TraceID(); traceId.IsValid() {
			input.CtxStr = "{" + traceId.String() + "}"
		}
		// Context values.
		if len(l.config.CtxKeys) > 0 {
			ctxStr := ""
			for _, key := range l.config.CtxKeys {
				if v := ctx.Value(key); v != nil {
					if ctxStr != "" {
						ctxStr += ", "
					}
					ctxStr += fmt.Sprintf("%s: %+v", key, v)
				}
			}
			if ctxStr != "" {
				input.CtxStr += "{" + ctxStr + "}"
			}
		}
	}
	var tempStr string
	for _, v := range values {
		tempStr = gconv.String(v)
		if len(input.Content) > 0 {
			if input.Content[len(input.Content)-1] == '\n' {
				// Remove one blank line(\n\n).
				if tempStr[0] == '\n' {
					input.Content += tempStr[1:]
				} else {
					input.Content += tempStr
				}
			} else {
				input.Content += " " + tempStr
			}
		} else {
			input.Content = tempStr
		}
	}
	if l.config.Flags&F_ASYNC > 0 {
		input.IsAsync = true
		err := asyncPool.Add(func() {
			input.Next()
		})
		if err != nil {
			intlog.Error(ctx, err)
		}
	} else {
		input.Next()
	}
}

// doPrint outputs the logging content according configuration.
func (l *Logger) doPrint(ctx context.Context, input *HandlerInput) {
	if l.config.Writer == nil {
		// Output content to disk file.
		if l.config.Path != "" {
			l.printToFile(ctx, input.Time, input)
		}
		// Allow output to stdout?
		if l.config.StdoutPrint {
			l.printToStdout(ctx, input)
		}
	} else {
		// Output to custom writer.
		l.printToWriter(ctx, input)
	}
}

// printToWriter writes buffer to writer.
func (l *Logger) printToWriter(ctx context.Context, input *HandlerInput) {
	if l.config.Writer != nil {
		var (
			buffer = input.getBuffer(l.config.WriterColorEnable)
		)
		if _, err := l.config.Writer.Write(buffer.Bytes()); err != nil {
			intlog.Error(ctx, err)
		}
	}
}

// printToStdout outputs logging content to stdout.
func (l *Logger) printToStdout(ctx context.Context, input *HandlerInput) {
	if l.config.StdoutPrint {
		// This will lose color in Windows os system.
		// if _, err := os.Stdout.Write(input.getBuffer(true).Bytes()); err != nil {
		// This will print color in Windows os system.
		if _, err := fmt.Fprintf(color.Output, input.getBuffer(true).String()); err != nil {
			intlog.Error(ctx, err)
		}
	}
}

// printToFile outputs logging content to disk file.
func (l *Logger) printToFile(ctx context.Context, t time.Time, input *HandlerInput) {
	var (
		buffer        = input.getBuffer(l.config.WriterColorEnable)
		logFilePath   = l.getFilePath(t)
		memoryLockKey = memoryLockPrefixForPrintingToFile + logFilePath
	)
	gmlock.Lock(memoryLockKey)
	defer gmlock.Unlock(memoryLockKey)

	// Rotation file size checks.
	if l.config.RotateSize > 0 {
		if gfile.Size(logFilePath) > l.config.RotateSize {
			l.rotateFileBySize(t)
		}
	}
	// Logging content outputting to disk file.
	if file := l.getFilePointer(ctx, logFilePath); file == nil {
		intlog.Errorf(ctx, `got nil file pointer for: %s`, logFilePath)
	} else {
		if _, err := file.Write(buffer.Bytes()); err != nil {
			intlog.Error(ctx, err)
		}
		if err := file.Close(); err != nil {
			intlog.Error(ctx, err)
		}
	}
}

// getFilePointer retrieves and returns a file pointer from file pool.
func (l *Logger) getFilePointer(ctx context.Context, path string) *gfpool.File {
	file, err := gfpool.Open(
		path,
		defaultFileFlags,
		defaultFilePerm,
		defaultFileExpire,
	)
	if err != nil {
		// panic(err)
		intlog.Error(ctx, err)
	}
	return file
}

// getCtx returns the context which is set through chaining operations.
// It returns an empty context if no context set previously.
func (l *Logger) getCtx() context.Context {
	if l.ctx != nil {
		return l.ctx
	}
	return context.TODO()
}

// printStd prints content <s> without stack.
func (l *Logger) printStd(level int, value ...interface{}) {
	l.print(l.getCtx(), level, value...)
}

// printStd prints content <s> with stack check.
func (l *Logger) printErr(level int, value ...interface{}) {
	if l.config.StStatus == 1 {
		if s := l.GetStack(); s != "" {
			value = append(value, "\nStack:\n"+s)
		}
	}
	// In matter of sequence, do not use stderr here, but use the same stdout.
	l.print(l.getCtx(), level, value...)
}

// format formats <values> using fmt.Sprintf.
func (l *Logger) format(format string, value ...interface{}) string {
	return fmt.Sprintf(format, value...)
}

// PrintStack prints the caller stack,
// the optional parameter <skip> specify the skipped stack offset from the end point.
func (l *Logger) PrintStack(skip ...int) {
	if s := l.GetStack(skip...); s != "" {
		l.Println("Stack:\n" + s)
	} else {
		l.Println()
	}
}

// GetStack returns the caller stack content,
// the optional parameter <skip> specify the skipped stack offset from the end point.
func (l *Logger) GetStack(skip ...int) string {
	stackSkip := l.config.StSkip
	if len(skip) > 0 {
		stackSkip += skip[0]
	}
	filters := []string{pathFilterKey}
	if l.config.StFilter != "" {
		filters = append(filters, l.config.StFilter)
	}
	return gdebug.StackWithFilters(filters, stackSkip)
}
<|endoftext|>


<commit_before>// +build integration

package etcd

import (
	"testing"
)

func Test_CreateExistsSearch(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	key := "my-key"
	val := "my-val"

	// There should be no key/value pair being stored initially.
	ok, err := newStorage.Exists(key)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if ok {
		t.Fatal("expected", false, "got", true)
	}

	// Creating the key/value pair should work.
	err = newStorage.Create(key, val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	// There should be the created key/value pair.
	ok, err = newStorage.Exists(key)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if !ok {
		t.Fatal("expected", true, "got", false)
	}
}

func Test_List(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	val := "my-val"

	err = newStorage.Create("key/one", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	err = newStorage.Create("key/two", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	values, err := newStorage.List("key")
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if len(values) != 2 {
		t.Fatal("expected", 2, "got", len(values))
	}
}

func Test_List_Invalid(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	val := "my-val"

	err = newStorage.Create("key/one", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	err = newStorage.Create("key/two", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	_, err = newStorage.List("ke")
	if !IsNotFound(err) {
		t.Fatal("expected", true, "got", false)
	}
}
<commit_msg>added tests for storage.Service.List (#27)<commit_after>// +build integration

package etcd

import (
	"testing"
)

func Test_CreateExistsSearch(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	key := "my-key"
	val := "my-val"

	// There should be no key/value pair being stored initially.
	ok, err := newStorage.Exists(key)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if ok {
		t.Fatal("expected", false, "got", true)
	}

	// Creating the key/value pair should work.
	err = newStorage.Create(key, val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	// There should be the created key/value pair.
	ok, err = newStorage.Exists(key)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if !ok {
		t.Fatal("expected", true, "got", false)
	}
}

func Test_List(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	val := "my-val"

	err = newStorage.Create("key/one", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	err = newStorage.Create("key/two", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	values, err := newStorage.List("key")
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if len(values) != 2 {
		t.Fatal("expected", 2, "got", len(values))
	}
	if values[0] != "one" {
		t.Fatal("expected", "one", "got", values[0])
	}
	if values[1] != "two" {
		t.Fatal("expected", "two", "got", values[1])
	}
}

func Test_List_Nesting(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	val := "my-val"

	err = newStorage.Create("tokend/token/some/scope/id1", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	err = newStorage.Create("tokend/token/some/other/scope/id34", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	values, err := newStorage.List("tokend/token")
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	if len(values) != 2 {
		t.Fatal("expected", 2, "got", len(values))
	}
	if values[0] != "some/other/scope/id34" {
		t.Fatal("expected", "some/other/scope/id34", "got", values[0])
	}
	if values[1] != "some/scope/id1" {
		t.Fatal("expected", "some/scope/id1", "got", values[1])
	}
}

func Test_List_Invalid_Key(t *testing.T) {
	config := DefaultConfig()
	config.Prefix = "foo"
	newStorage, err := New(config)
	if err != nil {
		panic(err)
	}

	val := "my-val"

	err = newStorage.Create("key/one", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}
	err = newStorage.Create("key/two", val)
	if err != nil {
		t.Fatal("expected", nil, "got", err)
	}

	_, err = newStorage.List("ke")
	if !IsNotFound(err) {
		t.Fatal("expected", true, "got", false)
	}
}
<|endoftext|>


<commit_before>package memory

import (
	"encoding/binary"
	"errors"
	"net"
	"runtime"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/chihaya/chihaya/bittorrent"
	"github.com/chihaya/chihaya/storage"
)

// ErrInvalidGCInterval is returned for a GarbageCollectionInterval that is
// less than or equal to zero.
var ErrInvalidGCInterval = errors.New("invalid garbage collection interval")

// Config holds the configuration of a memory PeerStore.
type Config struct {
	GarbageCollectionInterval time.Duration `yaml:"gc_interval"`
	PeerLifetime              time.Duration `yaml:"peer_lifetime"`
	ShardCount                int           `yaml:"shard_count"`
	MaxNumWant                int           `yaml:"max_numwant"`
}

// New creates a new PeerStore backed by memory.
func New(cfg Config) (storage.PeerStore, error) {
	shardCount := 1
	if cfg.ShardCount > 0 {
		shardCount = cfg.ShardCount
	}

	if cfg.GarbageCollectionInterval <= 0 {
		return nil, ErrInvalidGCInterval
	}

	ps := &peerStore{
		shards:     make([]*peerShard, shardCount*2),
		closed:     make(chan struct{}),
		maxNumWant: cfg.MaxNumWant,
	}

	for i := 0; i < shardCount*2; i++ {
		ps.shards[i] = &peerShard{swarms: make(map[bittorrent.InfoHash]swarm)}
	}

	go func() {
		for {
			select {
			case <-ps.closed:
				return
			case <-time.After(cfg.GarbageCollectionInterval):
				before := time.Now().Add(-cfg.PeerLifetime)
				log.Debugln("memory: purging peers with no announces since", before)
				ps.collectGarbage(before)
			}
		}
	}()

	return ps, nil
}

type serializedPeer string

type peerShard struct {
	swarms map[bittorrent.InfoHash]swarm
	sync.RWMutex
}

type swarm struct {
	// map serialized peer to mtime
	seeders  map[serializedPeer]int64
	leechers map[serializedPeer]int64
}

type peerStore struct {
	shards     []*peerShard
	closed     chan struct{}
	maxNumWant int
}

var _ storage.PeerStore = &peerStore{}

func (s *peerStore) shardIndex(infoHash bittorrent.InfoHash, p bittorrent.Peer) uint32 {
	idx := binary.BigEndian.Uint32(infoHash[:4]) % uint32(len(s.shards))
	if len(p.IP) == net.IPv6len {
		idx += idx + uint32(len(s.shards)/2)
	}
	return idx
}

func newPeerKey(p bittorrent.Peer) serializedPeer {
	b := make([]byte, 20+2+len(p.IP))
	copy(b[:20], p.ID[:])
	binary.BigEndian.PutUint16(b[20:22], p.Port)
	copy(b[22:], p.IP)

	return serializedPeer(b)
}

func decodePeerKey(pk serializedPeer) bittorrent.Peer {
	return bittorrent.Peer{
		ID:   bittorrent.PeerIDFromString(string(pk[:20])),
		Port: binary.BigEndian.Uint16([]byte(pk[20:22])),
		IP:   net.IP(pk[22:]),
	}
}

func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, p)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.swarms[ih] = swarm{
			seeders:  make(map[serializedPeer]int64),
			leechers: make(map[serializedPeer]int64),
		}
	}

	shard.swarms[ih].seeders[pk] = time.Now().UnixNano()

	shard.Unlock()
	return nil
}

func (s *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, p)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	if _, ok := shard.swarms[ih].seeders[pk]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	delete(shard.swarms[ih].seeders, pk)

	if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 {
		delete(shard.swarms, ih)
	}

	shard.Unlock()
	return nil
}

func (s *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, p)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.swarms[ih] = swarm{
			seeders:  make(map[serializedPeer]int64),
			leechers: make(map[serializedPeer]int64),
		}
	}

	shard.swarms[ih].leechers[pk] = time.Now().UnixNano()

	shard.Unlock()
	return nil
}

func (s *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, p)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	if _, ok := shard.swarms[ih].leechers[pk]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	delete(shard.swarms[ih].leechers, pk)

	if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 {
		delete(shard.swarms, ih)
	}

	shard.Unlock()
	return nil
}

func (s *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, p)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.swarms[ih] = swarm{
			seeders:  make(map[serializedPeer]int64),
			leechers: make(map[serializedPeer]int64),
		}
	}

	delete(shard.swarms[ih].leechers, pk)

	shard.swarms[ih].seeders[pk] = time.Now().UnixNano()

	shard.Unlock()
	return nil
}

func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant int, announcer bittorrent.Peer) (peers []bittorrent.Peer, err error) {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	if numWant > s.maxNumWant {
		numWant = s.maxNumWant
	}

	shard := s.shards[s.shardIndex(ih, announcer)]
	shard.RLock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.RUnlock()
		return nil, storage.ErrResourceDoesNotExist
	}

	if seeder {
		// Append leechers as possible.
		leechers := shard.swarms[ih].leechers
		for p := range leechers {
			decodedPeer := decodePeerKey(p)
			if numWant == 0 {
				break
			}

			peers = append(peers, decodedPeer)
			numWant--
		}
	} else {
		// Append as many seeders as possible.
		seeders := shard.swarms[ih].seeders
		for p := range seeders {
			decodedPeer := decodePeerKey(p)
			if numWant == 0 {
				break
			}

			peers = append(peers, decodedPeer)
			numWant--
		}

		// Append leechers until we reach numWant.
		leechers := shard.swarms[ih].leechers
		if numWant > 0 {
			for p := range leechers {
				decodedPeer := decodePeerKey(p)
				if numWant == 0 {
					break
				}

				if decodedPeer.Equal(announcer) {
					continue
				}
				peers = append(peers, decodedPeer)
				numWant--
			}
		}
	}

	shard.RUnlock()
	return
}

// collectGarbage deletes all Peers from the PeerStore which are older than the
// cutoff time.
//
// This function must be able to execute while other methods on this interface
// are being executed in parallel.
func (s *peerStore) collectGarbage(cutoff time.Time) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	cutoffUnix := cutoff.UnixNano()
	for _, shard := range s.shards {
		shard.RLock()
		var infohashes []bittorrent.InfoHash
		for ih := range shard.swarms {
			infohashes = append(infohashes, ih)
		}
		shard.RUnlock()
		runtime.Gosched()

		for _, ih := range infohashes {
			shard.Lock()

			if _, stillExists := shard.swarms[ih]; !stillExists {
				shard.Unlock()
				runtime.Gosched()
				continue
			}

			for pk, mtime := range shard.swarms[ih].leechers {
				if mtime <= cutoffUnix {
					delete(shard.swarms[ih].leechers, pk)
				}
			}

			for pk, mtime := range shard.swarms[ih].seeders {
				if mtime <= cutoffUnix {
					delete(shard.swarms[ih].seeders, pk)
				}
			}

			if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 {
				delete(shard.swarms, ih)
			}

			shard.Unlock()
			runtime.Gosched()
		}

		runtime.Gosched()
	}

	return nil
}

func (s *peerStore) Stop() <-chan error {
	toReturn := make(chan error)
	go func() {
		shards := make([]*peerShard, len(s.shards))
		for i := 0; i < len(s.shards); i++ {
			shards[i] = &peerShard{swarms: make(map[bittorrent.InfoHash]swarm)}
		}
		s.shards = shards
		close(s.closed)
		close(toReturn)
	}()
	return toReturn
}
<commit_msg>memory: add ScrapeSwarm method<commit_after>package memory

import (
	"encoding/binary"
	"errors"
	"net"
	"runtime"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/chihaya/chihaya/bittorrent"
	"github.com/chihaya/chihaya/storage"
)

// ErrInvalidGCInterval is returned for a GarbageCollectionInterval that is
// less than or equal to zero.
var ErrInvalidGCInterval = errors.New("invalid garbage collection interval")

// Config holds the configuration of a memory PeerStore.
type Config struct {
	GarbageCollectionInterval time.Duration `yaml:"gc_interval"`
	PeerLifetime              time.Duration `yaml:"peer_lifetime"`
	ShardCount                int           `yaml:"shard_count"`
	MaxNumWant                int           `yaml:"max_numwant"`
}

// New creates a new PeerStore backed by memory.
func New(cfg Config) (storage.PeerStore, error) {
	shardCount := 1
	if cfg.ShardCount > 0 {
		shardCount = cfg.ShardCount
	}

	if cfg.GarbageCollectionInterval <= 0 {
		return nil, ErrInvalidGCInterval
	}

	ps := &peerStore{
		shards:     make([]*peerShard, shardCount*2),
		closed:     make(chan struct{}),
		maxNumWant: cfg.MaxNumWant,
	}

	for i := 0; i < shardCount*2; i++ {
		ps.shards[i] = &peerShard{swarms: make(map[bittorrent.InfoHash]swarm)}
	}

	go func() {
		for {
			select {
			case <-ps.closed:
				return
			case <-time.After(cfg.GarbageCollectionInterval):
				before := time.Now().Add(-cfg.PeerLifetime)
				log.Debugln("memory: purging peers with no announces since", before)
				ps.collectGarbage(before)
			}
		}
	}()

	return ps, nil
}

type serializedPeer string

type peerShard struct {
	swarms map[bittorrent.InfoHash]swarm
	sync.RWMutex
}

type swarm struct {
	// map serialized peer to mtime
	seeders  map[serializedPeer]int64
	leechers map[serializedPeer]int64
}

type peerStore struct {
	shards     []*peerShard
	closed     chan struct{}
	maxNumWant int
}

var _ storage.PeerStore = &peerStore{}

func (s *peerStore) shardIndex(infoHash bittorrent.InfoHash, v6 bool) uint32 {
	// There are twice the amount of shards specified by the user, the first
	// half is dedicated to IPv4 swarms and the second half is dedicated to
	// IPv6 swarms.
	idx := binary.BigEndian.Uint32(infoHash[:4]) % (uint32(len(s.shards)) / 2)
	if v6 {
		idx += uint32(len(s.shards) / 2)
	}
	return idx
}

func newPeerKey(p bittorrent.Peer) serializedPeer {
	b := make([]byte, 20+2+len(p.IP))
	copy(b[:20], p.ID[:])
	binary.BigEndian.PutUint16(b[20:22], p.Port)
	copy(b[22:], p.IP)

	return serializedPeer(b)
}

func decodePeerKey(pk serializedPeer) bittorrent.Peer {
	return bittorrent.Peer{
		ID:   bittorrent.PeerIDFromString(string(pk[:20])),
		Port: binary.BigEndian.Uint16([]byte(pk[20:22])),
		IP:   net.IP(pk[22:]),
	}
}

func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.swarms[ih] = swarm{
			seeders:  make(map[serializedPeer]int64),
			leechers: make(map[serializedPeer]int64),
		}
	}

	shard.swarms[ih].seeders[pk] = time.Now().UnixNano()

	shard.Unlock()
	return nil
}

func (s *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	if _, ok := shard.swarms[ih].seeders[pk]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	delete(shard.swarms[ih].seeders, pk)

	if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 {
		delete(shard.swarms, ih)
	}

	shard.Unlock()
	return nil
}

func (s *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.swarms[ih] = swarm{
			seeders:  make(map[serializedPeer]int64),
			leechers: make(map[serializedPeer]int64),
		}
	}

	shard.swarms[ih].leechers[pk] = time.Now().UnixNano()

	shard.Unlock()
	return nil
}

func (s *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	if _, ok := shard.swarms[ih].leechers[pk]; !ok {
		shard.Unlock()
		return storage.ErrResourceDoesNotExist
	}

	delete(shard.swarms[ih].leechers, pk)

	if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 {
		delete(shard.swarms, ih)
	}

	shard.Unlock()
	return nil
}

func (s *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	pk := newPeerKey(p)

	shard := s.shards[s.shardIndex(ih, len(p.IP) == net.IPv6len)]
	shard.Lock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.swarms[ih] = swarm{
			seeders:  make(map[serializedPeer]int64),
			leechers: make(map[serializedPeer]int64),
		}
	}

	delete(shard.swarms[ih].leechers, pk)

	shard.swarms[ih].seeders[pk] = time.Now().UnixNano()

	shard.Unlock()
	return nil
}

func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant int, announcer bittorrent.Peer) (peers []bittorrent.Peer, err error) {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	if numWant > s.maxNumWant {
		numWant = s.maxNumWant
	}

	shard := s.shards[s.shardIndex(ih, len(announcer.IP) == net.IPv6len)]
	shard.RLock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.RUnlock()
		return nil, storage.ErrResourceDoesNotExist
	}

	if seeder {
		// Append leechers as possible.
		leechers := shard.swarms[ih].leechers
		for p := range leechers {
			decodedPeer := decodePeerKey(p)
			if numWant == 0 {
				break
			}

			peers = append(peers, decodedPeer)
			numWant--
		}
	} else {
		// Append as many seeders as possible.
		seeders := shard.swarms[ih].seeders
		for p := range seeders {
			decodedPeer := decodePeerKey(p)
			if numWant == 0 {
				break
			}

			peers = append(peers, decodedPeer)
			numWant--
		}

		// Append leechers until we reach numWant.
		leechers := shard.swarms[ih].leechers
		if numWant > 0 {
			for p := range leechers {
				decodedPeer := decodePeerKey(p)
				if numWant == 0 {
					break
				}

				if decodedPeer.Equal(announcer) {
					continue
				}
				peers = append(peers, decodedPeer)
				numWant--
			}
		}
	}

	shard.RUnlock()
	return
}

func (s *peerStore) ScrapeSwarm(ih bittorrent.InfoHash, v6 bool) (resp bittorrent.Scrape) {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	shard := s.shards[s.shardIndex(ih, v6)]
	shard.RLock()

	if _, ok := shard.swarms[ih]; !ok {
		shard.RUnlock()
		return
	}

	resp.Incomplete = uint32(len(shard.swarms[ih].leechers))
	resp.Complete = uint32(len(shard.swarms[ih].seeders))
	shard.RUnlock()

	return
}

// collectGarbage deletes all Peers from the PeerStore which are older than the
// cutoff time.
//
// This function must be able to execute while other methods on this interface
// are being executed in parallel.
func (s *peerStore) collectGarbage(cutoff time.Time) error {
	select {
	case <-s.closed:
		panic("attempted to interact with stopped memory store")
	default:
	}

	cutoffUnix := cutoff.UnixNano()
	for _, shard := range s.shards {
		shard.RLock()
		var infohashes []bittorrent.InfoHash
		for ih := range shard.swarms {
			infohashes = append(infohashes, ih)
		}
		shard.RUnlock()
		runtime.Gosched()

		for _, ih := range infohashes {
			shard.Lock()

			if _, stillExists := shard.swarms[ih]; !stillExists {
				shard.Unlock()
				runtime.Gosched()
				continue
			}

			for pk, mtime := range shard.swarms[ih].leechers {
				if mtime <= cutoffUnix {
					delete(shard.swarms[ih].leechers, pk)
				}
			}

			for pk, mtime := range shard.swarms[ih].seeders {
				if mtime <= cutoffUnix {
					delete(shard.swarms[ih].seeders, pk)
				}
			}

			if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 {
				delete(shard.swarms, ih)
			}

			shard.Unlock()
			runtime.Gosched()
		}

		runtime.Gosched()
	}

	return nil
}

func (s *peerStore) Stop() <-chan error {
	toReturn := make(chan error)
	go func() {
		shards := make([]*peerShard, len(s.shards))
		for i := 0; i < len(s.shards); i++ {
			shards[i] = &peerShard{swarms: make(map[bittorrent.InfoHash]swarm)}
		}
		s.shards = shards
		close(s.closed)
		close(toReturn)
	}()
	return toReturn
}
<|endoftext|>


<commit_before>package zazu

import (
	"reflect"
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"
)

// TestNextJobs tests the getNextJobs function, which queries the database to find
// the next queued jobs, in order of their priority.
func TestGetNextJobs(t *testing.T) {
	// TODO: consider edge cases and make this test more rigorous.
	// e.g.:
	// 	1. What happens when there are no queued jobs?
	//		2. What happens when n > len(queued jobs)?
	//		3. Is the job status correct at every stage?
	//		4. Is a given job gauranteed to only be returned by getNextJobs() once?
	flushdb()
	jobTypes = map[string]*JobType{}

	// Create a test job with high priority
	highPriorityJob, err := createTestJob()
	if err != nil {
		t.Errorf("Unexpected error creating test job: %s", err.Error())
	}
	highPriorityJob.priority = 1000
	highPriorityJob.id = "highPriorityJob"
	if err := highPriorityJob.save(); err != nil {
		t.Errorf("Unexpected error saving test job: %s", err.Error())
	}
	if err := highPriorityJob.Enqueue(); err != nil {
		t.Errorf("Unexpected error enqueuing test job: %s", err.Error())
	}

	// Create more tests with lower priorities
	for i := 0; i < 10; i++ {
		job, err := createTestJob()
		if err != nil {
			t.Errorf("Unexpected error creating test job: %s", err.Error())
		}
		job.priority = 100
		job.id = "lowPriorityJob" + strconv.Itoa(i)
		if err := job.save(); err != nil {
			t.Errorf("Unexpected error saving test job: %s", err.Error())
		}
		if err := job.Enqueue(); err != nil {
			t.Errorf("Unexpected error enqueuing test job: %s", err.Error())
		}
	}

	// Call getNextJobs with n = 1. We expect the one job returned to be the
	// highpriority one, but the status should now be executing
	jobs, err := getNextJobs(1)
	if err != nil {
		t.Errorf("Unexpected error from getNextJobs: %s", err.Error())
	}
	if len(jobs) != 1 {
		t.Errorf("Length of jobs was incorrect. Expected 1 but got %d", len(jobs))
	}
	gotJob := jobs[0]
	expectedJob := &Job{}
	(*expectedJob) = *highPriorityJob
	expectedJob.status = StatusExecuting
	if !reflect.DeepEqual(expectedJob, gotJob) {
		t.Errorf("Job returned by getNextJobs was incorrect.\n\tExpected: %+v\n\tBut got:  %+v", expectedJob, gotJob)
	}
}

// TestJobsWithHigherPriorityExecutedFirst creates two sets of jobs: one with lower priorities
// and one with higher priorities. Then it starts the worker pool and runs for exactly one iteration.
// Then it makes sure that the jobs with higher priorities were executed, and the lower priority ones
// were not.
func TestJobsWithHigherPriorityExecutedFirst(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	// Register some jobs which will simply set one of the values in data
	data := make([]string, 8)
	setStringJob, err := RegisterJobType("setString", func(i int) {
		data[i] = "ok"
	})
	if err != nil {
		t.Errorf("Unexpected error in RegisterJobType: %s", err.Error())
	}

	// Queue up some jobs
	queuedJobs := make([]*Job, len(data))
	for i := 0; i < len(data); i++ {
		// Lower indexes have higher priority and should be completed first
		job, err := setStringJob.Enqueue(8-i, time.Now(), i)
		if err != nil {
			t.Errorf("Unexpected error in Enqueue: %s", err.Error())
		}
		queuedJobs[i] = job
	}

	// Start the pool with 4 workers
	runtime.GOMAXPROCS(4)
	NumWorkers = 4
	BatchSize = 4
	Pool.Start()

	// Immediately stop the pool to stop the workers from doing more jobs
	Pool.Close()

	// Wait for the workers to finish
	Pool.Wait()

	// Check that the first 4 values of data were set to "ok"
	// This would mean that the first 4 jobs (in order of priority)
	// were successfully executed.
	for i, datum := range data[0:4] {
		if datum != "ok" {
			t.Errorf(`Expected data[%d] to be set to "ok" but got: "%s"`, i, datum)
		}
	}

	// Make sure all the other values of data are still blank
	for i, datum := range data[4:] {
		if datum != "" {
			t.Errorf(`Expected data[%d] to be set to "" but got: "%s"`, i, datum)
		}
	}

	// Make sure the first four jobs we queued are marked as finished
	for _, job := range queuedJobs[0:4] {
		// Since we don't have a fresh copy, set the status manually. I.e. there is
		// a difference between the reference we have to the job and what actually exists
		// in the database. The database is what we care about.
		// assertJobStatusEquals will check that the job is correct in the database.
		job.status = StatusFinished
		assertJobStatusEquals(t, job, StatusFinished)
	}

	// Make sure the next four jobs we queued are marked as queued
	for _, job := range queuedJobs[4:] {
		// Since we don't have a fresh copy, set the status manually. I.e. there is
		// a difference between the reference we have to the job and what actually exists
		// in the database. The database is what we care about.
		// assertJobStatusEquals will check that the job is correct in the database.
		job.status = StatusQueued
		assertJobStatusEquals(t, job, StatusQueued)
	}
}

// TestJobsOnlyExecutedOnce creates a few jobs that increment a counter (each job
// has its own counter). Then it starts the pool and runs the query loop for at most two
// iterations. Then it checks that each job was executed only once by observing the counters.
func TestJobsOnlyExecutedOnce(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	// Register some jobs which will simply increment one of the values in data
	data := make([]int, 4)
	waitForJobs := sync.WaitGroup{}
	incrementJob, err := RegisterJobType("increment", func(i int) {
		data[i] += 1
		waitForJobs.Done()
	})
	if err != nil {
		t.Errorf("Unexpected error in RegisterJobType: %s", err.Error())
	}

	// Queue up some jobs
	for i := 0; i < len(data); i++ {
		waitForJobs.Add(1)
		if _, err := incrementJob.Enqueue(100, time.Now(), i); err != nil {
			t.Errorf("Unexpected error in Enqueue: %s", err.Error())
		}
	}

	// Start the pool with 4 workers
	runtime.GOMAXPROCS(4)
	NumWorkers = 4
	BatchSize = 4
	Pool.Start()

	// Wait for the wait group, which tells us each job was executed at least once
	waitForJobs.Wait()
	// Close the pool, allowing for a max of one more iteration
	Pool.Close()
	// Wait for the workers to finish
	Pool.Wait()

	// Check that each value in data equals 1.
	// This would mean that each job was only executed once
	for i, datum := range data {
		if datum != 1 {
			t.Errorf(`Expected data[%d] to be 1 but got: %d`, i, datum)
		}
	}
}

// TestAllJobsExecuted creates many more jobs than workers. Then it starts
// the pool and continuously checks if every job was executed, it which case
// it exits successfully. If some of the jobs have not been executed after 1
// second, it breaks and reports an error. 1 second should be plenty of time
// to execute the jobs.
func TestAllJobsExecuted(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	defer func() {
		// Close the pool and wait for workers to finish
		Pool.Close()
		Pool.Wait()
	}()

	// Register some jobs which will simply increment one of the values in data
	data := make([]string, 100)
	setStringJob, err := RegisterJobType("setString", func(i int) {
		data[i] = "ok"
	})
	if err != nil {
		t.Errorf("Unexpected error in RegisterJobType: %s", err.Error())
	}

	// Queue up some jobs
	for i := 0; i < len(data); i++ {
		if _, err := setStringJob.Enqueue(100, time.Now(), i); err != nil {
			t.Errorf("Unexpected error in Enqueue: %s", err.Error())
		}
	}

	// Start the pool with 4 workers
	runtime.GOMAXPROCS(4)
	NumWorkers = 4
	BatchSize = 4
	Pool.Start()

	// Continuously check the data every 10 milliseconds. Eventually
	// we hope to see that everything was set to "ok". If 1 second has
	// passed, assume something went wrong.
	timeout := time.After(1 * time.Second)
	interval := time.Tick(10 * time.Millisecond)
	remainingJobs := len(data)
	for {
		select {
		case <-timeout:
			// More than 1 second has passed. Assume something went wrong.
			t.Errorf("1 second passed and %d jobs were not executed.", remainingJobs)
			break
		case <-interval:
			// Count the number of elements in data that equal "ok".
			// Anything that doesn't equal ok represents a job that hasn't been executed yet
			remainingJobs = len(data)
			for _, datum := range data {
				if datum == "ok" {
					remainingJobs -= 1
				}
			}
			if remainingJobs == 0 {
				// Each item in data was set to "ok", so all the jobs were executed correctly.
				return
			}
		}
	}
}
<commit_msg>Remove TODO from getNextJobs since the cases are covered in other tests.<commit_after>package zazu

import (
	"reflect"
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"
)

// TestNextJobs tests the getNextJobs function, which queries the database to find
// the next queued jobs, in order of their priority.
func TestGetNextJobs(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	// Create a test job with high priority
	highPriorityJob, err := createTestJob()
	if err != nil {
		t.Errorf("Unexpected error creating test job: %s", err.Error())
	}
	highPriorityJob.priority = 1000
	highPriorityJob.id = "highPriorityJob"
	if err := highPriorityJob.save(); err != nil {
		t.Errorf("Unexpected error saving test job: %s", err.Error())
	}
	if err := highPriorityJob.Enqueue(); err != nil {
		t.Errorf("Unexpected error enqueuing test job: %s", err.Error())
	}

	// Create more tests with lower priorities
	for i := 0; i < 10; i++ {
		job, err := createTestJob()
		if err != nil {
			t.Errorf("Unexpected error creating test job: %s", err.Error())
		}
		job.priority = 100
		job.id = "lowPriorityJob" + strconv.Itoa(i)
		if err := job.save(); err != nil {
			t.Errorf("Unexpected error saving test job: %s", err.Error())
		}
		if err := job.Enqueue(); err != nil {
			t.Errorf("Unexpected error enqueuing test job: %s", err.Error())
		}
	}

	// Call getNextJobs with n = 1. We expect the one job returned to be the
	// highpriority one, but the status should now be executing
	jobs, err := getNextJobs(1)
	if err != nil {
		t.Errorf("Unexpected error from getNextJobs: %s", err.Error())
	}
	if len(jobs) != 1 {
		t.Errorf("Length of jobs was incorrect. Expected 1 but got %d", len(jobs))
	}
	gotJob := jobs[0]
	expectedJob := &Job{}
	(*expectedJob) = *highPriorityJob
	expectedJob.status = StatusExecuting
	if !reflect.DeepEqual(expectedJob, gotJob) {
		t.Errorf("Job returned by getNextJobs was incorrect.\n\tExpected: %+v\n\tBut got:  %+v", expectedJob, gotJob)
	}
}

// TestJobsWithHigherPriorityExecutedFirst creates two sets of jobs: one with lower priorities
// and one with higher priorities. Then it starts the worker pool and runs for exactly one iteration.
// Then it makes sure that the jobs with higher priorities were executed, and the lower priority ones
// were not.
func TestJobsWithHigherPriorityExecutedFirst(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	// Register some jobs which will simply set one of the values in data
	data := make([]string, 8)
	setStringJob, err := RegisterJobType("setString", func(i int) {
		data[i] = "ok"
	})
	if err != nil {
		t.Errorf("Unexpected error in RegisterJobType: %s", err.Error())
	}

	// Queue up some jobs
	queuedJobs := make([]*Job, len(data))
	for i := 0; i < len(data); i++ {
		// Lower indexes have higher priority and should be completed first
		job, err := setStringJob.Enqueue(8-i, time.Now(), i)
		if err != nil {
			t.Errorf("Unexpected error in Enqueue: %s", err.Error())
		}
		queuedJobs[i] = job
	}

	// Start the pool with 4 workers
	runtime.GOMAXPROCS(4)
	NumWorkers = 4
	BatchSize = 4
	Pool.Start()

	// Immediately stop the pool to stop the workers from doing more jobs
	Pool.Close()

	// Wait for the workers to finish
	Pool.Wait()

	// Check that the first 4 values of data were set to "ok"
	// This would mean that the first 4 jobs (in order of priority)
	// were successfully executed.
	for i, datum := range data[0:4] {
		if datum != "ok" {
			t.Errorf(`Expected data[%d] to be set to "ok" but got: "%s"`, i, datum)
		}
	}

	// Make sure all the other values of data are still blank
	for i, datum := range data[4:] {
		if datum != "" {
			t.Errorf(`Expected data[%d] to be set to "" but got: "%s"`, i, datum)
		}
	}

	// Make sure the first four jobs we queued are marked as finished
	for _, job := range queuedJobs[0:4] {
		// Since we don't have a fresh copy, set the status manually. I.e. there is
		// a difference between the reference we have to the job and what actually exists
		// in the database. The database is what we care about.
		// assertJobStatusEquals will check that the job is correct in the database.
		job.status = StatusFinished
		assertJobStatusEquals(t, job, StatusFinished)
	}

	// Make sure the next four jobs we queued are marked as queued
	for _, job := range queuedJobs[4:] {
		// Since we don't have a fresh copy, set the status manually. I.e. there is
		// a difference between the reference we have to the job and what actually exists
		// in the database. The database is what we care about.
		// assertJobStatusEquals will check that the job is correct in the database.
		job.status = StatusQueued
		assertJobStatusEquals(t, job, StatusQueued)
	}
}

// TestJobsOnlyExecutedOnce creates a few jobs that increment a counter (each job
// has its own counter). Then it starts the pool and runs the query loop for at most two
// iterations. Then it checks that each job was executed only once by observing the counters.
func TestJobsOnlyExecutedOnce(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	// Register some jobs which will simply increment one of the values in data
	data := make([]int, 4)
	waitForJobs := sync.WaitGroup{}
	incrementJob, err := RegisterJobType("increment", func(i int) {
		data[i] += 1
		waitForJobs.Done()
	})
	if err != nil {
		t.Errorf("Unexpected error in RegisterJobType: %s", err.Error())
	}

	// Queue up some jobs
	for i := 0; i < len(data); i++ {
		waitForJobs.Add(1)
		if _, err := incrementJob.Enqueue(100, time.Now(), i); err != nil {
			t.Errorf("Unexpected error in Enqueue: %s", err.Error())
		}
	}

	// Start the pool with 4 workers
	runtime.GOMAXPROCS(4)
	NumWorkers = 4
	BatchSize = 4
	Pool.Start()

	// Wait for the wait group, which tells us each job was executed at least once
	waitForJobs.Wait()
	// Close the pool, allowing for a max of one more iteration
	Pool.Close()
	// Wait for the workers to finish
	Pool.Wait()

	// Check that each value in data equals 1.
	// This would mean that each job was only executed once
	for i, datum := range data {
		if datum != 1 {
			t.Errorf(`Expected data[%d] to be 1 but got: %d`, i, datum)
		}
	}
}

// TestAllJobsExecuted creates many more jobs than workers. Then it starts
// the pool and continuously checks if every job was executed, it which case
// it exits successfully. If some of the jobs have not been executed after 1
// second, it breaks and reports an error. 1 second should be plenty of time
// to execute the jobs.
func TestAllJobsExecuted(t *testing.T) {
	flushdb()
	jobTypes = map[string]*JobType{}

	defer func() {
		// Close the pool and wait for workers to finish
		Pool.Close()
		Pool.Wait()
	}()

	// Register some jobs which will simply increment one of the values in data
	data := make([]string, 100)
	setStringJob, err := RegisterJobType("setString", func(i int) {
		data[i] = "ok"
	})
	if err != nil {
		t.Errorf("Unexpected error in RegisterJobType: %s", err.Error())
	}

	// Queue up some jobs
	for i := 0; i < len(data); i++ {
		if _, err := setStringJob.Enqueue(100, time.Now(), i); err != nil {
			t.Errorf("Unexpected error in Enqueue: %s", err.Error())
		}
	}

	// Start the pool with 4 workers
	runtime.GOMAXPROCS(4)
	NumWorkers = 4
	BatchSize = 4
	Pool.Start()

	// Continuously check the data every 10 milliseconds. Eventually
	// we hope to see that everything was set to "ok". If 1 second has
	// passed, assume something went wrong.
	timeout := time.After(1 * time.Second)
	interval := time.Tick(10 * time.Millisecond)
	remainingJobs := len(data)
	for {
		select {
		case <-timeout:
			// More than 1 second has passed. Assume something went wrong.
			t.Errorf("1 second passed and %d jobs were not executed.", remainingJobs)
			break
		case <-interval:
			// Count the number of elements in data that equal "ok".
			// Anything that doesn't equal ok represents a job that hasn't been executed yet
			remainingJobs = len(data)
			for _, datum := range data {
				if datum == "ok" {
					remainingJobs -= 1
				}
			}
			if remainingJobs == 0 {
				// Each item in data was set to "ok", so all the jobs were executed correctly.
				return
			}
		}
	}
}
<|endoftext|>