<commit_before>// +build integration

package zk

import (
	"bytes"
	"fmt"
	"os"
	"testing"
	"time"

	stdzk "github.com/samuel/go-zookeeper/zk"
)

var (
	host []string
)

func TestMain(m *testing.M) {
	fmt.Println("ZooKeeper Integration Test Initializing. Starting ZooKeeper Server...")
	ts, err := stdzk.StartTestCluster(1, nil, nil)
	if err != nil {
		fmt.Printf("Unable to start ZooKeeper Server: %v\n", err)
		os.Exit(-1)
	}
	defer ts.Stop()
	host = []string{fmt.Sprintf("localhost:%d", ts.Servers[0].Port)}
	code := m.Run()
	os.Exit(code)
}

func TestCreateParentNodesOnServer(t *testing.T) {
	payload := [][]byte{[]byte("Payload"), []byte("Test")}
	c1, err := NewClient(host, logger, Payload(payload))
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}
	if c1 == nil {
		t.Fatal("Expected pointer to client, got nil")
	}
	defer c1.Stop()

	p, err := NewPublisher(c1, path, newFactory(""), logger)
	if err != nil {
		t.Fatalf("Unable to create Publisher: %v", err)
	}
	defer p.Stop()

	endpoints, err := p.Endpoints()
	if err != nil {
		t.Fatal(err)
	}
	if want, have := 0, len(endpoints); want != have {
		t.Errorf("want %d, have %d", want, have)
	}

	c2, err := NewClient(host, logger)
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}
	defer c2.Stop()
	data, _, err := c2.(*client).Get(path)
	if err != nil {
		t.Fatal(err)
	}
	// test Client implementation of CreateParentNodes. It should have created
	// our payload
	if bytes.Compare(data, payload[1]) != 0 {
		t.Errorf("want %s, have %s", payload[1], data)
	}

}

func TestCreateBadParentNodesOnServer(t *testing.T) {
	c, _ := NewClient(host, logger)
	defer c.Stop()

	_, err := NewPublisher(c, "invalid/path", newFactory(""), logger)

	if want, have := stdzk.ErrInvalidPath, err; want != have {
		t.Errorf("want %v, have %v", want, have)
	}
}

func TestCredentials1(t *testing.T) {
	acl := stdzk.DigestACL(stdzk.PermAll, "user", "secret")
	c, _ := NewClient(host, logger, ACL(acl), Credentials("user", "secret"))
	defer c.Stop()

	_, err := NewPublisher(c, "/acl-issue-test", newFactory(""), logger)

	if err != nil {
		t.Fatal(err)
	}
}

func TestCredentials2(t *testing.T) {
	acl := stdzk.DigestACL(stdzk.PermAll, "user", "secret")
	c, _ := NewClient(host, logger, ACL(acl))
	defer c.Stop()

	_, err := NewPublisher(c, "/acl-issue-test", newFactory(""), logger)

	if err != stdzk.ErrNoAuth {
		t.Errorf("want %v, have %v", stdzk.ErrNoAuth, err)
	}
}

func TestConnection(t *testing.T) {
	c, _ := NewClient(host, logger)
	c.Stop()

	_, err := NewPublisher(c, "/acl-issue-test", newFactory(""), logger)

	if err != ErrClientClosed {
		t.Errorf("want %v, have %v", ErrClientClosed, err)
	}
}

func TestGetEntriesOnServer(t *testing.T) {
	var instancePayload = "protocol://hostname:port/routing"

	c1, err := NewClient(host, logger)
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}

	defer c1.Stop()

	c2, err := NewClient(host, logger)
	p, err := NewPublisher(c2, path, newFactory(""), logger)
	if err != nil {
		t.Fatal(err)
	}
	defer c2.Stop()

	c2impl, _ := c2.(*client)
	_, err = c2impl.Create(
		path+"/instance1",
		[]byte(instancePayload),
		stdzk.FlagEphemeral|stdzk.FlagSequence,
		stdzk.WorldACL(stdzk.PermAll),
	)
	if err != nil {
		t.Fatalf("Unable to create test ephemeral znode 1: %v", err)
	}
	_, err = c2impl.Create(
		path+"/instance2",
		[]byte(instancePayload+"2"),
		stdzk.FlagEphemeral|stdzk.FlagSequence,
		stdzk.WorldACL(stdzk.PermAll),
	)
	if err != nil {
		t.Fatalf("Unable to create test ephemeral znode 2: %v", err)
	}

	time.Sleep(50 * time.Millisecond)

	endpoints, err := p.Endpoints()
	if err != nil {
		t.Fatal(err)
	}
	if want, have := 2, len(endpoints); want != have {
		t.Errorf("want %d, have %d", want, have)
	}
}

func TestGetEntriesPayloadOnServer(t *testing.T) {
	c, err := NewClient(host, logger)
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}
	_, eventc, err := c.GetEntries(path)
	if err != nil {
		t.Fatal(err)
	}
	_, err = c.(*client).Create(
		path+"/instance3",
		[]byte("just some payload"),
		stdzk.FlagEphemeral|stdzk.FlagSequence,
		stdzk.WorldACL(stdzk.PermAll),
	)
	if err != nil {
		t.Fatalf("Unable to create test ephemeral znode: %v", err)
	}
	select {
	case event := <-eventc:
		if want, have := stdzk.EventNodeChildrenChanged.String(), event.Type.String(); want != have {
			t.Errorf("want %s, have %s", want, have)
		}
	case <-time.After(20 * time.Millisecond):
		t.Errorf("expected incoming watch event, timeout occurred")
	}

}
<commit_msg>Minor fixes to zk integration test<commit_after>// +build integration

package zk

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	stdzk "github.com/samuel/go-zookeeper/zk"
)

var (
	host []string
)

func TestMain(m *testing.M) {
	flag.Parse()

	fmt.Println("Starting ZooKeeper server...")

	ts, err := stdzk.StartTestCluster(1, nil, nil)
	if err != nil {
		fmt.Printf("ZooKeeper server error: %v\n", err)
		os.Exit(1)
	}

	host = []string{fmt.Sprintf("localhost:%d", ts.Servers[0].Port)}
	code := m.Run()

	ts.Stop()
	os.Exit(code)
}

func TestCreateParentNodesOnServer(t *testing.T) {
	payload := [][]byte{[]byte("Payload"), []byte("Test")}
	c1, err := NewClient(host, logger, Payload(payload))
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}
	if c1 == nil {
		t.Fatal("Expected pointer to client, got nil")
	}
	defer c1.Stop()

	p, err := NewPublisher(c1, path, newFactory(""), logger)
	if err != nil {
		t.Fatalf("Unable to create Publisher: %v", err)
	}
	defer p.Stop()

	endpoints, err := p.Endpoints()
	if err != nil {
		t.Fatal(err)
	}
	if want, have := 0, len(endpoints); want != have {
		t.Errorf("want %d, have %d", want, have)
	}

	c2, err := NewClient(host, logger)
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}
	defer c2.Stop()
	data, _, err := c2.(*client).Get(path)
	if err != nil {
		t.Fatal(err)
	}
	// test Client implementation of CreateParentNodes. It should have created
	// our payload
	if bytes.Compare(data, payload[1]) != 0 {
		t.Errorf("want %s, have %s", payload[1], data)
	}

}

func TestCreateBadParentNodesOnServer(t *testing.T) {
	c, _ := NewClient(host, logger)
	defer c.Stop()

	_, err := NewPublisher(c, "invalid/path", newFactory(""), logger)

	if want, have := stdzk.ErrInvalidPath, err; want != have {
		t.Errorf("want %v, have %v", want, have)
	}
}

func TestCredentials1(t *testing.T) {
	acl := stdzk.DigestACL(stdzk.PermAll, "user", "secret")
	c, _ := NewClient(host, logger, ACL(acl), Credentials("user", "secret"))
	defer c.Stop()

	_, err := NewPublisher(c, "/acl-issue-test", newFactory(""), logger)

	if err != nil {
		t.Fatal(err)
	}
}

func TestCredentials2(t *testing.T) {
	acl := stdzk.DigestACL(stdzk.PermAll, "user", "secret")
	c, _ := NewClient(host, logger, ACL(acl))
	defer c.Stop()

	_, err := NewPublisher(c, "/acl-issue-test", newFactory(""), logger)

	if err != stdzk.ErrNoAuth {
		t.Errorf("want %v, have %v", stdzk.ErrNoAuth, err)
	}
}

func TestConnection(t *testing.T) {
	c, _ := NewClient(host, logger)
	c.Stop()

	_, err := NewPublisher(c, "/acl-issue-test", newFactory(""), logger)

	if err != ErrClientClosed {
		t.Errorf("want %v, have %v", ErrClientClosed, err)
	}
}

func TestGetEntriesOnServer(t *testing.T) {
	var instancePayload = "protocol://hostname:port/routing"

	c1, err := NewClient(host, logger)
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}

	defer c1.Stop()

	c2, err := NewClient(host, logger)
	p, err := NewPublisher(c2, path, newFactory(""), logger)
	if err != nil {
		t.Fatal(err)
	}
	defer c2.Stop()

	c2impl, _ := c2.(*client)
	_, err = c2impl.Create(
		path+"/instance1",
		[]byte(instancePayload),
		stdzk.FlagEphemeral|stdzk.FlagSequence,
		stdzk.WorldACL(stdzk.PermAll),
	)
	if err != nil {
		t.Fatalf("Unable to create test ephemeral znode 1: %v", err)
	}
	_, err = c2impl.Create(
		path+"/instance2",
		[]byte(instancePayload+"2"),
		stdzk.FlagEphemeral|stdzk.FlagSequence,
		stdzk.WorldACL(stdzk.PermAll),
	)
	if err != nil {
		t.Fatalf("Unable to create test ephemeral znode 2: %v", err)
	}

	time.Sleep(50 * time.Millisecond)

	endpoints, err := p.Endpoints()
	if err != nil {
		t.Fatal(err)
	}
	if want, have := 2, len(endpoints); want != have {
		t.Errorf("want %d, have %d", want, have)
	}
}

func TestGetEntriesPayloadOnServer(t *testing.T) {
	c, err := NewClient(host, logger)
	if err != nil {
		t.Fatalf("Connect returned error: %v", err)
	}
	_, eventc, err := c.GetEntries(path)
	if err != nil {
		t.Fatal(err)
	}
	_, err = c.(*client).Create(
		path+"/instance3",
		[]byte("just some payload"),
		stdzk.FlagEphemeral|stdzk.FlagSequence,
		stdzk.WorldACL(stdzk.PermAll),
	)
	if err != nil {
		t.Fatalf("Unable to create test ephemeral znode: %v", err)
	}
	select {
	case event := <-eventc:
		if want, have := stdzk.EventNodeChildrenChanged.String(), event.Type.String(); want != have {
			t.Errorf("want %s, have %s", want, have)
		}
	case <-time.After(20 * time.Millisecond):
		t.Errorf("expected incoming watch event, timeout occurred")
	}

}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Tigera, Inc. All rights reserved.
//
// 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 calc

import (
	"strings"

	"github.com/projectcalico/felix/dispatcher"
	"github.com/projectcalico/felix/proto"
	"github.com/projectcalico/libcalico-go/lib/backend/api"
	"github.com/projectcalico/libcalico-go/lib/backend/k8s/conversion"
	"github.com/projectcalico/libcalico-go/lib/backend/model"

	log "github.com/sirupsen/logrus"
)

type kind int

const (
	KindUnknown kind = iota
	KindServiceAccount
	KindNamespace
)

// ProfileDecoder takes updates from a dispatcher, determines if the profile is a Kubernetes Service Account or
// Kubernetes Namespace, and if it is, generates a dataplane update or remove for it.
type ProfileDecoder struct {
	callbacks passthruCallbacks
	converter conversion.Converter
}

