<commit_before>// Copyright © 2019 The Knative 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 service

import (
	"errors"
	"fmt"
	"io"

	"github.com/knative/client/pkg/kn/commands"
	"github.com/knative/client/pkg/serving/v1alpha1"

	serving_v1alpha1_api "github.com/knative/serving/pkg/apis/serving/v1alpha1"
	"github.com/spf13/cobra"

	corev1 "k8s.io/api/core/v1"
	api_errors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func NewServiceCreateCommand(p *commands.KnParams) *cobra.Command {
	var editFlags ConfigurationEditFlags
	var waitFlags commands.WaitFlags

	serviceCreateCommand := &cobra.Command{
		Use:   "create NAME --image IMAGE",
		Short: "Create a service.",
		Example: `
  # Create a service 'mysvc' using image at dev.local/ns/image:latest
  kn service create mysvc --image dev.local/ns/image:latest

  # Create a service with multiple environment variables
  kn service create mysvc --env KEY1=VALUE1 --env KEY2=VALUE2 --image dev.local/ns/image:latest

  # Create or replace a service 's1' with image dev.local/ns/image:v2 using --force flag
  # if service 's1' doesn't exist, it's just a normal create operation
  kn service create --force s1 --image dev.local/ns/image:v2

  # Create or replace environment variables of service 's1' using --force flag
  kn service create --force s1 --env KEY1=NEW_VALUE1 --env NEW_KEY2=NEW_VALUE2 --image dev.local/ns/image:v1

  # Create service 'mysvc' with port 80
  kn service create mysvc --port 80 --image dev.local/ns/image:latest

  # Create or replace default resources of a service 's1' using --force flag
  # (earlier configured resource requests and limits will be replaced with default)
  # (earlier configured environment variables will be cleared too if any)
  kn service create --force s1 --image dev.local/ns/image:v1`,

		RunE: func(cmd *cobra.Command, args []string) (err error) {
			if len(args) != 1 {
				return errors.New("'service create' requires the service name given as single argument")
			}
			name := args[0]
			if editFlags.Image == "" {
				return errors.New("'service create' requires the image name to run provided with the --image option")
			}

			namespace, err := p.GetNamespace(cmd)
			if err != nil {
				return err
			}

			service, err := constructService(cmd, editFlags, name, namespace)
			if err != nil {
				return err
			}

			client, err := p.NewClient(namespace)
			if err != nil {
				return err
			}

			serviceExists, err := serviceExists(client, name, namespace)
			if err != nil {
				return err
			}

			if serviceExists {
				if !editFlags.ForceCreate {
					return fmt.Errorf(
						"cannot create service '%s' in namespace '%s' "+
							"because the service already exists and no --force option was given", name, namespace)
				}
				err = replaceService(client, service, namespace, cmd.OutOrStdout())
			} else {
				err = createService(client, service, namespace, cmd.OutOrStdout())
			}
			if err != nil {
				return err
			}

			if !waitFlags.Async {
				out := cmd.OutOrStdout()
				err := waitForService(client, name, out, waitFlags.TimeoutInSeconds)
				if err != nil {
					return err
				}
				return showUrl(client, name, namespace, out)
			}

			return nil
		},
	}
	commands.AddNamespaceFlags(serviceCreateCommand.Flags(), false)
	editFlags.AddCreateFlags(serviceCreateCommand)
	waitFlags.AddConditionWaitFlags(serviceCreateCommand, 60, "Create", "service")
	return serviceCreateCommand
}

// Duck type for writers having a flush
type flusher interface {
	Flush() error
}

func flush(out io.Writer) {
	if flusher, ok := out.(flusher); ok {
		flusher.Flush()
	}
}

func createService(client v1alpha1.KnClient, service *serving_v1alpha1_api.Service, namespace string, out io.Writer) error {
	err := client.CreateService(service)
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "Service '%s' successfully created in namespace '%s'.\n", service.Name, namespace)
	return nil
}

func replaceService(client v1alpha1.KnClient, service *serving_v1alpha1_api.Service, namespace string, out io.Writer) error {
	var retries = 0
	for {
		existingService, err := client.GetService(service.Name)
		if err != nil {
			return err
		}
		service.ResourceVersion = existingService.ResourceVersion
		err = client.UpdateService(service)
		if err != nil {
			// Retry to update when a resource version conflict exists
			if api_errors.IsConflict(err) && retries < MaxUpdateRetries {
				retries++
				continue
			}
			return err
		}
		fmt.Fprintf(out, "Service '%s' successfully replaced in namespace '%s'.\n", service.Name, namespace)
		return nil
	}
}

