<commit_before>// Copyright 2019 The Cloud Robotics 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.

// The CR syncer syncs custom resources between a remote Kubernetes cluster
// and the local Kubernetes cluster. The spec part is copied from upstream to
// downstream, and the status part is copied from downstream to upstream.
//
// The behaviour can be customized by annotations on the CRDs.
// cr-syncer.cloudrobotics.com/filter-by-robot-name: <bool>
//   If true, only sync CRs that have a label 'cloudrobotics.com/robot-name: <robot-name>'
//   that matches the robot-name arg given on the command line.
// cr-syncer.cloudrobotics.com/status-subtree: <string>
//   If specified, only sync the given subtree of the Status field. This is useful
//   if resources have a shared status.
// cr-syncer.cloudrobotics.com/spec-source: <string>
//   If set to "cloud", the source of truth for object existence and specs (upstream) is
//   the remote cluster and for status it's local (downstream). If set to "robot", the roles
//   are reversed.
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/motemen/go-loghttp"
	"go.opencensus.io/exporter/prometheus"
	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/tag"
	"go.opencensus.io/zpages"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	crdtypes "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	crdclientset "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	crdinformer "k8s.io/apiextensions-apiserver/pkg/client/informers/externalversions"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/klog"
)

const (
	// Resync informers every 5 minutes. This will cause all current resources
	// to be sent as updates once again, which will trigger reconciliation on those
	// objects and thus fix any potential drift.
	resyncPeriod = 5 * time.Minute
)

var (
	remoteServer = flag.String("remote-server", "", "Remote Kubernetes server")
	robotName    = flag.String("robot-name", "", "Robot we are running on, can be used for selective syncing")
	verbose      = flag.Bool("verbose", false, "Enable verbose logging")
	listenAddr   = flag.String("listen-address", ":80", "HTTP listen address")

	sizeDistribution    = view.Distribution(0, 1024, 2048, 4096, 16384, 65536, 262144, 1048576, 4194304, 33554432)
	latencyDistribution = view.Distribution(0, 1, 2, 5, 10, 15, 25, 50, 100, 200, 400, 800, 1500, 3000, 6000)

	tagLocation = mustNewTagKey("location")
)

func init() {
	if err := view.Register(
		&view.View{
			Name:        ochttp.ClientRequestCount.Name(),
			Description: ochttp.ClientRequestCount.Description(),
			Measure:     ochttp.ClientRequestCount,
			TagKeys:     []tag.Key{ochttp.Method, tagLocation},
			Aggregation: view.Count(),
		},
		&view.View{
			Name:        ochttp.ClientRequestBytes.Name(),
			Description: ochttp.ClientRequestBytes.Description(),
			Measure:     ochttp.ClientRequestBytes,
			TagKeys:     []tag.Key{ochttp.Method, ochttp.StatusCode, tagLocation},
			Aggregation: sizeDistribution,
		},
		&view.View{
			Name:        ochttp.ClientResponseBytes.Name(),
			Description: ochttp.ClientResponseBytes.Description(),
			Measure:     ochttp.ClientResponseBytes,
			TagKeys:     []tag.Key{ochttp.Method, ochttp.StatusCode, tagLocation},
			Aggregation: sizeDistribution,
		},
		&view.View{
			Name:        ochttp.ClientLatency.Name(),
			Description: ochttp.ClientLatency.Description(),
			Measure:     ochttp.ClientLatency,
			TagKeys:     []tag.Key{ochttp.Method, ochttp.StatusCode, tagLocation},
			Aggregation: latencyDistribution,
		},
	); err != nil {
		panic(err)
	}
}

// PrefixingRoundtripper is a HTTP roundtripper that adds a specified prefix to
// all HTTP requests. We need to use it instead of setting APIPath because
// autogenerated and dynamic Kubernetes clients overwrite the REST config's
// APIPath.
type PrefixingRoundtripper struct {
	Prefix string
	Base   http.RoundTripper
}

func (pr *PrefixingRoundtripper) RoundTrip(r *http.Request) (*http.Response, error) {
	// Avoid an extra roundtrip for the protocol upgrade
	r.URL.Scheme = "https"
	if !strings.HasPrefix(r.URL.Path, pr.Prefix+"/") {
		r.URL.Path = pr.Prefix + r.URL.Path
	}
	resp, err := pr.Base.RoundTrip(r)
	return resp, err
}

// ctxRoundTripper injects a fixed context into all requests. This is used to
// provide static OpenCensus tags as Kubernetes' client-go provides no context hooks.
type ctxRoundTripper struct {
	base http.RoundTripper
	ctx  context.Context
}

func (r *ctxRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	return r.base.RoundTrip(req.WithContext(r.ctx))
}

// restConfigForRemote assembles the K8s REST config for the remote server.
func restConfigForRemote(ctx context.Context) (*rest.Config, error) {
	tokenSource, err := google.DefaultTokenSource(ctx, "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		return nil, err
	}
	ctx, err = tag.New(ctx, tag.Insert(tagLocation, "remote"))
	if err != nil {
		return nil, err
	}
	transport := func(base http.RoundTripper) (rt http.RoundTripper) {
		rt = &oauth2.Transport{
			Source: tokenSource,
			Base:   base,
		}
		rt = &PrefixingRoundtripper{
			Prefix: "/apis/core.kubernetes",
			Base:   rt,
		}
		if *verbose {
			rt = &loghttp.Transport{Transport: rt}
		}
		rt = &ochttp.Transport{Base: rt}
		return &ctxRoundTripper{base: rt, ctx: ctx}
	}
	return &rest.Config{
		Host:          *remoteServer,
		APIPath:       "/apis",
		WrapTransport: transport,
	}, nil
}

type CrdChange struct {
	Type watch.EventType
	CRD  *crdtypes.CustomResourceDefinition
}

func streamCrds(done <-chan struct{}, clientset crdclientset.Interface, crds chan<- CrdChange) error {
	factory := crdinformer.NewSharedInformerFactory(clientset, 0)
	informer := factory.Apiextensions().V1beta1().CustomResourceDefinitions().Informer()

	go informer.Run(done)

	log.Printf("Syncing cache for CRDs")
	ok := cache.WaitForCacheSync(done, informer.HasSynced)
	if !ok {
		return fmt.Errorf("WaitForCacheSync failed")
	}

	informer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			crds <- CrdChange{Type: watch.Added, CRD: obj.(*crdtypes.CustomResourceDefinition)}
		},
		UpdateFunc: func(oldObj, newObj interface{}) {
			crds <- CrdChange{Type: watch.Modified, CRD: newObj.(*crdtypes.CustomResourceDefinition)}
		},
		DeleteFunc: func(obj interface{}) {
			crds <- CrdChange{Type: watch.Deleted, CRD: obj.(*crdtypes.CustomResourceDefinition)}
		},
	})
	return nil
}

func main() {
	klog.InitFlags(nil)
	flag.Parse()
	ctx := context.Background()

	localConfig, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal(err)
	}
	localCtx, err := tag.New(ctx, tag.Insert(tagLocation, "local"))
	if err != nil {
		log.Fatal(err)
	}
	localConfig.WrapTransport = func(base http.RoundTripper) http.RoundTripper {
		if *verbose {
			base = &loghttp.Transport{Transport: base}
		}
		base = &ochttp.Transport{Base: base}
		return &ctxRoundTripper{base: base, ctx: localCtx}
	}
	local, err := dynamic.NewForConfig(localConfig)
	if err != nil {
		log.Fatal(err)
	}
	remoteConfig, err := restConfigForRemote(ctx)
	if err != nil {
		log.Fatal(err)
	}
	remote, err := dynamic.NewForConfig(remoteConfig)
	if err != nil {
		log.Fatal(err)
	}

	exporter, err := prometheus.NewExporter(prometheus.Options{})
	if err != nil {
		log.Fatal(err)
	}
	view.RegisterExporter(exporter)
	view.SetReportingPeriod(time.Second)
	zpages.Handle(nil, "/debug")
	http.Handle("/metrics", exporter)

	go func() {
		if err := http.ListenAndServe(*listenAddr, nil); err != nil {
			log.Fatalln(err)
		}
	}()

	crds := make(chan CrdChange)
	if err := streamCrds(ctx.Done(), crdclientset.NewForConfigOrDie(localConfig), crds); err != nil {
		log.Fatalf("Unable to stream CRDs from local Kubernetes: %v", err)
	}
	syncers := make(map[string]*crSyncer)
	for crd := range crds {
		name := crd.CRD.GetName()

		if cur, ok := syncers[name]; ok {
			if crd.Type == watch.Added {
				log.Printf("Warning: Already had a running sync for freshly added %s", name)
			}
			cur.stop()
			delete(syncers, name)
		}
		if crd.Type == watch.Added || crd.Type == watch.Modified {
			// The modify procedure is very heavyweight: We throw away
			// the informer for the CRD (read: all cached data) on every
			// modification and recreate it. If that ever turns out to
			// be a problem, we should use a shared informer cache
			// instead.
			s, err := newCRSyncer(*crd.CRD, local, remote, *robotName)
			if err != nil {
				log.Printf("skipping custom resource %s: %s", name, err)
				continue
			}
			syncers[name] = s
			go s.run()
		}
	}
}

func mustNewTagKey(s string) tag.Key {
	k, err := tag.NewKey(s)
	if err != nil {
		panic(err)
	}
	return k
}
<commit_msg>Update the package godoc for cr-syncer.<commit_after>// Copyright 2019 The Cloud Robotics 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.

// The CR syncer syncs custom resources between a remote Kubernetes cluster
// and the local Kubernetes cluster. The spec part is copied from upstream to
// downstream, and the status part is copied from downstream to upstream.
//
// The behaviour can be customized by annotations on the CRDs.
//
// cr-syncer.cloudrobotics.com/filter-by-robot-name: <bool>
//
// If true, only sync CRs that have a label 'cloudrobotics.com/robot-name: <robot-name>'
// that matches the robot-name arg given on the command line.
//
// cr-syncer.cloudrobotics.com/status-subtree: <string>
//
// If specified, only sync the given subtree of the Status field. This is useful
// if resources have a shared status.
//
// cr-syncer.cloudrobotics.com/spec-source: <string>
//
// If set to "cloud", the source of truth for object existence and specs (upstream) is
// the remote cluster and for status it's local (downstream). If set to "robot", the roles
// are reversed. Otherwise, eg when using the empty string "", synchronization is disabled.
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/motemen/go-loghttp"
	"go.opencensus.io/exporter/prometheus"
	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/tag"
	"go.opencensus.io/zpages"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	crdtypes "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	crdclientset "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	crdinformer "k8s.io/apiextensions-apiserver/pkg/client/informers/externalversions"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/klog"
)

const (
	// Resync informers every 5 minutes. This will cause all current resources
	// to be sent as updates once again, which will trigger reconciliation on those
	// objects and thus fix any potential drift.
	resyncPeriod = 5 * time.Minute
)

var (
	remoteServer = flag.String("remote-server", "", "Remote Kubernetes server")
	robotName    = flag.String("robot-name", "", "Robot we are running on, can be used for selective syncing")
	verbose      = flag.Bool("verbose", false, "Enable verbose logging")
	listenAddr   = flag.String("listen-address", ":80", "HTTP listen address")

	sizeDistribution    = view.Distribution(0, 1024, 2048, 4096, 16384, 65536, 262144, 1048576, 4194304, 33554432)
	latencyDistribution = view.Distribution(0, 1, 2, 5, 10, 15, 25, 50, 100, 200, 400, 800, 1500, 3000, 6000)

	tagLocation = mustNewTagKey("location")
)

func init() {
	if err := view.Register(
		&view.View{
			Name:        ochttp.ClientRequestCount.Name(),
			Description: ochttp.ClientRequestCount.Description(),
			Measure:     ochttp.ClientRequestCount,
			TagKeys:     []tag.Key{ochttp.Method, tagLocation},
			Aggregation: view.Count(),
		},
		&view.View{
			Name:        ochttp.ClientRequestBytes.Name(),
			Description: ochttp.ClientRequestBytes.Description(),
			Measure:     ochttp.ClientRequestBytes,
			TagKeys:     []tag.Key{ochttp.Method, ochttp.StatusCode, tagLocation},
			Aggregation: sizeDistribution,
		},
		&view.View{
			Name:        ochttp.ClientResponseBytes.Name(),
			Description: ochttp.ClientResponseBytes.Description(),
			Measure:     ochttp.ClientResponseBytes,
			TagKeys:     []tag.Key{ochttp.Method, ochttp.StatusCode, tagLocation},
			Aggregation: sizeDistribution,
		},
		&view.View{
			Name:        ochttp.ClientLatency.Name(),
			Description: ochttp.ClientLatency.Description(),
			Measure:     ochttp.ClientLatency,
			TagKeys:     []tag.Key{ochttp.Method, ochttp.StatusCode, tagLocation},
			Aggregation: latencyDistribution,
		},
	); err != nil {
		panic(err)
	}
}