func NewProfileDecoder(callbacks passthruCallbacks) *ProfileDecoder {
	return &ProfileDecoder{callbacks: callbacks, converter: conversion.Converter{true}}
}

func (p *ProfileDecoder) RegisterWith(d *dispatcher.Dispatcher) {
	d.Register(model.ProfileLabelsKey{}, p.OnUpdate)
}

func (p *ProfileDecoder) OnUpdate(update api.Update) (filterOut bool) {
	// This type assertion is safe because we only registered for ProfileLabels updates.
	key := update.Key.(model.ProfileLabelsKey)
	log.WithField("key", key.String()).Debug("Decoding ProfileLabels")
	id, knd := p.classifyProfile(key)
	switch knd {
	case KindUnknown:
		log.WithField("key", key.String()).Debug("Ignoring ProfileLabels")
	case KindServiceAccount:
		id := id.(proto.ServiceAccountID)
		if update.Value == nil {
			p.callbacks.OnServiceAccountRemove(id)
		} else {
			labels := update.Value.(map[string]string)
			msg := proto.ServiceAccountUpdate{Id: &id, Labels: decodeLabels(knd, labels)}
			p.callbacks.OnServiceAccountUpdate(&msg)
		}
	case KindNamespace:
		id := id.(proto.NamespaceID)
		if update.Value == nil {
			p.callbacks.OnNamespaceRemove(id)
		} else {
			labels := update.Value.(map[string]string)
			msg := proto.NamespaceUpdate{Id: &id, Labels: decodeLabels(knd, labels)}
			p.callbacks.OnNamespaceUpdate(&msg)
		}
	}
	return false
}

func (p *ProfileDecoder) classifyProfile(key model.ProfileLabelsKey) (interface{}, kind) {
	namespace, name, err := p.converter.ProfileNameToServiceAccount(key.Name)
	if err == nil {
		return proto.ServiceAccountID{Name: name, Namespace: namespace}, KindServiceAccount
	}
	name, err = p.converter.ProfileNameToNamespace(key.Name)
	if err == nil {
		return proto.NamespaceID{Name: name}, KindNamespace
	}
	return nil, KindUnknown
}

// decodeLabels strips the special prefix we add to Profile labels when converting. This gives us the original labels on
// the ServiceAccount or Namespace object.
func decodeLabels(k kind, in map[string]string) map[string]string {
	var prefix string
	if k == KindServiceAccount {
		prefix = conversion.ServiceAccountLabelPrefix
	} else if k == KindNamespace {
		prefix = conversion.NamespaceLabelPrefix
	}
	out := make(map[string]string)
	for k, v := range in {
		k = strings.TrimPrefix(k, prefix)
		out[k] = v
	}
	return out
}
<commit_msg>Refactor profile_decoder.go to use type switch<commit_after>// Copyright (c) 2018 Tigera, Inc. All rights reserved.
//
// 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 calc

import (
	"strings"

	"github.com/projectcalico/felix/dispatcher"
	"github.com/projectcalico/felix/proto"
	"github.com/projectcalico/libcalico-go/lib/backend/api"
	"github.com/projectcalico/libcalico-go/lib/backend/k8s/conversion"
	"github.com/projectcalico/libcalico-go/lib/backend/model"

	log "github.com/sirupsen/logrus"
)

// ProfileDecoder takes updates from a dispatcher, determines if the profile is a Kubernetes Service Account or
// Kubernetes Namespace, and if it is, generates a dataplane update or remove for it.
type ProfileDecoder struct {
	callbacks passthruCallbacks
	converter conversion.Converter
}

func NewProfileDecoder(callbacks passthruCallbacks) *ProfileDecoder {
	return &ProfileDecoder{callbacks: callbacks, converter: conversion.Converter{true}}
}

func (p *ProfileDecoder) RegisterWith(d *dispatcher.Dispatcher) {
	d.Register(model.ProfileLabelsKey{}, p.OnUpdate)
}

func (p *ProfileDecoder) OnUpdate(update api.Update) (filterOut bool) {
	// This type assertion is safe because we only registered for ProfileLabels updates.
	key := update.Key.(model.ProfileLabelsKey)
	log.WithField("key", key.String()).Debug("Decoding ProfileLabels")
	idInterface := p.classifyProfile(key)
	switch id := idInterface.(type) {
	case nil:
		log.WithField("key", key.String()).Debug("Ignoring ProfileLabels")
	case proto.ServiceAccountID:
		if update.Value == nil {
			p.callbacks.OnServiceAccountRemove(id)
		} else {
			labels := update.Value.(map[string]string)
			msg := proto.ServiceAccountUpdate{
				Id: &id, Labels: decodeLabels(conversion.ServiceAccountLabelPrefix, labels)}
			p.callbacks.OnServiceAccountUpdate(&msg)
		}
	case proto.NamespaceID:
		if update.Value == nil {
			p.callbacks.OnNamespaceRemove(id)
		} else {
			labels := update.Value.(map[string]string)
			msg := proto.NamespaceUpdate{
				Id: &id, Labels: decodeLabels(conversion.NamespaceLabelPrefix, labels)}
			p.callbacks.OnNamespaceUpdate(&msg)
		}
	}
	return false
}

func (p *ProfileDecoder) classifyProfile(key model.ProfileLabelsKey) interface{} {
	namespace, name, err := p.converter.ProfileNameToServiceAccount(key.Name)
	if err == nil {
		return proto.ServiceAccountID{Name: name, Namespace: namespace}
	}
	name, err = p.converter.ProfileNameToNamespace(key.Name)
	if err == nil {
		return proto.NamespaceID{Name: name}
	}
	return nil
}

