<commit_before>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"go/token"
	"sort"

	"golang.org/x/tools/go/packages"
)

type nodeLoader struct {
	wd   string
	fset *token.FileSet
}

func (l nodeLoader) typed(args []string, recurse bool) ([]*packages.Package, error) {
	cfg := &packages.Config{
		Mode:  packages.LoadSyntax,
		Dir:   l.wd,
		Fset:  l.fset,
		Tests: true,
	}
	if recurse {
		// we'll need the syntax trees for the dependencies too
		cfg.Mode = packages.LoadAllSyntax
	}
	pkgs, err := packages.Load(cfg, args...)
	if err != nil {
		return nil, err
	}

	// Make a sorted list of the packages, including transitive dependencies
	// if recurse is true.
	byPath := make(map[string]*packages.Package)
	var addDeps func(*packages.Package)
	addDeps = func(pkg *packages.Package) {
		for _, imp := range pkg.Imports {
			byPath[imp.PkgPath] = imp
			addDeps(imp)
		}
	}
	for _, pkg := range pkgs {
		byPath[pkg.PkgPath] = pkg
		if recurse {
			// add all dependencies once
			addDeps(pkg)
		}
	}
	pkgs = pkgs[:0]
	for _, pkg := range byPath {
		pkgs = append(pkgs, pkg)
	}
	sort.Slice(pkgs, func(i, j int) bool {
		return pkgs[i].PkgPath < pkgs[j].PkgPath
	})
	return pkgs, nil
}
<commit_msg>don't load recursive test dependencies with -r<commit_after>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"go/token"
	"sort"
	"strings"

	"golang.org/x/tools/go/packages"
)

type nodeLoader struct {
	wd   string
	fset *token.FileSet
}

func (l nodeLoader) typed(args []string, recurse bool) ([]*packages.Package, error) {
	cfg := &packages.Config{
		Mode:  packages.LoadSyntax,
		Dir:   l.wd,
		Fset:  l.fset,
		Tests: true,
	}
	if recurse {
		// we'll need the syntax trees for the dependencies too
		cfg.Mode = packages.LoadAllSyntax
	}
	pkgs, err := packages.Load(cfg, args...)
	if err != nil {
		return nil, err
	}

	// Make a sorted list of the packages, including transitive dependencies
	// if recurse is true.
	byPath := make(map[string]*packages.Package)
	var addDeps func(*packages.Package)
	addDeps = func(pkg *packages.Package) {
		if strings.HasSuffix(pkg.PkgPath, ".test") {
			// don't add recursive test deps
			return
		}
		for _, imp := range pkg.Imports {
			if _, ok := byPath[imp.PkgPath]; ok {
				continue // seen; avoid recursive call
			}
			byPath[imp.PkgPath] = imp
			addDeps(imp)
		}
	}
	for _, pkg := range pkgs {
		byPath[pkg.PkgPath] = pkg
		if recurse {
			// add all dependencies once
			addDeps(pkg)
		}
	}
	pkgs = pkgs[:0]
	for _, pkg := range byPath {
		pkgs = append(pkgs, pkg)
	}
	sort.Slice(pkgs, func(i, j int) bool {
		return pkgs[i].PkgPath < pkgs[j].PkgPath
	})
	return pkgs, nil
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"errors"
	"fmt"
	"os/exec"
	"strings"
	"sync"
	"time"

	"../config"

	etcd "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

var (
	ctx     = context.Background()
	etcdAPI etcd.KeysAPI
	// Config is a pointer need to be set to the main configuration
	Config      *config.ConfigurationInfo
	writeToEtcd sync.Mutex
)

// Run starts watching for commands to execute
func Run() {
	setUpEtcd()
	go watchForNewCommands()
}

func watchForNewCommands() {
	w := etcdAPI.Watcher(fmt.Sprintf("/dispatch/commands/%s/", Config.Zone), &etcd.WatcherOptions{Recursive: true})
	for {
		r, err := w.Next(ctx)
		if err != nil {
			go watchForNewCommands()
			return
		}
		if pathParts := strings.Split(r.Node.Key, "/"); r.Action == "set" && pathParts[len(pathParts)-1] == "command" {
			go executeAndReturnResults(strings.Join(pathParts[:len(pathParts)-1], "/"), r.Node.Value)
		}
	}
}

func executeAndReturnResults(key, command string) {
	fmt.Println(key) // TO DO execute and push info to etcd
	fmt.Println(command)

	etcdAPI.Set(ctx, fmt.Sprintf("%s/machines/%s/output", key, Config.MachineName), "", &etcd.SetOptions{})

	commandParts := strings.Split(command, " ")
	commandoProcess := exec.Command(commandParts[0], commandParts[1:]...)
	stdoutPipe, _ := commandoProcess.StdoutPipe()
	stderrPipe, _ := commandoProcess.StderrPipe()
	stdoutReader := bufio.NewReader(stdoutPipe)
	stderrReader := bufio.NewReader(stderrPipe)

	go readSdtToEtcd(key, stdoutReader)
	go readSdtToEtcd(key, stderrReader)

	err := commandoProcess.Run() // blocking this thread
	if err == nil {
		err = errors.New("ok")
	}
	fmt.Println("done")
	etcdAPI.Set(ctx, fmt.Sprintf("%s/machines/%s/result", key, Config.MachineName), err.Error(), &etcd.SetOptions{})
}

func readSdtToEtcd(key string, std *bufio.Reader) {
	outputPath := fmt.Sprintf("%s/machines/%s/output", key, Config.MachineName)
	for {
		line, _, err := std.ReadLine()
		if err != nil {
			return // end of stream
		}
		writeToEtcd.Lock()
		response, _ := etcdAPI.Get(ctx, outputPath, &etcd.GetOptions{})
		output := response.Node.Value + string(line[:]) + "\n"
		etcdAPI.Set(ctx, outputPath, output, &etcd.SetOptions{})
		writeToEtcd.Unlock()
	}
}

func setUpEtcd() {
	c, err := etcd.New(etcd.Config{
		Endpoints:               []string{Config.EtcdAddress},
		Transport:               etcd.DefaultTransport,
		HeaderTimeoutPerRequest: 10 * time.Second,
	})
	if err != nil {
		panic(err)
	}
	etcdAPI = etcd.NewKeysAPI(c)
}
<commit_msg>rename writeToEtcd to etcdMutex<commit_after>package command

import (
	"bufio"
	"errors"
	"fmt"
	"os/exec"
	"strings"
	"sync"
	"time"

	"../config"

	etcd "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

var (
	ctx     = context.Background()
	etcdAPI etcd.KeysAPI
	// Config is a pointer need to be set to the main configuration
	Config    *config.ConfigurationInfo
	etcdMutex sync.Mutex
)

// Run starts watching for commands to execute
func Run() {
	setUpEtcd()
	go watchForNewCommands()
}

func watchForNewCommands() {
	w := etcdAPI.Watcher(fmt.Sprintf("/dispatch/commands/%s/", Config.Zone), &etcd.WatcherOptions{Recursive: true})
	for {
		r, err := w.Next(ctx)
		if err != nil {
			go watchForNewCommands()
			return
		}
		if pathParts := strings.Split(r.Node.Key, "/"); r.Action == "set" && pathParts[len(pathParts)-1] == "command" {
			go executeAndReturnResults(strings.Join(pathParts[:len(pathParts)-1], "/"), r.Node.Value)
		}
	}
}

func executeAndReturnResults(key, command string) {
	fmt.Println(key) // TO DO execute and push info to etcd
	fmt.Println(command)

	etcdAPI.Set(ctx, fmt.Sprintf("%s/machines/%s/output", key, Config.MachineName), "", &etcd.SetOptions{})

	commandParts := strings.Split(command, " ")
	commandoProcess := exec.Command(commandParts[0], commandParts[1:]...)
	stdoutPipe, _ := commandoProcess.StdoutPipe()
	stderrPipe, _ := commandoProcess.StderrPipe()
	stdoutReader := bufio.NewReader(stdoutPipe)
	stderrReader := bufio.NewReader(stderrPipe)

	go readSdtToEtcd(key, stdoutReader)
	go readSdtToEtcd(key, stderrReader)

	err := commandoProcess.Run() // blocking this thread
	if err == nil {
		err = errors.New("ok")
	}
	fmt.Println("done")
	etcdAPI.Set(ctx, fmt.Sprintf("%s/machines/%s/result", key, Config.MachineName), err.Error(), &etcd.SetOptions{})
}

func readSdtToEtcd(key string, std *bufio.Reader) {
	outputPath := fmt.Sprintf("%s/machines/%s/output", key, Config.MachineName)
	for {
		line, _, err := std.ReadLine()
		if err != nil {
			return // end of stream
		}
		etcdMutex.Lock()
		response, _ := etcdAPI.Get(ctx, outputPath, &etcd.GetOptions{})
		output := response.Node.Value + string(line[:]) + "\n"
		etcdAPI.Set(ctx, outputPath, output, &etcd.SetOptions{})
		etcdMutex.Unlock()
	}
}

func setUpEtcd() {
	c, err := etcd.New(etcd.Config{
		Endpoints:               []string{Config.EtcdAddress},
		Transport:               etcd.DefaultTransport,
		HeaderTimeoutPerRequest: 10 * time.Second,
	})
	if err != nil {
		panic(err)
	}
	etcdAPI = etcd.NewKeysAPI(c)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package gcp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"k8s.io/apimachinery/pkg/util/net"
	"k8s.io/apimachinery/pkg/util/yaml"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/util/jsonpath"
)

func init() {
	if err := restclient.RegisterAuthProviderPlugin("gcp", newGCPAuthProvider); err != nil {
		glog.Fatalf("Failed to register gcp auth plugin: %v", err)
	}
}

// Stubbable for testing
var execCommand = exec.Command

// gcpAuthProvider is an auth provider plugin that uses GCP credentials to provide
// tokens for kubectl to authenticate itself to the apiserver. A sample json config
// is provided below with all recognized options described.
//
// {
//   'auth-provider': {
//     # Required
//     "name": "gcp",
//
//     'config': {
//       # Caching options
//
//       # Raw string data representing cached access token.
//       "access-token": "ya29.CjWdA4GiBPTt",
//       # RFC3339Nano expiration timestamp for cached access token.
//       "expiry": "2016-10-31 22:31:9.123",
//
//       # Command execution options
//       # These options direct the plugin to execute a specified command and parse
//       # token and expiry time from the output of the command.
//
//       # Command to execute for access token. Command output will be parsed as JSON.
//       # If "cmd-args" is not present, this value will be split on whitespace, with
//       # the first element interpreted as the command, remaining elements as args.
//       "cmd-path": "/usr/bin/gcloud",
//
//       # Arguments to pass to command to execute for access token.
//       "cmd-args": "config config-helper --output=json"
//
//       # JSONPath to the string field that represents the access token in
//       # command output. If omitted, defaults to "{.access_token}".
//       "token-key": "{.credential.access_token}",
//
//       # JSONPath to the string field that represents expiration timestamp
//       # of the access token in the command output. If omitted, defaults to
//       # "{.token_expiry}"
//       "expiry-key": ""{.credential.token_expiry}",
//
//       # golang reference time in the format that the expiration timestamp uses.
//       # If omitted, defaults to time.RFC3339Nano
//       "time-fmt": "2006-01-02 15:04:05.999999999"
//     }
//   }
// }
//
type gcpAuthProvider struct {
	tokenSource oauth2.TokenSource
	persister   restclient.AuthProviderConfigPersister
}

func newGCPAuthProvider(_ string, gcpConfig map[string]string, persister restclient.AuthProviderConfigPersister) (restclient.AuthProvider, error) {
	var ts oauth2.TokenSource
	var err error
	if cmd, useCmd := gcpConfig["cmd-path"]; useCmd {
		if len(cmd) == 0 {
			return nil, fmt.Errorf("missing access token cmd")
		}
		var args []string
		if cmdArgs, ok := gcpConfig["cmd-args"]; ok {
			args = strings.Fields(cmdArgs)
		} else {
			fields := strings.Fields(cmd)
			cmd = fields[0]
			args = fields[1:]
		}
		ts = newCmdTokenSource(cmd, args, gcpConfig["token-key"], gcpConfig["expiry-key"], gcpConfig["time-fmt"])
	} else {
		ts, err = google.DefaultTokenSource(context.Background(), "https://www.googleapis.com/auth/cloud-platform")
	}
	if err != nil {
		return nil, err
	}
	cts, err := newCachedTokenSource(gcpConfig["access-token"], gcpConfig["expiry"], persister, ts, gcpConfig)
	if err != nil {
		return nil, err
	}
	return &gcpAuthProvider{cts, persister}, nil
}

func (g *gcpAuthProvider) WrapTransport(rt http.RoundTripper) http.RoundTripper {
	return &conditionalTransport{&oauth2.Transport{Source: g.tokenSource, Base: rt}, g.persister}
}

func (g *gcpAuthProvider) Login() error { return nil }

type cachedTokenSource struct {
	lk          sync.Mutex
	source      oauth2.TokenSource
	accessToken string
	expiry      time.Time
	persister   restclient.AuthProviderConfigPersister
	cache       map[string]string
}

func newCachedTokenSource(accessToken, expiry string, persister restclient.AuthProviderConfigPersister, ts oauth2.TokenSource, cache map[string]string) (*cachedTokenSource, error) {
	var expiryTime time.Time
	if parsedTime, err := time.Parse(time.RFC3339Nano, expiry); err == nil {
		expiryTime = parsedTime
	}
	if cache == nil {
		cache = make(map[string]string)
	}
	return &cachedTokenSource{
		source:      ts,
		accessToken: accessToken,
		expiry:      expiryTime,
		persister:   persister,
		cache:       cache,
	}, nil
}

func (t *cachedTokenSource) Token() (*oauth2.Token, error) {
	tok := t.cachedToken()
	if tok.Valid() && !tok.Expiry.IsZero() {
		return tok, nil
	}
	tok, err := t.source.Token()
	if err != nil {
		return nil, err
	}
	cache := t.update(tok)
	if t.persister != nil {
		if err := t.persister.Persist(cache); err != nil {
			glog.V(4).Infof("Failed to persist token: %v", err)
		}
	}
	return tok, nil
}

func (t *cachedTokenSource) cachedToken() *oauth2.Token {
	t.lk.Lock()
	defer t.lk.Unlock()
	return &oauth2.Token{
		AccessToken: t.accessToken,
		TokenType:   "Bearer",
		Expiry:      t.expiry,
	}
}

func (t *cachedTokenSource) update(tok *oauth2.Token) map[string]string {
	t.lk.Lock()
	defer t.lk.Unlock()
	t.accessToken = tok.AccessToken
	t.expiry = tok.Expiry
	ret := map[string]string{}
	for k, v := range t.cache {
		ret[k] = v
	}
	ret["access-token"] = t.accessToken
	ret["expiry"] = t.expiry.Format(time.RFC3339Nano)
	return ret
}

type commandTokenSource struct {
	cmd       string
	args      []string
	tokenKey  string
	expiryKey string
	timeFmt   string
}

func newCmdTokenSource(cmd string, args []string, tokenKey, expiryKey, timeFmt string) *commandTokenSource {
	if len(timeFmt) == 0 {
		timeFmt = time.RFC3339Nano
	}
	if len(tokenKey) == 0 {
		tokenKey = "{.access_token}"
	}
	if len(expiryKey) == 0 {
		expiryKey = "{.token_expiry}"
	}
	return &commandTokenSource{
		cmd:       cmd,
		args:      args,
		tokenKey:  tokenKey,
		expiryKey: expiryKey,
		timeFmt:   timeFmt,
	}
}

func (c *commandTokenSource) Token() (*oauth2.Token, error) {
	fullCmd := strings.Join(append([]string{c.cmd}, c.args...), " ")
	cmd := execCommand(c.cmd, c.args...)
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("error executing access token command %q: err=%v output=%s", fullCmd, err, output)
	}
	token, err := c.parseTokenCmdOutput(output)
	if err != nil {
		return nil, fmt.Errorf("error parsing output for access token command %q: %v", fullCmd, err)
	}
	return token, nil
}

func (c *commandTokenSource) parseTokenCmdOutput(output []byte) (*oauth2.Token, error) {
	output, err := yaml.ToJSON(output)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := json.Unmarshal(output, &data); err != nil {
		return nil, err
	}

	accessToken, err := parseJSONPath(data, "token-key", c.tokenKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing token-key %q from %q: %v", c.tokenKey, string(output), err)
	}
	expiryStr, err := parseJSONPath(data, "expiry-key", c.expiryKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing expiry-key %q from %q: %v", c.expiryKey, string(output), err)
	}
	var expiry time.Time
	if t, err := time.Parse(c.timeFmt, expiryStr); err != nil {
		glog.V(4).Infof("Failed to parse token expiry from %s (fmt=%s): %v", expiryStr, c.timeFmt, err)
	} else {
		expiry = t
	}

	return &oauth2.Token{
		AccessToken: accessToken,
		TokenType:   "Bearer",
		Expiry:      expiry,
	}, nil
}

func parseJSONPath(input interface{}, name, template string) (string, error) {
	j := jsonpath.New(name)
	buf := new(bytes.Buffer)
	if err := j.Parse(template); err != nil {
		return "", err
	}
	if err := j.Execute(buf, input); err != nil {
		return "", err
	}
	return buf.String(), nil
}

type conditionalTransport struct {
	oauthTransport *oauth2.Transport
	persister      restclient.AuthProviderConfigPersister
}

var _ net.RoundTripperWrapper = &conditionalTransport{}

func (t *conditionalTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if len(req.Header.Get("Authorization")) != 0 {
		return t.oauthTransport.Base.RoundTrip(req)
	}

	res, err := t.oauthTransport.RoundTrip(req)

	if err != nil {
		return nil, err
	}

	if res.StatusCode == 401 {
		glog.V(4).Infof("The credentials that were supplied are invalid for the target cluster")
		emptyCache := make(map[string]string)
		t.persister.Persist(emptyCache)
	}

	return res, nil
}

func (t *conditionalTransport) WrappedRoundTripper() http.RoundTripper { return t.oauthTransport.Base }
<commit_msg>auth/gcp: capture stderr from gcloud config-helper<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package gcp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"k8s.io/apimachinery/pkg/util/net"
	"k8s.io/apimachinery/pkg/util/yaml"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/util/jsonpath"
)

