<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubeinformers "k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	corev1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/clock"

	"github.com/jetstack/cert-manager/cmd/controller/app/options"
	"github.com/jetstack/cert-manager/pkg/acme/accounts"
	clientset "github.com/jetstack/cert-manager/pkg/client/clientset/versioned"
	intscheme "github.com/jetstack/cert-manager/pkg/client/clientset/versioned/scheme"
	informers "github.com/jetstack/cert-manager/pkg/client/informers/externalversions"
	"github.com/jetstack/cert-manager/pkg/controller"
	"github.com/jetstack/cert-manager/pkg/controller/clusterissuers"
	dnsutil "github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/jetstack/cert-manager/pkg/metrics"
	"github.com/jetstack/cert-manager/pkg/util"
)

const controllerAgentName = "cert-manager"

func Run(opts *options.ControllerOptions, stopCh <-chan struct{}) {
	rootCtx := util.ContextWithStopCh(context.Background(), stopCh)
	rootCtx = logf.NewContext(rootCtx, nil, "controller")
	log := logf.FromContext(rootCtx)

	ctx, kubeCfg, err := buildControllerContext(rootCtx, stopCh, opts)

	if err != nil {
		log.Error(err, "error building controller context", "options", opts)
		os.Exit(1)
	}

	metricsServer, err := ctx.Metrics.Start(opts.MetricsListenAddress)
	if err != nil {
		log.Error(err, "failed to listen on prometheus address", "address", opts.MetricsListenAddress)
		os.Exit(1)
	}

	var wg sync.WaitGroup
	run := func(_ context.Context) {
		for n, fn := range controller.Known() {
			log := log.WithValues("controller", n)

			// only run a controller if it's been enabled
			if !util.Contains(opts.EnabledControllers, n) {
				log.V(logf.InfoLevel).Info("not starting controller as it's disabled")
				continue
			}

			// don't run clusterissuers controller if scoped to a single namespace
			if ctx.Namespace != "" && n == clusterissuers.ControllerName {
				log.V(logf.InfoLevel).Info("not starting controller as cert-manager has been scoped to a single namespace")
				continue
			}

			wg.Add(1)
			iface, err := fn(ctx)
			if err != nil {
				log.Error(err, "error starting controller")
				os.Exit(1)
			}
			go func(n string, fn controller.Interface) {
				defer wg.Done()
				log.V(logf.InfoLevel).Info("starting controller")

				workers := 5
				err := fn.Run(workers, stopCh)

				if err != nil {
					log.Error(err, "error starting controller")
					os.Exit(1)
				}
			}(n, iface)
		}

		log.V(logf.DebugLevel).Info("starting shared informer factories")
		ctx.SharedInformerFactory.Start(stopCh)
		ctx.KubeSharedInformerFactory.Start(stopCh)
		wg.Wait()
		log.V(logf.InfoLevel).Info("control loops exited")
		ctx.Metrics.Shutdown(metricsServer)
		os.Exit(0)
	}

	if !opts.LeaderElect {
		run(context.TODO())
		return
	}

	log.V(logf.InfoLevel).Info("starting leader election")
	leaderElectionClient, err := kubernetes.NewForConfig(rest.AddUserAgent(kubeCfg, "leader-election"))
	if err != nil {
		log.Error(err, "error creating leader election client")
		os.Exit(1)
	}

	startLeaderElection(rootCtx, opts, leaderElectionClient, ctx.Recorder, run)
}

func buildControllerContext(ctx context.Context, stopCh <-chan struct{}, opts *options.ControllerOptions) (*controller.Context, *rest.Config, error) {
	log := logf.FromContext(ctx, "build-context")
	// Load the users Kubernetes config
	kubeCfg, err := clientcmd.BuildConfigFromFlags(opts.APIServerHost, opts.Kubeconfig)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating rest config: %s", err.Error())
	}

	// Add User-Agent to client
	kubeCfg = rest.AddUserAgent(kubeCfg, util.CertManagerUserAgent)

	// Create a cert-manager api client
	intcl, err := clientset.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating internal group client: %s", err.Error())
	}

	// Create a Kubernetes api client
	cl, err := kubernetes.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating kubernetes client: %s", err.Error())
	}

	nameservers := opts.DNS01RecursiveNameservers
	if len(nameservers) == 0 {
		nameservers = dnsutil.RecursiveNameservers
	}
	log.V(logf.InfoLevel).WithValues("nameservers", nameservers).Info("configured acme dns01 nameservers")

	HTTP01SolverResourceRequestCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestCPU: %s", err.Error())
	}

	HTTP01SolverResourceRequestMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestMemory: %s", err.Error())
	}

	HTTP01SolverResourceLimitsCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsCPU: %s", err.Error())
	}

	HTTP01SolverResourceLimitsMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsMemory: %s", err.Error())
	}

	// Create event broadcaster
	// Add cert-manager types to the default Kubernetes Scheme so Events can be
	// logged properly
	intscheme.AddToScheme(scheme.Scheme)
	log.V(logf.DebugLevel).Info("creating event broadcaster")
	eventBroadcaster := record.NewBroadcaster()
	eventBroadcaster.StartLogging(log.V(logf.DebugLevel).Info)
	eventBroadcaster.StartRecordingToSink(&corev1.EventSinkImpl{Interface: cl.CoreV1().Events("")})
	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, v1.EventSource{Component: controllerAgentName})

	sharedInformerFactory := informers.NewSharedInformerFactoryWithOptions(intcl, time.Second*30, informers.WithNamespace(opts.Namespace))
	kubeSharedInformerFactory := kubeinformers.NewSharedInformerFactoryWithOptions(cl, time.Second*30, kubeinformers.WithNamespace(opts.Namespace))

	acmeAccountRegistry := accounts.NewDefaultRegistry()

	return &controller.Context{
		RootContext:               ctx,
		StopCh:                    stopCh,
		RESTConfig:                kubeCfg,
		Client:                    cl,
		CMClient:                  intcl,
		Recorder:                  recorder,
		KubeSharedInformerFactory: kubeSharedInformerFactory,
		SharedInformerFactory:     sharedInformerFactory,
		Namespace:                 opts.Namespace,
		Clock:                     clock.RealClock{},
		Metrics:                   metrics.New(log),
		ACMEOptions: controller.ACMEOptions{
			HTTP01SolverImage:                 opts.ACMEHTTP01SolverImage,
			HTTP01SolverResourceRequestCPU:    HTTP01SolverResourceRequestCPU,
			HTTP01SolverResourceRequestMemory: HTTP01SolverResourceRequestMemory,
			HTTP01SolverResourceLimitsCPU:     HTTP01SolverResourceLimitsCPU,
			HTTP01SolverResourceLimitsMemory:  HTTP01SolverResourceLimitsMemory,
			DNS01CheckAuthoritative:           !opts.DNS01RecursiveNameserversOnly,
			DNS01Nameservers:                  nameservers,
			AccountRegistry:                   acmeAccountRegistry,
		},
		IssuerOptions: controller.IssuerOptions{
			ClusterIssuerAmbientCredentials: opts.ClusterIssuerAmbientCredentials,
			IssuerAmbientCredentials:        opts.IssuerAmbientCredentials,
			ClusterResourceNamespace:        opts.ClusterResourceNamespace,
			RenewBeforeExpiryDuration:       opts.RenewBeforeExpiryDuration,
		},
		IngressShimOptions: controller.IngressShimOptions{
			DefaultIssuerName:                 opts.DefaultIssuerName,
			DefaultIssuerKind:                 opts.DefaultIssuerKind,
			DefaultIssuerGroup:                opts.DefaultIssuerGroup,
			DefaultAutoCertificateAnnotations: opts.DefaultAutoCertificateAnnotations,
		},
		CertificateOptions: controller.CertificateOptions{
			EnableOwnerRef: opts.EnableCertificateOwnerRef,
		},
		SchedulerOptions: controller.SchedulerOptions{
			MaxConcurrentChallenges: opts.MaxConcurrentChallenges,
		},
	}, kubeCfg, nil
}

func startLeaderElection(ctx context.Context, opts *options.ControllerOptions, leaderElectionClient kubernetes.Interface, recorder record.EventRecorder, run func(context.Context)) {
	log := logf.FromContext(ctx, "leader-election")

	// Identity used to distinguish between multiple controller manager instances
	id, err := os.Hostname()
	if err != nil {
		log.Error(err, "error getting hostname")
		os.Exit(1)
	}

	// Lock required for leader election
	rl := resourcelock.ConfigMapLock{
		ConfigMapMeta: metav1.ObjectMeta{
			Namespace: opts.LeaderElectionNamespace,
			Name:      "cert-manager-controller",
		},
		Client: leaderElectionClient.CoreV1(),
		LockConfig: resourcelock.ResourceLockConfig{
			Identity:      id + "-external-cert-manager-controller",
			EventRecorder: recorder,
		},
	}

	// Try and become the leader and start controller manager loops
	leaderelection.RunOrDie(ctx, leaderelection.LeaderElectionConfig{
		Lock:          &rl,
		LeaseDuration: opts.LeaderElectionLeaseDuration,
		RenewDeadline: opts.LeaderElectionRenewDeadline,
		RetryPeriod:   opts.LeaderElectionRetryPeriod,
		Callbacks: leaderelection.LeaderCallbacks{
			OnStartedLeading: run,
			OnStoppedLeading: func() {
				log.V(logf.ErrorLevel).Info("leader election lost")
				os.Exit(1)
			},
		},
	})
}
<commit_msg>Rename import to make way for corev1 and v1 api imports later<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubeinformers "k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	clientv1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/clock"

	"github.com/jetstack/cert-manager/cmd/controller/app/options"
	"github.com/jetstack/cert-manager/pkg/acme/accounts"
	clientset "github.com/jetstack/cert-manager/pkg/client/clientset/versioned"
	intscheme "github.com/jetstack/cert-manager/pkg/client/clientset/versioned/scheme"
	informers "github.com/jetstack/cert-manager/pkg/client/informers/externalversions"
	"github.com/jetstack/cert-manager/pkg/controller"
	"github.com/jetstack/cert-manager/pkg/controller/clusterissuers"
	dnsutil "github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/jetstack/cert-manager/pkg/metrics"
	"github.com/jetstack/cert-manager/pkg/util"
)

const controllerAgentName = "cert-manager"