// decodeLabels strips the special prefix we add to Profile labels when converting. This gives us the original labels on
// the ServiceAccount or Namespace object.
func decodeLabels(prefix string, in map[string]string) map[string]string {
	out := make(map[string]string)
	for k, v := range in {
		k = strings.TrimPrefix(k, prefix)
		out[k] = v
	}
	return out
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Niklas Wolber
// This file is licensed under the MIT license.
// See the LICENSE file for more information.

package xCUTEr

import (
	"bytes"
	"context"
	"io"
	"log"
	"sync"
	"time"

	sched "github.com/nwolber/cron"
	"github.com/nwolber/xCUTEr/flunc"
	"github.com/nwolber/xCUTEr/job"
	"github.com/nwolber/xCUTEr/telemetry"
)

type jobInfo struct {
	file string
	f    flunc.Flunc
	c    *job.Config

	telemetry bool
	events    *telemetry.EventStore
}

type schedInfo struct {
	id string
	j  *jobInfo
}

func (info *schedInfo) Config() *job.Config {
	return info.j.c
}

// runInfo holds information about a single run of a job.
type runInfo struct {
	// The executor the job runs on.
	e *executor
	// jobInfo about the running job.
	j *jobInfo
	// Function to cancel the job early.
	cancel context.CancelFunc
	// Start and finish time of the job.
	start, stop time.Time
	// Job output.
	output bytes.Buffer
}

// Config returns the running Config .
func (info *runInfo) Config() *job.Config {
	return info.j.c
}

// Start time.
func (info *runInfo) Start() time.Time {
	return info.start
}

// Finish time.
func (info *runInfo) Stop() time.Time {
	return info.stop
}

// Job output.
func (info *runInfo) Output() string {
	return info.output.String()
}

func (info *runInfo) run() {
	ctx, cancel := context.WithCancel(info.e.mainCtx)

	output, ok := ctx.Value(outputKey).(io.Writer)
	if ok {
		ctx = context.WithValue(ctx, outputKey, io.MultiWriter(output, &info.output))
	} else {
		// ctx = context.WithValue(ctx, "output", &info.output)
	}

	info.cancel = cancel

	if !info.e.addRunning(info) {
		log.Printf("another instance of %q is still running, consider adding/lowering the timeout", info.j.c.Name)
		return
	}

	defer func() {
		// release resources
		cancel()
		info.stop = time.Now()
		info.e.removeRunning(info)
		info.e.addComplete(info)
	}()
	info.start = time.Now()
	info.j.f(ctx)
}

type executor struct {
	// The main context used to cancel all jobs at once.
	mainCtx context.Context
	// Whether jobs need to be activated manually.
	manualActive bool
	// Number of completed runInfos kept.
	maxCompleted uint32
	// Functions to start and stop the scheduler.
	Start, Stop func()
	// Function to run a runInfo.
	run func(info *runInfo)
	// Function to schedule a new runInfo.
	schedule func(schedule string, f func()) (string, error)
	// Function to remove an existing runInfo.
	remove func(string)

	// List of inactive scheduled jobs. Only used if manualActive is true.
	inactive  map[string]*schedInfo
	mInactive sync.Mutex

	// List of scheduled jobs.
	scheduled map[string]*schedInfo
	mSched    sync.Mutex

	// List of currently running jobs.
	running map[string]*runInfo
	mRun    sync.Mutex

	// List of completed runInfos. A maximum of maxCompleted runInfos is kept.
	completed  []*runInfo
	mCompleted sync.Mutex
}

func newExecutor(ctx context.Context) *executor {
	run := func(info *runInfo) { info.run() }

	cron := sched.New()
	return &executor{
		mainCtx:      ctx,
		maxCompleted: 10,
		Start:        cron.Start,
		Stop:         cron.Stop,
		run:          run,
		schedule:     cron.AddFunc,
		remove:       cron.Remove,
		inactive:     make(map[string]*schedInfo),
		scheduled:    make(map[string]*schedInfo),
		running:      make(map[string]*runInfo),
	}
}

// parse parses the given file and stores the execution tree in the jobInfo.
func parse(file string) (*jobInfo, error) {
	c, err := job.ReadConfig(file)
	if err != nil {
		return nil, err
	}

	useTelemetry := c.Telemetry

	start := time.Now()
	var (
		f      flunc.Flunc
		events *telemetry.EventStore
	)
	if useTelemetry {
		f, events, err = telemetry.Instrument(c)
	} else {
		f, err = c.ExecutionTree()
	}

	if err != nil {
		return nil, err
	}

	stop := time.Now()
	log.Println("job preparation took", stop.Sub(start))

	return &jobInfo{
		file: file,
		c:    c,
		f:    f,

		telemetry: useTelemetry,
		events:    events,
	}, nil
}

// scheduleBody returns a function that can be used by the cron
// scheduler to execute the Job.
func scheduleBody(e *executor, j *jobInfo) func() {
	return func() {
		log.Println(j.c.Name, "woke up")
		info := &runInfo{
			e: e,
			j: j,
		}
		e.run(info)
		log.Println(j.c.Name, "finished in", time.Now().Sub(info.start))
	}
}

// Run either executes the job directly if either the job's schedule
// is "once" or the once parameter is true. Otherwise the job is
// scheduled as if Add would have been called.
func (e *executor) Run(j *jobInfo, once bool) {
	if j.c.Schedule == "once" || once {
		info := &runInfo{
			e: e,
			j: j,
		}
		e.run(info)
	} else {
		e.Add(j)
	}
}

// Add schedules the job.
func (e *executor) Add(j *jobInfo) error {
	if j.c.Schedule == "once" {
		info := &runInfo{
			e: e,
			j: j,
		}
		go e.run(info)
	} else {
		if e.manualActive {
			e.addInactive(&schedInfo{
				j: j,
			})
			return nil
		}

		id, err := e.schedule(j.c.Schedule, scheduleBody(e, j))

		if err != nil {
			return err
		}

		s := &schedInfo{
			id: id,
			j:  j,
		}
		log.Println(j.c.Name, "scheduled")
		e.addScheduled(s)
	}

	return nil
}

// Remove removes all resources associated with the given job file.
func (e *executor) Remove(file string) {
	log.Println("remove", file)
	if info := e.isRunning(file); info != nil {
		e.removeRunning(info)
		info.cancel()
		log.Println("found running", info.j.c.Name)
	}

	if info := e.isScheduled(file); info != nil {
		e.removeScheduled(info)
		e.remove(info.id)

		log.Println("found scheduled", info.j.c.Name)
	}

	if info := e.isInactive(file); info != nil {
		e.removeInactive(info)
		log.Println("found inactive", info.j.c.Name)
	}
}

// Activates the job associated with the job file.
func (e *executor) Activate(file string) {
	log.Println("activate", file)
	if info := e.isInactive(file); info != nil {
		e.removeInactive(info)
		e.schedule(info.j.c.Schedule, scheduleBody(e, info.j))

		e.addScheduled(info)
	} else {
		log.Println("didn't find ", file, "in inactive list")
	}
}

// Deactivates the job associated with the job file.
func (e *executor) Deactivate(file string) {
	log.Println("deactivate", file)

	if info := e.isScheduled(file); info != nil {
		e.removeScheduled(info)
		e.remove(info.id)
		e.addInactive(info)
	} else {
		log.Println("didn't find ", file, "in active list")
	}
}

// Start runs a job. If the return value is false,
// another instance of this job is still running.
func (e *executor) addRunning(info *runInfo) bool {
	e.mRun.Lock()
	defer e.mRun.Unlock()

	if _, ok := e.running[info.j.file]; ok {
		return false
	}

	e.running[info.j.file] = info
	return true
}

// Stop halts execution of a job.
func (e *executor) removeRunning(info *runInfo) {
	e.mRun.Lock()
	defer e.mRun.Unlock()

	delete(e.running, info.j.file)
}

// GetRunning returns a runInfo, if there is a running job.
func (e *executor) isRunning(file string) *runInfo {
	e.mRun.Lock()
	defer e.mRun.Unlock()
	return e.running[file]
}

func (e *executor) GetRunning() []*runInfo {
	e.mRun.Lock()
	defer e.mRun.Unlock()

	running := make([]*runInfo, len(e.running))
	i := 0
	for _, info := range e.running {
		running[i] = info
		i++
	}
	return running
}

func (e *executor) addComplete(info *runInfo) {
	e.mCompleted.Lock()
	defer e.mCompleted.Unlock()

	e.completed = append(e.completed, info)

	if l, max := len(e.completed), int(e.maxCompleted); max > 0 && max < l {
		e.completed = e.completed[l-max:]
	}
}

func (e *executor) GetCompleted() []*runInfo {
	e.mCompleted.Lock()
	defer e.mCompleted.Unlock()

	completed := make([]*runInfo, len(e.completed))
	copy(completed, e.completed)
	return completed
}

func (e *executor) addScheduled(info *schedInfo) {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	e.scheduled[info.j.file] = info
}

func (e *executor) removeScheduled(info *schedInfo) {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	delete(e.scheduled, info.j.file)
}

func (e *executor) isScheduled(file string) *schedInfo {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	return e.scheduled[file]
}

func (e *executor) GetScheduled() []*schedInfo {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	scheduled := make([]*schedInfo, len(e.scheduled))
	i := 0
	for _, info := range e.scheduled {
		scheduled[i] = info
	}
	return scheduled
}

func (e *executor) addInactive(info *schedInfo) {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()

	e.inactive[info.j.file] = info
}

func (e *executor) removeInactive(info *schedInfo) {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()

	delete(e.inactive, info.j.file)
}

func (e *executor) isInactive(file string) *schedInfo {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()
	return e.inactive[file]
}

func (e *executor) GetInactive() []*schedInfo {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()

	inactive := make([]*schedInfo, len(e.inactive))
	i := 0
	for _, info := range e.inactive {
		inactive[i] = info
	}
	return inactive
}
<commit_msg>Reset telemety events after job completed<commit_after>// Copyright (c) 2016 Niklas Wolber
// This file is licensed under the MIT license.
// See the LICENSE file for more information.

package xCUTEr

import (
	"bytes"
	"context"
	"io"
	"log"
	"sync"
	"time"

	sched "github.com/nwolber/cron"
	"github.com/nwolber/xCUTEr/flunc"
	"github.com/nwolber/xCUTEr/job"
	"github.com/nwolber/xCUTEr/telemetry"
)

type jobInfo struct {
	file string
	f    flunc.Flunc
	c    *job.Config

	telemetry bool
	events    *telemetry.EventStore
}

type schedInfo struct {
	id string
	j  *jobInfo
}

func (info *schedInfo) Config() *job.Config {
	return info.j.c
}

// runInfo holds information about a single run of a job.
type runInfo struct {
	// The executor the job runs on.
	e *executor
	// jobInfo about the running job.
	j *jobInfo
	// Function to cancel the job early.
	cancel context.CancelFunc
	// Start and finish time of the job.
	start, stop time.Time
	// Job output.
	output bytes.Buffer
}

// Config returns the running Config .
func (info *runInfo) Config() *job.Config {
	return info.j.c
}

// Start time.
func (info *runInfo) Start() time.Time {
	return info.start
}

// Finish time.
func (info *runInfo) Stop() time.Time {
	return info.stop
}

// Job output.
func (info *runInfo) Output() string {
	return info.output.String()
}

func (info *runInfo) run() {
	ctx, cancel := context.WithCancel(info.e.mainCtx)

	output, ok := ctx.Value(outputKey).(io.Writer)
	if ok {
		ctx = context.WithValue(ctx, outputKey, io.MultiWriter(output, &info.output))
	} else {
		// ctx = context.WithValue(ctx, "output", &info.output)
	}

	info.cancel = cancel

	if !info.e.addRunning(info) {
		log.Printf("another instance of %q is still running, consider adding/lowering the timeout", info.j.c.Name)
		return
	}

	defer func() {
		// release resources
		cancel()
		info.stop = time.Now()
		info.j.events.Reset()
		info.e.removeRunning(info)
		info.e.addComplete(info)
	}()
	info.start = time.Now()
	info.j.f(ctx)
}

type executor struct {
	// The main context used to cancel all jobs at once.
	mainCtx context.Context
	// Whether jobs need to be activated manually.
	manualActive bool
	// Number of completed runInfos kept.
	maxCompleted uint32
	// Functions to start and stop the scheduler.
	Start, Stop func()
	// Function to run a runInfo.
	run func(info *runInfo)
	// Function to schedule a new runInfo.
	schedule func(schedule string, f func()) (string, error)
	// Function to remove an existing runInfo.
	remove func(string)

	// List of inactive scheduled jobs. Only used if manualActive is true.
	inactive  map[string]*schedInfo
	mInactive sync.Mutex

	// List of scheduled jobs.
	scheduled map[string]*schedInfo
	mSched    sync.Mutex

	// List of currently running jobs.
	running map[string]*runInfo
	mRun    sync.Mutex

	// List of completed runInfos. A maximum of maxCompleted runInfos is kept.
	completed  []*runInfo
	mCompleted sync.Mutex
}

func newExecutor(ctx context.Context) *executor {
	run := func(info *runInfo) { info.run() }

	cron := sched.New()
	return &executor{
		mainCtx:      ctx,
		maxCompleted: 10,
		Start:        cron.Start,
		Stop:         cron.Stop,
		run:          run,
		schedule:     cron.AddFunc,
		remove:       cron.Remove,
		inactive:     make(map[string]*schedInfo),
		scheduled:    make(map[string]*schedInfo),
		running:      make(map[string]*runInfo),
	}
}

// parse parses the given file and stores the execution tree in the jobInfo.
func parse(file string) (*jobInfo, error) {
	c, err := job.ReadConfig(file)
	if err != nil {
		return nil, err
	}

	useTelemetry := c.Telemetry

	start := time.Now()
	var (
		f      flunc.Flunc
		events *telemetry.EventStore
	)
	if useTelemetry {
		f, events, err = telemetry.Instrument(c)
	} else {
		f, err = c.ExecutionTree()
	}

	if err != nil {
		return nil, err
	}

	stop := time.Now()
	log.Println("job preparation took", stop.Sub(start))

	return &jobInfo{
		file: file,
		c:    c,
		f:    f,

		telemetry: useTelemetry,
		events:    events,
	}, nil
}

// scheduleBody returns a function that can be used by the cron
// scheduler to execute the Job.
func scheduleBody(e *executor, j *jobInfo) func() {
	return func() {
		log.Println(j.c.Name, "woke up")
		info := &runInfo{
			e: e,
			j: j,
		}
		e.run(info)
		log.Println(j.c.Name, "finished in", time.Now().Sub(info.start))
	}
}

// Run either executes the job directly if either the job's schedule
// is "once" or the once parameter is true. Otherwise the job is
// scheduled as if Add would have been called.
func (e *executor) Run(j *jobInfo, once bool) {
	if j.c.Schedule == "once" || once {
		info := &runInfo{
			e: e,
			j: j,
		}
		e.run(info)
	} else {
		e.Add(j)
	}
}

// Add schedules the job.
func (e *executor) Add(j *jobInfo) error {
	if j.c.Schedule == "once" {
		info := &runInfo{
			e: e,
			j: j,
		}
		go e.run(info)
	} else {
		if e.manualActive {
			e.addInactive(&schedInfo{
				j: j,
			})
			return nil
		}

		id, err := e.schedule(j.c.Schedule, scheduleBody(e, j))

		if err != nil {
			return err
		}

		s := &schedInfo{
			id: id,
			j:  j,
		}
		log.Println(j.c.Name, "scheduled")
		e.addScheduled(s)
	}

	return nil
}

// Remove removes all resources associated with the given job file.
func (e *executor) Remove(file string) {
	log.Println("remove", file)
	if info := e.isRunning(file); info != nil {
		e.removeRunning(info)
		info.cancel()
		log.Println("found running", info.j.c.Name)
	}

	if info := e.isScheduled(file); info != nil {
		e.removeScheduled(info)
		e.remove(info.id)

		log.Println("found scheduled", info.j.c.Name)
	}

	if info := e.isInactive(file); info != nil {
		e.removeInactive(info)
		log.Println("found inactive", info.j.c.Name)
	}
}

// Activates the job associated with the job file.
func (e *executor) Activate(file string) {
	log.Println("activate", file)
	if info := e.isInactive(file); info != nil {
		e.removeInactive(info)
		e.schedule(info.j.c.Schedule, scheduleBody(e, info.j))

		e.addScheduled(info)
	} else {
		log.Println("didn't find ", file, "in inactive list")
	}
}

// Deactivates the job associated with the job file.
func (e *executor) Deactivate(file string) {
	log.Println("deactivate", file)

	if info := e.isScheduled(file); info != nil {
		e.removeScheduled(info)
		e.remove(info.id)
		e.addInactive(info)
	} else {
		log.Println("didn't find ", file, "in active list")
	}
}

// Start runs a job. If the return value is false,
// another instance of this job is still running.
func (e *executor) addRunning(info *runInfo) bool {
	e.mRun.Lock()
	defer e.mRun.Unlock()

	if _, ok := e.running[info.j.file]; ok {
		return false
	}

	e.running[info.j.file] = info
	return true
}

// Stop halts execution of a job.
func (e *executor) removeRunning(info *runInfo) {
	e.mRun.Lock()
	defer e.mRun.Unlock()

	delete(e.running, info.j.file)
}

// GetRunning returns a runInfo, if there is a running job.
func (e *executor) isRunning(file string) *runInfo {
	e.mRun.Lock()
	defer e.mRun.Unlock()
	return e.running[file]
}

func (e *executor) GetRunning() []*runInfo {
	e.mRun.Lock()
	defer e.mRun.Unlock()

	running := make([]*runInfo, len(e.running))
	i := 0
	for _, info := range e.running {
		running[i] = info
		i++
	}
	return running
}

func (e *executor) addComplete(info *runInfo) {
	e.mCompleted.Lock()
	defer e.mCompleted.Unlock()

	e.completed = append(e.completed, info)

	if l, max := len(e.completed), int(e.maxCompleted); max > 0 && max < l {
		e.completed = e.completed[l-max:]
	}
}

func (e *executor) GetCompleted() []*runInfo {
	e.mCompleted.Lock()
	defer e.mCompleted.Unlock()

	completed := make([]*runInfo, len(e.completed))
	copy(completed, e.completed)
	return completed
}

func (e *executor) addScheduled(info *schedInfo) {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	e.scheduled[info.j.file] = info
}

func (e *executor) removeScheduled(info *schedInfo) {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	delete(e.scheduled, info.j.file)
}

func (e *executor) isScheduled(file string) *schedInfo {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	return e.scheduled[file]
}

func (e *executor) GetScheduled() []*schedInfo {
	e.mSched.Lock()
	defer e.mSched.Unlock()

	scheduled := make([]*schedInfo, len(e.scheduled))
	i := 0
	for _, info := range e.scheduled {
		scheduled[i] = info
	}
	return scheduled
}

func (e *executor) addInactive(info *schedInfo) {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()

	e.inactive[info.j.file] = info
}

func (e *executor) removeInactive(info *schedInfo) {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()

	delete(e.inactive, info.j.file)
}

func (e *executor) isInactive(file string) *schedInfo {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()
	return e.inactive[file]
}

func (e *executor) GetInactive() []*schedInfo {
	e.mInactive.Lock()
	defer e.mInactive.Unlock()

	inactive := make([]*schedInfo, len(e.inactive))
	i := 0
	for _, info := range e.inactive {
		inactive[i] = info
	}
	return inactive
}
<|endoftext|>


<commit_before>// +build windows

package main

import (
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	maxScrapeDuration time.Duration
	collectors        map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system,textfile"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
	scrapeTimeoutDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_timeout"),
		"wmi_exporter: Whether the collector timed out.",
		[]string{"collector"},
		nil,
	)
	snapshotDuration = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "perflib_snapshot_duration_seconds"),
		"Duration of perflib snapshot capture",
		nil,
		nil,
	)

	// This can be removed when client_golang exposes this on Windows
	// (See https://github.com/prometheus/client_golang/issues/376)
	startTime     = float64(time.Now().Unix())
	startTimeDesc = prometheus.NewDesc(
		"process_start_time_seconds",
		"Start time of the process since unix epoch in seconds.",
		nil,
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

type collectorOutcome int

const (
	pending collectorOutcome = iota
	success
	failed
)

// Collect sends the collected metrics from each of the collectors to
// prometheus.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	ch <- prometheus.MustNewConstMetric(
		startTimeDesc,
		prometheus.CounterValue,
		startTime,
	)

	t := time.Now()
	cs := make([]string, 0, len(coll.collectors))
	for name := range coll.collectors {
		cs = append(cs, name)
	}
	scrapeContext, err := collector.PrepareScrapeContext(cs)
	ch <- prometheus.MustNewConstMetric(
		snapshotDuration,
		prometheus.GaugeValue,
		time.Since(t).Seconds(),
	)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(scrapeSuccessDesc, fmt.Errorf("failed to prepare scrape: %v", err))
		return
	}

	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	collectorOutcomes := make(map[string]collectorOutcome)
	for name := range coll.collectors {
		collectorOutcomes[name] = pending
	}

	metricsBuffer := make(chan prometheus.Metric)
	l := sync.Mutex{}
	finished := false
	go func() {
		for m := range metricsBuffer {
			l.Lock()
			if !finished {
				ch <- m
			}
			l.Unlock()
		}
	}()

	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			defer wg.Done()
			outcome := execute(name, c, scrapeContext, metricsBuffer)
			l.Lock()
			if !finished {
				collectorOutcomes[name] = outcome
			}
			l.Unlock()
		}(name, c)
	}

	allDone := make(chan struct{})
	go func() {
		wg.Wait()
		close(allDone)
		close(metricsBuffer)
	}()

	// Wait until either all collectors finish, or timeout expires
	select {
	case <-allDone:
	case <-time.After(coll.maxScrapeDuration):
	}

	l.Lock()
	finished = true

	remainingCollectorNames := make([]string, 0)
	for name, outcome := range collectorOutcomes {
		var successValue, timeoutValue float64
		if outcome == pending {
			timeoutValue = 1.0
			remainingCollectorNames = append(remainingCollectorNames, name)
		}
		if outcome == success {
			successValue = 1.0
		}

		ch <- prometheus.MustNewConstMetric(
			scrapeSuccessDesc,
			prometheus.GaugeValue,
			successValue,
			name,
		)
		ch <- prometheus.MustNewConstMetric(
			scrapeTimeoutDesc,
			prometheus.GaugeValue,
			timeoutValue,
			name,
		)
	}

	if len(remainingCollectorNames) > 0 {
		log.Warn("Collection timed out, still waiting for ", remainingCollectorNames)
	}

	l.Unlock()
}

