<commit_before>/*
Copyright 2016 The Kubernetes Authors 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 getter

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"strings"

	"k8s.io/helm/pkg/tlsutil"
	"k8s.io/helm/pkg/urlutil"
	"k8s.io/helm/pkg/version"
)

//httpGetter is the efault HTTP(/S) backend handler
type httpGetter struct {
	client *http.Client
}

//Get performs a Get from repo.Getter and returns the body.
func (g *httpGetter) Get(href string) (*bytes.Buffer, error) {
	buf := bytes.NewBuffer(nil)

	// Set a helm specific user agent so that a repo server and metrics can
	// separate helm calls from other tools interacting with repos.
	req, err := http.NewRequest("GET", href, nil)
	if err != nil {
		return buf, err
	}
	req.Header.Set("User-Agent", "Helm/"+strings.TrimPrefix(version.GetVersion(), "v"))

	resp, err := g.client.Do(req)
	if err != nil {
		return buf, err
	}
	if resp.StatusCode != 200 {
		return buf, fmt.Errorf("Failed to fetch %s : %s", href, resp.Status)
	}

	_, err = io.Copy(buf, resp.Body)
	resp.Body.Close()
	return buf, err
}

// newHTTPGetter constructs a valid http/https client as Getter
func newHTTPGetter(URL, CertFile, KeyFile, CAFile string) (Getter, error) {
	var client httpGetter
	if CertFile != "" && KeyFile != "" {
		tlsConf, err := tlsutil.NewClientTLS(CertFile, KeyFile, CAFile)
		if err != nil {
			return nil, fmt.Errorf("can't create TLS config for client: %s", err.Error())
		}
		tlsConf.BuildNameToCertificate()

		sni, err := urlutil.ExtractHostname(URL)
		if err != nil {
			return nil, err
		}
		tlsConf.ServerName = sni

		client.client = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: tlsConf,
			},
		}
	} else {
		client.client = http.DefaultClient
	}
	return &client, nil
}
<commit_msg>Respect env proxy settings when using custom Transport<commit_after>/*
Copyright 2016 The Kubernetes Authors 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 getter

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"strings"

	"k8s.io/helm/pkg/tlsutil"
	"k8s.io/helm/pkg/urlutil"
	"k8s.io/helm/pkg/version"
)

//httpGetter is the efault HTTP(/S) backend handler
type httpGetter struct {
	client *http.Client
}

//Get performs a Get from repo.Getter and returns the body.
func (g *httpGetter) Get(href string) (*bytes.Buffer, error) {
	buf := bytes.NewBuffer(nil)

	// Set a helm specific user agent so that a repo server and metrics can
	// separate helm calls from other tools interacting with repos.
	req, err := http.NewRequest("GET", href, nil)
	if err != nil {
		return buf, err
	}
	req.Header.Set("User-Agent", "Helm/"+strings.TrimPrefix(version.GetVersion(), "v"))

	resp, err := g.client.Do(req)
	if err != nil {
		return buf, err
	}
	if resp.StatusCode != 200 {
		return buf, fmt.Errorf("Failed to fetch %s : %s", href, resp.Status)
	}

	_, err = io.Copy(buf, resp.Body)
	resp.Body.Close()
	return buf, err
}

// newHTTPGetter constructs a valid http/https client as Getter
func newHTTPGetter(URL, CertFile, KeyFile, CAFile string) (Getter, error) {
	var client httpGetter
	if CertFile != "" && KeyFile != "" {
		tlsConf, err := tlsutil.NewClientTLS(CertFile, KeyFile, CAFile)
		if err != nil {
			return nil, fmt.Errorf("can't create TLS config for client: %s", err.Error())
		}
		tlsConf.BuildNameToCertificate()

		sni, err := urlutil.ExtractHostname(URL)
		if err != nil {
			return nil, err
		}
		tlsConf.ServerName = sni

		client.client = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: tlsConf,
				Proxy:           http.ProxyFromEnvironment,
			},
		}
	} else {
		client.client = http.DefaultClient
	}
	return &client, nil
}
<|endoftext|>


<commit_before>package gubled

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

	"fmt"
	"github.com/smancke/guble/client"
	"github.com/smancke/guble/protocol"
	"io/ioutil"
	"log"
	"os"
	"time"
)

type testgroup struct {
	t                *testing.T
	groupId          int
	addr             string
	done             chan bool
	messagesToSend   int
	client1, client2 client.Client
	topic            string
}

func newTestgroup(t *testing.T, groupId int, addr string, messagesToSend int) *testgroup {
	return &testgroup{
		t:              t,
		groupId:        groupId,
		addr:           addr,
		done:           make(chan bool),
		messagesToSend: messagesToSend,
	}
}

func TestThroughput(t *testing.T) {
	//defer enableDebugForMethod()()
	dir, _ := ioutil.TempDir("", "guble_benchmark_test")
	defer os.RemoveAll(dir)

	service := StartService(Args{Listen: "localhost:0", KVBackend: "memory", MSBackend: "file", StoragePath: dir})
	defer func() {
		service.Stop()
	}()
	time.Sleep(time.Millisecond * 10)

	testgroupCount := 4
	messagesPerGroup := 100
	log.Printf("init the %v testgroups", testgroupCount)
	testgroups := make([]*testgroup, testgroupCount, testgroupCount)
	for i := range testgroups {
		testgroups[i] = newTestgroup(t, i, service.WebServer().GetAddr(), messagesPerGroup)
	}

	// init test
	log.Print("init the testgroups")
	for i := range testgroups {
		testgroups[i].Init()
	}

	defer func() {
		// cleanup tests
		log.Print("cleanup the testgroups")
		for i := range testgroups {
			testgroups[i].Clean()
		}
	}()

	// start test
	log.Print("start the testgroups")
	start := time.Now()
	for i := range testgroups {
		go testgroups[i].Start()
	}

	log.Print("wait for finishing")
	timeout := time.After(time.Second * 60)
	for i, test := range testgroups {
		//fmt.Printf("wating for test %v\n", i)
		select {
		case successFlag := <-test.done:
			if !successFlag {
				t.Logf("testgroup %v returned with error", i)
				t.FailNow()
				return
			}
		case <-timeout:
			t.Log("timeout. testgroups not ready before timeout")
			t.Fail()
			return
		}
	}

	end := time.Now()
	totalMessages := testgroupCount * messagesPerGroup
	throughput := float64(totalMessages) / end.Sub(start).Seconds()
	log.Printf("finished! Throughput: %v/sec (%v message in %v)", int(throughput), totalMessages, end.Sub(start))
}

func (test *testgroup) Init() {
	test.topic = fmt.Sprintf("/%v-foo", test.groupId)
	var err error
	location := "ws://" + test.addr + "/stream/user/xy"
	//location := "ws://gathermon.mancke.net:8080/stream/"
	//location := "ws://127.0.0.1:8080/stream/"
	test.client1, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}
	test.client2, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}

	test.expectStatusMessage(protocol.SUCCESS_CONNECTED, "You are connected to the server.")

	test.client1.Subscribe(test.topic)
	time.Sleep(time.Millisecond * 1)
	//test.expectStatusMessage(protocol.SUCCESS_SUBSCRIBED_TO, test.topic)
}

func (test *testgroup) expectStatusMessage(name string, arg string) {
	select {
	case notify := <-test.client1.StatusMessages():
		assert.Equal(test.t, name, notify.Name)
		assert.Equal(test.t, arg, notify.Arg)
	case <-time.After(time.Second * 1):
		test.t.Logf("[%v] no notification of type %s after 1 second", test.groupId, name)
		test.done <- false
		test.t.Fail()
		return
	}
}

func (test *testgroup) Start() {
	go func() {
		for i := 0; i < test.messagesToSend; i++ {
			body := fmt.Sprintf("Hallo-%v", i)
			test.client2.Send(test.topic, body, "")
		}
	}()

	for i := 0; i < test.messagesToSend; i++ {
		body := fmt.Sprintf("Hallo-%v", i)

		select {
		case msg := <-test.client1.Messages():
			assert.Equal(test.t, body, msg.BodyAsString())
			assert.Equal(test.t, test.topic, string(msg.Path))
		case msg := <-test.client1.Errors():
			test.t.Logf("[%v] received error: %v", test.groupId, msg)
			test.done <- false
			test.t.Fail()
			return
		case <-time.After(time.Second * 5):
			test.t.Logf("[%v] no message received for 5 seconds, expected message %v", test.groupId, i)
			test.done <- false
			test.t.Fail()
			return
		}
	}
	test.done <- true
}

func (test *testgroup) Clean() {
	test.client1.Close()
	test.client2.Close()
}
<commit_msg>Linting<commit_after>package gubled

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

	"fmt"
	"github.com/smancke/guble/client"
	"github.com/smancke/guble/protocol"
	"io/ioutil"
	"log"
	"os"
	"time"
)

type testgroup struct {
	t                *testing.T
	groupID          int
	addr             string
	done             chan bool
	messagesToSend   int
	client1, client2 client.Client
	topic            string
}

func newTestgroup(t *testing.T, groupID int, addr string, messagesToSend int) *testgroup {
	return &testgroup{
		t:              t,
		groupID:        groupID,
		addr:           addr,
		done:           make(chan bool),
		messagesToSend: messagesToSend,
	}
}

func TestThroughput(t *testing.T) {
	//defer enableDebugForMethod()()
	dir, _ := ioutil.TempDir("", "guble_benchmark_test")
	defer os.RemoveAll(dir)

	service := StartService(Args{Listen: "localhost:0", KVBackend: "memory", MSBackend: "file", StoragePath: dir})
	defer func() {
		service.Stop()
	}()
	time.Sleep(time.Millisecond * 10)

	testgroupCount := 4
	messagesPerGroup := 100
	log.Printf("init the %v testgroups", testgroupCount)
	testgroups := make([]*testgroup, testgroupCount, testgroupCount)
	for i := range testgroups {
		testgroups[i] = newTestgroup(t, i, service.WebServer().GetAddr(), messagesPerGroup)
	}

	// init test
	log.Print("init the testgroups")
	for i := range testgroups {
		testgroups[i].Init()
	}

	defer func() {
		// cleanup tests
		log.Print("cleanup the testgroups")
		for i := range testgroups {
			testgroups[i].Clean()
		}
	}()

	// start test
	log.Print("start the testgroups")
	start := time.Now()
	for i := range testgroups {
		go testgroups[i].Start()
	}

	log.Print("wait for finishing")
	timeout := time.After(time.Second * 60)
	for i, test := range testgroups {
		//fmt.Printf("wating for test %v\n", i)
		select {
		case successFlag := <-test.done:
			if !successFlag {
				t.Logf("testgroup %v returned with error", i)
				t.FailNow()
				return
			}
		case <-timeout:
			t.Log("timeout. testgroups not ready before timeout")
			t.Fail()
			return
		}
	}

	end := time.Now()
	totalMessages := testgroupCount * messagesPerGroup
	throughput := float64(totalMessages) / end.Sub(start).Seconds()
	log.Printf("finished! Throughput: %v/sec (%v message in %v)", int(throughput), totalMessages, end.Sub(start))
}

func (test *testgroup) Init() {
	test.topic = fmt.Sprintf("/%v-foo", test.groupID)
	var err error
	location := "ws://" + test.addr + "/stream/user/xy"
	//location := "ws://gathermon.mancke.net:8080/stream/"
	//location := "ws://127.0.0.1:8080/stream/"
	test.client1, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}
	test.client2, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}

	test.expectStatusMessage(protocol.SUCCESS_CONNECTED, "You are connected to the server.")

	test.client1.Subscribe(test.topic)
	time.Sleep(time.Millisecond * 1)
	//test.expectStatusMessage(protocol.SUCCESS_SUBSCRIBED_TO, test.topic)
}

func (test *testgroup) expectStatusMessage(name string, arg string) {
	select {
	case notify := <-test.client1.StatusMessages():
		assert.Equal(test.t, name, notify.Name)
		assert.Equal(test.t, arg, notify.Arg)
	case <-time.After(time.Second * 1):
		test.t.Logf("[%v] no notification of type %s after 1 second", test.groupID, name)
		test.done <- false
		test.t.Fail()
		return
	}
}

func (test *testgroup) Start() {
	go func() {
		for i := 0; i < test.messagesToSend; i++ {
			body := fmt.Sprintf("Hallo-%v", i)
			test.client2.Send(test.topic, body, "")
		}
	}()

	for i := 0; i < test.messagesToSend; i++ {
		body := fmt.Sprintf("Hallo-%v", i)

		select {
		case msg := <-test.client1.Messages():
			assert.Equal(test.t, body, msg.BodyAsString())
			assert.Equal(test.t, test.topic, string(msg.Path))
		case msg := <-test.client1.Errors():
			test.t.Logf("[%v] received error: %v", test.groupID, msg)
			test.done <- false
			test.t.Fail()
			return
		case <-time.After(time.Second * 5):
			test.t.Logf("[%v] no message received for 5 seconds, expected message %v", test.groupID, i)
			test.done <- false
			test.t.Fail()
			return
		}
	}
	test.done <- true
}

func (test *testgroup) Clean() {
	test.client1.Close()
	test.client2.Close()
}
<|endoftext|>


<commit_before>// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package playground registers an HTTP handler at "/compile" that
// proxies requests to the golang.org playground service.
package playground

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
)

const runURL = "http://golang.org/compile"

func init() {
	http.HandleFunc("/compile", compile)
}

func compile(w http.ResponseWriter, r *http.Request) {
	b := new(bytes.Buffer)
	if err := passThru(b, r); err != nil {
		http.Error(w, "Compile server error.", http.StatusInternalServerError)
		report(r, err)
		return
	}
	io.Copy(w, b)
}

func passThru(w io.Writer, req *http.Request) error {
	defer req.Body.Close()
	r, err := client(req).Post(runURL, req.Header.Get("Content-type"), req.Body)
	if err != nil {
		return fmt.Errorf("making POST request: %v", err)
	}
	defer r.Body.Close()
	if _, err := io.Copy(w, r.Body); err != nil {
		return fmt.Errorf("copying response Body: %v", err)
	}
	return nil
}
<commit_msg>go.talks/pkg/playground: proxy share handler too<commit_after>// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package playground registers an HTTP handler at "/compile" that
// proxies requests to the golang.org playground service.
package playground

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
)

const baseURL = "http://play.golang.org"

func init() {
	http.HandleFunc("/compile", bounce)
	http.HandleFunc("/share", bounce)
}

func bounce(w http.ResponseWriter, r *http.Request) {
	b := new(bytes.Buffer)
	if err := passThru(b, r); err != nil {
		http.Error(w, "Server error.", http.StatusInternalServerError)
		report(r, err)
		return
	}
	io.Copy(w, b)
}

func passThru(w io.Writer, req *http.Request) error {
	defer req.Body.Close()
	url := baseURL + req.URL.Path
	r, err := client(req).Post(url, req.Header.Get("Content-type"), req.Body)
	if err != nil {
		return fmt.Errorf("making POST request: %v", err)
	}
	defer r.Body.Close()
	if _, err := io.Copy(w, r.Body); err != nil {
		return fmt.Errorf("copying response Body: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package settings

import "github.com/drone/drone/Godeps/_workspace/src/github.com/BurntSushi/toml"

// Service represents the configuration details required
// to connect to the revision control system (ie GitHub, Bitbucket)
type Service struct {
	// Name defines the name of the plugin. Possible values
	// may be github, gitlab, bitbucket, or gogs.
	Name string `toml:"name"`

	// Address defines the address (uri) of the plugin for
	// communication via the net/rpc package.
	Address string `toml:"address"`

	// Base defines the base URL for the service. For example:
	// https://github.com
	// https://bitbucket.org
	// https://gitlab.drone.io
	Base string `toml:"base"`

	// Indicates registration is open. If true any user
	// will be able to setup an account. If false, the
	// system administrator will need to provision accounts.
	Open bool `toml:"open"`

	// Orgs defines a list of organizations the user
	// must belong to in order to register. This will
	// take precedence over the `Open` paramter.
	Orgs []string `toml:"orgs"`

	// PrivateMode should be set to true if the
	// remote system requires authentication for
	// cloning public (open source) repositories.
	PrivateMode bool `toml:"private_mode"`

	// SkipVerify instructs the client to skip SSL verification.
	// This may be used with self-signed certificates, however,
	// is not recommended for security reasons.
	SkipVerify bool `toml:"skip_verify"`

	// OAuth configuration data. If nil or empty, Drone may
	// assume basic authentication via username and password.
	OAuth *OAuth `toml:"oauth"`
}

// OAuth defines how a user should autheticate with the service.
// This supports OAuth2 and OAuth1 protocols.
type OAuth struct {
	Client       string   `toml:"client"`
	Secret       string   `toml:"secret"`
	Authorize    string   `toml:"authorize"`
	AccessToken  string   `toml:"access_token"`
	RequestToken string   `toml:"request_token"`
	Scope        []string `toml:"scope"`
}

// Server represents the web server configuration details
// used to server HTTP requests.
type Server struct {
	Base string `toml:"base"`
	Addr string `toml:"addr"`
	Cert string `toml:"cert"`
	Key  string `toml:"key"`

	Scheme   string `toml:"scheme"`
	Hostname string `toml:"hostname"`
}

// Session represents the session configuration details
// used to generate, validate and expire authentication
// sessions.
type Session struct {
	Secret  string `toml:"secret"`
	Expires int64  `toml:"expires"`
}

// Docker represents the configuration details used
// to connect to the Docker daemon when scheduling
// and executing builds in containers.
type Docker struct {
	Cert  string   `toml:"cert"`
	Key   string   `toml:"key"`
	Nodes []string `toml:"nodes"`
}

// Database represents the configuration details used
// to connect to the embedded Bolt database.
type Database struct {
	Driver     string `toml:"driver"`
	Datasource string `toml:"datasource"`
}

type Agents struct {
	Secret string `toml:"secret"`
}

// Settings defines global settings for the Drone system.
type Settings struct {
	Database *Database `toml:"database"`
	Docker   *Docker   `toml:"docker"`
	Service  *Service  `toml:"service"`
	Server   *Server   `toml:"server"`
	Session  *Session  `toml:"session"`
	Agents   *Agents   `toml:"agents"`

	Plugins map[string]interface{} `toml:"plugins"`
}

// Parse parses the Drone settings file at the specified path
// and unmarshals to a Settings structure.
func Parse(path string) (*Settings, error) {
	s := &Settings{}
	_, err := toml.DecodeFile(path, s)
	return applyDefaults(s), err
}

// ParseString parses the Drone settings string and unmarshals
// to a Settings structure.
func ParseString(data string) (*Settings, error) {
	s := &Settings{}
	_, err := toml.Decode(data, s)
	return applyDefaults(s), err
}

func applyDefaults(s *Settings) *Settings {
	if s.Session == nil {
		s.Session = &Session{}
	}
	// if no session token is provided we can
	// instead use the client secret to sign
	// our sessions and tokens.
	if len(s.Session.Secret) == 0 {
		s.Session.Secret = s.Service.OAuth.Secret
	}
	return s
}
<commit_msg>Remove unused pkg/settings<commit_after><|endoftext|>


<commit_before>package helpers

import (
	"fmt"
	"testing"

	. "github.com/onsi/ginkgo"
	ginkgoconfig "github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/reporters"
	. "github.com/onsi/gomega"
	"github.com/vito/cmdtest"
	. "github.com/vito/cmdtest/matchers"

	"github.com/pivotal-cf-experimental/cf-test-helpers/cf"
)

func setUpSpaceWithUserAccess(uc cf.UserContext, sname string) {
	Expect(cf.Cf("create-space", "-o", uc.Org, sname)).To(ExitWith(0))
	Expect(cf.Cf("set-space-role", uc.Username, uc.Org, sname, "SpaceManager")).To(ExitWith(0))
	Expect(cf.Cf("set-space-role", uc.Username, uc.Org, sname, "SpaceDeveloper")).To(ExitWith(0))
	Expect(cf.Cf("set-space-role", uc.Username, uc.Org, sname, "SpaceAuditor")).To(ExitWith(0))
}

var RegularUserContext cf.UserContext

func GinkgoBootstrap(t *testing.T, suiteName string) {
	RegisterFailHandler(Fail)

	adminContext := AdminUserContext
	defer func() {
		cf.Cf("delete-user", "-f", RegularUserContext.Username)
		cf.Cf("delete-org", "-f", RegularUserContext.Org)
	}()

	var originalCfHomeDir, currentCfHomeDir string

	BeforeEach(func() {
		RegularUserContext = NewRegularUserContext()
		cf.AsUser(adminContext, func() {
			Expect(cf.Cf("create-user", RegularUserContext.Username, RegularUserContext.Password)).To(SayBranches(
				cmdtest.ExpectBranch{"OK", func() {}},
				cmdtest.ExpectBranch{"scim_resource_already_exists", func() {}},
			))
			Expect(cf.Cf("create-org", RegularUserContext.Org)).To(ExitWith(0))

			setUpSpaceWithUserAccess(RegularUserContext, RegularUserContext.Space)
			setUpSpaceWithUserAccess(RegularUserContext, "persistent-space")
		})

		originalCfHomeDir, currentCfHomeDir = cf.InitiateUserContext(RegularUserContext)
		cf.TargetSpace(RegularUserContext)
	})

	AfterEach(func() {
		cf.RestoreUserContext(RegularUserContext, originalCfHomeDir, currentCfHomeDir)
	})

	outputFile := fmt.Sprintf("../results/%s-junit_%d.xml", suiteName, ginkgoconfig.GinkgoConfig.ParallelNode)
	RunSpecsWithDefaultAndCustomReporters(t, suiteName, []Reporter{reporters.NewJUnitReporter(outputFile)})
}
<commit_msg>Delete space after run, not org<commit_after>package helpers

import (
	"fmt"
	"testing"

	. "github.com/onsi/ginkgo"
	ginkgoconfig "github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/reporters"
	. "github.com/onsi/gomega"
	"github.com/vito/cmdtest"
	. "github.com/vito/cmdtest/matchers"

	"github.com/pivotal-cf-experimental/cf-test-helpers/cf"
)

func setUpSpaceWithUserAccess(uc cf.UserContext, sname string) {
	Expect(cf.Cf("create-space", "-o", uc.Org, sname)).To(ExitWith(0))
	Expect(cf.Cf("set-space-role", uc.Username, uc.Org, sname, "SpaceManager")).To(ExitWith(0))
	Expect(cf.Cf("set-space-role", uc.Username, uc.Org, sname, "SpaceDeveloper")).To(ExitWith(0))
	Expect(cf.Cf("set-space-role", uc.Username, uc.Org, sname, "SpaceAuditor")).To(ExitWith(0))
}

var RegularUserContext cf.UserContext

func GinkgoBootstrap(t *testing.T, suiteName string) {
	RegisterFailHandler(Fail)

	adminContext := AdminUserContext
	defer func() {
		cf.Cf("delete-user", "-f", RegularUserContext.Username)
		cf.Cf("delete-space", "-f", RegularUserContext.Space)
	}()

	var originalCfHomeDir, currentCfHomeDir string

	BeforeEach(func() {
		RegularUserContext = NewRegularUserContext()
		cf.AsUser(adminContext, func() {
			Expect(cf.Cf("create-user", RegularUserContext.Username, RegularUserContext.Password)).To(SayBranches(
				cmdtest.ExpectBranch{"OK", func() {}},
				cmdtest.ExpectBranch{"scim_resource_already_exists", func() {}},
			))
			Expect(cf.Cf("create-org", RegularUserContext.Org)).To(ExitWith(0))

			setUpSpaceWithUserAccess(RegularUserContext, RegularUserContext.Space)
			setUpSpaceWithUserAccess(RegularUserContext, "persistent-space")
		})

		originalCfHomeDir, currentCfHomeDir = cf.InitiateUserContext(RegularUserContext)
		cf.TargetSpace(RegularUserContext)
	})

	AfterEach(func() {
		cf.RestoreUserContext(RegularUserContext, originalCfHomeDir, currentCfHomeDir)
	})

	outputFile := fmt.Sprintf("../results/%s-junit_%d.xml", suiteName, ginkgoconfig.GinkgoConfig.ParallelNode)
	RunSpecsWithDefaultAndCustomReporters(t, suiteName, []Reporter{reporters.NewJUnitReporter(outputFile)})
}
<|endoftext|>


<commit_before>package amqp

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey, // queue binding key
		nil, // exchange declare args
		nil, // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, err
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occured on queue: %s. Reconnecting", queueName)
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	connection, err := b.GetOrOpenConnection(signature.RoutingKey,
		b.GetConfig().AMQP.BindingKey, // queue binding key
		nil, // exchange declare args
		nil, // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	errorsChan := make(chan error)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		delivery.Nack(multiple, requeue)
		return nil
	}

	log.INFO.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	connection, err := b.GetOrOpenConnection(queueName,
		queueName,                                       // queue binding key
		nil,                                             // exchange declare args
		declareQueueArgs,                                // queue declare arghs
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct" {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<commit_msg>Fixing connections when is a delay queue<commit_after>package amqp

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey,   // queue binding key
		nil,                             // exchange declare args
		nil,                             // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, err
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occured on queue: %s. Reconnecting", queueName)
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	connection, err := b.GetOrOpenConnection(signature.RoutingKey,
		b.GetConfig().AMQP.BindingKey, // queue binding key
		nil,                           // exchange declare args
		nil,                           // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	errorsChan := make(chan error)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		delivery.Nack(multiple, requeue)
		return nil
	}

	log.INFO.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	conn, channel, _, _, _, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		queueName,                       // queue binding key
		nil,                             // exchange declare args
		declareQueueArgs,                // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	defer b.Close(channel, conn)

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct" {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	_ "image/jpeg"
	"image/png"
	"io/ioutil"
	"math"
	"os"

	"github.com/animenotifier/arn"
)

func main() {
	data, err := ioutil.ReadFile("input.jpg")

	if err != nil {
		panic(err)
	}

	img, format, err := image.Decode(bytes.NewReader(data))

	if err != nil {
		panic(err)
	}

	fmt.Println(img.Bounds().Dx(), img.Bounds().Dy(), format)
	improved := ImproveQuality(img)
	f, err := os.Create("output.png")

	if err != nil {
		panic(err)
	}

	defer f.Close()
	png.Encode(f, improved)
}

const max = float64(65535)

// Pixel ...
type Pixel struct {
	X     int
	Y     int
	Color arn.HSLColor
}

// Area ...
type Area struct {
	AverageColor color.Color
	Pixels       []Pixel
}

// Add ...
func (area *Area) Add(x, y int, hsl arn.HSLColor) {
	area.Pixels = append(area.Pixels, Pixel{
		X:     x,
		Y:     y,
		Color: hsl,
	})
}

const (
	hueTolerance        = 0.1
	lightnessTolerance  = 0.1
	saturationTolerance = 0.1
)

// ImproveQuality returns the average color of an image in HSL format.
func ImproveQuality(img image.Image) *image.NRGBA {
	width := img.Bounds().Dx()
	height := img.Bounds().Dy()
	clone := image.NewNRGBA(image.Rect(0, 0, width, height))
	hueAreas := []Area{}

	for x := 0; x < width; x++ {
		for y := 0; y < height; y++ {
			color := img.At(x, y)
			rUint, gUint, bUint, _ := color.RGBA()
			r := float64(rUint) / max
			g := float64(gUint) / max
			b := float64(bUint) / max
			h, s, l := arn.RGBToHSL(r, g, b)
			areaIndex := -1

			// Find similar area
			for i := 0; i < len(hueAreas); i++ {
				area := hueAreas[i]

				// Is the pixel close to any pixel in the area we're checking?
				for _, pixel := range area.Pixels {
					xDist := x - pixel.X
					yDist := y - pixel.Y

					if xDist < 0 {
						xDist = -xDist
					}

					if yDist < 0 {
						yDist = -yDist
					}

					if xDist <= 1 && yDist <= 1 {
						// Is the color similar?
						if math.Abs(h-pixel.Color.Hue) <= hueTolerance && math.Abs(s-pixel.Color.Saturation) <= saturationTolerance && math.Abs(l-pixel.Color.Lightness) <= lightnessTolerance {
							areaIndex = i
							break
						}
					}
				}

				if areaIndex != -1 {
					break
				}
			}

			// Insert new area
			if areaIndex == -1 {
				areaIndex = len(hueAreas)
				hueAreas = append(hueAreas, Area{})
			}

			hueAreas[areaIndex].Add(x, y, arn.HSLColor{
				Hue:        h,
				Saturation: s,
				Lightness:  l,
			})
		}
	}

	fmt.Println(len(hueAreas), "areas")

	// Build image from areas
	for _, area := range hueAreas {
		totalR := uint64(0)
		totalG := uint64(0)
		totalB := uint64(0)

		// Calculate area average color
		for _, pixel := range area.Pixels {
			col := img.At(pixel.X, pixel.Y)
			r, g, b, _ := col.RGBA()
			totalR += uint64(r)
			totalG += uint64(g)
			totalB += uint64(b)
		}

		averageR := float64(totalR/uint64(len(area.Pixels))) / max
		averageG := float64(totalG/uint64(len(area.Pixels))) / max
		averageB := float64(totalB/uint64(len(area.Pixels))) / max

		area.AverageColor = color.RGBA{
			R: uint8(averageR * 255),
			G: uint8(averageG * 255),
			B: uint8(averageB * 255),
			A: 255,
		}

		for _, pixel := range area.Pixels {
			clone.Set(pixel.X, pixel.Y, area.AverageColor)
		}
	}

	return clone
}
<commit_msg>Simplified area recognition<commit_after>package main

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	_ "image/jpeg"
	"image/png"
	"io/ioutil"
	"os"
)

func main() {
	data, err := ioutil.ReadFile("input.jpg")

	if err != nil {
		panic(err)
	}

	img, format, err := image.Decode(bytes.NewReader(data))

	if err != nil {
		panic(err)
	}

	fmt.Println(img.Bounds().Dx(), img.Bounds().Dy(), format)
	improved := ImproveQuality(img)
	f, err := os.Create("output.png")

	if err != nil {
		panic(err)
	}

	defer f.Close()
	png.Encode(f, improved)
}

const max = float64(65535)

// Pixel ...
type Pixel struct {
	X int
	Y int
}

// Area ...
type Area struct {
	color.Color
	Pixels []Pixel
	totalR uint64
	totalG uint64
	totalB uint64
	totalA uint64
}

// Add ...
func (area *Area) Add(x, y int, r, g, b, a uint32) {
	area.Pixels = append(area.Pixels, Pixel{
		X: x,
		Y: y,
	})

	area.totalR += uint64(r)
	area.totalG += uint64(g)
	area.totalB += uint64(b)
	area.totalA += uint64(a)
}

// AverageColor ...
func (area *Area) AverageColor() color.Color {
	return color.RGBA64{
		R: uint16(area.totalR / uint64(len(area.Pixels))),
		G: uint16(area.totalG / uint64(len(area.Pixels))),
		B: uint16(area.totalB / uint64(len(area.Pixels))),
		A: uint16(area.totalA / uint64(len(area.Pixels))),
	}
}

const (
	tolerance           = uint32(3000)
	hueTolerance        = 0.1
	lightnessTolerance  = 0.1
	saturationTolerance = 0.1
)

func diffAbs(a uint32, b uint32) uint32 {
	if a > b {
		return a - b
	}

	return b - a
}

// ImproveQuality returns the average color of an image in HSL format.
func ImproveQuality(img image.Image) *image.NRGBA {
	width := img.Bounds().Dx()
	height := img.Bounds().Dy()
	clone := image.NewNRGBA(image.Rect(0, 0, width, height))
	hueAreas := []Area{}

	for x := 0; x < width; x++ {
		for y := 0; y < height; y++ {
			color := img.At(x, y)
			r, g, b, a := color.RGBA()
			areaIndex := -1

			// Find similar area
			for i := 0; i < len(hueAreas); i++ {
				area := hueAreas[i]
				avgR, avgG, avgB, _ := area.AverageColor().RGBA()

				// Is the color similar?
				if diffAbs(r, avgR) <= tolerance && diffAbs(g, avgG) <= tolerance && diffAbs(b, avgB) <= tolerance {
					areaIndex = i
					break
				}
			}

			// Insert new area
			if areaIndex == -1 {
				areaIndex = len(hueAreas)
				hueAreas = append(hueAreas, Area{})
			}

			hueAreas[areaIndex].Add(x, y, r, g, b, a)
		}
	}

	fmt.Println(len(hueAreas), "areas")

	// Build image from areas
	for _, area := range hueAreas {
		avgColor := area.AverageColor()

		for _, pixel := range area.Pixels {
			clone.Set(pixel.X, pixel.Y, avgColor)
		}
	}

	return clone
}
<|endoftext|>


<commit_before>package vmware

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"strings"
)

// Workstation9Driver is a driver that can run VMware Workstation 9
// on non-Windows platforms.
type Workstation9Driver struct {
	AppPath          string
	VdiskManagerPath string
	VmrunPath        string
}

func (d *Workstation9Driver) CompactDisk(diskPath string) error {
	defragCmd := exec.Command(d.VdiskManagerPath, "-d", diskPath)
	if _, _, err := d.runAndLog(defragCmd); err != nil {
		return err
	}

	shrinkCmd := exec.Command(d.VdiskManagerPath, "-k", diskPath)
	if _, _, err := d.runAndLog(shrinkCmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) CreateDisk(output string, size string) error {
	cmd := exec.Command(d.VdiskManagerPath, "-c", "-s", size, "-a", "lsilogic", "-t", "1", output)
	if _, _, err := d.runAndLog(cmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) IsRunning(vmxPath string) (bool, error) {
	vmxPath, err := filepath.Abs(vmxPath)
	if err != nil {
		return false, err
	}

	cmd := exec.Command(d.VmrunPath, "-T", "ws", "list")
	stdout, _, err := d.runAndLog(cmd)
	if err != nil {
		return false, err
	}

	for _, line := range strings.Split(stdout, "\n") {
		if line == vmxPath {
			return true, nil
		}
	}

	return false, nil
}

func (d *Workstation9Driver) Start(vmxPath string, headless bool) error {
	guiArgument := "gui"
	if headless {
		guiArgument = "nogui"
	}

	cmd := exec.Command(d.VmrunPath, "-T", "ws", "start", vmxPath, guiArgument)
	if _, _, err := d.runAndLog(cmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) Stop(vmxPath string) error {
	cmd := exec.Command(d.VmrunPath, "-T", "ws", "stop", vmxPath, "hard")
	if _, _, err := d.runAndLog(cmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) Verify() error {
	var err error
	if d.AppPath == "" {
		if d.AppPath, err = workstationFindVMware(); err != nil {
			return err
		}
	}

	if d.VmrunPath == "" {
		if d.VmrunPath, err = workstationFindVmrun(); err != nil {
			return err
		}
	}

	if d.VdiskManagerPath == "" {
		if d.VdiskManagerPath, err = workstationFindVdiskManager(); err != nil {
			return err
		}
	}

	// Check to see if it APPEARS to be licensed.
	if err := workstationCheckLicense(); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) ToolsIsoPath(flavor string) string {
	return workstationToolsIsoPath(flavor)
}

func (d *Workstation9Driver) DhcpLeasesPath(device string) string {
	return workstationDhcpLeasesPath(device)
}

func (d *Workstation9Driver) runAndLog(cmd *exec.Cmd) (string, string, error) {
	var stdout, stderr bytes.Buffer

	log.Printf("Executing: %s %v", cmd.Path, cmd.Args[1:])
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	stdoutString := strings.TrimSpace(stdout.String())
	stderrString := strings.TrimSpace(stderr.String())

	if _, ok := err.(*exec.ExitError); ok {
		err = fmt.Errorf("VMware error: %s", stderrString)
	}

	log.Printf("stdout: %s", stdoutString)
	log.Printf("stderr: %s", stderrString)

	return stdout.String(), stderr.String(), err
}
<commit_msg>builder/vmware: verify various paths exist for Workstation<commit_after>package vmware

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Workstation9Driver is a driver that can run VMware Workstation 9
// on non-Windows platforms.
type Workstation9Driver struct {
	AppPath          string
	VdiskManagerPath string
	VmrunPath        string
}

func (d *Workstation9Driver) CompactDisk(diskPath string) error {
	defragCmd := exec.Command(d.VdiskManagerPath, "-d", diskPath)
	if _, _, err := d.runAndLog(defragCmd); err != nil {
		return err
	}

	shrinkCmd := exec.Command(d.VdiskManagerPath, "-k", diskPath)
	if _, _, err := d.runAndLog(shrinkCmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) CreateDisk(output string, size string) error {
	cmd := exec.Command(d.VdiskManagerPath, "-c", "-s", size, "-a", "lsilogic", "-t", "1", output)
	if _, _, err := d.runAndLog(cmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) IsRunning(vmxPath string) (bool, error) {
	vmxPath, err := filepath.Abs(vmxPath)
	if err != nil {
		return false, err
	}

	cmd := exec.Command(d.VmrunPath, "-T", "ws", "list")
	stdout, _, err := d.runAndLog(cmd)
	if err != nil {
		return false, err
	}

	for _, line := range strings.Split(stdout, "\n") {
		if line == vmxPath {
			return true, nil
		}
	}

	return false, nil
}

func (d *Workstation9Driver) Start(vmxPath string, headless bool) error {
	guiArgument := "gui"
	if headless {
		guiArgument = "nogui"
	}

	cmd := exec.Command(d.VmrunPath, "-T", "ws", "start", vmxPath, guiArgument)
	if _, _, err := d.runAndLog(cmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) Stop(vmxPath string) error {
	cmd := exec.Command(d.VmrunPath, "-T", "ws", "stop", vmxPath, "hard")
	if _, _, err := d.runAndLog(cmd); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) Verify() error {
	var err error
	if d.AppPath == "" {
		if d.AppPath, err = workstationFindVMware(); err != nil {
			return err
		}
	}

	if d.VmrunPath == "" {
		if d.VmrunPath, err = workstationFindVmrun(); err != nil {
			return err
		}
	}

	if d.VdiskManagerPath == "" {
		if d.VdiskManagerPath, err = workstationFindVdiskManager(); err != nil {
			return err
		}
	}

	if _, err := os.Stat(d.AppPath); err != nil {
		return fmt.Errorf("VMware application not found: %s", d.AppPath)
	}

	if _, err := os.Stat(d.VmrunPath); err != nil {
		return fmt.Errorf("'vmrun' application not found: %s", d.VmrunPath)
	}

	if _, err := os.Stat(d.VdiskManagerPath); err != nil {
		return fmt.Errorf("'vmrun' application not found: %s", d.VdiskManagerPath)
	}

	// Check to see if it APPEARS to be licensed.
	if err := workstationCheckLicense(); err != nil {
		return err
	}

	return nil
}

func (d *Workstation9Driver) ToolsIsoPath(flavor string) string {
	return workstationToolsIsoPath(flavor)
}

func (d *Workstation9Driver) DhcpLeasesPath(device string) string {
	return workstationDhcpLeasesPath(device)
}

func (d *Workstation9Driver) runAndLog(cmd *exec.Cmd) (string, string, error) {
	var stdout, stderr bytes.Buffer

	log.Printf("Executing: %s %v", cmd.Path, cmd.Args[1:])
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	stdoutString := strings.TrimSpace(stdout.String())
	stderrString := strings.TrimSpace(stderr.String())

	if _, ok := err.(*exec.ExitError); ok {
		err = fmt.Errorf("VMware error: %s", stderrString)
	}

	log.Printf("stdout: %s", stdoutString)
	log.Printf("stderr: %s", stderrString)

	return stdout.String(), stderr.String(), err
}
<|endoftext|>


<commit_before>package graph

// Not that I'm falling prey to NIH syndrome, but implementing my own
// heap will speed up MST and Search algorithms because...
//
// Although asserting is relatively cheap, I also need update methods
// and can save time with an explicit update function.
// For reference, with one set of data:
//  - using containers/heap took 425 microseconds,
//  - switching to my own version of heap took 330 microseconds average,
//  - adding an update function takes 255 microseconds average.
// So this was worth it.

type nodeSlice []*node

func (n nodeSlice) less(i, j int) bool {
	return n[i].state < n[j].state
}
func (n nodeSlice) swap(i, j int) {
	n[j].data, n[i].data = n[i].data, n[j].data // swap data containing indices
	n[j], n[i] = n[i], n[j]
}

func (n nodeSlice) shuffleUp(index int) {
	for {
		parent := (index - 1) / 2
		if parent == index || n.less(parent, index) {
			break
		}
		n.swap(parent, index)
		index = parent
	}
}

func (n nodeSlice) shuffleDown(elem, to int) {
	for {
		minchild := elem*2 + 1 // left child: elem * 2 + 1
		if minchild >= to {
			return
		}
		rchild := minchild + 1 // left child + 1
		if rchild < to && n.less(rchild, minchild) {
			minchild = rchild
		}
		if !n.less(minchild, elem) {
			return
		}
		n.swap(minchild, elem)
		elem = minchild
	}
}

func (n nodeSlice) heapInit() {
	length := len(n)
	for i := length/2 - 1; i >= 0; i-- {
		n.shuffleDown(i, length)
	}
}

func (p *nodeSlice) remove(index int) *node {
	n := *p
	length := len(n) - 1
	if length != index {
		n.swap(length, index)
		n.shuffleDown(index, length)
		n.shuffleUp(index)
	}
	popped := n[length]
	popped.data = -1
	n = n[0:length]
	*p = n
	return popped
}

// I don't need shuffleDown call because all updates update to a
// smaller state (for now)
func (n nodeSlice) update(index, newState int) {
	n[index].state = newState
	n.shuffleUp(index)
}

func (p *nodeSlice) push(x *node) {
	n := *p
	x.data = len(n) // index into heap
	n = append(n, x)
	n.shuffleUp(len(n) - 1)
	*p = n
}

func (p *nodeSlice) pop() *node {
	return p.remove(0)
}
func (n nodeSlice) heapContains(node *node) bool { // extend heap interface
	return node.data > dequeued
}
<commit_msg>change integer to constant value dequeued<commit_after>package graph

// Not that I'm falling prey to NIH syndrome, but implementing my own
// heap will speed up MST and Search algorithms because...
//
// Although asserting is relatively cheap, I also need update methods
// and can save time with an explicit update function.
// For reference, with one set of data:
//  - using containers/heap took 425 microseconds,
//  - switching to my own version of heap took 330 microseconds average,
//  - adding an update function takes 255 microseconds average.
// So this was worth it.

type nodeSlice []*node

func (n nodeSlice) less(i, j int) bool {
	return n[i].state < n[j].state
}
func (n nodeSlice) swap(i, j int) {
	n[j].data, n[i].data = n[i].data, n[j].data // swap data containing indices
	n[j], n[i] = n[i], n[j]
}

func (n nodeSlice) shuffleUp(index int) {
	for {
		parent := (index - 1) / 2
		if parent == index || n.less(parent, index) {
			break
		}
		n.swap(parent, index)
		index = parent
	}
}

func (n nodeSlice) shuffleDown(elem, to int) {
	for {
		minchild := elem*2 + 1 // left child: elem * 2 + 1
		if minchild >= to {
			return
		}
		rchild := minchild + 1 // left child + 1
		if rchild < to && n.less(rchild, minchild) {
			minchild = rchild
		}
		if !n.less(minchild, elem) {
			return
		}
		n.swap(minchild, elem)
		elem = minchild
	}
}

func (n nodeSlice) heapInit() {
	length := len(n)
	for i := length/2 - 1; i >= 0; i-- {
		n.shuffleDown(i, length)
	}
}

func (p *nodeSlice) remove(index int) *node {
	n := *p
	length := len(n) - 1
	if length != index {
		n.swap(length, index)
		n.shuffleDown(index, length)
		n.shuffleUp(index)
	}
	popped := n[length]
	popped.data = dequeued
	n = n[0:length]
	*p = n
	return popped
}

// I don't need shuffleDown call because all updates update to a
// smaller state (for now)
func (n nodeSlice) update(index, newState int) {
	n[index].state = newState
	n.shuffleUp(index)
}

func (p *nodeSlice) push(x *node) {
	n := *p
	x.data = len(n) // index into heap
	n = append(n, x)
	n.shuffleUp(len(n) - 1)
	*p = n
}

func (p *nodeSlice) pop() *node {
	return p.remove(0)
}
func (n nodeSlice) heapContains(node *node) bool { // extend heap interface
	return node.data > dequeued
}
<|endoftext|>


<commit_before>// Package main is complete tool for the go command line
package main

import "github.com/posener/complete"

var (
	ellipsis   = complete.PredictSet("./...")
	anyPackage = complete.PredictFunc(predictPackages)
	goFiles    = complete.PredictFiles("*.go")
	anyFile    = complete.PredictFiles("*")
	anyGo      = complete.PredictOr(goFiles, anyPackage, ellipsis)
)

func main() {
	build := complete.Command{
		Flags: complete.Flags{
			"-o": anyFile,
			"-i": complete.PredictNothing,

			"-a":             complete.PredictNothing,
			"-n":             complete.PredictNothing,
			"-p":             complete.PredictAnything,
			"-race":          complete.PredictNothing,
			"-msan":          complete.PredictNothing,
			"-v":             complete.PredictNothing,
			"-work":          complete.PredictNothing,
			"-x":             complete.PredictNothing,
			"-asmflags":      complete.PredictAnything,
			"-buildmode":     complete.PredictAnything,
			"-compiler":      complete.PredictAnything,
			"-gccgoflags":    complete.PredictAnything,
			"-gcflags":       complete.PredictAnything,
			"-installsuffix": complete.PredictAnything,
			"-ldflags":       complete.PredictAnything,
			"-linkshared":    complete.PredictNothing,
			"-pkgdir":        anyPackage,
			"-tags":          complete.PredictAnything,
			"-toolexec":      complete.PredictAnything,
		},
		Args: anyGo,
	}

	run := complete.Command{
		Flags: complete.Flags{
			"-exec": complete.PredictAnything,
		},
		Args: goFiles,
	}

	test := complete.Command{
		Flags: complete.Flags{
			"-args": complete.PredictAnything,
			"-c":    complete.PredictNothing,
			"-exec": complete.PredictAnything,

			"-bench":     predictBenchmark,
			"-benchtime": complete.PredictAnything,
			"-count":     complete.PredictAnything,
			"-cover":     complete.PredictNothing,
			"-covermode": complete.PredictSet("set", "count", "atomic"),
			"-coverpkg":  complete.PredictDirs("*"),
			"-cpu":       complete.PredictAnything,
			"-run":       predictTest,
			"-short":     complete.PredictNothing,
			"-timeout":   complete.PredictAnything,

			"-benchmem":             complete.PredictNothing,
			"-blockprofile":         complete.PredictFiles("*.out"),
			"-blockprofilerate":     complete.PredictAnything,
			"-coverprofile":         complete.PredictFiles("*.out"),
			"-cpuprofile":           complete.PredictFiles("*.out"),
			"-memprofile":           complete.PredictFiles("*.out"),
			"-memprofilerate":       complete.PredictAnything,
			"-mutexprofile":         complete.PredictFiles("*.out"),
			"-mutexprofilefraction": complete.PredictAnything,
			"-outputdir":            complete.PredictDirs("*"),
			"-trace":                complete.PredictFiles("*.out"),
		},
		Args: anyGo,
	}

	fmt := complete.Command{
		Flags: complete.Flags{
			"-n": complete.PredictNothing,
			"-x": complete.PredictNothing,
		},
		Args: anyGo,
	}

	get := complete.Command{
		Flags: complete.Flags{
			"-d":        complete.PredictNothing,
			"-f":        complete.PredictNothing,
			"-fix":      complete.PredictNothing,
			"-insecure": complete.PredictNothing,
			"-t":        complete.PredictNothing,
			"-u":        complete.PredictNothing,
		},
		Args: anyGo,
	}

	generate := complete.Command{
		Flags: complete.Flags{
			"-n":   complete.PredictNothing,
			"-x":   complete.PredictNothing,
			"-v":   complete.PredictNothing,
			"-run": complete.PredictAnything,
		},
		Args: anyGo,
	}

	vet := complete.Command{
		Flags: complete.Flags{
			"-n": complete.PredictNothing,
			"-x": complete.PredictNothing,
		},
		Args: anyGo,
	}

	list := complete.Command{
		Flags: complete.Flags{
			"-e":    complete.PredictNothing,
			"-f":    complete.PredictAnything,
			"-json": complete.PredictNothing,
		},
		Args: complete.PredictOr(anyPackage, ellipsis),
	}

	tool := complete.Command{
		Flags: complete.Flags{
			"-n": complete.PredictNothing,
		},
		Args: complete.PredictAnything,
	}

	clean := complete.Command{
		Flags: complete.Flags{
			"-i": complete.PredictNothing,
			"-r": complete.PredictNothing,
			"-n": complete.PredictNothing,
			"-x": complete.PredictNothing,
		},
		Args: complete.PredictOr(anyPackage, ellipsis),
	}

	env := complete.Command{
		Args: complete.PredictAnything,
	}

	bug := complete.Command{}
	version := complete.Command{}

	fix := complete.Command{
		Args: anyGo,
	}

	// commands that also accepts the build flags
	for name, options := range build.Flags {
		test.Flags[name] = options
		run.Flags[name] = options
		list.Flags[name] = options
		vet.Flags[name] = options
		get.Flags[name] = options
	}

	gogo := complete.Command{
		Sub: complete.Commands{
			"build":    build,
			"install":  build, // install and build have the same flags
			"run":      run,
			"test":     test,
			"fmt":      fmt,
			"get":      get,
			"generate": generate,
			"vet":      vet,
			"list":     list,
			"tool":     tool,
			"clean":    clean,
			"env":      env,
			"bug":      bug,
			"fix":      fix,
			"version":  version,
		},
		GlobalFlags: complete.Flags{
			"-h": complete.PredictNothing,
		},
	}

	complete.New("go", gogo).Run()
}
<commit_msg>gocomplete: add go tool subcommands<commit_after>// Package main is complete tool for the go command line
package main

import "github.com/posener/complete"

var (
	ellipsis   = complete.PredictSet("./...")
	anyPackage = complete.PredictFunc(predictPackages)
	goFiles    = complete.PredictFiles("*.go")
	anyFile    = complete.PredictFiles("*")
	anyGo      = complete.PredictOr(goFiles, anyPackage, ellipsis)
)

func main() {
	build := complete.Command{
		Flags: complete.Flags{
			"-o": anyFile,
			"-i": complete.PredictNothing,

			"-a":             complete.PredictNothing,
			"-n":             complete.PredictNothing,
			"-p":             complete.PredictAnything,
			"-race":          complete.PredictNothing,
			"-msan":          complete.PredictNothing,
			"-v":             complete.PredictNothing,
			"-work":          complete.PredictNothing,
			"-x":             complete.PredictNothing,
			"-asmflags":      complete.PredictAnything,
			"-buildmode":     complete.PredictAnything,
			"-compiler":      complete.PredictAnything,
			"-gccgoflags":    complete.PredictAnything,
			"-gcflags":       complete.PredictAnything,
			"-installsuffix": complete.PredictAnything,
			"-ldflags":       complete.PredictAnything,
			"-linkshared":    complete.PredictNothing,
			"-pkgdir":        anyPackage,
			"-tags":          complete.PredictAnything,
			"-toolexec":      complete.PredictAnything,
		},
		Args: anyGo,
	}

	run := complete.Command{
		Flags: complete.Flags{
			"-exec": complete.PredictAnything,
		},
		Args: goFiles,
	}

	test := complete.Command{
		Flags: complete.Flags{
			"-args": complete.PredictAnything,
			"-c":    complete.PredictNothing,
			"-exec": complete.PredictAnything,

			"-bench":     predictBenchmark,
			"-benchtime": complete.PredictAnything,
			"-count":     complete.PredictAnything,
			"-cover":     complete.PredictNothing,
			"-covermode": complete.PredictSet("set", "count", "atomic"),
			"-coverpkg":  complete.PredictDirs("*"),
			"-cpu":       complete.PredictAnything,
			"-run":       predictTest,
			"-short":     complete.PredictNothing,
			"-timeout":   complete.PredictAnything,

			"-benchmem":             complete.PredictNothing,
			"-blockprofile":         complete.PredictFiles("*.out"),
			"-blockprofilerate":     complete.PredictAnything,
			"-coverprofile":         complete.PredictFiles("*.out"),
			"-cpuprofile":           complete.PredictFiles("*.out"),
			"-memprofile":           complete.PredictFiles("*.out"),
			"-memprofilerate":       complete.PredictAnything,
			"-mutexprofile":         complete.PredictFiles("*.out"),
			"-mutexprofilefraction": complete.PredictAnything,
			"-outputdir":            complete.PredictDirs("*"),
			"-trace":                complete.PredictFiles("*.out"),
		},
		Args: anyGo,
	}

	fmt := complete.Command{
		Flags: complete.Flags{
			"-n": complete.PredictNothing,
			"-x": complete.PredictNothing,
		},
		Args: anyGo,
	}

	get := complete.Command{
		Flags: complete.Flags{
			"-d":        complete.PredictNothing,
			"-f":        complete.PredictNothing,
			"-fix":      complete.PredictNothing,
			"-insecure": complete.PredictNothing,
			"-t":        complete.PredictNothing,
			"-u":        complete.PredictNothing,
		},
		Args: anyGo,
	}

	generate := complete.Command{
		Flags: complete.Flags{
			"-n":   complete.PredictNothing,
			"-x":   complete.PredictNothing,
			"-v":   complete.PredictNothing,
			"-run": complete.PredictAnything,
		},
		Args: anyGo,
	}

	vet := complete.Command{
		Flags: complete.Flags{
			"-n": complete.PredictNothing,
			"-x": complete.PredictNothing,
		},
		Args: anyGo,
	}

	list := complete.Command{
		Flags: complete.Flags{
			"-e":    complete.PredictNothing,
			"-f":    complete.PredictAnything,
			"-json": complete.PredictNothing,
		},
		Args: complete.PredictOr(anyPackage, ellipsis),
	}

	doc := complete.Command{
		Flags: complete.Flags{
			"-c":   complete.PredictNothing,
			"-cmd": complete.PredictNothing,
			"-u":   complete.PredictNothing,
		},
		Args: anyPackage,
	}

	tool := complete.Command{
		Flags: complete.Flags{
			"-n": complete.PredictNothing,
		},
		Sub: complete.Commands{
			"addr2line": {
				Args: anyFile,
			},
			"asm": {
				Flags: complete.Flags{
					"-D":        complete.PredictAnything,
					"-I":        complete.PredictDirs("*"),
					"-S":        complete.PredictNothing,
					"-debug":    complete.PredictNothing,
					"-dynlink":  complete.PredictNothing,
					"-e":        complete.PredictNothing,
					"-o":        anyFile,
					"-shared":   complete.PredictNothing,
					"-trimpath": complete.PredictNothing,
				},
				Args: complete.PredictFiles("*.s"),
			},
			"cgo": {
				Flags: complete.Flags{
					"-debug-define":      complete.PredictNothing,
					"debug-gcc":          complete.PredictNothing,
					"dynimport":          anyFile,
					"dynlinker":          complete.PredictNothing,
					"dynout":             anyFile,
					"dynpackage":         anyPackage,
					"exportheader":       complete.PredictDirs("*"),
					"gccgo":              complete.PredictNothing,
					"gccgopkgpath":       complete.PredictDirs("*"),
					"gccgoprefix":        complete.PredictAnything,
					"godefs":             complete.PredictNothing,
					"import_runtime_cgo": complete.PredictNothing,
					"import_syscall":     complete.PredictNothing,
					"importpath":         complete.PredictDirs("*"),
					"objdir":             complete.PredictDirs("*"),
					"srcdir":             complete.PredictDirs("*"),
				},
				Args: goFiles,
			},
			"compile": {
				Flags: complete.Flags{
					"-%":              complete.PredictNothing,
					"-+":              complete.PredictNothing,
					"-B":              complete.PredictNothing,
					"-D":              complete.PredictDirs("*"),
					"-E":              complete.PredictNothing,
					"-I":              complete.PredictDirs("*"),
					"-K":              complete.PredictNothing,
					"-N":              complete.PredictNothing,
					"-S":              complete.PredictNothing,
					"-V":              complete.PredictNothing,
					"-W":              complete.PredictNothing,
					"-asmhdr":         anyFile,
					"-bench":          anyFile,
					"-buildid":        complete.PredictNothing,
					"-complete":       complete.PredictNothing,
					"-cpuprofile":     anyFile,
					"-d":              complete.PredictNothing,
					"-dynlink":        complete.PredictNothing,
					"-e":              complete.PredictNothing,
					"-f":              complete.PredictNothing,
					"-h":              complete.PredictNothing,
					"-i":              complete.PredictNothing,
					"-importmap":      complete.PredictAnything,
					"-installsuffix":  complete.PredictAnything,
					"-j":              complete.PredictNothing,
					"-l":              complete.PredictNothing,
					"-largemodel":     complete.PredictNothing,
					"-linkobj":        anyFile,
					"-live":           complete.PredictNothing,
					"-m":              complete.PredictNothing,
					"-memprofile":     complete.PredictNothing,
					"-memprofilerate": complete.PredictAnything,
					"-msan":           complete.PredictNothing,
					"-nolocalimports": complete.PredictNothing,
					"-o":              anyFile,
					"-p":              complete.PredictDirs("*"),
					"-pack":           complete.PredictNothing,
					"-r":              complete.PredictNothing,
					"-race":           complete.PredictNothing,
					"-s":              complete.PredictNothing,
					"-shared":         complete.PredictNothing,
					"-traceprofile":   anyFile,
					"-trimpath":       complete.PredictAnything,
					"-u":              complete.PredictNothing,
					"-v":              complete.PredictNothing,
					"-w":              complete.PredictNothing,
					"-wb":             complete.PredictNothing,
				},
				Args: goFiles,
			},
			"cover": {
				Flags: complete.Flags{
					"-func": complete.PredictAnything,
					"-html": complete.PredictAnything,
					"-mode": complete.PredictSet("set", "count", "atomic"),
					"-o":    anyFile,
					"-var":  complete.PredictAnything,
				},
				Args: anyFile,
			},
			"dist": {
				Sub: complete.Commands{
					"banner":    {Flags: complete.Flags{"-v": complete.PredictNothing}},
					"bootstrap": {Flags: complete.Flags{"-v": complete.PredictNothing}},
					"clean":     {Flags: complete.Flags{"-v": complete.PredictNothing}},
					"env":       {Flags: complete.Flags{"-v": complete.PredictNothing, "-p": complete.PredictNothing}},
					"install":   {Flags: complete.Flags{"-v": complete.PredictNothing}, Args: complete.PredictDirs("*")},
					"list":      {Flags: complete.Flags{"-v": complete.PredictNothing, "-json": complete.PredictNothing}},
					"test":      {Flags: complete.Flags{"-v": complete.PredictNothing, "-h": complete.PredictNothing}},
					"version":   {Flags: complete.Flags{"-v": complete.PredictNothing}},
				},
			},
			"doc": doc,
			"fix": {
				Flags: complete.Flags{
					"-diff":  complete.PredictNothing,
					"-force": complete.PredictAnything,
					"-r":     complete.PredictSet("context", "gotypes", "netipv6zone", "printerconfig"),
				},
				Args: anyGo,
			},
			"link": {},
			"nm": {
				Flags: complete.Flags{
					"-n":    complete.PredictNothing,
					"-size": complete.PredictNothing,
					"-sort": complete.PredictAnything,
					"-type": complete.PredictNothing,
				},
				Args: anyGo,
			},
			"objdump": {
				Flags: complete.Flags{
					"-s": complete.PredictAnything,
				},
				Args: anyFile,
			},
			"pack": {},
			"pprof": {
				Flags: complete.Flags{
					"-callgrind":     complete.PredictNothing,
					"-disasm":        complete.PredictAnything,
					"-dot":           complete.PredictNothing,
					"-eog":           complete.PredictNothing,
					"-evince":        complete.PredictNothing,
					"-gif":           complete.PredictNothing,
					"-gv":            complete.PredictNothing,
					"-list":          complete.PredictAnything,
					"-pdf":           complete.PredictNothing,
					"-peek":          complete.PredictAnything,
					"-png":           complete.PredictNothing,
					"-proto":         complete.PredictNothing,
					"-ps":            complete.PredictNothing,
					"-raw":           complete.PredictNothing,
					"-svg":           complete.PredictNothing,
					"-tags":          complete.PredictNothing,
					"-text":          complete.PredictNothing,
					"-top":           complete.PredictNothing,
					"-tree":          complete.PredictNothing,
					"-web":           complete.PredictNothing,
					"-weblist":       complete.PredictAnything,
					"-output":        anyFile,
					"-functions":     complete.PredictNothing,
					"-files":         complete.PredictNothing,
					"-lines":         complete.PredictNothing,
					"-addresses":     complete.PredictNothing,
					"-base":          complete.PredictAnything,
					"-drop_negative": complete.PredictNothing,
					"-cum":           complete.PredictNothing,
					"-seconds":       complete.PredictAnything,
					"-nodecount":     complete.PredictAnything,
					"-nodefraction":  complete.PredictAnything,
					"-edgefraction":  complete.PredictAnything,
					"-sample_index":  complete.PredictNothing,
					"-mean":          complete.PredictNothing,
					"-inuse_space":   complete.PredictNothing,
					"-inuse_objects": complete.PredictNothing,
					"-alloc_space":   complete.PredictNothing,
					"-alloc_objects": complete.PredictNothing,
					"-total_delay":   complete.PredictNothing,
					"-contentions":   complete.PredictNothing,
					"-mean_delay":    complete.PredictNothing,
					"-runtime":       complete.PredictNothing,
					"-focus":         complete.PredictAnything,
					"-ignore":        complete.PredictAnything,
					"-tagfocus":      complete.PredictAnything,
					"-tagignore":     complete.PredictAnything,
					"-call_tree":     complete.PredictNothing,
					"-unit":          complete.PredictAnything,
					"-divide_by":     complete.PredictAnything,
					"-buildid":       complete.PredictAnything,
					"-tools":         complete.PredictDirs("*"),
					"-help":          complete.PredictNothing,
				},
				Args: anyFile,
			},
			"tour": {
				Flags: complete.Flags{
					"-http":        complete.PredictAnything,
					"-openbrowser": complete.PredictNothing,
				},
			},
			"trace": {
				Flags: complete.Flags{
					"-http":  complete.PredictAnything,
					"-pprof": complete.PredictSet("net", "sync", "syscall", "sched"),
				},
				Args: anyFile,
			},
			"vet": {
				Flags: complete.Flags{
					"-all":                 complete.PredictNothing,
					"-asmdecl":             complete.PredictNothing,
					"-assign":              complete.PredictNothing,
					"-atomic":              complete.PredictNothing,
					"-bool":                complete.PredictNothing,
					"-buildtags":           complete.PredictNothing,
					"-cgocall":             complete.PredictNothing,
					"-composites":          complete.PredictNothing,
					"-compositewhitelist":  complete.PredictNothing,
					"-copylocks":           complete.PredictNothing,
					"-httpresponse":        complete.PredictNothing,
					"-lostcancel":          complete.PredictNothing,
					"-methods":             complete.PredictNothing,
					"-nilfunc":             complete.PredictNothing,
					"-printf":              complete.PredictNothing,
					"-printfuncs":          complete.PredictAnything,
					"-rangeloops":          complete.PredictNothing,
					"-shadow":              complete.PredictNothing,
					"-shadowstrict":        complete.PredictNothing,
					"-shift":               complete.PredictNothing,
					"-structtags":          complete.PredictNothing,
					"-tags":                complete.PredictAnything,
					"-tests":               complete.PredictNothing,
					"-unreachable":         complete.PredictNothing,
					"-unsafeptr":           complete.PredictNothing,
					"-unusedfuncs":         complete.PredictAnything,
					"-unusedresult":        complete.PredictNothing,
					"-unusedstringmethods": complete.PredictAnything,
					"-v": complete.PredictNothing,
				},
				Args: anyGo,
			},
		},
	}

	clean := complete.Command{
		Flags: complete.Flags{
			"-i": complete.PredictNothing,
			"-r": complete.PredictNothing,
			"-n": complete.PredictNothing,
			"-x": complete.PredictNothing,
		},
		Args: complete.PredictOr(anyPackage, ellipsis),
	}

	env := complete.Command{
		Args: complete.PredictAnything,
	}

	bug := complete.Command{}
	version := complete.Command{}

	fix := complete.Command{
		Args: anyGo,
	}

	// commands that also accepts the build flags
	for name, options := range build.Flags {
		test.Flags[name] = options
		run.Flags[name] = options
		list.Flags[name] = options
		vet.Flags[name] = options
		get.Flags[name] = options
	}

	gogo := complete.Command{
		Sub: complete.Commands{
			"build":    build,
			"install":  build, // install and build have the same flags
			"run":      run,
			"test":     test,
			"fmt":      fmt,
			"get":      get,
			"generate": generate,
			"vet":      vet,
			"list":     list,
			"doc":      doc,
			"tool":     tool,
			"clean":    clean,
			"env":      env,
			"bug":      bug,
			"fix":      fix,
			"version":  version,
		},
		GlobalFlags: complete.Flags{
			"-h": complete.PredictNothing,
		},
	}

	complete.New("go", gogo).Run()
}
<|endoftext|>