<commit_before>/*
Copyright 2018 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 main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	prowjobset "k8s.io/test-infra/prow/client/clientset/versioned"
	prowjobinfo "k8s.io/test-infra/prow/client/informers/externalversions"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/kube"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/pjutil"

	buildv1alpha1 "github.com/knative/build/pkg/apis/build/v1alpha1"
	"github.com/sirupsen/logrus"
	apimeta "k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp" // support gcp users in .kube/config
	"k8s.io/client-go/rest"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	ctrlruntimeclient "sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/manager"
)

type options struct {
	allContexts  bool
	buildCluster string
	config       string
	kubeconfig   string
	totURL       string

	// This is a termporary flag which gates the usage of plank.allow_cancellations config value
	// for build aborter.
	// TODO remove this flag and use directly the config flag.
	useAllowCancellations bool
}

func parseOptions() options {
	var o options
	if err := o.parse(flag.CommandLine, os.Args[1:]); err != nil {
		logrus.Fatalf("Invalid flags: %v", err)
	}
	return o
}

func (o *options) parse(flags *flag.FlagSet, args []string) error {
	flags.BoolVar(&o.allContexts, "all-contexts", false, "Monitor all cluster contexts, not just default")
	flags.StringVar(&o.totURL, "tot-url", "", "Tot URL")
	flags.StringVar(&o.kubeconfig, "kubeconfig", "", "Path to kubeconfig. Only required if out of cluster")
	flags.StringVar(&o.config, "config", "", "Path to prow config.yaml")
	flags.StringVar(&o.buildCluster, "build-cluster", "", "Path to file containing a YAML-marshalled kube.Cluster object. If empty, uses the local cluster.")
	flags.BoolVar(&o.useAllowCancellations, "use-allow-cancellations", false, "Gates the usage of plank.allow_cancellations config flag for build aborter")
	if err := flags.Parse(args); err != nil {
		return fmt.Errorf("parse flags: %v", err)
	}
	if o.kubeconfig != "" && o.buildCluster != "" {
		return errors.New("deprecated --build-cluster may not be used with --kubeconfig")
	}
	if o.buildCluster != "" {
		// TODO(fejta): change to warn and add a term date after plank migration
		logrus.Infof("--build-cluster is deprecated, please switch to --kubeconfig")
	}
	return nil
}

// stopper returns a channel that remains open until an interrupt is received.
func stopper() chan struct{} {
	stop := make(chan struct{})
	c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		logrus.Warn("Interrupt received, attempting clean shutdown...")
		close(stop)
		<-c
		logrus.Error("Second interrupt received, force exiting...")
		os.Exit(1)
	}()
	return stop
}

type buildConfig struct {
	client ctrlruntimeclient.Client
	// Only use the informer to add EventHandlers, for getting
	// objects use the client instead, its Reader interface is
	// backed by the cache
	informer cache.Informer
}

// newBuildConfig returns a client and informer capable of mutating and monitoring the specified config.
func newBuildConfig(cfg rest.Config, stop chan struct{}) (*buildConfig, error) {
	// Assume watches receive updates, but resync every 30m in case something wonky happens
	resyncInterval := 30 * time.Minute
	// We construct a manager because it has a client whose Reader interface is backed by its cache, which
	// is really nice to use, but the corresponding code is not exported.
	mgr, err := manager.New(&cfg, manager.Options{SyncPeriod: &resyncInterval, MetricsBindAddress: "0"})
	if err != nil {
		return nil, err
	}

	// Ensure the knative-build CRD is deployed
	// TODO(fejta): probably a better way to do this
	buildList := &buildv1alpha1.BuildList{}
	opts := &ctrlruntimeclient.ListOptions{Raw: &metav1.ListOptions{Limit: 1}}
	if err := mgr.GetClient().List(context.TODO(), buildList, ctrlruntimeclient.UseListOptions(opts)); err != nil {
		return nil, err
	}
	cache := mgr.GetCache()
	informer, err := cache.GetInformer(&buildv1alpha1.Build{})
	if err != nil {
		return nil, fmt.Errorf("failed to get cache for buildv1alpha1.Build: %v", err)
	}
	go cache.Start(stop)
	return &buildConfig{
		client:   mgr.GetClient(),
		informer: informer,
	}, nil
}

func main() {
	logrusutil.ComponentInit("build")

	o := parseOptions()

	pjutil.ServePProf()

	if err := buildv1alpha1.AddToScheme(scheme.Scheme); err != nil {
		logrus.WithError(err).Fatal("failed to add buildv1alpha1 to scheme")
	}

	configAgent := &config.Agent{}
	if o.config != "" {
		const ignoreJobConfig = ""
		if err := configAgent.Start(o.config, ignoreJobConfig); err != nil {
			logrus.WithError(err).Fatal("failed to load prow config")
		}
	}

	configs, err := kube.LoadClusterConfigs(o.kubeconfig, o.buildCluster)
	if err != nil {
		logrus.WithError(err).Fatal("Error building client configs")
	}

	local := configs[kube.InClusterContext]
	if !o.allContexts {
		logrus.Warn("Truncating to default context")
		configs = map[string]rest.Config{
			kube.DefaultClusterAlias: configs[kube.DefaultClusterAlias],
		}
	} else {
		// the InClusterContext is always mapped to DefaultClusterAlias in the controller, so there is no need to watch for this config.
		delete(configs, kube.InClusterContext)
	}

	stop := stopper()

	kc, err := kubernetes.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create local kubernetes client")
	}
	pjc, err := prowjobset.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create prowjob client")
	}
	pjif := prowjobinfo.NewSharedInformerFactory(pjc, 30*time.Minute)
	pjif.Prow().V1().ProwJobs().Lister()
	go pjif.Start(stop)

	buildConfigs := map[string]buildConfig{}
	for context, cfg := range configs {
		var bc *buildConfig
		bc, err = newBuildConfig(cfg, stop)
		if apimeta.IsNoMatchError(err) {
			logrus.WithError(err).Warnf("Ignoring %s: knative build CRD not deployed", context)
			continue
		}
		if err != nil {
			logrus.WithError(err).Fatalf("Failed to create %s build client", context)
		}
		buildConfigs[context] = *bc
	}

	opts := controllerOptions{
		kc:                    kc,
		pjc:                   pjc,
		pji:                   pjif.Prow().V1().ProwJobs(),
		buildConfigs:          buildConfigs,
		totURL:                o.totURL,
		prowConfig:            configAgent.Config,
		rl:                    kube.RateLimiter(controllerName),
		useAllowCancellations: o.useAllowCancellations,
	}
	controller, err := newController(opts)
	if err != nil {
		logrus.WithError(err).Fatal("Error creating controller")
	}
	if err := controller.Run(2, stop); err != nil {
		logrus.WithError(err).Fatal("Error running controller")
	}
	logrus.Info("Finished")
}
<commit_msg>Update controller-runtime list usage<commit_after>/*
Copyright 2018 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 main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	prowjobset "k8s.io/test-infra/prow/client/clientset/versioned"
	prowjobinfo "k8s.io/test-infra/prow/client/informers/externalversions"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/kube"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/pjutil"

	buildv1alpha1 "github.com/knative/build/pkg/apis/build/v1alpha1"
	"github.com/sirupsen/logrus"
	apimeta "k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp" // support gcp users in .kube/config
	"k8s.io/client-go/rest"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	ctrlruntimeclient "sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/manager"
)

type options struct {
	allContexts  bool
	buildCluster string
	config       string
	kubeconfig   string
	totURL       string

	// This is a termporary flag which gates the usage of plank.allow_cancellations config value
	// for build aborter.
	// TODO remove this flag and use directly the config flag.
	useAllowCancellations bool
}

func parseOptions() options {
	var o options
	if err := o.parse(flag.CommandLine, os.Args[1:]); err != nil {
		logrus.Fatalf("Invalid flags: %v", err)
	}
	return o
}

func (o *options) parse(flags *flag.FlagSet, args []string) error {
	flags.BoolVar(&o.allContexts, "all-contexts", false, "Monitor all cluster contexts, not just default")
	flags.StringVar(&o.totURL, "tot-url", "", "Tot URL")
	flags.StringVar(&o.kubeconfig, "kubeconfig", "", "Path to kubeconfig. Only required if out of cluster")
	flags.StringVar(&o.config, "config", "", "Path to prow config.yaml")
	flags.StringVar(&o.buildCluster, "build-cluster", "", "Path to file containing a YAML-marshalled kube.Cluster object. If empty, uses the local cluster.")
	flags.BoolVar(&o.useAllowCancellations, "use-allow-cancellations", false, "Gates the usage of plank.allow_cancellations config flag for build aborter")
	if err := flags.Parse(args); err != nil {
		return fmt.Errorf("parse flags: %v", err)
	}
	if o.kubeconfig != "" && o.buildCluster != "" {
		return errors.New("deprecated --build-cluster may not be used with --kubeconfig")
	}
	if o.buildCluster != "" {
		// TODO(fejta): change to warn and add a term date after plank migration
		logrus.Infof("--build-cluster is deprecated, please switch to --kubeconfig")
	}
	return nil
}

// stopper returns a channel that remains open until an interrupt is received.
func stopper() chan struct{} {
	stop := make(chan struct{})
	c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		logrus.Warn("Interrupt received, attempting clean shutdown...")
		close(stop)
		<-c
		logrus.Error("Second interrupt received, force exiting...")
		os.Exit(1)
	}()
	return stop
}

type buildConfig struct {
	client ctrlruntimeclient.Client
	// Only use the informer to add EventHandlers, for getting
	// objects use the client instead, its Reader interface is
	// backed by the cache
	informer cache.Informer
}

type listSingleItem struct{}

func (_ listSingleItem) ApplyToList(opts *ctrlruntimeclient.ListOptions) {
	if opts.Raw == nil {
		opts.Raw = &metav1.ListOptions{}
	}
	opts.Raw.Limit = 1
}

// newBuildConfig returns a client and informer capable of mutating and monitoring the specified config.
func newBuildConfig(cfg rest.Config, stop chan struct{}) (*buildConfig, error) {
	// Assume watches receive updates, but resync every 30m in case something wonky happens
	resyncInterval := 30 * time.Minute
	// We construct a manager because it has a client whose Reader interface is backed by its cache, which
	// is really nice to use, but the corresponding code is not exported.
	mgr, err := manager.New(&cfg, manager.Options{SyncPeriod: &resyncInterval, MetricsBindAddress: "0"})
	if err != nil {
		return nil, err
	}

	// Ensure the knative-build CRD is deployed
	// TODO(fejta): probably a better way to do this
	buildList := &buildv1alpha1.BuildList{}
	if err := mgr.GetClient().List(context.TODO(), buildList, listSingleItem{}); err != nil {
		return nil, err
	}
	cache := mgr.GetCache()
	informer, err := cache.GetInformer(&buildv1alpha1.Build{})
	if err != nil {
		return nil, fmt.Errorf("failed to get cache for buildv1alpha1.Build: %v", err)
	}
	go cache.Start(stop)
	return &buildConfig{
		client:   mgr.GetClient(),
		informer: informer,
	}, nil
}

func main() {
	logrusutil.ComponentInit("build")

	o := parseOptions()

	pjutil.ServePProf()

	if err := buildv1alpha1.AddToScheme(scheme.Scheme); err != nil {
		logrus.WithError(err).Fatal("failed to add buildv1alpha1 to scheme")
	}

	configAgent := &config.Agent{}
	if o.config != "" {
		const ignoreJobConfig = ""
		if err := configAgent.Start(o.config, ignoreJobConfig); err != nil {
			logrus.WithError(err).Fatal("failed to load prow config")
		}
	}

	configs, err := kube.LoadClusterConfigs(o.kubeconfig, o.buildCluster)
	if err != nil {
		logrus.WithError(err).Fatal("Error building client configs")
	}

	local := configs[kube.InClusterContext]
	if !o.allContexts {
		logrus.Warn("Truncating to default context")
		configs = map[string]rest.Config{
			kube.DefaultClusterAlias: configs[kube.DefaultClusterAlias],
		}
	} else {
		// the InClusterContext is always mapped to DefaultClusterAlias in the controller, so there is no need to watch for this config.
		delete(configs, kube.InClusterContext)
	}

	stop := stopper()

	kc, err := kubernetes.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create local kubernetes client")
	}
	pjc, err := prowjobset.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create prowjob client")
	}
	pjif := prowjobinfo.NewSharedInformerFactory(pjc, 30*time.Minute)
	pjif.Prow().V1().ProwJobs().Lister()
	go pjif.Start(stop)

	buildConfigs := map[string]buildConfig{}
	for context, cfg := range configs {
		var bc *buildConfig
		bc, err = newBuildConfig(cfg, stop)
		if apimeta.IsNoMatchError(err) {
			logrus.WithError(err).Warnf("Ignoring %s: knative build CRD not deployed", context)
			continue
		}
		if err != nil {
			logrus.WithError(err).Fatalf("Failed to create %s build client", context)
		}
		buildConfigs[context] = *bc
	}

	opts := controllerOptions{
		kc:                    kc,
		pjc:                   pjc,
		pji:                   pjif.Prow().V1().ProwJobs(),
		buildConfigs:          buildConfigs,
		totURL:                o.totURL,
		prowConfig:            configAgent.Config,
		rl:                    kube.RateLimiter(controllerName),
		useAllowCancellations: o.useAllowCancellations,
	}
	controller, err := newController(opts)
	if err != nil {
		logrus.WithError(err).Fatal("Error creating controller")
	}
	if err := controller.Run(2, stop); err != nil {
		logrus.WithError(err).Fatal("Error running controller")
	}
	logrus.Info("Finished")
}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	prowjobset "k8s.io/test-infra/prow/client/clientset/versioned"
	prowjobinfo "k8s.io/test-infra/prow/client/informers/externalversions"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/kube"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/pjutil"

	buildset "github.com/knative/build/pkg/client/clientset/versioned"
	buildinfo "github.com/knative/build/pkg/client/informers/externalversions"
	buildinfov1alpha1 "github.com/knative/build/pkg/client/informers/externalversions/build/v1alpha1"
	"github.com/sirupsen/logrus"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp" // support gcp users in .kube/config
	"k8s.io/client-go/rest"
)

type options struct {
	allContexts  bool
	buildCluster string
	config       string
	kubeconfig   string
	totURL       string

	// This is a termporary flag which gates the usage of plank.allow_cancellations config value
	// for build aborter.
	// TODO remove this flag and use directly the config flag.
	useAllowCancellations bool
}

func parseOptions() options {
	var o options
	if err := o.parse(flag.CommandLine, os.Args[1:]); err != nil {
		logrus.Fatalf("Invalid flags: %v", err)
	}
	return o
}

func (o *options) parse(flags *flag.FlagSet, args []string) error {
	flags.BoolVar(&o.allContexts, "all-contexts", false, "Monitor all cluster contexts, not just default")
	flags.StringVar(&o.totURL, "tot-url", "", "Tot URL")
	flags.StringVar(&o.kubeconfig, "kubeconfig", "", "Path to kubeconfig. Only required if out of cluster")
	flags.StringVar(&o.config, "config", "", "Path to prow config.yaml")
	flags.StringVar(&o.buildCluster, "build-cluster", "", "Path to file containing a YAML-marshalled kube.Cluster object. If empty, uses the local cluster.")
	flags.BoolVar(&o.useAllowCancellations, "use-allow-cancellations", false, "Gates the usage of plank.allow_cancellations config flag for build aborter")
	if err := flags.Parse(args); err != nil {
		return fmt.Errorf("parse flags: %v", err)
	}
	if o.kubeconfig != "" && o.buildCluster != "" {
		return errors.New("deprecated --build-cluster may not be used with --kubeconfig")
	}
	if o.buildCluster != "" {
		// TODO(fejta): change to warn and add a term date after plank migration
		logrus.Infof("--build-cluster is deprecated, please switch to --kubeconfig")
	}
	return nil
}

// stopper returns a channel that remains open until an interrupt is received.
func stopper() chan struct{} {
	stop := make(chan struct{})
	c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		logrus.Warn("Interrupt received, attempting clean shutdown...")
		close(stop)
		<-c
		logrus.Error("Second interrupt received, force exiting...")
		os.Exit(1)
	}()
	return stop
}

type buildConfig struct {
	client   buildset.Interface
	informer buildinfov1alpha1.BuildInformer
}

// newBuildConfig returns a client and informer capable of mutating and monitoring the specified config.
func newBuildConfig(cfg rest.Config, stop chan struct{}) (*buildConfig, error) {
	bc, err := buildset.NewForConfig(&cfg)
	if err != nil {
		return nil, err
	}

	// Ensure the knative-build CRD is deployed
	// TODO(fejta): probably a better way to do this
	_, err = bc.BuildV1alpha1().Builds("").List(metav1.ListOptions{Limit: 1})
	if err != nil {
		return nil, err
	}
	// Assume watches receive updates, but resync every 30m in case something wonky happens
	bif := buildinfo.NewSharedInformerFactory(bc, 30*time.Minute)
	bif.Build().V1alpha1().Builds().Lister()
	go bif.Start(stop)
	return &buildConfig{
		client:   bc,
		informer: bif.Build().V1alpha1().Builds(),
	}, nil
}

func main() {
	o := parseOptions()
	logrusutil.NewDefaultFieldsFormatter(nil, logrus.Fields{"component": "build"})

	pjutil.ServePProf()

	configAgent := &config.Agent{}
	if o.config != "" {
		const ignoreJobConfig = ""
		if err := configAgent.Start(o.config, ignoreJobConfig); err != nil {
			logrus.WithError(err).Fatal("failed to load prow config")
		}
	}

	configs, err := kube.LoadClusterConfigs(o.kubeconfig, o.buildCluster)
	if err != nil {
		logrus.WithError(err).Fatal("Error building client configs")
	}

	local := configs[kube.InClusterContext]
	if !o.allContexts {
		logrus.Warn("Truncating to default context")
		configs = map[string]rest.Config{
			kube.DefaultClusterAlias: configs[kube.DefaultClusterAlias],
		}
	}

	stop := stopper()

	kc, err := kubernetes.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create local kubernetes client")
	}
	pjc, err := prowjobset.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create prowjob client")
	}
	pjif := prowjobinfo.NewSharedInformerFactory(pjc, 30*time.Minute)
	pjif.Prow().V1().ProwJobs().Lister()
	go pjif.Start(stop)

	buildConfigs := map[string]buildConfig{}
	for context, cfg := range configs {
		var bc *buildConfig
		bc, err = newBuildConfig(cfg, stop)
		if apierrors.IsNotFound(err) {
			logrus.WithError(err).Warnf("Ignoring %s: knative build CRD not deployed", context)
			continue
		}
		if err != nil {
			logrus.WithError(err).Fatalf("Failed to create %s build client", context)
		}
		buildConfigs[context] = *bc
	}

	opts := controllerOptions{
		kc:                    kc,
		pjc:                   pjc,
		pji:                   pjif.Prow().V1().ProwJobs(),
		buildConfigs:          buildConfigs,
		totURL:                o.totURL,
		prowConfig:            configAgent.Config,
		rl:                    kube.RateLimiter(controllerName),
		useAllowCancellations: o.useAllowCancellations,
	}
	controller, err := newController(opts)
	if err != nil {
		logrus.WithError(err).Fatal("Error creating controller")
	}
	if err := controller.Run(2, stop); err != nil {
		logrus.WithError(err).Fatal("Error running controller")
	}
	logrus.Info("Finished")
}
<commit_msg>Make sure in the build controller that the in-cluster and default cluster alias do not point to the same cluster when all contexts are allowed<commit_after>/*
Copyright 2018 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 main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	prowjobset "k8s.io/test-infra/prow/client/clientset/versioned"
	prowjobinfo "k8s.io/test-infra/prow/client/informers/externalversions"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/kube"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/pjutil"

	buildset "github.com/knative/build/pkg/client/clientset/versioned"
	buildinfo "github.com/knative/build/pkg/client/informers/externalversions"
	buildinfov1alpha1 "github.com/knative/build/pkg/client/informers/externalversions/build/v1alpha1"
	"github.com/sirupsen/logrus"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp" // support gcp users in .kube/config
	"k8s.io/client-go/rest"
)

type options struct {
	allContexts  bool
	buildCluster string
	config       string
	kubeconfig   string
	totURL       string

	// This is a termporary flag which gates the usage of plank.allow_cancellations config value
	// for build aborter.
	// TODO remove this flag and use directly the config flag.
	useAllowCancellations bool
}

func parseOptions() options {
	var o options
	if err := o.parse(flag.CommandLine, os.Args[1:]); err != nil {
		logrus.Fatalf("Invalid flags: %v", err)
	}
	return o
}

func (o *options) parse(flags *flag.FlagSet, args []string) error {
	flags.BoolVar(&o.allContexts, "all-contexts", false, "Monitor all cluster contexts, not just default")
	flags.StringVar(&o.totURL, "tot-url", "", "Tot URL")
	flags.StringVar(&o.kubeconfig, "kubeconfig", "", "Path to kubeconfig. Only required if out of cluster")
	flags.StringVar(&o.config, "config", "", "Path to prow config.yaml")
	flags.StringVar(&o.buildCluster, "build-cluster", "", "Path to file containing a YAML-marshalled kube.Cluster object. If empty, uses the local cluster.")
	flags.BoolVar(&o.useAllowCancellations, "use-allow-cancellations", false, "Gates the usage of plank.allow_cancellations config flag for build aborter")
	if err := flags.Parse(args); err != nil {
		return fmt.Errorf("parse flags: %v", err)
	}
	if o.kubeconfig != "" && o.buildCluster != "" {
		return errors.New("deprecated --build-cluster may not be used with --kubeconfig")
	}
	if o.buildCluster != "" {
		// TODO(fejta): change to warn and add a term date after plank migration
		logrus.Infof("--build-cluster is deprecated, please switch to --kubeconfig")
	}
	return nil
}

// stopper returns a channel that remains open until an interrupt is received.
func stopper() chan struct{} {
	stop := make(chan struct{})
	c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		logrus.Warn("Interrupt received, attempting clean shutdown...")
		close(stop)
		<-c
		logrus.Error("Second interrupt received, force exiting...")
		os.Exit(1)
	}()
	return stop
}

type buildConfig struct {
	client   buildset.Interface
	informer buildinfov1alpha1.BuildInformer
}

// newBuildConfig returns a client and informer capable of mutating and monitoring the specified config.
func newBuildConfig(cfg rest.Config, stop chan struct{}) (*buildConfig, error) {
	bc, err := buildset.NewForConfig(&cfg)
	if err != nil {
		return nil, err
	}

	// Ensure the knative-build CRD is deployed
	// TODO(fejta): probably a better way to do this
	_, err = bc.BuildV1alpha1().Builds("").List(metav1.ListOptions{Limit: 1})
	if err != nil {
		return nil, err
	}
	// Assume watches receive updates, but resync every 30m in case something wonky happens
	bif := buildinfo.NewSharedInformerFactory(bc, 30*time.Minute)
	bif.Build().V1alpha1().Builds().Lister()
	go bif.Start(stop)
	return &buildConfig{
		client:   bc,
		informer: bif.Build().V1alpha1().Builds(),
	}, nil
}

func main() {
	o := parseOptions()
	logrusutil.NewDefaultFieldsFormatter(nil, logrus.Fields{"component": "build"})

	pjutil.ServePProf()

	configAgent := &config.Agent{}
	if o.config != "" {
		const ignoreJobConfig = ""
		if err := configAgent.Start(o.config, ignoreJobConfig); err != nil {
			logrus.WithError(err).Fatal("failed to load prow config")
		}
	}

	configs, err := kube.LoadClusterConfigs(o.kubeconfig, o.buildCluster)
	if err != nil {
		logrus.WithError(err).Fatal("Error building client configs")
	}

	local := configs[kube.InClusterContext]
	if !o.allContexts {
		logrus.Warn("Truncating to default context")
		configs = map[string]rest.Config{
			kube.DefaultClusterAlias: configs[kube.DefaultClusterAlias],
		}
	} else {
		if configs[kube.DefaultClusterAlias].Host == configs[kube.InClusterContext].Host {
			delete(configs, kube.InClusterContext)
		}
	}

	stop := stopper()

	kc, err := kubernetes.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create local kubernetes client")
	}
	pjc, err := prowjobset.NewForConfig(&local)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to create prowjob client")
	}
	pjif := prowjobinfo.NewSharedInformerFactory(pjc, 30*time.Minute)
	pjif.Prow().V1().ProwJobs().Lister()
	go pjif.Start(stop)

	buildConfigs := map[string]buildConfig{}
	for context, cfg := range configs {
		var bc *buildConfig
		bc, err = newBuildConfig(cfg, stop)
		if apierrors.IsNotFound(err) {
			logrus.WithError(err).Warnf("Ignoring %s: knative build CRD not deployed", context)
			continue
		}
		if err != nil {
			logrus.WithError(err).Fatalf("Failed to create %s build client", context)
		}
		buildConfigs[context] = *bc
	}

	opts := controllerOptions{
		kc:                    kc,
		pjc:                   pjc,
		pji:                   pjif.Prow().V1().ProwJobs(),
		buildConfigs:          buildConfigs,
		totURL:                o.totURL,
		prowConfig:            configAgent.Config,
		rl:                    kube.RateLimiter(controllerName),
		useAllowCancellations: o.useAllowCancellations,
	}
	controller, err := newController(opts)
	if err != nil {
		logrus.WithError(err).Fatal("Error creating controller")
	}
	if err := controller.Run(2, stop); err != nil {
		logrus.WithError(err).Fatal("Error running controller")
	}
	logrus.Info("Finished")
}
<|endoftext|>


<commit_before>//////////////////////////////////////////
// archivex.go
// Jhonathan Paulo Banczek - 2014
// jpbanczek@gmail.com - jhoonb.com
//////////////////////////////////////////

package archivex

import (
	"archive/tar"
	"archive/zip"
	"io"
	"io/ioutil"
	// "log"
	"os"
	"path"
	"strings"
)

// interface
type Archivex interface {
	Create(name string) error
	Add(name string, file []byte) error
	AddFile(name string) error
	AddAll(dir string, includeCurrentFolder bool) error
	Close() error
}

// ArchiveWriteFunc is the closure used by an archive's AddAll method to actually put a file into an archive
// Note that for directory entries, this func will be called with a nil 'file' param
type ArchiveWriteFunc func(info os.FileInfo, file io.Reader, entryName string) (err error)

// ZipFile implement *zip.Writer
type ZipFile struct {
	Writer *zip.Writer
	Name   string
}

// TarFile implement *tar.Writer
type TarFile struct {
	Writer *tar.Writer
	Name   string
}

// Create new file zip
func (z *ZipFile) Create(name string) error {
	// check extension .zip
	if strings.HasSuffix(name, ".zip") != true {
		if strings.HasSuffix(name, ".tar.gz") == true {
			name = strings.Replace(name, ".tar.gz", ".zip", -1)
		} else {
			name = name + ".zip"
		}
	}
	z.Name = name
	file, err := os.Create(z.Name)
	if err != nil {
		return err
	}
	z.Writer = zip.NewWriter(file)
	return nil
}

// Add add byte in archive zip
func (z *ZipFile) Add(name string, file []byte) error {

	iow, err := z.Writer.Create(name)
	if err != nil {
		return err
	}
	_, err = iow.Write(file)
	return err
}

// AddFile add file from dir in archive
func (z *ZipFile) AddFile(name string) error {
	bytearq, err := ioutil.ReadFile(name)
	if err != nil {
		return err
	}
	filep, err := z.Writer.Create(name)
	if err != nil {
		return err
	}
	_, err = filep.Write(bytearq)
	if err != nil {
		return err
	}
	return nil
}

// AddAll adds all files from dir in archive, recursively.
// Directories receive a zero-size entry in the archive, with a trailing slash in the header name, and no compression
func (z *ZipFile) AddAll(dir string, includeCurrentFolder bool) error {
	dir = path.Clean(dir)
	return addAll(dir, dir, includeCurrentFolder, func(info os.FileInfo, file io.Reader, entryName string) (err error) {

		// Create a header based off of the fileinfo
		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		// If it's a file, set the compression method to deflate (leave directories uncompressed)
		if !info.IsDir() {
			header.Method = zip.Deflate
		}

		// Set the header's name to what we want--it may not include the top folder
		header.Name = entryName

		// Add a trailing slash if the entry is a directory
		if info.IsDir() {
			header.Name += string(os.PathSeparator)
		}

		// Get a writer in the archive based on our header
		writer, err := z.Writer.CreateHeader(header)
		if err != nil {
			return err
		}

		// If we have a file to write (i.e., not a directory) then pipe the file into the archive writer
		if file != nil {
			if _, err := io.Copy(writer, file); err != nil {
				return err
			}
		}

		return nil
	})
}

func (z *ZipFile) Close() error {
	err := z.Writer.Close()
	return err
}

// Create new Tar file
func (t *TarFile) Create(name string) error {
	// check extension .zip
	if strings.HasSuffix(name, ".tar.gz") != true {
		if strings.HasSuffix(name, ".zip") == true {
			name = strings.Replace(name, ".zip", ".tar.gz", -1)
		} else {
			name = name + ".tar.gz"
		}
	}
	t.Name = name
	file, err := os.Create(t.Name)
	if err != nil {
		return err
	}
	t.Writer = tar.NewWriter(file)
	return nil
}

// Add add byte in archive tar
func (t *TarFile) Add(name string, file []byte) error {

	hdr := &tar.Header{Name: name, Size: int64(len(file)), Mode: 0777}
	if err := t.Writer.WriteHeader(hdr); err != nil {
		return err
	}
	_, err := t.Writer.Write(file)
	return err
}

// AddFile add file from dir in archive tar
func (t *TarFile) AddFile(name string) error {
	bytearq, err := ioutil.ReadFile(name)
	if err != nil {
		return err
	}

	info, err := os.Stat(name)
	if err != nil {
		return err
	}

	header, err := tar.FileInfoHeader(info, "")
	if err != nil {
		return err
	}

	err = t.Writer.WriteHeader(header)
	if err != nil {
		return err
	}
	_, err = t.Writer.Write(bytearq)
	if err != nil {
		return err
	}
	return nil

}

// AddAll adds all files from dir in archive
// Tar does not support directories
func (t *TarFile) AddAll(dir string, includeCurrentFolder bool) error {
	dir = path.Clean(dir)
	return addAll(dir, dir, includeCurrentFolder, func(info os.FileInfo, file io.Reader, entryName string) (err error) {

		// Skip directory entries
		if file == nil {
			return nil
		}

		// Create a header based off of the fileinfo
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}

		// Set the header's name to what we want--it may not include the top folder
		header.Name = entryName

		// Write the header into the tar file
		if err := t.Writer.WriteHeader(header); err != nil {
			return err
		}

		// Pipe the file into the tar
		if _, err := io.Copy(t.Writer, file); err != nil {
			return err
		}

		return nil
	})
}

// Close the file Tar
func (t *TarFile) Close() error {
	err := t.Writer.Close()
	return err
}

func getSubDir(dir string, rootDir string, includeCurrentFolder bool) (subDir string) {

	subDir = strings.Replace(dir, rootDir, "", 1)

	if includeCurrentFolder {
		parts := strings.Split(rootDir, string(os.PathSeparator))
		subDir = path.Join(parts[len(parts)-1], subDir)
	}

	return
}

// addAll is used to recursively go down through directories and add each file and directory to an archive, based on an ArchiveWriteFunc given to it
func addAll(dir string, rootDir string, includeCurrentFolder bool, writerFunc ArchiveWriteFunc) error {

	// Get a list of all entries in the directory, as []os.FileInfo
	fileInfos, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}

	// Loop through all entries
	for _, info := range fileInfos {

		full := path.Join(dir, info.Name())

		// If the entry is a file, get an io.Reader for it
		var file io.Reader
		if !info.IsDir() {
			file, err = os.Open(full)
			if err != nil {
				return err
			}
		}

		// Write the entry into the archive
		subDir := getSubDir(dir, rootDir, includeCurrentFolder)
		entryName := path.Join(subDir, info.Name())
		if err := writerFunc(info, file, entryName); err != nil {
			return err
		}

		// If the entry is a directory, recurse into it
		if info.IsDir() {
			addAll(full, rootDir, includeCurrentFolder, writerFunc)
		}
	}

	return nil
}
<commit_msg>changes 0777 to 0666, which is more sensible<commit_after>//////////////////////////////////////////
// archivex.go
// Jhonathan Paulo Banczek - 2014
// jpbanczek@gmail.com - jhoonb.com
//////////////////////////////////////////

package archivex

import (
	"archive/tar"
	"archive/zip"
	"io"
	"io/ioutil"
	// "log"
	"os"
	"path"
	"strings"
)

// interface
type Archivex interface {
	Create(name string) error
	Add(name string, file []byte) error
	AddFile(name string) error
	AddAll(dir string, includeCurrentFolder bool) error
	Close() error
}

// ArchiveWriteFunc is the closure used by an archive's AddAll method to actually put a file into an archive
// Note that for directory entries, this func will be called with a nil 'file' param
type ArchiveWriteFunc func(info os.FileInfo, file io.Reader, entryName string) (err error)

// ZipFile implement *zip.Writer
type ZipFile struct {
	Writer *zip.Writer
	Name   string
}

// TarFile implement *tar.Writer
type TarFile struct {
	Writer *tar.Writer
	Name   string
}

// Create new file zip
func (z *ZipFile) Create(name string) error {
	// check extension .zip
	if strings.HasSuffix(name, ".zip") != true {
		if strings.HasSuffix(name, ".tar.gz") == true {
			name = strings.Replace(name, ".tar.gz", ".zip", -1)
		} else {
			name = name + ".zip"
		}
	}
	z.Name = name
	file, err := os.Create(z.Name)
	if err != nil {
		return err
	}
	z.Writer = zip.NewWriter(file)
	return nil
}

// Add add byte in archive zip
func (z *ZipFile) Add(name string, file []byte) error {

	iow, err := z.Writer.Create(name)
	if err != nil {
		return err
	}
	_, err = iow.Write(file)
	return err
}

// AddFile add file from dir in archive
func (z *ZipFile) AddFile(name string) error {
	bytearq, err := ioutil.ReadFile(name)
	if err != nil {
		return err
	}
	filep, err := z.Writer.Create(name)
	if err != nil {
		return err
	}
	_, err = filep.Write(bytearq)
	if err != nil {
		return err
	}
	return nil
}

// AddAll adds all files from dir in archive, recursively.
// Directories receive a zero-size entry in the archive, with a trailing slash in the header name, and no compression
func (z *ZipFile) AddAll(dir string, includeCurrentFolder bool) error {
	dir = path.Clean(dir)
	return addAll(dir, dir, includeCurrentFolder, func(info os.FileInfo, file io.Reader, entryName string) (err error) {

		// Create a header based off of the fileinfo
		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		// If it's a file, set the compression method to deflate (leave directories uncompressed)
		if !info.IsDir() {
			header.Method = zip.Deflate
		}

		// Set the header's name to what we want--it may not include the top folder
		header.Name = entryName

		// Add a trailing slash if the entry is a directory
		if info.IsDir() {
			header.Name += string(os.PathSeparator)
		}

		// Get a writer in the archive based on our header
		writer, err := z.Writer.CreateHeader(header)
		if err != nil {
			return err
		}

		// If we have a file to write (i.e., not a directory) then pipe the file into the archive writer
		if file != nil {
			if _, err := io.Copy(writer, file); err != nil {
				return err
			}
		}

		return nil
	})
}

func (z *ZipFile) Close() error {
	err := z.Writer.Close()
	return err
}

// Create new Tar file
func (t *TarFile) Create(name string) error {
	// check extension .zip
	if strings.HasSuffix(name, ".tar.gz") != true {
		if strings.HasSuffix(name, ".zip") == true {
			name = strings.Replace(name, ".zip", ".tar.gz", -1)
		} else {
			name = name + ".tar.gz"
		}
	}
	t.Name = name
	file, err := os.Create(t.Name)
	if err != nil {
		return err
	}
	t.Writer = tar.NewWriter(file)
	return nil
}

// Add add byte in archive tar
func (t *TarFile) Add(name string, file []byte) error {

	hdr := &tar.Header{Name: name, Size: int64(len(file)), Mode: 0666}
	if err := t.Writer.WriteHeader(hdr); err != nil {
		return err
	}
	_, err := t.Writer.Write(file)
	return err
}

// AddFile add file from dir in archive tar
func (t *TarFile) AddFile(name string) error {
	bytearq, err := ioutil.ReadFile(name)
	if err != nil {
		return err
	}

	info, err := os.Stat(name)
	if err != nil {
		return err
	}

	header, err := tar.FileInfoHeader(info, "")
	if err != nil {
		return err
	}

	err = t.Writer.WriteHeader(header)
	if err != nil {
		return err
	}
	_, err = t.Writer.Write(bytearq)
	if err != nil {
		return err
	}
	return nil

}

// AddAll adds all files from dir in archive
// Tar does not support directories
func (t *TarFile) AddAll(dir string, includeCurrentFolder bool) error {
	dir = path.Clean(dir)
	return addAll(dir, dir, includeCurrentFolder, func(info os.FileInfo, file io.Reader, entryName string) (err error) {

		// Skip directory entries
		if file == nil {
			return nil
		}

		// Create a header based off of the fileinfo
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}

		// Set the header's name to what we want--it may not include the top folder
		header.Name = entryName

		// Write the header into the tar file
		if err := t.Writer.WriteHeader(header); err != nil {
			return err
		}

		// Pipe the file into the tar
		if _, err := io.Copy(t.Writer, file); err != nil {
			return err
		}

		return nil
	})
}

// Close the file Tar
func (t *TarFile) Close() error {
	err := t.Writer.Close()
	return err
}

func getSubDir(dir string, rootDir string, includeCurrentFolder bool) (subDir string) {

	subDir = strings.Replace(dir, rootDir, "", 1)

	if includeCurrentFolder {
		parts := strings.Split(rootDir, string(os.PathSeparator))
		subDir = path.Join(parts[len(parts)-1], subDir)
	}

	return
}

// addAll is used to recursively go down through directories and add each file and directory to an archive, based on an ArchiveWriteFunc given to it
func addAll(dir string, rootDir string, includeCurrentFolder bool, writerFunc ArchiveWriteFunc) error {

	// Get a list of all entries in the directory, as []os.FileInfo
	fileInfos, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}

	// Loop through all entries
	for _, info := range fileInfos {

		full := path.Join(dir, info.Name())

		// If the entry is a file, get an io.Reader for it
		var file io.Reader
		if !info.IsDir() {
			file, err = os.Open(full)
			if err != nil {
				return err
			}
		}

		// Write the entry into the archive
		subDir := getSubDir(dir, rootDir, includeCurrentFolder)
		entryName := path.Join(subDir, info.Name())
		if err := writerFunc(info, file, entryName); err != nil {
			return err
		}

		// If the entry is a directory, recurse into it
		if info.IsDir() {
			addAll(full, rootDir, includeCurrentFolder, writerFunc)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package consumergroup

import (
	"errors"
	"fmt"
	"runtime/debug"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/samuel/go-zookeeper/zk"
	"github.com/sirupsen/logrus"
)

const (
	cgInit = iota
	cgStart
	cgStopped
)

const (
	restartEvent = iota
	quitEvent
)

// ConsumerGroup consume message from Kafka with rebalancing supports
type ConsumerGroup struct {
	name           string
	storage        groupStorage
	topicConsumers map[string]*topicConsumer
	saramaConsumer sarama.Consumer

	id        string
	state     int
	wg        sync.WaitGroup
	stopCh    chan struct{}
	triggerCh chan int
	stopOnce  *sync.Once
	owners    map[string]map[int32]string

	config *Config
	logger *logrus.Logger

	onLoad, onClose []func()
}

// NewConsumerGroup create the ConsumerGroup instance with config
func NewConsumerGroup(config *Config) (*ConsumerGroup, error) {
	if config == nil {
		return nil, errors.New("config can't be empty")
	}
	err := config.validate()
	if err != nil {
		return nil, fmt.Errorf("vaildate config failed, as %s", err)
	}

	cg := new(ConsumerGroup)
	cg.state = cgInit
	cg.config = config
	cg.id = config.ConsumerID
	if cg.id == "" {
		cg.id = genConsumerID()
	}
	cg.name = config.GroupID
	cg.stopOnce = new(sync.Once)
	cg.triggerCh = make(chan int)
	cg.topicConsumers = make(map[string]*topicConsumer)
	cg.onLoad = make([]func(), 0)
	cg.onClose = make([]func(), 0)
	cg.storage = newZKGroupStorage(config.ZkList, config.ZkSessionTimeout)
	cg.logger = logrus.New()
	if _, ok := cg.storage.(*zkGroupStorage); ok {
		cg.storage.(*zkGroupStorage).Chroot(config.Chroot)
	}

	err = cg.initSaramaConsumer()
	if err != nil {
		return nil, fmt.Errorf("init sarama consumer, as %s", err)
	}
	cg.owners = make(map[string]map[int32]string)
	for _, topic := range config.TopicList {
		cg.topicConsumers[topic] = newTopicConsumer(cg, topic)
		cg.owners[topic] = make(map[int32]string)
	}
	return cg, nil
}

func (cg *ConsumerGroup) initSaramaConsumer() error {
	brokerList, err := cg.storage.getBrokerList()
	if err != nil {
		return err
	}
	if len(brokerList) == 0 {
		return errors.New("no broker alive")
	}
	cg.saramaConsumer, err = sarama.NewConsumer(brokerList, cg.config.SaramaConfig)
	return err
}

// Start would register ConsumerGroup, and rebalance would be triggered.
// ConsumerGroup computes the partitions which should be consumed by consumer's num, and start fetching message.
func (cg *ConsumerGroup) Start() error {
	// exit when failed to register the consumer
	err := cg.storage.registerConsumer(cg.name, cg.id, nil)
	if err != nil && err != zk.ErrNodeExists {
		return err
	}
	cg.wg.Add(1)
	go cg.start()
	return nil
}

// Stop would unregister ConsumerGroup, and rebalance would be triggered.
// The partitions which consumed by this ConsumerGroup would be assigned to others.
func (cg *ConsumerGroup) Stop() {
	cg.stop()
	cg.wg.Wait()
}

// SetLogger use to set the user's logger the consumer group
func (cg *ConsumerGroup) SetLogger(l *logrus.Logger) {
	if l != nil {
		cg.logger = l
	}
}

// IsStopped return whether the ConsumerGroup was stopped or not.
func (cg *ConsumerGroup) IsStopped() bool {
	return cg.state == cgStopped
}

func (cg *ConsumerGroup) triggerRebalance() {
	cg.triggerCh <- restartEvent
}

func (cg *ConsumerGroup) callRecover() {
	if err := recover(); err != nil {
		cg.logger.WithFields(logrus.Fields{
			"group": cg.name,
			"err":   err,
			"stack": string(debug.Stack()),
		}).Error("Recover panic")
		cg.stop()
	}
}

func (cg *ConsumerGroup) start() {
	var wg sync.WaitGroup

	defer cg.callRecover()
	defer func() {
		cg.state = cgStopped
		err := cg.storage.deleteConsumer(cg.name, cg.id)
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to delete consumer from zk")
		}
		for _, tc := range cg.topicConsumers {
			close(tc.messages)
			close(tc.errors)
		}
		cg.wg.Done()
	}()

CONSUME_TOPIC_LOOP:
	for {
		cg.logger.WithField("group", cg.name).Info("Consumer group started")
		cg.stopCh = make(chan struct{})

		err := cg.watchRebalance()
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to watch rebalance")
			cg.stop()
			return
		}
		wg.Add(1)
		go func() {
			defer cg.callRecover()
			defer wg.Done()
			cg.autoReconnect(cg.storage.(*zkGroupStorage).sessionTimeout / 3)
		}()
		for _, consumer := range cg.topicConsumers {
			wg.Add(1)
			consumer.start()
			go func(tc *topicConsumer) {
				defer cg.callRecover()
				defer wg.Done()
				tc.wg.Wait()
				cg.logger.WithFields(logrus.Fields{
					"group": tc.group,
					"topic": tc.name,
				}).Info("Stop the topic consumer")
			}(consumer)
		}
		cg.state = cgStart
		for _, onLoadFunc := range cg.onLoad {
			onLoadFunc()
		}
		msg := <-cg.triggerCh
		for _, onCloseFunc := range cg.onClose {
			onCloseFunc()
		}
		switch msg {
		case restartEvent:
			close(cg.stopCh)
			// The stop channel was used to notify partition's consumer to stop consuming when rebalance is triggered.
			// So we should reinit when rebalance was triggered, as it would be closed.
			wg.Wait()
			continue CONSUME_TOPIC_LOOP
		case quitEvent:
			close(cg.stopCh)
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup is stopping")
			wg.Wait()
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup was stopped")
			return
		}
	}
}

func (cg *ConsumerGroup) stop() {
	cg.stopOnce.Do(func() { cg.triggerCh <- quitEvent })
}

func (cg *ConsumerGroup) getPartitionConsumer(topic string, partition int32, nextOffset int64) (sarama.PartitionConsumer, error) {
	consumer, err := cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	if err == sarama.ErrOffsetOutOfRange {
		nextOffset = cg.config.OffsetAutoReset
		consumer, err = cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	}
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

// GetMessages was used to get a unbuffered message's channel from specified topic
func (cg *ConsumerGroup) GetMessages(topic string) (<-chan *sarama.ConsumerMessage, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.messages, true
	}
	return nil, false
}

// GetErrors was used to get a unbuffered error's channel from specified topic
func (cg *ConsumerGroup) GetErrors(topic string) (<-chan *sarama.ConsumerError, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.errors, true
	}
	return nil, false
}

// OnLoad load callback function that runs after startup
func (cg *ConsumerGroup) OnLoad(cb func()) {
	cg.onLoad = append(cg.onLoad, cb)
}

// OnClose load callback function that runs before the end
func (cg *ConsumerGroup) OnClose(cb func()) {
	cg.onClose = append(cg.onClose, cb)
}

func (cg *ConsumerGroup) autoReconnect(interval time.Duration) {
	timer := time.NewTimer(interval)
	cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was started")
	defer cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was stopped")
	for {
		select {
		case <-cg.stopCh:
			return
		case <-timer.C:
			timer.Reset(interval)
			exist, err := cg.storage.existsConsumer(cg.name, cg.id)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to check consumer existence")
				break
			}
			if exist {
				break
			}
			err = cg.storage.registerConsumer(cg.name, cg.id, nil)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to re-register consumer")
			}
		}
	}
}

func (cg *ConsumerGroup) watchRebalance() error {
	consumerListChange, err := cg.storage.watchConsumerList(cg.name)
	if err != nil {
		return err
	}
	go func() {
		defer cg.callRecover()
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was started")
		select {
		case <-consumerListChange:
			cg.triggerRebalance()
			cg.logger.WithField("group", cg.name).Info("Trigger rebalance while consumers was changed")
		case <-cg.stopCh:
		}
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was exited")
	}()
	return nil
}

// CommitOffset is used to commit offset when auto commit was disabled.
func (cg *ConsumerGroup) CommitOffset(topic string, partition int32, offset int64) error {
	if cg.config.OffsetAutoCommitEnable {
		return errors.New("commit offset take effect when offset auto commit was disabled")
	}
	return cg.storage.commitOffset(cg.name, topic, partition, offset)
}

// GetOffsets return the offset in memory for debug
func (cg *ConsumerGroup) GetOffsets() map[string]interface{} {
	topics := make(map[string]interface{})
	for topic, tc := range cg.topicConsumers {
		topics[topic] = tc.getOffsets()
	}
	return topics
}

// Owners return owners of all partitions
func (cg *ConsumerGroup) Owners() map[string]map[int32]string {
	return cg.owners
}
<commit_msg>FIX: send event to trigger channel more than once<commit_after>package consumergroup

import (
	"errors"
	"fmt"
	"runtime/debug"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/samuel/go-zookeeper/zk"
	"github.com/sirupsen/logrus"
)

const (
	cgInit = iota
	cgStart
	cgStopped
)

const (
	restartEvent = iota
	quitEvent
)

// ConsumerGroup consume message from Kafka with rebalancing supports
type ConsumerGroup struct {
	name           string
	storage        groupStorage
	topicConsumers map[string]*topicConsumer
	saramaConsumer sarama.Consumer

	id          string
	state       int
	wg          sync.WaitGroup
	stopCh      chan struct{}
	triggerCh   chan int
	triggerOnce *sync.Once
	owners      map[string]map[int32]string

	config *Config
	logger *logrus.Logger

	onLoad, onClose []func()
}

// NewConsumerGroup create the ConsumerGroup instance with config
func NewConsumerGroup(config *Config) (*ConsumerGroup, error) {
	if config == nil {
		return nil, errors.New("config can't be empty")
	}
	err := config.validate()
	if err != nil {
		return nil, fmt.Errorf("vaildate config failed, as %s", err)
	}

	cg := new(ConsumerGroup)
	cg.state = cgInit
	cg.config = config
	cg.id = config.ConsumerID
	if cg.id == "" {
		cg.id = genConsumerID()
	}
	cg.name = config.GroupID
	cg.triggerCh = make(chan int)
	cg.topicConsumers = make(map[string]*topicConsumer)
	cg.onLoad = make([]func(), 0)
	cg.onClose = make([]func(), 0)
	cg.storage = newZKGroupStorage(config.ZkList, config.ZkSessionTimeout)
	cg.logger = logrus.New()
	if _, ok := cg.storage.(*zkGroupStorage); ok {
		cg.storage.(*zkGroupStorage).Chroot(config.Chroot)
	}

	err = cg.initSaramaConsumer()
	if err != nil {
		return nil, fmt.Errorf("init sarama consumer, as %s", err)
	}
	cg.owners = make(map[string]map[int32]string)
	for _, topic := range config.TopicList {
		cg.topicConsumers[topic] = newTopicConsumer(cg, topic)
		cg.owners[topic] = make(map[int32]string)
	}
	return cg, nil
}

func (cg *ConsumerGroup) initSaramaConsumer() error {
	brokerList, err := cg.storage.getBrokerList()
	if err != nil {
		return err
	}
	if len(brokerList) == 0 {
		return errors.New("no broker alive")
	}
	cg.saramaConsumer, err = sarama.NewConsumer(brokerList, cg.config.SaramaConfig)
	return err
}

// Start would register ConsumerGroup, and rebalance would be triggered.
// ConsumerGroup computes the partitions which should be consumed by consumer's num, and start fetching message.
func (cg *ConsumerGroup) Start() error {
	// exit when failed to register the consumer
	err := cg.storage.registerConsumer(cg.name, cg.id, nil)
	if err != nil && err != zk.ErrNodeExists {
		return err
	}
	cg.wg.Add(1)
	go cg.start()
	return nil
}

// Stop would unregister ConsumerGroup, and rebalance would be triggered.
// The partitions which consumed by this ConsumerGroup would be assigned to others.
func (cg *ConsumerGroup) Stop() {
	cg.stop()
	cg.wg.Wait()
}

// SetLogger use to set the user's logger the consumer group
func (cg *ConsumerGroup) SetLogger(l *logrus.Logger) {
	if l != nil {
		cg.logger = l
	}
}

// IsStopped return whether the ConsumerGroup was stopped or not.
func (cg *ConsumerGroup) IsStopped() bool {
	return cg.state == cgStopped
}

func (cg *ConsumerGroup) callRecover() {
	if err := recover(); err != nil {
		cg.logger.WithFields(logrus.Fields{
			"group": cg.name,
			"err":   err,
			"stack": string(debug.Stack()),
		}).Error("Recover panic")
		cg.stop()
	}
}

func (cg *ConsumerGroup) start() {
	var wg sync.WaitGroup

	defer cg.callRecover()
	defer func() {
		cg.state = cgStopped
		err := cg.storage.deleteConsumer(cg.name, cg.id)
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to delete consumer from zk")
		}
		for _, tc := range cg.topicConsumers {
			close(tc.messages)
			close(tc.errors)
		}
		cg.wg.Done()
	}()

CONSUME_TOPIC_LOOP:
	for {
		cg.logger.WithField("group", cg.name).Info("Consumer group started")
		cg.triggerOnce = new(sync.Once)
		cg.stopCh = make(chan struct{})

		err := cg.watchRebalance()
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to watch rebalance")
			cg.stop()
			return
		}
		wg.Add(1)
		go func() {
			defer cg.callRecover()
			defer wg.Done()
			cg.autoReconnect(cg.storage.(*zkGroupStorage).sessionTimeout / 3)
		}()
		for _, consumer := range cg.topicConsumers {
			wg.Add(1)
			consumer.start()
			go func(tc *topicConsumer) {
				defer cg.callRecover()
				defer wg.Done()
				tc.wg.Wait()
				cg.logger.WithFields(logrus.Fields{
					"group": tc.group,
					"topic": tc.name,
				}).Info("Stop the topic consumer")
			}(consumer)
		}
		cg.state = cgStart
		for _, onLoadFunc := range cg.onLoad {
			onLoadFunc()
		}
		msg := <-cg.triggerCh
		for _, onCloseFunc := range cg.onClose {
			onCloseFunc()
		}
		switch msg {
		case restartEvent:
			close(cg.stopCh)
			// The stop channel was used to notify partition's consumer to stop consuming when rebalance is triggered.
			// So we should reinit when rebalance was triggered, as it would be closed.
			wg.Wait()
			continue CONSUME_TOPIC_LOOP
		case quitEvent:
			close(cg.stopCh)
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup is stopping")
			wg.Wait()
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup was stopped")
			return
		}
	}
}

func (cg *ConsumerGroup) stop() {
	cg.triggerOnce.Do(func() { cg.triggerCh <- quitEvent })
}

func (cg *ConsumerGroup) triggerRebalance() {
	cg.triggerOnce.Do(func() { cg.triggerCh <- restartEvent })
}

func (cg *ConsumerGroup) getPartitionConsumer(topic string, partition int32, nextOffset int64) (sarama.PartitionConsumer, error) {
	consumer, err := cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	if err == sarama.ErrOffsetOutOfRange {
		nextOffset = cg.config.OffsetAutoReset
		consumer, err = cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	}
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

// GetMessages was used to get a unbuffered message's channel from specified topic
func (cg *ConsumerGroup) GetMessages(topic string) (<-chan *sarama.ConsumerMessage, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.messages, true
	}
	return nil, false
}

// GetErrors was used to get a unbuffered error's channel from specified topic
func (cg *ConsumerGroup) GetErrors(topic string) (<-chan *sarama.ConsumerError, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.errors, true
	}
	return nil, false
}

// OnLoad load callback function that runs after startup
func (cg *ConsumerGroup) OnLoad(cb func()) {
	cg.onLoad = append(cg.onLoad, cb)
}

// OnClose load callback function that runs before the end
func (cg *ConsumerGroup) OnClose(cb func()) {
	cg.onClose = append(cg.onClose, cb)
}

func (cg *ConsumerGroup) autoReconnect(interval time.Duration) {
	timer := time.NewTimer(interval)
	cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was started")
	defer cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was stopped")
	for {
		select {
		case <-cg.stopCh:
			return
		case <-timer.C:
			timer.Reset(interval)
			exist, err := cg.storage.existsConsumer(cg.name, cg.id)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to check consumer existence")
				break
			}
			if exist {
				break
			}
			err = cg.storage.registerConsumer(cg.name, cg.id, nil)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to re-register consumer")
			}
		}
	}
}

func (cg *ConsumerGroup) watchRebalance() error {
	consumerListChange, err := cg.storage.watchConsumerList(cg.name)
	if err != nil {
		return err
	}
	go func() {
		defer cg.callRecover()
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was started")
		select {
		case <-consumerListChange:
			cg.triggerRebalance()
			cg.logger.WithField("group", cg.name).Info("Trigger rebalance while consumers was changed")
		case <-cg.stopCh:
		}
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was exited")
	}()
	return nil
}

// CommitOffset is used to commit offset when auto commit was disabled.
func (cg *ConsumerGroup) CommitOffset(topic string, partition int32, offset int64) error {
	if cg.config.OffsetAutoCommitEnable {
		return errors.New("commit offset take effect when offset auto commit was disabled")
	}
	return cg.storage.commitOffset(cg.name, topic, partition, offset)
}

// GetOffsets return the offset in memory for debug
func (cg *ConsumerGroup) GetOffsets() map[string]interface{} {
	topics := make(map[string]interface{})
	for topic, tc := range cg.topicConsumers {
		topics[topic] = tc.getOffsets()
	}
	return topics
}

// Owners return owners of all partitions
func (cg *ConsumerGroup) Owners() map[string]map[int32]string {
	return cg.owners
}
<|endoftext|>


<commit_before>// Copyright 2015,2016,2017,2018,2019 SeukWon Kang (kasworld@gmail.com)
// 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 basiclog

import (
	"bytes"
	"fmt"
	"math/bits"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/kasworld/log/logdestination_file"
	"github.com/kasworld/log/logdestination_stdio"
	"github.com/kasworld/log/logdestinationgroup"
	"github.com/kasworld/log/logdestinationi"
	"github.com/kasworld/log/logflagi"
	"github.com/kasworld/log/logflags"
)

type LL_Type uint64

func (ll LL_Type) ToShiftedNum() int {
	return bits.Len(uint(ll)) - 1
}

func (ll LL_Type) LevelsString() string {
	var buff bytes.Buffer

	buff.WriteString("LL_Type[")
	for i := LL_Type(1); i < LL_END; i <<= 1 {
		if ll.IsLevel(i) {
			fmt.Fprintf(&buff, "%s, ", i)
		}
	}
	buff.WriteString("]")
	return buff.String()
}

func (ll LL_Type) IsLevel(level LL_Type) bool {
	return ll&level != 0
}

func (i LL_Type) String() string {
	if str, ok := leveldata[i]; ok {
		return str
	}
	return "LL_Type(" + strconv.FormatInt(int64(i), 10) + ")"
}

///////////////////////////////////////////////////////////////////////////////////////////

var (
	OutputStdout = logdestination_stdio.NewStdOut()
	OutputStderr = logdestination_stdio.NewStdErr()
)

type logDestInfo struct {
	refCntByLogLv int // count referenced by each loglv
	dest          logdestinationi.LogDestinationI
}

type LogBase struct {
	mutex sync.RWMutex

	flag     logflagi.LogFlagI // properties
	prefix   string            // prefix to write at beginning of each line
	loglevel LL_Type

	ltype2destgrp     []*logdestinationgroup.LogDestinationGroup
	allDestInfoByName map[string]*logDestInfo
}

func New(prefix string, lf logflagi.LogFlagI, lv LL_Type) *LogBase {

	maxlen := LL_END.ToShiftedNum()
	dstgrp := make([]*logdestinationgroup.LogDestinationGroup, maxlen)
	for i := 0; i < maxlen; i++ {
		dstgrp[i] = logdestinationgroup.New()
	}

	return &LogBase{
		ltype2destgrp:     dstgrp,
		allDestInfoByName: make(map[string]*logDestInfo),
		flag:              lf,
		prefix:            prefix,
		loglevel:          lv,
	}
}

func makeLogFilename(logdir string, ll string) string {
	basename := filepath.Base(logdir)
	filename := fmt.Sprintf("%s.%s.%s", basename, ll, "log")
	return filepath.Join(logdir, filename)
}

func NewWithDstDir(prefix string, logdir string, lf logflagi.LogFlagI,
	loglevel LL_Type, splitLogLevel LL_Type) (*LogBase, error) {
	logdir = strings.TrimSpace(logdir)
	if logdir == "" {
		return nil, fmt.Errorf("logdir empty")
	}
	if err := os.MkdirAll(logdir, os.ModePerm); err != nil {
		return nil, err
	}
	newlg := New(prefix, lf, loglevel)
	newDestForOther, err := logdestination_file.New(
		makeLogFilename(logdir, "Other"))
	if err != nil {
		return nil, err
	}
	newlg.AddDestination(LL_All^splitLogLevel, newDestForOther)
	for ll := LL_Type(1); ll < LL_END; ll <<= 1 {
		if splitLogLevel&ll == ll {
			newDestForLL, serr := logdestination_file.New(
				makeLogFilename(logdir, ll.String()))
			if serr != nil {
				return nil, serr
			}
			newlg.AddDestination(ll, newDestForLL)
		}
	}
	newlg.AddDestination(LL_Fatal, OutputStderr)
	return newlg, nil
}

func (lg *LogBase) AddDestination(
	ll LL_Type, o logdestinationi.LogDestinationI) {

	lg.mutex.Lock()
	defer lg.mutex.Unlock()

	for i := 0; i < len(lg.ltype2destgrp); i++ {
		s := LL_Type(1 << uint(i))
		if ll&s == 0 {
			continue
		}
		lg.addDestination1DestGrp(i, o)
	}
}

func (lg *LogBase) addDestination1DestGrp(
	i int, o logdestinationi.LogDestinationI) {

	added := lg.ltype2destgrp[i].AddDestination(o)
	if !added {
		if _, ok := lg.allDestInfoByName[o.Name()]; !ok {
			panic(fmt.Sprintf(
				"%v failed to AddDestination to destgroup index:%v, abnormal state",
				lg,
				i))
		}
		fmt.Printf("%v not added to destgroup index:%v\n", o, i)
		return
	}

	if dstinfo, ok := lg.allDestInfoByName[o.Name()]; ok {
		dstinfo.refCntByLogLv++
	} else {
		lg.allDestInfoByName[o.Name()] = &logDestInfo{
			refCntByLogLv: 1,
			dest:          o,
		}
	}
}

func (lg *LogBase) DelDestination(
	ll LL_Type, o logdestinationi.LogDestinationI) {

	lg.mutex.Lock()
	defer lg.mutex.Unlock()

	for i := 0; i < len(lg.ltype2destgrp); i++ {
		s := LL_Type(1 << uint(i))
		if ll&s == 0 {
			continue
		}
		lg.delDestinationFrom1DestGrp(i, o)
	}
}

func (lg *LogBase) delDestinationFrom1DestGrp(
	i int, o logdestinationi.LogDestinationI) {

	deleted := lg.ltype2destgrp[i].DelDestination(o)
	if !deleted {
		fmt.Printf("%v not deleted from destgroup index:%v\n", o, i)
		return
	}

	if dstinfo, ok := lg.allDestInfoByName[o.Name()]; ok {
		dstinfo.refCntByLogLv--
		if dstinfo.refCntByLogLv <= 0 {
			delete(lg.allDestInfoByName, o.Name())
		}
	} else {
		panic(fmt.Sprintf(
			"%v failed to DelDestination %v from destgroup index:%v, abnormal state",
			lg,
			o,
			i,
		))
	}
}

func (lg *LogBase) Reload() error {
	lg.mutex.RLock()
	defer lg.mutex.RUnlock()

	for _, v := range lg.allDestInfoByName {
		if err := v.dest.Reload(); err != nil {
			fmt.Println(err)
		}
	}
	return nil
}

func (lg *LogBase) LogPrintf(
	calldepth int, ll LL_Type,
	format string, v ...interface{}) ([]byte, error) {
	s := lg.Format2Bytes(calldepth+1, ll, format, v...)
	err := lg.Output(ll, s)
	return s, err
}

func (lg *LogBase) Format2Bytes(
	calldepth int, ll LL_Type,
	format string, v ...interface{}) []byte {

	if !lg.IsLevel(ll) {
		return nil
	}
	s := fmt.Sprintf(format, v...)

	var buf []byte
	llinfo := fmt.Sprintf("%s", ll)
	lg.flag.FormatHeader(&buf, calldepth+2, time.Now(), lg.prefix, llinfo)
	buf = append(buf, s...)
	if len(s) == 0 || s[len(s)-1] != '\n' {
		buf = append(buf, '\n')
	}
	return buf
}

func (lg *LogBase) Output(ll LL_Type, b []byte) error {
	i := ll.ToShiftedNum()
	return lg.ltype2destgrp[i].Write(b)
}

func (lg *LogBase) Panic(format string, v ...interface{}) error {
	s, err := lg.LogPrintf(2, LL_Fatal, format, v...)
	panic(string(s))
	return err
}

func (lg *LogBase) AddLevel(level LL_Type) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.loglevel |= level
}

func (lg *LogBase) SetLevel(level LL_Type) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.loglevel = level
}

func (lg *LogBase) DelLevel(level LL_Type) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.loglevel &= ^level
}

func (lg *LogBase) IsLevel(level LL_Type) bool {
	return lg.loglevel&level != 0
}

func (lg *LogBase) FlagString() string {
	return lg.flag.FlagString()
}

func (lg *LogBase) LevelString() string {
	return lg.loglevel.LevelsString()
}

func (lg *LogBase) SetPrefix(p string) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.prefix = p
}

// Prefix returns the output prefix for the logger.
func (lg *LogBase) GetPrefix() string {
	return lg.prefix
}

// Flags returns the output flags for the logger.
func (lg *LogBase) GetFlags() logflagi.LogFlagI {
	return lg.flag
}

// SetFlags sets the output flags for the logger.
func (lg *LogBase) SetFlags(flag logflagi.LogFlagI) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.flag = flag
}

func (lg *LogBase) String() string {
	return fmt.Sprintf("LogBase[%v %v]",
		lg.FlagString(), lg.LevelString(),
	)
}

//////////////////////////////////////////////////////////////////

var GlobalLogger = New("", logflags.DefaultValue(false), LL_All)

func init() {
	GlobalLogger.AddDestination(LL_All, OutputStderr)
}

func Reload() error {
	return GlobalLogger.Reload()
}

func GetLogger() *LogBase {
	return GlobalLogger
}
func SetLogger(l *LogBase) {
	GlobalLogger = l
}

func LevelString() string {
	return GlobalLogger.String()
}

func AddLevel(level LL_Type) {
	GlobalLogger.AddLevel(level)
}

func SetLevel(level LL_Type) {
	GlobalLogger.SetLevel(level)
}

func DelLevel(level LL_Type) {
	GlobalLogger.DelLevel(level)
}

func IsLevel(level LL_Type) bool {
	return GlobalLogger.IsLevel(level)
}

func SetPrefix(p string) {
	GlobalLogger.SetPrefix(p)
}

// Prefix returns the output prefix for the GlobalLogger.
func GetPrefix() string {
	return GlobalLogger.GetPrefix()
}

// Flags returns the output flags for the GlobalLogger.
func GetFlags() logflagi.LogFlagI {
	return GlobalLogger.GetFlags()
}

// SetFlags sets the output flags for the GlobalLogger.
func SetFlags(flag logflagi.LogFlagI) {
	GlobalLogger.SetFlags(flag)
}

func Panic(format string, v ...interface{}) error {
	s, err := GlobalLogger.LogPrintf(2, LL_Fatal, format, v...)
	panic(string(s))
	return err
}

func SetLog(logdir string, loglevel LL_Type, splitLogLevel LL_Type) error {
	if logdir != "" {
		newlg, err := NewWithDstDir(
			"",
			logdir,
			logflags.DefaultValue(false),
			loglevel,
			splitLogLevel,
		)
		if err != nil {
			return err
		}
		SetLogger(newlg)
	} else {
		SetLevel(loglevel)
	}
	return nil
}
<commit_msg>update logbase<commit_after>// Copyright 2015,2016,2017,2018,2019 SeukWon Kang (kasworld@gmail.com)
// 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 basiclog

import (
	"bytes"
	"fmt"
	"math/bits"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/kasworld/log/logdestination_file"
	"github.com/kasworld/log/logdestination_stdio"
	"github.com/kasworld/log/logdestinationgroup"
	"github.com/kasworld/log/logdestinationi"
	"github.com/kasworld/log/logflagi"
	"github.com/kasworld/log/logflags"
)

type LL_Type uint64

func (ll LL_Type) ToShiftedNum() int {
	return bits.Len(uint(ll)) - 1
}

func (ll LL_Type) LevelsString() string {
	var buff bytes.Buffer

	buff.WriteString("LL_Type[")
	for i := LL_Type(1); i < LL_END; i <<= 1 {
		if ll.IsLevel(i) {
			fmt.Fprintf(&buff, "%s, ", i)
		}
	}
	buff.WriteString("]")
	return buff.String()
}

func (i LL_Type) String() string {
	if str, ok := leveldata[i]; ok {
		return str
	}
	return "LL_Type(" + strconv.FormatInt(int64(i), 10) + ")"
}

func (ll LL_Type) IsLevel(level LL_Type) bool {
	return ll&level != 0
}

func (ll LL_Type) AllLevel() LL_Type {
	return LL_All
}

func (ll LL_Type) StartLevel() LL_Type {
	return LL_Fatal
}

func (ll LL_Type) IsLastLevel() bool {
	return ll == LL_END
}

func (ll LL_Type) NextLevel(n uint) LL_Type {
	return ll << n
}

func (ll LL_Type) PreLevel(n uint) LL_Type {
	return ll >> n
}

func (ll LL_Type) BitAnd(l2 LL_Type) LL_Type {
	return ll & l2
}

func (ll LL_Type) BitOr(l2 LL_Type) LL_Type {
	return ll | l2
}

func (ll LL_Type) BitXor(l2 LL_Type) LL_Type {
	return ll ^ l2
}

func (ll LL_Type) BitClear(l2 LL_Type) LL_Type {
	return ll &^ l2
}

func (ll LL_Type) BitTest(l2 LL_Type) bool {
	return ll&l2 == 0
}

///////////////////////////////////////////////////////////////////////////////////////////

var (
	OutputStdout = logdestination_stdio.NewStdOut()
	OutputStderr = logdestination_stdio.NewStdErr()
)

type logDestInfo struct {
	refCntByLogLv int // count referenced by each loglv
	dest          logdestinationi.LogDestinationI
}

type LogBase struct {
	mutex sync.RWMutex

	flag     logflagi.LogFlagI // properties
	prefix   string            // prefix to write at beginning of each line
	loglevel LL_Type

	ltype2destgrp     []*logdestinationgroup.LogDestinationGroup
	allDestInfoByName map[string]*logDestInfo
}

func New(prefix string, lf logflagi.LogFlagI, lv LL_Type) *LogBase {

	maxlen := LL_END.ToShiftedNum()
	dstgrp := make([]*logdestinationgroup.LogDestinationGroup, maxlen)
	for i := 0; i < maxlen; i++ {
		dstgrp[i] = logdestinationgroup.New()
	}

	return &LogBase{
		ltype2destgrp:     dstgrp,
		allDestInfoByName: make(map[string]*logDestInfo),
		flag:              lf,
		prefix:            prefix,
		loglevel:          lv,
	}
}

func makeLogFilename(logdir string, ll string) string {
	basename := filepath.Base(logdir)
	filename := fmt.Sprintf("%s.%s.%s", basename, ll, "log")
	return filepath.Join(logdir, filename)
}

func NewWithDstDir(prefix string, logdir string, lf logflagi.LogFlagI,
	loglevel LL_Type, splitLogLevel LL_Type) (*LogBase, error) {
	logdir = strings.TrimSpace(logdir)
	if logdir == "" {
		return nil, fmt.Errorf("logdir empty")
	}
	if err := os.MkdirAll(logdir, os.ModePerm); err != nil {
		return nil, err
	}
	newlg := New(prefix, lf, loglevel)
	newDestForOther, err := logdestination_file.New(
		makeLogFilename(logdir, "Other"))
	if err != nil {
		return nil, err
	}
	newlg.AddDestination(LL_All^splitLogLevel, newDestForOther)
	for ll := LL_Type(1); !ll.IsLastLevel(); ll = ll.NextLevel(1) {
		if splitLogLevel.IsLevel(ll) {
			newDestForLL, serr := logdestination_file.New(
				makeLogFilename(logdir, ll.String()))
			if serr != nil {
				return nil, serr
			}
			newlg.AddDestination(ll, newDestForLL)
		}
	}
	newlg.AddDestination(LL_Fatal, OutputStderr)
	return newlg, nil
}

func (lg *LogBase) AddDestination(
	ll LL_Type, o logdestinationi.LogDestinationI) {

	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	i := 0
	for l2 := LL_Type(1); !l2.IsLastLevel(); l2 = l2.NextLevel(1) {
		if ll.IsLevel(l2) {
			lg.addDestination1DestGrp(i, o)
		}
		i++
	}
}

func (lg *LogBase) addDestination1DestGrp(
	i int, o logdestinationi.LogDestinationI) {

	added := lg.ltype2destgrp[i].AddDestination(o)
	if !added {
		if _, ok := lg.allDestInfoByName[o.Name()]; !ok {
			panic(fmt.Sprintf(
				"%v failed to AddDestination to destgroup index:%v, abnormal state",
				lg, i))
		}
		fmt.Printf("%v not added to destgroup index:%v\n", o, i)
		return
	}

	if dstinfo, ok := lg.allDestInfoByName[o.Name()]; ok {
		dstinfo.refCntByLogLv++
	} else {
		lg.allDestInfoByName[o.Name()] = &logDestInfo{
			refCntByLogLv: 1,
			dest:          o,
		}
	}
}

func (lg *LogBase) DelDestination(
	ll LL_Type, o logdestinationi.LogDestinationI) {

	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	i := 0
	for l2 := LL_Type(1); !l2.IsLastLevel(); l2 = l2.NextLevel(1) {
		if ll.IsLevel(l2) {
			lg.delDestinationFrom1DestGrp(i, o)
		}
		i++
	}
}

func (lg *LogBase) delDestinationFrom1DestGrp(
	i int, o logdestinationi.LogDestinationI) {

	deleted := lg.ltype2destgrp[i].DelDestination(o)
	if !deleted {
		fmt.Printf("%v not deleted from destgroup index:%v\n", o, i)
		return
	}

	if dstinfo, ok := lg.allDestInfoByName[o.Name()]; ok {
		dstinfo.refCntByLogLv--
		if dstinfo.refCntByLogLv <= 0 {
			delete(lg.allDestInfoByName, o.Name())
		}
	} else {
		panic(fmt.Sprintf(
			"%v failed to DelDestination %v from destgroup index:%v, abnormal state",
			lg, o, i,
		))
	}
}

func (lg *LogBase) Reload() error {
	lg.mutex.RLock()
	defer lg.mutex.RUnlock()

	for _, v := range lg.allDestInfoByName {
		if err := v.dest.Reload(); err != nil {
			fmt.Println(err)
		}
	}
	return nil
}

func (lg *LogBase) LogPrintf(
	calldepth int, ll LL_Type,
	format string, v ...interface{}) ([]byte, error) {
	s := lg.Format2Bytes(calldepth+1, ll, format, v...)
	err := lg.Output(ll, s)
	return s, err
}

func (lg *LogBase) Format2Bytes(
	calldepth int, ll LL_Type,
	format string, v ...interface{}) []byte {

	if !lg.IsLevel(ll) {
		return nil
	}
	s := fmt.Sprintf(format, v...)

	var buf []byte
	llinfo := fmt.Sprintf("%s", ll)
	lg.flag.FormatHeader(&buf, calldepth+2, time.Now(), lg.prefix, llinfo)
	buf = append(buf, s...)
	if len(s) == 0 || s[len(s)-1] != '\n' {
		buf = append(buf, '\n')
	}
	return buf
}

func (lg *LogBase) Output(ll LL_Type, b []byte) error {
	i := ll.ToShiftedNum()
	return lg.ltype2destgrp[i].Write(b)
}

func (lg *LogBase) Panic(format string, v ...interface{}) error {
	s, err := lg.LogPrintf(2, LL_Fatal, format, v...)
	panic(string(s))
	return err
}

func (lg *LogBase) AddLevel(level LL_Type) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.loglevel = lg.loglevel.BitOr(level)
}

func (lg *LogBase) SetLevel(level LL_Type) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.loglevel = level
}

func (lg *LogBase) DelLevel(level LL_Type) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.loglevel = lg.loglevel.BitClear(level)
}

func (lg *LogBase) IsLevel(level LL_Type) bool {
	return lg.loglevel.IsLevel(level)
}

func (lg *LogBase) FlagString() string {
	return lg.flag.FlagString()
}

func (lg *LogBase) LevelString() string {
	return lg.loglevel.LevelsString()
}

func (lg *LogBase) SetPrefix(p string) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.prefix = p
}

// Prefix returns the output prefix for the logger.
func (lg *LogBase) GetPrefix() string {
	return lg.prefix
}

// Flags returns the output flags for the logger.
func (lg *LogBase) GetFlags() logflagi.LogFlagI {
	return lg.flag
}

// SetFlags sets the output flags for the logger.
func (lg *LogBase) SetFlags(flag logflagi.LogFlagI) {
	lg.mutex.Lock()
	defer lg.mutex.Unlock()
	lg.flag = flag
}

func (lg *LogBase) String() string {
	return fmt.Sprintf("LogBase[%v %v]",
		lg.FlagString(), lg.LevelString(),
	)
}

//////////////////////////////////////////////////////////////////

var GlobalLogger = New("", logflags.DefaultValue(false), LL_All)

func init() {
	GlobalLogger.AddDestination(LL_All, OutputStderr)
}

func Reload() error {
	return GlobalLogger.Reload()
}

func GetLogger() *LogBase {
	return GlobalLogger
}
func SetLogger(l *LogBase) {
	GlobalLogger = l
}

func LevelString() string {
	return GlobalLogger.String()
}

func AddLevel(level LL_Type) {
	GlobalLogger.AddLevel(level)
}

func SetLevel(level LL_Type) {
	GlobalLogger.SetLevel(level)
}

func DelLevel(level LL_Type) {
	GlobalLogger.DelLevel(level)
}

func IsLevel(level LL_Type) bool {
	return GlobalLogger.IsLevel(level)
}

func SetPrefix(p string) {
	GlobalLogger.SetPrefix(p)
}

// Prefix returns the output prefix for the GlobalLogger.
func GetPrefix() string {
	return GlobalLogger.GetPrefix()
}

// Flags returns the output flags for the GlobalLogger.
func GetFlags() logflagi.LogFlagI {
	return GlobalLogger.GetFlags()
}

// SetFlags sets the output flags for the GlobalLogger.
func SetFlags(flag logflagi.LogFlagI) {
	GlobalLogger.SetFlags(flag)
}

func Panic(format string, v ...interface{}) error {
	s, err := GlobalLogger.LogPrintf(2, LL_Fatal, format, v...)
	panic(string(s))
	return err
}

func SetLog(logdir string, loglevel LL_Type, splitLogLevel LL_Type) error {
	if logdir != "" {
		newlg, err := NewWithDstDir(
			"",
			logdir,
			logflags.DefaultValue(false),
			loglevel,
			splitLogLevel,
		)
		if err != nil {
			return err
		}
		SetLogger(newlg)
	} else {
		SetLevel(loglevel)
	}
	return nil
}
<|endoftext|>


<commit_before>package geoip

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func startTestServer() *httptest.Server {
	s := make(chan bool)
	var ts *httptest.Server
	go func() {
		ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			p := r.URL.Path
			switch p {
			case "/2a02:2770::21a:4aff:feb3:2ee":
				response := `{"ip":"2a02:2770::21a:4aff:feb3:2ee",
							"country_code":"NL",
							"country_name":"Netherlands",
							"region_code":"",
							"region_name":"",
							"city":"",
							"zip_code":"",
							"time_zone":"Europe/Amsterdam",
							"latitude":52.25,
							"longitude":5.75,
							"metro_code":0}`
				fmt.Fprintf(w, response)
			case "/66.102.15.255":
				response := `{"ip":"66.102.15.255",
							"country_code":"US",
							"country_name":"United States",
							"region_code":"CA",
							"region_name":"California",
							"city":"Mountain View",
							"zip_code":"94043",
							"time_zone":"America/Los_Angeles",
							"latitude":37.4192,
							"longitude":-122.0574,
							"metro_code":807}`
				fmt.Fprintf(w, response)
			case "/":
				response := `{"ip":"217.140.98.70",
							"country_code":"GB",
							"country_name":"United Kingdom",
							"region_code":"ENG",
							"region_name":"England",
							"city":"Saint Neots",
							"zip_code":"CB5",
							"time_zone":"Europe/London",
							"longitude":0.1167,
							"latitude":52.2,
							"metro_code":0}`
				fmt.Fprintf(w, response)
			case "/github.com":
				response := `{"ip":"192.30.252.130",
							"country_code":"US",
							"country_name":"United States",
							"region_code":"CA",
							"region_name":"California",
							"city":"San Francisco",
							"zip_code":"94107",
							"time_zone":"America/Los_Angeles",
							"latitude":37.7697,
							"longitude":-122.3933,
							"metro_code":807}`
				fmt.Fprintf(w, response)
			case "/abcxyz":
				response := `404 page not found`
				fmt.Fprintf(w, response)
			}
		}))
		baseURI = ts.URL + "/"
		s <- true
	}()
	_ = <-s
	return ts
}

// TestLocationWithIP6 test location method using a IPv6 address.
func TestLocationWithIP6(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "2a02:2770::21a:4aff:feb3:2ee"

	want := map[string]string{
		"ip":           "2a02:2770::21a:4aff:feb3:2ee",
		"country_code": "NL",
		"country_name": "Netherlands",
		"region_code":  "",
		"region_name":  "",
		"city":         "",
		"zip_code":     "",
		"time_zone":    "Europe/Amsterdam",
		"latitude":     "52.25",
		"longitude":    "5.75",
		"metro_code":   "0",
	}

	got, err := Location(input)

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}

}

// TestLocationWithIP6 test location method using a IPv4 address.
func TestLocationWithIP4(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "66.102.15.255"

	want := map[string]string{
		"ip":           "66.102.15.255",
		"country_code": "US",
		"country_name": "United States",
		"region_code":  "CA",
		"region_name":  "California",
		"city":         "Mountain View",
		"zip_code":     "94043",
		"time_zone":    "America/Los_Angeles",
		"latitude":     "37.4192",
		"longitude":    "-122.0574",
		"metro_code":   "807",
	}

	got, err := Location(input)

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}
}

// TestLocationWithIP6 test location method using a blank IP address.
func TestLocationWithNoIPProvided(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	want := map[string]string{
		"ip":           "217.140.98.70",
		"country_code": "GB",
		"country_name": "United Kingdom",
		"region_code":  "ENG",
		"region_name":  "England",
		"city":         "Saint Neots",
		"zip_code":     "CB5",
		"time_zone":    "Europe/London",
		"longitude":    "0.1167",
		"latitude":     "52.2",
		"metro_code":   "0",
	}

	got, err := Location("")

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}
}

// TestLocationWithDomain test location with a domain i.e. github.com
func TestLocationWithDomain(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "github.com"

	want := map[string]string{
		"ip":           "192.30.252.130",
		"country_code": "US",
		"country_name": "United States",
		"region_code":  "CA",
		"region_name":  "California",
		"city":         "San Francisco",
		"zip_code":     "94107",
		"time_zone":    "America/Los_Angeles",
		"latitude":     "37.7697",
		"longitude":    "-122.3933",
		"metro_code":   "807",
	}

	got, err := Location(input)

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}

}

// TestLocation404 test location with bad host or ip
func TestLocation404(t *testing.T) {

}
<commit_msg>Add test to handle 404 response'<commit_after>package geoip

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func startTestServer() *httptest.Server {
	s := make(chan bool)
	var ts *httptest.Server
	go func() {
		ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			p := r.URL.Path
			switch p {
			case "/2a02:2770::21a:4aff:feb3:2ee":
				response := `{"ip":"2a02:2770::21a:4aff:feb3:2ee",
							"country_code":"NL",
							"country_name":"Netherlands",
							"region_code":"",
							"region_name":"",
							"city":"",
							"zip_code":"",
							"time_zone":"Europe/Amsterdam",
							"latitude":52.25,
							"longitude":5.75,
							"metro_code":0}`
				fmt.Fprintf(w, response)
			case "/66.102.15.255":
				response := `{"ip":"66.102.15.255",
							"country_code":"US",
							"country_name":"United States",
							"region_code":"CA",
							"region_name":"California",
							"city":"Mountain View",
							"zip_code":"94043",
							"time_zone":"America/Los_Angeles",
							"latitude":37.4192,
							"longitude":-122.0574,
							"metro_code":807}`
				fmt.Fprintf(w, response)
			case "/":
				response := `{"ip":"217.140.98.70",
							"country_code":"GB",
							"country_name":"United Kingdom",
							"region_code":"ENG",
							"region_name":"England",
							"city":"Saint Neots",
							"zip_code":"CB5",
							"time_zone":"Europe/London",
							"longitude":0.1167,
							"latitude":52.2,
							"metro_code":0}`
				fmt.Fprintf(w, response)
			case "/github.com":
				response := `{"ip":"192.30.252.130",
							"country_code":"US",
							"country_name":"United States",
							"region_code":"CA",
							"region_name":"California",
							"city":"San Francisco",
							"zip_code":"94107",
							"time_zone":"America/Los_Angeles",
							"latitude":37.7697,
							"longitude":-122.3933,
							"metro_code":807}`
				fmt.Fprintf(w, response)
			case "/abcxyz":
				response := `404 page not found`
				w.WriteHeader(http.StatusNotFound)
				fmt.Fprintf(w, response)
			}
		}))
		baseURI = ts.URL + "/"
		s <- true
	}()
	_ = <-s
	return ts
}

// TestLocationWithIP6 test location method using a IPv6 address.
func TestLocationWithIP6(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "2a02:2770::21a:4aff:feb3:2ee"

	want := map[string]string{
		"ip":           "2a02:2770::21a:4aff:feb3:2ee",
		"country_code": "NL",
		"country_name": "Netherlands",
		"region_code":  "",
		"region_name":  "",
		"city":         "",
		"zip_code":     "",
		"time_zone":    "Europe/Amsterdam",
		"latitude":     "52.25",
		"longitude":    "5.75",
		"metro_code":   "0",
	}

	got, err := Location(input)

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}

}

// TestLocationWithIP6 test location method using a IPv4 address.
func TestLocationWithIP4(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "66.102.15.255"

	want := map[string]string{
		"ip":           "66.102.15.255",
		"country_code": "US",
		"country_name": "United States",
		"region_code":  "CA",
		"region_name":  "California",
		"city":         "Mountain View",
		"zip_code":     "94043",
		"time_zone":    "America/Los_Angeles",
		"latitude":     "37.4192",
		"longitude":    "-122.0574",
		"metro_code":   "807",
	}

	got, err := Location(input)

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}
}

// TestLocationWithIP6 test location method using a blank IP address.
func TestLocationWithNoIPProvided(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	want := map[string]string{
		"ip":           "217.140.98.70",
		"country_code": "GB",
		"country_name": "United Kingdom",
		"region_code":  "ENG",
		"region_name":  "England",
		"city":         "Saint Neots",
		"zip_code":     "CB5",
		"time_zone":    "Europe/London",
		"longitude":    "0.1167",
		"latitude":     "52.2",
		"metro_code":   "0",
	}

	got, err := Location("")

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}
}

// TestLocationWithDomain test location with a domain i.e. github.com
func TestLocationWithDomain(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "github.com"

	want := map[string]string{
		"ip":           "192.30.252.130",
		"country_code": "US",
		"country_name": "United States",
		"region_code":  "CA",
		"region_name":  "California",
		"city":         "San Francisco",
		"zip_code":     "94107",
		"time_zone":    "America/Los_Angeles",
		"latitude":     "37.7697",
		"longitude":    "-122.3933",
		"metro_code":   "807",
	}

	got, err := Location(input)

	if err != nil {
		t.Errorf("Location error: method returned error %s", err)
	}

	for k, v := range got {
		if want[k] != v {
			t.Errorf("Location data: got %s - %s want %s - %s", k, v, k, want[k])
		}
	}

}

// TestLocation404 test location with bad host or ip
func TestLocation404(t *testing.T) {
	ts := startTestServer()
	defer ts.Close()

	input := "abcxyz"

	got, err := Location(input)

	if err != ErrReq {
		t.Errorf("Location error: method returned error %s", err)
	}

	if got != nil {
		t.Errorf("Location data: got %v want %v", got, nil)
	}
}
<|endoftext|>