<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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 annotate

import (
	"bytes"
	"fmt"
	"io"

	jsonpatch "github.com/evanphx/json-patch"
	"github.com/spf13/cobra"
	"k8s.io/klog"

	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/json"

	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/polymorphichelpers"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

// AnnotateOptions have the data required to perform the annotate operation
type AnnotateOptions struct {
	PrintFlags *genericclioptions.PrintFlags
	PrintObj   printers.ResourcePrinterFunc

	// Filename options
	resource.FilenameOptions
	RecordFlags *genericclioptions.RecordFlags

	// Common user flags
	overwrite       bool
	local           bool
	dryrun          bool
	all             bool
	resourceVersion string
	selector        string
	fieldSelector   string
	outputFormat    string

	// results of arg parsing
	resources                    []string
	newAnnotations               map[string]string
	removeAnnotations            []string
	Recorder                     genericclioptions.Recorder
	namespace                    string
	enforceNamespace             bool
	builder                      *resource.Builder
	unstructuredClientForMapping func(mapping *meta.RESTMapping) (resource.RESTClient, error)

	genericclioptions.IOStreams
}

var (
	annotateLong = templates.LongDesc(`
		Update the annotations on one or more resources

		All Kubernetes objects support the ability to store additional data with the object as
		annotations. Annotations are key/value pairs that can be larger than labels and include
		arbitrary string values such as structured JSON. Tools and system extensions may use
		annotations to store their own data.

		Attempting to set an annotation that already exists will fail unless --overwrite is set.
		If --resource-version is specified and does not match the current resource version on
		the server the command will fail.`)

	annotateExample = templates.Examples(i18n.T(`
    # Update pod 'foo' with the annotation 'description' and the value 'my frontend'.
    # If the same annotation is set multiple times, only the last value will be applied
    kubectl annotate pods foo description='my frontend'

    # Update a pod identified by type and name in "pod.json"
    kubectl annotate -f pod.json description='my frontend'

    # Update pod 'foo' with the annotation 'description' and the value 'my frontend running nginx', overwriting any existing value.
    kubectl annotate --overwrite pods foo description='my frontend running nginx'

    # Update all pods in the namespace
    kubectl annotate pods --all description='my frontend running nginx'

    # Update pod 'foo' only if the resource is unchanged from version 1.
    kubectl annotate pods foo description='my frontend running nginx' --resource-version=1

    # Update pod 'foo' by removing an annotation named 'description' if it exists.
    # Does not require the --overwrite flag.
    kubectl annotate pods foo description-`))
)

// NewAnnotateOptions creates the options for annotate
func NewAnnotateOptions(ioStreams genericclioptions.IOStreams) *AnnotateOptions {
	return &AnnotateOptions{
		PrintFlags: genericclioptions.NewPrintFlags("annotated").WithTypeSetter(scheme.Scheme),

		RecordFlags: genericclioptions.NewRecordFlags(),
		Recorder:    genericclioptions.NoopRecorder{},
		IOStreams:   ioStreams,
	}
}

// NewCmdAnnotate creates the `annotate` command
func NewCmdAnnotate(parent string, f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewAnnotateOptions(ioStreams)

	cmd := &cobra.Command{
		Use:                   "annotate [--overwrite] (-f FILENAME | TYPE NAME) KEY_1=VAL_1 ... KEY_N=VAL_N [--resource-version=version]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Update the annotations on a resource"),
		Long:                  annotateLong + "\n\n" + cmdutil.SuggestAPIResources(parent),
		Example:               annotateExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunAnnotate())
		},
	}

	// bind flag structs
	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)

	cmd.Flags().BoolVar(&o.overwrite, "overwrite", o.overwrite, "If true, allow annotations to be overwritten, otherwise reject annotation updates that overwrite existing annotations.")
	cmd.Flags().BoolVar(&o.local, "local", o.local, "If true, annotation will NOT contact api-server but run locally.")
	cmd.Flags().StringVarP(&o.selector, "selector", "l", o.selector, "Selector (label query) to filter on, not including uninitialized ones, supports '=', '==', and '!='.(e.g. -l key1=value1,key2=value2).")
	cmd.Flags().StringVar(&o.fieldSelector, "field-selector", o.fieldSelector, "Selector (field query) to filter on, supports '=', '==', and '!='.(e.g. --field-selector key1=value1,key2=value2). The server only supports a limited number of field queries per type.")
	cmd.Flags().BoolVar(&o.all, "all", o.all, "Select all resources, including uninitialized ones, in the namespace of the specified resource types.")
	cmd.Flags().StringVar(&o.resourceVersion, "resource-version", o.resourceVersion, i18n.T("If non-empty, the annotation update will only succeed if this is the current resource-version for the object. Only valid when specifying a single resource."))
	usage := "identifying the resource to update the annotation"
	cmdutil.AddFilenameOptionFlags(cmd, &o.FilenameOptions, usage)
	cmdutil.AddDryRunFlag(cmd)

	return cmd
}

// Complete adapts from the command line args and factory to the data required.
func (o *AnnotateOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error

	o.RecordFlags.Complete(cmd)
	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	o.outputFormat = cmdutil.GetFlagString(cmd, "output")
	o.dryrun = cmdutil.GetDryRunFlag(cmd)

	if o.dryrun {
		o.PrintFlags.Complete("%s (dry run)")
	}
	printer, err := o.PrintFlags.ToPrinter()
	if err != nil {
		return err
	}
	o.PrintObj = func(obj runtime.Object, out io.Writer) error {
		return printer.PrintObj(obj, out)
	}

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.builder = f.NewBuilder()
	o.unstructuredClientForMapping = f.UnstructuredClientForMapping

	// retrieves resource and annotation args from args
	// also checks args to verify that all resources are specified before annotations
	resources, annotationArgs, err := cmdutil.GetResourcesAndPairs(args, "annotation")
	if err != nil {
		return err
	}
	o.resources = resources
	o.newAnnotations, o.removeAnnotations, err = parseAnnotations(annotationArgs)
	if err != nil {
		return err
	}

	return nil
}

// Validate checks to the AnnotateOptions to see if there is sufficient information run the command.
func (o AnnotateOptions) Validate() error {
	if o.all && len(o.selector) > 0 {
		return fmt.Errorf("cannot set --all and --selector at the same time")
	}
	if o.all && len(o.fieldSelector) > 0 {
		return fmt.Errorf("cannot set --all and --field-selector at the same time")
	}
	if len(o.resources) < 1 && cmdutil.IsFilenameSliceEmpty(o.Filenames, o.Kustomize) {
		return fmt.Errorf("one or more resources must be specified as <resource> <name> or <resource>/<name>")
	}
	if len(o.newAnnotations) < 1 && len(o.removeAnnotations) < 1 {
		return fmt.Errorf("at least one annotation update is required")
	}
	return validateAnnotations(o.removeAnnotations, o.newAnnotations)
}

// RunAnnotate does the work
func (o AnnotateOptions) RunAnnotate() error {
	b := o.builder.
		Unstructured().
		LocalParam(o.local).
		ContinueOnError().
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, &o.FilenameOptions).
		Flatten()

	if !o.local {
		b = b.LabelSelectorParam(o.selector).
			FieldSelectorParam(o.fieldSelector).
			ResourceTypeOrNameArgs(o.all, o.resources...).
			Latest()
	}

	r := b.Do()
	if err := r.Err(); err != nil {
		return err
	}

	var singleItemImpliedResource bool
	r.IntoSingleItemImplied(&singleItemImpliedResource)

	// only apply resource version locking on a single resource.
	// we must perform this check after o.builder.Do() as
	// []o.resources can not accurately return the proper number
	// of resources when they are not passed in "resource/name" format.
	if !singleItemImpliedResource && len(o.resourceVersion) > 0 {
		return fmt.Errorf("--resource-version may only be used with a single resource")
	}

	return r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}

		var outputObj runtime.Object
		obj := info.Object

		if o.dryrun || o.local {
			if err := o.updateAnnotations(obj); err != nil {
				return err
			}
			outputObj = obj
		} else {
			name, namespace := info.Name, info.Namespace

			if len(o.resourceVersion) != 0 {
				// ensure resourceVersion is always sent in the patch by clearing it from the starting JSON
				accessor, err := meta.Accessor(obj)
				if err != nil {
					return err
				}
				accessor.SetResourceVersion("")
			}

			oldData, err := json.Marshal(obj)
			if err != nil {
				return err
			}
			if err := o.Recorder.Record(info.Object); err != nil {
				klog.V(4).Infof("error recording current command: %v", err)
			}
			if err := o.updateAnnotations(obj); err != nil {
				return err
			}
			newData, err := json.Marshal(obj)
			if err != nil {
				return err
			}
			patchBytes, err := jsonpatch.CreateMergePatch(oldData, newData)
			createdPatch := err == nil
			if err != nil {
				klog.V(2).Infof("couldn't compute patch: %v", err)
			}

			mapping := info.ResourceMapping()
			client, err := o.unstructuredClientForMapping(mapping)
			if err != nil {
				return err
			}
			helper := resource.NewHelper(client, mapping)

			if createdPatch {
				outputObj, err = helper.Patch(namespace, name, types.MergePatchType, patchBytes, nil)
			} else {
				outputObj, err = helper.Replace(namespace, name, false, obj)
			}
			if err != nil {
				return err
			}
		}

		return o.PrintObj(outputObj, o.Out)
	})
}

// parseAnnotations retrieves new and remove annotations from annotation args
func parseAnnotations(annotationArgs []string) (map[string]string, []string, error) {
	return cmdutil.ParsePairs(annotationArgs, "annotation", true)
}

// validateAnnotations checks the format of annotation args and checks removed annotations aren't in the new annotations map
func validateAnnotations(removeAnnotations []string, newAnnotations map[string]string) error {
	var modifyRemoveBuf bytes.Buffer
	for _, removeAnnotation := range removeAnnotations {
		if _, found := newAnnotations[removeAnnotation]; found {
			if modifyRemoveBuf.Len() > 0 {
				modifyRemoveBuf.WriteString(", ")
			}
			modifyRemoveBuf.WriteString(fmt.Sprintf(removeAnnotation))
		}
	}
	if modifyRemoveBuf.Len() > 0 {
		return fmt.Errorf("can not both modify and remove the following annotation(s) in the same command: %s", modifyRemoveBuf.String())
	}

	return nil
}

// validateNoAnnotationOverwrites validates that when overwrite is false, to-be-updated annotations don't exist in the object annotation map (yet)
func validateNoAnnotationOverwrites(accessor metav1.Object, annotations map[string]string) error {
	var buf bytes.Buffer
	for key := range annotations {
		// change-cause annotation can always be overwritten
		if key == polymorphichelpers.ChangeCauseAnnotation {
			continue
		}
		if value, found := accessor.GetAnnotations()[key]; found {
			if buf.Len() > 0 {
				buf.WriteString("; ")
			}
			buf.WriteString(fmt.Sprintf("'%s' already has a value (%s)", key, value))
		}
	}
	if buf.Len() > 0 {
		return fmt.Errorf("--overwrite is false but found the following declared annotation(s): %s", buf.String())
	}
	return nil
}

// updateAnnotations updates annotations of obj
func (o AnnotateOptions) updateAnnotations(obj runtime.Object) error {
	accessor, err := meta.Accessor(obj)
	if err != nil {
		return err
	}
	if !o.overwrite {
		if err := validateNoAnnotationOverwrites(accessor, o.newAnnotations); err != nil {
			return err
		}
	}

	annotations := accessor.GetAnnotations()
	if annotations == nil {
		annotations = make(map[string]string)
	}

	for key, value := range o.newAnnotations {
		annotations[key] = value
	}
	for _, annotation := range o.removeAnnotations {
		delete(annotations, annotation)
	}
	accessor.SetAnnotations(annotations)

	if len(o.resourceVersion) != 0 {
		accessor.SetResourceVersion(o.resourceVersion)
	}
	return nil
}
<commit_msg>fix static check in kubectl/pkg/cmd/annotate.<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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 annotate