func init() {
	if err := restclient.RegisterAuthProviderPlugin("gcp", newGCPAuthProvider); err != nil {
		glog.Fatalf("Failed to register gcp auth plugin: %v", err)
	}
}

// Stubbable for testing
var execCommand = exec.Command

// gcpAuthProvider is an auth provider plugin that uses GCP credentials to provide
// tokens for kubectl to authenticate itself to the apiserver. A sample json config
// is provided below with all recognized options described.
//
// {
//   'auth-provider': {
//     # Required
//     "name": "gcp",
//
//     'config': {
//       # Caching options
//
//       # Raw string data representing cached access token.
//       "access-token": "ya29.CjWdA4GiBPTt",
//       # RFC3339Nano expiration timestamp for cached access token.
//       "expiry": "2016-10-31 22:31:9.123",
//
//       # Command execution options
//       # These options direct the plugin to execute a specified command and parse
//       # token and expiry time from the output of the command.
//
//       # Command to execute for access token. Command output will be parsed as JSON.
//       # If "cmd-args" is not present, this value will be split on whitespace, with
//       # the first element interpreted as the command, remaining elements as args.
//       "cmd-path": "/usr/bin/gcloud",
//
//       # Arguments to pass to command to execute for access token.
//       "cmd-args": "config config-helper --output=json"
//
//       # JSONPath to the string field that represents the access token in
//       # command output. If omitted, defaults to "{.access_token}".
//       "token-key": "{.credential.access_token}",
//
//       # JSONPath to the string field that represents expiration timestamp
//       # of the access token in the command output. If omitted, defaults to
//       # "{.token_expiry}"
//       "expiry-key": ""{.credential.token_expiry}",
//
//       # golang reference time in the format that the expiration timestamp uses.
//       # If omitted, defaults to time.RFC3339Nano
//       "time-fmt": "2006-01-02 15:04:05.999999999"
//     }
//   }
// }
//
type gcpAuthProvider struct {
	tokenSource oauth2.TokenSource
	persister   restclient.AuthProviderConfigPersister
}

func newGCPAuthProvider(_ string, gcpConfig map[string]string, persister restclient.AuthProviderConfigPersister) (restclient.AuthProvider, error) {
	var ts oauth2.TokenSource
	var err error
	if cmd, useCmd := gcpConfig["cmd-path"]; useCmd {
		if len(cmd) == 0 {
			return nil, fmt.Errorf("missing access token cmd")
		}
		var args []string
		if cmdArgs, ok := gcpConfig["cmd-args"]; ok {
			args = strings.Fields(cmdArgs)
		} else {
			fields := strings.Fields(cmd)
			cmd = fields[0]
			args = fields[1:]
		}
		ts = newCmdTokenSource(cmd, args, gcpConfig["token-key"], gcpConfig["expiry-key"], gcpConfig["time-fmt"])
	} else {
		ts, err = google.DefaultTokenSource(context.Background(), "https://www.googleapis.com/auth/cloud-platform")
	}
	if err != nil {
		return nil, err
	}
	cts, err := newCachedTokenSource(gcpConfig["access-token"], gcpConfig["expiry"], persister, ts, gcpConfig)
	if err != nil {
		return nil, err
	}
	return &gcpAuthProvider{cts, persister}, nil
}

func (g *gcpAuthProvider) WrapTransport(rt http.RoundTripper) http.RoundTripper {
	return &conditionalTransport{&oauth2.Transport{Source: g.tokenSource, Base: rt}, g.persister}
}

func (g *gcpAuthProvider) Login() error { return nil }

type cachedTokenSource struct {
	lk          sync.Mutex
	source      oauth2.TokenSource
	accessToken string
	expiry      time.Time
	persister   restclient.AuthProviderConfigPersister
	cache       map[string]string
}