// PrefixingRoundtripper is a HTTP roundtripper that adds a specified prefix to
// all HTTP requests. We need to use it instead of setting APIPath because
// autogenerated and dynamic Kubernetes clients overwrite the REST config's
// APIPath.
type PrefixingRoundtripper struct {
	Prefix string
	Base   http.RoundTripper
}

func (pr *PrefixingRoundtripper) RoundTrip(r *http.Request) (*http.Response, error) {
	// Avoid an extra roundtrip for the protocol upgrade
	r.URL.Scheme = "https"
	if !strings.HasPrefix(r.URL.Path, pr.Prefix+"/") {
		r.URL.Path = pr.Prefix + r.URL.Path
	}
	resp, err := pr.Base.RoundTrip(r)
	return resp, err
}

// ctxRoundTripper injects a fixed context into all requests. This is used to
// provide static OpenCensus tags as Kubernetes' client-go provides no context hooks.
type ctxRoundTripper struct {
	base http.RoundTripper
	ctx  context.Context
}

func (r *ctxRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	return r.base.RoundTrip(req.WithContext(r.ctx))
}

// restConfigForRemote assembles the K8s REST config for the remote server.
func restConfigForRemote(ctx context.Context) (*rest.Config, error) {
	tokenSource, err := google.DefaultTokenSource(ctx, "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		return nil, err
	}
	ctx, err = tag.New(ctx, tag.Insert(tagLocation, "remote"))
	if err != nil {
		return nil, err
	}
	transport := func(base http.RoundTripper) (rt http.RoundTripper) {
		rt = &oauth2.Transport{
			Source: tokenSource,
			Base:   base,
		}
		rt = &PrefixingRoundtripper{
			Prefix: "/apis/core.kubernetes",
			Base:   rt,
		}
		if *verbose {
			rt = &loghttp.Transport{Transport: rt}
		}
		rt = &ochttp.Transport{Base: rt}
		return &ctxRoundTripper{base: rt, ctx: ctx}
	}
	return &rest.Config{
		Host:          *remoteServer,
		APIPath:       "/apis",
		WrapTransport: transport,
	}, nil
}

type CrdChange struct {
	Type watch.EventType
	CRD  *crdtypes.CustomResourceDefinition
}

func streamCrds(done <-chan struct{}, clientset crdclientset.Interface, crds chan<- CrdChange) error {
	factory := crdinformer.NewSharedInformerFactory(clientset, 0)
	informer := factory.Apiextensions().V1beta1().CustomResourceDefinitions().Informer()

	go informer.Run(done)

	log.Printf("Syncing cache for CRDs")
	ok := cache.WaitForCacheSync(done, informer.HasSynced)
	if !ok {
		return fmt.Errorf("WaitForCacheSync failed")
	}

	informer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			crds <- CrdChange{Type: watch.Added, CRD: obj.(*crdtypes.CustomResourceDefinition)}
		},
		UpdateFunc: func(oldObj, newObj interface{}) {
			crds <- CrdChange{Type: watch.Modified, CRD: newObj.(*crdtypes.CustomResourceDefinition)}
		},
		DeleteFunc: func(obj interface{}) {
			crds <- CrdChange{Type: watch.Deleted, CRD: obj.(*crdtypes.CustomResourceDefinition)}
		},
	})
	return nil
}

func main() {
	klog.InitFlags(nil)
	flag.Parse()
	ctx := context.Background()

	localConfig, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal(err)
	}
	localCtx, err := tag.New(ctx, tag.Insert(tagLocation, "local"))
	if err != nil {
		log.Fatal(err)
	}
	localConfig.WrapTransport = func(base http.RoundTripper) http.RoundTripper {
		if *verbose {
			base = &loghttp.Transport{Transport: base}
		}
		base = &ochttp.Transport{Base: base}
		return &ctxRoundTripper{base: base, ctx: localCtx}
	}
	local, err := dynamic.NewForConfig(localConfig)
	if err != nil {
		log.Fatal(err)
	}
	remoteConfig, err := restConfigForRemote(ctx)
	if err != nil {
		log.Fatal(err)
	}
	remote, err := dynamic.NewForConfig(remoteConfig)
	if err != nil {
		log.Fatal(err)
	}

	exporter, err := prometheus.NewExporter(prometheus.Options{})
	if err != nil {
		log.Fatal(err)
	}
	view.RegisterExporter(exporter)
	view.SetReportingPeriod(time.Second)
	zpages.Handle(nil, "/debug")
	http.Handle("/metrics", exporter)

	go func() {
		if err := http.ListenAndServe(*listenAddr, nil); err != nil {
			log.Fatalln(err)
		}
	}()

	crds := make(chan CrdChange)
	if err := streamCrds(ctx.Done(), crdclientset.NewForConfigOrDie(localConfig), crds); err != nil {
		log.Fatalf("Unable to stream CRDs from local Kubernetes: %v", err)
	}
	syncers := make(map[string]*crSyncer)
	for crd := range crds {
		name := crd.CRD.GetName()

		if cur, ok := syncers[name]; ok {
			if crd.Type == watch.Added {
				log.Printf("Warning: Already had a running sync for freshly added %s", name)
			}
			cur.stop()
			delete(syncers, name)
		}
		if crd.Type == watch.Added || crd.Type == watch.Modified {
			// The modify procedure is very heavyweight: We throw away
			// the informer for the CRD (read: all cached data) on every
			// modification and recreate it. If that ever turns out to
			// be a problem, we should use a shared informer cache
			// instead.
			s, err := newCRSyncer(*crd.CRD, local, remote, *robotName)
			if err != nil {
				log.Printf("skipping custom resource %s: %s", name, err)
				continue
			}
			syncers[name] = s
			go s.run()
		}
	}
}