import (
	"bytes"
	"fmt"
	"io"

	jsonpatch "github.com/evanphx/json-patch"
	"github.com/spf13/cobra"
	"k8s.io/klog"

	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/json"

	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/polymorphichelpers"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

// AnnotateOptions have the data required to perform the annotate operation
type AnnotateOptions struct {
	PrintFlags *genericclioptions.PrintFlags
	PrintObj   printers.ResourcePrinterFunc

	// Filename options
	resource.FilenameOptions
	RecordFlags *genericclioptions.RecordFlags

	// Common user flags
	overwrite       bool
	local           bool
	dryrun          bool
	all             bool
	resourceVersion string
	selector        string
	fieldSelector   string
	outputFormat    string

	// results of arg parsing
	resources                    []string
	newAnnotations               map[string]string
	removeAnnotations            []string
	Recorder                     genericclioptions.Recorder
	namespace                    string
	enforceNamespace             bool
	builder                      *resource.Builder
	unstructuredClientForMapping func(mapping *meta.RESTMapping) (resource.RESTClient, error)

	genericclioptions.IOStreams
}

var (
	annotateLong = templates.LongDesc(`
		Update the annotations on one or more resources

		All Kubernetes objects support the ability to store additional data with the object as
		annotations. Annotations are key/value pairs that can be larger than labels and include
		arbitrary string values such as structured JSON. Tools and system extensions may use
		annotations to store their own data.

		Attempting to set an annotation that already exists will fail unless --overwrite is set.
		If --resource-version is specified and does not match the current resource version on
		the server the command will fail.`)

	annotateExample = templates.Examples(i18n.T(`
    # Update pod 'foo' with the annotation 'description' and the value 'my frontend'.
    # If the same annotation is set multiple times, only the last value will be applied
    kubectl annotate pods foo description='my frontend'

    # Update a pod identified by type and name in "pod.json"
    kubectl annotate -f pod.json description='my frontend'

    # Update pod 'foo' with the annotation 'description' and the value 'my frontend running nginx', overwriting any existing value.
    kubectl annotate --overwrite pods foo description='my frontend running nginx'

    # Update all pods in the namespace
    kubectl annotate pods --all description='my frontend running nginx'

    # Update pod 'foo' only if the resource is unchanged from version 1.
    kubectl annotate pods foo description='my frontend running nginx' --resource-version=1

    # Update pod 'foo' by removing an annotation named 'description' if it exists.
    # Does not require the --overwrite flag.
    kubectl annotate pods foo description-`))
)

// NewAnnotateOptions creates the options for annotate
func NewAnnotateOptions(ioStreams genericclioptions.IOStreams) *AnnotateOptions {
	return &AnnotateOptions{
		PrintFlags: genericclioptions.NewPrintFlags("annotated").WithTypeSetter(scheme.Scheme),

		RecordFlags: genericclioptions.NewRecordFlags(),
		Recorder:    genericclioptions.NoopRecorder{},
		IOStreams:   ioStreams,
	}
}

// NewCmdAnnotate creates the `annotate` command
func NewCmdAnnotate(parent string, f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewAnnotateOptions(ioStreams)

	cmd := &cobra.Command{
		Use:                   "annotate [--overwrite] (-f FILENAME | TYPE NAME) KEY_1=VAL_1 ... KEY_N=VAL_N [--resource-version=version]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Update the annotations on a resource"),
		Long:                  annotateLong + "\n\n" + cmdutil.SuggestAPIResources(parent),
		Example:               annotateExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunAnnotate())
		},
	}

	// bind flag structs
	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)

	cmd.Flags().BoolVar(&o.overwrite, "overwrite", o.overwrite, "If true, allow annotations to be overwritten, otherwise reject annotation updates that overwrite existing annotations.")
	cmd.Flags().BoolVar(&o.local, "local", o.local, "If true, annotation will NOT contact api-server but run locally.")
	cmd.Flags().StringVarP(&o.selector, "selector", "l", o.selector, "Selector (label query) to filter on, not including uninitialized ones, supports '=', '==', and '!='.(e.g. -l key1=value1,key2=value2).")
	cmd.Flags().StringVar(&o.fieldSelector, "field-selector", o.fieldSelector, "Selector (field query) to filter on, supports '=', '==', and '!='.(e.g. --field-selector key1=value1,key2=value2). The server only supports a limited number of field queries per type.")
	cmd.Flags().BoolVar(&o.all, "all", o.all, "Select all resources, including uninitialized ones, in the namespace of the specified resource types.")
	cmd.Flags().StringVar(&o.resourceVersion, "resource-version", o.resourceVersion, i18n.T("If non-empty, the annotation update will only succeed if this is the current resource-version for the object. Only valid when specifying a single resource."))
	usage := "identifying the resource to update the annotation"
	cmdutil.AddFilenameOptionFlags(cmd, &o.FilenameOptions, usage)
	cmdutil.AddDryRunFlag(cmd)

	return cmd
}

// Complete adapts from the command line args and factory to the data required.
func (o *AnnotateOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error

	o.RecordFlags.Complete(cmd)
	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	o.outputFormat = cmdutil.GetFlagString(cmd, "output")
	o.dryrun = cmdutil.GetDryRunFlag(cmd)

	if o.dryrun {
		o.PrintFlags.Complete("%s (dry run)")
	}
	printer, err := o.PrintFlags.ToPrinter()
	if err != nil {
		return err
	}
	o.PrintObj = func(obj runtime.Object, out io.Writer) error {
		return printer.PrintObj(obj, out)
	}

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.builder = f.NewBuilder()
	o.unstructuredClientForMapping = f.UnstructuredClientForMapping

	// retrieves resource and annotation args from args
	// also checks args to verify that all resources are specified before annotations
	resources, annotationArgs, err := cmdutil.GetResourcesAndPairs(args, "annotation")
	if err != nil {
		return err
	}
	o.resources = resources
	o.newAnnotations, o.removeAnnotations, err = parseAnnotations(annotationArgs)
	if err != nil {
		return err
	}

	return nil
}

// Validate checks to the AnnotateOptions to see if there is sufficient information run the command.
func (o AnnotateOptions) Validate() error {
	if o.all && len(o.selector) > 0 {
		return fmt.Errorf("cannot set --all and --selector at the same time")
	}
	if o.all && len(o.fieldSelector) > 0 {
		return fmt.Errorf("cannot set --all and --field-selector at the same time")
	}
	if len(o.resources) < 1 && cmdutil.IsFilenameSliceEmpty(o.Filenames, o.Kustomize) {
		return fmt.Errorf("one or more resources must be specified as <resource> <name> or <resource>/<name>")
	}
	if len(o.newAnnotations) < 1 && len(o.removeAnnotations) < 1 {
		return fmt.Errorf("at least one annotation update is required")
	}
	return validateAnnotations(o.removeAnnotations, o.newAnnotations)
}

// RunAnnotate does the work
func (o AnnotateOptions) RunAnnotate() error {
	b := o.builder.
		Unstructured().
		LocalParam(o.local).
		ContinueOnError().
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, &o.FilenameOptions).
		Flatten()

	if !o.local {
		b = b.LabelSelectorParam(o.selector).
			FieldSelectorParam(o.fieldSelector).
			ResourceTypeOrNameArgs(o.all, o.resources...).
			Latest()
	}

	r := b.Do()
	if err := r.Err(); err != nil {
		return err
	}

	var singleItemImpliedResource bool
	r.IntoSingleItemImplied(&singleItemImpliedResource)

	// only apply resource version locking on a single resource.
	// we must perform this check after o.builder.Do() as
	// []o.resources can not accurately return the proper number
	// of resources when they are not passed in "resource/name" format.
	if !singleItemImpliedResource && len(o.resourceVersion) > 0 {
		return fmt.Errorf("--resource-version may only be used with a single resource")
	}

	return r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}

		var outputObj runtime.Object
		obj := info.Object

		if o.dryrun || o.local {
			if err := o.updateAnnotations(obj); err != nil {
				return err
			}
			outputObj = obj
		} else {
			name, namespace := info.Name, info.Namespace

			if len(o.resourceVersion) != 0 {
				// ensure resourceVersion is always sent in the patch by clearing it from the starting JSON
				accessor, err := meta.Accessor(obj)
				if err != nil {
					return err
				}
				accessor.SetResourceVersion("")
			}

			oldData, err := json.Marshal(obj)
			if err != nil {
				return err
			}
			if err := o.Recorder.Record(info.Object); err != nil {
				klog.V(4).Infof("error recording current command: %v", err)
			}
			if err := o.updateAnnotations(obj); err != nil {
				return err
			}
			newData, err := json.Marshal(obj)
			if err != nil {
				return err
			}
			patchBytes, err := jsonpatch.CreateMergePatch(oldData, newData)
			createdPatch := err == nil
			if err != nil {
				klog.V(2).Infof("couldn't compute patch: %v", err)
			}

			mapping := info.ResourceMapping()
			client, err := o.unstructuredClientForMapping(mapping)
			if err != nil {
				return err
			}
			helper := resource.NewHelper(client, mapping)

			if createdPatch {
				outputObj, err = helper.Patch(namespace, name, types.MergePatchType, patchBytes, nil)
			} else {
				outputObj, err = helper.Replace(namespace, name, false, obj)
			}
			if err != nil {
				return err
			}
		}

		return o.PrintObj(outputObj, o.Out)
	})
}

// parseAnnotations retrieves new and remove annotations from annotation args
func parseAnnotations(annotationArgs []string) (map[string]string, []string, error) {
	return cmdutil.ParsePairs(annotationArgs, "annotation", true)
}

// validateAnnotations checks the format of annotation args and checks removed annotations aren't in the new annotations map
func validateAnnotations(removeAnnotations []string, newAnnotations map[string]string) error {
	var modifyRemoveBuf bytes.Buffer
	for _, removeAnnotation := range removeAnnotations {
		if _, found := newAnnotations[removeAnnotation]; found {
			if modifyRemoveBuf.Len() > 0 {
				modifyRemoveBuf.WriteString(", ")
			}
			modifyRemoveBuf.WriteString(fmt.Sprint(removeAnnotation))
		}
	}
	if modifyRemoveBuf.Len() > 0 {
		return fmt.Errorf("can not both modify and remove the following annotation(s) in the same command: %s", modifyRemoveBuf.String())
	}

	return nil
}

// validateNoAnnotationOverwrites validates that when overwrite is false, to-be-updated annotations don't exist in the object annotation map (yet)
func validateNoAnnotationOverwrites(accessor metav1.Object, annotations map[string]string) error {
	var buf bytes.Buffer
	for key := range annotations {
		// change-cause annotation can always be overwritten
		if key == polymorphichelpers.ChangeCauseAnnotation {
			continue
		}
		if value, found := accessor.GetAnnotations()[key]; found {
			if buf.Len() > 0 {
				buf.WriteString("; ")
			}
			buf.WriteString(fmt.Sprintf("'%s' already has a value (%s)", key, value))
		}
	}
	if buf.Len() > 0 {
		return fmt.Errorf("--overwrite is false but found the following declared annotation(s): %s", buf.String())
	}
	return nil
}

// updateAnnotations updates annotations of obj
func (o AnnotateOptions) updateAnnotations(obj runtime.Object) error {
	accessor, err := meta.Accessor(obj)
	if err != nil {
		return err
	}
	if !o.overwrite {
		if err := validateNoAnnotationOverwrites(accessor, o.newAnnotations); err != nil {
			return err
		}
	}

	annotations := accessor.GetAnnotations()
	if annotations == nil {
		annotations = make(map[string]string)
	}

	for key, value := range o.newAnnotations {
		annotations[key] = value
	}
	for _, annotation := range o.removeAnnotations {
		delete(annotations, annotation)
	}
	accessor.SetAnnotations(annotations)

	if len(o.resourceVersion) != 0 {
		accessor.SetResourceVersion(o.resourceVersion)
	}
	return nil
}
<|endoftext|>