func newCachedTokenSource(accessToken, expiry string, persister restclient.AuthProviderConfigPersister, ts oauth2.TokenSource, cache map[string]string) (*cachedTokenSource, error) {
	var expiryTime time.Time
	if parsedTime, err := time.Parse(time.RFC3339Nano, expiry); err == nil {
		expiryTime = parsedTime
	}
	if cache == nil {
		cache = make(map[string]string)
	}
	return &cachedTokenSource{
		source:      ts,
		accessToken: accessToken,
		expiry:      expiryTime,
		persister:   persister,
		cache:       cache,
	}, nil
}

func (t *cachedTokenSource) Token() (*oauth2.Token, error) {
	tok := t.cachedToken()
	if tok.Valid() && !tok.Expiry.IsZero() {
		return tok, nil
	}
	tok, err := t.source.Token()
	if err != nil {
		return nil, err
	}
	cache := t.update(tok)
	if t.persister != nil {
		if err := t.persister.Persist(cache); err != nil {
			glog.V(4).Infof("Failed to persist token: %v", err)
		}
	}
	return tok, nil
}

func (t *cachedTokenSource) cachedToken() *oauth2.Token {
	t.lk.Lock()
	defer t.lk.Unlock()
	return &oauth2.Token{
		AccessToken: t.accessToken,
		TokenType:   "Bearer",
		Expiry:      t.expiry,
	}
}

func (t *cachedTokenSource) update(tok *oauth2.Token) map[string]string {
	t.lk.Lock()
	defer t.lk.Unlock()
	t.accessToken = tok.AccessToken
	t.expiry = tok.Expiry
	ret := map[string]string{}
	for k, v := range t.cache {
		ret[k] = v
	}
	ret["access-token"] = t.accessToken
	ret["expiry"] = t.expiry.Format(time.RFC3339Nano)
	return ret
}

type commandTokenSource struct {
	cmd       string
	args      []string
	tokenKey  string
	expiryKey string
	timeFmt   string
}

func newCmdTokenSource(cmd string, args []string, tokenKey, expiryKey, timeFmt string) *commandTokenSource {
	if len(timeFmt) == 0 {
		timeFmt = time.RFC3339Nano
	}
	if len(tokenKey) == 0 {
		tokenKey = "{.access_token}"
	}
	if len(expiryKey) == 0 {
		expiryKey = "{.token_expiry}"
	}
	return &commandTokenSource{
		cmd:       cmd,
		args:      args,
		tokenKey:  tokenKey,
		expiryKey: expiryKey,
		timeFmt:   timeFmt,
	}
}

func (c *commandTokenSource) Token() (*oauth2.Token, error) {
	fullCmd := strings.Join(append([]string{c.cmd}, c.args...), " ")
	cmd := execCommand(c.cmd, c.args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("error executing access token command %q: err=%v output=%s stderr=%s", fullCmd, err, output, string(stderr.Bytes()))
	}
	token, err := c.parseTokenCmdOutput(output)
	if err != nil {
		return nil, fmt.Errorf("error parsing output for access token command %q: %v", fullCmd, err)
	}
	return token, nil
}

func (c *commandTokenSource) parseTokenCmdOutput(output []byte) (*oauth2.Token, error) {
	output, err := yaml.ToJSON(output)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := json.Unmarshal(output, &data); err != nil {
		return nil, err
	}

	accessToken, err := parseJSONPath(data, "token-key", c.tokenKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing token-key %q from %q: %v", c.tokenKey, string(output), err)
	}
	expiryStr, err := parseJSONPath(data, "expiry-key", c.expiryKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing expiry-key %q from %q: %v", c.expiryKey, string(output), err)
	}
	var expiry time.Time
	if t, err := time.Parse(c.timeFmt, expiryStr); err != nil {
		glog.V(4).Infof("Failed to parse token expiry from %s (fmt=%s): %v", expiryStr, c.timeFmt, err)
	} else {
		expiry = t
	}

	return &oauth2.Token{
		AccessToken: accessToken,
		TokenType:   "Bearer",
		Expiry:      expiry,
	}, nil
}

func parseJSONPath(input interface{}, name, template string) (string, error) {
	j := jsonpath.New(name)
	buf := new(bytes.Buffer)
	if err := j.Parse(template); err != nil {
		return "", err
	}
	if err := j.Execute(buf, input); err != nil {
		return "", err
	}
	return buf.String(), nil
}

type conditionalTransport struct {
	oauthTransport *oauth2.Transport
	persister      restclient.AuthProviderConfigPersister
}

var _ net.RoundTripperWrapper = &conditionalTransport{}

func (t *conditionalTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if len(req.Header.Get("Authorization")) != 0 {
		return t.oauthTransport.Base.RoundTrip(req)
	}

	res, err := t.oauthTransport.RoundTrip(req)

	if err != nil {
		return nil, err
	}

	if res.StatusCode == 401 {
		glog.V(4).Infof("The credentials that were supplied are invalid for the target cluster")
		emptyCache := make(map[string]string)
		t.persister.Persist(emptyCache)
	}

	return res, nil
}

func (t *conditionalTransport) WrappedRoundTripper() http.RoundTripper { return t.oauthTransport.Base }
<|endoftext|>


<commit_before>package player

import "math"

func computeScore(keywords1, keywords2 []string) float64 {
	var total float64 = 0

	for _, k1 := range keywords1 {
		for _, k2 := range keywords2 {
			d := inverseLevenshteinDistance(k1, k2)
			if math.IsInf(d, 0) {
				total = total + 10
			} else {
				total = total + d
			}
		}
	}

	return total
}

func inverseLevenshteinDistance(word1, word2 string) float64 {
	// Code from http://rosettacode.org/wiki/Levenshtein_distance#Go
	d := make([][]int, len(word1)+1)
	for i := range d {
		d[i] = make([]int, len(word2)+1)
	}
	for i := range d {
		d[i][0] = i
	}
	for j := range d[0] {
		d[0][j] = j
	}
	for j := 1; j <= len(word2); j++ {
		for i := 1; i <= len(word1); i++ {
			if word1[i-1] == word2[j-1] {
				d[i][j] = d[i-1][j-1]
			} else {
				min := d[i-1][j]
				if d[i][j-1] < min {
					min = d[i][j-1]
				}
				if d[i-1][j-1] < min {
					min = d[i-1][j-1]
				}
				d[i][j] = min + 1
			}
		}

	}
	return 1 / float64(d[len(word1)][len(word2)])
}
<commit_msg>fix(score): compute score once per word<commit_after>package player

import "math"

func computeScore(keywords1, keywords2 []string) float64 {
	var total float64 = 0
	var seen map[string]bool

	for _, k1 := range keywords1 {
		if seen[k1] {
			for _, k2 := range keywords2 {
				d := inverseLevenshteinDistance(k1, k2)
				if math.IsInf(d, 0) {
					total = total + 10
				} else {
					total = total + d
				}
			}
			seen[k1] = true
		}
	}

	return total
}

func inverseLevenshteinDistance(word1, word2 string) float64 {
	// Code from http://rosettacode.org/wiki/Levenshtein_distance#Go
	d := make([][]int, len(word1)+1)
	for i := range d {
		d[i] = make([]int, len(word2)+1)
	}
	for i := range d {
		d[i][0] = i
	}
	for j := range d[0] {
		d[0][j] = j
	}
	for j := 1; j <= len(word2); j++ {
		for i := 1; i <= len(word1); i++ {
			if word1[i-1] == word2[j-1] {
				d[i][j] = d[i-1][j-1]
			} else {
				min := d[i-1][j]
				if d[i][j-1] < min {
					min = d[i][j-1]
				}
				if d[i-1][j-1] < min {
					min = d[i-1][j-1]
				}
				d[i][j] = min + 1
			}
		}

	}
	return 1 / float64(d[len(word1)][len(word2)])
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/inteliquent/casatunes"
	"github.com/nlopes/slack"
	"log"
	"os"
	"strings"
)

var (
	message string
)

func casa_PlayerAction(slack_api *slack.Client, ev *slack.MessageEvent) {
	CASA_ENDPOINT := os.Getenv("CASA_ENDPOINT")
	casa_api := casatunes.New(CASA_ENDPOINT)

	message_parameters := slack.NewPostMessageParameters()
	message_parameters.AsUser = true

	channelID := ev.Channel

	playerAction := regexp_playeraction.FindStringSubmatch(ev.Text)[1]

	err := casa_api.SourcesPlayerAction("0", playerAction)

	if err != nil {
		log.Println(err)
		message = fmt.Sprintf("Oops! %s", err)
	} else {
		user, err := slack_api.GetUserInfo(ev.User)

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

		switch strings.ToLower(playerAction) {
		case "play":
			message = fmt.Sprintf("%s has resumed playback!", user.RealName)

		case "pause":
			message = fmt.Sprintf("%s has paused playback!", user.RealName)

		case "previous":
			message = "_*BACK!!*_"

		case "next":
			message = "_*NEXT!!*_"
		}
	}

	slack_api.PostMessage(
		channelID,
		message,
		message_parameters,
	)
}
<commit_msg>PlayerAction now accepts casabot_command<commit_after>package main

import (
	"fmt"
	"github.com/inteliquent/casatunes"
	"github.com/nlopes/slack"
	"log"
	"os"
	"strings"
)

var (
	message string
)

func casa_PlayerAction(slack_api *slack.Client, command *casabot_command) {
	CASA_ENDPOINT := os.Getenv("CASA_ENDPOINT")
	casa_api := casatunes.New(CASA_ENDPOINT)

	message_parameters := slack.NewPostMessageParameters()
	message_parameters.AsUser = true

	channelID := command.ChannelID

	playerAction := command.Command

	if playerAction == "resume" {
		playerAction = "play"
	}

	err := casa_api.SourcesPlayerAction("0", playerAction)

	if err != nil {
		log.Println(err)
		message = fmt.Sprintf("Oops! %s", err)
	} else {
		if err != nil {
			log.Fatal(err)
		}

		switch strings.ToLower(playerAction) {
		case "play":
			message = "resumed playback!"

		case "pause":
			message = "paused playback!"

		case "previous":
			message = "_*BACK!!*_"

		case "next":
			message = "_*NEXT!!*_"
		}
	}

	slack_api.PostMessage(
		channelID,
		message,
		message_parameters,
	)
}
<|endoftext|>


<commit_before>package stager

import (
	"bytes"
	"fmt"
	"net/http/httputil"
	"net/url"
	"sync"
	"text/template"
)

type Backend struct {
	Port    int
	Name    string
	proxy   *httputil.ReverseProxy
	running bool
}

// This will do the setup of the backend at some point.
func (b *Backend) initialize() {
	b.running = true
}

type backendManager struct {
	backends     map[string]*Backend
	suffixLength int
	currentPort  int
	availPorts   []int
	lockPorts    sync.Mutex
	proxyPrefix  *template.Template
}