func execute(name string, c collector.Collector, ctx *collector.ScrapeContext, ch chan<- prometheus.Metric) collectorOutcome {
	t := time.Now()
	err := c.Collect(ctx, ch)
	duration := time.Since(t).Seconds()
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration,
		name,
	)

	if err != nil {
		log.Errorf("collector %s failed after %fs: %s", name, duration, err)
		return failed
	}
	log.Debugf("collector %s succeeded after %fs.", name, duration)
	return success
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		c, err := collector.Build(name)
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}

	return collectors, nil
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.AllowMissingFields = true
	wmi.DefaultClient.SWbemServicesClient = s
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"telemetry.addr",
			"host:port for WMI exporter.",
		).Default(":9182").String()
		metricsPath = kingpin.Flag(
			"telemetry.path",
			"URL path for surfacing collected metrics.",
		).Default("/metrics").String()
		enabledCollectors = kingpin.Flag(
			"collectors.enabled",
			"Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default.").
			Default(defaultCollectors).String()
		printCollectors = kingpin.Flag(
			"collectors.print",
			"If true, print available collectors and exit.",
		).Bool()
		timeoutMargin = kingpin.Flag(
			"scrape.timeout-margin",
			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
		).Default("0.5").Float64()
	)

	log.AddFlags(kingpin.CommandLine)
	kingpin.Version(version.Print("wmi_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *printCollectors {
		collectors := collector.Available()
		collectorNames := make(sort.StringSlice, 0, len(collectors))
		for _, n := range collectors {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go func() {
			err = svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
			if err != nil {
				log.Errorf("Failed to start service: %v", err)
			}
		}()
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	h := &metricsHandler{
		timeoutMargin: *timeoutMargin,
		collectorFactory: func(timeout time.Duration) *WmiCollector {
			return &WmiCollector{
				collectors:        collectors,
				maxScrapeDuration: timeout,
			}
		},
	}

	http.Handle(*metricsPath, h)
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, err := fmt.Fprintln(w, `{"status":"ok"}`)
	if err != nil {
		log.Debugf("Failed to write to stream: %v", err)
	}
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}

type metricsHandler struct {
	timeoutMargin    float64
	collectorFactory func(timeout time.Duration) *WmiCollector
}

func (mh *metricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	const defaultTimeout = 10.0

	var timeoutSeconds float64
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		var err error
		timeoutSeconds, err = strconv.ParseFloat(v, 64)
		if err != nil {
			log.Warnf("Couldn't parse X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %f", v, defaultTimeout)
		}
	}
	if timeoutSeconds == 0 {
		timeoutSeconds = defaultTimeout
	}
	timeoutSeconds = timeoutSeconds - mh.timeoutMargin

	reg := prometheus.NewRegistry()
	reg.MustRegister(mh.collectorFactory(time.Duration(timeoutSeconds * float64(time.Second))))
	reg.MustRegister(
		prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}),
		prometheus.NewGoCollector(),
		version.NewCollector("wmi_exporter"),
	)

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)
}
<commit_msg>Add option to limit concurrent requests<commit_after>// +build windows

package main

import (
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	maxScrapeDuration time.Duration
	collectors        map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system,textfile"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
	scrapeTimeoutDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_timeout"),
		"wmi_exporter: Whether the collector timed out.",
		[]string{"collector"},
		nil,
	)
	snapshotDuration = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "perflib_snapshot_duration_seconds"),
		"Duration of perflib snapshot capture",
		nil,
		nil,
	)

	// This can be removed when client_golang exposes this on Windows
	// (See https://github.com/prometheus/client_golang/issues/376)
	startTime     = float64(time.Now().Unix())
	startTimeDesc = prometheus.NewDesc(
		"process_start_time_seconds",
		"Start time of the process since unix epoch in seconds.",
		nil,
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

type collectorOutcome int

const (
	pending collectorOutcome = iota
	success
	failed
)

// Collect sends the collected metrics from each of the collectors to
// prometheus.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	ch <- prometheus.MustNewConstMetric(
		startTimeDesc,
		prometheus.CounterValue,
		startTime,
	)

	t := time.Now()
	cs := make([]string, 0, len(coll.collectors))
	for name := range coll.collectors {
		cs = append(cs, name)
	}
	scrapeContext, err := collector.PrepareScrapeContext(cs)
	ch <- prometheus.MustNewConstMetric(
		snapshotDuration,
		prometheus.GaugeValue,
		time.Since(t).Seconds(),
	)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(scrapeSuccessDesc, fmt.Errorf("failed to prepare scrape: %v", err))
		return
	}

	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	collectorOutcomes := make(map[string]collectorOutcome)
	for name := range coll.collectors {
		collectorOutcomes[name] = pending
	}

	metricsBuffer := make(chan prometheus.Metric)
	l := sync.Mutex{}
	finished := false
	go func() {
		for m := range metricsBuffer {
			l.Lock()
			if !finished {
				ch <- m
			}
			l.Unlock()
		}
	}()

	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			defer wg.Done()
			outcome := execute(name, c, scrapeContext, metricsBuffer)
			l.Lock()
			if !finished {
				collectorOutcomes[name] = outcome
			}
			l.Unlock()
		}(name, c)
	}

	allDone := make(chan struct{})
	go func() {
		wg.Wait()
		close(allDone)
		close(metricsBuffer)
	}()

	// Wait until either all collectors finish, or timeout expires
	select {
	case <-allDone:
	case <-time.After(coll.maxScrapeDuration):
	}

	l.Lock()
	finished = true

	remainingCollectorNames := make([]string, 0)
	for name, outcome := range collectorOutcomes {
		var successValue, timeoutValue float64
		if outcome == pending {
			timeoutValue = 1.0
			remainingCollectorNames = append(remainingCollectorNames, name)
		}
		if outcome == success {
			successValue = 1.0
		}

		ch <- prometheus.MustNewConstMetric(
			scrapeSuccessDesc,
			prometheus.GaugeValue,
			successValue,
			name,
		)
		ch <- prometheus.MustNewConstMetric(
			scrapeTimeoutDesc,
			prometheus.GaugeValue,
			timeoutValue,
			name,
		)
	}

	if len(remainingCollectorNames) > 0 {
		log.Warn("Collection timed out, still waiting for ", remainingCollectorNames)
	}

	l.Unlock()
}