<commit_before>package middleware

import (
	"net/http"

	"github.com/skygeario/skygear-server/pkg/core/logging"

	"github.com/skygeario/skygear-server/pkg/core/apiclientconfig"
	"github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authinfo"
	"github.com/skygeario/skygear-server/pkg/core/auth/session"
	"github.com/skygeario/skygear-server/pkg/core/config"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/model"
	"github.com/skygeario/skygear-server/pkg/core/skydb"
)

// AuthnMiddleware populate auth context information
type AuthnMiddleware struct {
	APIClientConfigurationProvider apiclientconfig.Provider `dependency:"APIClientConfigurationProvider"`
	AuthContextSetter              auth.ContextSetter       `dependency:"AuthContextSetter"`
	LoggerFactory                  logging.Factory          `dependency:"LoggerFactory"`
	SessionProvider                session.Provider         `dependency:"SessionProvider"`
	SessionWriter                  session.Writer           `dependency:"SessionWriter"`
	AuthInfoStore                  authinfo.Store           `dependency:"AuthInfoStore"`
	TxContext                      db.TxContext             `dependency:"TxContext"`
}

// AuthnMiddlewareFactory creates AuthnMiddleware per request.
type AuthnMiddlewareFactory struct{}

// NewInjectableMiddleware implements InjectableMiddlewareFactory.
func (f AuthnMiddlewareFactory) NewInjectableMiddleware() InjectableMiddleware {
	return &AuthnMiddleware{}
}

// Handle implements InjectableMiddleware.
func (m *AuthnMiddleware) Handle(next http.Handler) http.Handler {
	log := m.LoggerFactory.NewLogger("authn")
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var sess *auth.Session
		var authInfo *authinfo.AuthInfo
		var err error
		defer func() {
			m.AuthContextSetter.SetSessionAndAuthInfo(sess, authInfo, err)

			if err == model.ErrTokenConflict {
				// Clear session if token conflicts
				m.SessionWriter.ClearSession(w)
			} else if err == session.ErrSessionNotFound {
				// Clear session if session is not found
				m.SessionWriter.ClearSession(w)
			} else if err != nil {
				log.WithError(err).Debug("Cannot resolve session")
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			next.ServeHTTP(w, r)
		}()

		tenantConfig := config.GetTenantConfig(r)

		key := m.APIClientConfigurationProvider.GetAccessKeyByAPIKey(model.GetAPIKey(r))
		m.AuthContextSetter.SetAccessKey(key)

		accessToken, transport, err := model.GetAccessToken(r)
		if err != nil {
			return
		}

		// No access token found. Simply proceed.
		if accessToken == "" {
			return
		}

		if err = m.TxContext.BeginTx(); err != nil {
			return
		}
		defer m.TxContext.RollbackTx()

		sess, err = m.SessionProvider.GetByToken(accessToken, auth.SessionTokenKindAccessToken)
		if err != nil {
			return
		}

		if tenantConfig.UserConfig.Clients[sess.ClientID].SessionTransport != transport {
			err = session.ErrSessionNotFound
			return
		}

		ai := authinfo.AuthInfo{}
		err = m.AuthInfoStore.GetAuth(sess.UserID, &ai)
		if err != nil {
			if err == skydb.ErrUserNotFound {
				err = session.ErrSessionNotFound
			}
			return
		}
		authInfo = &ai

		// in case valid session is used, infer access key from session
		if !key.IsMasterKey() {
			key = model.NewAccessKey(sess.ClientID)
			m.AuthContextSetter.SetAccessKey(key)
		}

		// should not use new session data in context
		sessionCopy := *sess
		err = m.SessionProvider.Access(&sessionCopy)
		if err != nil {
			return
		}
	})
}
<commit_msg>Fix logging unexpected authn error in DEBUG level<commit_after>package middleware

import (
	"net/http"

	"github.com/skygeario/skygear-server/pkg/core/logging"

	"github.com/skygeario/skygear-server/pkg/core/apiclientconfig"
	"github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authinfo"
	"github.com/skygeario/skygear-server/pkg/core/auth/session"
	"github.com/skygeario/skygear-server/pkg/core/config"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/model"
	"github.com/skygeario/skygear-server/pkg/core/skydb"
)

// AuthnMiddleware populate auth context information
type AuthnMiddleware struct {
	APIClientConfigurationProvider apiclientconfig.Provider `dependency:"APIClientConfigurationProvider"`
	AuthContextSetter              auth.ContextSetter       `dependency:"AuthContextSetter"`
	LoggerFactory                  logging.Factory          `dependency:"LoggerFactory"`
	SessionProvider                session.Provider         `dependency:"SessionProvider"`
	SessionWriter                  session.Writer           `dependency:"SessionWriter"`
	AuthInfoStore                  authinfo.Store           `dependency:"AuthInfoStore"`
	TxContext                      db.TxContext             `dependency:"TxContext"`
}

// AuthnMiddlewareFactory creates AuthnMiddleware per request.
type AuthnMiddlewareFactory struct{}

// NewInjectableMiddleware implements InjectableMiddlewareFactory.
func (f AuthnMiddlewareFactory) NewInjectableMiddleware() InjectableMiddleware {
	return &AuthnMiddleware{}
}

// Handle implements InjectableMiddleware.
func (m *AuthnMiddleware) Handle(next http.Handler) http.Handler {
	log := m.LoggerFactory.NewLogger("authn")
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var sess *auth.Session
		var authInfo *authinfo.AuthInfo
		var err error
		defer func() {
			m.AuthContextSetter.SetSessionAndAuthInfo(sess, authInfo, err)

			if err == model.ErrTokenConflict {
				// Clear session if token conflicts
				m.SessionWriter.ClearSession(w)
			} else if err == session.ErrSessionNotFound {
				// Clear session if session is not found
				m.SessionWriter.ClearSession(w)
			} else if err != nil {
				log.WithError(err).Error("Cannot resolve session")
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			next.ServeHTTP(w, r)
		}()

		tenantConfig := config.GetTenantConfig(r)

		key := m.APIClientConfigurationProvider.GetAccessKeyByAPIKey(model.GetAPIKey(r))
		m.AuthContextSetter.SetAccessKey(key)

		accessToken, transport, err := model.GetAccessToken(r)
		if err != nil {
			return
		}

		// No access token found. Simply proceed.
		if accessToken == "" {
			return
		}

		if err = m.TxContext.BeginTx(); err != nil {
			return
		}
		defer m.TxContext.RollbackTx()

		sess, err = m.SessionProvider.GetByToken(accessToken, auth.SessionTokenKindAccessToken)
		if err != nil {
			return
		}

		if tenantConfig.UserConfig.Clients[sess.ClientID].SessionTransport != transport {
			err = session.ErrSessionNotFound
			return
		}

		ai := authinfo.AuthInfo{}
		err = m.AuthInfoStore.GetAuth(sess.UserID, &ai)
		if err != nil {
			if err == skydb.ErrUserNotFound {
				err = session.ErrSessionNotFound
			}
			return
		}
		authInfo = &ai

		// in case valid session is used, infer access key from session
		if !key.IsMasterKey() {
			key = model.NewAccessKey(sess.ClientID)
			m.AuthContextSetter.SetAccessKey(key)
		}

		// should not use new session data in context
		sessionCopy := *sess
		err = m.SessionProvider.Access(&sessionCopy)
		if err != nil {
			return
		}
	})
}
<|endoftext|>


<commit_before>package middleware

import (
	"encoding/json"
	"net/http"
	"time"

	"github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authinfo"
	"github.com/skygeario/skygear-server/pkg/core/auth/session"
	"github.com/skygeario/skygear-server/pkg/core/config"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/handler"
	coreHttp "github.com/skygeario/skygear-server/pkg/core/http"
	"github.com/skygeario/skygear-server/pkg/core/model"
	"github.com/skygeario/skygear-server/pkg/core/skydb"
	"github.com/skygeario/skygear-server/pkg/core/skyerr"
)

// AuthnMiddleware populate auth context information
type AuthnMiddleware struct {
	AuthContextSetter auth.ContextSetter `dependency:"AuthContextSetter"`
	SessionProvider   session.Provider   `dependency:"SessionProvider"`
	AuthInfoStore     authinfo.Store     `dependency:"AuthInfoStore"`
	TxContext         db.TxContext       `dependency:"TxContext"`
}

// AuthnMiddlewareFactory creates AuthnMiddleware per request.
type AuthnMiddlewareFactory struct{}

// NewInjectableMiddleware implements InjectableMiddlewareFactory.
func (f AuthnMiddlewareFactory) NewInjectableMiddleware() InjectableMiddleware {
	return &AuthnMiddleware{}
}

// Handle implements InjectableMiddleware.
func (m *AuthnMiddleware) Handle(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		defer func() {
			if err == nil {
				next.ServeHTTP(w, r)
			} else {
				// clear session cookie if error occurred
				cookie := &http.Cookie{
					Name:    coreHttp.CookieNameSession,
					Path:    "/",
					Expires: time.Unix(0, 0),
				}
				http.SetCookie(w, cookie)

				skyErr := skyerr.NewNotAuthenticatedErr()
				httpStatus := skyerr.ErrorDefaultStatusCode(skyErr)
				response := handler.APIResponse{Err: skyErr}
				encoder := json.NewEncoder(w)
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(httpStatus)
				encoder.Encode(response)
			}
		}()

		tenantConfig := config.GetTenantConfig(r)

		key := model.CheckAccessKey(tenantConfig, model.GetAPIKey(r))
		m.AuthContextSetter.SetAccessKey(key)

		accessToken, transport, err := model.GetAccessToken(r)
		if err != nil {
			// invalid session token -> must not proceed

			if err == model.ErrTokenConflict {
				// clear session cookie if session token is conflicted
				cookie := &http.Cookie{
					Name:    coreHttp.CookieNameSession,
					Path:    "/",
					Expires: time.Unix(0, 0),
				}
				http.SetCookie(w, cookie)
			}
			return
		}

		// No access token found. Simply proceed.
		if accessToken == "" {
			return
		}

		if err = m.TxContext.BeginTx(); err != nil {
			panic(err)
		}
		defer m.TxContext.RollbackTx()

		s, err := m.SessionProvider.GetByToken(accessToken, auth.SessionTokenKindAccessToken)
		if err != nil {
			// session not found -> treat as no access token is provided
			err = nil
			return
		}

		if tenantConfig.UserConfig.Clients[s.ClientID].SessionTransport != transport {
			// inconsistent session token transport -> must not proceed
			err = session.ErrSessionNotFound
			return
		}

		authInfo := authinfo.AuthInfo{}
		err = m.AuthInfoStore.GetAuth(s.UserID, &authInfo)
		if err != nil {
			if err == skydb.ErrUserNotFound {
				// user not found -> treat as no access token is provided
				err = nil
			}
			return
		}

		// in case valid session is used, infer access key from session
		key = model.NewAccessKey(s.ClientID)
		m.AuthContextSetter.SetAccessKey(key)

		// should not use current session in context
		sessionCopy := *s
		err = m.SessionProvider.Access(&sessionCopy)
		if err != nil {
			// cannot access session -> must not proceed
			return
		}

		m.AuthContextSetter.SetSession(s)
		m.AuthContextSetter.SetAuthInfo(&authInfo)
	})
}
<commit_msg>Do not return error in authn middleware treat errors as no session found<commit_after>package middleware

import (
	"net/http"
	"time"

	"github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authinfo"
	"github.com/skygeario/skygear-server/pkg/core/auth/session"
	"github.com/skygeario/skygear-server/pkg/core/config"
	"github.com/skygeario/skygear-server/pkg/core/db"
	coreHttp "github.com/skygeario/skygear-server/pkg/core/http"
	"github.com/skygeario/skygear-server/pkg/core/model"
)