func (m *backendManager) get(domain string) *Backend {
	name := domain[:len(domain)-m.suffixLength]
	if m.backends[name] == nil {
		port := m.allocatePort()
		b := &Backend{
			Port: port,
			Name: name,
		}
		m.backends[name] = b
		buf := &bytes.Buffer{}
		err := m.proxyPrefix.Execute(buf, b)
		if err != nil {
			panic(err)
		}
		rawurl := string(buf.Bytes())
		u, err := url.Parse(rawurl)
		if err != nil {
			panic(err)
		}
		fmt.Printf("making new instance %s on port %d with backend url %s\n", name, port, rawurl)

		b.proxy = httputil.NewSingleHostReverseProxy(u)

		go b.initialize()

	}
	return m.backends[name]
}

/* Allocate a port number to be used for another backend. */
func (m *backendManager) allocatePort() int {
	m.lockPorts.Lock()
	defer m.lockPorts.Unlock()
	l := len(m.availPorts)
	if l > 0 {
		port := m.availPorts[l-1]
		m.availPorts = m.availPorts[:l-1]
		return port
	} else {
		port := m.currentPort
		m.currentPort += 1
		return port
	}
}

/* Return a port number to the available ports slice */
func (m *backendManager) returnPort(portNum int) {
	m.lockPorts.Lock()
	m.availPorts = append(m.availPorts, portNum)
	m.lockPorts.Unlock()
}

func newBackendManager(config *Configuration) *backendManager {
	manager := &backendManager{
		backends:     make(map[string]*Backend),
		suffixLength: len(config.DomainSuffix),
		currentPort:  config.BasePort,
		proxyPrefix:  template.Must(template.New("p").Parse(config.ProxyFormat)),
	}
	return manager
}
<commit_msg>Update mutex to be part of the manager<commit_after>package stager

import (
	"bytes"
	"fmt"
	"net/http/httputil"
	"net/url"
	"sync"
	"text/template"
)

type Backend struct {
	Port    int
	Name    string
	proxy   *httputil.ReverseProxy
	running bool
}

// This will do the setup of the backend at some point.
func (b *Backend) initialize() {
	b.running = true
}

type backendManager struct {
	sync.Mutex
	backends     map[string]*Backend
	suffixLength int
	currentPort  int
	availPorts   []int
	proxyPrefix  *template.Template
}

func (m *backendManager) get(domain string) *Backend {
	name := domain[:len(domain)-m.suffixLength]
	if m.backends[name] == nil {
		port := m.allocatePort()
		b := &Backend{
			Port: port,
			Name: name,
		}
		m.backends[name] = b
		buf := &bytes.Buffer{}
		err := m.proxyPrefix.Execute(buf, b)
		if err != nil {
			panic(err)
		}
		rawurl := string(buf.Bytes())
		u, err := url.Parse(rawurl)
		if err != nil {
			panic(err)
		}
		fmt.Printf("making new instance %s on port %d with backend url %s\n", name, port, rawurl)

		b.proxy = httputil.NewSingleHostReverseProxy(u)

		go b.initialize()

	}
	return m.backends[name]
}

/* Allocate a port number to be used for another backend. */
func (m *backendManager) allocatePort() int {
	m.Lock()
	defer m.Unlock()
	l := len(m.availPorts)
	if l > 0 {
		port := m.availPorts[l-1]
		m.availPorts = m.availPorts[:l-1]
		return port
	} else {
		port := m.currentPort
		m.currentPort += 1
		return port
	}
}

/* Return a port number to the available ports slice */
func (m *backendManager) returnPort(portNum int) {
	m.Lock()
	m.availPorts = append(m.availPorts, portNum)
	m.Unlock()
}

func newBackendManager(config *Configuration) *backendManager {
	manager := &backendManager{
		backends:     make(map[string]*Backend),
		suffixLength: len(config.DomainSuffix),
		currentPort:  config.BasePort,
		proxyPrefix:  template.Must(template.New("p").Parse(config.ProxyFormat)),
	}
	return manager
}
<|endoftext|>


<commit_before>package optimize

import (
	"math"

	"github.com/gonum/floats"
)

const (
	iterationRestartFactor = 6
	angleRestartThreshold  = -0.9
)

// CGVariant calculates the scaling parameter, β, used for updating the
// conjugate direction in the nonlinear conjugate gradient (CG) method.
type CGVariant interface {
	// Init is called at the first iteration and provides a way to initialize
	// any internal state.
	Init(loc *Location)
	// Beta returns the value of the scaling parameter that is computed
	// according to the particular variant of the CG method.
	Beta(grad, gradPrev, dirPrev []float64) float64
}

// CG implements the nonlinear conjugate gradient method for solving nonlinear
// unconstrained optimization problems. It is a line search method that
// generates the search directions d_k according to the formula
//  d_{k+1} = -∇f_{k+1} + β_k*d_k,   d_0 = -∇f_0.
// Variants of the conjugate gradient method differ in the choice of the
// parameter β_k. The conjugate gradient method usually requires fewer function
// evaluations than the gradient descent method and no matrix storage, but
// L-BFGS is usually more efficient.
//
// CG implements a restart strategy that takes the steepest descent direction
// (i.e., d_{k+1} = -∇f_{k+1}) whenever any of the following conditions holds:
//
//  - A certain number of iterations has elapsed without a restart. This number
//    is controllable via IterationRestartFactor and if equal to 0, it is set to
//    a reasonable default based on the problem dimension.
//  - The angle between the gradients at two consecutive iterations ∇f_k and
//    ∇f_{k+1} is too large.
//  - The direction d_{k+1} is not a descent direction.
//  - β_k returned from CGVariant.Beta is equal to zero.
//
// The line search for CG must yield step sizes that satisfy the strong Wolfe
// conditions at every iteration, otherwise the generated search direction
// might fail to be a descent direction. The line search should be more
// stringent compared with those for Newton-like methods, which can be achieved
// by setting the gradient constant in the strong Wolfe conditions to a small
// value.
//
// See also William Hager, Hongchao Zhang, A survey of nonlinear conjugate
// gradient methods. Pacific Journal of Optimization, 2 (2006), pp. 35-58, and
// references therein.
type CG struct {
	// Linesearcher must satisfy the strong Wolfe conditions at every iteration.
	// If Linesearcher == nil, an appropriate default is chosen.
	Linesearcher Linesearcher
	// Variant implements the particular CG formula for computing β_k.
	// If Variant is nil, an appropriate default is chosen.
	Variant CGVariant
	// InitialStep estimates the initial line search step size, because the CG
	// method does not generate well-scaled search directions.
	// If InitialStep is nil, an appropriate default is chosen.
	InitialStep StepSizer

	// IterationRestartFactor determines the frequency of restarts based on the
	// problem dimension. The negative gradient direction is taken whenever
	// ceil(IterationRestartFactor*(problem dimension)) iterations have elapsed
	// without a restart. For medium and large-scale problems
	// IterationRestartFactor should be set to 1, low-dimensional problems a
	// larger value should be chosen. Note that if the ceil function returns 1,
	// CG will be identical to gradient descent.
	// If IterationRestartFactor is 0, it will be set to 6.
	// CG will panic if IterationRestartFactor is negative.
	IterationRestartFactor float64
	// AngleRestartThreshold sets the threshold angle for restart. The method
	// is restarted if the cosine of the angle between two consecutive
	// gradients is smaller than or equal to AngleRestartThreshold, that is, if
	//  ∇f_k·∇f_{k+1} / (|∇f_k| |∇f_{k+1}|) <= AngleRestartThreshold.
	// A value of AngleRestartThreshold closer to -1 (successive gradients in
	// exact opposite directions) will tend to reduce the number of restarts.
	// If AngleRestartThreshold is 0, it will be set to -0.9.
	// CG will panic if AngleRestartThreshold is not in the interval [-1, 0].
	AngleRestartThreshold float64

	ls *LinesearchMethod

	restartAfter    int
	iterFromRestart int

	dirPrev      []float64
	gradPrev     []float64
	gradPrevNorm float64
}

func (cg *CG) Init(loc *Location) (Operation, error) {
	if cg.IterationRestartFactor < 0 {
		panic("cg: IterationRestartFactor is negative")
	}
	if cg.AngleRestartThreshold < -1 || cg.AngleRestartThreshold > 0 {
		panic("cg: AngleRestartThreshold not in [-1, 0]")
	}

	if cg.Linesearcher == nil {
		cg.Linesearcher = &Bisection{GradConst: 0.1}
	}
	if cg.Variant == nil {
		cg.Variant = &HestenesStiefel{}
	}
	if cg.InitialStep == nil {
		cg.InitialStep = &FirstOrderStepSize{}
	}

	if cg.IterationRestartFactor == 0 {
		cg.IterationRestartFactor = iterationRestartFactor
	}
	if cg.AngleRestartThreshold == 0 {
		cg.AngleRestartThreshold = angleRestartThreshold
	}

	if cg.ls == nil {
		cg.ls = &LinesearchMethod{}
	}
	cg.ls.Linesearcher = cg.Linesearcher
	cg.ls.NextDirectioner = cg

	return cg.ls.Init(loc)
}

func (cg *CG) Iterate(loc *Location) (Operation, error) {
	return cg.ls.Iterate(loc)
}

func (cg *CG) InitDirection(loc *Location, dir []float64) (stepSize float64) {
	dim := len(loc.X)

	cg.restartAfter = int(math.Ceil(cg.IterationRestartFactor * float64(dim)))
	cg.iterFromRestart = 0

	// The initial direction is always the negative gradient.
	copy(dir, loc.Gradient)
	floats.Scale(-1, dir)

	cg.dirPrev = resize(cg.dirPrev, dim)
	copy(cg.dirPrev, dir)
	cg.gradPrev = resize(cg.gradPrev, dim)
	copy(cg.gradPrev, loc.Gradient)
	cg.gradPrevNorm = floats.Norm(loc.Gradient, 2)

	cg.Variant.Init(loc)
	return cg.InitialStep.Init(loc, dir)
}

func (cg *CG) NextDirection(loc *Location, dir []float64) (stepSize float64) {
	copy(dir, loc.Gradient)
	floats.Scale(-1, dir)

	cg.iterFromRestart++
	var restart bool
	if cg.iterFromRestart == cg.restartAfter {
		// Restart because too many iterations have been taken without a restart.
		restart = true
	}

	gDot := floats.Dot(loc.Gradient, cg.gradPrev)
	gNorm := floats.Norm(loc.Gradient, 2)
	if gDot <= cg.AngleRestartThreshold*gNorm*cg.gradPrevNorm {
		// Restart because the angle between the last two gradients is too large.
		restart = true
	}

	// Compute the scaling factor β_k even when restarting, because cg.Variant
	// may be keeping an inner state that needs to be updated at every iteration.
	beta := cg.Variant.Beta(loc.Gradient, cg.gradPrev, cg.dirPrev)
	if beta == 0 {
		// β_k == 0 means that the steepest descent direction will be taken, so
		// indicate that the method is in fact being restarted.
		restart = true
	}
	if !restart {
		// The method is not being restarted, so update the descent direction.
		floats.AddScaled(dir, beta, cg.dirPrev)
		if floats.Dot(loc.Gradient, dir) >= 0 {
			// Restart because the new direction is not a descent direction.
			restart = true
			copy(dir, loc.Gradient)
			floats.Scale(-1, dir)
		}
	}

	// Get the initial line search step size from the StepSizer even if the
	// method was restarted, because StepSizers need to see every iteration.
	stepSize = cg.InitialStep.StepSize(loc, dir)
	if restart {
		// The method was restarted and since the steepest descent direction is
		// not related to the previous direction, discard the estimated step
		// size from cg.InitialStep and use step size of 1 instead.
		stepSize = 1
		// Reset to 0 the counter of iterations taken since the last restart.
		cg.iterFromRestart = 0
	}

	copy(cg.gradPrev, loc.Gradient)
	copy(cg.dirPrev, dir)
	cg.gradPrevNorm = gNorm
	return stepSize
}