func Run(opts *options.ControllerOptions, stopCh <-chan struct{}) {
	rootCtx := util.ContextWithStopCh(context.Background(), stopCh)
	rootCtx = logf.NewContext(rootCtx, nil, "controller")
	log := logf.FromContext(rootCtx)

	ctx, kubeCfg, err := buildControllerContext(rootCtx, stopCh, opts)

	if err != nil {
		log.Error(err, "error building controller context", "options", opts)
		os.Exit(1)
	}

	metricsServer, err := ctx.Metrics.Start(opts.MetricsListenAddress)
	if err != nil {
		log.Error(err, "failed to listen on prometheus address", "address", opts.MetricsListenAddress)
		os.Exit(1)
	}

	var wg sync.WaitGroup
	run := func(_ context.Context) {
		for n, fn := range controller.Known() {
			log := log.WithValues("controller", n)

			// only run a controller if it's been enabled
			if !util.Contains(opts.EnabledControllers, n) {
				log.V(logf.InfoLevel).Info("not starting controller as it's disabled")
				continue
			}

			// don't run clusterissuers controller if scoped to a single namespace
			if ctx.Namespace != "" && n == clusterissuers.ControllerName {
				log.V(logf.InfoLevel).Info("not starting controller as cert-manager has been scoped to a single namespace")
				continue
			}

			wg.Add(1)
			iface, err := fn(ctx)
			if err != nil {
				log.Error(err, "error starting controller")
				os.Exit(1)
			}
			go func(n string, fn controller.Interface) {
				defer wg.Done()
				log.V(logf.InfoLevel).Info("starting controller")

				workers := 5
				err := fn.Run(workers, stopCh)

				if err != nil {
					log.Error(err, "error starting controller")
					os.Exit(1)
				}
			}(n, iface)
		}

		log.V(logf.DebugLevel).Info("starting shared informer factories")
		ctx.SharedInformerFactory.Start(stopCh)
		ctx.KubeSharedInformerFactory.Start(stopCh)
		wg.Wait()
		log.V(logf.InfoLevel).Info("control loops exited")
		ctx.Metrics.Shutdown(metricsServer)
		os.Exit(0)
	}

	if !opts.LeaderElect {
		run(context.TODO())
		return
	}

	log.V(logf.InfoLevel).Info("starting leader election")
	leaderElectionClient, err := kubernetes.NewForConfig(rest.AddUserAgent(kubeCfg, "leader-election"))
	if err != nil {
		log.Error(err, "error creating leader election client")
		os.Exit(1)
	}

	startLeaderElection(rootCtx, opts, leaderElectionClient, ctx.Recorder, run)
}

func buildControllerContext(ctx context.Context, stopCh <-chan struct{}, opts *options.ControllerOptions) (*controller.Context, *rest.Config, error) {
	log := logf.FromContext(ctx, "build-context")
	// Load the users Kubernetes config
	kubeCfg, err := clientcmd.BuildConfigFromFlags(opts.APIServerHost, opts.Kubeconfig)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating rest config: %s", err.Error())
	}

	// Add User-Agent to client
	kubeCfg = rest.AddUserAgent(kubeCfg, util.CertManagerUserAgent)

	// Create a cert-manager api client
	intcl, err := clientset.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating internal group client: %s", err.Error())
	}

	// Create a Kubernetes api client
	cl, err := kubernetes.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating kubernetes client: %s", err.Error())
	}

	nameservers := opts.DNS01RecursiveNameservers
	if len(nameservers) == 0 {
		nameservers = dnsutil.RecursiveNameservers
	}
	log.V(logf.InfoLevel).WithValues("nameservers", nameservers).Info("configured acme dns01 nameservers")

	HTTP01SolverResourceRequestCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestCPU: %s", err.Error())
	}

	HTTP01SolverResourceRequestMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestMemory: %s", err.Error())
	}

	HTTP01SolverResourceLimitsCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsCPU: %s", err.Error())
	}

	HTTP01SolverResourceLimitsMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsMemory: %s", err.Error())
	}

	// Create event broadcaster
	// Add cert-manager types to the default Kubernetes Scheme so Events can be
	// logged properly
	intscheme.AddToScheme(scheme.Scheme)
	log.V(logf.DebugLevel).Info("creating event broadcaster")
	eventBroadcaster := record.NewBroadcaster()
	eventBroadcaster.StartLogging(log.V(logf.DebugLevel).Info)
	eventBroadcaster.StartRecordingToSink(&clientv1.EventSinkImpl{Interface: cl.CoreV1().Events("")})
	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, v1.EventSource{Component: controllerAgentName})

	sharedInformerFactory := informers.NewSharedInformerFactoryWithOptions(intcl, time.Second*30, informers.WithNamespace(opts.Namespace))
	kubeSharedInformerFactory := kubeinformers.NewSharedInformerFactoryWithOptions(cl, time.Second*30, kubeinformers.WithNamespace(opts.Namespace))

	acmeAccountRegistry := accounts.NewDefaultRegistry()

	return &controller.Context{
		RootContext:               ctx,
		StopCh:                    stopCh,
		RESTConfig:                kubeCfg,
		Client:                    cl,
		CMClient:                  intcl,
		Recorder:                  recorder,
		KubeSharedInformerFactory: kubeSharedInformerFactory,
		SharedInformerFactory:     sharedInformerFactory,
		Namespace:                 opts.Namespace,
		Clock:                     clock.RealClock{},
		Metrics:                   metrics.New(log),
		ACMEOptions: controller.ACMEOptions{
			HTTP01SolverImage:                 opts.ACMEHTTP01SolverImage,
			HTTP01SolverResourceRequestCPU:    HTTP01SolverResourceRequestCPU,
			HTTP01SolverResourceRequestMemory: HTTP01SolverResourceRequestMemory,
			HTTP01SolverResourceLimitsCPU:     HTTP01SolverResourceLimitsCPU,
			HTTP01SolverResourceLimitsMemory:  HTTP01SolverResourceLimitsMemory,
			DNS01CheckAuthoritative:           !opts.DNS01RecursiveNameserversOnly,
			DNS01Nameservers:                  nameservers,
			AccountRegistry:                   acmeAccountRegistry,
		},
		IssuerOptions: controller.IssuerOptions{
			ClusterIssuerAmbientCredentials: opts.ClusterIssuerAmbientCredentials,
			IssuerAmbientCredentials:        opts.IssuerAmbientCredentials,
			ClusterResourceNamespace:        opts.ClusterResourceNamespace,
			RenewBeforeExpiryDuration:       opts.RenewBeforeExpiryDuration,
		},
		IngressShimOptions: controller.IngressShimOptions{
			DefaultIssuerName:                 opts.DefaultIssuerName,
			DefaultIssuerKind:                 opts.DefaultIssuerKind,
			DefaultIssuerGroup:                opts.DefaultIssuerGroup,
			DefaultAutoCertificateAnnotations: opts.DefaultAutoCertificateAnnotations,
		},
		CertificateOptions: controller.CertificateOptions{
			EnableOwnerRef: opts.EnableCertificateOwnerRef,
		},
		SchedulerOptions: controller.SchedulerOptions{
			MaxConcurrentChallenges: opts.MaxConcurrentChallenges,
		},
	}, kubeCfg, nil
}