func mustNewTagKey(s string) tag.Key {
	k, err := tag.NewKey(s)
	if err != nil {
		panic(err)
	}
	return k
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Google Inc.
 *
 * 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 stackdriver

import (
	"expvar"
	"fmt"
	"time"

	"cloud.google.com/go/compute/metadata"
	"github.com/cloudfoundry-community/stackdriver-tools/src/stackdriver-nozzle/telemetry"
	"github.com/cloudfoundry/lager"
	"github.com/golang/protobuf/ptypes/timestamp"
	labelpb "google.golang.org/genproto/googleapis/api/label"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	"google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

type telemetrySink struct {
	projectPath string
	labels      map[string]string
	resource    *monitoredres.MonitoredResource
	logger      lager.Logger
	client      MetricClient
	startTime   *timestamp.Timestamp
}

func now() *timestamp.Timestamp {
	now := time.Now()
	return &timestamp.Timestamp{
		Seconds: now.Unix(),
		Nanos:   int32(now.Nanosecond()),
	}
}

func detectMonitoredResource() (res *monitoredres.MonitoredResource) {
	res = &monitoredres.MonitoredResource{Type: "global"}

	if metadata.OnGCE() {
		projectId, err := metadata.ProjectID()
		if err != nil {
			return
		}
		instanceId, err := metadata.InstanceID()
		if err != nil {
			return
		}
		zone, err := metadata.Zone()
		if err != nil {
			return
		}

		res.Type = "gce_instance"
		res.Labels = map[string]string{"project_id": projectId, "instance_id": instanceId, "zone": zone}
	}
	return
}

// NewTelemetrySink provides a telemetry.Sink that writes metrics to Stackdriver Monitoring
func NewTelemetrySink(logger lager.Logger, client MetricClient, projectID, subscriptionId, foundation string) telemetry.Sink {
	return &telemetrySink{
		logger:      logger,
		client:      client,
		projectPath: fmt.Sprintf("projects/%s", projectID),
		labels:      map[string]string{"subscription_id": subscriptionId, "foundation": foundation},
		startTime:   now(),
		resource:    detectMonitoredResource()}
}

func (ts *telemetrySink) Init(registeredSeries []*expvar.KeyValue) {
	req := &monitoringpb.ListMetricDescriptorsRequest{
		Name:   ts.projectPath,
		Filter: fmt.Sprintf(`metric.type = starts_with("stackdriver-nozzle")`),
	}

	descriptors, err := ts.client.ListMetricDescriptors(req)
	if err != nil {
		ts.logger.Error("telemetrySink.ListMetricDescriptors", err, lager.Data{"req": req})
	}

	registered := map[string]bool{}
	for _, descriptor := range descriptors {
		registered[descriptor.Name] = true
	}

	for _, series := range registeredSeries {
		name := ts.metricDescriptorName(series.Key)
		if registered[name] {
			continue
		}

		labels := []*labelpb.LabelDescriptor{}
		for name := range ts.labels {
			labels = append(labels, &labelpb.LabelDescriptor{Key: name, ValueType: labelpb.LabelDescriptor_STRING})
		}

		if mapVal, ok := series.Value.(*telemetry.CounterMap); ok {
			labels = append(labels, &labelpb.LabelDescriptor{Key: mapVal.Category(), ValueType: labelpb.LabelDescriptor_STRING})
		}

		req := &monitoringpb.CreateMetricDescriptorRequest{
			Name: ts.projectPath,
			MetricDescriptor: &metricpb.MetricDescriptor{
				DisplayName: series.Key,
				Name:        name,
				Type:        ts.metricDescriptorType(series.Key),
				Labels:      labels,
				MetricKind:  metricpb.MetricDescriptor_CUMULATIVE,
				ValueType:   metricpb.MetricDescriptor_INT64,
				Description: "stackdriver-nozzle created custom metric.",
			},
		}
		err := ts.client.CreateMetricDescriptor(req)

		if err != nil {
			ts.logger.Error("telemetrySink.CreateMetricDescriptor", err, lager.Data{"req": req})
		}
	}
}

func (ts *telemetrySink) metricDescriptorName(key string) string {
	return fmt.Sprintf("%s/metricDescriptors/%s", ts.projectPath, ts.metricDescriptorType(key))
}

func (ts *telemetrySink) metricDescriptorType(key string) string {
	return fmt.Sprintf("custom.googleapis.com/%s", key)
}

const maxTimeSeries = 200

func (ts *telemetrySink) newRequest() *monitoringpb.CreateTimeSeriesRequest {
	return &monitoringpb.CreateTimeSeriesRequest{
		Name: ts.projectPath,
	}
}

func (ts *telemetrySink) Report(report []*expvar.KeyValue) {
	req := ts.newRequest()

	interval := &monitoringpb.TimeInterval{
		StartTime: ts.startTime,
		EndTime:   now(),
	}

	for _, data := range report {
		req.TimeSeries = append(req.TimeSeries, ts.timeSeries(ts.metricDescriptorType(data.Key), interval, data)...)

		if len(req.TimeSeries) == maxTimeSeries {
			ts.client.Post(req)
			req = ts.newRequest()
		}
	}

	if len(req.TimeSeries) != 0 {
		ts.client.Post(req)
	}
}

func (ts *telemetrySink) timeSeries(metricType string, interval *monitoringpb.TimeInterval, val *expvar.KeyValue) []*monitoringpb.TimeSeries {
	switch data := val.Value.(type) {
	case *telemetry.Counter:
		return []*monitoringpb.TimeSeries{ts.timeSeriesInt(metricType, interval, ts.labels, data.Value())}
	case *telemetry.CounterMap:
		series := []*monitoringpb.TimeSeries{}
		data.Do(func(value expvar.KeyValue) {
			if intVal, ok := value.Value.(*telemetry.Counter); ok {
				labels := duplicate(ts.labels)
				labels[data.Category()] = value.Key
				series = append(series, ts.timeSeriesInt(metricType, interval, labels, intVal.Value()))
			}
		})
		return series
	default:
		ts.logger.Error("telemetrySink.timeSeries", fmt.Errorf("unknown value type: %T", val), lager.Data{"value": val})
	}

	return nil
}

func duplicate(src map[string]string) map[string]string {
	dest := map[string]string{}
	for k, v := range src {
		dest[k] = v
	}
	return dest
}

func (ts *telemetrySink) timeSeriesInt(metricType string, interval *monitoringpb.TimeInterval, labels map[string]string, value int64) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type:   metricType,
			Labels: labels,
		},
		Points: []*monitoringpb.Point{{
			Interval: interval,
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{Int64Value: value},
			},
		}},
		Resource: ts.resource,
	}
}
<commit_msg>Provide ValueType and MetricKind so descriptor autovivication works correctly.<commit_after>/*
 * Copyright 2017 Google Inc.
 *
 * 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 stackdriver

import (
	"expvar"
	"fmt"
	"time"

	"cloud.google.com/go/compute/metadata"
	"github.com/cloudfoundry-community/stackdriver-tools/src/stackdriver-nozzle/telemetry"
	"github.com/cloudfoundry/lager"
	"github.com/golang/protobuf/ptypes/timestamp"
	labelpb "google.golang.org/genproto/googleapis/api/label"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	"google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

type telemetrySink struct {
	projectPath string
	labels      map[string]string
	resource    *monitoredres.MonitoredResource
	logger      lager.Logger
	client      MetricClient
	startTime   *timestamp.Timestamp
}

func now() *timestamp.Timestamp {
	now := time.Now()
	return &timestamp.Timestamp{
		Seconds: now.Unix(),
		Nanos:   int32(now.Nanosecond()),
	}
}

func detectMonitoredResource() (res *monitoredres.MonitoredResource) {
	res = &monitoredres.MonitoredResource{Type: "global"}

	if metadata.OnGCE() {
		projectId, err := metadata.ProjectID()
		if err != nil {
			return
		}
		instanceId, err := metadata.InstanceID()
		if err != nil {
			return
		}
		zone, err := metadata.Zone()
		if err != nil {
			return
		}

		res.Type = "gce_instance"
		res.Labels = map[string]string{"project_id": projectId, "instance_id": instanceId, "zone": zone}
	}
	return
}

// NewTelemetrySink provides a telemetry.Sink that writes metrics to Stackdriver Monitoring
func NewTelemetrySink(logger lager.Logger, client MetricClient, projectID, subscriptionId, foundation string) telemetry.Sink {
	return &telemetrySink{
		logger:      logger,
		client:      client,
		projectPath: fmt.Sprintf("projects/%s", projectID),
		labels:      map[string]string{"subscription_id": subscriptionId, "foundation": foundation},
		startTime:   now(),
		resource:    detectMonitoredResource()}
}

func (ts *telemetrySink) Init(registeredSeries []*expvar.KeyValue) {
	req := &monitoringpb.ListMetricDescriptorsRequest{
		Name:   ts.projectPath,
		Filter: fmt.Sprintf(`metric.type = starts_with("stackdriver-nozzle")`),
	}

	descriptors, err := ts.client.ListMetricDescriptors(req)
	if err != nil {
		ts.logger.Error("telemetrySink.ListMetricDescriptors", err, lager.Data{"req": req})
	}

	registered := map[string]bool{}
	for _, descriptor := range descriptors {
		registered[descriptor.Name] = true
	}

	for _, series := range registeredSeries {
		name := ts.metricDescriptorName(series.Key)
		if registered[name] {
			continue
		}

		labels := []*labelpb.LabelDescriptor{}
		for name := range ts.labels {
			labels = append(labels, &labelpb.LabelDescriptor{Key: name, ValueType: labelpb.LabelDescriptor_STRING})
		}

		if mapVal, ok := series.Value.(*telemetry.CounterMap); ok {
			labels = append(labels, &labelpb.LabelDescriptor{Key: mapVal.Category(), ValueType: labelpb.LabelDescriptor_STRING})
		}

		req := &monitoringpb.CreateMetricDescriptorRequest{
			Name: ts.projectPath,
			MetricDescriptor: &metricpb.MetricDescriptor{
				DisplayName: series.Key,
				Name:        name,
				Type:        ts.metricDescriptorType(series.Key),
				Labels:      labels,
				MetricKind:  metricpb.MetricDescriptor_CUMULATIVE,
				ValueType:   metricpb.MetricDescriptor_INT64,
				Description: "stackdriver-nozzle created custom metric.",
			},
		}
		err := ts.client.CreateMetricDescriptor(req)

		if err != nil {
			ts.logger.Error("telemetrySink.CreateMetricDescriptor", err, lager.Data{"req": req})
		}
	}
}

func (ts *telemetrySink) metricDescriptorName(key string) string {
	return fmt.Sprintf("%s/metricDescriptors/%s", ts.projectPath, ts.metricDescriptorType(key))
}

func (ts *telemetrySink) metricDescriptorType(key string) string {
	return fmt.Sprintf("custom.googleapis.com/%s", key)
}

const maxTimeSeries = 200

func (ts *telemetrySink) newRequest() *monitoringpb.CreateTimeSeriesRequest {
	return &monitoringpb.CreateTimeSeriesRequest{
		Name: ts.projectPath,
	}
}

func (ts *telemetrySink) Report(report []*expvar.KeyValue) {
	req := ts.newRequest()

	interval := &monitoringpb.TimeInterval{
		StartTime: ts.startTime,
		EndTime:   now(),
	}

	for _, data := range report {
		req.TimeSeries = append(req.TimeSeries, ts.timeSeries(ts.metricDescriptorType(data.Key), interval, data)...)

		if len(req.TimeSeries) == maxTimeSeries {
			ts.client.Post(req)
			req = ts.newRequest()
		}
	}

	if len(req.TimeSeries) != 0 {
		ts.client.Post(req)
	}
}

func (ts *telemetrySink) timeSeries(metricType string, interval *monitoringpb.TimeInterval, val *expvar.KeyValue) []*monitoringpb.TimeSeries {
	switch data := val.Value.(type) {
	case *telemetry.Counter:
		return []*monitoringpb.TimeSeries{ts.timeSeriesInt(metricType, interval, ts.labels, data.Value())}
	case *telemetry.CounterMap:
		series := []*monitoringpb.TimeSeries{}
		data.Do(func(value expvar.KeyValue) {
			if intVal, ok := value.Value.(*telemetry.Counter); ok {
				labels := duplicate(ts.labels)
				labels[data.Category()] = value.Key
				series = append(series, ts.timeSeriesInt(metricType, interval, labels, intVal.Value()))
			}
		})
		return series
	default:
		ts.logger.Error("telemetrySink.timeSeries", fmt.Errorf("unknown value type: %T", val), lager.Data{"value": val})
	}

	return nil
}

func duplicate(src map[string]string) map[string]string {
	dest := map[string]string{}
	for k, v := range src {
		dest[k] = v
	}
	return dest
}

func (ts *telemetrySink) timeSeriesInt(metricType string, interval *monitoringpb.TimeInterval, labels map[string]string, value int64) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		MetricKind: metricpb.MetricDescriptor_CUMULATIVE,
		ValueType:  metricpb.MetricDescriptor_INT64,
		Metric: &metricpb.Metric{
			Type:   metricType,
			Labels: labels,
		},
		Points: []*monitoringpb.Point{{
			Interval: interval,
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{Int64Value: value},
			},
		}},
		Resource: ts.resource,
	}
}
<|endoftext|>


<commit_before>package softlayer

import (
	datatypes "github.com/TheWeatherCompany/softlayer-go/data_types"
)

type SoftLayer_Dns_Domain_Record_Service interface {
	Service

	CreateObject(template datatypes.SoftLayer_Dns_Domain_Record_Template) (datatypes.SoftLayer_Dns_Domain_Record, error)
	CreateObjects(template []datatypes.SoftLayer_Dns_Domain_Record_Template) ([]datatypes.SoftLayer_Dns_Domain_Record, error)
	GetObject(recordId int) (datatypes.SoftLayer_Dns_Domain_Record, error)
	DeleteObject(recordId int) (bool, error)
	EditObject(recordId int, template datatypes.SoftLayer_Dns_Domain_Record) (bool, error)
}
<commit_msg>#8 removed unnecessary method<commit_after>package softlayer

import (
	datatypes "github.com/TheWeatherCompany/softlayer-go/data_types"
)

type SoftLayer_Dns_Domain_Record_Service interface {
	Service

	CreateObject(template datatypes.SoftLayer_Dns_Domain_Record_Template) (datatypes.SoftLayer_Dns_Domain_Record, error)
	GetObject(recordId int) (datatypes.SoftLayer_Dns_Domain_Record, error)
	DeleteObject(recordId int) (bool, error)
	EditObject(recordId int, template datatypes.SoftLayer_Dns_Domain_Record) (bool, error)
}
<|endoftext|>


<commit_before>package receiver

import (
	"fmt"
	"net"
	"net/url"

	"github.com/lomik/go-carbon/points"
)

type Receiver interface {
	Stop()
	Name() string
	Stat(func(metric string, value float64))
}

type Option func(Receiver) error

// PickleMaxMessageSize creates option for New contructor
func PickleMaxMessageSize(size uint32) Option {
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.maxPickleMessageSize = size
		}
		return nil
	}
}

// UDPLogIncomplete creates option for New contructor
func UDPLogIncomplete(enable bool) Option {
	return func(r Receiver) error {
		if t, ok := r.(*UDP); ok {
			t.logIncomplete = enable
		}
		return nil
	}
}

// OutChan creates option for New contructor
func OutChan(ch chan *points.Points) Option {
	out := func(p *points.Points) {
		ch <- p
	}
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.out = out
		}
		if t, ok := r.(*UDP); ok {
			t.out = out
		}
		return nil
	}
}

// OutFunc creates option for New contructor
func OutFunc(out func(*points.Points)) Option {
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.out = out
		}
		if t, ok := r.(*UDP); ok {
			t.out = out
		}
		return nil
	}
}

// Name creates option for New contructor
func Name(name string) Option {
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.name = name
		}
		if t, ok := r.(*UDP); ok {
			t.name = name
		}
		return nil
	}
}

func blackhole(p *points.Points) {}

// New creates udp, tcp, pickle receiver
func New(dsn string, opts ...Option) (Receiver, error) {
	u, err := url.Parse(dsn)
	if err != nil {
		return nil, err
	}

	if u.Scheme == "tcp" || u.Scheme == "pickle" {
		addr, err := net.ResolveTCPAddr("tcp", u.Host)
		if err != nil {
			return nil, err
		}

		r := &TCP{
			out:  blackhole,
			name: u.Scheme,
		}

		if u.Scheme == "pickle" {
			r.isPickle = true
			r.maxPickleMessageSize = 67108864 // 64Mb
		}

		for _, optApply := range opts {
			optApply(r)
		}

		if err = r.Listen(addr); err != nil {
			return nil, err
		}

		return r, err
	}

	if u.Scheme == "udp" {
		addr, err := net.ResolveUDPAddr("udp", u.Host)
		if err != nil {
			return nil, err
		}

		r := &UDP{
			out:  blackhole,
			name: u.Scheme,
		}

		for _, optApply := range opts {
			optApply(r)
		}

		err = r.Listen(addr)
		if err != nil {
			return nil, err
		}

		return r, err
	}

	return nil, fmt.Errorf("unknown proto %#v", u.Scheme)
}
<commit_msg>use OutFunc for OutChan option<commit_after>package receiver

import (
	"fmt"
	"net"
	"net/url"

	"github.com/lomik/go-carbon/points"
)

type Receiver interface {
	Stop()
	Name() string
	Stat(func(metric string, value float64))
}

type Option func(Receiver) error

// PickleMaxMessageSize creates option for New contructor
func PickleMaxMessageSize(size uint32) Option {
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.maxPickleMessageSize = size
		}
		return nil
	}
}

// UDPLogIncomplete creates option for New contructor
func UDPLogIncomplete(enable bool) Option {
	return func(r Receiver) error {
		if t, ok := r.(*UDP); ok {
			t.logIncomplete = enable
		}
		return nil
	}
}

// OutChan creates option for New contructor
func OutChan(ch chan *points.Points) Option {
	return OutFunc(func(p *points.Points) {
		ch <- p
	})
}

// OutFunc creates option for New contructor
func OutFunc(out func(*points.Points)) Option {
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.out = out
		}
		if t, ok := r.(*UDP); ok {
			t.out = out
		}
		return nil
	}
}

// Name creates option for New contructor
func Name(name string) Option {
	return func(r Receiver) error {
		if t, ok := r.(*TCP); ok {
			t.name = name
		}
		if t, ok := r.(*UDP); ok {
			t.name = name
		}
		return nil
	}
}

func blackhole(p *points.Points) {}

// New creates udp, tcp, pickle receiver
func New(dsn string, opts ...Option) (Receiver, error) {
	u, err := url.Parse(dsn)
	if err != nil {
		return nil, err
	}

	if u.Scheme == "tcp" || u.Scheme == "pickle" {
		addr, err := net.ResolveTCPAddr("tcp", u.Host)
		if err != nil {
			return nil, err
		}

		r := &TCP{
			out:  blackhole,
			name: u.Scheme,
		}

		if u.Scheme == "pickle" {
			r.isPickle = true
			r.maxPickleMessageSize = 67108864 // 64Mb
		}

		for _, optApply := range opts {
			optApply(r)
		}

		if err = r.Listen(addr); err != nil {
			return nil, err
		}

		return r, err
	}

	if u.Scheme == "udp" {
		addr, err := net.ResolveUDPAddr("udp", u.Host)
		if err != nil {
			return nil, err
		}

		r := &UDP{
			out:  blackhole,
			name: u.Scheme,
		}

		for _, optApply := range opts {
			optApply(r)
		}

		err = r.Listen(addr)
		if err != nil {
			return nil, err
		}

		return r, err
	}

	return nil, fmt.Errorf("unknown proto %#v", u.Scheme)
}
<|endoftext|>


<commit_before>package server

import (
    "time"
    "github.com/itsankoff/gotcha/common"
    "strconv"
)

type Group struct {
    Id              string
    Out             chan *common.Message
    outputs         map[string]chan<- *common.Message
}

func NewGroup() *Group {
    group := &Group{
        Id: strconv.FormatInt(time.Now().Unix(), 10),
        outputs: make(map[string]chan<- *common.Message),
    }

    go group.listen()
    return group
}

func (g *Group) listen() {
    for {
        select {
        case msg := <-g.Out:
            for _, out := range g.outputs {
                out<- msg
            }
        }
    }
}

func (g *Group) AddOutput(id string, output chan<- *common.Message) bool{
    g.outputs[id] = output
    return true
}

func (g *Group) RemoveOutput(id string) bool {
    _, ok := g.outputs[id]
    if ok {
        delete(g.outputs, id)
        return true
    }

    return false
}
<commit_msg>Initialize group Out channel<commit_after>package server

import (
    "time"
    "github.com/itsankoff/gotcha/common"
    "strconv"
)

type Group struct {
    Id              string
    Out             chan *common.Message
    outputs         map[string]chan<- *common.Message
}

func NewGroup() *Group {
    group := &Group{
        Id: strconv.FormatInt(time.Now().Unix(), 10),
        Out: make(chan *common.Message),
        outputs: make(map[string]chan<- *common.Message),
    }

    go group.listen()
    return group
}

func (g *Group) listen() {
    for {
        select {
        case msg := <-g.Out:
            for _, out := range g.outputs {
                out<- msg
            }
        }
    }
}

func (g *Group) AddOutput(id string, output chan<- *common.Message) bool{
    g.outputs[id] = output
    return true
}

func (g *Group) RemoveOutput(id string) bool {
    _, ok := g.outputs[id]
    if ok {
        delete(g.outputs, id)
        return true
    }

    return false
}
<|endoftext|>


<commit_before>/**
 * Copyright 2017 Yannick Roffin
 *
 *   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 teleinfo_service

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/spf13/viper"
	log "github.com/yroffin/jarvis-go-ext/logger"
)

// On linux
// apply: stty 1200 cs7 evenp cstopb -igncr -inlcr -brkint -icrnl -opost -isig -icanon -iexten -F /dev/ttyUSB0
// to configure teleinfo tty
// Cf. https://hallard.me/gestion-de-la-teleinfo-avec-un-raspberry-pi-et-une-carte-arduipi/

// Teleinfo : instance Teleinfo device struct
type TeleinfoService struct {
	Entries map[string]string
}

// Trame EDF
type TeleinfoTrame struct {
	etiquette string // ETIQUETTE (4 à 8 caractères)
	data      string // DATA (1 à 12 caractères)
	hecksum   string // CHECKSUM (caractère de contrôle : Somme de caractère)
}

var instance *TeleinfoService
var once sync.Once
var mutex = &sync.Mutex{}

// GetInstance : singleton instance
func Service() *TeleinfoService {
	once.Do(func() {
		instance = new(TeleinfoService)
		instance.init()
	})
	return instance
}

// declare canal
var canal = make(chan byte, 1024)

// handleReadFile : read file
func handleReadFile(device string) error {

	s, err := os.OpenFile(device, syscall.O_RDONLY, 0666)

	if err != nil {
		log.Default.Error("teleinfo", log.Fields{
			"error": err,
		})
	}

	log.Default.Info("teleinfo", log.Fields{
		"device": device,
	})

	buffer := make([]byte, 4096)
	reader := bufio.NewReader(s)
	var millis = time.Millisecond
	millis = 1

	// Receive reply
	for {
		if _, err := io.ReadFull(reader, buffer); err != nil {
			// sleep while no bytes
			// to avoid system flood read
			log.Default.Error("teleinfo", log.Fields{
				"Error": err,
			})
			fmt.Printf("e")
			time.Sleep(1000 * time.Millisecond)
		} else {
			// dispatch io
			for i := 0; i < len(buffer); i++ {
				canal <- buffer[i]
			}
			if len(buffer) < 3000 {
				fmt.Printf(">")
				millis = millis + 10
			} else {
				fmt.Printf("<")
				millis = millis - 10
				if millis == 0 {
					millis = 0
				}
			}
			time.Sleep(millis * time.Millisecond)
		}
	}

	log.Default.Info("teleinfo", log.Fields{
		"status": "done",
	})

	return nil
}

/**
* phase1: LF (0x0A)
* phase2: ETIQUETTE (4 à 8 caractères)
* phase3: SP (0x20)
* phase4: DATA (1 à 12 caractères)
* phase5: SP (0x20)
* phase6: CHECKSUM (caractère de contrôle : Somme de caractère)
* CR (0x0D)
 */