func (*CG) Needs() struct {
	Gradient bool
	Hessian  bool
} {
	return struct {
		Gradient bool
		Hessian  bool
	}{true, false}
}

// FletcherReeves implements the Fletcher-Reeves variant of the CG method that
// computes the scaling parameter β_k according to the formula
//  β_k = |∇f_{k+1}|^2 / |∇f_k|^2.
type FletcherReeves struct {
	prevNorm float64
}

func (fr *FletcherReeves) Init(loc *Location) {
	fr.prevNorm = floats.Norm(loc.Gradient, 2)
}

func (fr *FletcherReeves) Beta(grad, _, _ []float64) (beta float64) {
	norm := floats.Norm(grad, 2)
	beta = (norm / fr.prevNorm) * (norm / fr.prevNorm)
	fr.prevNorm = norm
	return beta
}

// PolakRibierePolyak implements the Polak-Ribiere-Polyak variant of the CG
// method that computes the scaling parameter β_k according to the formula
//  β_k = max(0, ∇f_{k+1}·y_k / |∇f_k|^2),
// where y_k = ∇f_{k+1} - ∇f_k.
type PolakRibierePolyak struct {
	prevNorm float64
}

func (pr *PolakRibierePolyak) Init(loc *Location) {
	pr.prevNorm = floats.Norm(loc.Gradient, 2)
}

func (pr *PolakRibierePolyak) Beta(grad, gradPrev, _ []float64) (beta float64) {
	norm := floats.Norm(grad, 2)
	dot := floats.Dot(grad, gradPrev)
	beta = (norm*norm - dot) / (pr.prevNorm * pr.prevNorm)
	pr.prevNorm = norm
	return math.Max(0, beta)
}

// HestenesStiefel implements the Hestenes-Stiefel variant of the CG method
// that computes the scaling parameter β_k according to the formula
//  β_k = max(0, ∇f_{k+1}·y_k / d_k·y_k),
// where y_k = ∇f_{k+1} - ∇f_k.
type HestenesStiefel struct {
	y []float64
}

func (hs *HestenesStiefel) Init(loc *Location) {
	hs.y = resize(hs.y, len(loc.Gradient))
}

func (hs *HestenesStiefel) Beta(grad, gradPrev, dirPrev []float64) (beta float64) {
	floats.SubTo(hs.y, grad, gradPrev)
	beta = floats.Dot(grad, hs.y) / floats.Dot(dirPrev, hs.y)
	return math.Max(0, beta)
}

// DaiYuan implements the Dai-Yuan variant of the CG method that computes the
// scaling parameter β_k according to the formula
//  β_k = |∇f_{k+1}|^2 / d_k·y_k,
// where y_k = ∇f_{k+1} - ∇f_k.
type DaiYuan struct {
	y []float64
}

func (dy *DaiYuan) Init(loc *Location) {
	dy.y = resize(dy.y, len(loc.Gradient))
}

func (dy *DaiYuan) Beta(grad, gradPrev, dirPrev []float64) (beta float64) {
	floats.SubTo(dy.y, grad, gradPrev)
	norm := floats.Norm(grad, 2)
	return norm * norm / floats.Dot(dirPrev, dy.y)
}

// HagerZhang implements the Hager-Zhang variant of the CG method that computes the
// scaling parameter β_k according to the formula
//  β_k = (y_k - 2 d_k |y_k|^2/(d_k·y_k))·∇f_{k+1} / (d_k·y_k),
// where y_k = ∇f_{k+1} - ∇f_k.
type HagerZhang struct {
	y []float64
}

func (hz *HagerZhang) Init(loc *Location) {
	hz.y = resize(hz.y, len(loc.Gradient))
}

func (hz *HagerZhang) Beta(grad, gradPrev, dirPrev []float64) (beta float64) {
	floats.SubTo(hz.y, grad, gradPrev)
	dirDotY := floats.Dot(dirPrev, hz.y)
	gDotY := floats.Dot(grad, hz.y)
	gDotDir := floats.Dot(grad, dirPrev)
	yNorm := floats.Norm(hz.y, 2)
	return (gDotY - 2*gDotDir*yNorm*yNorm/dirDotY) / dirDotY
}
<commit_msg>Use MoreThuente in CG as default Linesearcher<commit_after>package optimize

import (
	"math"

	"github.com/gonum/floats"
)

const (
	iterationRestartFactor = 6
	angleRestartThreshold  = -0.9
)

// CGVariant calculates the scaling parameter, β, used for updating the
// conjugate direction in the nonlinear conjugate gradient (CG) method.
type CGVariant interface {
	// Init is called at the first iteration and provides a way to initialize
	// any internal state.
	Init(loc *Location)
	// Beta returns the value of the scaling parameter that is computed
	// according to the particular variant of the CG method.
	Beta(grad, gradPrev, dirPrev []float64) float64
}

// CG implements the nonlinear conjugate gradient method for solving nonlinear
// unconstrained optimization problems. It is a line search method that
// generates the search directions d_k according to the formula
//  d_{k+1} = -∇f_{k+1} + β_k*d_k,   d_0 = -∇f_0.
// Variants of the conjugate gradient method differ in the choice of the
// parameter β_k. The conjugate gradient method usually requires fewer function
// evaluations than the gradient descent method and no matrix storage, but
// L-BFGS is usually more efficient.
//
// CG implements a restart strategy that takes the steepest descent direction
// (i.e., d_{k+1} = -∇f_{k+1}) whenever any of the following conditions holds:
//
//  - A certain number of iterations has elapsed without a restart. This number
//    is controllable via IterationRestartFactor and if equal to 0, it is set to
//    a reasonable default based on the problem dimension.
//  - The angle between the gradients at two consecutive iterations ∇f_k and
//    ∇f_{k+1} is too large.
//  - The direction d_{k+1} is not a descent direction.
//  - β_k returned from CGVariant.Beta is equal to zero.
//
// The line search for CG must yield step sizes that satisfy the strong Wolfe
// conditions at every iteration, otherwise the generated search direction
// might fail to be a descent direction. The line search should be more
// stringent compared with those for Newton-like methods, which can be achieved
// by setting the gradient constant in the strong Wolfe conditions to a small
// value.
//
// See also William Hager, Hongchao Zhang, A survey of nonlinear conjugate
// gradient methods. Pacific Journal of Optimization, 2 (2006), pp. 35-58, and
// references therein.
type CG struct {
	// Linesearcher must satisfy the strong Wolfe conditions at every iteration.
	// If Linesearcher == nil, an appropriate default is chosen.
	Linesearcher Linesearcher
	// Variant implements the particular CG formula for computing β_k.
	// If Variant is nil, an appropriate default is chosen.
	Variant CGVariant
	// InitialStep estimates the initial line search step size, because the CG
	// method does not generate well-scaled search directions.
	// If InitialStep is nil, an appropriate default is chosen.
	InitialStep StepSizer

	// IterationRestartFactor determines the frequency of restarts based on the
	// problem dimension. The negative gradient direction is taken whenever
	// ceil(IterationRestartFactor*(problem dimension)) iterations have elapsed
	// without a restart. For medium and large-scale problems
	// IterationRestartFactor should be set to 1, low-dimensional problems a
	// larger value should be chosen. Note that if the ceil function returns 1,
	// CG will be identical to gradient descent.
	// If IterationRestartFactor is 0, it will be set to 6.
	// CG will panic if IterationRestartFactor is negative.
	IterationRestartFactor float64
	// AngleRestartThreshold sets the threshold angle for restart. The method
	// is restarted if the cosine of the angle between two consecutive
	// gradients is smaller than or equal to AngleRestartThreshold, that is, if
	//  ∇f_k·∇f_{k+1} / (|∇f_k| |∇f_{k+1}|) <= AngleRestartThreshold.
	// A value of AngleRestartThreshold closer to -1 (successive gradients in
	// exact opposite directions) will tend to reduce the number of restarts.
	// If AngleRestartThreshold is 0, it will be set to -0.9.
	// CG will panic if AngleRestartThreshold is not in the interval [-1, 0].
	AngleRestartThreshold float64

	ls *LinesearchMethod

	restartAfter    int
	iterFromRestart int

	dirPrev      []float64
	gradPrev     []float64
	gradPrevNorm float64
}

func (cg *CG) Init(loc *Location) (Operation, error) {
	if cg.IterationRestartFactor < 0 {
		panic("cg: IterationRestartFactor is negative")
	}
	if cg.AngleRestartThreshold < -1 || cg.AngleRestartThreshold > 0 {
		panic("cg: AngleRestartThreshold not in [-1, 0]")
	}

	if cg.Linesearcher == nil {
		cg.Linesearcher = &MoreThuente{CurvatureFactor: 0.1}
	}
	if cg.Variant == nil {
		cg.Variant = &HestenesStiefel{}
	}
	if cg.InitialStep == nil {
		cg.InitialStep = &FirstOrderStepSize{}
	}

	if cg.IterationRestartFactor == 0 {
		cg.IterationRestartFactor = iterationRestartFactor
	}
	if cg.AngleRestartThreshold == 0 {
		cg.AngleRestartThreshold = angleRestartThreshold
	}

	if cg.ls == nil {
		cg.ls = &LinesearchMethod{}
	}
	cg.ls.Linesearcher = cg.Linesearcher
	cg.ls.NextDirectioner = cg

	return cg.ls.Init(loc)
}

func (cg *CG) Iterate(loc *Location) (Operation, error) {
	return cg.ls.Iterate(loc)
}

func (cg *CG) InitDirection(loc *Location, dir []float64) (stepSize float64) {
	dim := len(loc.X)

	cg.restartAfter = int(math.Ceil(cg.IterationRestartFactor * float64(dim)))
	cg.iterFromRestart = 0

	// The initial direction is always the negative gradient.
	copy(dir, loc.Gradient)
	floats.Scale(-1, dir)

	cg.dirPrev = resize(cg.dirPrev, dim)
	copy(cg.dirPrev, dir)
	cg.gradPrev = resize(cg.gradPrev, dim)
	copy(cg.gradPrev, loc.Gradient)
	cg.gradPrevNorm = floats.Norm(loc.Gradient, 2)

	cg.Variant.Init(loc)
	return cg.InitialStep.Init(loc, dir)
}