func startLeaderElection(ctx context.Context, opts *options.ControllerOptions, leaderElectionClient kubernetes.Interface, recorder record.EventRecorder, run func(context.Context)) {
	log := logf.FromContext(ctx, "leader-election")

	// Identity used to distinguish between multiple controller manager instances
	id, err := os.Hostname()
	if err != nil {
		log.Error(err, "error getting hostname")
		os.Exit(1)
	}

	// Lock required for leader election
	rl := resourcelock.ConfigMapLock{
		ConfigMapMeta: metav1.ObjectMeta{
			Namespace: opts.LeaderElectionNamespace,
			Name:      "cert-manager-controller",
		},
		Client: leaderElectionClient.CoreV1(),
		LockConfig: resourcelock.ResourceLockConfig{
			Identity:      id + "-external-cert-manager-controller",
			EventRecorder: recorder,
		},
	}

	// Try and become the leader and start controller manager loops
	leaderelection.RunOrDie(ctx, leaderelection.LeaderElectionConfig{
		Lock:          &rl,
		LeaseDuration: opts.LeaderElectionLeaseDuration,
		RenewDeadline: opts.LeaderElectionRenewDeadline,
		RetryPeriod:   opts.LeaderElectionRetryPeriod,
		Callbacks: leaderelection.LeaderCallbacks{
			OnStartedLeading: run,
			OnStoppedLeading: func() {
				log.V(logf.ErrorLevel).Info("leader election lost")
				os.Exit(1)
			},
		},
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubeinformers "k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	clientv1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/clock"

	"github.com/jetstack/cert-manager/cmd/controller/app/options"
	"github.com/jetstack/cert-manager/pkg/acme/accounts"
	clientset "github.com/jetstack/cert-manager/pkg/client/clientset/versioned"
	intscheme "github.com/jetstack/cert-manager/pkg/client/clientset/versioned/scheme"
	informers "github.com/jetstack/cert-manager/pkg/client/informers/externalversions"
	"github.com/jetstack/cert-manager/pkg/controller"
	"github.com/jetstack/cert-manager/pkg/controller/clusterissuers"
	dnsutil "github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/jetstack/cert-manager/pkg/metrics"
	"github.com/jetstack/cert-manager/pkg/util"
)

const controllerAgentName = "cert-manager"

func Run(opts *options.ControllerOptions, stopCh <-chan struct{}) {
	rootCtx := util.ContextWithStopCh(context.Background(), stopCh)
	rootCtx = logf.NewContext(rootCtx, nil, "controller")
	log := logf.FromContext(rootCtx)

	ctx, kubeCfg, err := buildControllerContext(rootCtx, stopCh, opts)

	if err != nil {
		log.Error(err, "error building controller context", "options", opts)
		os.Exit(1)
	}

	metricsServer, err := ctx.Metrics.Start(opts.MetricsListenAddress)
	if err != nil {
		log.Error(err, "failed to listen on prometheus address", "address", opts.MetricsListenAddress)
		os.Exit(1)
	}

	var wg sync.WaitGroup
	run := func(_ context.Context) {
		for n, fn := range controller.Known() {
			log := log.WithValues("controller", n)

			// only run a controller if it's been enabled
			if !util.Contains(opts.EnabledControllers, n) {
				log.V(logf.InfoLevel).Info("not starting controller as it's disabled")
				continue
			}

			// don't run clusterissuers controller if scoped to a single namespace
			if ctx.Namespace != "" && n == clusterissuers.ControllerName {
				log.V(logf.InfoLevel).Info("not starting controller as cert-manager has been scoped to a single namespace")
				continue
			}

			wg.Add(1)
			iface, err := fn(ctx)
			if err != nil {
				log.Error(err, "error starting controller")
				os.Exit(1)
			}
			go func(n string, fn controller.Interface) {
				defer wg.Done()
				log.V(logf.InfoLevel).Info("starting controller")

				workers := 5
				err := fn.Run(workers, stopCh)

				if err != nil {
					log.Error(err, "error starting controller")
					os.Exit(1)
				}
			}(n, iface)
		}

		log.V(logf.DebugLevel).Info("starting shared informer factories")
		ctx.SharedInformerFactory.Start(stopCh)
		ctx.KubeSharedInformerFactory.Start(stopCh)
		wg.Wait()
		log.V(logf.InfoLevel).Info("control loops exited")
		ctx.Metrics.Shutdown(metricsServer)
		os.Exit(0)
	}

	if !opts.LeaderElect {
		run(context.TODO())
		return
	}

	log.V(logf.InfoLevel).Info("starting leader election")
	leaderElectionClient, err := kubernetes.NewForConfig(rest.AddUserAgent(kubeCfg, "leader-election"))
	if err != nil {
		log.Error(err, "error creating leader election client")
		os.Exit(1)
	}

	startLeaderElection(rootCtx, opts, leaderElectionClient, ctx.Recorder, run)
}

func buildControllerContext(ctx context.Context, stopCh <-chan struct{}, opts *options.ControllerOptions) (*controller.Context, *rest.Config, error) {
	log := logf.FromContext(ctx, "build-context")
	// Load the users Kubernetes config
	kubeCfg, err := clientcmd.BuildConfigFromFlags(opts.APIServerHost, opts.Kubeconfig)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating rest config: %s", err.Error())
	}

	kubeCfg.QPS = opts.KubernetesAPIQPS
	kubeCfg.Burst = opts.KubernetesAPIBurst

	// Add User-Agent to client
	kubeCfg = rest.AddUserAgent(kubeCfg, util.CertManagerUserAgent)

	// Create a cert-manager api client
	intcl, err := clientset.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating internal group client: %s", err.Error())
	}

	// Create a Kubernetes api client
	cl, err := kubernetes.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating kubernetes client: %s", err.Error())
	}

	nameservers := opts.DNS01RecursiveNameservers
	if len(nameservers) == 0 {
		nameservers = dnsutil.RecursiveNameservers
	}
	log.V(logf.InfoLevel).WithValues("nameservers", nameservers).Info("configured acme dns01 nameservers")

	HTTP01SolverResourceRequestCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestCPU: %s", err.Error())
	}

	HTTP01SolverResourceRequestMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestMemory: %s", err.Error())
	}

	HTTP01SolverResourceLimitsCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsCPU: %s", err.Error())
	}

	HTTP01SolverResourceLimitsMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsMemory: %s", err.Error())
	}

	// Create event broadcaster
	// Add cert-manager types to the default Kubernetes Scheme so Events can be
	// logged properly
	intscheme.AddToScheme(scheme.Scheme)
	log.V(logf.DebugLevel).Info("creating event broadcaster")
	eventBroadcaster := record.NewBroadcaster()
	eventBroadcaster.StartLogging(logf.WithInfof(log.V(logf.DebugLevel)).Infof)
	eventBroadcaster.StartRecordingToSink(&clientv1.EventSinkImpl{Interface: cl.CoreV1().Events("")})
	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, corev1.EventSource{Component: controllerAgentName})

	sharedInformerFactory := informers.NewSharedInformerFactoryWithOptions(intcl, time.Second*30, informers.WithNamespace(opts.Namespace))
	kubeSharedInformerFactory := kubeinformers.NewSharedInformerFactoryWithOptions(cl, time.Second*30, kubeinformers.WithNamespace(opts.Namespace))

	acmeAccountRegistry := accounts.NewDefaultRegistry()

	return &controller.Context{
		RootContext:               ctx,
		StopCh:                    stopCh,
		RESTConfig:                kubeCfg,
		Client:                    cl,
		CMClient:                  intcl,
		Recorder:                  recorder,
		KubeSharedInformerFactory: kubeSharedInformerFactory,
		SharedInformerFactory:     sharedInformerFactory,
		Namespace:                 opts.Namespace,
		Clock:                     clock.RealClock{},
		Metrics:                   metrics.New(log),
		ACMEOptions: controller.ACMEOptions{
			HTTP01SolverImage:                 opts.ACMEHTTP01SolverImage,
			HTTP01SolverResourceRequestCPU:    HTTP01SolverResourceRequestCPU,
			HTTP01SolverResourceRequestMemory: HTTP01SolverResourceRequestMemory,
			HTTP01SolverResourceLimitsCPU:     HTTP01SolverResourceLimitsCPU,
			HTTP01SolverResourceLimitsMemory:  HTTP01SolverResourceLimitsMemory,
			DNS01CheckAuthoritative:           !opts.DNS01RecursiveNameserversOnly,
			DNS01Nameservers:                  nameservers,
			AccountRegistry:                   acmeAccountRegistry,
			DNS01CheckRetryPeriod:             opts.DNS01CheckRetryPeriod,
		},
		IssuerOptions: controller.IssuerOptions{
			ClusterIssuerAmbientCredentials: opts.ClusterIssuerAmbientCredentials,
			IssuerAmbientCredentials:        opts.IssuerAmbientCredentials,
			ClusterResourceNamespace:        opts.ClusterResourceNamespace,
			RenewBeforeExpiryDuration:       opts.RenewBeforeExpiryDuration,
		},
		IngressShimOptions: controller.IngressShimOptions{
			DefaultIssuerName:                 opts.DefaultIssuerName,
			DefaultIssuerKind:                 opts.DefaultIssuerKind,
			DefaultIssuerGroup:                opts.DefaultIssuerGroup,
			DefaultAutoCertificateAnnotations: opts.DefaultAutoCertificateAnnotations,
		},
		CertificateOptions: controller.CertificateOptions{
			EnableOwnerRef: opts.EnableCertificateOwnerRef,
		},
		SchedulerOptions: controller.SchedulerOptions{
			MaxConcurrentChallenges: opts.MaxConcurrentChallenges,
		},
	}, kubeCfg, nil
}

func startLeaderElection(ctx context.Context, opts *options.ControllerOptions, leaderElectionClient kubernetes.Interface, recorder record.EventRecorder, run func(context.Context)) {
	log := logf.FromContext(ctx, "leader-election")

	// Identity used to distinguish between multiple controller manager instances
	id, err := os.Hostname()
	if err != nil {
		log.Error(err, "error getting hostname")
		os.Exit(1)
	}

	// Lock required for leader election
	rl := resourcelock.ConfigMapLock{
		ConfigMapMeta: metav1.ObjectMeta{
			Namespace: opts.LeaderElectionNamespace,
			Name:      "cert-manager-controller",
		},
		Client: leaderElectionClient.CoreV1(),
		LockConfig: resourcelock.ResourceLockConfig{
			Identity:      id + "-external-cert-manager-controller",
			EventRecorder: recorder,
		},
	}

	// Try and become the leader and start controller manager loops
	leaderelection.RunOrDie(ctx, leaderelection.LeaderElectionConfig{
		Lock:          &rl,
		LeaseDuration: opts.LeaderElectionLeaseDuration,
		RenewDeadline: opts.LeaderElectionRenewDeadline,
		RetryPeriod:   opts.LeaderElectionRetryPeriod,
		Callbacks: leaderelection.LeaderCallbacks{
			OnStartedLeading: run,
			OnStoppedLeading: func() {
				log.V(logf.ErrorLevel).Info("leader election lost")
				os.Exit(1)
			},
		},
	})
}
<commit_msg>Set the resync period to 10 hours<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubeinformers "k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	clientv1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/clock"

	"github.com/jetstack/cert-manager/cmd/controller/app/options"
	"github.com/jetstack/cert-manager/pkg/acme/accounts"
	clientset "github.com/jetstack/cert-manager/pkg/client/clientset/versioned"
	intscheme "github.com/jetstack/cert-manager/pkg/client/clientset/versioned/scheme"
	informers "github.com/jetstack/cert-manager/pkg/client/informers/externalversions"
	"github.com/jetstack/cert-manager/pkg/controller"
	"github.com/jetstack/cert-manager/pkg/controller/clusterissuers"
	dnsutil "github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/jetstack/cert-manager/pkg/metrics"
	"github.com/jetstack/cert-manager/pkg/util"
)

const controllerAgentName = "cert-manager"

func Run(opts *options.ControllerOptions, stopCh <-chan struct{}) {
	rootCtx := util.ContextWithStopCh(context.Background(), stopCh)
	rootCtx = logf.NewContext(rootCtx, nil, "controller")
	log := logf.FromContext(rootCtx)

	ctx, kubeCfg, err := buildControllerContext(rootCtx, stopCh, opts)

	if err != nil {
		log.Error(err, "error building controller context", "options", opts)
		os.Exit(1)
	}

	metricsServer, err := ctx.Metrics.Start(opts.MetricsListenAddress)
	if err != nil {
		log.Error(err, "failed to listen on prometheus address", "address", opts.MetricsListenAddress)
		os.Exit(1)
	}

	var wg sync.WaitGroup
	run := func(_ context.Context) {
		for n, fn := range controller.Known() {
			log := log.WithValues("controller", n)

			// only run a controller if it's been enabled
			if !util.Contains(opts.EnabledControllers, n) {
				log.V(logf.InfoLevel).Info("not starting controller as it's disabled")
				continue
			}

			// don't run clusterissuers controller if scoped to a single namespace
			if ctx.Namespace != "" && n == clusterissuers.ControllerName {
				log.V(logf.InfoLevel).Info("not starting controller as cert-manager has been scoped to a single namespace")
				continue
			}

			wg.Add(1)
			iface, err := fn(ctx)
			if err != nil {
				log.Error(err, "error starting controller")
				os.Exit(1)
			}
			go func(n string, fn controller.Interface) {
				defer wg.Done()
				log.V(logf.InfoLevel).Info("starting controller")

				workers := 5
				err := fn.Run(workers, stopCh)

				if err != nil {
					log.Error(err, "error starting controller")
					os.Exit(1)
				}
			}(n, iface)
		}

		log.V(logf.DebugLevel).Info("starting shared informer factories")
		ctx.SharedInformerFactory.Start(stopCh)
		ctx.KubeSharedInformerFactory.Start(stopCh)
		wg.Wait()
		log.V(logf.InfoLevel).Info("control loops exited")
		ctx.Metrics.Shutdown(metricsServer)
		os.Exit(0)
	}

	if !opts.LeaderElect {
		run(context.TODO())
		return
	}

	log.V(logf.InfoLevel).Info("starting leader election")
	leaderElectionClient, err := kubernetes.NewForConfig(rest.AddUserAgent(kubeCfg, "leader-election"))
	if err != nil {
		log.Error(err, "error creating leader election client")
		os.Exit(1)
	}

	startLeaderElection(rootCtx, opts, leaderElectionClient, ctx.Recorder, run)
}