// submit trame
func (that *TeleinfoService) submit(trame TeleinfoTrame) {
	mutex.Lock()
	that.Entries[trame.etiquette] = trame.data
	mutex.Unlock()
}

// single trame
func (that *TeleinfoService) handleTrame(trame string) {
	var espace int
	var send TeleinfoTrame
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x20:
			espace++
			continue
		default:
			if espace == 0 {
				send.etiquette += string([]byte{trame[i]})
			}
			if espace == 1 {
				send.data += string([]byte{trame[i]})
			}
			continue
		}
	}
	// submit new value
	that.submit(send)
}

// all trames detection
func (that *TeleinfoService) handleTrames(trame string) {
	var send string
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x0A:
			send = ""
			continue
		case trame[i] == 0x0D:
			that.handleTrame(send)
			continue
		default:
			send += string([]byte{trame[i]})
		}
	}
}

// worker to consume file
func worker(that *TeleinfoService) {
	var trame string
	var etx bool

	// wait for ETX 0x003
	for i := 0; etx == false; i++ {
		var x = <-canal
		if x == 0x03 {
			etx = true
		}
	}

	// daemon
	for {
		var x = <-canal
		switch {
		case x == 0x03:
			// wait for ETX 0x003
			that.handleTrames(trame)
			break
		case x == 0x02:
			// wait for STX 0x002
			trame = ""
			break
		case x != 0x02 && x != 0x03:
			// other
			trame += string([]byte{x})
			break
		}
	}
}

// GetEntries load entries
func (that *TeleinfoService) GetEntries(entries map[string]string) map[string]string {
	mutex.Lock()
	for key, value := range that.Entries {
		entries[key] = value
	}
	mutex.Unlock()
	return entries
}

// Get get values
func (that *TeleinfoService) Get(key string) string {
	var value string
	mutex.Lock()
	i, ex := that.Entries[key]
	if ex {
		value = i
	}
	mutex.Unlock()
	return value
}

// init initialize this module
func (that *TeleinfoService) init() {
	// add map
	that.Entries = make(map[string]string)

	// start worker
	go handleReadFile(viper.GetString("jarvis.option.teleinfo.file"))
	go worker(that)

	// log information
	log.Default.Info("teleinfo", log.Fields{
		"teleinfoFile": viper.GetString("jarvis.option.teleinfo.file"),
	})
}
<commit_msg>correction<commit_after>/**
 * Copyright 2017 Yannick Roffin
 *
 *   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 teleinfo_service

import (
	"bufio"
	"fmt"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/spf13/viper"
	log "github.com/yroffin/jarvis-go-ext/logger"
)

// On linux
// apply: stty 1200 cs7 evenp cstopb -igncr -inlcr -brkint -icrnl -opost -isig -icanon -iexten -F /dev/ttyUSB0
// to configure teleinfo tty
// Cf. https://hallard.me/gestion-de-la-teleinfo-avec-un-raspberry-pi-et-une-carte-arduipi/

// Teleinfo : instance Teleinfo device struct
type TeleinfoService struct {
	Entries map[string]string
}

// Trame EDF
type TeleinfoTrame struct {
	etiquette string // ETIQUETTE (4 à 8 caractères)
	data      string // DATA (1 à 12 caractères)
	hecksum   string // CHECKSUM (caractère de contrôle : Somme de caractère)
}

var instance *TeleinfoService
var once sync.Once
var mutex = &sync.Mutex{}

// GetInstance : singleton instance
func Service() *TeleinfoService {
	once.Do(func() {
		instance = new(TeleinfoService)
		instance.init()
	})
	return instance
}

// declare canal
var canal = make(chan byte, 1024)

// handleReadFile : read file
func handleReadFile(device string) error {

	s, err := os.OpenFile(device, syscall.O_RDONLY, 0666)

	if err != nil {
		log.Default.Error("teleinfo", log.Fields{
			"error": err,
		})
	}

	log.Default.Info("teleinfo", log.Fields{
		"device": device,
	})

	buffer := make([]byte, 4096)
	reader := bufio.NewReader(s)
	var millis = time.Millisecond
	millis = 1

	// Receive reply
	for {
		if _, err := reader.Read(buffer); err != nil {
			// sleep while no bytes
			// to avoid system flood read
			log.Default.Error("teleinfo", log.Fields{
				"Error": err,
			})
			fmt.Printf("e")
			time.Sleep(1000 * time.Millisecond)
		} else {
			// dispatch io
			for i := 0; i < len(buffer); i++ {
				canal <- buffer[i]
			}
			if len(buffer) < 3000 {
				fmt.Printf(">")
				millis = millis + 10
			} else {
				fmt.Printf("<")
				millis = millis - 10
				if millis == 0 {
					millis = 0
				}
			}
			time.Sleep(millis * time.Millisecond)
		}
	}

	log.Default.Info("teleinfo", log.Fields{
		"status": "done",
	})

	return nil
}

/**
* phase1: LF (0x0A)
* phase2: ETIQUETTE (4 à 8 caractères)
* phase3: SP (0x20)
* phase4: DATA (1 à 12 caractères)
* phase5: SP (0x20)
* phase6: CHECKSUM (caractère de contrôle : Somme de caractère)
* CR (0x0D)
 */

// submit trame
func (that *TeleinfoService) submit(trame TeleinfoTrame) {
	mutex.Lock()
	that.Entries[trame.etiquette] = trame.data
	mutex.Unlock()
}

// single trame
func (that *TeleinfoService) handleTrame(trame string) {
	var espace int
	var send TeleinfoTrame
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x20:
			espace++
			continue
		default:
			if espace == 0 {
				send.etiquette += string([]byte{trame[i]})
			}
			if espace == 1 {
				send.data += string([]byte{trame[i]})
			}
			continue
		}
	}
	// submit new value
	that.submit(send)
}

// all trames detection
func (that *TeleinfoService) handleTrames(trame string) {
	var send string
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x0A:
			send = ""
			continue
		case trame[i] == 0x0D:
			that.handleTrame(send)
			continue
		default:
			send += string([]byte{trame[i]})
		}
	}
}

// worker to consume file
func worker(that *TeleinfoService) {
	var trame string
	var etx bool

	// wait for ETX 0x003
	for i := 0; etx == false; i++ {
		var x = <-canal
		if x == 0x03 {
			etx = true
		}
	}

	// daemon
	for {
		var x = <-canal
		switch {
		case x == 0x03:
			// wait for ETX 0x003
			that.handleTrames(trame)
			break
		case x == 0x02:
			// wait for STX 0x002
			trame = ""
			break
		case x != 0x02 && x != 0x03:
			// other
			trame += string([]byte{x})
			break
		}
	}
}

// GetEntries load entries
func (that *TeleinfoService) GetEntries(entries map[string]string) map[string]string {
	mutex.Lock()
	for key, value := range that.Entries {
		entries[key] = value
	}
	mutex.Unlock()
	return entries
}

// Get get values
func (that *TeleinfoService) Get(key string) string {
	var value string
	mutex.Lock()
	i, ex := that.Entries[key]
	if ex {
		value = i
	}
	mutex.Unlock()
	return value
}

// init initialize this module
func (that *TeleinfoService) init() {
	// add map
	that.Entries = make(map[string]string)

	// start worker
	go handleReadFile(viper.GetString("jarvis.option.teleinfo.file"))
	go worker(that)

	// log information
	log.Default.Info("teleinfo", log.Fields{
		"teleinfoFile": viper.GetString("jarvis.option.teleinfo.file"),
	})
}
<|endoftext|>