func execute(name string, c collector.Collector, ctx *collector.ScrapeContext, ch chan<- prometheus.Metric) collectorOutcome {
	t := time.Now()
	err := c.Collect(ctx, ch)
	duration := time.Since(t).Seconds()
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration,
		name,
	)

	if err != nil {
		log.Errorf("collector %s failed after %fs: %s", name, duration, err)
		return failed
	}
	log.Debugf("collector %s succeeded after %fs.", name, duration)
	return success
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		c, err := collector.Build(name)
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}

	return collectors, nil
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.AllowMissingFields = true
	wmi.DefaultClient.SWbemServicesClient = s
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"telemetry.addr",
			"host:port for WMI exporter.",
		).Default(":9182").String()
		metricsPath = kingpin.Flag(
			"telemetry.path",
			"URL path for surfacing collected metrics.",
		).Default("/metrics").String()
		maxRequests = kingpin.Flag(
			"telemetry.max-requests",
			"Maximum number of concurrent requests. 0 to disable.",
		).Default("5").Int()
		enabledCollectors = kingpin.Flag(
			"collectors.enabled",
			"Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default.").
			Default(defaultCollectors).String()
		printCollectors = kingpin.Flag(
			"collectors.print",
			"If true, print available collectors and exit.",
		).Bool()
		timeoutMargin = kingpin.Flag(
			"scrape.timeout-margin",
			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
		).Default("0.5").Float64()
	)

	log.AddFlags(kingpin.CommandLine)
	kingpin.Version(version.Print("wmi_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *printCollectors {
		collectors := collector.Available()
		collectorNames := make(sort.StringSlice, 0, len(collectors))
		for _, n := range collectors {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go func() {
			err = svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
			if err != nil {
				log.Errorf("Failed to start service: %v", err)
			}
		}()
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	h := &metricsHandler{
		timeoutMargin: *timeoutMargin,
		collectorFactory: func(timeout time.Duration) *WmiCollector {
			return &WmiCollector{
				collectors:        collectors,
				maxScrapeDuration: timeout,
			}
		},
	}

	http.HandleFunc(*metricsPath, withConcurrencyLimit(*maxRequests, h.ServeHTTP))
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, err := fmt.Fprintln(w, `{"status":"ok"}`)
	if err != nil {
		log.Debugf("Failed to write to stream: %v", err)
	}
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

func withConcurrencyLimit(n int, next http.HandlerFunc) http.HandlerFunc {
	if n <= 0 {
		return next
	}

	sem := make(chan struct{}, n)
	return func(w http.ResponseWriter, r *http.Request) {
		select {
		case sem <- struct{}{}:
			defer func() { <-sem }()
		default:
			w.WriteHeader(http.StatusServiceUnavailable)
			_, _ = w.Write([]byte("Too many concurrent requests"))
			return
		}
		next(w, r)
	}
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}

type metricsHandler struct {
	timeoutMargin    float64
	collectorFactory func(timeout time.Duration) *WmiCollector
}

func (mh *metricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	const defaultTimeout = 10.0

	var timeoutSeconds float64
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		var err error
		timeoutSeconds, err = strconv.ParseFloat(v, 64)
		if err != nil {
			log.Warnf("Couldn't parse X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %f", v, defaultTimeout)
		}
	}
	if timeoutSeconds == 0 {
		timeoutSeconds = defaultTimeout
	}
	timeoutSeconds = timeoutSeconds - mh.timeoutMargin

	reg := prometheus.NewRegistry()
	reg.MustRegister(mh.collectorFactory(time.Duration(timeoutSeconds * float64(time.Second))))
	reg.MustRegister(
		prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}),
		prometheus.NewGoCollector(),
		version.NewCollector("wmi_exporter"),
	)

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

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

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

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

package compose

import (
	"context"
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	ctxkube "github.com/docker/buildx/driver/kubernetes/context"
	"github.com/docker/buildx/store"
	"github.com/docker/buildx/store/storeutil"
	"github.com/docker/cli/cli/command"
	"github.com/docker/cli/cli/context/docker"
	ctxstore "github.com/docker/cli/cli/context/store"
	dockerclient "github.com/docker/docker/client"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/errgroup"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/docker/buildx/build"
	"github.com/docker/buildx/driver"
	_ "github.com/docker/buildx/driver/docker"           //nolint:blank-imports
	_ "github.com/docker/buildx/driver/docker-container" //nolint:blank-imports
	_ "github.com/docker/buildx/driver/kubernetes"       //nolint:blank-imports
	xprogress "github.com/docker/buildx/util/progress"
)

func (s *composeService) doBuildBuildkit(ctx context.Context, opts map[string]build.Options, mode string) (map[string]string, error) {
	dis, err := s.getDrivers(ctx)
	if err != nil {
		return nil, err
	}

	// Progress needs its own context that lives longer than the
	// build one otherwise it won't read all the messages from
	// build and will lock
	progressCtx, cancel := context.WithCancel(context.Background())
	defer cancel()
	w := xprogress.NewPrinter(progressCtx, s.stdout(), os.Stdout, mode)

	response, err := build.Build(ctx, dis, opts, &internalAPI{dockerCli: s.dockerCli}, filepath.Dir(s.configFile().Filename), w)
	errW := w.Wait()
	if err == nil {
		err = errW
	}
	if err != nil {
		return nil, WrapCategorisedComposeError(err, BuildFailure)
	}

	imagesBuilt := map[string]string{}
	for name, img := range response {
		if img == nil || len(img.ExporterResponse) == 0 {
			continue
		}
		digest, ok := img.ExporterResponse["containerimage.digest"]
		if !ok {
			continue
		}
		imagesBuilt[name] = digest
	}

	return imagesBuilt, err
}

func (s *composeService) getDrivers(ctx context.Context) ([]build.DriverInfo, error) { //nolint:gocyclo
	txn, release, err := storeutil.GetStore(s.dockerCli)
	if err != nil {
		return nil, err
	}
	defer release()

	ng, err := storeutil.GetCurrentInstance(txn, s.dockerCli)
	if err != nil {
		return nil, err
	}

	dis := make([]build.DriverInfo, len(ng.Nodes))
	var f driver.Factory
	if ng.Driver != "" {
		factories := driver.GetFactories(true)
		for _, fac := range factories {
			if fac.Name() == ng.Driver {
				f = fac
				continue
			}
		}
		if f == nil {
			if f, err = driver.GetFactory(ng.Driver, true); f == nil || err != nil {
				return nil, fmt.Errorf("failed to find buildx driver %q, error: %s", ng.Driver, err.Error())
			}
		}
	} else {
		ep := ng.Nodes[0].Endpoint
		dockerapi, err := clientForEndpoint(s.dockerCli, ep)
		if err != nil {
			return nil, err
		}
		f, err = driver.GetDefaultFactory(ctx, ep, dockerapi, false)
		if err != nil {
			return nil, err
		}
		ng.Driver = f.Name()
	}

	imageopt, err := storeutil.GetImageConfig(s.dockerCli, ng)
	if err != nil {
		return nil, err
	}

	eg, _ := errgroup.WithContext(ctx)
	for i, n := range ng.Nodes {
		func(i int, n store.Node) {
			eg.Go(func() error {
				di := build.DriverInfo{
					Name:        n.Name,
					Platform:    n.Platforms,
					ProxyConfig: storeutil.GetProxyConfig(s.dockerCli),
				}
				defer func() {
					dis[i] = di
				}()

				dockerapi, err := clientForEndpoint(s.dockerCli, n.Endpoint)
				if err != nil {
					di.Err = err
					return nil
				}
				// TODO: replace the following line with dockerclient.WithAPIVersionNegotiation option in clientForEndpoint
				dockerapi.NegotiateAPIVersion(ctx)

				contextStore := s.dockerCli.ContextStore()

				var kcc driver.KubeClientConfig
				kcc, err = configFromContext(n.Endpoint, contextStore)
				if err != nil {
					// err is returned if n.Endpoint is non-context name like "unix:///var/run/docker.sock".
					// try again with name="default".
					// FIXME: n should retain real context name.
					kcc, err = configFromContext("default", contextStore)
					if err != nil {
						logrus.Error(err)
					}
				}

				tryToUseKubeConfigInCluster := false
				if kcc == nil {
					tryToUseKubeConfigInCluster = true
				} else {
					if _, err := kcc.ClientConfig(); err != nil {
						tryToUseKubeConfigInCluster = true
					}
				}
				if tryToUseKubeConfigInCluster {
					kccInCluster := driver.KubeClientConfigInCluster{}
					if _, err := kccInCluster.ClientConfig(); err == nil {
						logrus.Debug("using kube config in cluster")
						kcc = kccInCluster
					}
				}

				d, err := driver.GetDriver(ctx, "buildx_buildkit_"+n.Name, f, n.Endpoint, dockerapi, imageopt.Auth, kcc, n.Flags, n.Files, n.DriverOpts, n.Platforms, "")
				if err != nil {
					di.Err = err
					return nil
				}
				di.Driver = d
				di.ImageOpt = imageopt
				return nil
			})
		}(i, n)
	}

	if err := eg.Wait(); err != nil {
		return nil, err
	}

	return dis, nil
}

func clientForEndpoint(dockerCli command.Cli, name string) (dockerclient.APIClient, error) {
	list, err := dockerCli.ContextStore().List()
	if err != nil {
		return nil, err
	}
	for _, l := range list {
		if l.Name != name {
			continue
		}
		dep, ok := l.Endpoints["docker"]
		if !ok {
			return nil, fmt.Errorf("context %q does not have a Docker endpoint", name)
		}
		epm, ok := dep.(docker.EndpointMeta)
		if !ok {
			return nil, fmt.Errorf("endpoint %q is not of type EndpointMeta, %T", dep, dep)
		}
		ep, err := docker.WithTLSData(dockerCli.ContextStore(), name, epm)
		if err != nil {
			return nil, err
		}
		clientOpts, err := ep.ClientOpts()
		if err != nil {
			return nil, err
		}
		return dockerclient.NewClientWithOpts(clientOpts...)
	}

	ep := docker.Endpoint{
		EndpointMeta: docker.EndpointMeta{
			Host: name,
		},
	}

	clientOpts, err := ep.ClientOpts()
	if err != nil {
		return nil, err
	}

	return dockerclient.NewClientWithOpts(clientOpts...)
}

func configFromContext(endpointName string, s ctxstore.Reader) (clientcmd.ClientConfig, error) {
	if strings.HasPrefix(endpointName, "kubernetes://") {
		u, _ := url.Parse(endpointName)
		if kubeconfig := u.Query().Get("kubeconfig"); kubeconfig != "" {
			_ = os.Setenv(clientcmd.RecommendedConfigPathEnvVar, kubeconfig)
		}
		rules := clientcmd.NewDefaultClientConfigLoadingRules()
		apiConfig, err := rules.Load()
		if err != nil {
			return nil, err
		}
		return clientcmd.NewDefaultClientConfig(*apiConfig, &clientcmd.ConfigOverrides{}), nil
	}
	return ctxkube.ConfigFromContext(endpointName, s)
}

type internalAPI struct {
	dockerCli command.Cli
}

func (a *internalAPI) DockerAPI(name string) (dockerclient.APIClient, error) {
	if name == "" {
		name = a.dockerCli.CurrentContext()
	}
	return clientForEndpoint(a.dockerCli, name)
}
<commit_msg>log the error object instead of the string message only<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

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

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

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

package compose

import (
	"context"
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	ctxkube "github.com/docker/buildx/driver/kubernetes/context"
	"github.com/docker/buildx/store"
	"github.com/docker/buildx/store/storeutil"
	"github.com/docker/cli/cli/command"
	"github.com/docker/cli/cli/context/docker"
	ctxstore "github.com/docker/cli/cli/context/store"
	dockerclient "github.com/docker/docker/client"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/errgroup"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/docker/buildx/build"
	"github.com/docker/buildx/driver"
	_ "github.com/docker/buildx/driver/docker"           //nolint:blank-imports
	_ "github.com/docker/buildx/driver/docker-container" //nolint:blank-imports
	_ "github.com/docker/buildx/driver/kubernetes"       //nolint:blank-imports
	xprogress "github.com/docker/buildx/util/progress"
)

func (s *composeService) doBuildBuildkit(ctx context.Context, opts map[string]build.Options, mode string) (map[string]string, error) {
	dis, err := s.getDrivers(ctx)
	if err != nil {
		return nil, err
	}

	// Progress needs its own context that lives longer than the
	// build one otherwise it won't read all the messages from
	// build and will lock
	progressCtx, cancel := context.WithCancel(context.Background())
	defer cancel()
	w := xprogress.NewPrinter(progressCtx, s.stdout(), os.Stdout, mode)

	response, err := build.Build(ctx, dis, opts, &internalAPI{dockerCli: s.dockerCli}, filepath.Dir(s.configFile().Filename), w)
	errW := w.Wait()
	if err == nil {
		err = errW
	}
	if err != nil {
		return nil, WrapCategorisedComposeError(err, BuildFailure)
	}

	imagesBuilt := map[string]string{}
	for name, img := range response {
		if img == nil || len(img.ExporterResponse) == 0 {
			continue
		}
		digest, ok := img.ExporterResponse["containerimage.digest"]
		if !ok {
			continue
		}
		imagesBuilt[name] = digest
	}

	return imagesBuilt, err
}

func (s *composeService) getDrivers(ctx context.Context) ([]build.DriverInfo, error) { //nolint:gocyclo
	txn, release, err := storeutil.GetStore(s.dockerCli)
	if err != nil {
		return nil, err
	}
	defer release()

	ng, err := storeutil.GetCurrentInstance(txn, s.dockerCli)
	if err != nil {
		return nil, err
	}

	dis := make([]build.DriverInfo, len(ng.Nodes))
	var f driver.Factory
	if ng.Driver != "" {
		factories := driver.GetFactories(true)
		for _, fac := range factories {
			if fac.Name() == ng.Driver {
				f = fac
				continue
			}
		}
		if f == nil {
			if f, err = driver.GetFactory(ng.Driver, true); f == nil || err != nil {
				return nil, fmt.Errorf("failed to find buildx driver %q, error: %w", ng.Driver, err)
			}
		}
	} else {
		ep := ng.Nodes[0].Endpoint
		dockerapi, err := clientForEndpoint(s.dockerCli, ep)
		if err != nil {
			return nil, err
		}
		f, err = driver.GetDefaultFactory(ctx, ep, dockerapi, false)
		if err != nil {
			return nil, err
		}
		ng.Driver = f.Name()
	}

	imageopt, err := storeutil.GetImageConfig(s.dockerCli, ng)
	if err != nil {
		return nil, err
	}

	eg, _ := errgroup.WithContext(ctx)
	for i, n := range ng.Nodes {
		func(i int, n store.Node) {
			eg.Go(func() error {
				di := build.DriverInfo{
					Name:        n.Name,
					Platform:    n.Platforms,
					ProxyConfig: storeutil.GetProxyConfig(s.dockerCli),
				}
				defer func() {
					dis[i] = di
				}()

				dockerapi, err := clientForEndpoint(s.dockerCli, n.Endpoint)
				if err != nil {
					di.Err = err
					return nil
				}
				// TODO: replace the following line with dockerclient.WithAPIVersionNegotiation option in clientForEndpoint
				dockerapi.NegotiateAPIVersion(ctx)

				contextStore := s.dockerCli.ContextStore()

				var kcc driver.KubeClientConfig
				kcc, err = configFromContext(n.Endpoint, contextStore)
				if err != nil {
					// err is returned if n.Endpoint is non-context name like "unix:///var/run/docker.sock".
					// try again with name="default".
					// FIXME: n should retain real context name.
					kcc, err = configFromContext("default", contextStore)
					if err != nil {
						logrus.Error(err)
					}
				}

				tryToUseKubeConfigInCluster := false
				if kcc == nil {
					tryToUseKubeConfigInCluster = true
				} else {
					if _, err := kcc.ClientConfig(); err != nil {
						tryToUseKubeConfigInCluster = true
					}
				}
				if tryToUseKubeConfigInCluster {
					kccInCluster := driver.KubeClientConfigInCluster{}
					if _, err := kccInCluster.ClientConfig(); err == nil {
						logrus.Debug("using kube config in cluster")
						kcc = kccInCluster
					}
				}

				d, err := driver.GetDriver(ctx, "buildx_buildkit_"+n.Name, f, n.Endpoint, dockerapi, imageopt.Auth, kcc, n.Flags, n.Files, n.DriverOpts, n.Platforms, "")
				if err != nil {
					di.Err = err
					return nil
				}
				di.Driver = d
				di.ImageOpt = imageopt
				return nil
			})
		}(i, n)
	}

	if err := eg.Wait(); err != nil {
		return nil, err
	}

	return dis, nil
}

func clientForEndpoint(dockerCli command.Cli, name string) (dockerclient.APIClient, error) {
	list, err := dockerCli.ContextStore().List()
	if err != nil {
		return nil, err
	}
	for _, l := range list {
		if l.Name != name {
			continue
		}
		dep, ok := l.Endpoints["docker"]
		if !ok {
			return nil, fmt.Errorf("context %q does not have a Docker endpoint", name)
		}
		epm, ok := dep.(docker.EndpointMeta)
		if !ok {
			return nil, fmt.Errorf("endpoint %q is not of type EndpointMeta, %T", dep, dep)
		}
		ep, err := docker.WithTLSData(dockerCli.ContextStore(), name, epm)
		if err != nil {
			return nil, err
		}
		clientOpts, err := ep.ClientOpts()
		if err != nil {
			return nil, err
		}
		return dockerclient.NewClientWithOpts(clientOpts...)
	}

	ep := docker.Endpoint{
		EndpointMeta: docker.EndpointMeta{
			Host: name,
		},
	}

	clientOpts, err := ep.ClientOpts()
	if err != nil {
		return nil, err
	}

	return dockerclient.NewClientWithOpts(clientOpts...)
}

func configFromContext(endpointName string, s ctxstore.Reader) (clientcmd.ClientConfig, error) {
	if strings.HasPrefix(endpointName, "kubernetes://") {
		u, _ := url.Parse(endpointName)
		if kubeconfig := u.Query().Get("kubeconfig"); kubeconfig != "" {
			_ = os.Setenv(clientcmd.RecommendedConfigPathEnvVar, kubeconfig)
		}
		rules := clientcmd.NewDefaultClientConfigLoadingRules()
		apiConfig, err := rules.Load()
		if err != nil {
			return nil, err
		}
		return clientcmd.NewDefaultClientConfig(*apiConfig, &clientcmd.ConfigOverrides{}), nil
	}
	return ctxkube.ConfigFromContext(endpointName, s)
}

type internalAPI struct {
	dockerCli command.Cli
}

func (a *internalAPI) DockerAPI(name string) (dockerclient.APIClient, error) {
	if name == "" {
		name = a.dockerCli.CurrentContext()
	}
	return clientForEndpoint(a.dockerCli, name)
}
<|endoftext|>


<commit_before>package vm

import (
	"fmt"

	"github.com/ethereum/go-ethereum/common"
)

type Memory struct {
	store []byte
}

func NewMemory() *Memory {
	return &Memory{nil}
}

func (m *Memory) Set(offset, size uint64, value []byte) {
	// length of store may never be less than offset + size.
	// The store should be resized PRIOR to setting the memory
	if size > uint64(len(m.store)) {
		panic("INVALID memory: store empty")
	}

	// It's possible the offset is greater than 0 and size equals 0. This is because
	// the calcMemSize (common.go) could potentially return 0 when size is zero (NO-OP)
	if size > 0 {
		copy(m.store[offset:offset+size], common.RightPadBytes(value, int(size)))
	}
}

func (m *Memory) Resize(size uint64) {
	if uint64(m.Len()) < size {
		m.store = append(m.store, make([]byte, size-uint64(m.Len()))...)
	}
}

func (self *Memory) Get(offset, size int64) (cpy []byte) {
	if size == 0 {
		return nil
	}

	if len(self.store) > int(offset) {
		cpy = make([]byte, size)
		copy(cpy, self.store[offset:offset+size])

		return
	}

	return
}

func (self *Memory) GetPtr(offset, size int64) []byte {
	if size == 0 {
		return nil
	}

	if len(self.store) > int(offset) {
		return self.store[offset : offset+size]
	}

	return nil
}

func (m *Memory) Len() int {
	return len(m.store)
}

func (m *Memory) Data() []byte {
	return m.store
}

func (m *Memory) Print() {
	fmt.Printf("### mem %d bytes ###\n", len(m.store))
	if len(m.store) > 0 {
		addr := 0
		for i := 0; i+32 <= len(m.store); i += 32 {
			fmt.Printf("%03d: % x\n", addr, m.store[i:i+32])
			addr++
		}
	} else {
		fmt.Println("-- empty --")
	}
	fmt.Println("####################")
}
<commit_msg>core/vm: Fork fix. Removal of appending 0 bytes in memset<commit_after>package vm

import "fmt"

type Memory struct {
	store []byte
}

func NewMemory() *Memory {
	return &Memory{nil}
}

func (m *Memory) Set(offset, size uint64, value []byte) {
	// length of store may never be less than offset + size.
	// The store should be resized PRIOR to setting the memory
	if size > uint64(len(m.store)) {
		panic("INVALID memory: store empty")
	}

	// It's possible the offset is greater than 0 and size equals 0. This is because
	// the calcMemSize (common.go) could potentially return 0 when size is zero (NO-OP)
	if size > 0 {
		copy(m.store[offset:offset+size], value)
	}
}

func (m *Memory) Resize(size uint64) {
	if uint64(m.Len()) < size {
		m.store = append(m.store, make([]byte, size-uint64(m.Len()))...)
	}
}

func (self *Memory) Get(offset, size int64) (cpy []byte) {
	if size == 0 {
		return nil
	}

	if len(self.store) > int(offset) {
		cpy = make([]byte, size)
		copy(cpy, self.store[offset:offset+size])

		return
	}

	return
}

func (self *Memory) GetPtr(offset, size int64) []byte {
	if size == 0 {
		return nil
	}

	if len(self.store) > int(offset) {
		return self.store[offset : offset+size]
	}

	return nil
}

func (m *Memory) Len() int {
	return len(m.store)
}

func (m *Memory) Data() []byte {
	return m.store
}

func (m *Memory) Print() {
	fmt.Printf("### mem %d bytes ###\n", len(m.store))
	if len(m.store) > 0 {
		addr := 0
		for i := 0; i+32 <= len(m.store); i += 32 {
			fmt.Printf("%03d: % x\n", addr, m.store[i:i+32])
			addr++
		}
	} else {
		fmt.Println("-- empty --")
	}
	fmt.Println("####################")
}
<|endoftext|>


<commit_before>package channel

import (
	"bytes"
	"regexp"
	"time"

	"github.com/scrapli/scrapligo/logging"
)

type authPatterns struct {
	telnetLoginPattern *regexp.Regexp
	passwordPattern    *regexp.Regexp
	passphrasePattern  *regexp.Regexp
}

var authPatternsInstance *authPatterns //nolint:gochecknoglobals

func getAuthPatterns() *authPatterns {
	if authPatternsInstance == nil {
		authPatternsInstance = &authPatterns{
			telnetLoginPattern: regexp.MustCompile(`(?im)^(.*username:)|(.*login:)\s?$`),
			passwordPattern:    regexp.MustCompile(`(?im)^password:\s?$`),
			passphrasePattern:  regexp.MustCompile(`(?i)enter passphrase for key`),
		}
	}

	return authPatternsInstance
}

func (c *Channel) authenticateSSH(authPassword, authPassphrase []byte) *channelResult {
	logging.LogDebug(c.FormatLogMessage("debug", "attempting in channel ssh authentication"))

	patterns := getAuthPatterns()
	passwordPattern := patterns.passwordPattern
	passphrasePattern := patterns.passphrasePattern

	var passwordCount = 0

	var passphraseCount = 0

	var b []byte

	for {
		chunk, err := c.Read()

		if err != nil {
			return &channelResult{
				result: b,
				error:  err,
			}
		}

		processedChunk := bytes.ToLower(bytes.Trim(chunk, "\x00"))

		b = append(b, processedChunk...)

		passwordMatch := passwordPattern.Match(b)
		passphraseMatch := passphrasePattern.Match(b)

		if passwordMatch { //nolint:nestif
			b = []byte{}
			passwordCount++

			if passwordCount > passwordSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassword,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found password prompt, sending password"))

			err = c.WriteAndReturn(authPassword, true)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		} else if passphraseMatch {
			b = []byte{}
			passphraseCount++

			if passwordCount > passphraseSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassphrase,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found passphrase prompt, sending passphrase"))

			err = c.WriteAndReturn(authPassphrase, true)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		}

		promptMatch := c.CommsPromptPattern.Match(b)
		if promptMatch {
			logging.LogDebug(c.FormatLogMessage("debug", "ssh authentication complete"))

			return &channelResult{
				result: b,
				error:  nil,
			}
		}
	}
}

// AuthenticateSSH Handles "in channel" SSH authentication (for "system" transport).
func (c *Channel) AuthenticateSSH(authPassword, authPassphrase string) ([]byte, error) {
	var _c = make(chan *channelResult)

	go func() {
		r := c.authenticateSSH([]byte(authPassword), []byte(authPassphrase))
		_c <- r
		close(_c)
	}()

	timer := time.NewTimer(c.DetermineOperationTimeout(*c.TimeoutOps))

	select {
	case r := <-_c:
		return r.result, r.error
	case <-timer.C:
		logging.LogError(c.FormatLogMessage("error", "timed out during ssh authentication"))

		return []byte{}, ErrAuthTimeout
	}
}

func (c *Channel) authenticateTelnet(authUsername, authPassword []byte) *channelResult {
	logging.LogDebug(c.FormatLogMessage("debug", "attempting in channel telnet authentication"))

	patterns := getAuthPatterns()
	usernamePattern := patterns.telnetLoginPattern
	passwordPattern := patterns.passwordPattern

	var usernameCount = 0

	var passwordCount = 0

	var b []byte

	for {
		chunk, err := c.Read()

		if err != nil {
			return &channelResult{
				result: b,
				error:  err,
			}
		}

		processedChunk := bytes.ToLower(bytes.Trim(chunk, "\x00"))

		b = append(b, processedChunk...)

		usernameMatch := usernamePattern.Match(b)
		passwordMatch := passwordPattern.Match(b)

		if usernameMatch { //nolint:nestif
			b = []byte{}
			usernameCount++

			if usernameCount > loginSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassword,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found login prompt, sending username"))

			err = c.WriteAndReturn(authUsername, false)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		} else if passwordMatch {
			b = []byte{}
			passwordCount++

			if passwordCount > passwordSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassword,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found password prompt, sending password"))

			err = c.WriteAndReturn(authPassword, true)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		}

		promptMatch := c.CommsPromptPattern.Match(b)
		if promptMatch {
			logging.LogDebug(c.FormatLogMessage("debug", "telnet authentication complete"))

			return &channelResult{
				result: b,
				error:  nil,
			}
		}
	}
}

// AuthenticateTelnet Handles "in channel" Telnet authentication (for "telnet" transport).
func (c *Channel) AuthenticateTelnet(authUsername, authPassword string) ([]byte, error) {
	var _c = make(chan *channelResult)

	go func() {
		r := c.authenticateTelnet([]byte(authUsername), []byte(authPassword))
		_c <- r
		close(_c)
	}()

	timer := time.NewTimer(c.DetermineOperationTimeout(*c.TimeoutOps))

	select {
	case r := <-_c:
		return r.result, r.error
	case <-timer.C:
		logging.LogError(c.FormatLogMessage("error", "timed out during telnet authentication"))

		return []byte{}, ErrAuthTimeout
	}
}
<commit_msg>fix password pattern regression<commit_after>package channel

import (
	"bytes"
	"regexp"
	"time"

	"github.com/scrapli/scrapligo/logging"
)

type authPatterns struct {
	telnetLoginPattern *regexp.Regexp
	passwordPattern    *regexp.Regexp
	passphrasePattern  *regexp.Regexp
}

var authPatternsInstance *authPatterns //nolint:gochecknoglobals

func getAuthPatterns() *authPatterns {
	if authPatternsInstance == nil {
		authPatternsInstance = &authPatterns{
			telnetLoginPattern: regexp.MustCompile(`(?im)^(.*username:)|(.*login:)\s?$`),
			passwordPattern:    regexp.MustCompile(`(?im)^(.*@.*)?password:\s?$`),
			passphrasePattern:  regexp.MustCompile(`(?i)enter passphrase for key`),
		}
	}

	return authPatternsInstance
}

func (c *Channel) authenticateSSH(authPassword, authPassphrase []byte) *channelResult {
	logging.LogDebug(c.FormatLogMessage("debug", "attempting in channel ssh authentication"))

	patterns := getAuthPatterns()
	passwordPattern := patterns.passwordPattern
	passphrasePattern := patterns.passphrasePattern

	var passwordCount = 0

	var passphraseCount = 0

	var b []byte

	for {
		chunk, err := c.Read()

		if err != nil {
			return &channelResult{
				result: b,
				error:  err,
			}
		}

		processedChunk := bytes.ToLower(bytes.Trim(chunk, "\x00"))

		b = append(b, processedChunk...)

		passwordMatch := passwordPattern.Match(b)
		passphraseMatch := passphrasePattern.Match(b)

		if passwordMatch { //nolint:nestif
			b = []byte{}
			passwordCount++

			if passwordCount > passwordSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassword,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found password prompt, sending password"))

			err = c.WriteAndReturn(authPassword, true)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		} else if passphraseMatch {
			b = []byte{}
			passphraseCount++

			if passwordCount > passphraseSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassphrase,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found passphrase prompt, sending passphrase"))

			err = c.WriteAndReturn(authPassphrase, true)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		}

		promptMatch := c.CommsPromptPattern.Match(b)
		if promptMatch {
			logging.LogDebug(c.FormatLogMessage("debug", "ssh authentication complete"))

			return &channelResult{
				result: b,
				error:  nil,
			}
		}
	}
}

// AuthenticateSSH Handles "in channel" SSH authentication (for "system" transport).
func (c *Channel) AuthenticateSSH(authPassword, authPassphrase string) ([]byte, error) {
	var _c = make(chan *channelResult)

	go func() {
		r := c.authenticateSSH([]byte(authPassword), []byte(authPassphrase))
		_c <- r
		close(_c)
	}()

	timer := time.NewTimer(c.DetermineOperationTimeout(*c.TimeoutOps))

	select {
	case r := <-_c:
		return r.result, r.error
	case <-timer.C:
		logging.LogError(c.FormatLogMessage("error", "timed out during ssh authentication"))

		return []byte{}, ErrAuthTimeout
	}
}

func (c *Channel) authenticateTelnet(authUsername, authPassword []byte) *channelResult {
	logging.LogDebug(c.FormatLogMessage("debug", "attempting in channel telnet authentication"))

	patterns := getAuthPatterns()
	usernamePattern := patterns.telnetLoginPattern
	passwordPattern := patterns.passwordPattern

	var usernameCount = 0

	var passwordCount = 0

	var b []byte

	for {
		chunk, err := c.Read()

		if err != nil {
			return &channelResult{
				result: b,
				error:  err,
			}
		}

		processedChunk := bytes.ToLower(bytes.Trim(chunk, "\x00"))

		b = append(b, processedChunk...)

		usernameMatch := usernamePattern.Match(b)
		passwordMatch := passwordPattern.Match(b)

		if usernameMatch { //nolint:nestif
			b = []byte{}
			usernameCount++

			if usernameCount > loginSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassword,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found login prompt, sending username"))

			err = c.WriteAndReturn(authUsername, false)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		} else if passwordMatch {
			b = []byte{}
			passwordCount++

			if passwordCount > passwordSeenMax {
				return &channelResult{
					result: []byte{},
					error:  ErrAuthFailedPassword,
				}
			}

			logging.LogDebug(c.FormatLogMessage("debug", "found password prompt, sending password"))

			err = c.WriteAndReturn(authPassword, true)
			if err != nil {
				return &channelResult{
					result: []byte{},
					error:  err,
				}
			}
		}

		promptMatch := c.CommsPromptPattern.Match(b)
		if promptMatch {
			logging.LogDebug(c.FormatLogMessage("debug", "telnet authentication complete"))

			return &channelResult{
				result: b,
				error:  nil,
			}
		}
	}
}

// AuthenticateTelnet Handles "in channel" Telnet authentication (for "telnet" transport).
func (c *Channel) AuthenticateTelnet(authUsername, authPassword string) ([]byte, error) {
	var _c = make(chan *channelResult)

	go func() {
		r := c.authenticateTelnet([]byte(authUsername), []byte(authPassword))
		_c <- r
		close(_c)
	}()

	timer := time.NewTimer(c.DetermineOperationTimeout(*c.TimeoutOps))

	select {
	case r := <-_c:
		return r.result, r.error
	case <-timer.C:
		logging.LogError(c.FormatLogMessage("error", "timed out during telnet authentication"))

		return []byte{}, ErrAuthTimeout
	}
}
<|endoftext|>


<commit_before>package charmap

import "testing"

func TestLanguageOf(t *testing.T) {

	// inArray := []rune{'ಅ', `ḥ`, `uː`, 'ഊ'}
	inArray := []rune{'ಅ', 'ഊ'}
	// outArray := []string{"kn_IN", "ISO15919", "IPA", "ml_IN"}
	outArray := []string{"kn_IN", "ml_IN"}

	for index, value := range inArray {
		if x, output := LanguageOf(value), outArray[index]; x != output {
			t.Errorf("LanguageOf(%v) = %s we need %s", value, x, output)
		}
	}

}

func TestCharCompare(t *testing.T) {
	// inArray := []interface{}{'ँ', 'అ', `aː`, 'ಆ'}
	inArray := []rune{'ँ', 'అ', 'ಆ'}
	// outArray := []interface{}{'ಁ', 'അ', 'ആ', `ā`}
	outArray := []rune{'ಁ', 'അ', 'ആ'}

	for index, value := range inArray {
		if x := CharCompare(value, outArray[index]); !x {
			t.Errorf("CharCompare(%v, %v) = %v we need %v", value, outArray[index], x, true)
		}
	}

}

func TestSoundexCode(t *testing.T) {
	inArray := []interface{}{'b', 't', 'ಇ', 'B', 'A'}
	outArray := []interface{}{'1', '3', 'B', '1', '0'}

	for index, value := range inArray {
		if x, err := SoundexCode(value); err != nil && x != outArray[index] {
			t.Errorf("SoundexCode(%v) = %v was expecting %v", value, x, outArray[index])
		}
	}
}
<commit_msg>Reorganized tests added IPA ISO15919 data for testing.<commit_after>package charmap

import "testing"

func testRuneInput(input runeSequence, output []string, t *testing.T) {

	for index, value := range input {
		if x, op := LanguageOf(value), output[index]; x != op {
			t.Errorf("LanguageOf(%v) = %s we need %s", value, x, op)
		}
	}
}

func testStringInput(input unicodeSequence, output []string, t *testing.T) {
	for index, value := range input {
		if x, op := LanguageOf(value), output[index]; x != op {
			t.Errorf("LanguageOf(%v) = %s we need %s", value, x, op)
		}
	}
}

func TestLanguageOf(t *testing.T) {

	inArray := []rune{'ಅ', 'ഊ'}
	outArray := []string{"kn_IN", "ml_IN"}

	inSplArray := []string{`ḥ`, `uː`}
	outSplArray := []string{"ISO15919", "IPA"}

	testRuneInput(inArray, outArray, t)
	testStringInput(inSplArray, outSplArray, t)

}

func TestCharCompare(t *testing.T) {
	inArray := []rune{'ँ', 'అ', 'ಆ'}
	outArray := []rune{'ಁ', 'അ', 'ആ'}

	for index, value := range inArray {
		if x := CharCompare(value, outArray[index]); !x {
			t.Errorf("CharCompare(%v, %v) = %v we need %v",
				value, outArray[index], x, true)
		}
	}

}

func TestSoundexCode(t *testing.T) {
	inArray := []interface{}{'b', 't', 'ಇ', 'B', 'A'}
	outArray := []interface{}{'1', '3', 'B', '1', '0'}

	for index, value := range inArray {
		if x, err := SoundexCode(value); err != nil && x !=
			outArray[index] {
			t.Errorf("SoundexCode(%v) = %v was expecting %v",
				value, x, outArray[index])
		}
	}
}
<|endoftext|>


<commit_before>package database

import (
	"fmt"

	"database/sql"

	"github.com/pivotal-golang/lager"
)

// Sprintf thinks the % sign is a formatting directive, so we escape it with a second %.
const revokeQuery = `REVOKE INSERT, UPDATE, CREATE ON %s.* FROM '%s'@'%%'`

const grantQuery = `GRANT INSERT, UPDATE, CREATE ON %s.* TO '%s'@'%%'`

type Database interface {
	Name() string
	GrantPrivileges() error
	RevokePrivileges() error
	KillActiveConnections() error
}

type database struct {
	name   string
	user   string
	db     *sql.DB
	logger lager.Logger
}

func New(name, user string, db *sql.DB, logger lager.Logger) Database {
	return &database{
		name:   name,
		user:   user,
		db:     db,
		logger: logger,
	}
}

func (d database) Name() string {
	return d.name
}

func (d database) RevokePrivileges() error {
	d.logger.Info(fmt.Sprintf("Revoking privileges to db '%s', user '%s'", d.name, d.user))
	d.logger.Info(fmt.Sprintf(revokeQuery, d.name, d.user))
	result, err := d.db.Exec(fmt.Sprintf(revokeQuery, d.name, d.user))
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to revoke privileges: %s", d.name, d.user, err.Error())
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to revoke privileges: Getting rows affected: %s", d.name, d.user, err.Error())
	}

	d.logger.Info(fmt.Sprintf("Updating db '%s', user '%s' to revoke privileges: Rows affected: %d", d.name, d.user, rowsAffected))

	_, err = d.db.Exec("FLUSH PRIVILEGES")
	if err != nil {
		return fmt.Errorf("Flushing privileges: %s", err.Error())
	}

	return nil
}