func buildControllerContext(ctx context.Context, stopCh <-chan struct{}, opts *options.ControllerOptions) (*controller.Context, *rest.Config, error) {
	log := logf.FromContext(ctx, "build-context")
	// Load the users Kubernetes config
	kubeCfg, err := clientcmd.BuildConfigFromFlags(opts.APIServerHost, opts.Kubeconfig)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating rest config: %s", err.Error())
	}

	kubeCfg.QPS = opts.KubernetesAPIQPS
	kubeCfg.Burst = opts.KubernetesAPIBurst

	// Add User-Agent to client
	kubeCfg = rest.AddUserAgent(kubeCfg, util.CertManagerUserAgent)

	// Create a cert-manager api client
	intcl, err := clientset.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating internal group client: %s", err.Error())
	}

	// Create a Kubernetes api client
	cl, err := kubernetes.NewForConfig(kubeCfg)
	if err != nil {
		return nil, nil, fmt.Errorf("error creating kubernetes client: %s", err.Error())
	}

	nameservers := opts.DNS01RecursiveNameservers
	if len(nameservers) == 0 {
		nameservers = dnsutil.RecursiveNameservers
	}
	log.V(logf.InfoLevel).WithValues("nameservers", nameservers).Info("configured acme dns01 nameservers")

	HTTP01SolverResourceRequestCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestCPU: %s", err.Error())
	}

	HTTP01SolverResourceRequestMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceRequestMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceRequestMemory: %s", err.Error())
	}

	HTTP01SolverResourceLimitsCPU, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsCPU)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsCPU: %s", err.Error())
	}

	HTTP01SolverResourceLimitsMemory, err := resource.ParseQuantity(opts.ACMEHTTP01SolverResourceLimitsMemory)
	if err != nil {
		return nil, nil, fmt.Errorf("error parsing ACMEHTTP01SolverResourceLimitsMemory: %s", err.Error())
	}

	// Create event broadcaster
	// Add cert-manager types to the default Kubernetes Scheme so Events can be
	// logged properly
	intscheme.AddToScheme(scheme.Scheme)
	log.V(logf.DebugLevel).Info("creating event broadcaster")
	eventBroadcaster := record.NewBroadcaster()
	eventBroadcaster.StartLogging(logf.WithInfof(log.V(logf.DebugLevel)).Infof)
	eventBroadcaster.StartRecordingToSink(&clientv1.EventSinkImpl{Interface: cl.CoreV1().Events("")})
	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, corev1.EventSource{Component: controllerAgentName})

	sharedInformerFactory := informers.NewSharedInformerFactoryWithOptions(intcl, 10*time.Hour, informers.WithNamespace(opts.Namespace))
	kubeSharedInformerFactory := kubeinformers.NewSharedInformerFactoryWithOptions(cl, 10*time.Hour, kubeinformers.WithNamespace(opts.Namespace))

	acmeAccountRegistry := accounts.NewDefaultRegistry()

	return &controller.Context{
		RootContext:               ctx,
		StopCh:                    stopCh,
		RESTConfig:                kubeCfg,
		Client:                    cl,
		CMClient:                  intcl,
		Recorder:                  recorder,
		KubeSharedInformerFactory: kubeSharedInformerFactory,
		SharedInformerFactory:     sharedInformerFactory,
		Namespace:                 opts.Namespace,
		Clock:                     clock.RealClock{},
		Metrics:                   metrics.New(log),
		ACMEOptions: controller.ACMEOptions{
			HTTP01SolverImage:                 opts.ACMEHTTP01SolverImage,
			HTTP01SolverResourceRequestCPU:    HTTP01SolverResourceRequestCPU,
			HTTP01SolverResourceRequestMemory: HTTP01SolverResourceRequestMemory,
			HTTP01SolverResourceLimitsCPU:     HTTP01SolverResourceLimitsCPU,
			HTTP01SolverResourceLimitsMemory:  HTTP01SolverResourceLimitsMemory,
			DNS01CheckAuthoritative:           !opts.DNS01RecursiveNameserversOnly,
			DNS01Nameservers:                  nameservers,
			AccountRegistry:                   acmeAccountRegistry,
			DNS01CheckRetryPeriod:             opts.DNS01CheckRetryPeriod,
		},
		IssuerOptions: controller.IssuerOptions{
			ClusterIssuerAmbientCredentials: opts.ClusterIssuerAmbientCredentials,
			IssuerAmbientCredentials:        opts.IssuerAmbientCredentials,
			ClusterResourceNamespace:        opts.ClusterResourceNamespace,
			RenewBeforeExpiryDuration:       opts.RenewBeforeExpiryDuration,
		},
		IngressShimOptions: controller.IngressShimOptions{
			DefaultIssuerName:                 opts.DefaultIssuerName,
			DefaultIssuerKind:                 opts.DefaultIssuerKind,
			DefaultIssuerGroup:                opts.DefaultIssuerGroup,
			DefaultAutoCertificateAnnotations: opts.DefaultAutoCertificateAnnotations,
		},
		CertificateOptions: controller.CertificateOptions{
			EnableOwnerRef: opts.EnableCertificateOwnerRef,
		},
		SchedulerOptions: controller.SchedulerOptions{
			MaxConcurrentChallenges: opts.MaxConcurrentChallenges,
		},
	}, kubeCfg, nil
}

func startLeaderElection(ctx context.Context, opts *options.ControllerOptions, leaderElectionClient kubernetes.Interface, recorder record.EventRecorder, run func(context.Context)) {
	log := logf.FromContext(ctx, "leader-election")

	// Identity used to distinguish between multiple controller manager instances
	id, err := os.Hostname()
	if err != nil {
		log.Error(err, "error getting hostname")
		os.Exit(1)
	}

	// Lock required for leader election
	rl := resourcelock.ConfigMapLock{
		ConfigMapMeta: metav1.ObjectMeta{
			Namespace: opts.LeaderElectionNamespace,
			Name:      "cert-manager-controller",
		},
		Client: leaderElectionClient.CoreV1(),
		LockConfig: resourcelock.ResourceLockConfig{
			Identity:      id + "-external-cert-manager-controller",
			EventRecorder: recorder,
		},
	}

	// Try and become the leader and start controller manager loops
	leaderelection.RunOrDie(ctx, leaderelection.LeaderElectionConfig{
		Lock:          &rl,
		LeaseDuration: opts.LeaderElectionLeaseDuration,
		RenewDeadline: opts.LeaderElectionRenewDeadline,
		RetryPeriod:   opts.LeaderElectionRetryPeriod,
		Callbacks: leaderelection.LeaderCallbacks{
			OnStartedLeading: run,
			OnStoppedLeading: func() {
				log.V(logf.ErrorLevel).Info("leader election lost")
				os.Exit(1)
			},
		},
	})
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"os"
	"testing"

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

func testAccAwsOrganizationsAccount_basic(t *testing.T) {
	var account organizations.Account

	test_email, ok := os.LookupEnv("TEST_AWS_ORGANIZATION_ACCOUNT_EMAIL")

	if !ok {
		t.Skip("'TEST_AWS_ORGANIZATION_ACCOUNT_EMAIL' not set, skipping test.")
	}

	name := "my_new_account"
	email := test_email

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAwsOrganizationsAccountDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAwsOrganizationsAccountConfig(name, email),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsOrganizationsAccountExists("aws_organizations_account.test", &account),
				),
			},
			{
				ResourceName:      "aws_organizations_account.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func testAccCheckAwsOrganizationsAccountDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).organizationsconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_organizations_account" {
			continue
		}

		params := &organizations.DescribeAccountInput{
			AccountId: &rs.Primary.ID,
		}

		resp, err := conn.DescribeAccount(params)

		if err != nil {
			return err
		}

		if resp == nil && resp.Account != nil {
			return fmt.Errorf("Bad: Account still exists: %q", rs.Primary.ID)
		}
	}

	return nil

}

func testAccCheckAwsOrganizationsAccountExists(n string, a *organizations.Account) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).organizationsconn
		params := &organizations.DescribeAccountInput{
			AccountId: &rs.Primary.ID,
		}

		resp, err := conn.DescribeAccount(params)

		if err != nil {
			return err
		}

		if resp == nil || resp.Account == nil {
			return fmt.Errorf("Account %q does not exist", rs.Primary.ID)
		}

		a = resp.Account

		return nil
	}
}

func testAccAwsOrganizationsAccountConfig(name, email string) string {
	return fmt.Sprintf(`
resource "aws_organizations_account" "test" {
  name = "%s"
  email = "%s"
}
`, name, email)
}
<commit_msg>Randomize the name and email inputs for acceptance test account creation.<commit_after>package aws

import (
	"fmt"
	"os"
	"testing"

	"github.com/aws/aws-sdk-go/service/organizations"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func testAccAwsOrganizationsAccount_basic(t *testing.T) {
	var account organizations.Account

	orgsEmailDomain, ok := os.LookupEnv("TEST_AWS_ORGANIZATION_ACCOUNT_EMAIL_DOMAIN")

	if !ok {
		t.Skip("'TEST_AWS_ORGANIZATION_ACCOUNT_EMAIL_DOMAIN' not set, skipping test.")
	}

	rInt := acctest.RandInt()
	name := fmt.Sprintf("tf_acctest_%s", rInt)
	email := fmt.Sprintf("tf-acctest+%d@%s", rInt, orgsEmailDomain)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAwsOrganizationsAccountDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAwsOrganizationsAccountConfig(name, email),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsOrganizationsAccountExists("aws_organizations_account.test", &account),
				),
			},
			{
				ResourceName:      "aws_organizations_account.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func testAccCheckAwsOrganizationsAccountDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).organizationsconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_organizations_account" {
			continue
		}

		params := &organizations.DescribeAccountInput{
			AccountId: &rs.Primary.ID,
		}

		resp, err := conn.DescribeAccount(params)

		if err != nil {
			return err
		}

		if resp == nil && resp.Account != nil {
			return fmt.Errorf("Bad: Account still exists: %q", rs.Primary.ID)
		}
	}

	return nil

}

func testAccCheckAwsOrganizationsAccountExists(n string, a *organizations.Account) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).organizationsconn
		params := &organizations.DescribeAccountInput{
			AccountId: &rs.Primary.ID,
		}

		resp, err := conn.DescribeAccount(params)

		if err != nil {
			return err
		}

		if resp == nil || resp.Account == nil {
			return fmt.Errorf("Account %q does not exist", rs.Primary.ID)
		}

		a = resp.Account

		return nil
	}
}

func testAccAwsOrganizationsAccountConfig(name, email string) string {
	return fmt.Sprintf(`
resource "aws_organizations_account" "test" {
  name = "%s"
  email = "%s"
}
`, name, email)
}
<|endoftext|>


<commit_before>package jobsupervisor_test