<commit_before>package jaeger

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestBaggageIterator(t *testing.T) {
	tracer, closer := NewTracer("DOOP", NewConstSampler(true), NewNullReporter())
	defer closer.Close()

	sp1 := tracer.StartSpan("s1").(*span)
	sp1.SetBaggageItem("Some_Key", "12345")
	sp1.SetBaggageItem("Some-other-key", "42")

	b := make(map[string]string)
	sp1.Context().ForeachBaggageItem(func(k, v string) bool {
		b[k] = v
		return true
	})
	assert.Equal(t, map[string]string{"some-key": "12345", "some-other-key": "42"}, b)

	b = make(map[string]string)
	sp1.Context().ForeachBaggageItem(func(k, v string) bool {
		b[k] = v
		return false // break out early
	})
	assert.Equal(t, 1, len(b), "only one baggage item should be extracted")
}

func TestSpanProperties(t *testing.T) {
	tracer, closer := NewTracer("DOOP", NewConstSampler(true), NewNullReporter())
	defer closer.Close()

	sp1 := tracer.StartSpan("s1").(*span)
	assert.Equal(t, tracer, sp1.Tracer())
	assert.NotNil(t, sp1.Context())
}
<commit_msg>Verify that baggage is passed through to the child span (#44)<commit_after>package jaeger

import (
	"testing"

	"github.com/opentracing/opentracing-go"
	"github.com/stretchr/testify/assert"
)

func TestBaggageIterator(t *testing.T) {
	tracer, closer := NewTracer("DOOP", NewConstSampler(true), NewNullReporter())
	defer closer.Close()

	sp1 := tracer.StartSpan("s1").(*span)
	sp1.SetBaggageItem("Some_Key", "12345")
	sp1.SetBaggageItem("Some-other-key", "42")
	expectedBaggage := map[string]string{"some-key": "12345", "some-other-key": "42"}
	assertBaggage(t, sp1, expectedBaggage)

	b := extractBaggage(sp1, false) // break out early
	assert.Equal(t, 1, len(b), "only one baggage item should be extracted")

	sp2 := tracer.StartSpan("s2", opentracing.ChildOf(sp1.Context()))
	assertBaggage(t, sp2, expectedBaggage) // child inherits the same baggage
}

func assertBaggage(t *testing.T, sp opentracing.Span, expected map[string]string) {
	b := extractBaggage(sp, true)
	assert.Equal(t, expected, b)
}

func extractBaggage(sp opentracing.Span, allItems bool) map[string]string {
	b := make(map[string]string)
	sp.Context().ForeachBaggageItem(func(k, v string) bool {
		b[k] = v
		return allItems
	})
	return b
}

func TestSpanProperties(t *testing.T) {
	tracer, closer := NewTracer("DOOP", NewConstSampler(true), NewNullReporter())
	defer closer.Close()

	sp1 := tracer.StartSpan("s1").(*span)
	assert.Equal(t, tracer, sp1.Tracer())
	assert.NotNil(t, sp1.Context())
}
<|endoftext|>


<commit_before>package dnode

import (
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
)

// Partial is the type of "arguments" field in dnode.Message.
type Partial struct {
	Raw           []byte
	CallbackSpecs []CallbackSpec
}

// MarshalJSON returns the raw bytes of the Partial.
func (p *Partial) MarshalJSON() ([]byte, error) {
	return p.Raw, nil
}

// UnmarshalJSON puts the data into Partial.Raw.
func (p *Partial) UnmarshalJSON(data []byte) error {
	if p == nil {
		return errors.New("json.Partial: UnmarshalJSON on nil pointer")
	}

	p.Raw = make([]byte, len(data))
	copy(p.Raw, data)
	return nil
}

// Unmarshal unmarshals the raw data (p.Raw) into v and prepares callbacks.
// v must be a struct that is the type of expected arguments.
func (p *Partial) Unmarshal(v interface{}) error {
	if p == nil {
		return fmt.Errorf("Cannot unmarshal nil argument")
	}

	value := reflect.ValueOf(v)
	if value.Kind() != reflect.Ptr {
		panic("v must be a pointer")
	}

	if err := json.Unmarshal(p.Raw, &v); err != nil {
		return fmt.Errorf("%s. Data: %s", err.Error(), string(p.Raw))
	}

	for _, spec := range p.CallbackSpecs {
		if err := setCallback(value, spec.Path, spec.Function.Caller.(functionReceived)); err != nil {
			return err
		}
	}

	return nil
}

func (p *Partial) MustUnmarshal(v interface{}) {
	err := p.Unmarshal(v)
	checkError(err)
}

//-------------------------------------------
// Helper methods for unmarshaling JSON types
//-------------------------------------------

// Slice is a helper method to unmarshal a JSON Array.
func (p *Partial) Slice() (a []*Partial, err error) {
	err = p.Unmarshal(&a)
	return
}

// SliceOfLength is a helper method to unmarshal a JSON Array with specified length.
func (p *Partial) SliceOfLength(length int) (a []*Partial, err error) {
	err = p.Unmarshal(&a)
	if err != nil {
		return
	}

	if len(a) != length {
		err = errors.New("Invalid array length")
	}

	return
}

// Map is a helper method to unmarshal to a JSON Object.
func (p *Partial) Map() (m map[string]*Partial, err error) {
	err = p.Unmarshal(&m)
	return
}

// String is a helper to unmarshal a JSON String.
func (p *Partial) String() (s string, err error) {
	err = p.Unmarshal(&s)
	return
}

// Float64 is a helper to unmarshal a JSON Number.
func (p *Partial) Float64() (f float64, err error) {
	err = p.Unmarshal(&f)
	return
}

// Bool is a helper to unmarshal a JSON Boolean.
func (p *Partial) Bool() (b bool, err error) {
	err = p.Unmarshal(&b)
	return
}

// Function is a helper to unmarshal a callback function.
func (p *Partial) Function() (f Function, err error) {
	err = p.Unmarshal(&f)
	return
}

//----------------------------------------------------------------
// Helper methods for unmarshaling JSON types that panic on errors
//----------------------------------------------------------------

func checkError(err error) {
	if err != nil {
		panic(&ArgumentError{err.Error()})
	}
}

func (p *Partial) MustSlice() []*Partial {
	a, err := p.Slice()
	checkError(err)
	return a
}

func (p *Partial) MustSliceOfLength(length int) []*Partial {
	a, err := p.SliceOfLength(length)
	checkError(err)
	return a
}

func (p *Partial) One() *Partial {
	return p.MustSliceOfLength(1)[0]
}

func (p *Partial) MustMap() map[string]*Partial {
	m, err := p.Map()
	checkError(err)
	return m
}

func (p *Partial) MustString() string {
	s, err := p.String()
	checkError(err)
	return s
}

func (p *Partial) MustFloat64() float64 {
	f, err := p.Float64()
	checkError(err)
	return f
}

func (p *Partial) MustBool() bool {
	b, err := p.Bool()
	checkError(err)
	return b
}

func (p *Partial) MustFunction() Function {
	f, err := p.Function()
	checkError(err)
	return f
}
<commit_msg>dnode/partial: checking is done in Unmarshall() already<commit_after>package dnode

import (
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
)

// Partial is the type of "arguments" field in dnode.Message.
type Partial struct {
	Raw           []byte
	CallbackSpecs []CallbackSpec
}

// MarshalJSON returns the raw bytes of the Partial.
func (p *Partial) MarshalJSON() ([]byte, error) {
	return p.Raw, nil
}

// UnmarshalJSON puts the data into Partial.Raw.
func (p *Partial) UnmarshalJSON(data []byte) error {
	if p == nil {
		return errors.New("json.Partial: UnmarshalJSON on nil pointer")
	}

	p.Raw = make([]byte, len(data))
	copy(p.Raw, data)
	return nil
}

// Unmarshal unmarshals the raw data (p.Raw) into v and prepares callbacks.
// v must be a struct that is the type of expected arguments.
func (p *Partial) Unmarshal(v interface{}) error {
	if p == nil {
		return fmt.Errorf("Cannot unmarshal nil argument")
	}

	if err := json.Unmarshal(p.Raw, &v); err != nil {
		return fmt.Errorf("%s. Data: %s", err.Error(), string(p.Raw))
	}

	value := reflect.ValueOf(v)

	for _, spec := range p.CallbackSpecs {
		if err := setCallback(value, spec.Path, spec.Function.Caller.(functionReceived)); err != nil {
			return err
		}
	}

	return nil
}

func (p *Partial) MustUnmarshal(v interface{}) {
	err := p.Unmarshal(v)
	checkError(err)
}

//-------------------------------------------
// Helper methods for unmarshaling JSON types
//-------------------------------------------

// Slice is a helper method to unmarshal a JSON Array.
func (p *Partial) Slice() (a []*Partial, err error) {
	err = p.Unmarshal(&a)
	return
}

// SliceOfLength is a helper method to unmarshal a JSON Array with specified length.
func (p *Partial) SliceOfLength(length int) (a []*Partial, err error) {
	err = p.Unmarshal(&a)
	if err != nil {
		return
	}

	if len(a) != length {
		err = errors.New("Invalid array length")
	}

	return
}

// Map is a helper method to unmarshal to a JSON Object.
func (p *Partial) Map() (m map[string]*Partial, err error) {
	err = p.Unmarshal(&m)
	return
}

// String is a helper to unmarshal a JSON String.
func (p *Partial) String() (s string, err error) {
	err = p.Unmarshal(&s)
	return
}

// Float64 is a helper to unmarshal a JSON Number.
func (p *Partial) Float64() (f float64, err error) {
	err = p.Unmarshal(&f)
	return
}

// Bool is a helper to unmarshal a JSON Boolean.
func (p *Partial) Bool() (b bool, err error) {
	err = p.Unmarshal(&b)
	return
}

// Function is a helper to unmarshal a callback function.
func (p *Partial) Function() (f Function, err error) {
	err = p.Unmarshal(&f)
	return
}

//----------------------------------------------------------------
// Helper methods for unmarshaling JSON types that panic on errors
//----------------------------------------------------------------

func checkError(err error) {
	if err != nil {
		panic(&ArgumentError{err.Error()})
	}
}

func (p *Partial) MustSlice() []*Partial {
	a, err := p.Slice()
	checkError(err)
	return a
}

func (p *Partial) MustSliceOfLength(length int) []*Partial {
	a, err := p.SliceOfLength(length)
	checkError(err)
	return a
}

func (p *Partial) One() *Partial {
	return p.MustSliceOfLength(1)[0]
}

func (p *Partial) MustMap() map[string]*Partial {
	m, err := p.Map()
	checkError(err)
	return m
}

func (p *Partial) MustString() string {
	s, err := p.String()
	checkError(err)
	return s
}

func (p *Partial) MustFloat64() float64 {
	f, err := p.Float64()
	checkError(err)
	return f
}

func (p *Partial) MustBool() bool {
	b, err := p.Bool()
	checkError(err)
	return b
}

func (p *Partial) MustFunction() Function {
	f, err := p.Function()
	checkError(err)
	return f
}
<|endoftext|>


<commit_before>package ltick

import (
	"context"
	"reflect"
	"strings"

	"github.com/fatih/structs"
	"github.com/juju/errors"
	"github.com/ltick/tick-framework/config"
	"github.com/ltick/tick-framework/database"
	"github.com/ltick/tick-framework/filesystem"
	"github.com/ltick/tick-framework/kvstore"
	"github.com/ltick/tick-framework/logger"
	"github.com/ltick/tick-framework/queue"
	"github.com/ltick/tick-framework/session"
	"github.com/ltick/tick-framework/utility"
)

var (
	errComponentExists              = "ltick: component '%s' exists"
	errComponentNotExists           = "ltick: component '%s' not exists"
	errRegisterComponent            = "ltick: register component '%s' error"
	errUnregisterComponent          = "ltick: unregister component '%s' error"
	errInjectComponent              = "ltick: inject component '%s' field '%s' error"
	errInjectComponentTo            = "ltick: inject component '%s' field '%s' error"
	errUseComponent                 = "ltick: use component '%s' error"
	errValueExists                  = "ltick: value '%s' exists"
	errValueNotExists               = "ltick: value '%s' not exists"
	errConfigureComponentFileConfig = "ltick: configure component '%s' file config  error"
)

func (r *Registry) GetComponentMap() map[string]interface{} {
	return r.ComponentMap
}

type ComponentInterface interface {
	Initiate(ctx context.Context) (context.Context, error)
	OnStartup(ctx context.Context) (context.Context, error)
	OnShutdown(ctx context.Context) (context.Context, error)
}

type Components []*Component

func (cs Components) Get(name string) *Component {
	for _, c := range cs {
		canonicalName := strings.ToUpper(name[0:1]) + name[1:]
		if canonicalName == c.Name {
			return c
		}
	}
	return nil
}

type Component struct {
	Name          string
	Component     ComponentInterface
	ConfigurePath string
	Dependencies  []*Component
}

var (
	OptionalComponents = Components{
		&Component{Name: "Log", Component: &log.Logger{}, ConfigurePath: "components.log"},
		&Component{Name: "Database", Component: &database.Database{}, ConfigurePath: "components.database"},
		&Component{Name: "Kvstore", Component: &kvstore.Kvstore{}, ConfigurePath: "components.kvstore"},
		&Component{Name: "Queue", Component: &queue.Queue{}, ConfigurePath: "components.queue"},
		&Component{Name: "Filesystem", Component: &filesystem.Filesystem{}, ConfigurePath: "components.filesystem"},
		&Component{Name: "Session", Component: &session.Session{}, ConfigurePath: "components.session"},
	}
)

/**************** Component ****************/
func (r *Registry) UseComponent(componentNames ...string) error {
	var err error
	components := make([]*Component, 0)
	for _, componentName := range componentNames {
		canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
		component := OptionalComponents.Get(canonicalComponentName)
		if component != nil {
			components = append(components, component)
			err = r.RegisterComponent(strings.ToLower(component.Name[0:1])+component.Name[1:], component.Component, true)
			if err != nil {
				return errors.Annotatef(err, errUseComponent, component.Name)
			}
		} else {
			return errors.Annotatef(err, errComponentNotExists, canonicalComponentName)
		}
	}
	sortedComponents := SortComponent(components)
	for _, name := range sortedComponents {
		component, err := r.GetComponentByName(name)
		if err != nil {
			return errors.Annotatef(err, errUseComponent, name)
		}
		err = r.InjectComponentTo([]interface{}{component})
		if err != nil {
			return errors.Annotatef(err, errUseComponent, name)
		}
	}
	return nil
}

// Register As Component
func (r *Registry) RegisterComponent(componentName string, component ComponentInterface, ignoreIfExistses ...bool) error {
	canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
	ignoreIfExists := false
	if len(ignoreIfExistses) > 0 {
		ignoreIfExists = ignoreIfExistses[0]
	}
	if _, ok := r.ComponentMap[canonicalComponentName]; ok {
		if !ignoreIfExists {
			return errors.Errorf(errComponentExists, canonicalComponentName)
		}
		err := r.UnregisterComponent(canonicalComponentName)
		if err != nil {
			return errors.Annotatef(err, errRegisterComponent+": %s", canonicalComponentName)
		}
	}
	r.Components = append(r.Components, component)
	r.ComponentMap[canonicalComponentName] = component
	r.SortedComponentName = append(r.SortedComponentName, canonicalComponentName)
	return nil
}

// Unregister As Component
func (r *Registry) UnregisterComponent(componentNames ...string) error {
	if len(componentNames) > 0 {
		for _, componentName := range componentNames {
			canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
			// r.ComponentMap
			delete(r.ComponentMap, canonicalComponentName)
			// r.SortedComponentName
			for index, sortedComponentName := range r.SortedComponentName {
				if canonicalComponentName == sortedComponentName {
					r.SortedComponentName = append(r.SortedComponentName[:index], r.SortedComponentName[index+1:]...)
				}
			}
			// r.Components
			for index, c := range r.Components {
				if component, ok := c.(*Component); ok {
					if canonicalComponentName == component.Name {
						r.Components = append(r.Components[:index], r.Components[index+1:]...)
					}
				}
			}
		}
	}
	return nil
}

func (r *Registry) GetComponentByName(componentName string) (interface{}, error) {
	canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
	if _, ok := r.ComponentMap[canonicalComponentName]; !ok {
		return nil, errors.Errorf(errComponentNotExists, canonicalComponentName)
	}
	return r.ComponentMap[canonicalComponentName], nil
}

func (r *Registry) GetSortedComponents(reverses ...bool) []interface{} {
	components := make([]interface{}, len(r.Components))
	if len(r.SortedComponentName) > 0 {
		index := 0
		reverse := false
		if len(reverses) > 0 {
			reverse = reverses[0]
		}
		if reverse {
			for i := len(r.SortedComponentName) - 1; i >= 0; i-- {
				if component, ok := r.ComponentMap[r.SortedComponentName[i]]; ok {
					components[index] = component
					index++
				}
			}
		} else {
			for i := 0; i < len(r.SortedComponentName); i++ {
				if component, ok := r.ComponentMap[r.SortedComponentName[i]]; ok {
					components[index] = component
					index++
				}
			}
		}
	}
	return components
}

func (r *Registry) GetSortedComponentName() []string {
	return r.SortedComponentName
}

// Register As Value
func (r *Registry) RegisterValue(key string, value interface{}, forceOverwrites ...bool) error {
	Key := strings.ToUpper(key[0:1]) + key[1:]
	forceOverwrite := false
	if len(forceOverwrites) > 0 {
		forceOverwrite = forceOverwrites[0]
	}
	if _, ok := r.Values[Key]; ok && !forceOverwrite {
		return errors.Errorf(errValueExists, Key)
	}
	r.Values[Key] = value
	return nil
}

// Unregister As Value
func (r *Registry) UnregisterValue(keys ...string) error {
	if len(keys) > 0 {
		for _, key := range keys {
			Key := strings.ToUpper(key[0:1]) + key[1:]
			if _, ok := r.Values[Key]; !ok {
				return errors.Errorf(errValueNotExists, Key)
			}
			delete(r.Values, Key)
		}
	}
	return nil
}

func (r *Registry) GetValue(key string) (interface{}, error) {
	Key := strings.ToUpper(key[0:1]) + key[1:]
	if _, ok := r.Values[Key]; !ok {
		return nil, errors.Errorf(errValueNotExists, Key)
	}
	return r.Values[Key], nil
}

func (r *Registry) GetValues() map[string]interface{} {
	return r.Values
}

func (r *Registry) InjectComponent() error {
	return r.InjectComponentTo(r.GetSortedComponents())
}

func (r *Registry) InjectComponentByName(componentNames []string) error {
	componentMap := r.GetComponentMap()
	injectTargets := make([]interface{}, 0)
	for _, componentName := range componentNames {
		canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
		if injectTarget, ok := componentMap[canonicalComponentName]; ok {
			injectTargets = append(injectTargets, injectTarget)
		}
	}
	return r.InjectComponentTo(injectTargets)
}

func (r *Registry) InjectComponentTo(injectTargets []interface{}) error {
	for _, injectTarget := range injectTargets {
		injectTargetValue := reflect.ValueOf(injectTarget)
		for injectTargetValue.Kind() == reflect.Ptr {
			injectTargetValue = injectTargetValue.Elem()
		}
		if injectTargetValue.Kind() != reflect.Struct {
			continue
		}
		s := structs.New(injectTarget)
		componentType := reflect.TypeOf((*ComponentInterface)(nil)).Elem()
		for _, f := range s.Fields() {
			if f.IsExported() && f.Tag(INJECT_TAG) == "true" {
				if reflect.TypeOf(f.Value()).Implements(componentType) {
					if _, ok := f.Value().(ComponentInterface); ok {
						fieldInjected := false
						if _, ok := r.ComponentMap[f.Name()]; ok {
							err := f.Set(r.ComponentMap[f.Name()])
							if err != nil {
								return errors.Annotatef(err, errInjectComponentTo, injectTargetValue.String(), f.Name())
							}
							fieldInjected = true
						}
						if !fieldInjected {
							return errors.Errorf(errInjectComponentTo+": component or key not exists", injectTargetValue.String(), f.Name())
						}
					}
				}
				if _, ok := r.Values[f.Name()]; ok {
					err := f.Set(r.Values[f.Name()])
					if err != nil {
						return errors.Annotatef(err, errInjectComponentTo, injectTargetValue.String(), f.Name())
					}
				}
			}
		}
	}
	return nil
}

/******** Component Dependency Manage ********/
// SortComponent - Sort user components.
func SortComponent(components []*Component) []string {
	// 初始化依赖关系
	for _, c := range components {
		if c.Dependencies == nil {
			c.Dependencies = make([]*Component, 0)
		}
		s := structs.New(c.Component)
		componentType := reflect.TypeOf((*ComponentInterface)(nil)).Elem()
		for _, f := range s.Fields() {
			if f.IsExported() && f.Tag(INJECT_TAG) == "true" {
				if reflect.TypeOf(f.Value()).Implements(componentType) {
					if dc, ok := f.Value().(ComponentInterface); ok {
						c.Dependencies = append(c.Dependencies, &Component{
							Name:      f.Name(),
							Component: dc,
						})
					}
				}
			}
		}
	}
	// root components
	roots := []*Component{}
	for {
		rootDependenciesCount := 0
		for _, c := range components {
			if len(c.Dependencies) == rootDependenciesCount {
				roots = append(roots, c)
			}
		}
		if len(roots) > 0 {
			break
		} else {
			rootDependenciesCount++
		}
	}

	sortedComponents := make([]string, 0)
	return sortComponent(components, roots, sortedComponents)
}

// sortComponent
func sortComponent(components []*Component, currentComponents []*Component, sortedComponents []string) []string {
	if components != nil && currentComponents != nil && sortedComponents != nil {
		// 没有下级依赖
		if len(currentComponents) == 0 {
			return sortedComponents
		}
		for _, currentComponent := range currentComponents {
			// 当前层级组件
			index := utility.InArrayString(currentComponent.Name, sortedComponents, false)
			if index == nil {
				sortedComponents = append(sortedComponents, currentComponent.Name)
			} else {
				sortedComponents = append(sortedComponents[:*index], append(sortedComponents[*index+1:], sortedComponents[*index])...)
			}
			// 依赖当前级别组件的组件
			componentsNextLevel := make([]*Component, 0)
			for _, component := range components {
				for _, componentDependencie := range component.Dependencies {
					if strings.Compare(componentDependencie.Name, currentComponent.Name) == 0 {
						componentsNextLevel = append(componentsNextLevel, component)
					}
				}
			}
			sortedComponents = sortComponent(components, componentsNextLevel, sortedComponents)
		}
	}
	return sortedComponents
}

func (r *Registry) ConfigureComponentFileConfig(componentName string, configFile string, configProviders map[string]interface{}, configTag ...string) (err error) {
	canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
	configComponent, err := r.GetComponentByName("Config")
	if err != nil {
		return errors.Annotatef(err, errConfigureComponentFileConfig, canonicalComponentName)
	}
	configer, ok := configComponent.(*config.Config)
	if !ok {
		return errors.Annotatef(errors.Errorf("invalid 'Config' component type"), errConfigureComponentFileConfig, canonicalComponentName)
	}
	// configer
	for componentName, component := range r.ComponentMap {
		canonicalExistsComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
		if canonicalComponentName == canonicalExistsComponentName {
			if len(configTag) > 0 {
				// create a Config object
				err = configer.ConfigureFileConfig(component, configFile, configProviders, configTag...)
				if err != nil {
					return errors.Annotatef(err, errConfigureComponentFileConfig, canonicalComponentName)
				}
			}
		}
	}
	return nil
}
<commit_msg>component<commit_after>package ltick

import (
	"context"
	"reflect"
	"strings"

	"github.com/fatih/structs"
	"github.com/juju/errors"
	"github.com/ltick/tick-framework/config"
	"github.com/ltick/tick-framework/database"
	"github.com/ltick/tick-framework/filesystem"
	"github.com/ltick/tick-framework/kvstore"
	"github.com/ltick/tick-framework/logger"
	"github.com/ltick/tick-framework/queue"
	"github.com/ltick/tick-framework/session"
	"github.com/ltick/tick-framework/utility"
)

var (
	errComponentExists              = "ltick: component '%s' exists"
	errComponentNotExists           = "ltick: component '%s' not exists"
	errRegisterComponent            = "ltick: register component '%s' error"
	errUnregisterComponent          = "ltick: unregister component '%s' error"
	errInjectComponent              = "ltick: inject component '%s' field '%s' error"
	errInjectComponentTo            = "ltick: inject component '%s' field '%s' error"
	errUseComponent                 = "ltick: use component '%s' error"
	errValueExists                  = "ltick: value '%s' exists"
	errValueNotExists               = "ltick: value '%s' not exists"
	errConfigureComponentFileConfig = "ltick: configure component '%s' file config  error"
)

func (r *Registry) GetComponentMap() map[string]interface{} {
	return r.ComponentMap
}

type ComponentInterface interface {
	Initiate(ctx context.Context) (context.Context, error)
	OnStartup(ctx context.Context) (context.Context, error)
	OnShutdown(ctx context.Context) (context.Context, error)
}

type Components []*Component

func (cs Components) Get(name string) *Component {
	for _, c := range cs {
		canonicalName := strings.ToUpper(name[0:1]) + name[1:]
		if canonicalName == c.Name {
			return c
		}
	}
	return nil
}

type Component struct {
	Name          string
	Component     ComponentInterface
	ConfigurePath string
	Dependencies  []*Component
}

var (
	OptionalComponents = Components{
		&Component{Name: "Log", Component: &log.Logger{}, ConfigurePath: "components.log"},
		&Component{Name: "Database", Component: &database.Database{}, ConfigurePath: "components.database"},
		&Component{Name: "Kvstore", Component: &kvstore.Kvstore{}, ConfigurePath: "components.kvstore"},
		&Component{Name: "Queue", Component: &queue.Queue{}, ConfigurePath: "components.queue"},
		&Component{Name: "Filesystem", Component: &filesystem.Filesystem{}, ConfigurePath: "components.filesystem"},
		&Component{Name: "Session", Component: &session.Session{}, ConfigurePath: "components.session"},
	}
)

/**************** Component ****************/
func (r *Registry) UseComponent(componentNames ...string) error {
	var err error
	components := make([]*Component, 0)
	for _, componentName := range componentNames {
		canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
		component := OptionalComponents.Get(canonicalComponentName)
		if component != nil {
			components = append(components, component)
			err = r.RegisterComponent(strings.ToLower(component.Name[0:1])+component.Name[1:], component.Component, true)
			if err != nil {
				return errors.Annotatef(err, errUseComponent, component.Name)
			}
		} else {
			return errors.Annotatef(err, errComponentNotExists, canonicalComponentName)
		}
	}
	sortedComponents := SortComponent(components)
	for _, name := range sortedComponents {
		component, err := r.GetComponentByName(name)
		if err != nil {
			return errors.Annotatef(err, errUseComponent, name)
		}
		err = r.InjectComponentTo([]interface{}{component})
		if err != nil {
			return errors.Annotatef(err, errUseComponent, name)
		}
	}
	return nil
}

// Register As Component
func (r *Registry) RegisterComponent(componentName string, component ComponentInterface, ignoreIfExistses ...bool) error {
	canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
	ignoreIfExists := false
	if len(ignoreIfExistses) > 0 {
		ignoreIfExists = ignoreIfExistses[0]
	}
	if _, ok := r.ComponentMap[canonicalComponentName]; ok {
		if !ignoreIfExists {
			return errors.Errorf(errComponentExists, canonicalComponentName)
		}
		err := r.UnregisterComponent(canonicalComponentName)
		if err != nil {
			return errors.Annotatef(err, errRegisterComponent+": %s", canonicalComponentName)
		}
	}
	r.Components = append(r.Components, component)
	r.ComponentMap[canonicalComponentName] = component
	r.SortedComponentName = append(r.SortedComponentName, canonicalComponentName)
	return nil
}

// Unregister As Component
func (r *Registry) UnregisterComponent(componentNames ...string) error {
	if len(componentNames) > 0 {
		for _, componentName := range componentNames {
			canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
			// r.ComponentMap
			delete(r.ComponentMap, canonicalComponentName)
			// r.SortedComponentName
			for index, sortedComponentName := range r.SortedComponentName {
				if canonicalComponentName == sortedComponentName {
					r.SortedComponentName = append(r.SortedComponentName[:index], r.SortedComponentName[index+1:]...)
				}
			}
			// r.Components
			for index, c := range r.Components {
				if component, ok := c.(*Component); ok {
					if canonicalComponentName == component.Name {
						r.Components = append(r.Components[:index], r.Components[index+1:]...)
					}
				}
			}
		}
	}
	return nil
}

func (r *Registry) GetComponentByName(componentName string) (interface{}, error) {
	canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
	if _, ok := r.ComponentMap[canonicalComponentName]; !ok {
		return nil, errors.Errorf(errComponentNotExists, canonicalComponentName)
	}
	return r.ComponentMap[canonicalComponentName], nil
}

func (r *Registry) GetSortedComponents(reverses ...bool) []interface{} {
	components := make([]interface{}, len(r.Components))
	if len(r.SortedComponentName) > 0 {
		index := 0
		reverse := false
		if len(reverses) > 0 {
			reverse = reverses[0]
		}
		if reverse {
			for i := len(r.SortedComponentName) - 1; i >= 0; i-- {
				if component, ok := r.ComponentMap[r.SortedComponentName[i]]; ok {
					components[index] = component
					index++
				}
			}
		} else {
			for i := 0; i < len(r.SortedComponentName); i++ {
				if component, ok := r.ComponentMap[r.SortedComponentName[i]]; ok {
					components[index] = component
					index++
				}
			}
		}
	}
	return components
}

func (r *Registry) GetSortedComponentName() []string {
	return r.SortedComponentName
}

// Register As Value
func (r *Registry) RegisterValue(key string, value interface{}, forceOverwrites ...bool) error {
	Key := strings.ToUpper(key[0:1]) + key[1:]
	forceOverwrite := false
	if len(forceOverwrites) > 0 {
		forceOverwrite = forceOverwrites[0]
	}
	if _, ok := r.Values[Key]; ok && !forceOverwrite {
		return errors.Errorf(errValueExists, Key)
	}
	r.Values[Key] = value
	return nil
}

// Unregister As Value
func (r *Registry) UnregisterValue(keys ...string) error {
	if len(keys) > 0 {
		for _, key := range keys {
			Key := strings.ToUpper(key[0:1]) + key[1:]
			if _, ok := r.Values[Key]; !ok {
				return errors.Errorf(errValueNotExists, Key)
			}
			delete(r.Values, Key)
		}
	}
	return nil
}

func (r *Registry) GetValue(key string) (interface{}, error) {
	Key := strings.ToUpper(key[0:1]) + key[1:]
	if _, ok := r.Values[Key]; !ok {
		return nil, errors.Errorf(errValueNotExists, Key)
	}
	return r.Values[Key], nil
}

func (r *Registry) GetValues() map[string]interface{} {
	return r.Values
}

func (r *Registry) InjectComponent() error {
	return r.InjectComponentTo(r.GetSortedComponents())
}

func (r *Registry) InjectComponentByName(componentNames []string) error {
	componentMap := r.GetComponentMap()
	injectTargets := make([]interface{}, 0)
	for _, componentName := range componentNames {
		canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
		if injectTarget, ok := componentMap[canonicalComponentName]; ok {
			injectTargets = append(injectTargets, injectTarget)
		}
	}
	return r.InjectComponentTo(injectTargets)
}

func (r *Registry) InjectComponentTo(injectTargets []interface{}) error {
	for _, injectTarget := range injectTargets {
		injectTargetValue := reflect.ValueOf(injectTarget)
		for injectTargetValue.Kind() == reflect.Ptr {
			injectTargetValue = injectTargetValue.Elem()
		}
		if injectTargetValue.Kind() != reflect.Struct {
			continue
		}
		s := structs.New(injectTarget)
		componentType := reflect.TypeOf((*ComponentInterface)(nil)).Elem()
		for _, f := range s.Fields() {
			if f.IsExported() && f.Tag(INJECT_TAG) == "true" {
				if reflect.TypeOf(f.Value()).Implements(componentType) {
					if _, ok := f.Value().(ComponentInterface); ok {
						fieldInjected := false
						if _, ok := r.ComponentMap[f.Name()]; ok {
							err := f.Set(r.ComponentMap[f.Name()])
							if err != nil {
								return errors.Annotatef(err, errInjectComponentTo, injectTargetValue.String(), f.Name())
							}
							fieldInjected = true
						}
						if !fieldInjected {
							return errors.Errorf(errInjectComponentTo+": component or key not exists", injectTargetValue.String(), f.Name())
						}
					}
				}
				if _, ok := r.Values[f.Name()]; ok {
					err := f.Set(r.Values[f.Name()])
					if err != nil {
						return errors.Annotatef(err, errInjectComponentTo, injectTargetValue.String(), f.Name())
					}
				}
			}
		}
	}
	return nil
}

/******** Component Dependency Manage ********/
// SortComponent - Sort user components.
func SortComponent(components []*Component) []string {
	// 初始化依赖关系
	for _, c := range components {
		if c.Dependencies == nil {
			c.Dependencies = make([]*Component, 0)
		}
		s := structs.New(c.Component)
		componentType := reflect.TypeOf((*ComponentInterface)(nil)).Elem()
		for _, f := range s.Fields() {
			if f.IsExported() && f.Tag(INJECT_TAG) == "true" {
				if reflect.TypeOf(f.Value()).Implements(componentType) {
					if dc, ok := f.Value().(ComponentInterface); ok {
						c.Dependencies = append(c.Dependencies, &Component{
							Name:      f.Name(),
							Component: dc,
						})
					}
				}
			}
		}
	}
	// root components
	roots := []*Component{}
	rootDependenciesCount := 0
	for {
		for _, c := range components {
			if len(c.Dependencies) == rootDependenciesCount {
				roots = append(roots, c)
			}
		}
		if len(roots) > 0 {
			break
		} else {
			rootDependenciesCount++
		}
	}

	sortedComponents := make([]string, 0)
	return sortComponent(components, roots, sortedComponents)
}

// sortComponent
func sortComponent(components []*Component, currentComponents []*Component, sortedComponents []string) []string {
	if components != nil && currentComponents != nil && sortedComponents != nil {
		// 没有下级依赖
		if len(currentComponents) == 0 {
			return sortedComponents
		}
		for _, currentComponent := range currentComponents {
			// 当前层级组件
			index := utility.InArrayString(currentComponent.Name, sortedComponents, false)
			if index == nil {
				sortedComponents = append(sortedComponents, currentComponent.Name)
			} else {
				sortedComponents = append(sortedComponents[:*index], append(sortedComponents[*index+1:], sortedComponents[*index])...)
			}
			// 依赖当前级别组件的组件
			componentsNextLevel := make([]*Component, 0)
			for _, component := range components {
				for _, componentDependencie := range component.Dependencies {
					if strings.Compare(componentDependencie.Name, currentComponent.Name) == 0 {
						componentsNextLevel = append(componentsNextLevel, component)
					}
				}
			}
			sortedComponents = sortComponent(components, componentsNextLevel, sortedComponents)
		}
	}
	return sortedComponents
}

func (r *Registry) ConfigureComponentFileConfig(componentName string, configFile string, configProviders map[string]interface{}, configTag ...string) (err error) {
	canonicalComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
	configComponent, err := r.GetComponentByName("Config")
	if err != nil {
		return errors.Annotatef(err, errConfigureComponentFileConfig, canonicalComponentName)
	}
	configer, ok := configComponent.(*config.Config)
	if !ok {
		return errors.Annotatef(errors.Errorf("invalid 'Config' component type"), errConfigureComponentFileConfig, canonicalComponentName)
	}
	// configer
	for componentName, component := range r.ComponentMap {
		canonicalExistsComponentName := strings.ToUpper(componentName[0:1]) + componentName[1:]
		if canonicalComponentName == canonicalExistsComponentName {
			if len(configTag) > 0 {
				// create a Config object
				err = configer.ConfigureFileConfig(component, configFile, configProviders, configTag...)
				if err != nil {
					return errors.Annotatef(err, errConfigureComponentFileConfig, canonicalComponentName)
				}
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

// merges multiple frequency count files into one
// this is could use a merge sort and be smart

import (
	"bufio"
	"compress/gzip"
	"flag"
	"fmt"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"
)

// freqCount is mapping of string->count
type freqCount map[string]int

// make a new counter with some minor preallocation
//  each month has about 2.2M uniques
func newFreqCount() freqCount {
	return make(freqCount, 3000000)
}

// LoadCSV loads in a CSV in form of WORD,COUNT
func loadCSV(counts freqCount, fname string) error {
	fi, err := os.Open(fname)
	if err != nil {
		return err
	}
	defer fi.Close()

	fizip, err := gzip.NewReader(fi)
	if err != nil {
		return err
	}
	defer fizip.Close()

	scanner := bufio.NewScanner(fizip)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.SplitN(line, ",", 2)
		if len(parts) != 2 {
			return fmt.Errorf("Got extra junk in line: %q", line)
		}
		c, err := strconv.Atoi(parts[1])
		if err != nil {
			return fmt.Errorf("Number conversion failed: %q", line)
		}

		counts[parts[0]] += c
	}
	return scanner.Err()
}

func main() {
	outfile := flag.String("o", "", "output file name")
	mincount := flag.Int("mincount", 0, "only output if freqcount greater than this, 0 = all")
	minlen := flag.Int("minlen", 0, "only output if word is >, 0 = all")
	flag.Parse()
	if *outfile == "" {
		log.Fatalf("Must specificy outfile")
	}
	args := flag.Args()
	counts := newFreqCount()
	for _, arg := range args {
		log.Printf("Loading %s", arg)
		err := loadCSV(counts, arg)
		if err != nil {
			log.Fatalf("%s: %s", arg, err)
		}
	}
	fo, err := os.Create(*outfile)
	if err != nil {
		log.Fatalf("OH NO, unable to write: %s", err)
	}
	fout := gzip.NewWriter(fo)

	keys := make([]string, 0, len(counts))
	total := 0
	for k, v := range counts {
		total += v
		if v > *mincount && len(k) > *minlen {
			keys = append(keys, k)
		}
	}
	sort.Strings(keys)
	for _, k := range keys {
		fout.Write([]byte(fmt.Sprintf("%s,%d\n", k, counts[k])))
	}

	fout.Close()
	fo.Close()
	log.Printf("DONE: wrote %s got %d unique words from %d", *outfile, len(keys), total)
}
<commit_msg>make flag more clear<commit_after>package main

// merges multiple frequency count files into one
// this is could use a merge sort and be smart

import (
	"bufio"
	"compress/gzip"
	"flag"
	"fmt"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"
)

// freqCount is mapping of string->count
type freqCount map[string]int

// make a new counter with some minor preallocation
//  each month has about 2.2M uniques
func newFreqCount() freqCount {
	return make(freqCount, 3000000)
}

// LoadCSV loads in a CSV in form of WORD,COUNT
func loadCSV(counts freqCount, fname string) error {
	fi, err := os.Open(fname)
	if err != nil {
		return err
	}
	defer fi.Close()

	fizip, err := gzip.NewReader(fi)
	if err != nil {
		return err
	}
	defer fizip.Close()

	scanner := bufio.NewScanner(fizip)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.SplitN(line, ",", 2)
		if len(parts) != 2 {
			return fmt.Errorf("Got extra junk in line: %q", line)
		}
		c, err := strconv.Atoi(parts[1])
		if err != nil {
			return fmt.Errorf("Number conversion failed: %q", line)
		}

		counts[parts[0]] += c
	}
	return scanner.Err()
}

func main() {
	outfile := flag.String("o", "", "output file name")
	mincount := flag.Int("mincount", 0, "only output if freqcount greater than this, 0 = all")
	minlen := flag.Int("minlen", 0, "only output if word is >, 0 = all")
	flag.Parse()
	if *outfile == "" {
		log.Fatalf("Must specificy outfile")
	}
	args := flag.Args()
	counts := newFreqCount()
	for _, arg := range args {
		log.Printf("Loading %s", arg)
		err := loadCSV(counts, arg)
		if err != nil {
			log.Fatalf("%s: %s", arg, err)
		}
	}
	fo, err := os.Create(*outfile)
	if err != nil {
		log.Fatalf("OH NO, unable to write: %s", err)
	}
	fout := gzip.NewWriter(fo)

	keys := make([]string, 0, len(counts))
	total := 0
	for k, v := range counts {
		total += v
		if v >= *mincount && len(k) >= *minlen {
			keys = append(keys, k)
		}
	}
	sort.Strings(keys)
	for _, k := range keys {
		fout.Write([]byte(fmt.Sprintf("%s,%d\n", k, counts[k])))
	}

	fout.Close()
	fo.Close()
	log.Printf("DONE: wrote %s got %d unique words from %d", *outfile, len(keys), total)
}
<|endoftext|>


<commit_before>package portmapper

import (
	"errors"
	"fmt"
	"net"
	"sync"

	"github.com/dotcloud/docker/daemon/networkdriver/portallocator"
	"github.com/dotcloud/docker/pkg/iptables"
	"github.com/dotcloud/docker/pkg/proxy"
)

type mapping struct {
	proto         string
	userlandProxy proxy.Proxy
	host          net.Addr
	container     net.Addr
}

var (
	chain *iptables.Chain
	lock  sync.Mutex

	// udp:ip:port
	currentMappings = make(map[string]*mapping)
	newProxy        = proxy.NewProxy
)

var (
	ErrUnknownBackendAddressType = errors.New("unknown container address type not supported")
	ErrPortMappedForIP           = errors.New("port is already mapped to ip")
	ErrPortNotMapped             = errors.New("port is not mapped")
)

func SetIptablesChain(c *iptables.Chain) {
	chain = c
}

func Map(container net.Addr, hostIP net.IP, hostPort int) (net.Addr, error) {
	lock.Lock()
	defer lock.Unlock()

	var (
		m                 *mapping
		err               error
		proto             string
		allocatedHostPort int
	)

	// release the port on any error during return.
	defer func() {
		if err != nil {
			portallocator.ReleasePort(hostIP, proto, allocatedHostPort)
		}
	}()

	switch container.(type) {
	case *net.TCPAddr:
		proto = "tcp"
		if allocatedHostPort, err = portallocator.RequestPort(hostIP, proto, hostPort); err != nil {
			return nil, err
		}
		m = &mapping{
			proto:     proto,
			host:      &net.TCPAddr{IP: hostIP, Port: allocatedHostPort},
			container: container,
		}
	case *net.UDPAddr:
		proto = "udp"
		if allocatedHostPort, err = portallocator.RequestPort(hostIP, proto, hostPort); err != nil {
			return nil, err
		}
		m = &mapping{
			proto:     proto,
			host:      &net.UDPAddr{IP: hostIP, Port: allocatedHostPort},
			container: container,
		}
	default:
		err = ErrUnknownBackendAddressType
		return nil, err
	}

	key := getKey(m.host)
	if _, exists := currentMappings[key]; exists {
		err = ErrPortMappedForIP
		return nil, err
	}

	containerIP, containerPort := getIPAndPort(m.container)
	if err := forward(iptables.Add, m.proto, hostIP, allocatedHostPort, containerIP.String(), containerPort); err != nil {
		return nil, err
	}

	p, err := newProxy(m.host, m.container)
	if err != nil {
		// need to undo the iptables rules before we return
		forward(iptables.Delete, m.proto, hostIP, allocatedHostPort, containerIP.String(), containerPort)
		return nil, err
	}

	m.userlandProxy = p
	currentMappings[key] = m

	go p.Run()

	return m.host, nil
}

func Unmap(host net.Addr) error {
	lock.Lock()
	defer lock.Unlock()

	key := getKey(host)
	data, exists := currentMappings[key]
	if !exists {
		return ErrPortNotMapped
	}

	data.userlandProxy.Close()
	delete(currentMappings, key)

	containerIP, containerPort := getIPAndPort(data.container)
	hostIP, hostPort := getIPAndPort(data.host)
	if err := forward(iptables.Delete, data.proto, hostIP, hostPort, containerIP.String(), containerPort); err != nil {
		return err
	}

	switch a := host.(type) {
	case *net.TCPAddr:
		if err := portallocator.ReleasePort(a.IP, "tcp", a.Port); err != nil {
			return err
		}
	case *net.UDPAddr:
		if err := portallocator.ReleasePort(a.IP, "udp", a.Port); err != nil {
			return err
		}
	}

	return nil
}

func getKey(a net.Addr) string {
	switch t := a.(type) {
	case *net.TCPAddr:
		return fmt.Sprintf("%s:%d/%s", t.IP.String(), t.Port, "tcp")
	case *net.UDPAddr:
		return fmt.Sprintf("%s:%d/%s", t.IP.String(), t.Port, "udp")
	}
	return ""
}

func getIPAndPort(a net.Addr) (net.IP, int) {
	switch t := a.(type) {
	case *net.TCPAddr:
		return t.IP, t.Port
	case *net.UDPAddr:
		return t.IP, t.Port
	}
	return nil, 0
}

func forward(action iptables.Action, proto string, sourceIP net.IP, sourcePort int, containerIP string, containerPort int) error {
	if chain == nil {
		return nil
	}
	return chain.Forward(action, sourceIP, sourcePort, proto, containerIP, containerPort)
}
<commit_msg>Make sure err never gets masked<commit_after>package portmapper

import (
	"errors"
	"fmt"
	"net"
	"sync"

	"github.com/dotcloud/docker/daemon/networkdriver/portallocator"
	"github.com/dotcloud/docker/pkg/iptables"
	"github.com/dotcloud/docker/pkg/proxy"
)

type mapping struct {
	proto         string
	userlandProxy proxy.Proxy
	host          net.Addr
	container     net.Addr
}

var (
	chain *iptables.Chain
	lock  sync.Mutex

	// udp:ip:port
	currentMappings = make(map[string]*mapping)
	newProxy        = proxy.NewProxy
)

var (
	ErrUnknownBackendAddressType = errors.New("unknown container address type not supported")
	ErrPortMappedForIP           = errors.New("port is already mapped to ip")
	ErrPortNotMapped             = errors.New("port is not mapped")
)

func SetIptablesChain(c *iptables.Chain) {
	chain = c
}

func Map(container net.Addr, hostIP net.IP, hostPort int) (host net.Addr, err error) {
	lock.Lock()
	defer lock.Unlock()

	var (
		m                 *mapping
		proto             string
		allocatedHostPort int
	)

	switch container.(type) {
	case *net.TCPAddr:
		proto = "tcp"
		if allocatedHostPort, err = portallocator.RequestPort(hostIP, proto, hostPort); err != nil {
			return nil, err
		}
		m = &mapping{
			proto:     proto,
			host:      &net.TCPAddr{IP: hostIP, Port: allocatedHostPort},
			container: container,
		}
	case *net.UDPAddr:
		proto = "udp"
		if allocatedHostPort, err = portallocator.RequestPort(hostIP, proto, hostPort); err != nil {
			return nil, err
		}
		m = &mapping{
			proto:     proto,
			host:      &net.UDPAddr{IP: hostIP, Port: allocatedHostPort},
			container: container,
		}
	default:
		return nil, ErrUnknownBackendAddressType
	}

	// release the allocated port on any further error during return.
	defer func() {
		if err != nil {
			portallocator.ReleasePort(hostIP, proto, allocatedHostPort)
		}
	}()

	key := getKey(m.host)
	if _, exists := currentMappings[key]; exists {
		return nil, ErrPortMappedForIP
	}

	containerIP, containerPort := getIPAndPort(m.container)
	if err := forward(iptables.Add, m.proto, hostIP, allocatedHostPort, containerIP.String(), containerPort); err != nil {
		return nil, err
	}

	p, err := newProxy(m.host, m.container)
	if err != nil {
		// need to undo the iptables rules before we return
		forward(iptables.Delete, m.proto, hostIP, allocatedHostPort, containerIP.String(), containerPort)
		return nil, err
	}

	m.userlandProxy = p
	currentMappings[key] = m

	go p.Run()

	return m.host, nil
}

func Unmap(host net.Addr) error {
	lock.Lock()
	defer lock.Unlock()

	key := getKey(host)
	data, exists := currentMappings[key]
	if !exists {
		return ErrPortNotMapped
	}

	data.userlandProxy.Close()
	delete(currentMappings, key)

	containerIP, containerPort := getIPAndPort(data.container)
	hostIP, hostPort := getIPAndPort(data.host)
	if err := forward(iptables.Delete, data.proto, hostIP, hostPort, containerIP.String(), containerPort); err != nil {
		return err
	}

	switch a := host.(type) {
	case *net.TCPAddr:
		if err := portallocator.ReleasePort(a.IP, "tcp", a.Port); err != nil {
			return err
		}
	case *net.UDPAddr:
		if err := portallocator.ReleasePort(a.IP, "udp", a.Port); err != nil {
			return err
		}
	}

	return nil
}

func getKey(a net.Addr) string {
	switch t := a.(type) {
	case *net.TCPAddr:
		return fmt.Sprintf("%s:%d/%s", t.IP.String(), t.Port, "tcp")
	case *net.UDPAddr:
		return fmt.Sprintf("%s:%d/%s", t.IP.String(), t.Port, "udp")
	}
	return ""
}

func getIPAndPort(a net.Addr) (net.IP, int) {
	switch t := a.(type) {
	case *net.TCPAddr:
		return t.IP, t.Port
	case *net.UDPAddr:
		return t.IP, t.Port
	}
	return nil, 0
}

func forward(action iptables.Action, proto string, sourceIP net.IP, sourcePort int, containerIP string, containerPort int) error {
	if chain == nil {
		return nil
	}
	return chain.Forward(action, sourceIP, sourcePort, proto, containerIP, containerPort)
}
<|endoftext|>