func (d database) GrantPrivileges() error {
	d.logger.Info(fmt.Sprintf("Granting privileges to db '%s', user '%s'", d.name, d.user))
	d.logger.Info(fmt.Sprintf(grantQuery, d.name, d.user))
	result, err := d.db.Exec(fmt.Sprintf(grantQuery, d.name, d.user))
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to grant privileges: %s", d.name, d.user, err.Error())
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to grant privileges: Getting rows affected: %s", d.name, d.user, err.Error())
	}

	d.logger.Info(fmt.Sprintf("Updating db '%s', user '%s' to grant privileges: Rows affected: %d", d.name, d.user, rowsAffected))

	_, err = d.db.Exec("FLUSH PRIVILEGES")
	if err != nil {
		return fmt.Errorf("Flushing privileges: %s", err.Error())
	}

	return nil
}

// ResetActivePrivileges flushes the privileges and kills all active connections to this database.
// New connections will get the new privileges.
func (d database) KillActiveConnections() error {
	d.logger.Info(fmt.Sprintf("Killing active connections to database '%s'", d.name))

	rows, err := d.db.Query("SELECT ID FROM INFORMATION_SCHEMA.PROCESSLIST WHERE DB = ? AND USER <> 'root'", d.name)
	if err != nil {
		return fmt.Errorf("Getting list of open connections to database '%s': %s", d.name, err.Error())
	}
	//TODO: untested Close, due to limitation of sqlmock: https://github.com/DATA-DOG/go-sqlmock/issues/15
	defer rows.Close()
	for rows.Next() {
		var connectionID int64
		if err := rows.Scan(&connectionID); err != nil {
			//TODO: untested error case, due to limitation of sqlmock: https://github.com/DATA-DOG/go-sqlmock/issues/13
			return fmt.Errorf("Scanning open connections to database '%s': %s", d.name, err.Error())
		}

		d.logger.Debug(fmt.Sprintf("Killing active connection %d to database '%s'", connectionID, d.name))
		_, err := d.db.Exec("KILL CONNECTION ?", connectionID)
		if err != nil {
			d.logger.Error(fmt.Sprintf("Failed to kill active connection %d to database '%s'", connectionID, d.name), err)
		}
	}
	//TODO: untested error case, due to limitation of sqlmock: https://github.com/DATA-DOG/go-sqlmock/issues/13
	if err := rows.Err(); err != nil {
		return fmt.Errorf("Reading open connections to database '%s': %s", d.name, err.Error())
	}

	return nil
}
<commit_msg>Remove debug log lines<commit_after>package database

