<commit_before>package orm

import (
	"net/url"
	"strings"

	"github.com/go-pg/pg/types"
)

// URLFilter is used with Query.Apply to add WHERE clauses from the URL values:
//   - ?foo=bar - Where(`"foo" = 'bar'`)
//   - ?foo=hello&foo=world - Where(`"foo" IN ('hello','world')`)
//   - ?foo__exclude=bar - Where(`"foo" != 'bar'`)
//   - ?foo__ieq=bar - Where(`"foo" ILIKE 'bar'`)
//   - ?foo__match=bar - Where(`"foo" SIMILAR TO 'bar'`)
//   - ?foo__gt=42 - Where(`"foo" > 42`)
//   - ?foo__gte=42 - Where(`"foo" >= 42`)
//   - ?foo__lt=42 - Where(`"foo" < 42`)
//   - ?foo__lte=42 - Where(`"foo" <= 42`)
type URLFilter struct {
	values  URLValues
	allowed map[string]struct{}
}

func NewURLFilter(values url.Values) *URLFilter {
	return &URLFilter{
		values: URLValues(values),
	}
}

// Values returns URL values.
func (f *URLFilter) Values() URLValues {
	return f.values
}

func (f *URLFilter) Allow(filter string) {
	if f.allowed == nil {
		f.allowed = make(map[string]struct{})
	}
	f.allowed[filter] = struct{}{}
}

func (f *URLFilter) isAllowed(filter string) bool {
	if len(f.allowed) == 0 {
		return true
	}
	_, ok := f.allowed[filter]
	return ok
}

func (f *URLFilter) Filters(q *Query) (*Query, error) {
	if f == nil {
		return q, nil
	}

	for filter, values := range f.values {
		if strings.HasSuffix(filter, "[]") {
			filter = filter[:len(filter)-2]
		}

		if !f.isAllowed(filter) {
			continue
		}

		var operation string
		if i := strings.Index(filter, "__"); i != -1 {
			filter, operation = filter[:i], filter[i+2:]
		}

		if q.model.Table().HasField(filter) {
			q = addOperator(q, filter, operation, values)
		}
	}
	return q, nil
}

// URLFilters is a shortcut for NewURLFilter(urlValues).Filters.
func URLFilters(urlValues url.Values) func(*Query) (*Query, error) {
	return NewURLFilter(urlValues).Filters
}

func addOperator(q *Query, field, operator string, values []string) *Query {
	switch operator {
	case "gt":
		q = forEachValue(q, field, values, "? > ?")
	case "gte":
		q = forEachValue(q, field, values, "? >= ?")
	case "lt":
		q = forEachValue(q, field, values, "? < ?")
	case "lte":
		q = forEachValue(q, field, values, "? <= ?")
	case "ieq":
		q = forEachValue(q, field, values, "? ILIKE ?")
	case "match":
		q = forEachValue(q, field, values, "? SIMILAR TO ?")
	case "exclude":
		q = forAllValues(q, field, values, "? != ?", "? NOT IN (?)")
	case "", "include":
		q = forAllValues(q, field, values, "? = ?", "? IN (?)")
	}
	return q
}

func forEachValue(q *Query, field string, values []string, queryTemplate string) *Query {
	for _, value := range values {
		q = q.Where(queryTemplate, types.F(field), value)
	}
	return q
}

func forAllValues(q *Query, field string, values []string, queryTemplate, queryArrayTemplate string) *Query {
	if len(values) > 1 {
		q = q.Where(queryArrayTemplate, types.F(field), types.InSlice(values))
	} else {
		q = q.Where(queryTemplate, types.F(field), values[0])
	}
	return q
}
<commit_msg>orm: accept multiple filters<commit_after>package orm

import (
	"net/url"
	"strings"

	"github.com/go-pg/pg/types"
)

// URLFilter is used with Query.Apply to add WHERE clauses from the URL values:
//   - ?foo=bar - Where(`"foo" = 'bar'`)
//   - ?foo=hello&foo=world - Where(`"foo" IN ('hello','world')`)
//   - ?foo__exclude=bar - Where(`"foo" != 'bar'`)
//   - ?foo__ieq=bar - Where(`"foo" ILIKE 'bar'`)
//   - ?foo__match=bar - Where(`"foo" SIMILAR TO 'bar'`)
//   - ?foo__gt=42 - Where(`"foo" > 42`)
//   - ?foo__gte=42 - Where(`"foo" >= 42`)
//   - ?foo__lt=42 - Where(`"foo" < 42`)
//   - ?foo__lte=42 - Where(`"foo" <= 42`)
type URLFilter struct {
	values  URLValues
	allowed map[string]struct{}
}

func NewURLFilter(values url.Values) *URLFilter {
	return &URLFilter{
		values: URLValues(values),
	}
}

// Values returns URL values.
func (f *URLFilter) Values() URLValues {
	return f.values
}

func (f *URLFilter) Allow(filters ...string) {
	if f.allowed == nil {
		f.allowed = make(map[string]struct{})
	}
	for _, filter := range filters {
		f.allowed[filter] = struct{}{}
	}
}

func (f *URLFilter) isAllowed(filter string) bool {
	if len(f.allowed) == 0 {
		return true
	}
	_, ok := f.allowed[filter]
	return ok
}

func (f *URLFilter) Filters(q *Query) (*Query, error) {
	if f == nil {
		return q, nil
	}

	for filter, values := range f.values {
		if strings.HasSuffix(filter, "[]") {
			filter = filter[:len(filter)-2]
		}

		if !f.isAllowed(filter) {
			continue
		}

		var operation string
		if i := strings.Index(filter, "__"); i != -1 {
			filter, operation = filter[:i], filter[i+2:]
		}

		if q.model.Table().HasField(filter) {
			q = addOperator(q, filter, operation, values)
		}
	}
	return q, nil
}

// URLFilters is a shortcut for NewURLFilter(urlValues).Filters.
func URLFilters(urlValues url.Values) func(*Query) (*Query, error) {
	return NewURLFilter(urlValues).Filters
}

func addOperator(q *Query, field, operator string, values []string) *Query {
	switch operator {
	case "gt":
		q = forEachValue(q, field, values, "? > ?")
	case "gte":
		q = forEachValue(q, field, values, "? >= ?")
	case "lt":
		q = forEachValue(q, field, values, "? < ?")
	case "lte":
		q = forEachValue(q, field, values, "? <= ?")
	case "ieq":
		q = forEachValue(q, field, values, "? ILIKE ?")
	case "match":
		q = forEachValue(q, field, values, "? SIMILAR TO ?")
	case "exclude":
		q = forAllValues(q, field, values, "? != ?", "? NOT IN (?)")
	case "", "include":
		q = forAllValues(q, field, values, "? = ?", "? IN (?)")
	}
	return q
}

func forEachValue(q *Query, field string, values []string, queryTemplate string) *Query {
	for _, value := range values {
		q = q.Where(queryTemplate, types.F(field), value)
	}
	return q
}

func forAllValues(q *Query, field string, values []string, queryTemplate, queryArrayTemplate string) *Query {
	if len(values) > 1 {
		q = q.Where(queryArrayTemplate, types.F(field), types.InSlice(values))
	} else {
		q = q.Where(queryTemplate, types.F(field), values[0])
	}
	return q
}
<|endoftext|>


<commit_before>package world

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"time"

	"github.com/cloudfoundry-incubator/inigo/fake_cc"
	"github.com/cloudfoundry-incubator/inigo/loggregator_runner"
	wardenrunner "github.com/cloudfoundry-incubator/warden-linux/integration/runner"
	gorouterconfig "github.com/cloudfoundry/gorouter/config"
	"github.com/fraenkel/candiedyaml"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

type BuiltExecutables map[string]string
type BuiltCircuses map[string]string

const CircusZipFilename = "some-circus.tgz"

type BuiltArtifacts struct {
	Executables BuiltExecutables
	Circuses    BuiltCircuses
}

type ComponentAddresses struct {
	NATS           string
	Etcd           string
	EtcdPeer       string
	LoggregatorIn  string
	LoggregatorOut string
	Executor       string
	Rep            string
	FakeCC         string
	FileServer     string
	Router         string
	TPS            string
	WardenLinux    string
}

type ComponentMaker struct {
	Artifacts BuiltArtifacts
	Addresses ComponentAddresses

	ExternalAddress string

	Stack string

	WardenBinPath    string
	WardenRootFSPath string
	WardenGraphPath  string
}