import (
	"path/filepath"
	"runtime"

	boshdirs "github.com/cloudfoundry/bosh-agent/settings/directories"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"

	. "github.com/cloudfoundry/bosh-agent/jobsupervisor"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("WindowsJobSupervisor", func() {
	BeforeEach(func() {
		if runtime.GOOS != "windows" {
			Skip("Pending on non-Windows")
		}
	})

	var (
		runner        boshsys.CmdRunner
		fs            boshsys.FileSystem
		jobSupervisor JobSupervisor
	)

	const (
		processConfigContents = `{
"processes": [
	{
		"name": "say-hello",
		"executable": "powershell",
		"args": ["/C", "Start-Sleep 10"]
	}
]
}
`
	)

	BeforeEach(func() {
		logger := boshlog.NewLogger(boshlog.LevelNone)
		dirProvider := boshdirs.NewProvider("/var/vcap/")

		fs = boshsys.NewOsFileSystem(logger)
		runner = boshsys.NewExecCmdRunner(logger)
		jobSupervisor = NewWindowsJobSupervisor(runner, dirProvider, fs, logger)
	})

	Context("when started", func() {
		BeforeEach(func() {
			jobSupervisor.Start()
		})

		It("reports running", func() {
			Expect(jobSupervisor.Status()).To(Equal("running"))
		})
	})

	Context("when stopped", func() {
		BeforeEach(func() {
			jobSupervisor.Stop()
		})

		It("reports stopped", func() {
			Expect(jobSupervisor.Status()).To(Equal("stopped"))
		})
	})

	Context("when unmonitored", func() {
		BeforeEach(func() {
			jobSupervisor.Unmonitor()
		})

		It("reports unmonitored", func() {
			Expect(jobSupervisor.Status()).To(Equal("unmonitored"))
		})
	})

	Context("with service", func() {
		var (
			jobDir string
		)

		BeforeEach(func() {
			var err error
			jobDir, err = fs.TempDir("testWindowsJobSupervisor")
			processConfigPath := filepath.Join(jobDir, "monit")

			err = fs.WriteFileString(processConfigPath, processConfigContents)
			Expect(err).ToNot(HaveOccurred())

			err = jobSupervisor.AddJob("say-hello", 0, processConfigPath)
			Expect(err).ToNot(HaveOccurred())
		})

		AfterEach(func() {
			jobSupervisor.Stop()
			jobSupervisor.RemoveAllJobs()
			fs.RemoveAll(jobDir)
		})

		Describe("AddJob", func() {
			It("creates a service with vcap description", func() {
				stdout, _, _, err := runner.RunCommand("powershell", "/C", "get-service", "say-hello")
				Expect(err).ToNot(HaveOccurred())
				Expect(stdout).To(ContainSubstring("say-hello"))
				Expect(stdout).To(ContainSubstring("Stopped"))
			})
		})

		Describe("Start", func() {
			It("will start all the services", func() {
				err := jobSupervisor.Start()
				Expect(err).ToNot(HaveOccurred())

				stdout, _, _, err := runner.RunCommand("powershell", "/C", "get-service", "say-hello")
				Expect(err).ToNot(HaveOccurred())
				Expect(stdout).To(ContainSubstring("say-hello"))
				Expect(stdout).To(ContainSubstring("Running"))
			})
		})

		Describe("Status", func() {
			Context("when running", func() {
				It("reports that the job is 'Running'", func() {
					err := jobSupervisor.Start()
					Expect(err).ToNot(HaveOccurred())

					Expect(jobSupervisor.Status()).To(Equal("running"))
				})
			})

			Context("when stopped", func() {
				It("reports that the job is 'Stopped'", func() {
					err := jobSupervisor.Start()
					Expect(err).ToNot(HaveOccurred())

					err = jobSupervisor.Stop()
					Expect(err).ToNot(HaveOccurred())

					Expect(jobSupervisor.Status()).To(Equal("stopped"))
				})
			})
		})
	})
})
<commit_msg>Fix windows job supervisor tests on windows<commit_after>package jobsupervisor_test

import (
	"path/filepath"
	"runtime"

	boshdirs "github.com/cloudfoundry/bosh-agent/settings/directories"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"

	. "github.com/cloudfoundry/bosh-agent/jobsupervisor"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("WindowsJobSupervisor", func() {
	BeforeEach(func() {
		if runtime.GOOS != "windows" {
			Skip("Pending on non-Windows")
		}
	})

	var (
		runner        boshsys.CmdRunner
		fs            boshsys.FileSystem
		jobSupervisor JobSupervisor
		jobDir        string
	)

	const (
		processConfigContents = `{
"processes": [
	{
		"name": "say-hello",
		"executable": "powershell",
		"args": ["/C", "Start-Sleep 10"]
	}
]
}
`
	)

	BeforeEach(func() {
		logger := boshlog.NewLogger(boshlog.LevelNone)
		dirProvider := boshdirs.NewProvider("C:/var/vcap/")

		fs = boshsys.NewOsFileSystem(logger)
		runner = boshsys.NewExecCmdRunner(logger)
		jobSupervisor = NewWindowsJobSupervisor(runner, dirProvider, fs, logger)
		jobSupervisor.RemoveAllJobs()

		var err error
		jobDir, err = fs.TempDir("testWindowsJobSupervisor")
		processConfigPath := filepath.Join(jobDir, "monit")

		err = fs.WriteFileString(processConfigPath, processConfigContents)
		Expect(err).ToNot(HaveOccurred())

		err = jobSupervisor.AddJob("say-hello", 0, processConfigPath)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		jobSupervisor.Stop()
		jobSupervisor.RemoveAllJobs()
		fs.RemoveAll(jobDir)
	})

	Describe("AddJob", func() {
		It("creates a service with vcap description", func() {
			stdout, _, _, err := runner.RunCommand("powershell", "/C", "get-service", "say-hello")
			Expect(err).ToNot(HaveOccurred())
			Expect(stdout).To(ContainSubstring("say-hello"))
			Expect(stdout).To(ContainSubstring("Stopped"))
		})
	})

	Describe("Start", func() {
		It("will start all the services", func() {
			err := jobSupervisor.Start()
			Expect(err).ToNot(HaveOccurred())

			stdout, _, _, err := runner.RunCommand("powershell", "/C", "get-service", "say-hello")
			Expect(err).ToNot(HaveOccurred())
			Expect(stdout).To(ContainSubstring("say-hello"))
			Expect(stdout).To(ContainSubstring("Running"))
		})
	})

	Describe("Status", func() {
		Context("when running", func() {
			It("reports that the job is 'Running'", func() {
				err := jobSupervisor.Start()
				Expect(err).ToNot(HaveOccurred())

				Expect(jobSupervisor.Status()).To(Equal("running"))
			})
		})

		Context("when stopped", func() {
			It("reports that the job is 'Stopped'", func() {
				err := jobSupervisor.Start()
				Expect(err).ToNot(HaveOccurred())

				err = jobSupervisor.Stop()
				Expect(err).ToNot(HaveOccurred())

				Expect(jobSupervisor.Status()).To(Equal("stopped"))
			})
		})
	})
})
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"io"
	"math"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
)

// DataSource is a Root execution unit.
type DataSource struct {
	UID   UnitID
	SID   StreamID
	Coder *coder.Coder
	Out   Node

	source   DataManager
	count    int64
	splitPos int64
	start    time.Time

	mu sync.Mutex
}

func (n *DataSource) ID() UnitID {
	return n.UID
}

func (n *DataSource) Up(ctx context.Context) error {
	return nil
}

func (n *DataSource) StartBundle(ctx context.Context, id string, data DataContext) error {
	n.mu.Lock()
	n.source = data.Data
	n.start = time.Now()
	n.count = 0
	n.splitPos = math.MaxInt64
	n.mu.Unlock()
	return n.Out.StartBundle(ctx, id, data)
}

func (n *DataSource) Process(ctx context.Context) error {
	r, err := n.source.OpenRead(ctx, n.SID)
	if err != nil {
		return err
	}
	defer r.Close()

	c := coder.SkipW(n.Coder)
	wc := MakeWindowDecoder(n.Coder.Window)

	switch {
	case coder.IsCoGBK(c):
		ck := MakeElementDecoder(c.Components[0])
		cv := MakeElementDecoder(c.Components[1])

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}
			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			// Decode key

			key, err := ck.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			key.Timestamp = t
			key.Windows = ws

			// TODO(herohde) 4/30/2017: the State API will be handle re-iterations
			// and only "small" value streams would be inline. Presumably, that
			// would entail buffering the whole stream. We do that for now.

			var buf []FullValue

			size, err := coder.DecodeInt32(r)
			if err != nil {
				return errors.Wrap(err, "stream size decoding failed")
			}

			if size > -1 {
				// Single chunk stream.

				// log.Printf("Fixed size=%v", size)
				for i := int32(0); i < size; i++ {
					value, err := cv.Decode(r)
					if err != nil {
						return errors.Wrap(err, "stream value decode failed")
					}
					buf = append(buf, *value)
				}
			} else {
				// Multi-chunked stream.

				for {
					chunk, err := coder.DecodeVarUint64(r)
					if err != nil {
						return errors.Wrap(err, "stream chunk size decoding failed")
					}

					// log.Printf("Chunk size=%v", chunk)

					if chunk == 0 {
						break
					}

					for i := uint64(0); i < chunk; i++ {
						value, err := cv.Decode(r)
						if err != nil {
							return errors.Wrap(err, "stream value decode failed")
						}
						buf = append(buf, *value)
					}
				}
			}

			values := &FixedReStream{Buf: buf}
			if err := n.Out.ProcessElement(ctx, key, values); err != nil {
				return err
			}
		}

	default:
		ec := MakeElementDecoder(c)

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}

			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			elm, err := ec.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			elm.Timestamp = t
			elm.Windows = ws

			// log.Printf("READ: %v %v", elm.Key.Type(), elm.Key.Interface())

			if err := n.Out.ProcessElement(ctx, elm); err != nil {
				return err
			}
		}
	}
}

func (n *DataSource) FinishBundle(ctx context.Context) error {
	n.mu.Lock()
	defer n.mu.Unlock()
	log.Infof(ctx, "DataSource: %d elements in %d ns", n.count, time.Now().Sub(n.start))
	n.source = nil
	err := n.Out.FinishBundle(ctx)
	n.count = 0
	n.splitPos = math.MaxInt64
	return err
}

func (n *DataSource) Down(ctx context.Context) error {
	n.source = nil
	return nil
}

func (n *DataSource) String() string {
	return fmt.Sprintf("DataSource[%v] Coder:%v Out:%v", n.SID, n.Coder, n.Out.ID())
}

// IncrementCountAndCheckSplit increments DataSource.count by one and checks if
// the caller should abort further element processing, and finish the bundle.
// Returns true if the new value of count is greater than or equal to the split
// point, and false otherwise.
func (n *DataSource) IncrementCountAndCheckSplit(ctx context.Context) bool {
	b := false
	n.mu.Lock()
	n.count++
	if n.count >= n.splitPos {
		b = true
	}
	n.mu.Unlock()
	return b
}

// ProgressReportSnapshot captures the progress reading an input source.
type ProgressReportSnapshot struct {
	ID, Name string
	Count    int64
}

// Progress returns a snapshot of the source's progress.
func (n *DataSource) Progress() ProgressReportSnapshot {
	if n == nil {
		return ProgressReportSnapshot{}
	}
	n.mu.Lock()
	c := n.count
	n.mu.Unlock()
	return ProgressReportSnapshot{n.SID.Target.ID, n.SID.Target.Name, c}
}