import (
	"fmt"

	"database/sql"

	"github.com/pivotal-golang/lager"
)

// Sprintf thinks the % sign is a formatting directive, so we escape it with a second %.
const revokeQuery = `REVOKE INSERT, UPDATE, CREATE ON %s.* FROM '%s'@'%%'`

const grantQuery = `GRANT INSERT, UPDATE, CREATE ON %s.* TO '%s'@'%%'`

type Database interface {
	Name() string
	GrantPrivileges() error
	RevokePrivileges() error
	KillActiveConnections() error
}

type database struct {
	name   string
	user   string
	db     *sql.DB
	logger lager.Logger
}

func New(name, user string, db *sql.DB, logger lager.Logger) Database {
	return &database{
		name:   name,
		user:   user,
		db:     db,
		logger: logger,
	}
}

func (d database) Name() string {
	return d.name
}

func (d database) RevokePrivileges() error {
	d.logger.Info(fmt.Sprintf("Revoking privileges to db '%s', user '%s'", d.name, d.user))
	result, err := d.db.Exec(fmt.Sprintf(revokeQuery, d.name, d.user))
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to revoke privileges: %s", d.name, d.user, err.Error())
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to revoke privileges: Getting rows affected: %s", d.name, d.user, err.Error())
	}

	d.logger.Info(fmt.Sprintf("Updating db '%s', user '%s' to revoke privileges: Rows affected: %d", d.name, d.user, rowsAffected))

	_, err = d.db.Exec("FLUSH PRIVILEGES")
	if err != nil {
		return fmt.Errorf("Flushing privileges: %s", err.Error())
	}

	return nil
}