func (cg *CG) NextDirection(loc *Location, dir []float64) (stepSize float64) {
	copy(dir, loc.Gradient)
	floats.Scale(-1, dir)

	cg.iterFromRestart++
	var restart bool
	if cg.iterFromRestart == cg.restartAfter {
		// Restart because too many iterations have been taken without a restart.
		restart = true
	}

	gDot := floats.Dot(loc.Gradient, cg.gradPrev)
	gNorm := floats.Norm(loc.Gradient, 2)
	if gDot <= cg.AngleRestartThreshold*gNorm*cg.gradPrevNorm {
		// Restart because the angle between the last two gradients is too large.
		restart = true
	}

	// Compute the scaling factor β_k even when restarting, because cg.Variant
	// may be keeping an inner state that needs to be updated at every iteration.
	beta := cg.Variant.Beta(loc.Gradient, cg.gradPrev, cg.dirPrev)
	if beta == 0 {
		// β_k == 0 means that the steepest descent direction will be taken, so
		// indicate that the method is in fact being restarted.
		restart = true
	}
	if !restart {
		// The method is not being restarted, so update the descent direction.
		floats.AddScaled(dir, beta, cg.dirPrev)
		if floats.Dot(loc.Gradient, dir) >= 0 {
			// Restart because the new direction is not a descent direction.
			restart = true
			copy(dir, loc.Gradient)
			floats.Scale(-1, dir)
		}
	}

	// Get the initial line search step size from the StepSizer even if the
	// method was restarted, because StepSizers need to see every iteration.
	stepSize = cg.InitialStep.StepSize(loc, dir)
	if restart {
		// The method was restarted and since the steepest descent direction is
		// not related to the previous direction, discard the estimated step
		// size from cg.InitialStep and use step size of 1 instead.
		stepSize = 1
		// Reset to 0 the counter of iterations taken since the last restart.
		cg.iterFromRestart = 0
	}

	copy(cg.gradPrev, loc.Gradient)
	copy(cg.dirPrev, dir)
	cg.gradPrevNorm = gNorm
	return stepSize
}

func (*CG) Needs() struct {
	Gradient bool
	Hessian  bool
} {
	return struct {
		Gradient bool
		Hessian  bool
	}{true, false}
}

// FletcherReeves implements the Fletcher-Reeves variant of the CG method that
// computes the scaling parameter β_k according to the formula
//  β_k = |∇f_{k+1}|^2 / |∇f_k|^2.
type FletcherReeves struct {
	prevNorm float64
}

func (fr *FletcherReeves) Init(loc *Location) {
	fr.prevNorm = floats.Norm(loc.Gradient, 2)
}

func (fr *FletcherReeves) Beta(grad, _, _ []float64) (beta float64) {
	norm := floats.Norm(grad, 2)
	beta = (norm / fr.prevNorm) * (norm / fr.prevNorm)
	fr.prevNorm = norm
	return beta
}

// PolakRibierePolyak implements the Polak-Ribiere-Polyak variant of the CG
// method that computes the scaling parameter β_k according to the formula
//  β_k = max(0, ∇f_{k+1}·y_k / |∇f_k|^2),
// where y_k = ∇f_{k+1} - ∇f_k.
type PolakRibierePolyak struct {
	prevNorm float64
}

func (pr *PolakRibierePolyak) Init(loc *Location) {
	pr.prevNorm = floats.Norm(loc.Gradient, 2)
}

func (pr *PolakRibierePolyak) Beta(grad, gradPrev, _ []float64) (beta float64) {
	norm := floats.Norm(grad, 2)
	dot := floats.Dot(grad, gradPrev)
	beta = (norm*norm - dot) / (pr.prevNorm * pr.prevNorm)
	pr.prevNorm = norm
	return math.Max(0, beta)
}

// HestenesStiefel implements the Hestenes-Stiefel variant of the CG method
// that computes the scaling parameter β_k according to the formula
//  β_k = max(0, ∇f_{k+1}·y_k / d_k·y_k),
// where y_k = ∇f_{k+1} - ∇f_k.
type HestenesStiefel struct {
	y []float64
}

func (hs *HestenesStiefel) Init(loc *Location) {
	hs.y = resize(hs.y, len(loc.Gradient))
}

func (hs *HestenesStiefel) Beta(grad, gradPrev, dirPrev []float64) (beta float64) {
	floats.SubTo(hs.y, grad, gradPrev)
	beta = floats.Dot(grad, hs.y) / floats.Dot(dirPrev, hs.y)
	return math.Max(0, beta)
}

// DaiYuan implements the Dai-Yuan variant of the CG method that computes the
// scaling parameter β_k according to the formula
//  β_k = |∇f_{k+1}|^2 / d_k·y_k,
// where y_k = ∇f_{k+1} - ∇f_k.
type DaiYuan struct {
	y []float64
}

func (dy *DaiYuan) Init(loc *Location) {
	dy.y = resize(dy.y, len(loc.Gradient))
}

func (dy *DaiYuan) Beta(grad, gradPrev, dirPrev []float64) (beta float64) {
	floats.SubTo(dy.y, grad, gradPrev)
	norm := floats.Norm(grad, 2)
	return norm * norm / floats.Dot(dirPrev, dy.y)
}

// HagerZhang implements the Hager-Zhang variant of the CG method that computes the
// scaling parameter β_k according to the formula
//  β_k = (y_k - 2 d_k |y_k|^2/(d_k·y_k))·∇f_{k+1} / (d_k·y_k),
// where y_k = ∇f_{k+1} - ∇f_k.
type HagerZhang struct {
	y []float64
}

func (hz *HagerZhang) Init(loc *Location) {
	hz.y = resize(hz.y, len(loc.Gradient))
}

func (hz *HagerZhang) Beta(grad, gradPrev, dirPrev []float64) (beta float64) {
	floats.SubTo(hz.y, grad, gradPrev)
	dirDotY := floats.Dot(dirPrev, hz.y)
	gDotY := floats.Dot(grad, hz.y)
	gDotDir := floats.Dot(grad, dirPrev)
	yNorm := floats.Norm(hz.y, 2)
	return (gDotY - 2*gDotDir*yNorm*yNorm/dirDotY) / dirDotY
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"

	log "github.com/couchbaselabs/clog"
)

// Compares two dotted versioning string and returns true when x >= y.
func VersionGTE(x, y string) bool {
	xa := strings.Split(x, ".")
	ya := strings.Split(y, ".")
	for i := range xa {
		if i >= len(ya) {
			return true
		}
		xv, err := strconv.Atoi(xa[i])
		if err != nil {
			return false
		}
		yv, err := strconv.Atoi(ya[i])
		if err != nil {
			return false
		}
		if xv < yv {
			return false
		}
	}
	return len(xa) >= len(ya)
}

func NewUUID() string {
	val1 := rand.Int63()
	val2 := rand.Int63()
	uuid := fmt.Sprintf("%x%x", val1, val2)
	return uuid[0:16]
}

// Calls f() in a loop, sleeping in an exponential backoff if needed.
// The provided f() function should return < 0 to stop the loop; >= 0
// to continue the loop, where > 0 means there was progress which
// allows an immediate retry of f() with no sleeping.  A return of < 0
// is useful when f() will never make any future progress.
func ExponentialBackoffLoop(name string,
	f func() int,
	startSleepMS int,
	backoffFactor float32,
	maxSleepMS int) {
	nextSleepMS := startSleepMS
	for {
		progress := f()
		if progress < 0 {
			return
		}
		if progress > 0 {
			// When there was some progress, we can reset nextSleepMS.
			log.Printf("backoff: %s, progress: %d", name, progress)
			nextSleepMS = startSleepMS
		} else {
			// If zero progress was made this cycle, then sleep.
			log.Printf("backoff: %s, sleep: %d (ms)", name, nextSleepMS)
			time.Sleep(time.Duration(nextSleepMS) * time.Millisecond)

			// Increase nextSleepMS in case next time also has 0 progress.
			nextSleepMS = int(float32(nextSleepMS) * backoffFactor)
			if nextSleepMS > maxSleepMS {
				nextSleepMS = maxSleepMS
			}
		}
	}
}
<commit_msg>comment on VersionGTE()<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"

	log "github.com/couchbaselabs/clog"
)

// Compares two dotted versioning strings, like "1.0.1" and "1.2.3".
// Returns true when x >= y.
func VersionGTE(x, y string) bool {
	xa := strings.Split(x, ".")
	ya := strings.Split(y, ".")
	for i := range xa {
		if i >= len(ya) {
			return true
		}
		xv, err := strconv.Atoi(xa[i])
		if err != nil {
			return false
		}
		yv, err := strconv.Atoi(ya[i])
		if err != nil {
			return false
		}
		if xv < yv {
			return false
		}
	}
	return len(xa) >= len(ya)
}

func NewUUID() string {
	val1 := rand.Int63()
	val2 := rand.Int63()
	uuid := fmt.Sprintf("%x%x", val1, val2)
	return uuid[0:16]
}

// Calls f() in a loop, sleeping in an exponential backoff if needed.
// The provided f() function should return < 0 to stop the loop; >= 0
// to continue the loop, where > 0 means there was progress which
// allows an immediate retry of f() with no sleeping.  A return of < 0
// is useful when f() will never make any future progress.
func ExponentialBackoffLoop(name string,
	f func() int,
	startSleepMS int,
	backoffFactor float32,
	maxSleepMS int) {
	nextSleepMS := startSleepMS
	for {
		progress := f()
		if progress < 0 {
			return
		}
		if progress > 0 {
			// When there was some progress, we can reset nextSleepMS.
			log.Printf("backoff: %s, progress: %d", name, progress)
			nextSleepMS = startSleepMS
		} else {
			// If zero progress was made this cycle, then sleep.
			log.Printf("backoff: %s, sleep: %d (ms)", name, nextSleepMS)
			time.Sleep(time.Duration(nextSleepMS) * time.Millisecond)

			// Increase nextSleepMS in case next time also has 0 progress.
			nextSleepMS = int(float32(nextSleepMS) * backoffFactor)
			if nextSleepMS > maxSleepMS {
				nextSleepMS = maxSleepMS
			}
		}
	}
}
<|endoftext|>


<commit_before>package jsonrpc2

import "context"

// AsyncHandler wraps a Handler such that it each request is handled in its
// own goroutine. It is a convenience wrapper.
func AsyncHandler(h Handler) Handler {
	return asyncHandler{h}
}

type asyncHandler struct {
	Handler
}

func (h asyncHandler) Handle(ctx context.Context, conn *Conn, req *Request) {
	go h.Handler.Handle(ctx, conn, req)
}
<commit_msg>typo in AsyncHandler docstring<commit_after>package jsonrpc2

import "context"

// AsyncHandler wraps a Handler such that each request is handled in its own
// goroutine. It is a convenience wrapper.
func AsyncHandler(h Handler) Handler {
	return asyncHandler{h}
}

type asyncHandler struct {
	Handler
}

func (h asyncHandler) Handle(ctx context.Context, conn *Conn, req *Request) {
	go h.Handler.Handle(ctx, conn, req)
}
<|endoftext|>