// Split takes a sorted set of potential split points, selects and actuates
// split on an appropriate split point, and returns the selected split point
// if successful. Returns an error when unable to split.
func (n *DataSource) Split(splits []int64, frac float32) (int64, error) {
	if splits == nil {
		return 0, fmt.Errorf("failed to split: requested splits were empty")
	}
	if n == nil {
		return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource not initialized", splits)
	}
	n.mu.Lock()
	c := n.count
	// Find the smallest split index that we haven't yet processed, and set
	// the promised split position to this value.
	for _, s := range splits {
		if s >= c {
			n.splitPos = s
			fs := n.splitPos
			n.mu.Unlock()
			return fs, nil
		}
	}
	n.mu.Unlock()
	// If we can't find a suitable split point from the requested choices,
	// return an error.
	return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource at index: %v", splits, c)
}
<commit_msg>Make the Go SDK more robust with respect to requested split points.<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"io"
	"math"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
)

// DataSource is a Root execution unit.
type DataSource struct {
	UID   UnitID
	SID   StreamID
	Coder *coder.Coder
	Out   Node

	source   DataManager
	count    int64
	splitPos int64
	start    time.Time

	mu sync.Mutex
}

func (n *DataSource) ID() UnitID {
	return n.UID
}

func (n *DataSource) Up(ctx context.Context) error {
	return nil
}

func (n *DataSource) StartBundle(ctx context.Context, id string, data DataContext) error {
	n.mu.Lock()
	n.source = data.Data
	n.start = time.Now()
	n.count = 0
	n.splitPos = math.MaxInt64
	n.mu.Unlock()
	return n.Out.StartBundle(ctx, id, data)
}

func (n *DataSource) Process(ctx context.Context) error {
	r, err := n.source.OpenRead(ctx, n.SID)
	if err != nil {
		return err
	}
	defer r.Close()

	c := coder.SkipW(n.Coder)
	wc := MakeWindowDecoder(n.Coder.Window)

	switch {
	case coder.IsCoGBK(c):
		ck := MakeElementDecoder(c.Components[0])
		cv := MakeElementDecoder(c.Components[1])

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}
			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			// Decode key

			key, err := ck.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			key.Timestamp = t
			key.Windows = ws

			// TODO(herohde) 4/30/2017: the State API will be handle re-iterations
			// and only "small" value streams would be inline. Presumably, that
			// would entail buffering the whole stream. We do that for now.

			var buf []FullValue

			size, err := coder.DecodeInt32(r)
			if err != nil {
				return errors.Wrap(err, "stream size decoding failed")
			}

			if size > -1 {
				// Single chunk stream.

				// log.Printf("Fixed size=%v", size)
				for i := int32(0); i < size; i++ {
					value, err := cv.Decode(r)
					if err != nil {
						return errors.Wrap(err, "stream value decode failed")
					}
					buf = append(buf, *value)
				}
			} else {
				// Multi-chunked stream.

				for {
					chunk, err := coder.DecodeVarUint64(r)
					if err != nil {
						return errors.Wrap(err, "stream chunk size decoding failed")
					}

					// log.Printf("Chunk size=%v", chunk)

					if chunk == 0 {
						break
					}

					for i := uint64(0); i < chunk; i++ {
						value, err := cv.Decode(r)
						if err != nil {
							return errors.Wrap(err, "stream value decode failed")
						}
						buf = append(buf, *value)
					}
				}
			}

			values := &FixedReStream{Buf: buf}
			if err := n.Out.ProcessElement(ctx, key, values); err != nil {
				return err
			}
		}

	default:
		ec := MakeElementDecoder(c)

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}

			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			elm, err := ec.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			elm.Timestamp = t
			elm.Windows = ws

			// log.Printf("READ: %v %v", elm.Key.Type(), elm.Key.Interface())

			if err := n.Out.ProcessElement(ctx, elm); err != nil {
				return err
			}
		}
	}
}

func (n *DataSource) FinishBundle(ctx context.Context) error {
	n.mu.Lock()
	defer n.mu.Unlock()
	log.Infof(ctx, "DataSource: %d elements in %d ns", n.count, time.Now().Sub(n.start))
	n.source = nil
	err := n.Out.FinishBundle(ctx)
	n.count = 0
	n.splitPos = math.MaxInt64
	return err
}

func (n *DataSource) Down(ctx context.Context) error {
	n.source = nil
	return nil
}

func (n *DataSource) String() string {
	return fmt.Sprintf("DataSource[%v] Coder:%v Out:%v", n.SID, n.Coder, n.Out.ID())
}

// IncrementCountAndCheckSplit increments DataSource.count by one and checks if
// the caller should abort further element processing, and finish the bundle.
// Returns true if the new value of count is greater than or equal to the split
// point, and false otherwise.
func (n *DataSource) IncrementCountAndCheckSplit(ctx context.Context) bool {
	b := false
	n.mu.Lock()
	n.count++
	if n.count >= n.splitPos {
		b = true
	}
	n.mu.Unlock()
	return b
}

// ProgressReportSnapshot captures the progress reading an input source.
type ProgressReportSnapshot struct {
	ID, Name string
	Count    int64
}

// Progress returns a snapshot of the source's progress.
func (n *DataSource) Progress() ProgressReportSnapshot {
	if n == nil {
		return ProgressReportSnapshot{}
	}
	n.mu.Lock()
	c := n.count
	n.mu.Unlock()
	return ProgressReportSnapshot{n.SID.Target.ID, n.SID.Target.Name, c}
}

// Split takes a sorted set of potential split points, selects and actuates
// split on an appropriate split point, and returns the selected split point
// if successful. Returns an error when unable to split.
func (n *DataSource) Split(splits []int64, frac float32) (int64, error) {
	if splits == nil {
		return 0, fmt.Errorf("failed to split: requested splits were empty")
	}
	if n == nil {
		return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource not initialized", splits)
	}
	n.mu.Lock()
	c := n.count
	// Find the smallest split index that we haven't yet processed, and set
	// the promised split position to this value.
	for _, s := range splits {
		if s >= c && s < n.splitPos {
			n.splitPos = s
			fs := n.splitPos
			n.mu.Unlock()
			return fs, nil
		}
	}
	n.mu.Unlock()
	// If we can't find a suitable split point from the requested choices,
	// return an error.
	return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource at index: %v", splits, c)
}
<|endoftext|>


<commit_before>package pg_query

import "encoding/json"
import "fmt"
import "reflect"

type Node interface {
	Deparse() string
}

func UnmarshalNodeJSON(input json.RawMessage) Node {
	var nodeMap map[string]json.RawMessage
	json.Unmarshal(input, &nodeMap)

	for nodeType, jsonText := range nodeMap {
		switch nodeType {
		case "SELECT":
			var selectStmt SelectStmt
			json.Unmarshal(jsonText, &selectStmt)
			return selectStmt
		case "RESTARGET":
			var resTarget ResTarget
			json.Unmarshal(jsonText, &resTarget)
			return resTarget
		case "A_CONST":
			var aConst A_Const
			json.Unmarshal(jsonText, &aConst)
			return aConst
		}
	}

	return nil
}

func UnmarshalNodeFieldJSON(input []byte, node Node) (err error) {
	var fields map[string]json.RawMessage

	err = json.Unmarshal([]byte(input), &fields)
	if err != nil {
		return
	}

	s := reflect.ValueOf(node).Elem()
	typeOfT := s.Type()

	for i := 0; i < s.NumField(); i++ {
		f := s.Field(i)
		f2 := typeOfT.Field(i)
		fieldName := typeOfT.Field(i).Tag.Get("json")
		nodeType := reflect.TypeOf((*Node)(nil)).Elem()

		// Null is the default, so we just skip the type magic here
		if string(fields[fieldName]) == "null" {
			continue
		}

		switch f.Interface().(type) {
		case int:
			var value int64
			json.Unmarshal(fields[fieldName], &value)
			f.SetInt(value)
		case Value:
			// FIXME: We shouldn't special case like this
			var value Value
			json.Unmarshal(fields[fieldName], &value)
			f.Set(reflect.ValueOf(value))
		case *string:
			var value string
			json.Unmarshal(fields[fieldName], &value)
			f.Set(reflect.ValueOf(&value))
		case []Node:
			var list []json.RawMessage
			var nodes []Node
			json.Unmarshal(fields[fieldName], &list)
			for _, node := range list {
				nodes = append(nodes, UnmarshalNodeJSON(node))
			}
			f.Set(reflect.ValueOf(nodes))
		default:
			if f2.Type.Implements(nodeType) {
				node := UnmarshalNodeJSON(fields[fieldName])
				if node != nil {
					f.Set(reflect.ValueOf(node))
				}
			} else {
				err = fmt.Errorf("Invalid type %v", f.Type())
			}
		}
	}

	return
}
<commit_msg>Node unmarshal: Fix null detection.<commit_after>package pg_query

import "encoding/json"
import "fmt"
import "reflect"

type Node interface {
	Deparse() string
}

func UnmarshalNodeJSON(input json.RawMessage) Node {
	var nodeMap map[string]json.RawMessage
	json.Unmarshal(input, &nodeMap)

	for nodeType, jsonText := range nodeMap {
		switch nodeType {
		case "SELECT":
			var selectStmt SelectStmt
			json.Unmarshal(jsonText, &selectStmt)
			return selectStmt
		case "RESTARGET":
			var resTarget ResTarget
			json.Unmarshal(jsonText, &resTarget)
			return resTarget
		case "A_CONST":
			var aConst A_Const
			json.Unmarshal(jsonText, &aConst)
			return aConst
		}
	}

	return nil
}