func (d database) GrantPrivileges() error {
	d.logger.Info(fmt.Sprintf("Granting privileges to db '%s', user '%s'", d.name, d.user))
	result, err := d.db.Exec(fmt.Sprintf(grantQuery, d.name, d.user))
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to grant privileges: %s", d.name, d.user, err.Error())
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("Updating db '%s', user '%s' to grant privileges: Getting rows affected: %s", d.name, d.user, err.Error())
	}

	d.logger.Info(fmt.Sprintf("Updating db '%s', user '%s' to grant privileges: Rows affected: %d", d.name, d.user, rowsAffected))

	_, err = d.db.Exec("FLUSH PRIVILEGES")
	if err != nil {
		return fmt.Errorf("Flushing privileges: %s", err.Error())
	}

	return nil
}

// ResetActivePrivileges flushes the privileges and kills all active connections to this database.
// New connections will get the new privileges.
func (d database) KillActiveConnections() error {
	d.logger.Info(fmt.Sprintf("Killing active connections to database '%s'", d.name))

	rows, err := d.db.Query("SELECT ID FROM INFORMATION_SCHEMA.PROCESSLIST WHERE DB = ? AND USER <> 'root'", d.name)
	if err != nil {
		return fmt.Errorf("Getting list of open connections to database '%s': %s", d.name, err.Error())
	}
	//TODO: untested Close, due to limitation of sqlmock: https://github.com/DATA-DOG/go-sqlmock/issues/15
	defer rows.Close()
	for rows.Next() {
		var connectionID int64
		if err := rows.Scan(&connectionID); err != nil {
			//TODO: untested error case, due to limitation of sqlmock: https://github.com/DATA-DOG/go-sqlmock/issues/13
			return fmt.Errorf("Scanning open connections to database '%s': %s", d.name, err.Error())
		}

		d.logger.Debug(fmt.Sprintf("Killing active connection %d to database '%s'", connectionID, d.name))
		_, err := d.db.Exec("KILL CONNECTION ?", connectionID)
		if err != nil {
			d.logger.Error(fmt.Sprintf("Failed to kill active connection %d to database '%s'", connectionID, d.name), err)
		}
	}
	//TODO: untested error case, due to limitation of sqlmock: https://github.com/DATA-DOG/go-sqlmock/issues/13
	if err := rows.Err(); err != nil {
		return fmt.Errorf("Reading open connections to database '%s': %s", d.name, err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package transports_test

import(
	"github.com/matiasinsaurralde/transports/marshalers"
	// "github.com/matiasinsaurralde/transports/marshalers/protos"

	"strings"
	"net/http"
	"net/url"
	"testing"
	"log"
)

type UnknownType struct {
	Field string
}

const TestRequestUrl string = "http://whatismyip.akamai.com/"

var request http.Request

func init() {

	log.Println("init")

	url, _ := url.Parse( TestRequestUrl)

	request = http.Request{
		Method: "GET",
		URL: url,
		Proto: "HTTP/1.0",
	}
}

func TestHttpRequestMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpRequest")
	}
}

func TestHttpResponseMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpResponse")
	}
}

func TestUnsupportedType( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}

	var v UnknownType
	v = UnknownType{"Value"}

	var i interface{}
	i = v

	err, _ := marshaler.Marshal(&i)

	exists := strings.Index(err.Error(), transports.MarshalerTypeNotSupportedError)

	if exists < 0 {
		t.Fatal("Unsupported type doesn't break the Protobuf marshaler")
	}
}

func TestNilInput( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	err, _ := marshaler.Marshal(nil)

	if strings.Index( err.Error(), transports.MarshalerNilTypeError ) < 0 {
		t.Fatal("Nil type doesn't break the Protobuf marshaler")
	}
}
<commit_msg>Removing UnknownType structure from ProtobufMarshaler<commit_after>package transports_test

import(
	"github.com/matiasinsaurralde/transports/marshalers"
	// "github.com/matiasinsaurralde/transports/marshalers/protos"

	"strings"
	"net/http"
	"net/url"
	"testing"
	"log"
)

const TestRequestUrl string = "http://whatismyip.akamai.com/"

var request http.Request

func init() {

	log.Println("init")

	url, _ := url.Parse( TestRequestUrl)

	request = http.Request{
		Method: "GET",
		URL: url,
		Proto: "HTTP/1.0",
	}
}

func TestHttpRequestMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpRequest")
	}
}

func TestHttpResponseMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpResponse")
	}
}

func TestUnsupportedType( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}

	var v UnknownType
	v = UnknownType{"Value"}

	var i interface{}
	i = v

	err, _ := marshaler.Marshal(&i)

	exists := strings.Index(err.Error(), transports.MarshalerTypeNotSupportedError)

	if exists < 0 {
		t.Fatal("Unsupported type doesn't break the Protobuf marshaler")
	}
}

func TestNilInput( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	err, _ := marshaler.Marshal(nil)

	if strings.Index( err.Error(), transports.MarshalerNilTypeError ) < 0 {
		t.Fatal("Nil type doesn't break the Protobuf marshaler")
	}
}
<|endoftext|>