<commit_before>// bulk_query_gen generates queries for various use cases. Its output will
// typically be consume by query_benchmarker.
//
// Output formats: InfluxDB, ElasticSearch.
//
// Query style use cases: Devops.
package main

import (
	"bufio"
	"encoding/gob"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"sort"
	"time"
)

// Output data format choices:
var formatChoices = []string{"influx-http", "es-http"}

// Use case choices:
var useCaseChoices = []string{"devops"}

// Program option vars:
var (
	format     string
	useCase    string
	scaleVar   int
	queryCount int

	dbName string

	timestampStartStr string
	timestampEndStr   string

	timestampStart time.Time
	timestampEnd   time.Time

	seed  int64
	debug int
)

// Parse args:
func init() {
	flag.StringVar(&format, "format", formatChoices[0], "Format to emit. (choices: influx-http, es-http)")
	flag.StringVar(&useCase, "use-case", useCaseChoices[0], "Use case to model. (choices: devops, iot)")
	flag.IntVar(&scaleVar, "scale-var", 1, "Scaling variable (must be the equal to the scalevar used for data generation).")
	flag.IntVar(&queryCount, "queries", 1000, "Number of queries to generate.")

	flag.StringVar(&dbName, "db", "benchmark_db", "Database for influx to use (ignored for elastic)")

	flag.StringVar(&timestampStartStr, "timestamp-start", "2016-01-01T00:00:00Z", "Beginning timestamp (RFC3339).")
	flag.StringVar(&timestampEndStr, "timestamp-end", "2016-01-01T06:00:00Z", "Ending timestamp (RFC3339).")

	flag.Int64Var(&seed, "seed", 0, "PRNG seed (default, or 0, uses the current timestamp).")
	flag.IntVar(&debug, "debug", 0, "Debug printing (choices: 0, 1) (default 0).")

	flag.Parse()

	validFormat := false
	for _, s := range formatChoices {
		if s == format {
			validFormat = true
			break
		}
	}
	if !validFormat {
		log.Fatal("invalid format specifier")
	}

	// the default seed is the current timestamp:
	if seed == 0 {
		seed = int64(time.Now().Nanosecond())
	}
	fmt.Fprintf(os.Stderr, "using random seed %d\n", seed)

	// Parse timestamps:
	var err error
	timestampStart, err = time.Parse(time.RFC3339, timestampStartStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampStart = timestampStart.UTC()
	timestampEnd, err = time.Parse(time.RFC3339, timestampEndStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampEnd = timestampEnd.UTC()
}

func main() {
	rand.Seed(seed)

	var generator QueryGenerator
	switch useCase {
	case "devops":
		switch format {
		case "influx-http":
			generator = NewInfluxDevops(dbName, timestampStart, timestampEnd)
		case "es-http":
			generator = NewElasticSearchDevops(timestampStart, timestampEnd)
		default:
			panic("invalid format")
		}
	default:
		panic("invalid use case")
	}

	stats := make(map[string]int64)

	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	// create request instances, serializing them to stdout and collecting
	// counts for each kind:
	enc := gob.NewEncoder(out)
	q := &Query{}
	for i := 0; i < queryCount; i++ {
		generator.Dispatch(i, q, scaleVar)
		err := enc.Encode(q)
		if err != nil {
			log.Fatal(err)
		}
		stats[string(q.HumanLabel)]++

		if debug == 1 {
			_, err := fmt.Fprintf(os.Stderr, "%s\n", q.HumanLabel)
			if err != nil {
				log.Fatal(err)
			}
		} else if debug == 2 {
			_, err := fmt.Fprintf(os.Stderr, "%s\n", q.HumanDescription)
			if err != nil {
				log.Fatal(err)
			}
		} else if debug >= 3 {
			_, err := fmt.Fprintf(os.Stderr, "%s\n", q.String())
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	// print stats:
	keys := []string{}
	for k, _ := range stats {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		_, err := fmt.Fprintf(os.Stderr, "%s: %d points\n", k, stats[k])
		if err != nil {
			log.Fatal(err)
		}
	}
}
<commit_msg>Update query gen use case names.<commit_after>// bulk_query_gen generates queries for various use cases. Its output will
// typically be consume by query_benchmarker.
//
// Output formats: InfluxDB, ElasticSearch.
//
// Query style use cases: Devops.
package main

import (
	"bufio"
	"encoding/gob"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"sort"
	"time"
)

// Output data format choices:
var formatChoices = []string{"influx-http", "es-http"}

// Use case choices:
var useCaseChoices = []string{"devops-single-host", "devops-full-dataset", "devops-groupby-hosts"}

// Program option vars:
var (
	format     string
	useCase    string
	scaleVar   int
	queryCount int

	dbName string

	timestampStartStr string
	timestampEndStr   string

	timestampStart time.Time
	timestampEnd   time.Time

	seed  int64
	debug int
)

// Parse args:
func init() {
	flag.StringVar(&format, "format", formatChoices[0], "Format to emit. (choices: influx-http, es-http)")
	flag.StringVar(&useCase, "use-case", useCaseChoices[0], fmt.Sprintf("Use case to model. (choices: %s)", useCaseChoices))
	flag.IntVar(&scaleVar, "scale-var", 1, "Scaling variable (must be the equal to the scalevar used for data generation).")
	flag.IntVar(&queryCount, "queries", 1000, "Number of queries to generate.")

	flag.StringVar(&dbName, "db", "benchmark_db", "Database for influx to use (ignored for elastic)")

	flag.StringVar(&timestampStartStr, "timestamp-start", "2016-01-01T00:00:00Z", "Beginning timestamp (RFC3339).")
	flag.StringVar(&timestampEndStr, "timestamp-end", "2016-01-01T06:00:00Z", "Ending timestamp (RFC3339).")

	flag.Int64Var(&seed, "seed", 0, "PRNG seed (default, or 0, uses the current timestamp).")
	flag.IntVar(&debug, "debug", 0, "Debug printing (choices: 0, 1) (default 0).")

	flag.Parse()

	validFormat := false
	for _, s := range formatChoices {
		if s == format {
			validFormat = true
			break
		}
	}
	if !validFormat {
		log.Fatal("invalid format specifier")
	}

	// the default seed is the current timestamp:
	if seed == 0 {
		seed = int64(time.Now().Nanosecond())
	}
	fmt.Fprintf(os.Stderr, "using random seed %d\n", seed)

	// Parse timestamps:
	var err error
	timestampStart, err = time.Parse(time.RFC3339, timestampStartStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampStart = timestampStart.UTC()
	timestampEnd, err = time.Parse(time.RFC3339, timestampEndStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampEnd = timestampEnd.UTC()
}

func main() {
	rand.Seed(seed)

	var generator QueryGenerator
	switch useCase {
	case "devops-single-host":
		switch format {
		case "influx-http":
			generator = NewInfluxDevops(dbName, timestampStart, timestampEnd)
		case "es-http":
			generator = NewElasticSearchDevops(timestampStart, timestampEnd)
		default:
			panic("invalid format")
		}
	default:
		panic("invalid use case")
	}

	stats := make(map[string]int64)

	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	// create request instances, serializing them to stdout and collecting
	// counts for each kind:
	enc := gob.NewEncoder(out)
	q := &Query{}
	for i := 0; i < queryCount; i++ {
		generator.Dispatch(i, q, scaleVar)
		err := enc.Encode(q)
		if err != nil {
			log.Fatal(err)
		}
		stats[string(q.HumanLabel)]++

		if debug == 1 {
			_, err := fmt.Fprintf(os.Stderr, "%s\n", q.HumanLabel)
			if err != nil {
				log.Fatal(err)
			}
		} else if debug == 2 {
			_, err := fmt.Fprintf(os.Stderr, "%s\n", q.HumanDescription)
			if err != nil {
				log.Fatal(err)
			}
		} else if debug >= 3 {
			_, err := fmt.Fprintf(os.Stderr, "%s\n", q.String())
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	// print stats:
	keys := []string{}
	for k, _ := range stats {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		_, err := fmt.Fprintf(os.Stderr, "%s: %d points\n", k, stats[k])
		if err != nil {
			log.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package stream

import (
	"archive/tar"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var (
	modeDir os.FileMode = 0755
)

// Untar untars the contents of a reader into dst
func Untar(r io.Reader, dst string) error {
	tr := tar.NewReader(r)
	root := dst

	// Special case the first file to figure out where to unpack stream.
	hdr, err := tr.Next()
	if err != nil {
		if err == io.EOF {
			return errors.New("received empty Tar file")
		}
		return err
	}

	var name string
	var rootReplace string

	// ferr is used to determine where to unpack the stream
	fi, ferr := os.Stat(dst)
	switch hdr.Typeflag {
	case tar.TypeSymlink:
		fallthrough
	case tar.TypeRegA:
		fallthrough
	case tar.TypeReg:
		if ferr == nil {
			// dst exists and user specified directory as dst
			// eg push /foo/bar/file.txt /baz/some_dir
			if fi.Mode().IsDir() {
				name = filepath.Join(dst, hdr.Name)
				root = dst
			} else {
				// dst exists and user specified file name as dst
				// eg push /foo/bar/file.txt /baz/some_dir/file.txt
				name = dst
				root = filepath.Dir(dst)
			}
			break
		}
		if os.IsNotExist(ferr) {
			// dst does not exist push to parent dir and rename to dst.
			root = filepath.Dir(dst)
			if ferr = os.MkdirAll(root, modeDir); ferr == nil {
				name = dst
				break
			}
		}
		return ferr
	case tar.TypeDir:
		if ferr == nil {
			if fi.Mode().IsDir() {
				// dst exists and is a directory. Create dir in directory
				name = filepath.Join(dst, hdr.Name)
				break
			}
			return fmt.Errorf("file %s exists and is not a directory", dst)
		} else if os.IsNotExist(ferr) {
			root = dst
			name = dst
			rootReplace = hdr.Name
			break
		}
		return ferr
	}

	for {
		switch hdr.Typeflag {
		case tar.TypeRegA:
			fallthrough
		case tar.TypeReg:
			os.RemoveAll(name)
			f, err := os.OpenFile(name, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644)
			if err != nil {
				return err
			}
			// Close file ASAP to avoid too many open files
			if _, err = io.Copy(f, tr); err != nil {
				return err
			}
			if err := f.Close(); err != nil {
				return err
			}
		case tar.TypeDir:
			if err := os.MkdirAll(name, modeDir); err != nil {
				return err
			}
		case tar.TypeSymlink:
			if err := os.Symlink(hdr.Linkname, name); err != nil {
				return err
			}
		}

		hdr, err = tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		name = filepath.Join(root, strings.TrimPrefix(hdr.Name, rootReplace))
	}
}

// UntarBytes untars a tar r to a dst writer. This assumes that the tar contains only a single
// file.
func UntarBytes(dst io.Writer, r io.Reader) error {
	tr := tar.NewReader(r)
	// Advance to the first and only file in the tar. If there are more files present in the
	// tar, return an error.
	hdr, err := tr.Next()
	if err != nil {
		return err
	}
	switch hdr.Typeflag {
	case tar.TypeReg:
		fallthrough
	case tar.TypeRegA:
		if _, err = io.Copy(dst, tr); err != nil {
			return err
		}
	default:
		return fmt.Errorf("trying to read a directory into writer")
	}
	if _, eof := tr.Next(); eof != io.EOF {
		return fmt.Errorf("too many files received when trying to untar file")
	}
	return nil
}

// Tar archive tars the contents of src into a writer stream
func Tar(writer io.Writer, src string) error {
	if _, err := os.Stat(src); err != nil {
		return fmt.Errorf("Unable to tar files - %v", err.Error())
	}

	// avoid making the directory relative to itself
	src = strings.TrimRight(src, "/")

	tw := tar.NewWriter(writer)
	defer tw.Close()

	return filepath.Walk(src, func(file string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		var hdr *tar.Header
		if fi.Mode()&os.ModeSymlink != 0 {
			t, err := filepath.EvalSymlinks(file)
			if err != nil {
				// Don't include broken symlinks
				log.Printf("Skipping broken link %s\n", file)
				return nil
			}

			// ok to ignore error file always lives inside src
			sr, _ := filepath.Rel(src, file)
			name := filepath.Join(filepath.Base(src), sr)
			if strings.HasPrefix(t, src) {
				// Pointee lives under same root as pointer, so write the symlink.
				if hdr, err = tar.FileInfoHeader(fi, t); err != nil {
					return err
				}
				hdr.Name = name
			} else {
				// Pointee lives outside pointer root, write the file pointed to.
				var ti os.FileInfo
				if ti, err = os.Stat(t); err != nil {
					// Skip broken links. Error out on everything else.
					if os.IsNotExist(err) {
						log.Printf("Skipping broken link %s -> %s\n", file, t)
						return nil
					}
					return err
				}
				if !ti.Mode().IsRegular() {
					return nil
				}
				if hdr, err = tar.FileInfoHeader(ti, ""); err != nil {
					return err
				}
				hdr.Name = name
			}
		} else {
			hdr, err = tar.FileInfoHeader(fi, fi.Name())
			if err != nil {
				return err
			}

			rel, err := filepath.Rel(filepath.Dir(src), file)
			if err != nil {
				return err
			}
			hdr.Name = rel
		}

		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}

		if (hdr.Typeflag != tar.TypeReg && hdr.Typeflag != tar.TypeRegA) || hdr.Size == 0 {
			return nil
		}

		f, err := os.Open(file)
		if err != nil {
			return err
		}
		defer f.Close()

		if _, err := io.Copy(tw, f); err != nil {
			return err
		}

		return nil
	})
}

// TarBytes tars bytes as a single logical file into a tar file writer.
func TarBytes(writer io.Writer, bytes []byte) error {
	tw := tar.NewWriter(writer)
	defer tw.Close()

	hdr := &tar.Header{
		Typeflag: tar.TypeReg,
		Name:     "bytefile",
		Size:     int64(len(bytes)),
	}
	if err := tw.WriteHeader(hdr); err != nil {
		return err
	}

	if _, err := tw.Write(bytes); err != nil {
		return err
	}
	return nil
}
<commit_msg>Improve Untar error msg when input is empty.<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package stream

import (
	"archive/tar"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var (
	modeDir os.FileMode = 0755
)

// Untar untars the contents of a reader into dst
func Untar(r io.Reader, dst string) error {
	tr := tar.NewReader(r)
	root := dst

	// Special case the first file to figure out where to unpack stream.
	hdr, err := tr.Next()
	if err != nil {
		if err == io.EOF {
			return errors.New("received empty Tar file")
		}
		return err
	}

	var name string
	var rootReplace string

	// ferr is used to determine where to unpack the stream
	fi, ferr := os.Stat(dst)
	switch hdr.Typeflag {
	case tar.TypeSymlink:
		fallthrough
	case tar.TypeRegA:
		fallthrough
	case tar.TypeReg:
		if ferr == nil {
			// dst exists and user specified directory as dst
			// eg push /foo/bar/file.txt /baz/some_dir
			if fi.Mode().IsDir() {
				name = filepath.Join(dst, hdr.Name)
				root = dst
			} else {
				// dst exists and user specified file name as dst
				// eg push /foo/bar/file.txt /baz/some_dir/file.txt
				name = dst
				root = filepath.Dir(dst)
			}
			break
		}
		if os.IsNotExist(ferr) {
			// dst does not exist push to parent dir and rename to dst.
			root = filepath.Dir(dst)
			if ferr = os.MkdirAll(root, modeDir); ferr == nil {
				name = dst
				break
			}
		}
		return ferr
	case tar.TypeDir:
		if ferr == nil {
			if fi.Mode().IsDir() {
				// dst exists and is a directory. Create dir in directory
				name = filepath.Join(dst, hdr.Name)
				break
			}
			return fmt.Errorf("file %s exists and is not a directory", dst)
		} else if os.IsNotExist(ferr) {
			root = dst
			name = dst
			rootReplace = hdr.Name
			break
		}
		return ferr
	}

	for {
		switch hdr.Typeflag {
		case tar.TypeRegA:
			fallthrough
		case tar.TypeReg:
			os.RemoveAll(name)
			f, err := os.OpenFile(name, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644)
			if err != nil {
				return err
			}
			// Close file ASAP to avoid too many open files
			if _, err = io.Copy(f, tr); err != nil {
				return err
			}
			if err := f.Close(); err != nil {
				return err
			}
		case tar.TypeDir:
			if err := os.MkdirAll(name, modeDir); err != nil {
				return err
			}
		case tar.TypeSymlink:
			if err := os.Symlink(hdr.Linkname, name); err != nil {
				return err
			}
		}

		hdr, err = tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		name = filepath.Join(root, strings.TrimPrefix(hdr.Name, rootReplace))
	}
}

// UntarBytes untars a tar r to a dst writer. This assumes that the tar contains only a single
// file.
func UntarBytes(dst io.Writer, r io.Reader) error {
	tr := tar.NewReader(r)
	// Advance to the first and only file in the tar. If there are more files present in the
	// tar, return an error.
	hdr, err := tr.Next()
	if err != nil {
		return err
	}
	switch hdr.Typeflag {
	case tar.TypeReg:
		fallthrough
	case tar.TypeRegA:
		if _, err = io.Copy(dst, tr); err != nil {
			return err
		}
	default:
		return fmt.Errorf("trying to read a directory into writer")
	}
	if _, eof := tr.Next(); eof != io.EOF {
		return fmt.Errorf("too many files received when trying to untar file")
	}
	return nil
}

// Tar archive tars the contents of src into a writer stream
func Tar(writer io.Writer, src string) error {
	if _, err := os.Stat(src); err != nil {
		return fmt.Errorf("Unable to tar files - %v", err.Error())
	}

	// avoid making the directory relative to itself
	src = strings.TrimRight(src, "/")

	tw := tar.NewWriter(writer)
	defer tw.Close()

	return filepath.Walk(src, func(file string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		var hdr *tar.Header
		if fi.Mode()&os.ModeSymlink != 0 {
			t, err := filepath.EvalSymlinks(file)
			if err != nil {
				// Don't include broken symlinks
				log.Printf("Skipping broken link %s\n", file)
				return nil
			}

			// ok to ignore error file always lives inside src
			sr, _ := filepath.Rel(src, file)
			name := filepath.Join(filepath.Base(src), sr)
			if strings.HasPrefix(t, src) {
				// Pointee lives under same root as pointer, so write the symlink.
				if hdr, err = tar.FileInfoHeader(fi, t); err != nil {
					return err
				}
				hdr.Name = name
			} else {
				// Pointee lives outside pointer root, write the file pointed to.
				var ti os.FileInfo
				if ti, err = os.Stat(t); err != nil {
					// Skip broken links. Error out on everything else.
					if os.IsNotExist(err) {
						log.Printf("Skipping broken link %s -> %s\n", file, t)
						return nil
					}
					return err
				}
				if !ti.Mode().IsRegular() {
					return nil
				}
				if hdr, err = tar.FileInfoHeader(ti, ""); err != nil {
					return err
				}
				hdr.Name = name
			}
		} else {
			hdr, err = tar.FileInfoHeader(fi, fi.Name())
			if err != nil {
				return err
			}

			rel, err := filepath.Rel(filepath.Dir(src), file)
			if err != nil {
				return err
			}
			hdr.Name = rel
		}

		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}

		if (hdr.Typeflag != tar.TypeReg && hdr.Typeflag != tar.TypeRegA) || hdr.Size == 0 {
			return nil
		}

		f, err := os.Open(file)
		if err != nil {
			return err
		}
		defer f.Close()

		if _, err := io.Copy(tw, f); err != nil {
			return err
		}

		return nil
	})
}

// TarBytes tars bytes as a single logical file into a tar file writer.
func TarBytes(writer io.Writer, bytes []byte) error {
	tw := tar.NewWriter(writer)
	defer tw.Close()

	hdr := &tar.Header{
		Typeflag: tar.TypeReg,
		Name:     "bytefile",
		Size:     int64(len(bytes)),
	}
	if err := tw.WriteHeader(hdr); err != nil {
		return err
	}

	if _, err := tw.Write(bytes); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package evaluator

import "monkey/ast"
import "monkey/object"

var (
	TRUE  = &object.Boolean{Value: true}
	FALSE = &object.Boolean{Value: false}
)

func Eval(node ast.Node) object.Object {
	switch node := node.(type) {
	// Statements
	case *ast.Program:
		return evalStatements(node.Statements)
	case *ast.ExpressionStatement:
		return Eval(node.Expression)
	// Expressions
	case *ast.IntegerLiteral:
		return &object.Integer{Value: node.Value}
	case *ast.Boolean:
		return nativeBoolToBooleanObject(node.Value)
	}

	return nil
}

func evalStatements(stmts []ast.Statement) object.Object {
	var result object.Object

	for _, statement := range stmts {
		result = Eval(statement)
	}

	return result
}

func nativeBoolToBooleanObject(input bool) *object.Boolean {
	if input {
		return TRUE
	}
	return FALSE
}
<commit_msg>added single null object to reference<commit_after>package evaluator

import "monkey/ast"
import "monkey/object"

var (
	NULL  = &object.Null{}
	TRUE  = &object.Boolean{Value: true}
	FALSE = &object.Boolean{Value: false}
)

func Eval(node ast.Node) object.Object {
	switch node := node.(type) {
	// Statements
	case *ast.Program:
		return evalStatements(node.Statements)
	case *ast.ExpressionStatement:
		return Eval(node.Expression)
	// Expressions
	case *ast.IntegerLiteral:
		return &object.Integer{Value: node.Value}
	case *ast.Boolean:
		return nativeBoolToBooleanObject(node.Value)
	}

	return nil
}

func evalStatements(stmts []ast.Statement) object.Object {
	var result object.Object

	for _, statement := range stmts {
		result = Eval(statement)
	}

	return result
}

func nativeBoolToBooleanObject(input bool) *object.Boolean {
	if input {
		return TRUE
	}
	return FALSE
}
<|endoftext|>