// AuthnMiddleware populate auth context information
type AuthnMiddleware struct {
	AuthContextSetter auth.ContextSetter `dependency:"AuthContextSetter"`
	SessionProvider   session.Provider   `dependency:"SessionProvider"`
	AuthInfoStore     authinfo.Store     `dependency:"AuthInfoStore"`
	TxContext         db.TxContext       `dependency:"TxContext"`
}

// AuthnMiddlewareFactory creates AuthnMiddleware per request.
type AuthnMiddlewareFactory struct{}

// NewInjectableMiddleware implements InjectableMiddlewareFactory.
func (f AuthnMiddlewareFactory) NewInjectableMiddleware() InjectableMiddleware {
	return &AuthnMiddleware{}
}

// Handle implements InjectableMiddleware.
func (m *AuthnMiddleware) Handle(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		defer func() {
			if err != nil {
				// clear session cookie if error occurred
				cookie := &http.Cookie{
					Name:    coreHttp.CookieNameSession,
					Path:    "/",
					Expires: time.Unix(0, 0),
				}
				http.SetCookie(w, cookie)
			}

			next.ServeHTTP(w, r)
		}()

		tenantConfig := config.GetTenantConfig(r)

		key := model.CheckAccessKey(tenantConfig, model.GetAPIKey(r))
		m.AuthContextSetter.SetAccessKey(key)

		accessToken, transport, err := model.GetAccessToken(r)
		if err != nil {
			return
		}

		// No access token found. Simply proceed.
		if accessToken == "" {
			return
		}

		if err = m.TxContext.BeginTx(); err != nil {
			return
		}
		defer m.TxContext.RollbackTx()

		s, err := m.SessionProvider.GetByToken(accessToken, auth.SessionTokenKindAccessToken)
		if err != nil {
			return
		}

		if tenantConfig.UserConfig.Clients[s.ClientID].SessionTransport != transport {
			err = session.ErrSessionNotFound
			return
		}

		authInfo := authinfo.AuthInfo{}
		err = m.AuthInfoStore.GetAuth(s.UserID, &authInfo)
		if err != nil {
			return
		}

		// in case valid session is used, infer access key from session
		key = model.NewAccessKey(s.ClientID)
		m.AuthContextSetter.SetAccessKey(key)

		// should not use new session data in context
		sessionCopy := *s
		err = m.SessionProvider.Access(&sessionCopy)
		if err != nil {
			return
		}

		m.AuthContextSetter.SetSession(s)
		m.AuthContextSetter.SetAuthInfo(&authInfo)
	})
}
<|endoftext|>


<commit_before>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package app

import (
	"fmt"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/queue"
	"github.com/globocom/tsuru/service"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"sync"
)

const (
	// queue actions
	regenerateApprc         = "regenerate-apprc"
	startApp                = "start-app"
	RegenerateApprcAndStart = "regenerate-apprc-start-app"
	bindService             = "bind-service"

	// name of the queue for internal messages.
	queueName = "tsuru-app"

	// Name of the queue for external messages.
	QueueName = "tsuru-app-public"
)

// ensureAppIsStarted make sure that the app and all units present in the given
// message are started.
func ensureAppIsStarted(msg *queue.Message) (App, error) {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		return a, fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	units := getUnits(&a, msg.Args[1:])
	if !a.Available() || !units.Started() {
		format := "Error handling %q for the app %q:"
		uState := units.State()
		if uState == "error" || uState == "down" {
			format += fmt.Sprintf(" units are in %q state.", uState)
			msg.Delete()
		} else {
			format += " all units must be started."
		}
		return a, fmt.Errorf(format, msg.Action, a.Name)
	}
	return a, nil
}

// bindUnit handles the bind-service message, binding a unit to all service
// instances binded to the app.
func bindUnit(msg *queue.Message) error {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		msg.Delete()
		return fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		return fmt.Errorf("Error handling %q: %s", msg.Action, err)
	}
	unit := getUnits(&a, msg.Args[1:])[0]
	var instances []service.ServiceInstance
	q := bson.M{"apps": bson.M{"$in": []string{msg.Args[0]}}}
	err = conn.ServiceInstances().Find(q).All(&instances)
	if err != nil {
		return err
	}
	for _, instance := range instances {
		_, err = instance.BindUnit(&unit)
		if err != nil {
			log.Printf("Error binding the unit %s with the service instance %s.", unit.Name, instance.Name)
		}
	}
	return nil
}

// handle is the function called by the queue handler on each message.
func handle(msg *queue.Message) {
	switch msg.Action {
	case RegenerateApprcAndStart:
		fallthrough
	case regenerateApprc:
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
			return
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
		app.serializeEnvVars()
		fallthrough
	case startApp:
		if msg.Action == regenerateApprc {
			break
		}
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		err = app.Restart(ioutil.Discard)
		if err != nil {
			log.Printf("Error handling %q. App failed to start:\n%s.", msg.Action, err)
			return
		}
		msg.Delete()
	case bindService:
		err := bindUnit(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
	default:
		log.Printf("Error handling %q: invalid action.", msg.Action)
		msg.Delete()
	}
}

// unitList is a simple slice of units, with special methods to handle state.
type unitList []Unit

// Started returns true if all units in the list is started.
func (l unitList) Started() bool {
	for _, unit := range l {
		if unit.State != string(provision.StatusStarted) {
			return false
		}
	}
	return true
}

// State returns a string if all units have the same state. Otherwise it
// returns an empty string.
func (l unitList) State() string {
	if len(l) == 0 {
		return ""
	}
	state := l[0].State
	for i := 1; i < len(l); i++ {
		if l[i].State != state {
			return ""
		}
	}
	return state
}

// getUnits builds a unitList from the given app and the names in the string
// slice.
func getUnits(app *App, names []string) unitList {
	var units []Unit
	if len(names) > 0 {
		units = make([]Unit, len(names))
		i := 0
		for _, unitName := range names {
			for _, appUnit := range app.Units {
				if appUnit.Name == unitName {
					units[i] = appUnit
					i++
					break
				}
			}
		}
	}
	return unitList(units)
}

var (
	qfactory queue.QFactory
	_queue   queue.Q
	_handler queue.Handler
	o        sync.Once
)

func setQueue() {
	var err error
	qfactory, err = queue.Factory()
	if err != nil {
		log.Fatalf("Failed to get the queue instance: %s", err)
	}
	_handler, err = qfactory.Handler(handle, queueName, QueueName)
	if err != nil {
		log.Fatalf("Failed to create the queue handler: %s", err)
	}
	_queue, err = qfactory.Get(queueName)
	if err != nil {
		log.Fatalf("Failed to get the queue instance: %s", err)
	}
}

func handler() queue.Handler {
	o.Do(setQueue)
	return _handler
}

func aqueue() queue.Q {
	o.Do(setQueue)
	return _queue
}

// enqueues the given messages and start the handler.
func enqueue(msgs ...queue.Message) {
	for _, msg := range msgs {
		copy := msg
		aqueue().Put(&copy, 0)
	}
	handler().Start()
}
<commit_msg>app: delete message from the queue when app is not found<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package app

import (
	"fmt"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/queue"
	"github.com/globocom/tsuru/service"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"sync"
)

const (
	// queue actions
	regenerateApprc         = "regenerate-apprc"
	startApp                = "start-app"
	RegenerateApprcAndStart = "regenerate-apprc-start-app"
	bindService             = "bind-service"

	// name of the queue for internal messages.
	queueName = "tsuru-app"

	// Name of the queue for external messages.
	QueueName = "tsuru-app-public"
)

// ensureAppIsStarted make sure that the app and all units present in the given
// message are started.
func ensureAppIsStarted(msg *queue.Message) (App, error) {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		msg.Delete()
		return a, fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	units := getUnits(&a, msg.Args[1:])
	if !a.Available() || !units.Started() {
		format := "Error handling %q for the app %q:"
		uState := units.State()
		if uState == "error" || uState == "down" {
			format += fmt.Sprintf(" units are in %q state.", uState)
			msg.Delete()
		} else {
			format += " all units must be started."
		}
		return a, fmt.Errorf(format, msg.Action, a.Name)
	}
	return a, nil
}

// bindUnit handles the bind-service message, binding a unit to all service
// instances binded to the app.
func bindUnit(msg *queue.Message) error {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		msg.Delete()
		return fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		return fmt.Errorf("Error handling %q: %s", msg.Action, err)
	}
	unit := getUnits(&a, msg.Args[1:])[0]
	var instances []service.ServiceInstance
	q := bson.M{"apps": bson.M{"$in": []string{msg.Args[0]}}}
	err = conn.ServiceInstances().Find(q).All(&instances)
	if err != nil {
		return err
	}
	for _, instance := range instances {
		_, err = instance.BindUnit(&unit)
		if err != nil {
			log.Printf("Error binding the unit %s with the service instance %s.", unit.Name, instance.Name)
		}
	}
	return nil
}