func serviceExists(client v1alpha1.KnClient, name string, namespace string) (bool, error) {
	_, err := client.GetService(name)
	if api_errors.IsNotFound(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

// Create service struct from provided options
func constructService(cmd *cobra.Command, editFlags ConfigurationEditFlags, name string, namespace string) (*serving_v1alpha1_api.Service,
	error) {

	service := serving_v1alpha1_api.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
	}

	// TODO: Should it always be `runLatest` ?
	service.Spec.DeprecatedRunLatest = &serving_v1alpha1_api.RunLatestType{
		Configuration: serving_v1alpha1_api.ConfigurationSpec{
			DeprecatedRevisionTemplate: &serving_v1alpha1_api.RevisionTemplateSpec{
				Spec: serving_v1alpha1_api.RevisionSpec{
					DeprecatedContainer: &corev1.Container{},
				},
			},
		},
	}

	err := editFlags.Apply(&service, cmd)
	if err != nil {
		return nil, err
	}
	return &service, nil
}

func showUrl(client v1alpha1.KnClient, serviceName string, namespace string, out io.Writer) error {
	service, err := client.GetService(serviceName)
	if err != nil {
		return fmt.Errorf("cannot fetch service '%s' in namespace '%s' for extracting the URL: %v", serviceName, namespace, err)
	}
	url := service.Status.URL.String()
	if url == "" {
		url = service.Status.DeprecatedDomain
	}
	fmt.Fprintln(out, "\nService URL:")
	fmt.Fprintf(out, "%s\n", url)
	return nil
}
<commit_msg>Add 'creator' annotation on create --force (#341)<commit_after>// Copyright © 2019 The Knative 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 service

import (
	"errors"
	"fmt"
	"io"

	"github.com/knative/client/pkg/kn/commands"
	"github.com/knative/client/pkg/serving/v1alpha1"
	"github.com/knative/serving/pkg/apis/serving"

	serving_v1alpha1_api "github.com/knative/serving/pkg/apis/serving/v1alpha1"
	"github.com/spf13/cobra"

	corev1 "k8s.io/api/core/v1"
	api_errors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func NewServiceCreateCommand(p *commands.KnParams) *cobra.Command {
	var editFlags ConfigurationEditFlags
	var waitFlags commands.WaitFlags

	serviceCreateCommand := &cobra.Command{
		Use:   "create NAME --image IMAGE",
		Short: "Create a service.",
		Example: `
  # Create a service 'mysvc' using image at dev.local/ns/image:latest
  kn service create mysvc --image dev.local/ns/image:latest

  # Create a service with multiple environment variables
  kn service create mysvc --env KEY1=VALUE1 --env KEY2=VALUE2 --image dev.local/ns/image:latest

  # Create or replace a service 's1' with image dev.local/ns/image:v2 using --force flag
  # if service 's1' doesn't exist, it's just a normal create operation
  kn service create --force s1 --image dev.local/ns/image:v2

  # Create or replace environment variables of service 's1' using --force flag
  kn service create --force s1 --env KEY1=NEW_VALUE1 --env NEW_KEY2=NEW_VALUE2 --image dev.local/ns/image:v1

  # Create service 'mysvc' with port 80
  kn service create mysvc --port 80 --image dev.local/ns/image:latest

  # Create or replace default resources of a service 's1' using --force flag
  # (earlier configured resource requests and limits will be replaced with default)
  # (earlier configured environment variables will be cleared too if any)
  kn service create --force s1 --image dev.local/ns/image:v1`,

		RunE: func(cmd *cobra.Command, args []string) (err error) {
			if len(args) != 1 {
				return errors.New("'service create' requires the service name given as single argument")
			}
			name := args[0]
			if editFlags.Image == "" {
				return errors.New("'service create' requires the image name to run provided with the --image option")
			}

			namespace, err := p.GetNamespace(cmd)
			if err != nil {
				return err
			}

			service, err := constructService(cmd, editFlags, name, namespace)
			if err != nil {
				return err
			}

			client, err := p.NewClient(namespace)
			if err != nil {
				return err
			}

			serviceExists, err := serviceExists(client, name, namespace)
			if err != nil {
				return err
			}

			if serviceExists {
				if !editFlags.ForceCreate {
					return fmt.Errorf(
						"cannot create service '%s' in namespace '%s' "+
							"because the service already exists and no --force option was given", name, namespace)
				}
				err = replaceService(client, service, namespace, cmd.OutOrStdout())
			} else {
				err = createService(client, service, namespace, cmd.OutOrStdout())
			}
			if err != nil {
				return err
			}

			if !waitFlags.Async {
				out := cmd.OutOrStdout()
				err := waitForService(client, name, out, waitFlags.TimeoutInSeconds)
				if err != nil {
					return err
				}
				return showUrl(client, name, namespace, out)
			}

			return nil
		},
	}
	commands.AddNamespaceFlags(serviceCreateCommand.Flags(), false)
	editFlags.AddCreateFlags(serviceCreateCommand)
	waitFlags.AddConditionWaitFlags(serviceCreateCommand, 60, "Create", "service")
	return serviceCreateCommand
}

// Duck type for writers having a flush
type flusher interface {
	Flush() error
}

func flush(out io.Writer) {
	if flusher, ok := out.(flusher); ok {
		flusher.Flush()
	}
}

func createService(client v1alpha1.KnClient, service *serving_v1alpha1_api.Service, namespace string, out io.Writer) error {
	err := client.CreateService(service)
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "Service '%s' successfully created in namespace '%s'.\n", service.Name, namespace)
	return nil
}

func replaceService(client v1alpha1.KnClient, service *serving_v1alpha1_api.Service, namespace string, out io.Writer) error {
	var retries = 0
	for {
		existingService, err := client.GetService(service.Name)
		if err != nil {
			return err
		}

		// Copy over some annotations that we want to keep around. Erase others
		copyList := []string{
			serving.CreatorAnnotation,
			serving.UpdaterAnnotation,
		}

		// If the target Annotation doesn't exist, create it even if
		// we don't end up copying anything over so that we erase all
		// existing annotations
		if service.Annotations == nil {
			service.Annotations = map[string]string{}
		}

		// Do the actual copy now, but only if it's in the source annotation
		for _, k := range copyList {
			if v, ok := existingService.Annotations[k]; ok {
				service.Annotations[k] = v
			}
		}

		service.ResourceVersion = existingService.ResourceVersion
		err = client.UpdateService(service)
		if err != nil {
			// Retry to update when a resource version conflict exists
			if api_errors.IsConflict(err) && retries < MaxUpdateRetries {
				retries++
				continue
			}
			return err
		}
		fmt.Fprintf(out, "Service '%s' successfully replaced in namespace '%s'.\n", service.Name, namespace)
		return nil
	}
}

func serviceExists(client v1alpha1.KnClient, name string, namespace string) (bool, error) {
	_, err := client.GetService(name)
	if api_errors.IsNotFound(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

// Create service struct from provided options
func constructService(cmd *cobra.Command, editFlags ConfigurationEditFlags, name string, namespace string) (*serving_v1alpha1_api.Service,
	error) {

	service := serving_v1alpha1_api.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
	}

	// TODO: Should it always be `runLatest` ?
	service.Spec.DeprecatedRunLatest = &serving_v1alpha1_api.RunLatestType{
		Configuration: serving_v1alpha1_api.ConfigurationSpec{
			DeprecatedRevisionTemplate: &serving_v1alpha1_api.RevisionTemplateSpec{
				Spec: serving_v1alpha1_api.RevisionSpec{
					DeprecatedContainer: &corev1.Container{},
				},
			},
		},
	}

	err := editFlags.Apply(&service, cmd)
	if err != nil {
		return nil, err
	}
	return &service, nil
}

func showUrl(client v1alpha1.KnClient, serviceName string, namespace string, out io.Writer) error {
	service, err := client.GetService(serviceName)
	if err != nil {
		return fmt.Errorf("cannot fetch service '%s' in namespace '%s' for extracting the URL: %v", serviceName, namespace, err)
	}
	url := service.Status.URL.String()
	if url == "" {
		url = service.Status.DeprecatedDomain
	}
	fmt.Fprintln(out, "\nService URL:")
	fmt.Fprintf(out, "%s\n", url)
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 rollout

import (
	"context"
	"fmt"
	"time"

	"github.com/spf13/cobra"

	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/genericclioptions/resource"
	watchtools "k8s.io/client-go/tools/watch"
	"k8s.io/kubernetes/pkg/kubectl"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/polymorphichelpers"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/util/interrupt"
)

var (
	status_long = templates.LongDesc(`
		Show the status of the rollout.

		By default 'rollout status' will watch the status of the latest rollout
		until it's done. If you don't want to wait for the rollout to finish then
		you can use --watch=false. Note that if a new rollout starts in-between, then
		'rollout status' will continue watching the latest revision. If you want to
		pin to a specific revision and abort if it is rolled over by another revision,
		use --revision=N where N is the revision you need to watch for.`)

	status_example = templates.Examples(`
		# Watch the rollout status of a deployment
		kubectl rollout status deployment/nginx`)
)

type RolloutStatusOptions struct {
	PrintFlags *genericclioptions.PrintFlags

	Namespace        string
	EnforceNamespace bool
	BuilderArgs      []string

	Watch    bool
	Revision int64

	StatusViewer func(*meta.RESTMapping) (kubectl.StatusViewer, error)
	Builder      func() *resource.Builder

	FilenameOptions *resource.FilenameOptions
	genericclioptions.IOStreams
}

func NewRolloutStatusOptions(streams genericclioptions.IOStreams) *RolloutStatusOptions {
	return &RolloutStatusOptions{
		PrintFlags:      genericclioptions.NewPrintFlags("").WithTypeSetter(scheme.Scheme),
		FilenameOptions: &resource.FilenameOptions{},
		IOStreams:       streams,
		Watch:           true,
	}
}

func NewCmdRolloutStatus(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	o := NewRolloutStatusOptions(streams)

	validArgs := []string{"deployment", "daemonset", "statefulset"}

	cmd := &cobra.Command{
		Use: "status (TYPE NAME | TYPE/NAME) [flags]",
		DisableFlagsInUseLine: true,
		Short:   i18n.T("Show the status of the rollout"),
		Long:    status_long,
		Example: status_example,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, args))
			cmdutil.CheckErr(o.Validate(cmd, args))
			cmdutil.CheckErr(o.Run())
		},
		ValidArgs: validArgs,
	}

	usage := "identifying the resource to get from a server."
	cmdutil.AddFilenameOptionFlags(cmd, o.FilenameOptions, usage)
	cmd.Flags().BoolVarP(&o.Watch, "watch", "w", o.Watch, "Watch the status of the rollout until it's done.")
	cmd.Flags().Int64Var(&o.Revision, "revision", o.Revision, "Pin to a specific revision for showing its status. Defaults to 0 (last revision).")

	return cmd
}

func (o *RolloutStatusOptions) Complete(f cmdutil.Factory, args []string) error {
	o.Builder = f.NewBuilder

	var err error
	o.Namespace, o.EnforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.BuilderArgs = args
	o.StatusViewer = func(mapping *meta.RESTMapping) (kubectl.StatusViewer, error) {
		return polymorphichelpers.StatusViewerFn(f, mapping)
	}
	return nil
}

func (o *RolloutStatusOptions) Validate(cmd *cobra.Command, args []string) error {
	if len(args) == 0 && cmdutil.IsFilenameSliceEmpty(o.FilenameOptions.Filenames) {
		return cmdutil.UsageErrorf(cmd, "Required resource not specified.")
	}

	if o.Revision < 0 {
		return fmt.Errorf("revision must be a positive integer: %v", o.Revision)
	}

	return nil
}

func (o *RolloutStatusOptions) Run() error {
	r := o.Builder().
		WithScheme(scheme.Scheme, scheme.Scheme.PrioritizedVersionsAllGroups()...).
		NamespaceParam(o.Namespace).DefaultNamespace().
		FilenameParam(o.EnforceNamespace, o.FilenameOptions).
		ResourceTypeOrNameArgs(true, o.BuilderArgs...).
		SingleResourceType().
		Latest().
		Do()
	err := r.Err()
	if err != nil {
		return err
	}

	infos, err := r.Infos()
	if err != nil {
		return err
	}
	if len(infos) != 1 {
		return fmt.Errorf("rollout status is only supported on individual resources and resource collections - %d resources were found", len(infos))
	}
	info := infos[0]
	mapping := info.ResourceMapping()

	obj, err := r.Object()
	if err != nil {
		return err
	}
	rv, err := meta.NewAccessor().ResourceVersion(obj)
	if err != nil {
		return err
	}

	statusViewer, err := o.StatusViewer(mapping)
	if err != nil {
		return err
	}

	// check if deployment's has finished the rollout
	status, done, err := statusViewer.Status(info.Namespace, info.Name, o.Revision)
	if err != nil {
		return err
	}
	fmt.Fprintf(o.Out, "%s", status)
	if done {
		return nil
	}

	shouldWatch := o.Watch
	if !shouldWatch {
		return nil
	}

	// watch for changes to the deployment
	w, err := r.Watch(rv)
	if err != nil {
		return err
	}

	// if the rollout isn't done yet, keep watching deployment status
	// TODO: expose timeout
	timeout := 0 * time.Second
	ctx, cancel := watchtools.ContextWithOptionalTimeout(context.Background(), timeout)
	defer cancel()
	intr := interrupt.New(nil, cancel)
	return intr.Run(func() error {
		_, err := watchtools.UntilWithoutRetry(ctx, w, func(e watch.Event) (bool, error) {
			// print deployment's status
			status, done, err := statusViewer.Status(info.Namespace, info.Name, o.Revision)
			if err != nil {
				return false, err
			}
			fmt.Fprintf(o.Out, "%s", status)
			// Quit waiting if the rollout is done
			if done {
				return true, nil
			}
			return false, nil
		})
		return err
	})
}
<commit_msg>Switch kubectl rollout status to UntilWithSync to avoid premature timeouts<commit_after>/*
Copyright 2016 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 rollout

import (
	"context"
	"fmt"
	"time"

	"github.com/spf13/cobra"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/genericclioptions/resource"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/tools/cache"
	watchtools "k8s.io/client-go/tools/watch"
	"k8s.io/kubernetes/pkg/kubectl"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/polymorphichelpers"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/util/interrupt"
)

var (
	status_long = templates.LongDesc(`
		Show the status of the rollout.

		By default 'rollout status' will watch the status of the latest rollout
		until it's done. If you don't want to wait for the rollout to finish then
		you can use --watch=false. Note that if a new rollout starts in-between, then
		'rollout status' will continue watching the latest revision. If you want to
		pin to a specific revision and abort if it is rolled over by another revision,
		use --revision=N where N is the revision you need to watch for.`)

	status_example = templates.Examples(`
		# Watch the rollout status of a deployment
		kubectl rollout status deployment/nginx`)
)

type RolloutStatusOptions struct {
	PrintFlags *genericclioptions.PrintFlags

	Namespace        string
	EnforceNamespace bool
	BuilderArgs      []string

	Watch    bool
	Revision int64
	Timeout  time.Duration

	StatusViewer  func(*meta.RESTMapping) (kubectl.StatusViewer, error)
	Builder       func() *resource.Builder
	DynamicClient dynamic.Interface

	FilenameOptions *resource.FilenameOptions
	genericclioptions.IOStreams
}

func NewRolloutStatusOptions(streams genericclioptions.IOStreams) *RolloutStatusOptions {
	return &RolloutStatusOptions{
		PrintFlags:      genericclioptions.NewPrintFlags("").WithTypeSetter(scheme.Scheme),
		FilenameOptions: &resource.FilenameOptions{},
		IOStreams:       streams,
		Watch:           true,
		Timeout:         0,
	}
}

func NewCmdRolloutStatus(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	o := NewRolloutStatusOptions(streams)

	validArgs := []string{"deployment", "daemonset", "statefulset"}

	cmd := &cobra.Command{
		Use: "status (TYPE NAME | TYPE/NAME) [flags]",
		DisableFlagsInUseLine: true,
		Short:   i18n.T("Show the status of the rollout"),
		Long:    status_long,
		Example: status_example,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, args))
			cmdutil.CheckErr(o.Validate(cmd, args))
			cmdutil.CheckErr(o.Run())
		},
		ValidArgs: validArgs,
	}

	usage := "identifying the resource to get from a server."
	cmdutil.AddFilenameOptionFlags(cmd, o.FilenameOptions, usage)
	cmd.Flags().BoolVarP(&o.Watch, "watch", "w", o.Watch, "Watch the status of the rollout until it's done.")
	cmd.Flags().Int64Var(&o.Revision, "revision", o.Revision, "Pin to a specific revision for showing its status. Defaults to 0 (last revision).")
	cmd.Flags().DurationVar(&o.Timeout, "timeout", o.Timeout, "The length of time to wait before ending watch, zero means never. Any other values should contain a corresponding time unit (e.g. 1s, 2m, 3h).")

	return cmd
}

func (o *RolloutStatusOptions) Complete(f cmdutil.Factory, args []string) error {
	o.Builder = f.NewBuilder

	var err error
	o.Namespace, o.EnforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.BuilderArgs = args
	o.StatusViewer = func(mapping *meta.RESTMapping) (kubectl.StatusViewer, error) {
		return polymorphichelpers.StatusViewerFn(f, mapping)
	}

	clientConfig, err := f.ToRESTConfig()
	if err != nil {
		return err
	}

	o.DynamicClient, err = dynamic.NewForConfig(clientConfig)
	if err != nil {
		return err
	}

	return nil
}

func (o *RolloutStatusOptions) Validate(cmd *cobra.Command, args []string) error {
	if len(args) == 0 && cmdutil.IsFilenameSliceEmpty(o.FilenameOptions.Filenames) {
		return cmdutil.UsageErrorf(cmd, "Required resource not specified.")
	}

	if o.Revision < 0 {
		return fmt.Errorf("revision must be a positive integer: %v", o.Revision)
	}

	return nil
}

func (o *RolloutStatusOptions) Run() error {
	r := o.Builder().
		WithScheme(scheme.Scheme, scheme.Scheme.PrioritizedVersionsAllGroups()...).
		NamespaceParam(o.Namespace).DefaultNamespace().
		FilenameParam(o.EnforceNamespace, o.FilenameOptions).
		ResourceTypeOrNameArgs(true, o.BuilderArgs...).
		SingleResourceType().
		Latest().
		Do()
	err := r.Err()
	if err != nil {
		return err
	}

	infos, err := r.Infos()
	if err != nil {
		return err
	}
	if len(infos) != 1 {
		return fmt.Errorf("rollout status is only supported on individual resources and resource collections - %d resources were found", len(infos))
	}
	info := infos[0]
	mapping := info.ResourceMapping()

	statusViewer, err := o.StatusViewer(mapping)
	if err != nil {
		return err
	}

	fieldSelector := fields.OneTermEqualSelector("metadata.name", info.Name).String()
	lw := &cache.ListWatch{
		ListFunc: func(options metav1.ListOptions) (runtime.Object, error) {
			options.FieldSelector = fieldSelector
			return o.DynamicClient.Resource(info.Mapping.Resource).Namespace(info.Namespace).List(options)
		},
		WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
			options.FieldSelector = fieldSelector
			return o.DynamicClient.Resource(info.Mapping.Resource).Namespace(info.Namespace).Watch(options)
		},
	}

	preconditionFunc := func(store cache.Store) (bool, error) {
		_, exists, err := store.Get(&metav1.ObjectMeta{Namespace: info.Namespace, Name: info.Name})
		if err != nil {
			return true, err
		}
		if !exists {
			// We need to make sure we see the object in the cache before we start waiting for events
			// or we would be waiting for the timeout if such object didn't exist.
			return true, apierrors.NewNotFound(mapping.Resource.GroupResource(), info.Name)
		}

		return false, nil
	}

	// if the rollout isn't done yet, keep watching deployment status
	ctx, cancel := watchtools.ContextWithOptionalTimeout(context.Background(), o.Timeout)
	intr := interrupt.New(nil, cancel)
	return intr.Run(func() error {
		_, err = watchtools.UntilWithSync(ctx, lw, &unstructured.Unstructured{}, preconditionFunc, func(e watch.Event) (bool, error) {
			switch t := e.Type; t {
			case watch.Added, watch.Modified:
				status, done, err := statusViewer.Status(e.Object.(runtime.Unstructured), o.Revision)
				if err != nil {
					return false, err
				}
				fmt.Fprintf(o.Out, "%s", status)
				// Quit waiting if the rollout is done
				if done {
					return true, nil
				}

				shouldWatch := o.Watch
				if !shouldWatch {
					return true, nil
				}

				return false, nil

			case watch.Deleted:
				// We need to abort to avoid cases of recreation and not to silently watch the wrong (new) object
				return true, fmt.Errorf("object has been deleted")

			default:
				return true, fmt.Errorf("internal error: unexpected event %#v", e)
			}
		})
		return err
	})
}
<|endoftext|>


<commit_before>package provisioning

import (
	"context"
	"fmt"
	"path"

	"github.com/grafana/grafana/pkg/registry"
	"github.com/grafana/grafana/pkg/services/provisioning/notifiers"
	"github.com/grafana/grafana/pkg/services/provisioning/dashboards"
	"github.com/grafana/grafana/pkg/services/provisioning/datasources"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	registry.RegisterService(&ProvisioningService{})
}

type ProvisioningService struct {
	Cfg *setting.Cfg `inject:""`
}

func (ps *ProvisioningService) Init() error {
	datasourcePath := path.Join(ps.Cfg.ProvisioningPath, "datasources")
	if err := datasources.Provision(datasourcePath); err != nil {
		return fmt.Errorf("Datasource provisioning error: %v", err)
	}

	alertNotificationsPath := path.Join(ps.Cfg.ProvisioningPath, "notifiers")
	if err := notifiers.Provision(alertNotificationsPath); err != nil {
		return fmt.Errorf("Alert notification provisioning error: %v", err)
	}

	return nil
}

func (ps *ProvisioningService) Run(ctx context.Context) error {
	dashboardPath := path.Join(ps.Cfg.ProvisioningPath, "dashboards")
	dashProvisioner := dashboards.NewDashboardProvisioner(dashboardPath)

	if err := dashProvisioner.Provision(ctx); err != nil {
		return err
	}

	<-ctx.Done()
	return ctx.Err()
}
<commit_msg>gofmt issue<commit_after>package provisioning

import (
	"context"
	"fmt"
	"path"

	"github.com/grafana/grafana/pkg/registry"
	"github.com/grafana/grafana/pkg/services/provisioning/dashboards"
	"github.com/grafana/grafana/pkg/services/provisioning/datasources"
	"github.com/grafana/grafana/pkg/services/provisioning/notifiers"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	registry.RegisterService(&ProvisioningService{})
}

type ProvisioningService struct {
	Cfg *setting.Cfg `inject:""`
}

func (ps *ProvisioningService) Init() error {
	datasourcePath := path.Join(ps.Cfg.ProvisioningPath, "datasources")
	if err := datasources.Provision(datasourcePath); err != nil {
		return fmt.Errorf("Datasource provisioning error: %v", err)
	}

	alertNotificationsPath := path.Join(ps.Cfg.ProvisioningPath, "notifiers")
	if err := notifiers.Provision(alertNotificationsPath); err != nil {
		return fmt.Errorf("Alert notification provisioning error: %v", err)
	}

	return nil
}

func (ps *ProvisioningService) Run(ctx context.Context) error {
	dashboardPath := path.Join(ps.Cfg.ProvisioningPath, "dashboards")
	dashProvisioner := dashboards.NewDashboardProvisioner(dashboardPath)

	if err := dashProvisioner.Provision(ctx); err != nil {
		return err
	}

	<-ctx.Done()
	return ctx.Err()
}
<|endoftext|>


<commit_before>/*
In order to keep track of active weave peers, we use annotations on the Kubernetes cluster.

Kubernetes uses etcd to distribute and synchronise these annotations so we don't have to.
*/
package main

import (
	"log"
	"time"

	"github.com/pkg/errors"

	v1 "k8s.io/api/core/v1"
	kubeErrors "k8s.io/apimachinery/pkg/api/errors"
	api "k8s.io/apimachinery/pkg/apis/meta/v1"
	wait "k8s.io/apimachinery/pkg/util/wait"
	kubernetes "k8s.io/client-go/kubernetes"
	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
)

type configMapAnnotations struct {
	ConfigMapName string
	Namespace     string
	Client        corev1client.ConfigMapsGetter
	cm            *v1.ConfigMap
}

func newConfigMapAnnotations(ns string, configMapName string, c kubernetes.Interface) *configMapAnnotations {
	return &configMapAnnotations{
		Namespace:     ns,
		ConfigMapName: configMapName,
		Client:        c.CoreV1(),
	}
}

const (
	retryPeriod  = time.Second * 2
	jitterFactor = 1.0

	// Prefix all our annotation keys with this string so they don't clash with anyone else's
	KubePeersPrefix = "kube-peers.weave.works/"
	// KubePeersAnnotationKey is the default annotation key
	KubePeersAnnotationKey = KubePeersPrefix + "peers"
)

func (cml *configMapAnnotations) Init() error {
	for {
		// Since it's potentially racy to GET, then CREATE if not found, we wrap in a check loop
		// so that if the configmap is created after our GET but before or CREATE, we'll gracefully
		// re-try to get the configmap.
		var err error
		cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Get(cml.ConfigMapName, api.GetOptions{})
		if err != nil {
			if !kubeErrors.IsNotFound(err) {
				return errors.Wrapf(err, "Unable to fetch ConfigMap %s/%s", cml.Namespace, cml.ConfigMapName)
			}
			cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Create(&v1.ConfigMap{
				ObjectMeta: api.ObjectMeta{
					Name:      cml.ConfigMapName,
					Namespace: cml.Namespace,
				},
			})
			if err != nil {
				if kubeErrors.IsAlreadyExists(err) {
					continue
				}
				return errors.Wrapf(err, "Unable to create ConfigMap %s/%s", cml.Namespace, cml.ConfigMapName)
			}
		}
		break
	}
	if cml.cm.Annotations == nil {
		cml.cm.Annotations = make(map[string]string)
	}
	return nil
}

// Clean up a string so it meets the Kubernetes requiremements for Annotation keys:
// name part must consist of alphanumeric characters, '-', '_' or '.', and must
// start and end with an alphanumeric character (e.g. 'MyName', or 'my.name', or '123-abc')
func cleanKey(key string) string {
	buf := []byte(key)
	for i, c := range buf {
		if (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-' || c == '_' || c == '.' || c == '/' {
			continue
		}
		buf[i] = '_'
	}
	return string(buf)
}

func (cml *configMapAnnotations) GetAnnotation(key string) (string, bool) {
	value, ok := cml.cm.Annotations[cleanKey(key)]
	return value, ok
}

func (cml *configMapAnnotations) UpdateAnnotation(key, value string) (err error) {
	if cml.cm == nil {
		return errors.New("endpoint not initialized, call Init first")
	}
	// speculatively change the state, then replace with whatever comes back
	// from Update(), which will be the latest on the server, or nil if error
	cml.cm.Annotations[cleanKey(key)] = value
	cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Update(cml.cm)
	return err
}

func (cml *configMapAnnotations) RemoveAnnotation(key string) (err error) {
	if cml.cm == nil {
		return errors.New("endpoint not initialized, call Init first")
	}
	// speculatively change the state, then replace with whatever comes back
	// from Update(), which will be the latest on the server, or nil if error
	delete(cml.cm.Annotations, cleanKey(key))
	cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Update(cml.cm)
	return err
}

func (cml *configMapAnnotations) RemoveAnnotationsWithValue(valueToRemove string) (err error) {
	if cml.cm == nil {
		return errors.New("endpoint not initialized, call Init first")
	}
	// speculatively change the state, then replace with whatever comes back
	// from Update(), which will be the latest on the server, or nil if error
	for key, value := range cml.cm.Annotations {
		if value == valueToRemove {
			delete(cml.cm.Annotations, key) // don't need to clean this key as it came from the map
		}
	}
	cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Update(cml.cm)
	return err
}

// Loop with jitter, fetching the cml data and calling f() until it
// doesn't get an optimistic locking conflict.
// If it succeeds or gets any other kind of error, stop the loop.
func (cml *configMapAnnotations) LoopUpdate(f func() error) error {
	stop := make(chan struct{})
	var err error
	wait.JitterUntil(func() {
		if err = cml.Init(); err != nil {
			close(stop)
			return
		}
		err = f()
		if err != nil && kubeErrors.IsConflict(err) {
			log.Printf("Optimistic locking conflict: trying again: %s", err)
			return
		}
		close(stop)
	}, retryPeriod, jitterFactor, true, stop)
	return err
}
<commit_msg>Avoid panic on nil map<commit_after>/*
In order to keep track of active weave peers, we use annotations on the Kubernetes cluster.

Kubernetes uses etcd to distribute and synchronise these annotations so we don't have to.
*/
package main

import (
	"log"
	"time"

	"github.com/pkg/errors"

	v1 "k8s.io/api/core/v1"
	kubeErrors "k8s.io/apimachinery/pkg/api/errors"
	api "k8s.io/apimachinery/pkg/apis/meta/v1"
	wait "k8s.io/apimachinery/pkg/util/wait"
	kubernetes "k8s.io/client-go/kubernetes"
	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
)

type configMapAnnotations struct {
	ConfigMapName string
	Namespace     string
	Client        corev1client.ConfigMapsGetter
	cm            *v1.ConfigMap
}

func newConfigMapAnnotations(ns string, configMapName string, c kubernetes.Interface) *configMapAnnotations {
	return &configMapAnnotations{
		Namespace:     ns,
		ConfigMapName: configMapName,
		Client:        c.CoreV1(),
	}
}

const (
	retryPeriod  = time.Second * 2
	jitterFactor = 1.0

	// Prefix all our annotation keys with this string so they don't clash with anyone else's
	KubePeersPrefix = "kube-peers.weave.works/"
	// KubePeersAnnotationKey is the default annotation key
	KubePeersAnnotationKey = KubePeersPrefix + "peers"
)

func (cml *configMapAnnotations) Init() error {
	for {
		// Since it's potentially racy to GET, then CREATE if not found, we wrap in a check loop
		// so that if the configmap is created after our GET but before or CREATE, we'll gracefully
		// re-try to get the configmap.
		var err error
		cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Get(cml.ConfigMapName, api.GetOptions{})
		if err != nil {
			if !kubeErrors.IsNotFound(err) {
				return errors.Wrapf(err, "Unable to fetch ConfigMap %s/%s", cml.Namespace, cml.ConfigMapName)
			}
			cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Create(&v1.ConfigMap{
				ObjectMeta: api.ObjectMeta{
					Name:      cml.ConfigMapName,
					Namespace: cml.Namespace,
				},
			})
			if err != nil {
				if kubeErrors.IsAlreadyExists(err) {
					continue
				}
				return errors.Wrapf(err, "Unable to create ConfigMap %s/%s", cml.Namespace, cml.ConfigMapName)
			}
		}
		break
	}
	if cml.cm.Annotations == nil {
		cml.cm.Annotations = make(map[string]string)
	}
	return nil
}

// Clean up a string so it meets the Kubernetes requiremements for Annotation keys:
// name part must consist of alphanumeric characters, '-', '_' or '.', and must
// start and end with an alphanumeric character (e.g. 'MyName', or 'my.name', or '123-abc')
func cleanKey(key string) string {
	buf := []byte(key)
	for i, c := range buf {
		if (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-' || c == '_' || c == '.' || c == '/' {
			continue
		}
		buf[i] = '_'
	}
	return string(buf)
}

func (cml *configMapAnnotations) GetAnnotation(key string) (string, bool) {
	value, ok := cml.cm.Annotations[cleanKey(key)]
	return value, ok
}

func (cml *configMapAnnotations) UpdateAnnotation(key, value string) (err error) {
	if cml.cm == nil || cml.cm.Annotations == nil {
		return errors.New("endpoint not initialized, call Init first")
	}
	// speculatively change the state, then replace with whatever comes back
	// from Update(), which will be the latest on the server, or nil if error
	cml.cm.Annotations[cleanKey(key)] = value
	cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Update(cml.cm)
	return err
}

func (cml *configMapAnnotations) RemoveAnnotation(key string) (err error) {
	if cml.cm == nil || cml.cm.Annotations == nil {
		return errors.New("endpoint not initialized, call Init first")
	}
	// speculatively change the state, then replace with whatever comes back
	// from Update(), which will be the latest on the server, or nil if error
	delete(cml.cm.Annotations, cleanKey(key))
	cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Update(cml.cm)
	return err
}

func (cml *configMapAnnotations) RemoveAnnotationsWithValue(valueToRemove string) (err error) {
	if cml.cm == nil || cml.cm.Annotations == nil {
		return errors.New("endpoint not initialized, call Init first")
	}
	// speculatively change the state, then replace with whatever comes back
	// from Update(), which will be the latest on the server, or nil if error
	for key, value := range cml.cm.Annotations {
		if value == valueToRemove {
			delete(cml.cm.Annotations, key) // don't need to clean this key as it came from the map
		}
	}
	cml.cm, err = cml.Client.ConfigMaps(cml.Namespace).Update(cml.cm)
	return err
}

// Loop with jitter, fetching the cml data and calling f() until it
// doesn't get an optimistic locking conflict.
// If it succeeds or gets any other kind of error, stop the loop.
func (cml *configMapAnnotations) LoopUpdate(f func() error) error {
	stop := make(chan struct{})
	var err error
	wait.JitterUntil(func() {
		if err = cml.Init(); err != nil {
			close(stop)
			return
		}
		err = f()
		if err != nil && kubeErrors.IsConflict(err) {
			log.Printf("Optimistic locking conflict: trying again: %s", err)
			return
		}
		close(stop)
	}, retryPeriod, jitterFactor, true, stop)
	return err
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/budgets"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsBudgetsBudget() *schema.Resource {
	return &schema.Resource{
		Schema: map[string]*schema.Schema{
			"account_id": {
				Type:         schema.TypeString,
				Optional:     true,
				ForceNew:     true,
				ValidateFunc: validateAwsAccountId,
			},
			"name": {
				Type:          schema.TypeString,
				Computed:      true,
				Optional:      true,
				ConflictsWith: []string{"name_prefix"},
			},
			"name_prefix": {
				Type:     schema.TypeString,
				Computed: true,
				Optional: true,
				ForceNew: true,
			},
			"budget_type": {
				Type:     schema.TypeString,
				Required: true,
			},
			"limit_amount": {
				Type:     schema.TypeString,
				Required: true,
			},
			"limit_unit": {
				Type:     schema.TypeString,
				Required: true,
			},
			"cost_types": {
				Type:     schema.TypeList,
				Optional: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"include_credit": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_other_subscription": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_recurring": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_refund": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_subscription": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_support": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_tax": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_upfront": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"use_blended": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
					},
				},
			},
			"time_period_start": {
				Type:     schema.TypeString,
				Required: true,
			},
			"time_period_end": {
				Type:     schema.TypeString,
				Optional: true,
				Default:  "2087-06-15_00:00",
			},
			"time_unit": {
				Type:     schema.TypeString,
				Required: true,
			},
			"cost_filters": {
				Type:     schema.TypeMap,
				Optional: true,
				Computed: true,
			},
		},
		Create: resourceAwsBudgetsBudgetCreate,
		Read:   resourceAwsBudgetsBudgetRead,
		Update: resourceAwsBudgetsBudgetUpdate,
		Delete: resourceAwsBudgetsBudgetDelete,
	}
}

func resourceAwsBudgetsBudgetCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).budgetconn
	var accountID string
	if v, ok := d.GetOk("account_id"); ok {
		accountID = v.(string)
	} else {
		accountID = meta.(*AWSClient).accountid
	}

	budget, err := expandBudgetsBudgetUnmarshal(d)
	if err != nil {
		return fmt.Errorf("failed creating budget: %v", err)
	}

	_, err = client.CreateBudget(&budgets.CreateBudgetInput{
		AccountId: aws.String(accountID),
		Budget:    budget,
	})
	if err != nil {
		return fmt.Errorf("create budget failed: %v", err)
	}

	d.SetId(fmt.Sprintf("%s:%s", accountID, *budget.BudgetName))
	return resourceAwsBudgetsBudgetRead(d, meta)
}

func resourceAwsBudgetsBudgetRead(d *schema.ResourceData, meta interface{}) error {
	accountID, budgetName, err := decodeBudgetsBudgetID(d.Id())
	if err != nil {
		return err
	}

	client := meta.(*AWSClient).budgetconn
	describeBudgetOutput, err := client.DescribeBudget(&budgets.DescribeBudgetInput{
		BudgetName: aws.String(budgetName),
		AccountId:  aws.String(accountID),
	})
	if isAWSErr(err, budgets.ErrCodeNotFoundException, "") {
		log.Printf("[WARN] Budget %s not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("describe budget failed: %v", err)
	}

	flattenedBudget, err := expandBudgetsBudgetFlatten(describeBudgetOutput.Budget)
	if err != nil {
		return fmt.Errorf("failed flattening budget output: %v", err)
	}

	if err := d.Set("name", flattenedBudget.name); err != nil {
		return err
	}

	for k, v := range map[string]interface{}{
		"account_id":        accountID,
		"budget_type":       flattenedBudget.budgetType,
		"time_unit":         flattenedBudget.timeUnit,
		"cost_filters":      convertCostFiltersToStringMap(flattenedBudget.costFilters),
		"limit_amount":      flattenedBudget.limitAmount,
		"limit_unit":        flattenedBudget.limitUnit,
		"cost_types":        []interface{}{flattenedBudget.costTypes},
		"time_period_start": flattenedBudget.timePeriodStart.Format("2006-01-02_15:04"),
		"time_period_end":   flattenedBudget.timePeriodEnd.Format("2006-01-02_15:04"),
	} {
		if _, ok := d.GetOk(k); ok {
			if err := d.Set(k, v); err != nil {
				return err
			}
		}
	}
	return nil
}

func resourceAwsBudgetsBudgetUpdate(d *schema.ResourceData, meta interface{}) error {
	accountID, _, err := decodeBudgetsBudgetID(d.Id())
	if err != nil {
		return err
	}

	client := meta.(*AWSClient).budgetconn
	budget, err := expandBudgetsBudgetUnmarshal(d)
	if err != nil {
		return fmt.Errorf("could not create budget: %v", err)
	}

	_, err = client.UpdateBudget(&budgets.UpdateBudgetInput{
		AccountId: aws.String(accountID),
		NewBudget: budget,
	})
	if err != nil {
		return fmt.Errorf("update budget failed: %v", err)
	}

	return resourceAwsBudgetsBudgetRead(d, meta)
}

func resourceAwsBudgetsBudgetDelete(d *schema.ResourceData, meta interface{}) error {
	accountID, budgetName, err := decodeBudgetsBudgetID(d.Id())
	if err != nil {
		return err
	}

	client := meta.(*AWSClient).budgetconn
	_, err = client.DescribeBudget(&budgets.DescribeBudgetInput{
		BudgetName: aws.String(budgetName),
		AccountId:  aws.String(accountID),
	})
	if isAWSErr(err, budgets.ErrCodeNotFoundException, "") {
		log.Printf("[INFO] budget %s could not be found. skipping delete.", d.Id())
		return nil
	}

	_, err = client.DeleteBudget(&budgets.DeleteBudgetInput{
		BudgetName: aws.String(budgetName),
		AccountId:  aws.String(accountID),
	})
	if err != nil {
		return fmt.Errorf("delete budget failed: %v", err)
	}

	return nil
}

type flattenedBudgetsBudget struct {
	name            *string
	budgetType      *string
	timeUnit        *string
	costFilters     map[string][]*string
	limitAmount     *string
	limitUnit       *string
	costTypes       map[string]bool
	timePeriodStart *time.Time
	timePeriodEnd   *time.Time
}

func expandBudgetsBudgetFlatten(budget *budgets.Budget) (*flattenedBudgetsBudget, error) {
	if budget == nil {
		return nil, fmt.Errorf("empty budget returned from budget output: %v", budget)
	}

	budgetLimit := budget.BudgetLimit
	if budgetLimit == nil {
		return nil, fmt.Errorf("empty limit in budget: %v", budget)
	}

	budgetCostTypes := budget.CostTypes
	if budgetCostTypes == nil {
		return nil, fmt.Errorf("empty CostTypes in budget: %v", budget)
	}

	costTypesMap := map[string]bool{
		"include_credit":             *budgetCostTypes.IncludeCredit,
		"include_other_subscription": *budgetCostTypes.IncludeOtherSubscription,
		"include_recurring":          *budgetCostTypes.IncludeRecurring,
		"include_refund":             *budgetCostTypes.IncludeRefund,
		"include_subscription":       *budgetCostTypes.IncludeSubscription,
		"include_support":            *budgetCostTypes.IncludeSupport,
		"include_tax":                *budgetCostTypes.IncludeTax,
		"include_upfront":            *budgetCostTypes.IncludeUpfront,
		"use_blended":                *budgetCostTypes.UseBlended,
	}
	budgetTimePeriod := budget.TimePeriod
	if budgetTimePeriod == nil {
		return nil, fmt.Errorf("empty TimePeriod in budget: %v", budget)
	}

	budgetTimePeriodStart := budgetTimePeriod.Start
	if budgetTimePeriodStart == nil {
		return nil, fmt.Errorf("empty TimePeriodStart in budget: %v", budget)
	}

	budgetTimePeriodEnd := budgetTimePeriod.End
	if budgetTimePeriodEnd == nil {
		return nil, fmt.Errorf("empty TimePeriodEnd in budget: %v", budget)
	}

	return &flattenedBudgetsBudget{
		name:            budget.BudgetName,
		budgetType:      budget.BudgetType,
		timeUnit:        budget.TimeUnit,
		costFilters:     budget.CostFilters,
		limitAmount:     budgetLimit.Amount,
		limitUnit:       budgetLimit.Unit,
		costTypes:       costTypesMap,
		timePeriodStart: budgetTimePeriodStart,
		timePeriodEnd:   budgetTimePeriodEnd,
	}, nil
}

func convertCostFiltersToStringMap(costFilters map[string][]*string) map[string]string {
	convertedCostFilters := make(map[string]string)
	for k, v := range costFilters {
		filterValues := make([]string, 0)
		for _, singleFilterValue := range v {
			filterValues = append(filterValues, *singleFilterValue)
		}

		convertedCostFilters[k] = strings.Join(filterValues, ",")
	}

	return convertedCostFilters
}

func expandBudgetsBudgetUnmarshal(d *schema.ResourceData) (*budgets.Budget, error) {
	var budgetName string
	if _, id, err := decodeBudgetsBudgetID(d.Id()); err == nil && id != "" {
		budgetName = id

	} else if v, ok := d.GetOk("name"); ok {
		budgetName = v.(string)

	} else if v, ok := d.GetOk("name_prefix"); ok {
		budgetName = resource.PrefixedUniqueId(v.(string))

	} else {
		budgetName = resource.UniqueId()
	}

	budgetType := d.Get("budget_type").(string)
	budgetLimitAmount := d.Get("limit_amount").(string)
	budgetLimitUnit := d.Get("limit_unit").(string)
	costTypes := expandBudgetsCostTypesUnmarshal(d.Get("cost_types").([]interface{}))
	budgetTimeUnit := d.Get("time_unit").(string)
	budgetCostFilters := make(map[string][]*string)
	for k, v := range d.Get("cost_filters").(map[string]interface{}) {
		filterValue := v.(string)
		budgetCostFilters[k] = append(budgetCostFilters[k], aws.String(filterValue))
	}

	budgetTimePeriodStart, err := time.Parse("2006-01-02_15:04", d.Get("time_period_start").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budgetTimePeriodEnd, err := time.Parse("2006-01-02_15:04", d.Get("time_period_end").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budget := &budgets.Budget{
		BudgetName: aws.String(budgetName),
		BudgetType: aws.String(budgetType),
		BudgetLimit: &budgets.Spend{
			Amount: aws.String(budgetLimitAmount),
			Unit:   aws.String(budgetLimitUnit),
		},
		CostTypes: costTypes,
		TimePeriod: &budgets.TimePeriod{
			End:   &budgetTimePeriodEnd,
			Start: &budgetTimePeriodStart,
		},
		TimeUnit:    aws.String(budgetTimeUnit),
		CostFilters: budgetCostFilters,
	}
	return budget, nil
}

func decodeBudgetsBudgetID(id string) (string, string, error) {
	parts := strings.Split(id, ":")
	if len(parts) != 2 {
		return "", "", fmt.Errorf("Unexpected format of ID (%q), expected AccountID:BudgetName", id)
	}
	return parts[0], parts[1], nil
}

func expandBudgetsCostTypesUnmarshal(budgetCostTypes []interface{}) *budgets.CostTypes {
	costTypes := &budgets.CostTypes{
		IncludeCredit:            aws.Bool(true),
		IncludeOtherSubscription: aws.Bool(true),
		IncludeRecurring:         aws.Bool(true),
		IncludeRefund:            aws.Bool(true),
		IncludeSubscription:      aws.Bool(true),
		IncludeSupport:           aws.Bool(true),
		IncludeTax:               aws.Bool(true),
		IncludeUpfront:           aws.Bool(true),
		UseBlended:               aws.Bool(false),
	}
	if len(budgetCostTypes) == 1 {
		costTypesMap := budgetCostTypes[0].(map[string]interface{})
		for k, v := range map[string]*bool{
			"include_credit":             costTypes.IncludeCredit,
			"include_other_subscription": costTypes.IncludeOtherSubscription,
			"include_recurring":          costTypes.IncludeRecurring,
			"include_refund":             costTypes.IncludeRefund,
			"include_subscription":       costTypes.IncludeSubscription,
			"include_support":            costTypes.IncludeSupport,
			"include_tax":                costTypes.IncludeTax,
			"include_upfront":            costTypes.IncludeUpfront,
			"use_blended":                costTypes.UseBlended,
		} {
			if val, ok := costTypesMap[k]; ok {
				*v = val.(bool)
			}
		}
	}

	return costTypes
}
<commit_msg>remove redundant call on resource delete<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/budgets"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsBudgetsBudget() *schema.Resource {
	return &schema.Resource{
		Schema: map[string]*schema.Schema{
			"account_id": {
				Type:         schema.TypeString,
				Optional:     true,
				ForceNew:     true,
				ValidateFunc: validateAwsAccountId,
			},
			"name": {
				Type:          schema.TypeString,
				Computed:      true,
				Optional:      true,
				ConflictsWith: []string{"name_prefix"},
			},
			"name_prefix": {
				Type:     schema.TypeString,
				Computed: true,
				Optional: true,
				ForceNew: true,
			},
			"budget_type": {
				Type:     schema.TypeString,
				Required: true,
			},
			"limit_amount": {
				Type:     schema.TypeString,
				Required: true,
			},
			"limit_unit": {
				Type:     schema.TypeString,
				Required: true,
			},
			"cost_types": {
				Type:     schema.TypeList,
				Optional: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"include_credit": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_other_subscription": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_recurring": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_refund": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_subscription": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_support": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_tax": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"include_upfront": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"use_blended": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
					},
				},
			},
			"time_period_start": {
				Type:     schema.TypeString,
				Required: true,
			},
			"time_period_end": {
				Type:     schema.TypeString,
				Optional: true,
				Default:  "2087-06-15_00:00",
			},
			"time_unit": {
				Type:     schema.TypeString,
				Required: true,
			},
			"cost_filters": {
				Type:     schema.TypeMap,
				Optional: true,
				Computed: true,
			},
		},
		Create: resourceAwsBudgetsBudgetCreate,
		Read:   resourceAwsBudgetsBudgetRead,
		Update: resourceAwsBudgetsBudgetUpdate,
		Delete: resourceAwsBudgetsBudgetDelete,
	}
}

func resourceAwsBudgetsBudgetCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).budgetconn
	var accountID string
	if v, ok := d.GetOk("account_id"); ok {
		accountID = v.(string)
	} else {
		accountID = meta.(*AWSClient).accountid
	}

	budget, err := expandBudgetsBudgetUnmarshal(d)
	if err != nil {
		return fmt.Errorf("failed creating budget: %v", err)
	}

	_, err = client.CreateBudget(&budgets.CreateBudgetInput{
		AccountId: aws.String(accountID),
		Budget:    budget,
	})
	if err != nil {
		return fmt.Errorf("create budget failed: %v", err)
	}

	d.SetId(fmt.Sprintf("%s:%s", accountID, *budget.BudgetName))
	return resourceAwsBudgetsBudgetRead(d, meta)
}

func resourceAwsBudgetsBudgetRead(d *schema.ResourceData, meta interface{}) error {
	accountID, budgetName, err := decodeBudgetsBudgetID(d.Id())
	if err != nil {
		return err
	}

	client := meta.(*AWSClient).budgetconn
	describeBudgetOutput, err := client.DescribeBudget(&budgets.DescribeBudgetInput{
		BudgetName: aws.String(budgetName),
		AccountId:  aws.String(accountID),
	})
	if isAWSErr(err, budgets.ErrCodeNotFoundException, "") {
		log.Printf("[WARN] Budget %s not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("describe budget failed: %v", err)
	}

	flattenedBudget, err := expandBudgetsBudgetFlatten(describeBudgetOutput.Budget)
	if err != nil {
		return fmt.Errorf("failed flattening budget output: %v", err)
	}

	if err := d.Set("name", flattenedBudget.name); err != nil {
		return err
	}

	for k, v := range map[string]interface{}{
		"account_id":        accountID,
		"budget_type":       flattenedBudget.budgetType,
		"time_unit":         flattenedBudget.timeUnit,
		"cost_filters":      convertCostFiltersToStringMap(flattenedBudget.costFilters),
		"limit_amount":      flattenedBudget.limitAmount,
		"limit_unit":        flattenedBudget.limitUnit,
		"cost_types":        []interface{}{flattenedBudget.costTypes},
		"time_period_start": flattenedBudget.timePeriodStart.Format("2006-01-02_15:04"),
		"time_period_end":   flattenedBudget.timePeriodEnd.Format("2006-01-02_15:04"),
	} {
		if _, ok := d.GetOk(k); ok {
			if err := d.Set(k, v); err != nil {
				return err
			}
		}
	}
	return nil
}

func resourceAwsBudgetsBudgetUpdate(d *schema.ResourceData, meta interface{}) error {
	accountID, _, err := decodeBudgetsBudgetID(d.Id())
	if err != nil {
		return err
	}

	client := meta.(*AWSClient).budgetconn
	budget, err := expandBudgetsBudgetUnmarshal(d)
	if err != nil {
		return fmt.Errorf("could not create budget: %v", err)
	}

	_, err = client.UpdateBudget(&budgets.UpdateBudgetInput{
		AccountId: aws.String(accountID),
		NewBudget: budget,
	})
	if err != nil {
		return fmt.Errorf("update budget failed: %v", err)
	}

	return resourceAwsBudgetsBudgetRead(d, meta)
}

func resourceAwsBudgetsBudgetDelete(d *schema.ResourceData, meta interface{}) error {
	accountID, budgetName, err := decodeBudgetsBudgetID(d.Id())
	if err != nil {
		return err
	}

	client := meta.(*AWSClient).budgetconn
	_, err = client.DeleteBudget(&budgets.DeleteBudgetInput{
		BudgetName: aws.String(budgetName),
		AccountId:  aws.String(accountID),
	})
	if err != nil {
		if isAWSErr(err, budgets.ErrCodeNotFoundException, "") {
			log.Printf("[INFO] budget %s could not be found. skipping delete.", d.Id())
			return nil
		}

		return fmt.Errorf("delete budget failed: %v", err)
	}

	return nil
}

type flattenedBudgetsBudget struct {
	name            *string
	budgetType      *string
	timeUnit        *string
	costFilters     map[string][]*string
	limitAmount     *string
	limitUnit       *string
	costTypes       map[string]bool
	timePeriodStart *time.Time
	timePeriodEnd   *time.Time
}

func expandBudgetsBudgetFlatten(budget *budgets.Budget) (*flattenedBudgetsBudget, error) {
	if budget == nil {
		return nil, fmt.Errorf("empty budget returned from budget output: %v", budget)
	}

	budgetLimit := budget.BudgetLimit
	if budgetLimit == nil {
		return nil, fmt.Errorf("empty limit in budget: %v", budget)
	}

	budgetCostTypes := budget.CostTypes
	if budgetCostTypes == nil {
		return nil, fmt.Errorf("empty CostTypes in budget: %v", budget)
	}

	costTypesMap := map[string]bool{
		"include_credit":             *budgetCostTypes.IncludeCredit,
		"include_other_subscription": *budgetCostTypes.IncludeOtherSubscription,
		"include_recurring":          *budgetCostTypes.IncludeRecurring,
		"include_refund":             *budgetCostTypes.IncludeRefund,
		"include_subscription":       *budgetCostTypes.IncludeSubscription,
		"include_support":            *budgetCostTypes.IncludeSupport,
		"include_tax":                *budgetCostTypes.IncludeTax,
		"include_upfront":            *budgetCostTypes.IncludeUpfront,
		"use_blended":                *budgetCostTypes.UseBlended,
	}
	budgetTimePeriod := budget.TimePeriod
	if budgetTimePeriod == nil {
		return nil, fmt.Errorf("empty TimePeriod in budget: %v", budget)
	}

	budgetTimePeriodStart := budgetTimePeriod.Start
	if budgetTimePeriodStart == nil {
		return nil, fmt.Errorf("empty TimePeriodStart in budget: %v", budget)
	}

	budgetTimePeriodEnd := budgetTimePeriod.End
	if budgetTimePeriodEnd == nil {
		return nil, fmt.Errorf("empty TimePeriodEnd in budget: %v", budget)
	}

	return &flattenedBudgetsBudget{
		name:            budget.BudgetName,
		budgetType:      budget.BudgetType,
		timeUnit:        budget.TimeUnit,
		costFilters:     budget.CostFilters,
		limitAmount:     budgetLimit.Amount,
		limitUnit:       budgetLimit.Unit,
		costTypes:       costTypesMap,
		timePeriodStart: budgetTimePeriodStart,
		timePeriodEnd:   budgetTimePeriodEnd,
	}, nil
}

func convertCostFiltersToStringMap(costFilters map[string][]*string) map[string]string {
	convertedCostFilters := make(map[string]string)
	for k, v := range costFilters {
		filterValues := make([]string, 0)
		for _, singleFilterValue := range v {
			filterValues = append(filterValues, *singleFilterValue)
		}

		convertedCostFilters[k] = strings.Join(filterValues, ",")
	}

	return convertedCostFilters
}

func expandBudgetsBudgetUnmarshal(d *schema.ResourceData) (*budgets.Budget, error) {
	var budgetName string
	if _, id, err := decodeBudgetsBudgetID(d.Id()); err == nil && id != "" {
		budgetName = id

	} else if v, ok := d.GetOk("name"); ok {
		budgetName = v.(string)

	} else if v, ok := d.GetOk("name_prefix"); ok {
		budgetName = resource.PrefixedUniqueId(v.(string))

	} else {
		budgetName = resource.UniqueId()
	}

	budgetType := d.Get("budget_type").(string)
	budgetLimitAmount := d.Get("limit_amount").(string)
	budgetLimitUnit := d.Get("limit_unit").(string)
	costTypes := expandBudgetsCostTypesUnmarshal(d.Get("cost_types").([]interface{}))
	budgetTimeUnit := d.Get("time_unit").(string)
	budgetCostFilters := make(map[string][]*string)
	for k, v := range d.Get("cost_filters").(map[string]interface{}) {
		filterValue := v.(string)
		budgetCostFilters[k] = append(budgetCostFilters[k], aws.String(filterValue))
	}

	budgetTimePeriodStart, err := time.Parse("2006-01-02_15:04", d.Get("time_period_start").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budgetTimePeriodEnd, err := time.Parse("2006-01-02_15:04", d.Get("time_period_end").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budget := &budgets.Budget{
		BudgetName: aws.String(budgetName),
		BudgetType: aws.String(budgetType),
		BudgetLimit: &budgets.Spend{
			Amount: aws.String(budgetLimitAmount),
			Unit:   aws.String(budgetLimitUnit),
		},
		CostTypes: costTypes,
		TimePeriod: &budgets.TimePeriod{
			End:   &budgetTimePeriodEnd,
			Start: &budgetTimePeriodStart,
		},
		TimeUnit:    aws.String(budgetTimeUnit),
		CostFilters: budgetCostFilters,
	}
	return budget, nil
}

func decodeBudgetsBudgetID(id string) (string, string, error) {
	parts := strings.Split(id, ":")
	if len(parts) != 2 {
		return "", "", fmt.Errorf("Unexpected format of ID (%q), expected AccountID:BudgetName", id)
	}
	return parts[0], parts[1], nil
}

func expandBudgetsCostTypesUnmarshal(budgetCostTypes []interface{}) *budgets.CostTypes {
	costTypes := &budgets.CostTypes{
		IncludeCredit:            aws.Bool(true),
		IncludeOtherSubscription: aws.Bool(true),
		IncludeRecurring:         aws.Bool(true),
		IncludeRefund:            aws.Bool(true),
		IncludeSubscription:      aws.Bool(true),
		IncludeSupport:           aws.Bool(true),
		IncludeTax:               aws.Bool(true),
		IncludeUpfront:           aws.Bool(true),
		UseBlended:               aws.Bool(false),
	}
	if len(budgetCostTypes) == 1 {
		costTypesMap := budgetCostTypes[0].(map[string]interface{})
		for k, v := range map[string]*bool{
			"include_credit":             costTypes.IncludeCredit,
			"include_other_subscription": costTypes.IncludeOtherSubscription,
			"include_recurring":          costTypes.IncludeRecurring,
			"include_refund":             costTypes.IncludeRefund,
			"include_subscription":       costTypes.IncludeSubscription,
			"include_support":            costTypes.IncludeSupport,
			"include_tax":                costTypes.IncludeTax,
			"include_upfront":            costTypes.IncludeUpfront,
			"use_blended":                costTypes.UseBlended,
		} {
			if val, ok := costTypesMap[k]; ok {
				*v = val.(bool)
			}
		}
	}

	return costTypes
}
<|endoftext|>


<commit_before>package mesh

import (
	"fmt"
	"net"
	"sync"
	"time"
)

type LocalPeer struct {
	sync.RWMutex
	*Peer
	router     *Router
	actionChan chan<- LocalPeerAction
}

type LocalPeerAction func()

func NewLocalPeer(name PeerName, nickName string, router *Router) *LocalPeer {
	actionChan := make(chan LocalPeerAction, ChannelSize)
	peer := &LocalPeer{
		Peer:       NewPeer(name, nickName, randomPeerUID(), 0, randomPeerShortID()),
		router:     router,
		actionChan: actionChan,
	}
	go peer.actorLoop(actionChan)
	return peer
}

func (peer *LocalPeer) Connections() ConnectionSet {
	connections := make(ConnectionSet)
	peer.RLock()
	defer peer.RUnlock()
	for _, conn := range peer.connections {
		connections[conn] = void
	}
	return connections
}

func (peer *LocalPeer) ConnectionTo(name PeerName) (Connection, bool) {
	peer.RLock()
	defer peer.RUnlock()
	conn, found := peer.connections[name]
	return conn, found // yes, you really can't inline that. FFS.
}

func (peer *LocalPeer) ConnectionsTo(names []PeerName) []Connection {
	conns := make([]Connection, 0, len(names))
	peer.RLock()
	defer peer.RUnlock()
	for _, name := range names {
		conn, found := peer.connections[name]
		// Again, !found could just be due to a race.
		if found {
			conns = append(conns, conn)
		}
	}
	return conns
}

func (peer *LocalPeer) CreateConnection(peerAddr string, acceptNewPeer bool) error {
	if err := peer.checkConnectionLimit(); err != nil {
		return err
	}
	tcpAddr, err := net.ResolveTCPAddr("tcp4", peerAddr)
	if err != nil {
		return err
	}
	tcpConn, err := net.DialTCP("tcp4", nil, tcpAddr)
	if err != nil {
		return err
	}
	connRemote := NewRemoteConnection(peer.Peer, nil, tcpConn.RemoteAddr().String(), true, false)
	StartLocalConnection(connRemote, tcpConn, peer.router, acceptNewPeer)
	return nil
}

// ACTOR client API

// Sync.
func (peer *LocalPeer) AddConnection(conn *LocalConnection) error {
	resultChan := make(chan error)
	peer.actionChan <- func() {
		resultChan <- peer.handleAddConnection(conn)
	}
	return <-resultChan
}

// Async.
func (peer *LocalPeer) ConnectionEstablished(conn *LocalConnection) {
	peer.actionChan <- func() {
		peer.handleConnectionEstablished(conn)
	}
}

// Sync.
func (peer *LocalPeer) DeleteConnection(conn *LocalConnection) {
	resultChan := make(chan interface{})
	peer.actionChan <- func() {
		peer.handleDeleteConnection(conn)
		resultChan <- nil
	}
	<-resultChan
}

// ACTOR server

func (peer *LocalPeer) actorLoop(actionChan <-chan LocalPeerAction) {
	gossipTimer := time.Tick(GossipInterval)
	for {
		select {
		case action := <-actionChan:
			action()
		case <-gossipTimer:
			peer.router.SendAllGossip()
		}
	}
}

func (peer *LocalPeer) handleAddConnection(conn Connection) error {
	if peer.Peer != conn.Local() {
		log.Fatal("Attempt made to add connection to peer where peer is not the source of connection")
	}
	if conn.Remote() == nil {
		log.Fatal("Attempt made to add connection to peer with unknown remote peer")
	}
	toName := conn.Remote().Name
	dupErr := fmt.Errorf("Multiple connections to %s added to %s", conn.Remote(), peer.String())
	// deliberately non symmetrical
	if dupConn, found := peer.connections[toName]; found {
		if dupConn == conn {
			return nil
		}
		switch conn.BreakTie(dupConn) {
		case TieBreakWon:
			dupConn.Shutdown(dupErr)
			peer.handleDeleteConnection(dupConn)
		case TieBreakLost:
			return dupErr
		case TieBreakTied:
			// oh good grief. Sod it, just kill both of them.
			dupConn.Shutdown(dupErr)
			peer.handleDeleteConnection(dupConn)
			return dupErr
		}
	}
	if err := peer.checkConnectionLimit(); err != nil {
		return err
	}
	_, isConnectedPeer := peer.router.Routes.Unicast(toName)
	peer.addConnection(conn)
	if isConnectedPeer {
		conn.Log("connection added")
	} else {
		conn.Log("connection added (new peer)")
		peer.router.SendAllGossipDown(conn)
	}

	peer.router.Routes.Recalculate()
	peer.broadcastPeerUpdate(conn.Remote())

	return nil
}

func (peer *LocalPeer) handleConnectionEstablished(conn Connection) {
	if peer.Peer != conn.Local() {
		log.Fatal("Peer informed of active connection where peer is not the source of connection")
	}
	if dupConn, found := peer.connections[conn.Remote().Name]; !found || conn != dupConn {
		conn.Shutdown(fmt.Errorf("Cannot set unknown connection active"))
		return
	}
	peer.connectionEstablished(conn)
	conn.Log("connection fully established")

	peer.router.Routes.Recalculate()
	peer.broadcastPeerUpdate()
}

func (peer *LocalPeer) handleDeleteConnection(conn Connection) {
	if peer.Peer != conn.Local() {
		log.Fatal("Attempt made to delete connection from peer where peer is not the source of connection")
	}
	if conn.Remote() == nil {
		log.Fatal("Attempt made to delete connection to peer with unknown remote peer")
	}
	toName := conn.Remote().Name
	if connFound, found := peer.connections[toName]; !found || connFound != conn {
		return
	}
	peer.deleteConnection(conn)
	conn.Log("connection deleted")
	// Must do garbage collection first to ensure we don't send out an
	// update with unreachable peers (can cause looping)
	peer.router.Peers.GarbageCollect()
	peer.router.Routes.Recalculate()
	peer.broadcastPeerUpdate()
}

// helpers

func (peer *LocalPeer) broadcastPeerUpdate(peers ...*Peer) {
	// Some tests run without a router.  This should be fixed so
	// that the relevant part of Router can be easily run in the
	// context of a test, but that will involve significant
	// reworking of tests.
	if peer.router != nil {
		peer.router.BroadcastTopologyUpdate(append(peers, peer.Peer))
	}
}

func (peer *LocalPeer) checkConnectionLimit() error {
	limit := peer.router.ConnLimit
	if 0 != limit && peer.connectionCount() >= limit {
		return fmt.Errorf("Connection limit reached (%v)", limit)
	}
	return nil
}

func (peer *LocalPeer) addConnection(conn Connection) {
	peer.Lock()
	defer peer.Unlock()
	peer.connections[conn.Remote().Name] = conn
	peer.Version++
}

func (peer *LocalPeer) deleteConnection(conn Connection) {
	peer.Lock()
	defer peer.Unlock()
	delete(peer.connections, conn.Remote().Name)
	peer.Version++
}

func (peer *LocalPeer) connectionEstablished(conn Connection) {
	peer.Lock()
	defer peer.Unlock()
	peer.Version++
}

func (peer *LocalPeer) connectionCount() int {
	peer.RLock()
	defer peer.RUnlock()
	return len(peer.connections)
}

func (peer *LocalPeer) setShortID(shortID PeerShortID) {
	peer.Lock()
	defer peer.Unlock()
	peer.ShortID = shortID
	peer.Version++
}

func (peer *LocalPeer) setVersionBeyond(version uint64) bool {
	peer.Lock()
	defer peer.Unlock()
	if version >= peer.Version {
		peer.Version = version + 1
		return true
	}
	return false
}
<commit_msg>don't broadcast-gossip other peer when adding a connection<commit_after>package mesh

import (
	"fmt"
	"net"
	"sync"
	"time"
)

type LocalPeer struct {
	sync.RWMutex
	*Peer
	router     *Router
	actionChan chan<- LocalPeerAction
}

type LocalPeerAction func()

func NewLocalPeer(name PeerName, nickName string, router *Router) *LocalPeer {
	actionChan := make(chan LocalPeerAction, ChannelSize)
	peer := &LocalPeer{
		Peer:       NewPeer(name, nickName, randomPeerUID(), 0, randomPeerShortID()),
		router:     router,
		actionChan: actionChan,
	}
	go peer.actorLoop(actionChan)
	return peer
}

func (peer *LocalPeer) Connections() ConnectionSet {
	connections := make(ConnectionSet)
	peer.RLock()
	defer peer.RUnlock()
	for _, conn := range peer.connections {
		connections[conn] = void
	}
	return connections
}

func (peer *LocalPeer) ConnectionTo(name PeerName) (Connection, bool) {
	peer.RLock()
	defer peer.RUnlock()
	conn, found := peer.connections[name]
	return conn, found // yes, you really can't inline that. FFS.
}

func (peer *LocalPeer) ConnectionsTo(names []PeerName) []Connection {
	conns := make([]Connection, 0, len(names))
	peer.RLock()
	defer peer.RUnlock()
	for _, name := range names {
		conn, found := peer.connections[name]
		// Again, !found could just be due to a race.
		if found {
			conns = append(conns, conn)
		}
	}
	return conns
}

func (peer *LocalPeer) CreateConnection(peerAddr string, acceptNewPeer bool) error {
	if err := peer.checkConnectionLimit(); err != nil {
		return err
	}
	tcpAddr, err := net.ResolveTCPAddr("tcp4", peerAddr)
	if err != nil {
		return err
	}
	tcpConn, err := net.DialTCP("tcp4", nil, tcpAddr)
	if err != nil {
		return err
	}
	connRemote := NewRemoteConnection(peer.Peer, nil, tcpConn.RemoteAddr().String(), true, false)
	StartLocalConnection(connRemote, tcpConn, peer.router, acceptNewPeer)
	return nil
}

// ACTOR client API

// Sync.
func (peer *LocalPeer) AddConnection(conn *LocalConnection) error {
	resultChan := make(chan error)
	peer.actionChan <- func() {
		resultChan <- peer.handleAddConnection(conn)
	}
	return <-resultChan
}

// Async.
func (peer *LocalPeer) ConnectionEstablished(conn *LocalConnection) {
	peer.actionChan <- func() {
		peer.handleConnectionEstablished(conn)
	}
}

// Sync.
func (peer *LocalPeer) DeleteConnection(conn *LocalConnection) {
	resultChan := make(chan interface{})
	peer.actionChan <- func() {
		peer.handleDeleteConnection(conn)
		resultChan <- nil
	}
	<-resultChan
}

// ACTOR server

func (peer *LocalPeer) actorLoop(actionChan <-chan LocalPeerAction) {
	gossipTimer := time.Tick(GossipInterval)
	for {
		select {
		case action := <-actionChan:
			action()
		case <-gossipTimer:
			peer.router.SendAllGossip()
		}
	}
}

func (peer *LocalPeer) handleAddConnection(conn Connection) error {
	if peer.Peer != conn.Local() {
		log.Fatal("Attempt made to add connection to peer where peer is not the source of connection")
	}
	if conn.Remote() == nil {
		log.Fatal("Attempt made to add connection to peer with unknown remote peer")
	}
	toName := conn.Remote().Name
	dupErr := fmt.Errorf("Multiple connections to %s added to %s", conn.Remote(), peer.String())
	// deliberately non symmetrical
	if dupConn, found := peer.connections[toName]; found {
		if dupConn == conn {
			return nil
		}
		switch conn.BreakTie(dupConn) {
		case TieBreakWon:
			dupConn.Shutdown(dupErr)
			peer.handleDeleteConnection(dupConn)
		case TieBreakLost:
			return dupErr
		case TieBreakTied:
			// oh good grief. Sod it, just kill both of them.
			dupConn.Shutdown(dupErr)
			peer.handleDeleteConnection(dupConn)
			return dupErr
		}
	}
	if err := peer.checkConnectionLimit(); err != nil {
		return err
	}
	_, isConnectedPeer := peer.router.Routes.Unicast(toName)
	peer.addConnection(conn)
	if isConnectedPeer {
		conn.Log("connection added")
	} else {
		conn.Log("connection added (new peer)")
		peer.router.SendAllGossipDown(conn)
	}

	peer.router.Routes.Recalculate()
	peer.broadcastPeerUpdate()

	return nil
}

func (peer *LocalPeer) handleConnectionEstablished(conn Connection) {
	if peer.Peer != conn.Local() {
		log.Fatal("Peer informed of active connection where peer is not the source of connection")
	}
	if dupConn, found := peer.connections[conn.Remote().Name]; !found || conn != dupConn {
		conn.Shutdown(fmt.Errorf("Cannot set unknown connection active"))
		return
	}
	peer.connectionEstablished(conn)
	conn.Log("connection fully established")

	peer.router.Routes.Recalculate()
	peer.broadcastPeerUpdate()
}

func (peer *LocalPeer) handleDeleteConnection(conn Connection) {
	if peer.Peer != conn.Local() {
		log.Fatal("Attempt made to delete connection from peer where peer is not the source of connection")
	}
	if conn.Remote() == nil {
		log.Fatal("Attempt made to delete connection to peer with unknown remote peer")
	}
	toName := conn.Remote().Name
	if connFound, found := peer.connections[toName]; !found || connFound != conn {
		return
	}
	peer.deleteConnection(conn)
	conn.Log("connection deleted")
	// Must do garbage collection first to ensure we don't send out an
	// update with unreachable peers (can cause looping)
	peer.router.Peers.GarbageCollect()
	peer.router.Routes.Recalculate()
	peer.broadcastPeerUpdate()
}

// helpers

func (peer *LocalPeer) broadcastPeerUpdate() {
	// Some tests run without a router.  This should be fixed so
	// that the relevant part of Router can be easily run in the
	// context of a test, but that will involve significant
	// reworking of tests.
	if peer.router != nil {
		peer.router.BroadcastTopologyUpdate([]*Peer{peer.Peer})
	}
}

func (peer *LocalPeer) checkConnectionLimit() error {
	limit := peer.router.ConnLimit
	if 0 != limit && peer.connectionCount() >= limit {
		return fmt.Errorf("Connection limit reached (%v)", limit)
	}
	return nil
}

func (peer *LocalPeer) addConnection(conn Connection) {
	peer.Lock()
	defer peer.Unlock()
	peer.connections[conn.Remote().Name] = conn
	peer.Version++
}

func (peer *LocalPeer) deleteConnection(conn Connection) {
	peer.Lock()
	defer peer.Unlock()
	delete(peer.connections, conn.Remote().Name)
	peer.Version++
}

func (peer *LocalPeer) connectionEstablished(conn Connection) {
	peer.Lock()
	defer peer.Unlock()
	peer.Version++
}

func (peer *LocalPeer) connectionCount() int {
	peer.RLock()
	defer peer.RUnlock()
	return len(peer.connections)
}

func (peer *LocalPeer) setShortID(shortID PeerShortID) {
	peer.Lock()
	defer peer.Unlock()
	peer.ShortID = shortID
	peer.Version++
}

func (peer *LocalPeer) setVersionBeyond(version uint64) bool {
	peer.Lock()
	defer peer.Unlock()
	if version >= peer.Version {
		peer.Version = version + 1
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package goetty

import (
	"net"
	"sync"
	"sync/atomic"
)

type syncClientMiddleware struct {
	BaseMiddleware

	bizDecoder                Decoder
	bizEncoder                Encoder
	writer                    func(IOSession, interface{}) error
	cached                    *simpleQueue
	localOffset, serverOffset uint64
	syncing                   bool
	maxReadTimeouts           int
	timeouts                  int
}

// NewSyncProtocolClientMiddleware return a middleware to process sync protocol
func NewSyncProtocolClientMiddleware(bizDecoder Decoder, bizEncoder Encoder, writer func(IOSession, interface{}) error, maxReadTimeouts int) Middleware {
	return &syncClientMiddleware{
		cached:          newSimpleQueue(),
		writer:          writer,
		bizDecoder:      bizDecoder,
		bizEncoder:      bizEncoder,
		maxReadTimeouts: maxReadTimeouts,
	}
}

func (sm *syncClientMiddleware) PreWrite(msg interface{}, conn IOSession) (bool, interface{}, error) {
	// The client side can only send notifySync,notifyRaw,notifyHB msg to the server,
	// wrap the raw biz msg to notifyRaw
	_, isSync := msg.(*notifySync)
	if !isSync {
		_, isSync = msg.(*notifyHB)
	}

	if !isSync {
		m := acquireNotifyRaw()
		err := sm.bizEncoder.Encode(msg, m.buf)
		if err != nil {
			return false, nil, err
		}

		return true, m, nil
	}

	return sm.BaseMiddleware.PreWrite(msg, conn)
}

func (sm *syncClientMiddleware) PreRead(conn IOSession) (bool, interface{}, error) {
	// If there is any biz msg in the queue, returned it and cancel read option
	if sm.cached.len() > 0 {
		return false, sm.cached.pop(), nil
	}

	return true, nil, nil
}

func (sm *syncClientMiddleware) PostRead(msg interface{}, conn IOSession) (bool, interface{}, error) {
	sm.timeouts = 0

	// if read notify msg from server, sync msg with server,
	// the client read option will block until sync complete and read the raw biz msg
	if nt, ok := msg.(*notify); ok {
		if nt.offset == 0 || nt.offset > sm.getLocalOffset() {
			sm.resetServerOffset(nt.offset)
			releaseNotify(nt)

			err := sm.sync(conn)
			if err != nil {
				return false, nil, err
			}
		}

		return false, nil, nil
	} else if rsp, ok := msg.(*notifySyncRsp); ok {
		sm.syncing = false
		sm.resetLocalOffset(rsp.offset)
		for i := byte(0); i < rsp.count; i++ {
			c, rawMsg, err := sm.bizDecoder.Decode(rsp.buf)
			if err != nil {
				return false, nil, err
			}
			if !c {
				panic("data cann't be enough")
			}

			sm.cached.push(rawMsg)
		}
		releaseNotifySyncRsp(rsp)

		if sm.getLocalOffset() < sm.getServerOffset() {
			err := sm.sync(conn)
			if err != nil {
				return false, nil, err
			}
		}

		return false, nil, nil
	}

	return true, msg, nil
}

func (sm *syncClientMiddleware) Closed(conn IOSession) {
	sm.cached = nil
	sm.localOffset = 0
	sm.serverOffset = 0
	sm.syncing = false
}

func (sm *syncClientMiddleware) Connected(conn IOSession) {
	sm.cached = newSimpleQueue()
	sm.localOffset = 0
	sm.serverOffset = 0
	sm.syncing = false
}

func (sm *syncClientMiddleware) ReadError(err error, conn IOSession) error {
	if netErr, ok := err.(*net.OpError); ok &&
		netErr.Timeout() &&
		sm.timeouts < sm.maxReadTimeouts {
		sm.timeouts++
		return sm.writer(conn, &notifyHB{
			offset: sm.getLocalOffset(),
		})
	}

	return err
}

func (sm *syncClientMiddleware) getLocalOffset() uint64 {
	return atomic.LoadUint64(&sm.localOffset)
}

func (sm *syncClientMiddleware) getServerOffset() uint64 {
	return sm.serverOffset
}

func (sm *syncClientMiddleware) resetLocalOffset(offset uint64) {
	atomic.StoreUint64(&sm.localOffset, offset)
}

func (sm *syncClientMiddleware) resetServerOffset(offset uint64) {
	sm.serverOffset = offset
}

func (sm *syncClientMiddleware) sync(conn IOSession) error {
	if !sm.syncing {
		req := acquireNotifySync()
		req.offset = sm.getLocalOffset()
		err := sm.writer(conn, req)
		if err != nil {
			return err
		}
		sm.syncing = true
	}

	return nil
}

type syncServerMiddleware struct {
	sync.RWMutex
	BaseMiddleware

	bizDecoder     Decoder
	bizEncoder     Encoder
	writer         func(IOSession, interface{}) error
	offsetQueueMap map[interface{}]*OffsetQueue
}

// NewSyncProtocolServerMiddleware return a middleware to process sync protocol
func NewSyncProtocolServerMiddleware(bizDecoder Decoder, bizEncoder Encoder, writer func(IOSession, interface{}) error) Middleware {
	return &syncServerMiddleware{
		bizDecoder:     bizDecoder,
		bizEncoder:     bizEncoder,
		writer:         writer,
		offsetQueueMap: make(map[interface{}]*OffsetQueue),
	}
}

func (sm *syncServerMiddleware) Connected(conn IOSession) {
	sm.Lock()
	sm.offsetQueueMap[conn.ID()] = newOffsetQueue()
	sm.Unlock()
}

func (sm *syncServerMiddleware) Closed(conn IOSession) {
	sm.Lock()
	delete(sm.offsetQueueMap, conn.ID())
	sm.Unlock()
}

func (sm *syncServerMiddleware) PreWrite(msg interface{}, conn IOSession) (bool, interface{}, error) {
	if _, ok := msg.(*notify); ok {
		return sm.BaseMiddleware.PreWrite(msg, conn)
	} else if _, ok := msg.(*notifySyncRsp); ok {
		return sm.BaseMiddleware.PreWrite(msg, conn)
	}

	// add biz msg to the offset queue, and send notify msg to client
	sm.RLock()
	q := sm.offsetQueueMap[conn.ID()]
	sm.RUnlock()

	if q == nil {
		panic("offset queue cann't be nil")
	}

	m := acquireNotify()
	m.offset = q.Add(msg)
	return true, m, nil
}

func (sm *syncServerMiddleware) PostRead(msg interface{}, conn IOSession) (bool, interface{}, error) {
	if m, ok := msg.(*notifySync); ok {
		sm.RLock()
		q := sm.offsetQueueMap[conn.ID()]
		sm.RUnlock()

		if q == nil {
			panic("offset queue cann't be nil")
		}

		// send biz msgs to client
		items, maxOffset := q.Get(m.offset)
		releaseNotifySync(m)

		rsp := acquireNotifySyncRsp()
		rsp.offset = maxOffset
		rsp.count = byte(len(items))
		for _, item := range items {
			err := sm.bizEncoder.Encode(item, rsp.buf)
			if err != nil {
				return false, nil, err
			}
		}
		err := sm.writer(conn, rsp)
		if err != nil {
			return false, nil, err
		}

		return false, nil, nil
	} else if m, ok := msg.(*notifyRaw); ok {
		c, biz, err := sm.bizDecoder.Decode(m.buf)
		releaseNotifyRaw(m)
		if err != nil {
			return false, nil, err
		}
		if !c {
			panic("bug, cann't missing data")
		}

		return true, biz, nil
	} else if m, ok := msg.(*notifyHB); ok {
		sm.RLock()
		q := sm.offsetQueueMap[conn.ID()]
		sm.RUnlock()

		if q == nil {
			panic("offset queue cann't be nil")
		}

		max := q.GetMaxOffset()
		if m.offset < max {
			nt := acquireNotify()
			nt.offset = q.GetMaxOffset()
			err := sm.writer(conn, nt)
			if err != nil {
				return false, nil, err
			}
		}

		return false, nil, nil
	}

	return true, msg, nil
}
<commit_msg>fix: sync protocol not syncing if loss net packket<commit_after>package goetty

import (
	"net"
	"sync"
	"sync/atomic"
)

type syncClientMiddleware struct {
	BaseMiddleware

	bizDecoder                Decoder
	bizEncoder                Encoder
	writer                    func(IOSession, interface{}) error
	cached                    *simpleQueue
	localOffset, serverOffset uint64
	syncing                   bool
	maxReadTimeouts           int
	timeouts                  int
}

// NewSyncProtocolClientMiddleware return a middleware to process sync protocol
func NewSyncProtocolClientMiddleware(bizDecoder Decoder, bizEncoder Encoder, writer func(IOSession, interface{}) error, maxReadTimeouts int) Middleware {
	return &syncClientMiddleware{
		cached:          newSimpleQueue(),
		writer:          writer,
		bizDecoder:      bizDecoder,
		bizEncoder:      bizEncoder,
		maxReadTimeouts: maxReadTimeouts,
	}
}

func (sm *syncClientMiddleware) PreWrite(msg interface{}, conn IOSession) (bool, interface{}, error) {
	// The client side can only send notifySync,notifyRaw,notifyHB msg to the server,
	// wrap the raw biz msg to notifyRaw
	_, isSync := msg.(*notifySync)
	if !isSync {
		_, isSync = msg.(*notifyHB)
	}

	if !isSync {
		m := acquireNotifyRaw()
		err := sm.bizEncoder.Encode(msg, m.buf)
		if err != nil {
			return false, nil, err
		}

		return true, m, nil
	}

	return sm.BaseMiddleware.PreWrite(msg, conn)
}

func (sm *syncClientMiddleware) PreRead(conn IOSession) (bool, interface{}, error) {
	// If there is any biz msg in the queue, returned it and cancel read option
	if sm.cached.len() > 0 {
		return false, sm.cached.pop(), nil
	}

	return true, nil, nil
}

func (sm *syncClientMiddleware) PostRead(msg interface{}, conn IOSession) (bool, interface{}, error) {
	sm.timeouts = 0

	// if read notify msg from server, sync msg with server,
	// the client read option will block until sync complete and read the raw biz msg
	if nt, ok := msg.(*notify); ok {
		if nt.offset == 0 || nt.offset > sm.getLocalOffset() {
			sm.resetServerOffset(nt.offset)
			releaseNotify(nt)

			err := sm.sync(conn)
			if err != nil {
				return false, nil, err
			}
		}

		return false, nil, nil
	} else if rsp, ok := msg.(*notifySyncRsp); ok {
		sm.syncing = false
		sm.resetLocalOffset(rsp.offset)
		for i := byte(0); i < rsp.count; i++ {
			c, rawMsg, err := sm.bizDecoder.Decode(rsp.buf)
			if err != nil {
				return false, nil, err
			}
			if !c {
				panic("data cann't be enough")
			}

			sm.cached.push(rawMsg)
		}
		releaseNotifySyncRsp(rsp)

		if sm.getLocalOffset() < sm.getServerOffset() {
			err := sm.sync(conn)
			if err != nil {
				return false, nil, err
			}
		}

		return false, nil, nil
	}

	return true, msg, nil
}

func (sm *syncClientMiddleware) Closed(conn IOSession) {
	sm.cached = nil
	sm.localOffset = 0
	sm.serverOffset = 0
	sm.syncing = false
	sm.timeouts = 0
}

func (sm *syncClientMiddleware) Connected(conn IOSession) {
	sm.cached = newSimpleQueue()
	sm.localOffset = 0
	sm.serverOffset = 0
	sm.syncing = false
	sm.timeouts = 0
}

func (sm *syncClientMiddleware) ReadError(err error, conn IOSession) error {
	if netErr, ok := err.(*net.OpError); ok &&
		netErr.Timeout() &&
		sm.timeouts < sm.maxReadTimeouts {
		sm.timeouts++
		sm.syncing = false
		return sm.writer(conn, &notifyHB{
			offset: sm.getLocalOffset(),
		})
	}

	return err
}

func (sm *syncClientMiddleware) getLocalOffset() uint64 {
	return atomic.LoadUint64(&sm.localOffset)
}

func (sm *syncClientMiddleware) getServerOffset() uint64 {
	return sm.serverOffset
}

func (sm *syncClientMiddleware) resetLocalOffset(offset uint64) {
	atomic.StoreUint64(&sm.localOffset, offset)
}

func (sm *syncClientMiddleware) resetServerOffset(offset uint64) {
	sm.serverOffset = offset
}

func (sm *syncClientMiddleware) sync(conn IOSession) error {
	if !sm.syncing {
		req := acquireNotifySync()
		req.offset = sm.getLocalOffset()
		err := sm.writer(conn, req)
		if err != nil {
			return err
		}
		sm.syncing = true
	}

	return nil
}

type syncServerMiddleware struct {
	sync.RWMutex
	BaseMiddleware

	bizDecoder     Decoder
	bizEncoder     Encoder
	writer         func(IOSession, interface{}) error
	offsetQueueMap map[interface{}]*OffsetQueue
}

// NewSyncProtocolServerMiddleware return a middleware to process sync protocol
func NewSyncProtocolServerMiddleware(bizDecoder Decoder, bizEncoder Encoder, writer func(IOSession, interface{}) error) Middleware {
	return &syncServerMiddleware{
		bizDecoder:     bizDecoder,
		bizEncoder:     bizEncoder,
		writer:         writer,
		offsetQueueMap: make(map[interface{}]*OffsetQueue),
	}
}

func (sm *syncServerMiddleware) Connected(conn IOSession) {
	sm.Lock()
	sm.offsetQueueMap[conn.ID()] = newOffsetQueue()
	sm.Unlock()
}

func (sm *syncServerMiddleware) Closed(conn IOSession) {
	sm.Lock()
	delete(sm.offsetQueueMap, conn.ID())
	sm.Unlock()
}

func (sm *syncServerMiddleware) PreWrite(msg interface{}, conn IOSession) (bool, interface{}, error) {
	if _, ok := msg.(*notify); ok {
		return sm.BaseMiddleware.PreWrite(msg, conn)
	} else if _, ok := msg.(*notifySyncRsp); ok {
		return sm.BaseMiddleware.PreWrite(msg, conn)
	}

	// add biz msg to the offset queue, and send notify msg to client
	sm.RLock()
	q := sm.offsetQueueMap[conn.ID()]
	sm.RUnlock()

	if q == nil {
		panic("offset queue cann't be nil")
	}

	m := acquireNotify()
	m.offset = q.Add(msg)
	return true, m, nil
}

func (sm *syncServerMiddleware) PostRead(msg interface{}, conn IOSession) (bool, interface{}, error) {
	if m, ok := msg.(*notifySync); ok {
		sm.RLock()
		q := sm.offsetQueueMap[conn.ID()]
		sm.RUnlock()

		if q == nil {
			panic("offset queue cann't be nil")
		}

		// send biz msgs to client
		items, maxOffset := q.Get(m.offset)
		releaseNotifySync(m)

		rsp := acquireNotifySyncRsp()
		rsp.offset = maxOffset
		rsp.count = byte(len(items))
		for _, item := range items {
			err := sm.bizEncoder.Encode(item, rsp.buf)
			if err != nil {
				return false, nil, err
			}
		}
		err := sm.writer(conn, rsp)
		if err != nil {
			return false, nil, err
		}

		return false, nil, nil
	} else if m, ok := msg.(*notifyRaw); ok {
		c, biz, err := sm.bizDecoder.Decode(m.buf)
		releaseNotifyRaw(m)
		if err != nil {
			return false, nil, err
		}
		if !c {
			panic("bug, cann't missing data")
		}

		return true, biz, nil
	} else if m, ok := msg.(*notifyHB); ok {
		sm.RLock()
		q := sm.offsetQueueMap[conn.ID()]
		sm.RUnlock()

		if q == nil {
			panic("offset queue cann't be nil")
		}

		max := q.GetMaxOffset()
		if m.offset < max {
			nt := acquireNotify()
			nt.offset = q.GetMaxOffset()
			err := sm.writer(conn, nt)
			if err != nil {
				return false, nil, err
			}
		}

		return false, nil, nil
	}

	return true, msg, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/fsouza/go-dockerclient"
)

type DockerDynoDriver struct {
	d     *Docker
	state DynoState

	cmd       *exec.Cmd
	container *docker.Container
	waiting   chan error
}

func NewDockerDynoDriver() *DockerDynoDriver {
	return &DockerDynoDriver{}
}

func (dd *DockerDynoDriver) State() DynoState {
	return dd.state
}

func (dd *DockerDynoDriver) Start(b *Bundle) error {
	if dd.d == nil {
		dd.d = &Docker{}
		if err := dd.d.Connect(); err != nil {
			dd.d = nil
			return err
		}
	}

	si, err := dd.d.StackStat("cedar-14")
	if err != nil {
		return err
	}

	log.Printf("StackImage %+v", si)
	imageName, err := dd.d.BuildSlugImage(si, b)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Built image successfully")

	// Fill environment vector from Heroku configuration.
	env := make([]string, 0)
	for k, v := range b.config {
		env = append(env, k+"="+v)
	}

	dd.container, err = dd.d.c.CreateContainer(docker.CreateContainerOptions{
		Name: fmt.Sprintf("%v-%v", imageName, int32(time.Now().Unix())),
		Config: &docker.Config{
			Cmd:   b.argv,
			Env:   env,
			Image: imageName,
		},
	})
	if err != nil {
		log.Fatal(err)
	}

	err = dd.d.c.StartContainer(dd.container.ID, &docker.HostConfig{})
	if err != nil {
		log.Fatal(err)
	}

	dd.state = Started

	return nil
}

func (dd *DockerDynoDriver) Stop() error {
	// If we could never start the process, don't worry about stopping it. May
	// occur in cases like if Docker was down.
	if dd.cmd == nil {
		return nil
	}

	p := dd.cmd.Process

	group, err := os.FindProcess(-1 * p.Pid)
	if err != nil {
		return err
	}

	// Begin graceful shutdown via SIGTERM.
	group.Signal(syscall.SIGTERM)

	for {
		select {
		case <-time.After(10 * time.Second):
			log.Println("sigkill", group)
			group.Signal(syscall.SIGKILL)
		case err := <-dd.waiting:
			log.Println("waited", group)
			dd.state = Stopped
			return err
		}
		log.Println("spin", group)
		time.Sleep(1)
	}
}
<commit_msg>Support stopping in the Docker driver<commit_after>package main

import (
	"fmt"
	"log"
	"os/exec"
	"time"

	"github.com/fsouza/go-dockerclient"
)

type DockerDynoDriver struct {
	d     *Docker
	state DynoState

	cmd       *exec.Cmd
	container *docker.Container
	waiting   chan error
}

func NewDockerDynoDriver() *DockerDynoDriver {
	return &DockerDynoDriver{}
}

func (dd *DockerDynoDriver) State() DynoState {
	return dd.state
}

func (dd *DockerDynoDriver) Start(b *Bundle) error {
	if dd.d == nil {
		dd.d = &Docker{}
		if err := dd.d.Connect(); err != nil {
			dd.d = nil
			return err
		}
	}

	si, err := dd.d.StackStat("cedar-14")
	if err != nil {
		return err
	}

	log.Printf("StackImage %+v", si)
	imageName, err := dd.d.BuildSlugImage(si, b)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Built image successfully")

	// Fill environment vector from Heroku configuration.
	env := make([]string, 0)
	for k, v := range b.config {
		env = append(env, k+"="+v)
	}

	dd.container, err = dd.d.c.CreateContainer(docker.CreateContainerOptions{
		Name: fmt.Sprintf("%v-%v", imageName, int32(time.Now().Unix())),
		Config: &docker.Config{
			Cmd:   b.argv,
			Env:   env,
			Image: imageName,
		},
	})
	if err != nil {
		log.Fatal(err)
	}

	err = dd.d.c.StartContainer(dd.container.ID, &docker.HostConfig{})
	if err != nil {
		return err
	}

	dd.state = Started

	return nil
}

func (dd *DockerDynoDriver) Stop() error {
	err := dd.d.c.StopContainer(dd.container.ID, 10)
	dd.state = Stopped
	return err
}
<|endoftext|>