func UnmarshalNodeFieldJSON(input []byte, node Node) (err error) {
	var fields map[string]json.RawMessage

	err = json.Unmarshal([]byte(input), &fields)
	if err != nil {
		return
	}

	s := reflect.ValueOf(node).Elem()
	typeOfT := s.Type()

	for i := 0; i < s.NumField(); i++ {
		f := s.Field(i)
		f2 := typeOfT.Field(i)
		fieldName := typeOfT.Field(i).Tag.Get("json")
		nodeType := reflect.TypeOf((*Node)(nil)).Elem()

		// Null is the default, so we just skip the type magic here
		if fields[fieldName] == nil {
			continue
		}

		switch f.Interface().(type) {
		case int:
			var value int64
			json.Unmarshal(fields[fieldName], &value)
			f.SetInt(value)
		case Value:
			// FIXME: We shouldn't special case like this
			var value Value
			json.Unmarshal(fields[fieldName], &value)
			f.Set(reflect.ValueOf(value))
		case *string:
			var value string
			json.Unmarshal(fields[fieldName], &value)
			f.Set(reflect.ValueOf(&value))
		case []Node:
			var list []json.RawMessage
			var nodes []Node
			json.Unmarshal(fields[fieldName], &list)
			for _, node := range list {
				nodes = append(nodes, UnmarshalNodeJSON(node))
			}
			f.Set(reflect.ValueOf(nodes))
		default:
			if f2.Type.Implements(nodeType) {
				node := UnmarshalNodeJSON(fields[fieldName])
				if node != nil {
					f.Set(reflect.ValueOf(node))
				}
			} else {
				err = fmt.Errorf("Invalid type %v", f.Type())
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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.

// Implements utils/archive-source buildpack.
// The archive-source buildpack archives user's source code.
package main

import (
	"os"
	"path/filepath"
	"strconv"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/env"
	gcp "github.com/GoogleCloudPlatform/buildpacks/pkg/gcpbuildpack"
)

const (
	archiveName = "source-code.tar.gz"
)

func main() {
	gcp.Main(detectFn, buildFn)
}

func detectFn(ctx *gcp.Context) error {
	return nil
}

func buildFn(ctx *gcp.Context) error {
	// Fail archiving source when users want to clear source from the final container.
	// TODO(https://github.com/buildpacks/lifecycle/issues/306): Move this logic to the detect phase when we can attribute failures to users.
	if cs, ok := os.LookupEnv(env.ClearSource); ok {
		c, err := strconv.ParseBool(cs)
		if err != nil {
			ctx.Warnf("Failed to parse %q: %v", env.ClearSource, err)
		} else if c {
			return gcp.UserErrorf("%s is not allowed in this environment", env.ClearSource)
		}
	}

	sl := ctx.Layer("src", gcp.LaunchLayer)

	sp := filepath.Join(sl.Path, archiveName)
	archiveSource(ctx, sp, ctx.ApplicationRoot())

	// Symlink the archive to /workspace/.googlebuild for a stable path.
	ctx.MkdirAll(".googlebuild", 0755)
	ctx.Symlink(sp, filepath.Join(ctx.ApplicationRoot(), ".googlebuild", archiveName))

	return nil
}

// archiveSource archives user's source code in a layer
func archiveSource(ctx *gcp.Context, fileName, dirName string) {
	ctx.Exec([]string{"tar",
		"--create", "--gzip", "--preserve-permissions",
		"--file=" + fileName,
		"--directory", dirName,
		"."}, gcp.WithUserTimingAttribution)
}
<commit_msg>The archive_source buildpack stores the archive location as a LABEL.<commit_after>// Copyright 2020 Google LLC
//
// 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.

// Implements utils/archive-source buildpack.
// The archive-source buildpack archives user's source code.
package main

import (
	"os"
	"path/filepath"
	"strconv"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/env"
	gcp "github.com/GoogleCloudPlatform/buildpacks/pkg/gcpbuildpack"
)

const (
	archiveName = "source-code.tar.gz"
)

func main() {
	gcp.Main(detectFn, buildFn)
}

func detectFn(ctx *gcp.Context) error {
	return nil
}

func buildFn(ctx *gcp.Context) error {
	// Fail archiving source when users want to clear source from the final container.
	// TODO(https://github.com/buildpacks/lifecycle/issues/306): Move this logic to the detect phase when we can attribute failures to users.
	if cs, ok := os.LookupEnv(env.ClearSource); ok {
		c, err := strconv.ParseBool(cs)
		if err != nil {
			ctx.Warnf("Failed to parse %q: %v", env.ClearSource, err)
		} else if c {
			return gcp.UserErrorf("%s is not allowed in this environment", env.ClearSource)
		}
	}

	sl := ctx.Layer("src", gcp.LaunchLayer)

	sp := filepath.Join(sl.Path, archiveName)
	archiveSource(ctx, sp, ctx.ApplicationRoot())

	// Symlink the archive to /workspace/.googlebuild for a stable path; add LABEL to container.
	ctx.MkdirAll(".googlebuild", 0755)
	stable := filepath.Join(ctx.ApplicationRoot(), ".googlebuild", archiveName)
	ctx.Symlink(sp, stable)
	ctx.AddLabel("source-archive", stable)

	return nil
}

// archiveSource archives user's source code in a layer
func archiveSource(ctx *gcp.Context, fileName, dirName string) {
	ctx.Exec([]string{"tar",
		"--create", "--gzip", "--preserve-permissions",
		"--file=" + fileName,
		"--directory", dirName,
		"."}, gcp.WithUserTimingAttribution)
}
<|endoftext|>


<commit_before>/*
 * Copyright 2013 Nan Deng
 *
 * 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 msgcenter

import (
	"errors"
	"fmt"
	"github.com/uniqush/uniqush-conn/config"
	"github.com/uniqush/uniqush-conn/proto/server"
	"github.com/uniqush/uniqush-conn/rpc"
	"io"
	"strings"
)

type serviceCenter struct {
	config     *config.ServiceConfig
	fwdChan    chan<- *rpc.ForwardRequest
	subReqChan chan *rpc.SubscribeRequest
	conns      connMap
	peer       rpc.UniqushConnPeer
}

func (self *serviceCenter) Stop() {
	self.conns.CloseAll()
	close(self.subReqChan)
}

func (self *serviceCenter) serveConn(conn server.Conn) {
	var reason error
	defer func() {
		self.config.OnLogout(conn, reason)
		self.conns.DelConn(conn)
		conn.Close()
	}()
	for {
		msg, err := conn.ReceiveMessage()
		if err != nil {
			if err != io.EOF {
				self.config.OnError(conn, err)
				reason = err
			}
			return
		}
		if msg != nil {
			self.config.OnMessage(conn, msg)
		}
	}
}

func (self *serviceCenter) NewConn(conn server.Conn) {
	if conn == nil {
		//self.config.OnError(conn, fmt.Errorf("Nil conn")
		return
	}
	usr := conn.Username()
	if len(usr) == 0 || strings.Contains(usr, ":") || strings.Contains(usr, "\n") {
		self.config.OnError(conn, fmt.Errorf("invalid username"))
		conn.Close()
		return
	}
	conn.SetMessageCache(self.config.Cache())
	conn.SetForwardRequestChannel(self.fwdChan)
	conn.SetSubscribeRequestChan(self.subReqChan)
	err := self.conns.AddConn(conn)
	if err != nil {
		self.config.OnError(conn, err)
		conn.Close()
		return
	}

	go self.serveConn(conn)
	return
}

func (self *serviceCenter) Send(req *rpc.SendRequest) *rpc.Result {
	ret := new(rpc.Result)

	if req == nil {
		ret.SetError(fmt.Errorf("invalid request"))
		return ret
	}
	if req.Message == nil || req.Message.IsEmpty() {
		ret.SetError(fmt.Errorf("invalid request: empty message"))
		return ret
	}
	if len(req.Receivers) == 0 {
		ret.SetError(fmt.Errorf("invalid request: no receiver"))
		return ret
	}

	shouldPush := !req.DontPush
	shouldCache := !req.DontCache
	shouldPropagate := !req.DontPropagate
	receivers := req.Receivers

	for _, recver := range receivers {
		mid := req.Id
		msg := req.Message

		if shouldCache {
			mc := &rpc.MessageContainer{
				Sender:        "",
				SenderService: "",
				Message:       msg,
			}
			var err error
			mid, err = self.config.CacheMessage(recver, mc, req.TTL)
			if err != nil {
				ret.SetError(err)
				return ret
			}
		}

		if len(mid) == 0 {
			ret.SetError(fmt.Errorf("undefined message Id"))
			return ret
		}

		n := 0

		conns := self.conns.GetConn(recver)
		conns.Traverse(func(conn server.Conn) error {
			err := conn.SendMessage(msg, mid, nil, !req.NeverDigest)
			ret.Append(conn, err)
			if err != nil {
				conn.Close()
				// We won't delete this connection here.
				// Instead, we close it and let the reader
				// goroutine detect the error and close it.
			} else {
				n++
			}
			// Instead of returning an error,
			// we wourld rather let the Traverse() move forward.
			return nil
		})

		// Don't propagate this request to other instances in the cluster.
		req.DontPropagate = true
		// Don't push the message. We will push it on this node.
		req.DontPush = true
		// Don't cache the message. We have already cached it.
		req.DontCache = true
		req.Id = mid
		req.Receivers = []string{recver}

		if shouldPropagate {
			r := self.peer.Send(req)
			n += r.NrSuccess()
			ret.Join(r)
		}

		if n == 0 && shouldPush {
			self.config.Push(recver, "", "", req.PushInfo, mid, msg.Size())
		}
	}
	return ret
}

func (self *serviceCenter) Forward(req *rpc.ForwardRequest) *rpc.Result {
	ret := new(rpc.Result)

	if req == nil {
		ret.SetError(fmt.Errorf("invalid request"))
		return ret
	}
	if req.Message == nil || req.Message.IsEmpty() {
		ret.SetError(fmt.Errorf("invalid request: empty message"))
		return ret
	}
	if len(req.Receivers) == 0 {
		ret.SetError(fmt.Errorf("invalid request: no receiver"))
		return ret
	}

	mid := req.Id
	msg := req.Message
	mc := &req.MessageContainer

	var pushInfo map[string]string
	var shouldForward bool
	shouldPush := !req.DontPush
	shouldCache := !req.DontCache
	shouldPropagate := !req.DontPropagate

	if !req.DontAsk {
		// We need to ask for permission to forward this message.
		// This means the forward request is generated directly from a user,
		// not from a uniqush-conn node in a cluster.

		mc.Id = ""
		shouldForward, shouldPush, pushInfo = self.config.ShouldForward(req)

		if !shouldForward {
			return nil
		}
	}

	receivers := req.Receivers

	for _, recver := range receivers {
		if shouldCache {
			var err error
			mid, err = self.config.CacheMessage(recver, mc, req.TTL)
			if err != nil {
				ret.SetError(err)
				return ret
			}
		}

		if len(mid) == 0 {
			ret.SetError(fmt.Errorf("undefined message Id"))
			return ret
		}

		n := 0

		conns := self.conns.GetConn(recver)
		conns.Traverse(func(conn server.Conn) error {
			err := conn.ForwardMessage(req.Sender, req.SenderService, msg, mid, !req.NeverDigest)
			ret.Append(conn, err)
			if err != nil {
				conn.Close()
				// We won't delete this connection here.
				// Instead, we close it and let the reader
				// goroutine detect the error and close it.
			} else {
				n++
			}
			// Instead of returning an error,
			// we wourld rather let the Traverse() move forward.
			return nil
		})

		// forward the message if possible,
		// Don't propagate this request to other instances in the cluster.
		req.DontPropagate = true
		// Don't ask the permission to forward (we have already got the permission)
		req.DontAsk = true
		// And don't push the message. We will push it on this node.
		req.DontPush = true
		// Dont' cache it
		req.DontCache = true
		req.Id = mid
		req.Receivers = []string{recver}

		if shouldPropagate {
			r := self.peer.Forward(req)
			n += r.NrSuccess()
			ret.Join(r)
		}

		if n == 0 && shouldPush {
			self.config.Push(recver, req.SenderService, req.Sender, pushInfo, mid, msg.Size())
		}
	}
	return ret
}

func (self *serviceCenter) Redirect(req *rpc.RedirectRequest) *rpc.Result {
	conns := self.conns.GetConn(req.Receiver)
	var sc server.Conn
	result := new(rpc.Result)
	conns.Traverse(func(conn server.Conn) error {
		if len(req.ConnId) == 0 || conn.UniqId() == req.ConnId {
			sc = conn
			result.Append(sc, nil)
			return errors.New("done")
		}
		return nil
	})
	if sc != nil {
		// self.conns.DelConn(sc)
		sc.Redirect(req.CandidateSersers...)
		sc.Close()
		return result
	}

	if req.DontPropagate {
		return result
	}
	req.DontPropagate = true
	return self.peer.Redirect(req)
}

func (self *serviceCenter) CheckUserStatus(req *rpc.UserStatusQuery) *rpc.Result {
	conns := self.conns.GetConn(req.Username)
	result := new(rpc.Result)
	conns.Traverse(func(conn server.Conn) error {
		result.Append(conn, nil)
		return nil
	})
	if req.DontPropagate {
		return result
	}
	req.DontPropagate = true
	r := self.peer.CheckUserStatus(req)
	result.Join(r)
	return result
}

func (self *serviceCenter) processSubscription() {
	for req := range self.subReqChan {
		if req == nil {
			return
		}
		go self.config.Subscribe(req)
	}
}

func (self *serviceCenter) NrConns() int {
	return self.conns.NrConns()
}

func (self *serviceCenter) NrUsers() int {
	return self.conns.NrUsers()
}

func (self *serviceCenter) AllUsernames() []string {
	return self.conns.AllUsernames()
}

func newServiceCenter(conf *config.ServiceConfig, fwdChan chan<- *rpc.ForwardRequest, peer rpc.UniqushConnPeer) *serviceCenter {
	if conf == nil || fwdChan == nil {
		return nil
	}
	ret := new(serviceCenter)
	ret.config = conf
	ret.conns = newTreeBasedConnMap(conf.MaxNrConns, conf.MaxNrUsers, conf.MaxNrConnsPerUser)
	ret.fwdChan = fwdChan
	ret.subReqChan = make(chan *rpc.SubscribeRequest)
	ret.peer = peer

	go ret.processSubscription()
	return ret
}
<commit_msg>fix #20.<commit_after>/*
 * Copyright 2013 Nan Deng
 *
 * 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 msgcenter

import (
	"errors"
	"fmt"
	"github.com/uniqush/uniqush-conn/config"
	"github.com/uniqush/uniqush-conn/proto/server"
	"github.com/uniqush/uniqush-conn/rpc"
	"io"
	"strings"
)

type serviceCenter struct {
	config     *config.ServiceConfig
	fwdChan    chan<- *rpc.ForwardRequest
	subReqChan chan *rpc.SubscribeRequest
	conns      connMap
	peer       rpc.UniqushConnPeer
}

func (self *serviceCenter) Stop() {
	self.conns.CloseAll()
	close(self.subReqChan)
}

func (self *serviceCenter) serveConn(conn server.Conn) {
	var reason error
	defer func() {
		self.config.OnLogout(conn, reason)
		self.conns.DelConn(conn)
		conn.Close()
	}()
	for {
		msg, err := conn.ReceiveMessage()
		if err != nil {
			if err != io.EOF {
				self.config.OnError(conn, err)
				reason = err
			}
			return
		}
		if msg != nil {
			self.config.OnMessage(conn, msg)
		}
	}
}

func (self *serviceCenter) NewConn(conn server.Conn) {
	if conn == nil {
		//self.config.OnError(conn, fmt.Errorf("Nil conn")
		return
	}
	usr := conn.Username()
	if len(usr) == 0 || strings.Contains(usr, ":") || strings.Contains(usr, "\n") {
		self.config.OnError(conn, fmt.Errorf("invalid username"))
		conn.Close()
		return
	}
	conn.SetMessageCache(self.config.Cache())
	conn.SetForwardRequestChannel(self.fwdChan)
	conn.SetSubscribeRequestChan(self.subReqChan)
	err := self.conns.AddConn(conn)
	if err != nil {
		self.config.OnError(conn, err)
		conn.Close()
		return
	}

	self.config.OnLogin(conn)
	go self.serveConn(conn)
	return
}

func (self *serviceCenter) Send(req *rpc.SendRequest) *rpc.Result {
	ret := new(rpc.Result)

	if req == nil {
		ret.SetError(fmt.Errorf("invalid request"))
		return ret
	}
	if req.Message == nil || req.Message.IsEmpty() {
		ret.SetError(fmt.Errorf("invalid request: empty message"))
		return ret
	}
	if len(req.Receivers) == 0 {
		ret.SetError(fmt.Errorf("invalid request: no receiver"))
		return ret
	}

	shouldPush := !req.DontPush
	shouldCache := !req.DontCache
	shouldPropagate := !req.DontPropagate
	receivers := req.Receivers

	for _, recver := range receivers {
		mid := req.Id
		msg := req.Message

		if shouldCache {
			mc := &rpc.MessageContainer{
				Sender:        "",
				SenderService: "",
				Message:       msg,
			}
			var err error
			mid, err = self.config.CacheMessage(recver, mc, req.TTL)
			if err != nil {
				ret.SetError(err)
				return ret
			}
		}

		if len(mid) == 0 {
			ret.SetError(fmt.Errorf("undefined message Id"))
			return ret
		}

		n := 0

		conns := self.conns.GetConn(recver)
		conns.Traverse(func(conn server.Conn) error {
			err := conn.SendMessage(msg, mid, nil, !req.NeverDigest)
			ret.Append(conn, err)
			if err != nil {
				conn.Close()
				// We won't delete this connection here.
				// Instead, we close it and let the reader
				// goroutine detect the error and close it.
			} else {
				n++
			}
			// Instead of returning an error,
			// we wourld rather let the Traverse() move forward.
			return nil
		})

		// Don't propagate this request to other instances in the cluster.
		req.DontPropagate = true
		// Don't push the message. We will push it on this node.
		req.DontPush = true
		// Don't cache the message. We have already cached it.
		req.DontCache = true
		req.Id = mid
		req.Receivers = []string{recver}

		if shouldPropagate {
			r := self.peer.Send(req)
			n += r.NrSuccess()
			ret.Join(r)
		}

		if n == 0 && shouldPush {
			self.config.Push(recver, "", "", req.PushInfo, mid, msg.Size())
		}
	}
	return ret
}

func (self *serviceCenter) Forward(req *rpc.ForwardRequest) *rpc.Result {
	ret := new(rpc.Result)

	if req == nil {
		ret.SetError(fmt.Errorf("invalid request"))
		return ret
	}
	if req.Message == nil || req.Message.IsEmpty() {
		ret.SetError(fmt.Errorf("invalid request: empty message"))
		return ret
	}
	if len(req.Receivers) == 0 {
		ret.SetError(fmt.Errorf("invalid request: no receiver"))
		return ret
	}

	mid := req.Id
	msg := req.Message
	mc := &req.MessageContainer

	var pushInfo map[string]string
	var shouldForward bool
	shouldPush := !req.DontPush
	shouldCache := !req.DontCache
	shouldPropagate := !req.DontPropagate

	if !req.DontAsk {
		// We need to ask for permission to forward this message.
		// This means the forward request is generated directly from a user,
		// not from a uniqush-conn node in a cluster.

		mc.Id = ""
		shouldForward, shouldPush, pushInfo = self.config.ShouldForward(req)

		if !shouldForward {
			return nil
		}
	}

	receivers := req.Receivers

	for _, recver := range receivers {
		if shouldCache {
			var err error
			mid, err = self.config.CacheMessage(recver, mc, req.TTL)
			if err != nil {
				ret.SetError(err)
				return ret
			}
		}

		if len(mid) == 0 {
			ret.SetError(fmt.Errorf("undefined message Id"))
			return ret
		}

		n := 0

		conns := self.conns.GetConn(recver)
		conns.Traverse(func(conn server.Conn) error {
			err := conn.ForwardMessage(req.Sender, req.SenderService, msg, mid, !req.NeverDigest)
			ret.Append(conn, err)
			if err != nil {
				conn.Close()
				// We won't delete this connection here.
				// Instead, we close it and let the reader
				// goroutine detect the error and close it.
			} else {
				n++
			}
			// Instead of returning an error,
			// we wourld rather let the Traverse() move forward.
			return nil
		})

		// forward the message if possible,
		// Don't propagate this request to other instances in the cluster.
		req.DontPropagate = true
		// Don't ask the permission to forward (we have already got the permission)
		req.DontAsk = true
		// And don't push the message. We will push it on this node.
		req.DontPush = true
		// Dont' cache it
		req.DontCache = true
		req.Id = mid
		req.Receivers = []string{recver}

		if shouldPropagate {
			r := self.peer.Forward(req)
			n += r.NrSuccess()
			ret.Join(r)
		}

		if n == 0 && shouldPush {
			self.config.Push(recver, req.SenderService, req.Sender, pushInfo, mid, msg.Size())
		}
	}
	return ret
}

func (self *serviceCenter) Redirect(req *rpc.RedirectRequest) *rpc.Result {
	conns := self.conns.GetConn(req.Receiver)
	var sc server.Conn
	result := new(rpc.Result)
	conns.Traverse(func(conn server.Conn) error {
		if len(req.ConnId) == 0 || conn.UniqId() == req.ConnId {
			sc = conn
			result.Append(sc, nil)
			return errors.New("done")
		}
		return nil
	})
	if sc != nil {
		// self.conns.DelConn(sc)
		sc.Redirect(req.CandidateSersers...)
		sc.Close()
		return result
	}

	if req.DontPropagate {
		return result
	}
	req.DontPropagate = true
	return self.peer.Redirect(req)
}

func (self *serviceCenter) CheckUserStatus(req *rpc.UserStatusQuery) *rpc.Result {
	conns := self.conns.GetConn(req.Username)
	result := new(rpc.Result)
	conns.Traverse(func(conn server.Conn) error {
		result.Append(conn, nil)
		return nil
	})
	if req.DontPropagate {
		return result
	}
	req.DontPropagate = true
	r := self.peer.CheckUserStatus(req)
	result.Join(r)
	return result
}

func (self *serviceCenter) processSubscription() {
	for req := range self.subReqChan {
		if req == nil {
			return
		}
		go self.config.Subscribe(req)
	}
}

func (self *serviceCenter) NrConns() int {
	return self.conns.NrConns()
}

func (self *serviceCenter) NrUsers() int {
	return self.conns.NrUsers()
}

func (self *serviceCenter) AllUsernames() []string {
	return self.conns.AllUsernames()
}

func newServiceCenter(conf *config.ServiceConfig, fwdChan chan<- *rpc.ForwardRequest, peer rpc.UniqushConnPeer) *serviceCenter {
	if conf == nil || fwdChan == nil {
		return nil
	}
	ret := new(serviceCenter)
	ret.config = conf
	ret.conns = newTreeBasedConnMap(conf.MaxNrConns, conf.MaxNrUsers, conf.MaxNrConnsPerUser)
	ret.fwdChan = fwdChan
	ret.subReqChan = make(chan *rpc.SubscribeRequest)
	ret.peer = peer

	go ret.processSubscription()
	return ret
}
<|endoftext|>