// handle is the function called by the queue handler on each message.
func handle(msg *queue.Message) {
	switch msg.Action {
	case RegenerateApprcAndStart:
		fallthrough
	case regenerateApprc:
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
			return
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
		app.serializeEnvVars()
		fallthrough
	case startApp:
		if msg.Action == regenerateApprc {
			break
		}
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		err = app.Restart(ioutil.Discard)
		if err != nil {
			log.Printf("Error handling %q. App failed to start:\n%s.", msg.Action, err)
			return
		}
		msg.Delete()
	case bindService:
		err := bindUnit(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
	default:
		log.Printf("Error handling %q: invalid action.", msg.Action)
		msg.Delete()
	}
}

// unitList is a simple slice of units, with special methods to handle state.
type unitList []Unit

// Started returns true if all units in the list is started.
func (l unitList) Started() bool {
	for _, unit := range l {
		if unit.State != string(provision.StatusStarted) {
			return false
		}
	}
	return true
}

// State returns a string if all units have the same state. Otherwise it
// returns an empty string.
func (l unitList) State() string {
	if len(l) == 0 {
		return ""
	}
	state := l[0].State
	for i := 1; i < len(l); i++ {
		if l[i].State != state {
			return ""
		}
	}
	return state
}

// getUnits builds a unitList from the given app and the names in the string
// slice.
func getUnits(app *App, names []string) unitList {
	var units []Unit
	if len(names) > 0 {
		units = make([]Unit, len(names))
		i := 0
		for _, unitName := range names {
			for _, appUnit := range app.Units {
				if appUnit.Name == unitName {
					units[i] = appUnit
					i++
					break
				}
			}
		}
	}
	return unitList(units)
}

var (
	qfactory queue.QFactory
	_queue   queue.Q
	_handler queue.Handler
	o        sync.Once
)

func setQueue() {
	var err error
	qfactory, err = queue.Factory()
	if err != nil {
		log.Fatalf("Failed to get the queue instance: %s", err)
	}
	_handler, err = qfactory.Handler(handle, queueName, QueueName)
	if err != nil {
		log.Fatalf("Failed to create the queue handler: %s", err)
	}
	_queue, err = qfactory.Get(queueName)
	if err != nil {
		log.Fatalf("Failed to get the queue instance: %s", err)
	}
}

func handler() queue.Handler {
	o.Do(setQueue)
	return _handler
}

func aqueue() queue.Q {
	o.Do(setQueue)
	return _queue
}

// enqueues the given messages and start the handler.
func enqueue(msgs ...queue.Message) {
	for _, msg := range msgs {
		copy := msg
		aqueue().Put(&copy, 0)
	}
	handler().Start()
}
<|endoftext|>


<commit_before>package platform

import (
	boshdpresolv "github.com/cloudfoundry/bosh-agent/infrastructure/devicepathresolver"
	boshcert "github.com/cloudfoundry/bosh-agent/platform/cert"
	boshnet "github.com/cloudfoundry/bosh-agent/platform/net"
	boshstats "github.com/cloudfoundry/bosh-agent/platform/stats"
	boshvitals "github.com/cloudfoundry/bosh-agent/platform/vitals"
	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	boshdir "github.com/cloudfoundry/bosh-agent/settings/directories"
	boshdirs "github.com/cloudfoundry/bosh-agent/settings/directories"
	boshcmd "github.com/cloudfoundry/bosh-utils/fileutil"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type WindowsPlatform struct {
	collector          boshstats.Collector
	fs                 boshsys.FileSystem
	cmdRunner          boshsys.CmdRunner
	compressor         boshcmd.Compressor
	copier             boshcmd.Copier
	dirProvider        boshdirs.Provider
	vitalsService      boshvitals.Service
	netManager         boshnet.Manager
	devicePathResolver boshdpresolv.DevicePathResolver
	certManager        boshcert.Manager
}

func NewWindowsPlatform(
	collector boshstats.Collector,
	fs boshsys.FileSystem,
	cmdRunner boshsys.CmdRunner,
	dirProvider boshdirs.Provider,
	netManager boshnet.Manager,
	devicePathResolver boshdpresolv.DevicePathResolver,
	logger boshlog.Logger,
) Platform {
	return &WindowsPlatform{
		fs:                 fs,
		cmdRunner:          cmdRunner,
		collector:          collector,
		compressor:         boshcmd.NewTarballCompressor(cmdRunner, fs),
		copier:             boshcmd.NewGenericCpCopier(fs, logger),
		dirProvider:        dirProvider,
		netManager:         netManager,
		devicePathResolver: devicePathResolver,
		vitalsService:      boshvitals.NewService(collector, dirProvider),
		certManager:        boshcert.NewDummyCertManager(fs, cmdRunner, logger),
	}
}

func (p WindowsPlatform) GetFs() (fs boshsys.FileSystem) {
	return p.fs
}

func (p WindowsPlatform) GetRunner() (runner boshsys.CmdRunner) {
	return p.cmdRunner
}

func (p WindowsPlatform) GetCompressor() (compressor boshcmd.Compressor) {
	return p.compressor
}

func (p WindowsPlatform) GetCopier() (copier boshcmd.Copier) {
	return p.copier
}

func (p WindowsPlatform) GetDirProvider() (dirProvider boshdir.Provider) {
	return p.dirProvider
}

func (p WindowsPlatform) GetVitalsService() (service boshvitals.Service) {
	return p.vitalsService
}

func (p WindowsPlatform) GetDevicePathResolver() (devicePathResolver boshdpresolv.DevicePathResolver) {
	return p.devicePathResolver
}

func (p WindowsPlatform) SetupRuntimeConfiguration() (err error) {
	return
}

func (p WindowsPlatform) CreateUser(username, password, basePath string) (err error) {
	return
}

func (p WindowsPlatform) AddUserToGroups(username string, groups []string) (err error) {
	return
}

func (p WindowsPlatform) DeleteEphemeralUsersMatching(regex string) (err error) {
	return
}

func (p WindowsPlatform) SetupRootDisk(ephemeralDiskPath string) (err error) {
	return
}

func (p WindowsPlatform) SetupSSH(publicKey, username string) (err error) {
	return
}

func (p WindowsPlatform) SetUserPassword(user, encryptedPwd string) (err error) {
	return
}

func (p WindowsPlatform) SetupHostname(hostname string) (err error) {
	return
}

func (p WindowsPlatform) SetupNetworking(networks boshsettings.Networks) (err error) {
	return p.netManager.SetupNetworking(networks, nil)
}

func (p WindowsPlatform) GetConfiguredNetworkInterfaces() (interfaces []string, err error) {
	return
}

func (p WindowsPlatform) GetCertManager() (certManager boshcert.Manager) {
	return p.certManager
}

func (p WindowsPlatform) SetupLogrotate(groupName, basePath, size string) (err error) {
	return
}

func (p WindowsPlatform) SetTimeWithNtpServers(servers []string) (err error) {
	return
}

func (p WindowsPlatform) SetupEphemeralDiskWithPath(devicePath string) (err error) {
	return
}

func (p WindowsPlatform) SetupRawEphemeralDisks(devices []boshsettings.DiskSettings) (err error) {
	return
}

func (p WindowsPlatform) SetupDataDir() error {
	return nil
}

func (p WindowsPlatform) SetupTmpDir() error {
	return nil
}

func (p WindowsPlatform) MountPersistentDisk(diskSettings boshsettings.DiskSettings, mountPoint string) (err error) {
	return
}

func (p WindowsPlatform) UnmountPersistentDisk(diskSettings boshsettings.DiskSettings) (didUnmount bool, err error) {
	return
}

func (p WindowsPlatform) GetEphemeralDiskPath(diskSettings boshsettings.DiskSettings) string {
	return ""
}

func (p WindowsPlatform) GetFileContentsFromCDROM(filePath string) (contents []byte, err error) {
	return p.fs.ReadFile("D:/" + filePath)
}

func (p WindowsPlatform) GetFilesContentsFromDisk(diskPath string, fileNames []string) (contents [][]byte, err error) {
	return
}

func (p WindowsPlatform) MigratePersistentDisk(fromMountPoint, toMountPoint string) (err error) {
	return
}

func (p WindowsPlatform) IsMountPoint(path string) (result bool, err error) {
	return
}

func (p WindowsPlatform) IsPersistentDiskMounted(diskSettings boshsettings.DiskSettings) (bool, error) {
	return true, nil
}

func (p WindowsPlatform) StartMonit() (err error) {
	return
}

func (p WindowsPlatform) SetupMonitUser() (err error) {
	return
}

func (p WindowsPlatform) GetMonitCredentials() (username, password string, err error) {
	return
}

func (p WindowsPlatform) PrepareForNetworkingChange() error {
	return nil
}

func (p WindowsPlatform) GetDefaultNetwork() (boshsettings.Network, error) {
	return boshsettings.Network{}, nil
}

func (p WindowsPlatform) GetHostPublicKey() (string, error) {
	return "", nil
}
<commit_msg>Update windows_platform to match new platform interface<commit_after>package platform

import (
	boshdpresolv "github.com/cloudfoundry/bosh-agent/infrastructure/devicepathresolver"
	boshcert "github.com/cloudfoundry/bosh-agent/platform/cert"
	boshnet "github.com/cloudfoundry/bosh-agent/platform/net"
	boshstats "github.com/cloudfoundry/bosh-agent/platform/stats"
	boshvitals "github.com/cloudfoundry/bosh-agent/platform/vitals"
	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	boshdir "github.com/cloudfoundry/bosh-agent/settings/directories"
	boshdirs "github.com/cloudfoundry/bosh-agent/settings/directories"
	boshcmd "github.com/cloudfoundry/bosh-utils/fileutil"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type WindowsPlatform struct {
	collector          boshstats.Collector
	fs                 boshsys.FileSystem
	cmdRunner          boshsys.CmdRunner
	compressor         boshcmd.Compressor
	copier             boshcmd.Copier
	dirProvider        boshdirs.Provider
	vitalsService      boshvitals.Service
	netManager         boshnet.Manager
	devicePathResolver boshdpresolv.DevicePathResolver
	certManager        boshcert.Manager
}

func NewWindowsPlatform(
	collector boshstats.Collector,
	fs boshsys.FileSystem,
	cmdRunner boshsys.CmdRunner,
	dirProvider boshdirs.Provider,
	netManager boshnet.Manager,
	devicePathResolver boshdpresolv.DevicePathResolver,
	logger boshlog.Logger,
) Platform {
	return &WindowsPlatform{
		fs:                 fs,
		cmdRunner:          cmdRunner,
		collector:          collector,
		compressor:         boshcmd.NewTarballCompressor(cmdRunner, fs),
		copier:             boshcmd.NewGenericCpCopier(fs, logger),
		dirProvider:        dirProvider,
		netManager:         netManager,
		devicePathResolver: devicePathResolver,
		vitalsService:      boshvitals.NewService(collector, dirProvider),
		certManager:        boshcert.NewDummyCertManager(fs, cmdRunner, 0, logger),
	}
}

func (p WindowsPlatform) GetFs() (fs boshsys.FileSystem) {
	return p.fs
}

func (p WindowsPlatform) GetRunner() (runner boshsys.CmdRunner) {
	return p.cmdRunner
}

func (p WindowsPlatform) GetCompressor() (compressor boshcmd.Compressor) {
	return p.compressor
}

func (p WindowsPlatform) GetCopier() (copier boshcmd.Copier) {
	return p.copier
}

func (p WindowsPlatform) GetDirProvider() (dirProvider boshdir.Provider) {
	return p.dirProvider
}

func (p WindowsPlatform) GetVitalsService() (service boshvitals.Service) {
	return p.vitalsService
}

func (p WindowsPlatform) GetDevicePathResolver() (devicePathResolver boshdpresolv.DevicePathResolver) {
	return p.devicePathResolver
}

func (p WindowsPlatform) SetupRuntimeConfiguration() (err error) {
	return
}

func (p WindowsPlatform) CreateUser(username, password, basePath string) (err error) {
	return
}

func (p WindowsPlatform) AddUserToGroups(username string, groups []string) (err error) {
	return
}

func (p WindowsPlatform) DeleteEphemeralUsersMatching(regex string) (err error) {
	return
}

func (p WindowsPlatform) SetupRootDisk(ephemeralDiskPath string) (err error) {
	return
}

func (p WindowsPlatform) SetupSSH(publicKey, username string) (err error) {
	return
}

func (p WindowsPlatform) SetUserPassword(user, encryptedPwd string) (err error) {
	return
}

func (p WindowsPlatform) SetupHostname(hostname string) (err error) {
	return
}

func (p WindowsPlatform) SetupNetworking(networks boshsettings.Networks) (err error) {
	return p.netManager.SetupNetworking(networks, nil)
}

func (p WindowsPlatform) GetConfiguredNetworkInterfaces() (interfaces []string, err error) {
	return
}

func (p WindowsPlatform) GetCertManager() (certManager boshcert.Manager) {
	return p.certManager
}

func (p WindowsPlatform) SetupLogrotate(groupName, basePath, size string) (err error) {
	return
}

func (p WindowsPlatform) SetTimeWithNtpServers(servers []string) (err error) {
	return
}

func (p WindowsPlatform) SetupEphemeralDiskWithPath(devicePath string) (err error) {
	return
}

func (p WindowsPlatform) SetupRawEphemeralDisks(devices []boshsettings.DiskSettings) (err error) {
	return
}

func (p WindowsPlatform) SetupDataDir() error {
	return nil
}

func (p WindowsPlatform) SetupTmpDir() error {
	return nil
}

func (p WindowsPlatform) MountPersistentDisk(diskSettings boshsettings.DiskSettings, mountPoint string) (err error) {
	return
}

func (p WindowsPlatform) UnmountPersistentDisk(diskSettings boshsettings.DiskSettings) (didUnmount bool, err error) {
	return
}

func (p WindowsPlatform) GetEphemeralDiskPath(diskSettings boshsettings.DiskSettings) string {
	return ""
}

func (p WindowsPlatform) GetFileContentsFromCDROM(filePath string) (contents []byte, err error) {
	return p.fs.ReadFile("D:/" + filePath)
}

func (p WindowsPlatform) GetFilesContentsFromDisk(diskPath string, fileNames []string) (contents [][]byte, err error) {
	return
}

func (p WindowsPlatform) MigratePersistentDisk(fromMountPoint, toMountPoint string) (err error) {
	return
}

func (p WindowsPlatform) IsMountPoint(path string) (result bool, err error) {
	return
}

func (p WindowsPlatform) IsPersistentDiskMounted(diskSettings boshsettings.DiskSettings) (bool, error) {
	return true, nil
}

func (p WindowsPlatform) StartMonit() (err error) {
	return
}

func (p WindowsPlatform) SetupMonitUser() (err error) {
	return
}

func (p WindowsPlatform) GetMonitCredentials() (username, password string, err error) {
	return
}

func (p WindowsPlatform) PrepareForNetworkingChange() error {
	return nil
}

func (p WindowsPlatform) CleanIPMacAddressCache(ip string) error {
	return nil
}

func (p WindowsPlatform) RemoveDevTools(packageFileListPath string) error {
	return nil
}

func (p WindowsPlatform) GetDefaultNetwork() (boshsettings.Network, error) {
	return boshsettings.Network{}, nil
}

func (p WindowsPlatform) GetHostPublicKey() (string, error) {
	return "", nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package deployer_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/agent/tools"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/testing"
	coretools "github.com/juju/juju/tools"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker/deployer"
)

type SimpleContextSuite struct {
	SimpleToolsFixture
}

var _ = gc.Suite(&SimpleContextSuite{})

func (s *SimpleContextSuite) SetUpTest(c *gc.C) {
	s.SimpleToolsFixture.SetUp(c, c.MkDir())
}

func (s *SimpleContextSuite) TearDownTest(c *gc.C) {
	s.SimpleToolsFixture.TearDown(c)
}

func (s *SimpleContextSuite) TestDeployRecall(c *gc.C) {
	mgr0 := s.getContext(c)
	units, err := mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	err = mgr0.DeployUnit("foo/123", "some-password")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.DeepEquals, []string{"foo/123"})
	s.assertUpstartCount(c, 1)
	s.checkUnitInstalled(c, "foo/123", "some-password")

	err = mgr0.RecallUnit("foo/123")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)
	s.checkUnitRemoved(c, "foo/123")
}

func (s *SimpleContextSuite) TestOldDeployedUnitsCanBeRecalled(c *gc.C) {
	// After r1347 deployer tag is no longer part of the upstart conf filenames,
	// now only the units' tags are used. This change is with the assumption only
	// one deployer will be running on a machine (in the machine agent as a task,
	// unlike before where there was one in the unit agent as well).
	// This test ensures units deployed previously (or their upstart confs more
	// specifically) can be detected and recalled by the deployer.

	manager := s.getContext(c)

	// No deployed units at first.
	units, err := manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	// Trying to recall any units will fail.
	err = manager.RecallUnit("principal/1")
	c.Assert(err, gc.ErrorMatches, `unit "principal/1" is not deployed`)

	// Simulate some previously deployed units with the old
	// upstart conf filename format (+deployer tags).
	s.injectUnit(c, "jujud-machine-0:unit-mysql-0.conf", "unit-mysql-0")
	s.assertUpstartCount(c, 1)
	s.injectUnit(c, "jujud-unit-wordpress-0:unit-nrpe-0.conf", "unit-nrpe-0")
	s.assertUpstartCount(c, 2)

	// Make sure we can discover them.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 2)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0"})

	// Deploy some units.
	err = manager.DeployUnit("principal/1", "some-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "principal/1", "some-password")
	s.assertUpstartCount(c, 3)
	err = manager.DeployUnit("subordinate/2", "fake-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "subordinate/2", "fake-password")
	s.assertUpstartCount(c, 4)

	// Verify the newly deployed units are also discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 4)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0", "principal/1", "subordinate/2"})

	// Recall all of them - should work ok.
	unitCount := 4
	for _, unitName := range units {
		err = manager.RecallUnit(unitName)
		c.Assert(err, gc.IsNil)
		unitCount--
		s.checkUnitRemoved(c, unitName)
		s.assertUpstartCount(c, unitCount)
	}

	// Verify they're no longer discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
}

type SimpleToolsFixture struct {
	dataDir  string
	logDir   string
	initDir  string
	origPath string
	binDir   string
}

var fakeJujud = "#!/bin/bash --norc\n# fake-jujud\nexit 0\n"

func (fix *SimpleToolsFixture) SetUp(c *gc.C, dataDir string) {
	fix.dataDir = dataDir
	fix.initDir = c.MkDir()
	fix.logDir = c.MkDir()
	toolsDir := tools.SharedToolsDir(fix.dataDir, version.Current)
	err := os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
	jujudPath := filepath.Join(toolsDir, "jujud")
	err = ioutil.WriteFile(jujudPath, []byte(fakeJujud), 0755)
	c.Assert(err, gc.IsNil)
	toolsPath := filepath.Join(toolsDir, "downloaded-tools.txt")
	testTools := coretools.Tools{Version: version.Current, URL: "http://testing.invalid/tools"}
	data, err := json.Marshal(testTools)
	c.Assert(err, gc.IsNil)
	err = ioutil.WriteFile(toolsPath, data, 0644)
	c.Assert(err, gc.IsNil)
	fix.binDir = c.MkDir()
	fix.origPath = os.Getenv("PATH")
	os.Setenv("PATH", fix.binDir+":"+fix.origPath)
	fix.makeBin(c, "status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "stopped-status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "started-status", `echo "blah start/running, process 666"`)
	fix.makeBin(c, "start", "cp $(which started-status) $(which status)")
	fix.makeBin(c, "stop", "cp $(which stopped-status) $(which status)")
}

func (fix *SimpleToolsFixture) TearDown(c *gc.C) {
	os.Setenv("PATH", fix.origPath)
}

func (fix *SimpleToolsFixture) makeBin(c *gc.C, name, script string) {
	path := filepath.Join(fix.binDir, name)
	err := ioutil.WriteFile(path, []byte("#!/bin/bash --norc\n"+script), 0755)
	c.Assert(err, gc.IsNil)
}

func (fix *SimpleToolsFixture) assertUpstartCount(c *gc.C, count int) {
	fis, err := ioutil.ReadDir(fix.initDir)
	c.Assert(err, gc.IsNil)
	c.Assert(fis, gc.HasLen, count)
}

func (fix *SimpleToolsFixture) getContext(c *gc.C) *deployer.SimpleContext {
	config := agentConfig("machine-tag", fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) getContextForMachine(c *gc.C, machineTag string) *deployer.SimpleContext {
	config := agentConfig(machineTag, fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) paths(tag string) (confPath, agentDir, toolsDir string) {
	confName := fmt.Sprintf("jujud-%s.conf", tag)
	confPath = filepath.Join(fix.initDir, confName)
	agentDir = agent.Dir(fix.dataDir, tag)
	toolsDir = tools.ToolsDir(fix.dataDir, tag)
	return
}

func (fix *SimpleToolsFixture) checkUnitInstalled(c *gc.C, name, password string) {
	tag := names.NewUnitTag(name).String()
	uconfPath, _, toolsDir := fix.paths(tag)
	uconfData, err := ioutil.ReadFile(uconfPath)
	c.Assert(err, gc.IsNil)
	uconf := string(uconfData)
	var execLine string
	for _, line := range strings.Split(uconf, "\n") {
		if strings.HasPrefix(line, "exec ") {
			execLine = line
			break
		}
	}
	if execLine == "" {
		c.Fatalf("no command found in %s:\n%s", uconfPath, uconf)
	}
	logPath := filepath.Join(fix.logDir, tag+".log")
	jujudPath := filepath.Join(toolsDir, "jujud")
	for _, pat := range []string{
		"^exec " + jujudPath + " unit ",
		" --unit-name " + name + " ",
		" >> " + logPath + " 2>&1$",
	} {
		match, err := regexp.MatchString(pat, execLine)
		c.Assert(err, gc.IsNil)
		if !match {
			c.Fatalf("failed to match:\n%s\nin:\n%s", pat, execLine)
		}
	}

	conf, err := agent.ReadConfig(agent.ConfigPath(fix.dataDir, tag))
	c.Assert(err, gc.IsNil)
	c.Assert(conf.Tag(), gc.Equals, tag)
	c.Assert(conf.DataDir(), gc.Equals, fix.dataDir)

	jujudData, err := ioutil.ReadFile(jujudPath)
	c.Assert(err, gc.IsNil)
	c.Assert(string(jujudData), gc.Equals, fakeJujud)
}

func (fix *SimpleToolsFixture) checkUnitRemoved(c *gc.C, name string) {
	tag := names.NewUnitTag(name).String()
	confPath, agentDir, toolsDir := fix.paths(tag)
	for _, path := range []string{confPath, agentDir, toolsDir} {
		_, err := ioutil.ReadFile(path)
		if err == nil {
			c.Log("Warning: %q not removed as expected", path)
		} else {
			c.Assert(err, jc.Satisfies, os.IsNotExist)
		}
	}
}

func (fix *SimpleToolsFixture) injectUnit(c *gc.C, upstartConf, unitTag string) {
	confPath := filepath.Join(fix.initDir, upstartConf)
	err := ioutil.WriteFile(confPath, []byte("#!/bin/bash --norc\necho $0"), 0644)
	c.Assert(err, gc.IsNil)
	toolsDir := filepath.Join(fix.dataDir, "tools", unitTag)
	err = os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
}

type mockConfig struct {
	agent.Config
	tag               string
	datadir           string
	logdir            string
	upgradedToVersion version.Number
	jobs              []params.MachineJob
}

func (mock *mockConfig) Tag() string {
	return mock.tag
}

func (mock *mockConfig) DataDir() string {
	return mock.datadir
}

func (mock *mockConfig) LogDir() string {
	return mock.logdir
}

func (mock *mockConfig) Jobs() []params.MachineJob {
	return mock.jobs
}

func (mock *mockConfig) UpgradedToVersion() version.Number {
	return mock.upgradedToVersion
}

func (mock *mockConfig) WriteUpgradedToVersion(newVersion version.Number) error {
	mock.upgradedToVersion = newVersion
	return nil
}

func (mock *mockConfig) CACert() string {
	return testing.CACert
}

func (mock *mockConfig) Value(_ string) string {
	return ""
}

func agentConfig(tag, datadir, logdir string) agent.Config {
	return &mockConfig{tag: tag, datadir: datadir, logdir: logdir}
}
<commit_msg>Corrected search for exec block to work with new script block.<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package deployer_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"

	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/agent/tools"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/testing"
	coretools "github.com/juju/juju/tools"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker/deployer"
)

type SimpleContextSuite struct {
	SimpleToolsFixture
}

var _ = gc.Suite(&SimpleContextSuite{})

func (s *SimpleContextSuite) SetUpTest(c *gc.C) {
	s.SimpleToolsFixture.SetUp(c, c.MkDir())
}

func (s *SimpleContextSuite) TearDownTest(c *gc.C) {
	s.SimpleToolsFixture.TearDown(c)
}

func (s *SimpleContextSuite) TestDeployRecall(c *gc.C) {
	mgr0 := s.getContext(c)
	units, err := mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	err = mgr0.DeployUnit("foo/123", "some-password")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.DeepEquals, []string{"foo/123"})
	s.assertUpstartCount(c, 1)
	s.checkUnitInstalled(c, "foo/123", "some-password")

	err = mgr0.RecallUnit("foo/123")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)
	s.checkUnitRemoved(c, "foo/123")
}

func (s *SimpleContextSuite) TestOldDeployedUnitsCanBeRecalled(c *gc.C) {
	// After r1347 deployer tag is no longer part of the upstart conf filenames,
	// now only the units' tags are used. This change is with the assumption only
	// one deployer will be running on a machine (in the machine agent as a task,
	// unlike before where there was one in the unit agent as well).
	// This test ensures units deployed previously (or their upstart confs more
	// specifically) can be detected and recalled by the deployer.

	manager := s.getContext(c)

	// No deployed units at first.
	units, err := manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	// Trying to recall any units will fail.
	err = manager.RecallUnit("principal/1")
	c.Assert(err, gc.ErrorMatches, `unit "principal/1" is not deployed`)

	// Simulate some previously deployed units with the old
	// upstart conf filename format (+deployer tags).
	s.injectUnit(c, "jujud-machine-0:unit-mysql-0.conf", "unit-mysql-0")
	s.assertUpstartCount(c, 1)
	s.injectUnit(c, "jujud-unit-wordpress-0:unit-nrpe-0.conf", "unit-nrpe-0")
	s.assertUpstartCount(c, 2)

	// Make sure we can discover them.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 2)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0"})

	// Deploy some units.
	err = manager.DeployUnit("principal/1", "some-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "principal/1", "some-password")
	s.assertUpstartCount(c, 3)
	err = manager.DeployUnit("subordinate/2", "fake-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "subordinate/2", "fake-password")
	s.assertUpstartCount(c, 4)

	// Verify the newly deployed units are also discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 4)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0", "principal/1", "subordinate/2"})

	// Recall all of them - should work ok.
	unitCount := 4
	for _, unitName := range units {
		err = manager.RecallUnit(unitName)
		c.Assert(err, gc.IsNil)
		unitCount--
		s.checkUnitRemoved(c, unitName)
		s.assertUpstartCount(c, unitCount)
	}

	// Verify they're no longer discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
}

type SimpleToolsFixture struct {
	dataDir  string
	logDir   string
	initDir  string
	origPath string
	binDir   string
}

var fakeJujud = "#!/bin/bash --norc\n# fake-jujud\nexit 0\n"

func (fix *SimpleToolsFixture) SetUp(c *gc.C, dataDir string) {
	fix.dataDir = dataDir
	fix.initDir = c.MkDir()
	fix.logDir = c.MkDir()
	toolsDir := tools.SharedToolsDir(fix.dataDir, version.Current)
	err := os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
	jujudPath := filepath.Join(toolsDir, "jujud")
	err = ioutil.WriteFile(jujudPath, []byte(fakeJujud), 0755)
	c.Assert(err, gc.IsNil)
	toolsPath := filepath.Join(toolsDir, "downloaded-tools.txt")
	testTools := coretools.Tools{Version: version.Current, URL: "http://testing.invalid/tools"}
	data, err := json.Marshal(testTools)
	c.Assert(err, gc.IsNil)
	err = ioutil.WriteFile(toolsPath, data, 0644)
	c.Assert(err, gc.IsNil)
	fix.binDir = c.MkDir()
	fix.origPath = os.Getenv("PATH")
	os.Setenv("PATH", fix.binDir+":"+fix.origPath)
	fix.makeBin(c, "status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "stopped-status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "started-status", `echo "blah start/running, process 666"`)
	fix.makeBin(c, "start", "cp $(which started-status) $(which status)")
	fix.makeBin(c, "stop", "cp $(which stopped-status) $(which status)")
}

func (fix *SimpleToolsFixture) TearDown(c *gc.C) {
	os.Setenv("PATH", fix.origPath)
}

func (fix *SimpleToolsFixture) makeBin(c *gc.C, name, script string) {
	path := filepath.Join(fix.binDir, name)
	err := ioutil.WriteFile(path, []byte("#!/bin/bash --norc\n"+script), 0755)
	c.Assert(err, gc.IsNil)
}

func (fix *SimpleToolsFixture) assertUpstartCount(c *gc.C, count int) {
	fis, err := ioutil.ReadDir(fix.initDir)
	c.Assert(err, gc.IsNil)
	c.Assert(fis, gc.HasLen, count)
}

func (fix *SimpleToolsFixture) getContext(c *gc.C) *deployer.SimpleContext {
	config := agentConfig("machine-tag", fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) getContextForMachine(c *gc.C, machineTag string) *deployer.SimpleContext {
	config := agentConfig(machineTag, fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) paths(tag string) (confPath, agentDir, toolsDir string) {
	confName := fmt.Sprintf("jujud-%s.conf", tag)
	confPath = filepath.Join(fix.initDir, confName)
	agentDir = agent.Dir(fix.dataDir, tag)
	toolsDir = tools.ToolsDir(fix.dataDir, tag)
	return
}

func (fix *SimpleToolsFixture) checkUnitInstalled(c *gc.C, name, password string) {
	tag := names.NewUnitTag(name).String()
	uconfPath, _, toolsDir := fix.paths(tag)
	uconfData, err := ioutil.ReadFile(uconfPath)
	c.Assert(err, gc.IsNil)
	uconf := string(uconfData)

	regex := regexp.MustCompile("(?m)(?:^\\s)*exec\\s.+$")
	execs := regex.FindAllString(uconf, -1)

	if nil == execs {
		c.Fatalf("no command found in %s:\n%s", uconfPath, uconf)
	} else if 1 > len(execs) {
		c.Fatalf("Test is not built to handle more than one exec line.")
	}

	logPath := filepath.Join(fix.logDir, tag+".log")
	jujudPath := filepath.Join(toolsDir, "jujud")

	for _, pat := range []string{
		"^exec " + jujudPath + " unit ",
		" --unit-name " + name + " ",
		" >> " + logPath + " 2>&1$",
	} {
		match, err := regexp.MatchString(pat, execs[0])
		c.Assert(err, gc.IsNil)
		if !match {
			c.Fatalf("failed to match:\n%s\nin:\n%s", pat, execs[0])
		}
	}

	conf, err := agent.ReadConfig(agent.ConfigPath(fix.dataDir, tag))
	c.Assert(err, gc.IsNil)
	c.Assert(conf.Tag(), gc.Equals, tag)
	c.Assert(conf.DataDir(), gc.Equals, fix.dataDir)

	jujudData, err := ioutil.ReadFile(jujudPath)
	c.Assert(err, gc.IsNil)
	c.Assert(string(jujudData), gc.Equals, fakeJujud)
}

func (fix *SimpleToolsFixture) checkUnitRemoved(c *gc.C, name string) {
	tag := names.NewUnitTag(name).String()
	confPath, agentDir, toolsDir := fix.paths(tag)
	for _, path := range []string{confPath, agentDir, toolsDir} {
		_, err := ioutil.ReadFile(path)
		if err == nil {
			c.Log("Warning: %q not removed as expected", path)
		} else {
			c.Assert(err, jc.Satisfies, os.IsNotExist)
		}
	}
}

func (fix *SimpleToolsFixture) injectUnit(c *gc.C, upstartConf, unitTag string) {
	confPath := filepath.Join(fix.initDir, upstartConf)
	err := ioutil.WriteFile(confPath, []byte("#!/bin/bash --norc\necho $0"), 0644)
	c.Assert(err, gc.IsNil)
	toolsDir := filepath.Join(fix.dataDir, "tools", unitTag)
	err = os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
}

type mockConfig struct {
	agent.Config
	tag               string
	datadir           string
	logdir            string
	upgradedToVersion version.Number
	jobs              []params.MachineJob
}

func (mock *mockConfig) Tag() string {
	return mock.tag
}

func (mock *mockConfig) DataDir() string {
	return mock.datadir
}

func (mock *mockConfig) LogDir() string {
	return mock.logdir
}

func (mock *mockConfig) Jobs() []params.MachineJob {
	return mock.jobs
}

func (mock *mockConfig) UpgradedToVersion() version.Number {
	return mock.upgradedToVersion
}

func (mock *mockConfig) WriteUpgradedToVersion(newVersion version.Number) error {
	mock.upgradedToVersion = newVersion
	return nil
}

func (mock *mockConfig) CACert() string {
	return testing.CACert
}

func (mock *mockConfig) Value(_ string) string {
	return ""
}

func agentConfig(tag, datadir, logdir string) agent.Config {
	return &mockConfig{tag: tag, datadir: datadir, logdir: logdir}
}
<|endoftext|>


<commit_before>package relationer

import (
	"launchpad.net/juju-core/state"
	"sort"
)

// HookInfo holds everything we need to know to execute a relation hook,
// assuming we have access to the relation itself and the unit executing
// the hook.
type HookInfo struct {
	Kind    string
	Unit    string
	Members map[string]map[string]interface{}
}

// unitInfo holds information about a unit that is, or has been, or will
// be, a member of the relation; and the next hook that will run and its
// position in the queue if appropriate.
type unitInfo struct {
	unit     string
	version  int
	settings map[string]interface{}
	hook     string
	prev     *unitInfo
	next     *unitInfo
}

// HookQueue accepts state.RelationUnitsChange events and converts them
// into HookInfos.
type HookQueue struct {
	// info holds the unitInfo for every unit known to be present in
	// the relation at any time from the "present" to the furthest
	// known "future".
	info map[string]*unitInfo

	// head and tail are the ends of the queue.
	head *unitInfo
	tail *unitInfo

	// members holds the names of the relation units known to the last
	// HookInfo emitted from the queue; from the point of view of the
	// unit agent, it is the "present" membership of the relation.
	members map[string]struct{}

	// inflight holds enough information to reconstruct the last HookInfo
	// returned from Next, until Done is called.
	inflight *unitInfo
}

// NewHookQueue returns an empty HookQueue.
func NewHookQueue() *HookQueue {
	return &HookQueue{
		members: map[string]struct{}{},
		info:    map[string]*unitInfo{},
	}
}

// Add updates the queue.
func (q *HookQueue) Add(ruc state.RelationUnitsChange) {
	// Enforce consistent addition order, mainly for testing purposes.
	changedUnits := []string{}
	for unit, _ := range ruc.Changed {
		changedUnits = append(changedUnits, unit)
	}
	sort.Strings(changedUnits)

	for _, unit := range changedUnits {
		settings := ruc.Changed[unit]
		if info, found := q.info[unit]; !found {
			// If not known, add to info and queue a join.
			q.info[unit] = &unitInfo{unit: unit}
			q.append(unit, "joined")
		} else if info.hook == "" {
			// If known, and not already in queue, queue a change.
			q.append(unit, "changed")
		} else if info.hook == "departed" {
			// If known, and queued for depart, clear out the old entry...
			q.remove(unit)
			if settings.Version != info.version {
				// ...and queue a change if justified.
				q.append(unit, "changed")
			}
		} // Otherwise, it's already queued for either join or change; ignore.

		// Always update the stored settings.
		q.info[unit].version = settings.Version
		q.info[unit].settings = settings.Settings
	}

	for _, unit := range ruc.Departed {
		if hook := q.info[unit].hook; hook == "" {
			// If not present in queue, queue for depart.
			q.append(unit, "departed")
		} else {
			// Otherwise, clear out the old entry...
			q.remove(unit)
			if hook == "changed" {
				// ...and schedule a deletion if justified.
				q.append(unit, "departed")
			}
		}
	}
}

// Next returns information about the next hook to fire. If the returned
// bool is false, the HookInfo is not valid and the queue is currently empty.
// If the HookInfo is valid, subsequent calls to Next will return the same
// HookInfo, until Done is called.
func (q *HookQueue) Next() (HookInfo, bool) {
	if q.inflight == nil {
		if q.head == nil {
			// The queue is empty.
			return HookInfo{}, false
		}

		// Clone the unit info and update the queue.
		info := *q.head
		if info.hook == "joined" {
			q.members[info.unit] = struct{}{}
			q.head.hook = "changed"
		} else {
			q.remove(info.unit)
		}

		// If we're departing, we can forget all about this unit.
		if info.hook == "departed" {
			delete(q.info, info.unit)
			delete(q.members, info.unit)
		}

		q.inflight = &info
	} else if q.inflight.hook == "changed" {
		// It may be the case that the settings for an inflight changed
		// hook have changed since we last tried to run it. If that is
		// so, *and* that unit has a queued changed hook, that queued hook
		// is now redundant and can be dropped from the queue. (If the
		// settings have changed, we will certainly have a queued hook
		// for that unit; but the unit might subsequently have departed,
		// so it's not necessarily a changed hook.)
		latest := q.info[q.inflight.unit]
		if q.inflight.version != latest.version && latest.hook == "changed" {
			q.remove(latest.unit)
		}
	}

	// Return a HookInfo created from the inflight unitInfo.
	result := HookInfo{q.inflight.hook, q.inflight.unit, make(map[string]map[string]interface{})}
	for m := range q.members {
		result.Members[m] = q.info[m].settings
	}
	return result, true
}

// Done informs the queue that the last HookInfo returned from Next has been
// handled, and can be safely forgotten. Done will panic if no change is
// in flight.
func (q *HookQueue) Done() {
	if q.inflight == nil {
		panic("can't call Done when no hook is inflight")
	}
	q.inflight = nil
}

// append places the named unit info at the tail of the queue. It will
// panic if the unit is not known.
func (q *HookQueue) append(unit, hook string) {
	// First make sure it's out of the queue.
	q.remove(unit)

	// Set the unit's next action.
	info := q.info[unit]
	info.hook = hook

	// If the queue is empty, place it at the start.
	if q.head == nil {
		q.head = info
	}

	// Place it at the end.
	info.prev = q.tail
	if q.tail != nil {
		q.tail.next = info
	}
	q.tail = info
}

// remove removes the named unit from the queue, but does not remove
// its information. It will panic if the unit is not known.
func (q *HookQueue) remove(unit string) {
	if q.head == nil {
		// The queue is empty, nothing to do.
		return
	}

	// Get the unit info and clear its next action.
	info := q.info[unit]
	info.hook = ""

	// If the info is at either end (or both ends) of the queue,
	// update the appropriate end(s) directly.
	if q.head == info {
		q.head = info.next
	}
	if q.tail == info {
		q.tail = info.prev
	}

	// Remove the unit from the list.
	if info.prev != nil {
		info.prev.next = info.next
	}
	if info.next != nil {
		info.next.prev = info.prev
	}
	info.prev = nil
	info.next = nil
}
<commit_msg>doc, fieldname, code structure tweaks<commit_after>package relationer

import (
	"launchpad.net/juju-core/state"
	"sort"
)

// HookInfo holds everything we need to know to execute a relation hook,
// assuming we have access to the relation itself and the unit executing
// the hook.
type HookInfo struct {
	HookName   string
	RemoteUnit string
	Members    map[string]map[string]interface{}
}

// unitInfo holds information about a unit that is, or has been, or will
// be, a member of the relation; and the next hook that will run and its
// position in the queue if appropriate.
type unitInfo struct {
	unit     string
	version  int
	settings map[string]interface{}
	hook     string
	prev     *unitInfo
	next     *unitInfo
}

// HookQueue accepts state.RelationUnitsChange events and converts them
// into HookInfos.
type HookQueue struct {
	// info holds the unitInfo for every unit known to be present in
	// the relation at any time from the "present" to the furthest
	// known "future".
	info map[string]*unitInfo

	// head and tail are the ends of the queue.
	head *unitInfo
	tail *unitInfo

	// members holds the names of the relation units known to the last
	// HookInfo emitted from the queue; from the point of view of the
	// unit agent, it is the "present" membership of the relation.
	members map[string]struct{}

	// inflight holds enough information to reconstruct the last HookInfo
	// returned from Next, until Done is called.
	inflight *unitInfo
}

// NewHookQueue returns an empty HookQueue.
func NewHookQueue() *HookQueue {
	return &HookQueue{
		members: map[string]struct{}{},
		info:    map[string]*unitInfo{},
	}
}

// Add updates the queue.
func (q *HookQueue) Add(ruc state.RelationUnitsChange) {
	// Enforce consistent addition order, mainly for testing purposes.
	changedUnits := []string{}
	for unit, _ := range ruc.Changed {
		changedUnits = append(changedUnits, unit)
	}
	sort.Strings(changedUnits)

	for _, unit := range changedUnits {
		settings := ruc.Changed[unit]
		if info, found := q.info[unit]; !found {
			// If not known, add to info and queue a join.
			q.info[unit] = &unitInfo{unit: unit}
			q.append(unit, "joined")
		} else if info.hook == "" {
			// If known, and not already in queue, queue a change.
			q.append(unit, "changed")
		} else if info.hook == "departed" {
			// If known, and queued for depart, clear out the old entry...
			q.remove(unit)
			if settings.Version != info.version {
				// ...and queue a change if justified.
				q.append(unit, "changed")
			}
		} // Otherwise, it's already queued for either join or change; ignore.

		// Always update the stored settings.
		q.info[unit].version = settings.Version
		q.info[unit].settings = settings.Settings
	}

	for _, unit := range ruc.Departed {
		if hook := q.info[unit].hook; hook == "" {
			// If not present in queue, queue for depart.
			q.append(unit, "departed")
		} else {
			// Otherwise, clear out the old entry...
			q.remove(unit)
			if hook == "changed" {
				// ...and schedule a deletion if justified.
				q.append(unit, "departed")
			}
		}
	}
}

// Next returns information about the next hook to fire. If the returned
// bool is false, the HookInfo is not valid and the queue is currently empty.
// If the HookInfo is valid, subsequent calls to Next will return the same
// HookInfo, until Done is called.
func (q *HookQueue) Next() (HookInfo, bool) {
	if q.inflight == nil {
		if q.head == nil {
			// The queue is empty.
			return HookInfo{}, false
		}
		// Clone the unit info and update our state.
		info := *q.head
		q.inflight = &info
		if info.hook == "joined" {
			q.members[info.unit] = struct{}{}
			q.head.hook = "changed"
		} else {
			q.remove(info.unit)
			if info.hook == "departed" {
				delete(q.info, info.unit)
				delete(q.members, info.unit)
			}
		}
	} else if q.inflight.hook == "changed" {
		// It may be the case that the settings for an inflight changed
		// hook have changed since we last tried to run it. If that is
		// so, *and* that unit has a queued changed hook, that queued hook
		// is now redundant and can be dropped from the queue. (If the
		// settings have changed, we will certainly have a queued hook
		// for that unit; but the unit might subsequently have departed,
		// so it's not necessarily a changed hook.)
		latest := q.info[q.inflight.unit]
		if q.inflight.version != latest.version && latest.hook == "changed" {
			q.remove(latest.unit)
		}
	}

	// Return a HookInfo created from the inflight unitInfo.
	result := HookInfo{
		HookName:   q.inflight.hook,
		RemoteUnit: q.inflight.unit,
		Members:    make(map[string]map[string]interface{}),
	}
	for m := range q.members {
		result.Members[m] = q.info[m].settings
	}
	return result, true
}

// Done informs the queue that the last HookInfo returned from Next has been
// handled, and can be safely forgotten. Done will panic if no change is
// in flight.
func (q *HookQueue) Done() {
	if q.inflight == nil {
		panic("can't call Done when no hook is inflight")
	}
	q.inflight = nil
}

// append places the named unit info at the tail of the queue. It will
// panic if the unit is not known.
func (q *HookQueue) append(unit, hook string) {
	// First make sure it's out of the queue.
	q.remove(unit)

	// Set the unit's next action.
	info := q.info[unit]
	info.hook = hook

	// If the queue is empty, place it at the start.
	if q.head == nil {
		q.head = info
	}

	// Place it at the end.
	info.prev = q.tail
	if q.tail != nil {
		q.tail.next = info
	}
	q.tail = info
}

// remove removes the named unit from the queue, but does not remove
// its information. It will panic if the unit is not known.
func (q *HookQueue) remove(unit string) {
	if q.head == nil {
		// The queue is empty, nothing to do.
		return
	}

	// Get the unit info and clear its next action.
	info := q.info[unit]
	info.hook = ""

	// If the info is at either end (or both ends) of the queue,
	// update the appropriate end(s) directly.
	if q.head == info {
		q.head = info.next
	}
	if q.tail == info {
		q.tail = info.prev
	}

	// Remove the unit from the list.
	if info.prev != nil {
		info.prev.next = info.next
	}
	if info.next != nil {
		info.next.prev = info.prev
	}
	info.prev = nil
	info.next = nil
}
<|endoftext|>


<commit_before>// Copyright 2018, OpenCensus Authors
//
// 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 ochttp_test

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"

	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/stats"
)

const reqCount = 5

func TestClient(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
		resp.Write([]byte("Hello, world!"))
	}))
	defer server.Close()

	views := []string{
		"opencensus.io/http/client/requests",
		"opencensus.io/http/client/latency",
		"opencensus.io/http/client/request_size",
		"opencensus.io/http/client/response_size",
	}
	for _, name := range views {
		v := stats.FindView(name)
		if v == nil {
			t.Errorf("view not found %q", name)
			continue
		}
		v.Subscribe()
	}

	var (
		w  sync.WaitGroup
		tr ochttp.Transport
	)
	w.Add(reqCount)
	for i := 0; i < reqCount; i++ {
		go func() {
			req, err := http.NewRequest("POST", server.URL, strings.NewReader("req-body"))
			if err != nil {
				t.Fatalf("error creating request: %v", err)
			}
			resp, err := tr.RoundTrip(req)
			if err != nil {
				t.Fatalf("response error: %v", err)
			}
			if got, want := resp.StatusCode, 200; got != want {
				t.Fatalf("resp.StatusCode=%d; wantCount %d", got, want)
			}
			ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			w.Done()
		}()
	}
	w.Wait()

	for _, viewName := range views {
		v := stats.FindView(viewName)
		if v == nil {
			t.Errorf("view not found %q", viewName)
			continue
		}
		rows, err := v.RetrieveData()
		if err != nil {
			t.Error(err)
			continue
		}
		if got, want := len(rows), 1; got != want {
			t.Errorf("len(%q) = %d; want %d", viewName, got, want)
			continue
		}
		data := rows[0].Data
		var count int64
		switch data := data.(type) {
		case *stats.CountData:
			count = *(*int64)(data)
		case *stats.DistributionData:
			count = data.Count
		default:
			t.Errorf("don't know how to handle data type: %v", data)
		}
		if got := count; got != reqCount {
			t.Fatalf("%s = %d; want %d", viewName, got, reqCount)
		}
	}
}
<commit_msg>Not to call t.Fatalf in another goroutine (#419)<commit_after>// Copyright 2018, OpenCensus Authors
//
// 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 ochttp_test

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"

	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/stats"
)

const reqCount = 5

func TestClient(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
		resp.Write([]byte("Hello, world!"))
	}))
	defer server.Close()

	views := []string{
		"opencensus.io/http/client/requests",
		"opencensus.io/http/client/latency",
		"opencensus.io/http/client/request_size",
		"opencensus.io/http/client/response_size",
	}
	for _, name := range views {
		v := stats.FindView(name)
		if v == nil {
			t.Errorf("view not found %q", name)
			continue
		}
		v.Subscribe()
	}

	var (
		w    sync.WaitGroup
		tr   ochttp.Transport
		errs = make(chan error, reqCount)
	)
	w.Add(reqCount)
	for i := 0; i < reqCount; i++ {
		go func() {
			defer w.Done()
			req, err := http.NewRequest("POST", server.URL, strings.NewReader("req-body"))
			if err != nil {
				errs <- fmt.Errorf("error creating request: %v", err)
			}
			resp, err := tr.RoundTrip(req)
			if err != nil {
				errs <- fmt.Errorf("response error: %v", err)
			}
			if err := resp.Body.Close(); err != nil {
				errs <- fmt.Errorf("error closing response body: %v", err)
			}
			if got, want := resp.StatusCode, 200; got != want {
				errs <- fmt.Errorf("resp.StatusCode=%d; wantCount %d", got, want)
			}
		}()
	}

	go func() {
		w.Wait()
		close(errs)
	}()

	for err := range errs {
		if err != nil {
			t.Fatal(err)
		}
	}

	for _, viewName := range views {
		v := stats.FindView(viewName)
		if v == nil {
			t.Errorf("view not found %q", viewName)
			continue
		}
		rows, err := v.RetrieveData()
		if err != nil {
			t.Error(err)
			continue
		}
		if got, want := len(rows), 1; got != want {
			t.Errorf("len(%q) = %d; want %d", viewName, got, want)
			continue
		}
		data := rows[0].Data
		var count int64
		switch data := data.(type) {
		case *stats.CountData:
			count = *(*int64)(data)
		case *stats.DistributionData:
			count = data.Count
		default:
			t.Errorf("don't know how to handle data type: %v", data)
		}
		if got := count; got != reqCount {
			t.Fatalf("%s = %d; want %d", viewName, got, reqCount)
		}
	}
}
<|endoftext|>