func (maker ComponentMaker) NATS(argv ...string) ifrit.Runner {
	host, port, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "gnatsd",
		AnsiColorCode:     "30",
		StartCheck:        "gnatsd is ready",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"gnatsd",
			append([]string{
				"--addr", host,
				"--port", port,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Etcd(argv ...string) ifrit.Runner {
	nodeName := fmt.Sprintf("etcd_%d", ginkgo.GinkgoParallelNode())
	dataDir := path.Join(os.TempDir(), nodeName)

	return &ginkgomon.Runner{
		Name:              "etcd",
		AnsiColorCode:     "30",
		StartCheck:        "leader changed",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"etcd",
			append([]string{
				"-data-dir", dataDir,
				"-addr", maker.Addresses.Etcd,
				"-peer-addr", maker.Addresses.EtcdPeer,
				"-name", nodeName,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(dataDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) WardenLinux(argv ...string) *wardenrunner.Runner {
	return wardenrunner.New(
		"tcp",
		maker.Addresses.WardenLinux,
		maker.Artifacts.Executables["warden-linux"],
		maker.WardenBinPath,
		maker.WardenRootFSPath,
		argv...,
	)
}

func (maker ComponentMaker) Executor(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "exec",
		AnsiColorCode:     "31",
		StartCheck:        "executor.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["exec"],
			append([]string{
				"-listenAddr", maker.Addresses.Executor,
				"-wardenNetwork", "tcp",
				"-wardenAddr", maker.Addresses.WardenLinux,
				"-loggregatorServer", maker.Addresses.LoggregatorIn,
				"-loggregatorSecret", "loggregator-secret",
				"-containerMaxCpuShares", "1024",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Rep(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "rep",
		AnsiColorCode:     "32",
		StartCheck:        "rep.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["rep"],
			append(
				[]string{
					"-stack", maker.Stack,
					"-lrpHost", maker.ExternalAddress,
					"-listenAddr", maker.Addresses.Rep,
					"-etcdCluster", "http://" + maker.Addresses.Etcd,
					"-natsAddresses", maker.Addresses.NATS,
					"-executorID", "the-executor-id-" + strconv.Itoa(ginkgo.GinkgoParallelNode()),
					"-executorURL", "http://" + maker.Addresses.Executor,
					"-heartbeatInterval", "1s",
				},
				argv...,
			)...,
		),
	}
}

func (maker ComponentMaker) Converger(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "converger",
		AnsiColorCode:     "32",
		StartCheck:        "converger.started",
		StartCheckTimeout: 5 * time.Second,

		Command: exec.Command(
			maker.Artifacts.Executables["converger"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Auctioneer(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "auctioneer",
		AnsiColorCode: "33",
		StartCheck:    "auctioneer.started",
		Command: exec.Command(
			maker.Artifacts.Executables["auctioneer"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,

				// we limit this to prevent overwhelming numbers of auctioneer logs.  it
				// should not impact the behavior of the tests.
				"-maxRounds", "3",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) RouteEmitter(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "route-emitter",
		AnsiColorCode: "35",
		StartCheck:    "route-emitter.started",
		Command: exec.Command(
			maker.Artifacts.Executables["route-emitter"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) TPS(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "tps",
		AnsiColorCode: "36",
		StartCheck:    "tps.started",
		Command: exec.Command(
			maker.Artifacts.Executables["tps"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-listenAddr", maker.Addresses.TPS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) NsyncListener(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "nsync-listener",
		AnsiColorCode: "37",
		StartCheck:    "nsync.listener.started",
		Command: exec.Command(
			maker.Artifacts.Executables["nsync-listener"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-repAddrRelativeToExecutor", maker.Addresses.Rep,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) FileServer(argv ...string) (ifrit.Runner, string) {
	servedFilesDir, err := ioutil.TempDir("", "file-server-files")
	Ω(err).ShouldNot(HaveOccurred())

	host, port, err := net.SplitHostPort(maker.Addresses.FileServer)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "file-server",
		AnsiColorCode: "31",
		StartCheck:    "file-server.ready",
		Command: exec.Command(
			maker.Artifacts.Executables["file-server"],
			append([]string{
				"-address", host,
				"-port", port,
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-ccAddress", "http://" + maker.Addresses.FakeCC,
				"-ccJobPollingInterval", "100ms",
				"-ccUsername", fake_cc.CC_USERNAME,
				"-ccPassword", fake_cc.CC_PASSWORD,
				"-staticDirectory", servedFilesDir,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(servedFilesDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}, servedFilesDir
}

func (maker ComponentMaker) Router() ifrit.Runner {
	_, routerPort, err := net.SplitHostPort(maker.Addresses.Router)
	Ω(err).ShouldNot(HaveOccurred())

	routerPortInt, err := strconv.Atoi(routerPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	routerConfig := &gorouterconfig.Config{
		Port: uint16(routerPortInt),

		PruneStaleDropletsIntervalInSeconds: 5,
		DropletStaleThresholdInSeconds:      10,
		PublishActiveAppsIntervalInSeconds:  0,
		StartResponseDelayIntervalInSeconds: 1,

		Nats: []gorouterconfig.NatsConfig{
			{
				Host: natsHost,
				Port: uint16(natsPortInt),
			},
		},
		Logging: gorouterconfig.LoggingConfig{
			File:  "/dev/stdout",
			Level: "info",
		},
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "router-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = candiedyaml.NewEncoder(configFile).Encode(routerConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "router",
		AnsiColorCode:     "32",
		StartCheck:        "router.started",
		StartCheckTimeout: 5 * time.Second, // it waits 1 second before listening. yep.
		Command: exec.Command(
			maker.Artifacts.Executables["router"],
			"-c", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) Loggregator() ifrit.Runner {
	_, inPort, err := net.SplitHostPort(maker.Addresses.LoggregatorIn)
	Ω(err).ShouldNot(HaveOccurred())

	_, outPort, err := net.SplitHostPort(maker.Addresses.LoggregatorOut)
	Ω(err).ShouldNot(HaveOccurred())

	inPortInt, err := strconv.Atoi(inPort)
	Ω(err).ShouldNot(HaveOccurred())

	outPortInt, err := strconv.Atoi(outPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	loggregatorConfig := loggregator_runner.Config{
		IncomingPort:           inPortInt,
		OutgoingPort:           outPortInt,
		MaxRetainedLogMessages: 1000,
		SharedSecret:           "loggregator-secret",
		NatsHost:               natsHost,
		NatsPort:               natsPortInt,
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "loggregator-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = json.NewEncoder(configFile).Encode(loggregatorConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "loggregator",
		AnsiColorCode: "33",
		StartCheck:    "Listening on port",
		Command: exec.Command(
			maker.Artifacts.Executables["loggregator"],
			"-config", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) FakeCC() *fake_cc.FakeCC {
	return fake_cc.New(maker.Addresses.FakeCC)
}

func (maker ComponentMaker) Stager(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "stager",
		AnsiColorCode: "34",
		StartCheck:    "Listening for staging requests!",
		Command: exec.Command(
			maker.Artifacts.Executables["stager"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}
<commit_msg>fix colors<commit_after>package world

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"time"

	"github.com/cloudfoundry-incubator/inigo/fake_cc"
	"github.com/cloudfoundry-incubator/inigo/loggregator_runner"
	wardenrunner "github.com/cloudfoundry-incubator/warden-linux/integration/runner"
	gorouterconfig "github.com/cloudfoundry/gorouter/config"
	"github.com/fraenkel/candiedyaml"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

type BuiltExecutables map[string]string
type BuiltCircuses map[string]string

const CircusZipFilename = "some-circus.tgz"

type BuiltArtifacts struct {
	Executables BuiltExecutables
	Circuses    BuiltCircuses
}

type ComponentAddresses struct {
	NATS           string
	Etcd           string
	EtcdPeer       string
	LoggregatorIn  string
	LoggregatorOut string
	Executor       string
	Rep            string
	FakeCC         string
	FileServer     string
	Router         string
	TPS            string
	WardenLinux    string
}

type ComponentMaker struct {
	Artifacts BuiltArtifacts
	Addresses ComponentAddresses

	ExternalAddress string

	Stack string

	WardenBinPath    string
	WardenRootFSPath string
	WardenGraphPath  string
}

func (maker ComponentMaker) NATS(argv ...string) ifrit.Runner {
	host, port, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "gnatsd",
		AnsiColorCode:     "30m",
		StartCheck:        "gnatsd is ready",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"gnatsd",
			append([]string{
				"--addr", host,
				"--port", port,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Etcd(argv ...string) ifrit.Runner {
	nodeName := fmt.Sprintf("etcd_%d", ginkgo.GinkgoParallelNode())
	dataDir := path.Join(os.TempDir(), nodeName)

	return &ginkgomon.Runner{
		Name:              "etcd",
		AnsiColorCode:     "31m",
		StartCheck:        "leader changed",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"etcd",
			append([]string{
				"-data-dir", dataDir,
				"-addr", maker.Addresses.Etcd,
				"-peer-addr", maker.Addresses.EtcdPeer,
				"-name", nodeName,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(dataDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) WardenLinux(argv ...string) *wardenrunner.Runner {
	return wardenrunner.New(
		"tcp",
		maker.Addresses.WardenLinux,
		maker.Artifacts.Executables["warden-linux"],
		maker.WardenBinPath,
		maker.WardenRootFSPath,
		argv...,
	)
}

func (maker ComponentMaker) Executor(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "executor",
		AnsiColorCode:     "91m",
		StartCheck:        "executor.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["exec"],
			append([]string{
				"-listenAddr", maker.Addresses.Executor,
				"-wardenNetwork", "tcp",
				"-wardenAddr", maker.Addresses.WardenLinux,
				"-loggregatorServer", maker.Addresses.LoggregatorIn,
				"-loggregatorSecret", "loggregator-secret",
				"-containerMaxCpuShares", "1024",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Rep(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "rep",
		AnsiColorCode:     "92m",
		StartCheck:        "rep.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["rep"],
			append(
				[]string{
					"-stack", maker.Stack,
					"-lrpHost", maker.ExternalAddress,
					"-listenAddr", maker.Addresses.Rep,
					"-etcdCluster", "http://" + maker.Addresses.Etcd,
					"-natsAddresses", maker.Addresses.NATS,
					"-executorID", "the-executor-id-" + strconv.Itoa(ginkgo.GinkgoParallelNode()),
					"-executorURL", "http://" + maker.Addresses.Executor,
					"-heartbeatInterval", "1s",
				},
				argv...,
			)...,
		),
	}
}

func (maker ComponentMaker) Converger(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "converger",
		AnsiColorCode:     "93m",
		StartCheck:        "converger.started",
		StartCheckTimeout: 5 * time.Second,

		Command: exec.Command(
			maker.Artifacts.Executables["converger"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Auctioneer(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "auctioneer",
		AnsiColorCode: "94m",
		StartCheck:    "auctioneer.started",
		Command: exec.Command(
			maker.Artifacts.Executables["auctioneer"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,

				// we limit this to prevent overwhelming numbers of auctioneer logs.  it
				// should not impact the behavior of the tests.
				"-maxRounds", "3",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) RouteEmitter(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "route-emitter",
		AnsiColorCode: "95m",
		StartCheck:    "route-emitter.started",
		Command: exec.Command(
			maker.Artifacts.Executables["route-emitter"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) TPS(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "tps",
		AnsiColorCode: "96m",
		StartCheck:    "tps.started",
		Command: exec.Command(
			maker.Artifacts.Executables["tps"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-listenAddr", maker.Addresses.TPS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) NsyncListener(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "nsync-listener",
		AnsiColorCode: "97m",
		StartCheck:    "nsync.listener.started",
		Command: exec.Command(
			maker.Artifacts.Executables["nsync-listener"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-repAddrRelativeToExecutor", maker.Addresses.Rep,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) FileServer(argv ...string) (ifrit.Runner, string) {
	servedFilesDir, err := ioutil.TempDir("", "file-server-files")
	Ω(err).ShouldNot(HaveOccurred())

	host, port, err := net.SplitHostPort(maker.Addresses.FileServer)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "file-server",
		AnsiColorCode: "90m",
		StartCheck:    "file-server.ready",
		Command: exec.Command(
			maker.Artifacts.Executables["file-server"],
			append([]string{
				"-address", host,
				"-port", port,
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-ccAddress", "http://" + maker.Addresses.FakeCC,
				"-ccJobPollingInterval", "100ms",
				"-ccUsername", fake_cc.CC_USERNAME,
				"-ccPassword", fake_cc.CC_PASSWORD,
				"-staticDirectory", servedFilesDir,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(servedFilesDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}, servedFilesDir
}

func (maker ComponentMaker) Router() ifrit.Runner {
	_, routerPort, err := net.SplitHostPort(maker.Addresses.Router)
	Ω(err).ShouldNot(HaveOccurred())

	routerPortInt, err := strconv.Atoi(routerPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	routerConfig := &gorouterconfig.Config{
		Port: uint16(routerPortInt),

		PruneStaleDropletsIntervalInSeconds: 5,
		DropletStaleThresholdInSeconds:      10,
		PublishActiveAppsIntervalInSeconds:  0,
		StartResponseDelayIntervalInSeconds: 1,

		Nats: []gorouterconfig.NatsConfig{
			{
				Host: natsHost,
				Port: uint16(natsPortInt),
			},
		},
		Logging: gorouterconfig.LoggingConfig{
			File:  "/dev/stdout",
			Level: "info",
		},
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "router-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = candiedyaml.NewEncoder(configFile).Encode(routerConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "router",
		AnsiColorCode:     "32m",
		StartCheck:        "router.started",
		StartCheckTimeout: 5 * time.Second, // it waits 1 second before listening. yep.
		Command: exec.Command(
			maker.Artifacts.Executables["router"],
			"-c", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) Loggregator() ifrit.Runner {
	_, inPort, err := net.SplitHostPort(maker.Addresses.LoggregatorIn)
	Ω(err).ShouldNot(HaveOccurred())

	_, outPort, err := net.SplitHostPort(maker.Addresses.LoggregatorOut)
	Ω(err).ShouldNot(HaveOccurred())

	inPortInt, err := strconv.Atoi(inPort)
	Ω(err).ShouldNot(HaveOccurred())

	outPortInt, err := strconv.Atoi(outPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	loggregatorConfig := loggregator_runner.Config{
		IncomingPort:           inPortInt,
		OutgoingPort:           outPortInt,
		MaxRetainedLogMessages: 1000,
		SharedSecret:           "loggregator-secret",
		NatsHost:               natsHost,
		NatsPort:               natsPortInt,
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "loggregator-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = json.NewEncoder(configFile).Encode(loggregatorConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "loggregator",
		AnsiColorCode: "33m",
		StartCheck:    "Listening on port",
		Command: exec.Command(
			maker.Artifacts.Executables["loggregator"],
			"-config", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) FakeCC() *fake_cc.FakeCC {
	return fake_cc.New(maker.Addresses.FakeCC)
}

func (maker ComponentMaker) Stager(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "stager",
		AnsiColorCode: "94m",
		StartCheck:    "Listening for staging requests!",
		Command: exec.Command(
			maker.Artifacts.Executables["stager"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"image"
	"image/color"
	"image/png"
	"math/cmplx"
	"os"
)

func getColor(c int) color.RGBA {
	rgb := uint8(c % 255)
	return color.RGBA{rgb, rgb, rgb, 255}
}

func computeMandel(cx, cy float64) color.RGBA {
	var i int
	var absq float64
	var z, c complex128
	c = complex(cx, cy)
	for {
		z = z*z + c
		absq = cmplx.Abs(z)
		if i >= 1024 {
			return getColor(255)
		}
		if absq > 4.0 {
			return getColor(i)
		}
		i = i + 1
	}
	return getColor(0)
}

func main() {
	const maxX, maxY = 640, 480

	const centerRe, centerIm float64 = -0.7435669, 0.1314023
	const hd float64 = 0.0022878
	const hv float64 = float64(maxY) / float64(maxX) * hd

	const hdStepX = hd / maxX
	const hdStepY = hv / maxY

	const upperLeftRe = centerRe - (float64(maxX/2) * hdStepX)
	const upperLeftIm = centerIm - (float64(maxY/2) * hdStepY)

	f, err := os.OpenFile("mandel.png", os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		panic(err)
	}
	m := image.NewRGBA(image.Rect(0, 0, maxY, maxY))
	for y := 0; y < maxY; y++ {
		for x := 0; x < maxX; x++ {
			dx := float64(x) * hdStepX
			dy := float64(y) * hdStepY
			c := computeMandel(upperLeftRe+dx, upperLeftIm+dy)
			m.Set(x, maxY-y, c)
		}
	}
	if err = png.Encode(f, m); err != nil {
		panic(err)
	}
}
<commit_msg>bigger image<commit_after>package main

import (
	"image"
	"image/color"
	"image/png"
	"math/cmplx"
	"os"
)

func getColor(c int) color.RGBA {
	rgb := uint8(c % 255)
	return color.RGBA{rgb, rgb, rgb, 255}
}

func computeMandel(cx, cy float64) color.RGBA {
	var i int
	var absq float64
	var z, c complex128
	c = complex(cx, cy)
	for {
		z = z*z + c
		absq = cmplx.Abs(z)
		if i >= 1024 {
			return getColor(255)
		}
		if absq > 4.0 {
			return getColor(i)
		}
		i = i + 1
	}
	return getColor(0)
}

func main() {
	const maxX = 1024
	const maxY = 800

	const centerRe, centerIm float64 = -0.7435669, 0.1314023
	const hd float64 = 0.0022878
	const hv float64 = float64(maxY) / float64(maxX) * hd

	const hdStepX = hd / maxX
	const hdStepY = hv / maxY

	const upperLeftRe = centerRe - (float64(maxX/2) * hdStepX)
	const upperLeftIm = centerIm - (float64(maxY/2) * hdStepY)

	f, err := os.Create("mandel.png")
	if err != nil {
		panic(err)
	}
	m := image.NewRGBA(image.Rect(0, 0, maxX, maxY))
	b := m.Bounds()
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			dx := float64(x) * hdStepX
			dy := float64(y) * hdStepY
			c := computeMandel(upperLeftRe+dx, upperLeftIm+dy)
			m.Set(x, maxY-y, c)
		}
	}

	if err = png.Encode(f, m); err != nil {
		panic(err)
	}
}
<|endoftext|>


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

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
	"github.com/sdboyer/gps"
)

// safeWriter transactionalizes writes of manifest, lock, and vendor dir, both
// individually and in any combination, into a pseudo-atomic action with
// transactional rollback.
//
// It is not impervious to errors (writing to disk is hard), but it should
// guard against non-arcane failure conditions.
type safeWriter struct {
	root          string
	nm            *manifest // the new manifest to write
	lock          *lock     // the old lock, if any
	nl            gps.Lock  // the new lock to write, if desired
	sm            gps.SourceManager
	mpath, vendor string
}

// writeAllSafe writes out some combination of config yaml, lock, and a vendor
// tree, to a temp dir, then moves them into place if and only if all the write
// operations succeeded. It also does its best to roll back if any moves fail.
//
// This mostly guarantees that dep cannot terminate with a partial write,
// resulting in an undefined disk state.
//
// - If a gw.conf is provided, it will be written to the standard manifest file
// name beneath gw.pr
// - If gw.lock is provided without a gw.nl, it will be written to
//   `glide.lock` in the parent dir of gw.vendor
// - If gw.lock and gw.nl are both provided and are not equivalent,
//   the nl will be written to the same location as above, and a vendor
//   tree will be written to gw.vendor
// - If gw.nl is provided and gw.lock is not, it will write both a lock
//   and vendor dir in the same way
//
// Any of the conf, lock, or result can be omitted; the grouped write operation
// will continue for whichever inputs are present.
func (gw safeWriter) writeAllSafe() error {
	// Decide which writes we need to do
	var writeM, writeL, writeV bool

	if gw.nm != nil {
		writeM = true
	}

	if gw.nl != nil {
		if gw.lock == nil {
			writeL, writeV = true, true
		} else {
			rlf := lockFromInterface(gw.nl)
			if !locksAreEquivalent(rlf, gw.lock) {
				writeL, writeV = true, true
			}
		}
	} else if gw.lock != nil {
		writeL = true
	}

	if !writeM && !writeL && !writeV {
		// nothing to do
		return nil
	}

	if writeM && gw.mpath == "" {
		return fmt.Errorf("Must provide a path if writing out a config yaml.")
	}

	if (writeL || writeV) && gw.vendor == "" {
		return fmt.Errorf("Must provide a vendor dir if writing out a lock or vendor dir.")
	}

	if writeV && gw.sm == nil {
		return fmt.Errorf("Must provide a SourceManager if writing out a vendor dir.")
	}

	td, err := ioutil.TempDir(os.TempDir(), "dep")
	if err != nil {
		return errors.Wrap(err, "error while creating temp dir for writing manifest/lock/vendor")
	}
	defer os.RemoveAll(td)

	if writeM {
		if err := writeFile(filepath.Join(td, manifestName), gw.nm); err != nil {
			return errors.Wrap(err, "failed to write manifest file to temp dir")
		}
	}

	if writeL {
		if gw.nl == nil {
			// the new lock is nil but the flag is on, so we must be writing
			// the other one
			if err := writeFile(filepath.Join(td, lockName), gw.lock); err != nil {
				return errors.Wrap(err, "failed to write lock file to temp dir")
			}
		} else {
			rlf := lockFromInterface(gw.nl)
			if err := writeFile(filepath.Join(td, lockName), rlf); err != nil {
				return errors.Wrap(err, "failed to write lock file to temp dir")
			}
		}
	}

	if writeV {
		err = gps.WriteDepTree(filepath.Join(td, "vendor"), gw.nl, gw.sm, true)
		if err != nil {
			return fmt.Errorf("Error while generating vendor tree: %s", err)
		}
	}

	// Move the existing files and dirs to the temp dir while we put the new
	// ones in, to provide insurance against errors for as long as possible
	var fail bool
	var failerr error
	type pathpair struct {
		from, to string
	}
	var restore []pathpair

	if writeM {
		if _, err := os.Stat(gw.mpath); err == nil {
			// move out the old one
			tmploc := filepath.Join(td, manifestName+".orig")
			failerr = os.Rename(gw.mpath, tmploc)
			if failerr != nil {
				fail = true
			} else {
				restore = append(restore, pathpair{from: tmploc, to: gw.mpath})
			}
		}

		// move in the new one
		failerr = os.Rename(filepath.Join(td, manifestName), gw.mpath)
		if failerr != nil {
			fail = true
		}
	}

	if !fail && writeL {
		tgt := filepath.Join(filepath.Dir(gw.vendor), lockName)
		if _, err := os.Stat(tgt); err == nil {
			// move out the old one
			tmploc := filepath.Join(td, lockName+".orig")

			failerr = os.Rename(tgt, tmploc)
			if failerr != nil {
				fail = true
			} else {
				restore = append(restore, pathpair{from: tmploc, to: tgt})
			}
		}

		// move in the new one
		failerr = renameElseCopy(filepath.Join(td, lockName), tgt)
		if failerr != nil {
			fail = true
		}
	}

	// have to declare out here so it's present later
	var vendorbak string
	if !fail && writeV {
		if _, err := os.Stat(gw.vendor); err == nil {
			// move out the old vendor dir. just do it into an adjacent dir, to
			// try to mitigate the possibility of a pointless cross-filesystem
			// move with a temp dir
			vendorbak = gw.vendor + ".orig"
			if _, err := os.Stat(vendorbak); err == nil {
				// If that does already exist bite the bullet and use a proper
				// tempdir
				vendorbak = filepath.Join(td, "vendor.orig")
			}

			failerr = renameElseCopy(gw.vendor, vendorbak)
			if failerr != nil {
				fail = true
			} else {
				restore = append(restore, pathpair{from: vendorbak, to: gw.vendor})
			}
		}

		// move in the new one
		failerr = renameElseCopy(filepath.Join(td, "vendor"), gw.vendor)
		if failerr != nil {
			fail = true
		}
	}

	// If we failed at any point, move all the things back into place, then bail
	if fail {
		for _, pair := range restore {
			// Nothing we can do on err here, we're already in recovery mode
			renameElseCopy(pair.from, pair.to)
		}
		return failerr
	}

	// Renames all went smoothly. The deferred os.RemoveAll will get the temp
	// dir, but if we wrote vendor, we have to clean that up directly

	if writeV {
		// Again, kinda nothing we can do about an error at this point
		os.RemoveAll(vendorbak)
	}

	return nil
}
<commit_msg>Numerous var renames and comment updates<commit_after>// Copyright 2016 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 main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
	"github.com/sdboyer/gps"
)

// safeWriter transactionalizes writes of manifest, lock, and vendor dir, both
// individually and in any combination, into a pseudo-atomic action with
// transactional rollback.
//
// It is not impervious to errors (writing to disk is hard), but it should
// guard against non-arcane failure conditions.
type safeWriter struct {
	root string    // absolute path of root dir in which to write
	m    *manifest // the manifest to write, if any
	l    *lock     // the old lock, if any
	nl   gps.Lock  // the new lock, if any
	sm   gps.SourceManager
}

// writeAllSafe writes out some combination of config yaml, lock, and a vendor
// tree, to a temp dir, then moves them into place if and only if all the write
// operations succeeded. It also does its best to roll back if any moves fail.
//
// This mostly guarantees that dep cannot exit with a partial write that would
// leave an undefined state on disk.
//
// - If a sw.m is provided, it will be written to the standard manifest file
//   name beneath sw.root
// - If sw.l is provided without an sw.nl, it will be unconditionally written
//   to the standard lock file name in the root dir
// - If sw.l and sw.nl are both provided and are equivalent, then neither lock
//   nor vendor will be written
// - If sw.l and sw.nl are both provided and are not equivalent,
//   the nl will be written to the same location as above, and a vendor
//   tree will be written to sw.root/vendor
// - If sw.nl is provided and sw.lock is not, it will write both a lock
//   and vendor dir in the same way
// - If the forceVendor param is true, then vendor will be unconditionally
//   written out based on sw.nl if present, else sw.l, else error.
//
// Any of m, l, or nl can be omitted; the grouped write operation will continue
// for whichever inputs are present. A SourceManager is only required if vendor
// is being written.
func (sw safeWriter) writeAllSafe(forceVendor bool) error {
	// Decide which writes we need to do
	var writeM, writeL, writeV bool

	if sw.m != nil {
		writeM = true
	}

	if sw.nl != nil {
		if sw.l == nil {
			writeL, writeV = true, true
		} else {
			rlf := lockFromInterface(sw.nl)
			if !locksAreEquivalent(rlf, sw.l) {
				writeL, writeV = true, true
			}
		}
	} else if sw.l != nil {
		writeL = true
	}

	if sw.root == "" {
		return errors.New("root path must be non-empty")
	}
	if is, err := isDir(sw.root); !is {
		if err != nil {
			return err
		}
		return fmt.Errorf("root path %q does not exist", sw.root)
	}

	mpath := filepath.Join(sw.root, manifestName)
	lpath := filepath.Join(sw.root, lockName)
	vpath := filepath.Join(sw.root, "vendor")

	if !writeM && !writeL && !writeV {
		// nothing to do
		return nil
	}

	if (writeL || writeV) && vpath == "" {
		return errors.New("must provide a vendor dir if writing out a lock or vendor dir.")
	}

	if writeV && sw.sm == nil {
		return errors.New("must provide a SourceManager if writing out a vendor dir.")
	}

	td, err := ioutil.TempDir(os.TempDir(), "dep")
	if err != nil {
		return errors.Wrap(err, "error while creating temp dir for writing manifest/lock/vendor")
	}
	defer os.RemoveAll(td)

	if writeM {
		if err := writeFile(filepath.Join(td, manifestName), sw.m); err != nil {
			return errors.Wrap(err, "failed to write manifest file to temp dir")
		}
	}

	if writeL {
		if sw.nl == nil {
			// the new lock is nil but the flag is on, so we must be writing
			// the other one
			if err := writeFile(filepath.Join(td, lockName), sw.l); err != nil {
				return errors.Wrap(err, "failed to write lock file to temp dir")
			}
		} else {
			rlf := lockFromInterface(sw.nl)
			if err := writeFile(filepath.Join(td, lockName), rlf); err != nil {
				return errors.Wrap(err, "failed to write lock file to temp dir")
			}
		}
	}

	if writeV {
		err = gps.WriteDepTree(filepath.Join(td, "vendor"), sw.nl, sw.sm, true)
		if err != nil {
			return errors.Wrap(err, "error while writing out vendor tree")
		}
	}

	// Move the existing files and dirs to the temp dir while we put the new
	// ones in, to provide insurance against errors for as long as possible
	var fail bool
	var failerr error
	type pathpair struct {
		from, to string
	}
	var restore []pathpair

	if writeM {
		if _, err := os.Stat(mpath); err == nil {
			// move out the old one
			tmploc := filepath.Join(td, manifestName+".orig")
			failerr = os.Rename(mpath, tmploc)
			if failerr != nil {
				fail = true
			} else {
				restore = append(restore, pathpair{from: tmploc, to: mpath})
			}
		}

		// move in the new one
		failerr = os.Rename(filepath.Join(td, manifestName), mpath)
		if failerr != nil {
			fail = true
		}
	}

	if !fail && writeL {
		if _, err := os.Stat(lpath); err == nil {
			// move out the old one
			tmploc := filepath.Join(td, lockName+".orig")

			failerr = os.Rename(lpath, tmploc)
			if failerr != nil {
				fail = true
			} else {
				restore = append(restore, pathpair{from: tmploc, to: lpath})
			}
		}

		// move in the new one
		failerr = renameElseCopy(filepath.Join(td, lockName), lpath)
		if failerr != nil {
			fail = true
		}
	}

	// have to declare out here so it's present later
	var vendorbak string
	if !fail && writeV {
		if _, err := os.Stat(vpath); err == nil {
			// move out the old vendor dir. just do it into an adjacent dir, to
			// try to mitigate the possibility of a pointless cross-filesystem
			// move with a temp dir
			vendorbak = vpath + ".orig"
			if _, err := os.Stat(vendorbak); err == nil {
				// If that does already exist bite the bullet and use a proper
				// tempdir
				vendorbak = filepath.Join(td, "vendor.orig")
			}

			failerr = renameElseCopy(vpath, vendorbak)
			if failerr != nil {
				fail = true
			} else {
				restore = append(restore, pathpair{from: vendorbak, to: vpath})
			}
		}

		// move in the new one
		failerr = renameElseCopy(filepath.Join(td, "vendor"), vpath)
		if failerr != nil {
			fail = true
		}
	}

	// If we failed at any point, move all the things back into place, then bail
	if fail {
		for _, pair := range restore {
			// Nothing we can do on err here, we're already in recovery mode
			renameElseCopy(pair.from, pair.to)
		}
		return failerr
	}

	// Renames all went smoothly. The deferred os.RemoveAll will get the temp
	// dir, but if we wrote vendor, we have to clean that up directly

	if writeV {
		// Again, kinda nothing we can do about an error at this point
		os.RemoveAll(vendorbak)
	}

	return nil
}
<|endoftext|>


<commit_before>package disruption

import (
	"encoding/xml"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	g "github.com/onsi/ginkgo"

	"k8s.io/kubernetes/test/e2e/chaosmonkey"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/framework/ginkgowrapper"
	"k8s.io/kubernetes/test/e2e/upgrades"
	"k8s.io/kubernetes/test/utils/junit"
)

// TestData is passed to the invariant tests executed during the upgrade. The default UpgradeType
// is MasterUpgrade.
type TestData struct {
	UpgradeType    upgrades.UpgradeType
	UpgradeContext upgrades.UpgradeContext
}

// Run executes the provided fn in a test context, ensuring that invariants are preserved while the
// test is being executed. Description is used to populate the JUnit suite name, and testname is
// used to define the overall test that will be run.
func Run(description, testname string, adapter TestData, invariants []upgrades.Test, fn func()) {
	testSuite := &junit.TestSuite{Name: description, Package: testname}
	test := &junit.TestCase{Name: testname, Classname: testname}
	testSuite.TestCases = append(testSuite.TestCases, test)
	cm := chaosmonkey.New(func() {
		start := time.Now()
		defer finalizeTest(start, test)
		fn()
	})
	runChaosmonkey(cm, adapter, invariants, testSuite)
}

func runChaosmonkey(
	cm *chaosmonkey.Chaosmonkey,
	testData TestData,
	tests []upgrades.Test,
	testSuite *junit.TestSuite,
) {
	testFrameworks := createTestFrameworks(tests)
	for _, t := range tests {
		testCase := &junit.TestCase{
			Name:      t.Name(),
			Classname: "disruption_tests",
		}
		testSuite.TestCases = append(testSuite.TestCases, testCase)

		f, ok := testFrameworks[t.Name()]
		if !ok {
			panic(fmt.Sprintf("can't find test framework for %q", t.Name()))
		}
		cma := chaosMonkeyAdapter{
			TestData:   testData,
			framework:  f,
			test:       t,
			testReport: testCase,
		}
		cm.Register(cma.Test)
	}

	start := time.Now()
	defer func() {
		testSuite.Update()
		testSuite.Time = time.Since(start).Seconds()
		if framework.TestContext.ReportDir != "" {
			fname := filepath.Join(framework.TestContext.ReportDir, fmt.Sprintf("junit_%s_%d.xml", testSuite.Package, time.Now().Unix()))
			f, err := os.Create(fname)
			if err != nil {
				return
			}
			defer f.Close()
			xml.NewEncoder(f).Encode(testSuite)
		}
	}()
	cm.Do()
}

type chaosMonkeyAdapter struct {
	TestData

	test       upgrades.Test
	testReport *junit.TestCase
	framework  *framework.Framework
}

func (cma *chaosMonkeyAdapter) Test(sem *chaosmonkey.Semaphore) {
	start := time.Now()
	var once sync.Once
	ready := func() {
		once.Do(func() {
			sem.Ready()
		})
	}
	defer finalizeTest(start, cma.testReport)
	defer g.GinkgoRecover()
	defer ready()
	if skippable, ok := cma.test.(upgrades.Skippable); ok && skippable.Skip(cma.UpgradeContext) {
		g.By("skipping test " + cma.test.Name())
		cma.testReport.Skipped = "skipping test " + cma.test.Name()
		return
	}

	cma.framework.BeforeEach()
	cma.test.Setup(cma.framework)
	defer cma.test.Teardown(cma.framework)
	ready()
	cma.test.Test(cma.framework, sem.StopCh, cma.UpgradeType)
}

func finalizeTest(start time.Time, tc *junit.TestCase) {
	tc.Time = time.Since(start).Seconds()
	r := recover()
	if r == nil {
		return
	}

	switch r := r.(type) {
	case ginkgowrapper.FailurePanic:
		tc.Failures = []*junit.Failure{
			{
				Message: r.Message,
				Type:    "Failure",
				Value:   fmt.Sprintf("%s\n\n%s", r.Message, r.FullStackTrace),
			},
		}
	case ginkgowrapper.SkipPanic:
		tc.Skipped = fmt.Sprintf("%s:%d %q", r.Filename, r.Line, r.Message)
	default:
		tc.Errors = []*junit.Error{
			{
				Message: fmt.Sprintf("%v", r),
				Type:    "Panic",
				Value:   fmt.Sprintf("%v\n\n%s", r, debug.Stack()),
			},
		}
	}
}

// TODO: accept a default framework
func createTestFrameworks(tests []upgrades.Test) map[string]*framework.Framework {
	nsFilter := regexp.MustCompile("[^[:word:]-]+") // match anything that's not a word character or hyphen
	testFrameworks := map[string]*framework.Framework{}
	for _, t := range tests {
		ns := nsFilter.ReplaceAllString(t.Name(), "-") // and replace with a single hyphen
		ns = strings.Trim(ns, "-")
		// identify tests that come from kube as strictly e2e tests so they get the correct semantics
		if strings.Contains(reflect.ValueOf(t).Elem().Type().PkgPath(), "/kubernetes/test/e2e/") {
			ns = "e2e-k8s-" + ns
		}
		testFrameworks[t.Name()] = &framework.Framework{
			BaseName:                 ns,
			AddonResourceConstraints: make(map[string]framework.ResourceConstraint),
			Options: framework.Options{
				ClientQPS:   20,
				ClientBurst: 50,
			},
		}
	}
	return testFrameworks
}
<commit_msg>test: Remove duplicate gingko recover from disruption tests<commit_after>package disruption

import (
	"encoding/xml"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	g "github.com/onsi/ginkgo"

	"k8s.io/kubernetes/test/e2e/chaosmonkey"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/framework/ginkgowrapper"
	"k8s.io/kubernetes/test/e2e/upgrades"
	"k8s.io/kubernetes/test/utils/junit"
)

// TestData is passed to the invariant tests executed during the upgrade. The default UpgradeType
// is MasterUpgrade.
type TestData struct {
	UpgradeType    upgrades.UpgradeType
	UpgradeContext upgrades.UpgradeContext
}

// Run executes the provided fn in a test context, ensuring that invariants are preserved while the
// test is being executed. Description is used to populate the JUnit suite name, and testname is
// used to define the overall test that will be run.
func Run(description, testname string, adapter TestData, invariants []upgrades.Test, fn func()) {
	testSuite := &junit.TestSuite{Name: description, Package: testname}
	test := &junit.TestCase{Name: testname, Classname: testname}
	testSuite.TestCases = append(testSuite.TestCases, test)
	cm := chaosmonkey.New(func() {
		start := time.Now()
		defer finalizeTest(start, test)
		fn()
	})
	runChaosmonkey(cm, adapter, invariants, testSuite)
}

func runChaosmonkey(
	cm *chaosmonkey.Chaosmonkey,
	testData TestData,
	tests []upgrades.Test,
	testSuite *junit.TestSuite,
) {
	testFrameworks := createTestFrameworks(tests)
	for _, t := range tests {
		testCase := &junit.TestCase{
			Name:      t.Name(),
			Classname: "disruption_tests",
		}
		testSuite.TestCases = append(testSuite.TestCases, testCase)

		f, ok := testFrameworks[t.Name()]
		if !ok {
			panic(fmt.Sprintf("can't find test framework for %q", t.Name()))
		}
		cma := chaosMonkeyAdapter{
			TestData:   testData,
			framework:  f,
			test:       t,
			testReport: testCase,
		}
		cm.Register(cma.Test)
	}

	start := time.Now()
	defer func() {
		testSuite.Update()
		testSuite.Time = time.Since(start).Seconds()
		if framework.TestContext.ReportDir != "" {
			fname := filepath.Join(framework.TestContext.ReportDir, fmt.Sprintf("junit_%s_%d.xml", testSuite.Package, time.Now().Unix()))
			f, err := os.Create(fname)
			if err != nil {
				return
			}
			defer f.Close()
			xml.NewEncoder(f).Encode(testSuite)
		}
	}()
	cm.Do()
}

type chaosMonkeyAdapter struct {
	TestData

	test       upgrades.Test
	testReport *junit.TestCase
	framework  *framework.Framework
}

func (cma *chaosMonkeyAdapter) Test(sem *chaosmonkey.Semaphore) {
	start := time.Now()
	var once sync.Once
	ready := func() {
		once.Do(func() {
			sem.Ready()
		})
	}
	defer finalizeTest(start, cma.testReport)
	defer ready()
	if skippable, ok := cma.test.(upgrades.Skippable); ok && skippable.Skip(cma.UpgradeContext) {
		g.By("skipping test " + cma.test.Name())
		cma.testReport.Skipped = "skipping test " + cma.test.Name()
		return
	}

	cma.framework.BeforeEach()
	cma.test.Setup(cma.framework)
	defer cma.test.Teardown(cma.framework)
	ready()
	cma.test.Test(cma.framework, sem.StopCh, cma.UpgradeType)
}

func finalizeTest(start time.Time, tc *junit.TestCase) {
	tc.Time = time.Since(start).Seconds()
	r := recover()
	if r == nil {
		return
	}

	switch r := r.(type) {
	case ginkgowrapper.FailurePanic:
		tc.Failures = []*junit.Failure{
			{
				Message: r.Message,
				Type:    "Failure",
				Value:   fmt.Sprintf("%s\n\n%s", r.Message, r.FullStackTrace),
			},
		}
	case ginkgowrapper.SkipPanic:
		tc.Skipped = fmt.Sprintf("%s:%d %q", r.Filename, r.Line, r.Message)
	default:
		tc.Errors = []*junit.Error{
			{
				Message: fmt.Sprintf("%v", r),
				Type:    "Panic",
				Value:   fmt.Sprintf("%v\n\n%s", r, debug.Stack()),
			},
		}
	}
}

// TODO: accept a default framework
func createTestFrameworks(tests []upgrades.Test) map[string]*framework.Framework {
	nsFilter := regexp.MustCompile("[^[:word:]-]+") // match anything that's not a word character or hyphen
	testFrameworks := map[string]*framework.Framework{}
	for _, t := range tests {
		ns := nsFilter.ReplaceAllString(t.Name(), "-") // and replace with a single hyphen
		ns = strings.Trim(ns, "-")
		// identify tests that come from kube as strictly e2e tests so they get the correct semantics
		if strings.Contains(reflect.ValueOf(t).Elem().Type().PkgPath(), "/kubernetes/test/e2e/") {
			ns = "e2e-k8s-" + ns
		}
		testFrameworks[t.Name()] = &framework.Framework{
			BaseName:                 ns,
			AddonResourceConstraints: make(map[string]framework.ResourceConstraint),
			Options: framework.Options{
				ClientQPS:   20,
				ClientBurst: 50,
			},
		}
	}
	return testFrameworks
}
<|endoftext|>


<commit_before>package bench

import (
	"fmt"

	"flag"
	"io/ioutil"

	"context"
	"math/rand"
	"time"
)

// ZipfSetBits sets bits randomly and deterministically based on a seed, according to the Zipf distribution
type ZipfSetBits struct {
	HasClient
	BaseBitmapID    int64
	BaseProfileID   int64
	BitmapIDRange   int64
	ProfileIDRange  int64
	Iterations      int // number of bits that will be set
	Seed            int64
	BitmapExponent  float64
	BitmapOffset    float64
	ProfileExponent float64
	ProfileOffset   float64
	DB              string // DB to use in pilosa.

}

func (b *ZipfSetBits) Usage() string {
	return `
zipf-set-bits sets random bits according to Zipf distribution

Usage: zipf-set-bits [arguments]

The following arguments are available:

	-base-bitmap-id int
		bits being set will all be greater than BaseBitmapID

	-bitmap-id-range int
		number of possible bitmap ids that can be set

	-base-profile-id int
		profile id num to start from

	-profile-id-range int
		number of possible profile ids that can be set

	-iterations int
		number of bits to set

	-seed int
		Seed for RNG

	-db string
		pilosa db to use

  -bitmap-exponent float64
    zipf exponent parameter for bitmap IDs

  -bitmap-offset float64
    zipf offset parameter for bitmap IDs

  -profile-exponent float64
    zipf exponent parameter for profile IDs

  -profile-offset float64
    zipf offset parameter for profile IDs

	-client-type string
		Can be 'single' (all agents hitting one host) or 'round_robin'
`[1:]
}

func (b *ZipfSetBits) ConsumeFlags(args []string) ([]string, error) {
	fs := flag.NewFlagSet("ZipfSetBits", flag.ContinueOnError)
	fs.SetOutput(ioutil.Discard)
	fs.Int64Var(&b.BaseBitmapID, "base-bitmap-id", 0, "")
	fs.Int64Var(&b.BitmapIDRange, "bitmap-id-range", 100000, "")
	fs.Int64Var(&b.BaseProfileID, "base-profile-id", 0, "")
	fs.Int64Var(&b.ProfileIDRange, "profile-id-range", 100000, "")
	fs.Int64Var(&b.Seed, "seed", 1, "")
	fs.IntVar(&b.Iterations, "iterations", 100, "")
	fs.StringVar(&b.DB, "db", "benchdb", "")
	fs.Float64Var(&b.BitmapExponent, "bitmap-exponent", 1.01, "")
	fs.Float64Var(&b.BitmapOffset, "bitmap-offset", 1, "")
	fs.Float64Var(&b.ProfileExponent, "profile-exponent", 1.01, "")
	fs.Float64Var(&b.ProfileOffset, "profile-offset", 1, "")
	fs.StringVar(&b.ClientType, "client-type", "single", "")

	if err := fs.Parse(args); err != nil {
		return nil, err
	}
	return fs.Args(), nil
}

// Run runs the ZipfSetBits benchmark
func (b *ZipfSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
	rnd := rand.New(rand.NewSource(b.Seed + int64(agentNum)))
	bitmapRng := rand.NewZipf(rnd, b.BitmapExponent, b.BitmapOffset, uint64(b.BitmapIDRange))
	profileRng := rand.NewZipf(rnd, b.ProfileExponent, b.ProfileOffset, uint64(b.ProfileIDRange))
	bitmapPerm := NewPermutationGenerator(b.BitmapIDRange, b.Seed)
	profilePerm := NewPermutationGenerator(b.ProfileIDRange, b.Seed)

	results := make(map[string]interface{})
	if b.cli == nil {
		results["error"] = fmt.Errorf("No client set for ZipfSetBits agent: %v", agentNum)
		return results
	}
	s := NewStats()
	var start time.Time
	for n := 0; n < b.Iterations; n++ {
		// generate IDs from Zipf distribution
		bitmapIDOriginal := bitmapRng.Uint64()
		profIDOriginal := profileRng.Uint64()
		// permute IDs randomly, but repeatably
		bitmapID := bitmapPerm.Next(int64(bitmapIDOriginal))
		profID := profilePerm.Next(int64(profIDOriginal))

		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+int64(bitmapID), b.BaseProfileID+int64(profID))
		start = time.Now()
		b.cli.ExecuteQuery(ctx, b.DB, query, true)
		s.Add(time.Now().Sub(start))
	}
	AddToResults(s, results)
	return results
}
<commit_msg>Put tabs in usage string<commit_after>package bench

import (
	"fmt"

	"flag"
	"io/ioutil"

	"context"
	"math/rand"
	"time"
)

// ZipfSetBits sets bits randomly and deterministically based on a seed, according to the Zipf distribution
type ZipfSetBits struct {
	HasClient
	BaseBitmapID    int64
	BaseProfileID   int64
	BitmapIDRange   int64
	ProfileIDRange  int64
	Iterations      int // number of bits that will be set
	Seed            int64
	BitmapExponent  float64
	BitmapOffset    float64
	ProfileExponent float64
	ProfileOffset   float64
	DB              string // DB to use in pilosa.

}

func (b *ZipfSetBits) Usage() string {
	return `
zipf-set-bits sets random bits according to Zipf distribution

Usage: zipf-set-bits [arguments]

The following arguments are available:

	-base-bitmap-id int
		bits being set will all be greater than BaseBitmapID

	-bitmap-id-range int
		number of possible bitmap ids that can be set

	-base-profile-id int
		profile id num to start from

	-profile-id-range int
		number of possible profile ids that can be set

	-iterations int
		number of bits to set

	-seed int
		Seed for RNG

	-db string
		pilosa db to use

	-bitmap-exponent float64
		zipf exponent parameter for bitmap IDs

	-bitmap-ratio float64
		zipf probability ratio parameter for bitmap IDs

	-profile-exponent float64
		zipf exponent parameter for profile IDs

	-profile-ratio float64
		zipf probability ratio parameter for profile IDs

	-client-type string
		Can be 'single' (all agents hitting one host) or 'round_robin'
`[1:]
}

func (b *ZipfSetBits) ConsumeFlags(args []string) ([]string, error) {
	fs := flag.NewFlagSet("ZipfSetBits", flag.ContinueOnError)
	fs.SetOutput(ioutil.Discard)
	fs.Int64Var(&b.BaseBitmapID, "base-bitmap-id", 0, "")
	fs.Int64Var(&b.BitmapIDRange, "bitmap-id-range", 100000, "")
	fs.Int64Var(&b.BaseProfileID, "base-profile-id", 0, "")
	fs.Int64Var(&b.ProfileIDRange, "profile-id-range", 100000, "")
	fs.Int64Var(&b.Seed, "seed", 1, "")
	fs.IntVar(&b.Iterations, "iterations", 100, "")
	fs.StringVar(&b.DB, "db", "benchdb", "")
	fs.Float64Var(&b.BitmapExponent, "bitmap-exponent", 1.01, "")
	fs.Float64Var(&b.BitmapOffset, "bitmap-offset", 1, "")
	fs.Float64Var(&b.ProfileExponent, "profile-exponent", 1.01, "")
	fs.Float64Var(&b.ProfileOffset, "profile-offset", 1, "")
	fs.StringVar(&b.ClientType, "client-type", "single", "")

	if err := fs.Parse(args); err != nil {
		return nil, err
	}
	return fs.Args(), nil
}

// Run runs the ZipfSetBits benchmark
func (b *ZipfSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
	rnd := rand.New(rand.NewSource(b.Seed + int64(agentNum)))
	bitmapRng := rand.NewZipf(rnd, b.BitmapExponent, b.BitmapOffset, uint64(b.BitmapIDRange))
	profileRng := rand.NewZipf(rnd, b.ProfileExponent, b.ProfileOffset, uint64(b.ProfileIDRange))
	bitmapPerm := NewPermutationGenerator(b.BitmapIDRange, b.Seed)
	profilePerm := NewPermutationGenerator(b.ProfileIDRange, b.Seed)

	results := make(map[string]interface{})
	if b.cli == nil {
		results["error"] = fmt.Errorf("No client set for ZipfSetBits agent: %v", agentNum)
		return results
	}
	s := NewStats()
	var start time.Time
	for n := 0; n < b.Iterations; n++ {
		// generate IDs from Zipf distribution
		bitmapIDOriginal := bitmapRng.Uint64()
		profIDOriginal := profileRng.Uint64()
		// permute IDs randomly, but repeatably
		bitmapID := bitmapPerm.Next(int64(bitmapIDOriginal))
		profID := profilePerm.Next(int64(profIDOriginal))

		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+int64(bitmapID), b.BaseProfileID+int64(profID))
		start = time.Now()
		b.cli.ExecuteQuery(ctx, b.DB, query, true)
		s.Add(time.Now().Sub(start))
	}
	AddToResults(s, results)
	return results
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var maxCount uint64
	b.Run("read-random-badger", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			var val badger.KVItem
			for pb.Next() {
				key := newKey()
				if err := bdb.Get(key, &val); err != nil {
					if len(val.Value()) != *flagValueSize {
						b.Logf("Value size for key %s - %d != %d", string(key), len(val.Value()), *flagValueSize)
						b.FailNow()
					}
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("badger: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var maxCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if v, err := rdb.Get(key); err == nil {

					if len(v.Data()) != *flagValueSize {
						b.Logf("Value size %d != %d", len(v.Data()), *flagValueSize)
						b.FailNow()
					}
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("rocks: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var maxCount uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					v, err := txn.Get(lmdbDBI, key)
					if err != nil {
						return err
					}
					if len(v) != *flagValueSize {
						b.Logf("Value size %d != %d", v, *flagValueSize)
						b.FailNow()
					}
					count++
					return nil
				})
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("lmdb: Maximum no. of keys found: %d", maxCount)
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				if len(v) != *flagValueSize {
					b.Logf("Value size %d != %d", v, *flagValueSize)
					b.FailNow()
				}

				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}
					if len(v1) != *flagValueSize {
						b.Logf("Value size %d != %d", v, *flagValueSize)
						b.FailNow()
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				if len(v) != *flagValueSize {
					b.Logf("Value size %d != %d", v, *flagValueSize)
					b.FailNow()
				}

				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<commit_msg>err == nil, not err != nil<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var maxCount uint64
	b.Run("read-random-badger", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			var val badger.KVItem
			for pb.Next() {
				key := newKey()
				if err := bdb.Get(key, &val); err == nil {
					if len(val.Value()) != *flagValueSize {
						b.Logf("Value size for key %s - %d != %d", string(key), len(val.Value()), *flagValueSize)
						b.FailNow()
					}
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("badger: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var maxCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if v, err := rdb.Get(key); err == nil {

					if len(v.Data()) != *flagValueSize {
						b.Logf("Value size %d != %d", len(v.Data()), *flagValueSize)
						b.FailNow()
					}
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("rocks: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var maxCount uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					v, err := txn.Get(lmdbDBI, key)
					if err != nil {
						return err
					}
					if len(v) != *flagValueSize {
						b.Logf("Value size %d != %d", v, *flagValueSize)
						b.FailNow()
					}
					count++
					return nil
				})
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("lmdb: Maximum no. of keys found: %d", maxCount)
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				if len(v) != *flagValueSize {
					b.Logf("Value size %d != %d", v, *flagValueSize)
					b.FailNow()
				}

				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}
					if len(v1) != *flagValueSize {
						b.Logf("Value size %d != %d", v, *flagValueSize)
						b.FailNow()
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				if len(v) != *flagValueSize {
					b.Logf("Value size %d != %d", v, *flagValueSize)
					b.FailNow()
				}

				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package gorpc

import (
	"bytes"
	"fmt"
	"math/rand"
	"runtime"
	"sync"
	"testing"
	"time"
)

func BenchmarkRealApp1Worker(b *testing.B) {
	simulateRealApp(b, 1)
}

func BenchmarkRealApp10Workers(b *testing.B) {
	simulateRealApp(b, 10)
}

func BenchmarkRealApp100Workers(b *testing.B) {
	simulateRealApp(b, 100)
}

func BenchmarkRealApp1000Workers(b *testing.B) {
	simulateRealApp(b, 1000)
}

func BenchmarkRealApp10000Workers(b *testing.B) {
	simulateRealApp(b, 10000)
}

func BenchmarkRealApp100000Workers(b *testing.B) {
	simulateRealApp(b, 100000)
}

func doRealWork() {
	time.Sleep(time.Duration(5+rand.Intn(10)) * time.Millisecond)
}

func simulateRealApp(b *testing.B, workersCount int) {
	addr := ":43245"
	s := NewTCPServer(addr, func(clientAddr string, request interface{}) interface{} {
		doRealWork()
		return request
	})
	s.PendingResponses = workersCount
	s.Start()
	defer s.Stop()

	c := NewTCPClient(addr)
	c.PendingRequests = workersCount
	c.Start()
	defer c.Stop()

	ch := make(chan int, b.N)
	for i := 0; i < b.N; i++ {
		ch <- i
	}

	type realMessage struct {
		FooString string
		FooInt    int
		Map       map[string]string
		Bytes     []byte
	}
	RegisterType(&realMessage{})

	b.ResetTimer()
	var wg sync.WaitGroup
	for i := 0; i < workersCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for n := range ch {
				doRealWork()
				req := &realMessage{
					FooString: fmt.Sprintf("aaa bbb ccc %d", n),
					FooInt:    n,
					Map: map[string]string{
						"aaa": "bbb",
						"xxx": fmt.Sprintf("cvv%d", n+4),
					},
					Bytes: []byte(",.zssdsdkl23 23"),
				}
				resp, err := c.Call(req)
				if err != nil {
					b.Fatalf("Unexpected response error: [%s]", err)
				}
				x, ok := resp.(*realMessage)
				if !ok {
					b.Fatalf("Unexpected response type %T", resp)
				}
				if x.FooString != req.FooString || x.FooInt != req.FooInt || x.Map["aaa"] != req.Map["aaa"] || x.Map["xxx"] != req.Map["xxx"] || !bytes.Equal(x.Bytes, req.Bytes) {
					b.Fatalf("Unexpected response value %v. Expected %v", x, req)
				}
			}
		}()
	}
	close(ch)
	wg.Wait()
}

func BenchmarkEchoInt1Worker(b *testing.B) {
	benchEchoInt(b, 1, false, false)
}

func BenchmarkEchoInt10Workers(b *testing.B) {
	benchEchoInt(b, 10, false, false)
}

func BenchmarkEchoInt100Workers(b *testing.B) {
	benchEchoInt(b, 100, false, false)
}

func BenchmarkEchoInt1000Workers(b *testing.B) {
	benchEchoInt(b, 1000, false, false)
}

func BenchmarkEchoInt10000Workers(b *testing.B) {
	benchEchoInt(b, 10000, false, false)
}

func BenchmarkEchoIntUnix1Worker(b *testing.B) {
	benchEchoInt(b, 1, false, true)
}

func BenchmarkEchoIntUnix10Workers(b *testing.B) {
	benchEchoInt(b, 10, false, true)
}

func BenchmarkEchoIntUnix100Workers(b *testing.B) {
	benchEchoInt(b, 100, false, true)
}

func BenchmarkEchoIntUnix1000Workers(b *testing.B) {
	benchEchoInt(b, 1000, false, true)
}

func BenchmarkEchoIntUnix10000Workers(b *testing.B) {
	benchEchoInt(b, 10000, false, true)
}

func BenchmarkEchoIntNocompress1Worker(b *testing.B) {
	benchEchoInt(b, 1, true, false)
}

func BenchmarkEchoIntNocompress10Workers(b *testing.B) {
	benchEchoInt(b, 10, true, false)
}

func BenchmarkEchoIntNocompress100Workers(b *testing.B) {
	benchEchoInt(b, 100, true, false)
}

func BenchmarkEchoIntNocompress1000Workers(b *testing.B) {
	benchEchoInt(b, 1000, true, false)
}

func BenchmarkEchoIntNocompress10000Workers(b *testing.B) {
	benchEchoInt(b, 10000, true, false)
}

func BenchmarkEchoString1Worker(b *testing.B) {
	benchEchoString(b, 1, false, false)
}

func BenchmarkEchoString10Workers(b *testing.B) {
	benchEchoString(b, 10, false, false)
}

func BenchmarkEchoString100Workers(b *testing.B) {
	benchEchoString(b, 100, false, false)
}

func BenchmarkEchoString1000Workers(b *testing.B) {
	benchEchoString(b, 1000, false, false)
}

func BenchmarkEchoString10000Workers(b *testing.B) {
	benchEchoString(b, 10000, false, false)
}

func BenchmarkEchoStringUnix1Worker(b *testing.B) {
	benchEchoString(b, 1, false, true)
}

func BenchmarkEchoStringUnix10Workers(b *testing.B) {
	benchEchoString(b, 10, false, true)
}

func BenchmarkEchoStringUnix100Workers(b *testing.B) {
	benchEchoString(b, 100, false, true)
}

func BenchmarkEchoStringUnix1000Workers(b *testing.B) {
	benchEchoString(b, 1000, false, true)
}

func BenchmarkEchoStringUnix10000Workers(b *testing.B) {
	benchEchoString(b, 10000, false, true)
}

func BenchmarkEchoStringNocompress1Worker(b *testing.B) {
	benchEchoString(b, 1, true, false)
}

func BenchmarkEchoStringNocompress10Workers(b *testing.B) {
	benchEchoString(b, 10, true, false)
}

func BenchmarkEchoStringNocompress100Workers(b *testing.B) {
	benchEchoString(b, 100, true, false)
}

func BenchmarkEchoStringNocompress1000Workers(b *testing.B) {
	benchEchoString(b, 1000, true, false)
}

func BenchmarkEchoStringNocompress10000Workers(b *testing.B) {
	benchEchoString(b, 10000, true, false)
}

func BenchmarkEchoStruct1Worker(b *testing.B) {
	benchEchoStruct(b, 1, false, false)
}

func BenchmarkEchoStruct10Workers(b *testing.B) {
	benchEchoStruct(b, 10, false, false)
}

func BenchmarkEchoStruct100Workers(b *testing.B) {
	benchEchoStruct(b, 100, false, false)
}

func BenchmarkEchoStruct1000Workers(b *testing.B) {
	benchEchoStruct(b, 1000, false, false)
}

func BenchmarkEchoStruct10000Workers(b *testing.B) {
	benchEchoStruct(b, 10000, false, false)
}

func BenchmarkEchoStructUnix1Worker(b *testing.B) {
	benchEchoStruct(b, 1, false, true)
}

func BenchmarkEchoStructUnix10Workers(b *testing.B) {
	benchEchoStruct(b, 10, false, true)
}

func BenchmarkEchoStructUnix100Workers(b *testing.B) {
	benchEchoStruct(b, 100, false, true)
}

func BenchmarkEchoStructUnix1000Workers(b *testing.B) {
	benchEchoStruct(b, 1000, false, true)
}

func BenchmarkEchoStructUnix10000Workers(b *testing.B) {
	benchEchoStruct(b, 10000, false, true)
}

func BenchmarkEchoStructNocompress1Worker(b *testing.B) {
	benchEchoStruct(b, 1, true, false)
}

func BenchmarkEchoStructNocompress10Workers(b *testing.B) {
	benchEchoStruct(b, 10, true, false)
}

func BenchmarkEchoStructNocompress100Workers(b *testing.B) {
	benchEchoStruct(b, 100, true, false)
}

func BenchmarkEchoStructNocompress1000Workers(b *testing.B) {
	benchEchoStruct(b, 1000, true, false)
}

func BenchmarkEchoStructNocompress10000Workers(b *testing.B) {
	benchEchoStruct(b, 10000, true, false)
}

func benchEchoInt(b *testing.B, workers int, disableCompression, isUnixTransport bool) {
	benchEchoFunc(b, workers, disableCompression, isUnixTransport, func(c *Client, n int) {
		resp, err := c.Call(n)
		if err != nil {
			b.Fatalf("Unexpected error: [%s]", err)
		}
		if resp == nil {
			b.Fatalf("Unexpected nil response")
		}
		x, ok := resp.(int)
		if !ok {
			b.Fatalf("Unexpected response type: %T. Expected int", resp)
		}
		if x != n {
			b.Fatalf("Unexpected value returned: %d. Expected %d", x, n)
		}
	})
}

func benchEchoString(b *testing.B, workers int, disableCompression, isUnixTransport bool) {
	benchEchoFunc(b, workers, disableCompression, isUnixTransport, func(c *Client, n int) {
		s := fmt.Sprintf("test string %d", n)
		resp, err := c.Call(s)
		if err != nil {
			b.Fatalf("Unexpected error: [%s]", err)
		}
		if resp == nil {
			b.Fatalf("Unexpected nil response")
		}
		x, ok := resp.(string)
		if !ok {
			b.Fatalf("Unexpected response type: %T. Expected string", resp)
		}
		if x != s {
			b.Fatalf("Unexpected value returned: %s. Expected %s", x, s)
		}
	})
}

func benchEchoStruct(b *testing.B, workers int, disableCompression, isUnixTransport bool) {
	type BenchEchoStruct struct {
		A int
		B string
		C []byte
	}

	RegisterType(&BenchEchoStruct{})

	benchEchoFunc(b, workers, disableCompression, isUnixTransport, func(c *Client, n int) {
		s := &BenchEchoStruct{
			A: n,
			B: fmt.Sprintf("test string %d", n),
			C: []byte(fmt.Sprintf("test bytes %d", n)),
		}
		resp, err := c.Call(s)
		if err != nil {
			b.Fatalf("Unexpected error: [%s]", err)
		}
		if resp == nil {
			b.Fatalf("Unexpected nil response")
		}
		x, ok := resp.(*BenchEchoStruct)
		if !ok {
			b.Fatalf("Unexpected response type: %T. Expected BenchEchoStruct", resp)
		}
		if x.A != s.A || x.B != s.B || !bytes.Equal(x.C, s.C) {
			b.Fatalf("Unexpected value returned: %+v. Expected %+v", x, s)
		}
	})
}

func benchEchoFunc(b *testing.B, workers int, disableCompression, isUnixTransport bool, f func(*Client, int)) {
	s, c := createEchoServerAndClient(b, disableCompression, workers, isUnixTransport)
	defer s.Stop()
	defer c.Stop()

	var wg sync.WaitGroup

	ch := make(chan int, 100000)

	for i := 0; i < workers; i++ {
		wg.Add(1)

		go func() {
			defer wg.Done()
			for n := range ch {
				f(c, n)
			}
		}()
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ch <- i
	}

	close(ch)

	wg.Wait()
}

func createEchoServerAndClient(b *testing.B, disableCompression bool, pendingMessages int, isUnixTransport bool) (s *Server, c *Client) {
	if isUnixTransport {
		addr := "./gorpc-bench.sock"
		s = NewUnixServer(addr, echoHandler)
		c = NewUnixClient(addr)
	} else {
		addr := fmt.Sprintf(":%d", rand.Intn(20000)+10000)
		s = &Server{
			Addr:    addr,
			Handler: echoHandler,
		}
		c = &Client{
			Addr: addr,
		}
	}

	if err := s.Start(); err != nil {
		b.Fatalf("Cannot start gorpc server: [%s]", err)
	}

	c.Conns = runtime.GOMAXPROCS(-1)
	c.DisableCompression = disableCompression
	c.PendingRequests = pendingMessages
	c.Start()

	return s, c
}
<commit_msg>Make the number of connections in 'real app' benchmark equivalent to GOMAXPROCS. This improves its' performance on multi-CPU systems with GOMAXPROCS=NumCPU<commit_after>package gorpc

import (
	"bytes"
	"fmt"
	"math/rand"
	"runtime"
	"sync"
	"testing"
	"time"
)

func BenchmarkRealApp1Worker(b *testing.B) {
	simulateRealApp(b, 1)
}

func BenchmarkRealApp10Workers(b *testing.B) {
	simulateRealApp(b, 10)
}

func BenchmarkRealApp100Workers(b *testing.B) {
	simulateRealApp(b, 100)
}

func BenchmarkRealApp1000Workers(b *testing.B) {
	simulateRealApp(b, 1000)
}

func BenchmarkRealApp10000Workers(b *testing.B) {
	simulateRealApp(b, 10000)
}

func BenchmarkRealApp100000Workers(b *testing.B) {
	simulateRealApp(b, 100000)
}

func doRealWork() {
	time.Sleep(time.Duration(5+rand.Intn(10)) * time.Millisecond)
}

func simulateRealApp(b *testing.B, workersCount int) {
	addr := ":43245"
	s := NewTCPServer(addr, func(clientAddr string, request interface{}) interface{} {
		doRealWork()
		return request
	})
	s.PendingResponses = workersCount
	s.Start()
	defer s.Stop()

	c := NewTCPClient(addr)
	c.Conns = runtime.GOMAXPROCS(-1)
	c.PendingRequests = workersCount
	c.Start()
	defer c.Stop()

	ch := make(chan int, b.N)
	for i := 0; i < b.N; i++ {
		ch <- i
	}

	type realMessage struct {
		FooString string
		FooInt    int
		Map       map[string]string
		Bytes     []byte
	}
	RegisterType(&realMessage{})

	b.ResetTimer()
	var wg sync.WaitGroup
	for i := 0; i < workersCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for n := range ch {
				doRealWork()
				req := &realMessage{
					FooString: fmt.Sprintf("aaa bbb ccc %d", n),
					FooInt:    n,
					Map: map[string]string{
						"aaa": "bbb",
						"xxx": fmt.Sprintf("cvv%d", n+4),
					},
					Bytes: []byte(",.zssdsdkl23 23"),
				}
				resp, err := c.Call(req)
				if err != nil {
					b.Fatalf("Unexpected response error: [%s]", err)
				}
				x, ok := resp.(*realMessage)
				if !ok {
					b.Fatalf("Unexpected response type %T", resp)
				}
				if x.FooString != req.FooString || x.FooInt != req.FooInt || x.Map["aaa"] != req.Map["aaa"] || x.Map["xxx"] != req.Map["xxx"] || !bytes.Equal(x.Bytes, req.Bytes) {
					b.Fatalf("Unexpected response value %v. Expected %v", x, req)
				}
			}
		}()
	}
	close(ch)
	wg.Wait()
}

func BenchmarkEchoInt1Worker(b *testing.B) {
	benchEchoInt(b, 1, false, false)
}

func BenchmarkEchoInt10Workers(b *testing.B) {
	benchEchoInt(b, 10, false, false)
}

func BenchmarkEchoInt100Workers(b *testing.B) {
	benchEchoInt(b, 100, false, false)
}

func BenchmarkEchoInt1000Workers(b *testing.B) {
	benchEchoInt(b, 1000, false, false)
}

func BenchmarkEchoInt10000Workers(b *testing.B) {
	benchEchoInt(b, 10000, false, false)
}

func BenchmarkEchoIntUnix1Worker(b *testing.B) {
	benchEchoInt(b, 1, false, true)
}

func BenchmarkEchoIntUnix10Workers(b *testing.B) {
	benchEchoInt(b, 10, false, true)
}

func BenchmarkEchoIntUnix100Workers(b *testing.B) {
	benchEchoInt(b, 100, false, true)
}

func BenchmarkEchoIntUnix1000Workers(b *testing.B) {
	benchEchoInt(b, 1000, false, true)
}

func BenchmarkEchoIntUnix10000Workers(b *testing.B) {
	benchEchoInt(b, 10000, false, true)
}

func BenchmarkEchoIntNocompress1Worker(b *testing.B) {
	benchEchoInt(b, 1, true, false)
}

func BenchmarkEchoIntNocompress10Workers(b *testing.B) {
	benchEchoInt(b, 10, true, false)
}

func BenchmarkEchoIntNocompress100Workers(b *testing.B) {
	benchEchoInt(b, 100, true, false)
}

func BenchmarkEchoIntNocompress1000Workers(b *testing.B) {
	benchEchoInt(b, 1000, true, false)
}

func BenchmarkEchoIntNocompress10000Workers(b *testing.B) {
	benchEchoInt(b, 10000, true, false)
}

func BenchmarkEchoString1Worker(b *testing.B) {
	benchEchoString(b, 1, false, false)
}

func BenchmarkEchoString10Workers(b *testing.B) {
	benchEchoString(b, 10, false, false)
}

func BenchmarkEchoString100Workers(b *testing.B) {
	benchEchoString(b, 100, false, false)
}

func BenchmarkEchoString1000Workers(b *testing.B) {
	benchEchoString(b, 1000, false, false)
}

func BenchmarkEchoString10000Workers(b *testing.B) {
	benchEchoString(b, 10000, false, false)
}

func BenchmarkEchoStringUnix1Worker(b *testing.B) {
	benchEchoString(b, 1, false, true)
}

func BenchmarkEchoStringUnix10Workers(b *testing.B) {
	benchEchoString(b, 10, false, true)
}

func BenchmarkEchoStringUnix100Workers(b *testing.B) {
	benchEchoString(b, 100, false, true)
}

func BenchmarkEchoStringUnix1000Workers(b *testing.B) {
	benchEchoString(b, 1000, false, true)
}

func BenchmarkEchoStringUnix10000Workers(b *testing.B) {
	benchEchoString(b, 10000, false, true)
}

func BenchmarkEchoStringNocompress1Worker(b *testing.B) {
	benchEchoString(b, 1, true, false)
}

func BenchmarkEchoStringNocompress10Workers(b *testing.B) {
	benchEchoString(b, 10, true, false)
}

func BenchmarkEchoStringNocompress100Workers(b *testing.B) {
	benchEchoString(b, 100, true, false)
}

func BenchmarkEchoStringNocompress1000Workers(b *testing.B) {
	benchEchoString(b, 1000, true, false)
}

func BenchmarkEchoStringNocompress10000Workers(b *testing.B) {
	benchEchoString(b, 10000, true, false)
}

func BenchmarkEchoStruct1Worker(b *testing.B) {
	benchEchoStruct(b, 1, false, false)
}

func BenchmarkEchoStruct10Workers(b *testing.B) {
	benchEchoStruct(b, 10, false, false)
}

func BenchmarkEchoStruct100Workers(b *testing.B) {
	benchEchoStruct(b, 100, false, false)
}

func BenchmarkEchoStruct1000Workers(b *testing.B) {
	benchEchoStruct(b, 1000, false, false)
}

func BenchmarkEchoStruct10000Workers(b *testing.B) {
	benchEchoStruct(b, 10000, false, false)
}

func BenchmarkEchoStructUnix1Worker(b *testing.B) {
	benchEchoStruct(b, 1, false, true)
}

func BenchmarkEchoStructUnix10Workers(b *testing.B) {
	benchEchoStruct(b, 10, false, true)
}

func BenchmarkEchoStructUnix100Workers(b *testing.B) {
	benchEchoStruct(b, 100, false, true)
}

func BenchmarkEchoStructUnix1000Workers(b *testing.B) {
	benchEchoStruct(b, 1000, false, true)
}

func BenchmarkEchoStructUnix10000Workers(b *testing.B) {
	benchEchoStruct(b, 10000, false, true)
}

func BenchmarkEchoStructNocompress1Worker(b *testing.B) {
	benchEchoStruct(b, 1, true, false)
}

func BenchmarkEchoStructNocompress10Workers(b *testing.B) {
	benchEchoStruct(b, 10, true, false)
}

func BenchmarkEchoStructNocompress100Workers(b *testing.B) {
	benchEchoStruct(b, 100, true, false)
}

func BenchmarkEchoStructNocompress1000Workers(b *testing.B) {
	benchEchoStruct(b, 1000, true, false)
}

func BenchmarkEchoStructNocompress10000Workers(b *testing.B) {
	benchEchoStruct(b, 10000, true, false)
}

func benchEchoInt(b *testing.B, workers int, disableCompression, isUnixTransport bool) {
	benchEchoFunc(b, workers, disableCompression, isUnixTransport, func(c *Client, n int) {
		resp, err := c.Call(n)
		if err != nil {
			b.Fatalf("Unexpected error: [%s]", err)
		}
		if resp == nil {
			b.Fatalf("Unexpected nil response")
		}
		x, ok := resp.(int)
		if !ok {
			b.Fatalf("Unexpected response type: %T. Expected int", resp)
		}
		if x != n {
			b.Fatalf("Unexpected value returned: %d. Expected %d", x, n)
		}
	})
}

func benchEchoString(b *testing.B, workers int, disableCompression, isUnixTransport bool) {
	benchEchoFunc(b, workers, disableCompression, isUnixTransport, func(c *Client, n int) {
		s := fmt.Sprintf("test string %d", n)
		resp, err := c.Call(s)
		if err != nil {
			b.Fatalf("Unexpected error: [%s]", err)
		}
		if resp == nil {
			b.Fatalf("Unexpected nil response")
		}
		x, ok := resp.(string)
		if !ok {
			b.Fatalf("Unexpected response type: %T. Expected string", resp)
		}
		if x != s {
			b.Fatalf("Unexpected value returned: %s. Expected %s", x, s)
		}
	})
}

func benchEchoStruct(b *testing.B, workers int, disableCompression, isUnixTransport bool) {
	type BenchEchoStruct struct {
		A int
		B string
		C []byte
	}

	RegisterType(&BenchEchoStruct{})

	benchEchoFunc(b, workers, disableCompression, isUnixTransport, func(c *Client, n int) {
		s := &BenchEchoStruct{
			A: n,
			B: fmt.Sprintf("test string %d", n),
			C: []byte(fmt.Sprintf("test bytes %d", n)),
		}
		resp, err := c.Call(s)
		if err != nil {
			b.Fatalf("Unexpected error: [%s]", err)
		}
		if resp == nil {
			b.Fatalf("Unexpected nil response")
		}
		x, ok := resp.(*BenchEchoStruct)
		if !ok {
			b.Fatalf("Unexpected response type: %T. Expected BenchEchoStruct", resp)
		}
		if x.A != s.A || x.B != s.B || !bytes.Equal(x.C, s.C) {
			b.Fatalf("Unexpected value returned: %+v. Expected %+v", x, s)
		}
	})
}

func benchEchoFunc(b *testing.B, workers int, disableCompression, isUnixTransport bool, f func(*Client, int)) {
	s, c := createEchoServerAndClient(b, disableCompression, workers, isUnixTransport)
	defer s.Stop()
	defer c.Stop()

	var wg sync.WaitGroup

	ch := make(chan int, 100000)

	for i := 0; i < workers; i++ {
		wg.Add(1)

		go func() {
			defer wg.Done()
			for n := range ch {
				f(c, n)
			}
		}()
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ch <- i
	}

	close(ch)

	wg.Wait()
}

func createEchoServerAndClient(b *testing.B, disableCompression bool, pendingMessages int, isUnixTransport bool) (s *Server, c *Client) {
	if isUnixTransport {
		addr := "./gorpc-bench.sock"
		s = NewUnixServer(addr, echoHandler)
		c = NewUnixClient(addr)
	} else {
		addr := fmt.Sprintf(":%d", rand.Intn(20000)+10000)
		s = &Server{
			Addr:    addr,
			Handler: echoHandler,
		}
		c = &Client{
			Addr: addr,
		}
	}

	if err := s.Start(); err != nil {
		b.Fatalf("Cannot start gorpc server: [%s]", err)
	}

	c.Conns = runtime.GOMAXPROCS(-1)
	c.DisableCompression = disableCompression
	c.PendingRequests = pendingMessages
	c.Start()

	return s, c
}
<|endoftext|>


<commit_before>package ics

import (
	"bytes"
	"io"
	"reflect"
	"testing"
	"time"
)

type typeTest struct {
	Params map[string]string
	Data   string
	Input  interface {
		Decode(map[string]string, string) error
	}
	Match interface {
		Encode(io.Writer)
	}
	Output string
	Error  error
}

func testType(t *testing.T, tests []typeTest) {
	var buf bytes.Buffer
	for n, test := range tests {
		err := test.Input.Decode(test.Params, test.Data)
		if err != test.Error {
			if !reflect.DeepEqual(err, test.Error) {
				t.Errorf("test %d: expecting error %s, got %s", n+1, test.Error, err)
				continue
			}
		}
		if !reflect.DeepEqual(test.Input, test.Match) {
			t.Errorf("test %d: input does not match expected", n+1)
			continue
		}
		test.Match.Encode(&buf)
		if str := buf.String(); str != test.Output {
			t.Errorf("test %d: expecting output string %q, got %q", n+1, test.Output, str)
		}
		buf.Reset()
	}
}

func TestBinary(t *testing.T) {
	testType(t, []typeTest{
		{
			Input: &Binary{},
			Match: &Binary{},
			Error: ErrInvalidEncoding,
		},
		{
			Params: map[string]string{"ENCODING": "BASE64"},
			Data:   "MTIzNDU=",
			Input:  &Binary{},
			Match:  &Binary{'1', '2', '3', '4', '5'},
			Output: "MTIzNDU=",
		},
	})
}

func TestBoolean(t *testing.T) {
	tr := new(Boolean)
	fa := new(Boolean)
	*tr = true
	testType(t, []typeTest{
		{
			Data:   "False",
			Input:  fa,
			Match:  fa,
			Output: "FALSE",
		},
		{
			Data:   "true",
			Input:  tr,
			Match:  tr,
			Output: "TRUE",
		},
		{
			Data:   "HotDog",
			Input:  fa,
			Match:  fa,
			Output: "FALSE",
			Error:  ErrInvalidBoolean,
		},
	})
}

func TestDate(t *testing.T) {
	testType(t, []typeTest{
		{
			Data:   "20011225",
			Input:  &Date{},
			Match:  &Date{time.Date(2001, 12, 25, 0, 0, 0, 0, time.UTC)},
			Output: "20011225",
		},
		{
			Data:   "20081111",
			Input:  &Date{},
			Match:  &Date{time.Date(2008, 11, 11, 0, 0, 0, 0, time.UTC)},
			Output: "20081111",
		},
	})
}

func TestDateTime(t *testing.T) {
	l, err := time.LoadLocation("America/New_York")
	if err != nil {
		t.Errorf("unexpected error: %s", err)
		return
	}
	testType(t, []typeTest{
		{
			Params: map[string]string{},
			Data:   "20011225T131415",
			Input:  &DateTime{},
			Match:  &DateTime{time.Date(2001, 12, 25, 13, 14, 15, 0, time.Local)},
			Output: "20011225T131415",
		},
		{
			Params: map[string]string{},
			Data:   "20011225T131415Z",
			Input:  &DateTime{},
			Match:  &DateTime{time.Date(2001, 12, 25, 13, 14, 15, 0, time.UTC)},
			Output: "20011225T131415Z",
		},
		{
			Params: map[string]string{"TZID": "America/New_York"},
			Data:   "20011225T131415",
			Input:  &DateTime{},
			Match:  &DateTime{time.Date(2001, 12, 25, 13, 14, 15, 0, l)},
			Output: "20011225T131415",
		},
	})
}
<commit_msg>added nil map check in testType<commit_after>package ics

import (
	"bytes"
	"io"
	"reflect"
	"testing"
	"time"
)

type typeTest struct {
	Params map[string]string
	Data   string
	Input  interface {
		Decode(map[string]string, string) error
	}
	Match interface {
		Encode(io.Writer)
	}
	Output string
	Error  error
}

var emptyMap = make(map[string]string)

func testType(t *testing.T, tests []typeTest) {
	var buf bytes.Buffer
	for n, test := range tests {
		if test.Params == nil {
			test.Params = emptyMap
		}
		err := test.Input.Decode(test.Params, test.Data)
		if err != test.Error {
			if !reflect.DeepEqual(err, test.Error) {
				t.Errorf("test %d: expecting error %s, got %s", n+1, test.Error, err)
				continue
			}
		}
		if !reflect.DeepEqual(test.Input, test.Match) {
			t.Errorf("test %d: input does not match expected", n+1)
			continue
		}
		test.Match.Encode(&buf)
		if str := buf.String(); str != test.Output {
			t.Errorf("test %d: expecting output string %q, got %q", n+1, test.Output, str)
		}
		buf.Reset()
	}
}

func TestBinary(t *testing.T) {
	testType(t, []typeTest{
		{
			Input: &Binary{},
			Match: &Binary{},
			Error: ErrInvalidEncoding,
		},
		{
			Params: map[string]string{"ENCODING": "BASE64"},
			Data:   "MTIzNDU=",
			Input:  &Binary{},
			Match:  &Binary{'1', '2', '3', '4', '5'},
			Output: "MTIzNDU=",
		},
	})
}

func TestBoolean(t *testing.T) {
	tr := new(Boolean)
	fa := new(Boolean)
	*tr = true
	testType(t, []typeTest{
		{
			Data:   "False",
			Input:  fa,
			Match:  fa,
			Output: "FALSE",
		},
		{
			Data:   "true",
			Input:  tr,
			Match:  tr,
			Output: "TRUE",
		},
		{
			Data:   "HotDog",
			Input:  fa,
			Match:  fa,
			Output: "FALSE",
			Error:  ErrInvalidBoolean,
		},
	})
}

func TestDate(t *testing.T) {
	testType(t, []typeTest{
		{
			Data:   "20011225",
			Input:  &Date{},
			Match:  &Date{time.Date(2001, 12, 25, 0, 0, 0, 0, time.UTC)},
			Output: "20011225",
		},
		{
			Data:   "20081111",
			Input:  &Date{},
			Match:  &Date{time.Date(2008, 11, 11, 0, 0, 0, 0, time.UTC)},
			Output: "20081111",
		},
	})
}

func TestDateTime(t *testing.T) {
	l, err := time.LoadLocation("America/New_York")
	if err != nil {
		t.Errorf("unexpected error: %s", err)
		return
	}
	testType(t, []typeTest{
		{
			Data:   "20011225T131415",
			Input:  &DateTime{},
			Match:  &DateTime{time.Date(2001, 12, 25, 13, 14, 15, 0, time.Local)},
			Output: "20011225T131415",
		},
		{
			Data:   "20011225T131415Z",
			Input:  &DateTime{},
			Match:  &DateTime{time.Date(2001, 12, 25, 13, 14, 15, 0, time.UTC)},
			Output: "20011225T131415Z",
		},
		{
			Params: map[string]string{"TZID": "America/New_York"},
			Data:   "20011225T131415",
			Input:  &DateTime{},
			Match:  &DateTime{time.Date(2001, 12, 25, 13, 14, 15, 0, l)},
			Output: "20011225T131415",
		},
	})
}
<|endoftext|>