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

package monitor

import (
	"time"

	"github.com/golang/protobuf/proto"
	"go.chromium.org/luci/common/tsmon/distribution"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/types"

	pb "go.chromium.org/luci/common/tsmon/ts_mon_proto"
)

type dataSetKey struct {
	targetHash uint64
	metricName string
}

// SerializeCells creates a MetricsCollection message from a slice of cells.
func SerializeCells(cells []types.Cell, now time.Time) []*pb.MetricsCollection {
	collections := map[uint64]*pb.MetricsCollection{}
	dataSets := map[dataSetKey]*pb.MetricsDataSet{}

	// TODO(1026140): the hash and proto of a Target object should be created
	// at the time of the object creation to avoid unnecessary invocation of
	// Target.Hash() and Target.PopulateProto()
	for _, c := range cells {
		// Find the collection, add it if it doesn't exist.
		targetHash := c.Target.Hash()
		collection, ok := collections[targetHash]
		if !ok {
			collection = &pb.MetricsCollection{}
			collections[targetHash] = collection
			c.Target.PopulateProto(collection)
		}

		// Find the data set, add it if it doesn't exist.
		key := dataSetKey{targetHash, c.Name}
		dataSet, ok := dataSets[key]
		if !ok {
			dataSet = SerializeDataSet(c)
			dataSets[key] = dataSet
			collection.MetricsDataSet = append(collection.MetricsDataSet, dataSet)
		}

		// Add the data to the data set.
		dataSet.Data = append(dataSet.Data, SerializeValue(c, now))
	}

	// Turn the hash into a list and return it.
	ret := make([]*pb.MetricsCollection, 0, len(collections))
	for _, collection := range collections {
		ret = append(ret, collection)
	}
	return ret
}

// SerializeDataSet creates a new MetricsDataSet without any data, but just with
// the metric metadata fields populated.
func SerializeDataSet(c types.Cell) *pb.MetricsDataSet {
	d := pb.MetricsDataSet{}
	d.MetricName = proto.String(metricNamePrefix + c.Name)
	d.FieldDescriptor = field.SerializeDescriptor(c.Fields)
	d.Description = proto.String(c.Description)

	if c.ValueType.IsCumulative() {
		d.StreamKind = pb.StreamKind_CUMULATIVE.Enum()
	} else {
		d.StreamKind = pb.StreamKind_GAUGE.Enum()
	}

	switch c.ValueType {
	case types.NonCumulativeIntType, types.CumulativeIntType:
		d.ValueType = pb.ValueType_INT64.Enum()
	case types.NonCumulativeFloatType, types.CumulativeFloatType:
		d.ValueType = pb.ValueType_DOUBLE.Enum()
	case types.NonCumulativeDistributionType, types.CumulativeDistributionType:
		d.ValueType = pb.ValueType_DISTRIBUTION.Enum()
	case types.StringType:
		d.ValueType = pb.ValueType_STRING.Enum()
	case types.BoolType:
		d.ValueType = pb.ValueType_BOOL.Enum()
	}

	if c.Units.IsSpecified() {
		d.Annotations.Unit = proto.String(string(c.Units))
	}
	return &d
}

func toTimestamp(t time.Time) *pb.Timestamp {
	return &pb.Timestamp{
		Seconds: proto.Int64(t.Unix()),
		Nanos:   proto.Int32(int32(t.UnixNano() % int64(time.Second))),
	}
}

// SerializeValue creates a new MetricsData representing this cell's value.
func SerializeValue(c types.Cell, now time.Time) *pb.MetricsData {
	d := pb.MetricsData{}
	d.Field = field.Serialize(c.Fields, c.FieldVals)

	if c.ValueType.IsCumulative() {
		d.StartTimestamp = toTimestamp(c.ResetTime)
	} else {
		d.StartTimestamp = toTimestamp(now)
	}
	d.EndTimestamp = toTimestamp(now)

	switch c.ValueType {
	case types.NonCumulativeIntType, types.CumulativeIntType:
		d.Value = &pb.MetricsData_Int64Value{c.Value.(int64)}
	case types.NonCumulativeFloatType, types.CumulativeFloatType:
		d.Value = &pb.MetricsData_DoubleValue{c.Value.(float64)}
	case types.CumulativeDistributionType, types.NonCumulativeDistributionType:
		d.Value = &pb.MetricsData_DistributionValue{serializeDistribution(c.Value.(*distribution.Distribution))}
	case types.StringType:
		d.Value = &pb.MetricsData_StringValue{c.Value.(string)}
	case types.BoolType:
		d.Value = &pb.MetricsData_BoolValue{c.Value.(bool)}
	}
	return &d
}

func serializeDistribution(d *distribution.Distribution) *pb.MetricsData_Distribution {
	ret := pb.MetricsData_Distribution{
		Count: proto.Int64(d.Count()),
	}

	if d.Count() > 0 {
		ret.Mean = proto.Float64(d.Sum() / float64(d.Count()))
	}

	// Copy the bucketer params.
	if d.Bucketer().Width() == 0 {
		ret.BucketOptions = &pb.MetricsData_Distribution_ExponentialBuckets{
			&pb.MetricsData_Distribution_ExponentialOptions{
				NumFiniteBuckets: proto.Int32(int32(d.Bucketer().NumFiniteBuckets())),
				GrowthFactor:     proto.Float64(d.Bucketer().GrowthFactor()),
				Scale:            proto.Float64(1.0),
			},
		}
	} else {
		ret.BucketOptions = &pb.MetricsData_Distribution_LinearBuckets{
			&pb.MetricsData_Distribution_LinearOptions{
				NumFiniteBuckets: proto.Int32(int32(d.Bucketer().NumFiniteBuckets())),
				Width:            proto.Float64(d.Bucketer().Width()),
				Offset:           proto.Float64(0.0),
			},
		}
	}

	// Copy the distribution bucket values.  Include the overflow buckets on
	// either end.
	ret.BucketCount = d.Buckets()

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

package monitor

import (
	"time"

	"github.com/golang/protobuf/proto"
	"go.chromium.org/luci/common/tsmon/distribution"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/types"

	pb "go.chromium.org/luci/common/tsmon/ts_mon_proto"
)

type dataSetKey struct {
	targetHash uint64
	metricName string
}

// SerializeCells creates a MetricsCollection message from a slice of cells.
func SerializeCells(cells []types.Cell, now time.Time) []*pb.MetricsCollection {
	collections := map[uint64]*pb.MetricsCollection{}
	dataSets := map[dataSetKey]*pb.MetricsDataSet{}

	// TODO(1026140): the hash and proto of a Target object should be created
	// at the time of the object creation to avoid unnecessary invocation of
	// Target.Hash() and Target.PopulateProto()
	for _, c := range cells {
		// Find the collection, add it if it doesn't exist.
		targetHash := c.Target.Hash()
		collection, ok := collections[targetHash]
		if !ok {
			collection = &pb.MetricsCollection{}
			collections[targetHash] = collection
			c.Target.PopulateProto(collection)
		}

		// Find the data set, add it if it doesn't exist.
		key := dataSetKey{targetHash, c.Name}
		dataSet, ok := dataSets[key]
		if !ok {
			dataSet = SerializeDataSet(c)
			dataSets[key] = dataSet
			collection.MetricsDataSet = append(collection.MetricsDataSet, dataSet)
		}

		// Add the data to the data set.
		dataSet.Data = append(dataSet.Data, SerializeValue(c, now))
	}

	// Turn the hash into a list and return it.
	ret := make([]*pb.MetricsCollection, 0, len(collections))
	for _, collection := range collections {
		ret = append(ret, collection)
	}
	return ret
}

// SerializeDataSet creates a new MetricsDataSet without any data, but just with
// the metric metadata fields populated.
func SerializeDataSet(c types.Cell) *pb.MetricsDataSet {
	d := pb.MetricsDataSet{}
	d.MetricName = proto.String(metricNamePrefix + c.Name)
	d.FieldDescriptor = field.SerializeDescriptor(c.Fields)
	d.Description = proto.String(c.Description)

	if c.ValueType.IsCumulative() {
		d.StreamKind = pb.StreamKind_CUMULATIVE.Enum()
	} else {
		d.StreamKind = pb.StreamKind_GAUGE.Enum()
	}

	switch c.ValueType {
	case types.NonCumulativeIntType, types.CumulativeIntType:
		d.ValueType = pb.ValueType_INT64.Enum()
	case types.NonCumulativeFloatType, types.CumulativeFloatType:
		d.ValueType = pb.ValueType_DOUBLE.Enum()
	case types.NonCumulativeDistributionType, types.CumulativeDistributionType:
		d.ValueType = pb.ValueType_DISTRIBUTION.Enum()
	case types.StringType:
		d.ValueType = pb.ValueType_STRING.Enum()
	case types.BoolType:
		d.ValueType = pb.ValueType_BOOL.Enum()
	}

	if c.Units.IsSpecified() {
		d.Annotations.Unit = proto.String(string(c.Units))
	}
	return &d
}

func toTimestamp(t time.Time) *pb.Timestamp {
	return &pb.Timestamp{
		Seconds: proto.Int64(t.Unix()),
		Nanos:   proto.Int32(int32(t.Nanosecond())),
	}
}

// SerializeValue creates a new MetricsData representing this cell's value.
func SerializeValue(c types.Cell, now time.Time) *pb.MetricsData {
	d := pb.MetricsData{}
	d.Field = field.Serialize(c.Fields, c.FieldVals)

	if c.ValueType.IsCumulative() {
		d.StartTimestamp = toTimestamp(c.ResetTime)
	} else {
		d.StartTimestamp = toTimestamp(now)
	}
	d.EndTimestamp = toTimestamp(now)

	switch c.ValueType {
	case types.NonCumulativeIntType, types.CumulativeIntType:
		d.Value = &pb.MetricsData_Int64Value{c.Value.(int64)}
	case types.NonCumulativeFloatType, types.CumulativeFloatType:
		d.Value = &pb.MetricsData_DoubleValue{c.Value.(float64)}
	case types.CumulativeDistributionType, types.NonCumulativeDistributionType:
		d.Value = &pb.MetricsData_DistributionValue{serializeDistribution(c.Value.(*distribution.Distribution))}
	case types.StringType:
		d.Value = &pb.MetricsData_StringValue{c.Value.(string)}
	case types.BoolType:
		d.Value = &pb.MetricsData_BoolValue{c.Value.(bool)}
	}
	return &d
}

func serializeDistribution(d *distribution.Distribution) *pb.MetricsData_Distribution {
	ret := pb.MetricsData_Distribution{
		Count: proto.Int64(d.Count()),
	}

	if d.Count() > 0 {
		ret.Mean = proto.Float64(d.Sum() / float64(d.Count()))
	}

	// Copy the bucketer params.
	if d.Bucketer().Width() == 0 {
		ret.BucketOptions = &pb.MetricsData_Distribution_ExponentialBuckets{
			&pb.MetricsData_Distribution_ExponentialOptions{
				NumFiniteBuckets: proto.Int32(int32(d.Bucketer().NumFiniteBuckets())),
				GrowthFactor:     proto.Float64(d.Bucketer().GrowthFactor()),
				Scale:            proto.Float64(1.0),
			},
		}
	} else {
		ret.BucketOptions = &pb.MetricsData_Distribution_LinearBuckets{
			&pb.MetricsData_Distribution_LinearOptions{
				NumFiniteBuckets: proto.Int32(int32(d.Bucketer().NumFiniteBuckets())),
				Width:            proto.Float64(d.Bucketer().Width()),
				Offset:           proto.Float64(0.0),
			},
		}
	}

	// Copy the distribution bucket values.  Include the overflow buckets on
	// either end.
	ret.BucketCount = d.Buckets()

	return &ret
}
<|endoftext|>


<commit_before>package provision

import (
	"bufio"
	"fmt"
	"net/http"
	"strings"

	"github.com/docker/machine/commands/mcndirs"
	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnutils"
	"github.com/docker/machine/libmachine/provision/pkgaction"
	"github.com/docker/machine/libmachine/provision/serviceaction"
	"github.com/docker/machine/libmachine/state"
	"github.com/docker/machine/libmachine/swarm"
)

const (
	versionsURL  = "http://releases.rancher.com/os/versions.yml"
	isoURL       = "https://github.com/rancherio/os/releases/download/%s/machine-rancheros.iso"
	hostnameTmpl = `sudo mkdir -p /var/lib/rancher/conf/cloud-config.d/
sudo tee /var/lib/rancher/conf/cloud-config.d/machine-hostname.yml << EOF
#cloud-config

hostname: %s
EOF
`
)

func init() {
	Register("RancherOS", &RegisteredProvisioner{
		New: NewRancherProvisioner,
	})
}

func NewRancherProvisioner(d drivers.Driver) Provisioner {
	return &RancherProvisioner{
		GenericProvisioner{
			SSHCommander:      GenericSSHCommander{Driver: d},
			DockerOptionsDir:  "/var/lib/rancher/conf",
			DaemonOptionsFile: "/var/lib/rancher/conf/docker",
			OsReleaseID:       "rancheros",
			Driver:            d,
		},
	}
}

type RancherProvisioner struct {
	GenericProvisioner
}

func (provisioner *RancherProvisioner) String() string {
	return "rancheros"
}

func (provisioner *RancherProvisioner) Service(name string, action serviceaction.ServiceAction) error {
	command := fmt.Sprintf("sudo system-docker %s %s", action.String(), name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) Package(name string, action pkgaction.PackageAction) error {
	var packageAction string

	if name == "docker" && action == pkgaction.Upgrade {
		return provisioner.upgrade()
	}

	switch action {
	case pkgaction.Install:
		packageAction = "enabled"
	case pkgaction.Remove:
		packageAction = "disable"
	case pkgaction.Upgrade:
		// TODO: support upgrade
		packageAction = "upgrade"
	}

	command := fmt.Sprintf("sudo rancherctl service %s %s", packageAction, name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) Provision(swarmOptions swarm.Options, authOptions auth.Options, engineOptions engine.Options) error {
	log.Debugf("Running RancherOS provisioner on %s", provisioner.Driver.GetMachineName())

	provisioner.SwarmOptions = swarmOptions
	provisioner.AuthOptions = authOptions
	provisioner.EngineOptions = engineOptions
	swarmOptions.Env = engineOptions.Env

	if provisioner.EngineOptions.StorageDriver == "" {
		provisioner.EngineOptions.StorageDriver = "overlay"
	} else if provisioner.EngineOptions.StorageDriver != "overlay" {
		return fmt.Errorf("Unsupported storage driver: %s", provisioner.EngineOptions.StorageDriver)
	}

	log.Debugf("Setting hostname %s", provisioner.Driver.GetMachineName())
	if err := provisioner.SetHostname(provisioner.Driver.GetMachineName()); err != nil {
		return err
	}

	for _, pkg := range provisioner.Packages {
		log.Debugf("Installing package %s", pkg)
		if err := provisioner.Package(pkg, pkgaction.Install); err != nil {
			return err
		}
	}

	log.Debugf("Selecting docker engine: %s", engineOptions.InstallURL)
	if err := selectDocker(provisioner, engineOptions.InstallURL); err != nil {
		return err
	}

	log.Debugf("Preparing certificates")
	provisioner.AuthOptions = setRemoteAuthOptions(provisioner)

	log.Debugf("Setting up certificates")
	if err := ConfigureAuth(provisioner); err != nil {
		return err
	}

	log.Debugf("Configuring swarm")
	if err := configureSwarm(provisioner, swarmOptions, provisioner.AuthOptions); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) SetHostname(hostname string) error {
	// /etc/hosts is bind mounted from Docker, this is hack to that the generic provisioner doesn't try to mv /etc/hosts
	if _, err := provisioner.SSHCommand("sed /127.0.1.1/d /etc/hosts > /tmp/hosts && cat /tmp/hosts | sudo tee /etc/hosts"); err != nil {
		return err
	}

	if err := provisioner.GenericProvisioner.SetHostname(hostname); err != nil {
		return err
	}

	if _, err := provisioner.SSHCommand(fmt.Sprintf(hostnameTmpl, hostname)); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) upgrade() error {
	switch provisioner.Driver.DriverName() {
	case "virtualbox":
		return provisioner.upgradeIso()
	default:
		log.Infof("Running upgrade")
		if _, err := provisioner.SSHCommand("sudo rancherctl os upgrade -f --no-reboot"); err != nil {
			return err
		}

		log.Infof("Upgrade succeeded, rebooting")
		// ignore errors here because the SSH connection will close
		provisioner.SSHCommand("sudo reboot")

		return nil
	}
}

func (provisioner *RancherProvisioner) upgradeIso() error {
	// Largely copied from Boot2Docker provisioner, we should find a way to share this code
	log.Info("Stopping machine to do the upgrade...")

	if err := provisioner.Driver.Stop(); err != nil {
		return err
	}

	if err := mcnutils.WaitFor(drivers.MachineInState(provisioner.Driver, state.Stopped)); err != nil {
		return err
	}

	machineName := provisioner.GetDriver().GetMachineName()

	log.Infof("Upgrading machine %s...", machineName)

	// TODO: Ideally, we should not read from mcndirs directory at all.
	// The driver should be able to communicate how and where to place the
	// relevant files.
	b2dutils := mcnutils.NewB2dUtils(mcndirs.GetBaseDir())

	url, err := provisioner.getLatestISOURL()
	if err != nil {
		return err
	}

	if err := b2dutils.DownloadISOFromURL(url); err != nil {
		return err
	}

	// Copy the latest version of boot2docker ISO to the machine's directory
	if err := b2dutils.CopyIsoToMachineDir("", machineName); err != nil {
		return err
	}

	log.Infof("Starting machine back up...")

	if err := provisioner.Driver.Start(); err != nil {
		return err
	}

	return mcnutils.WaitFor(drivers.MachineInState(provisioner.Driver, state.Running))
}

func (provisioner *RancherProvisioner) getLatestISOURL() (string, error) {
	log.Debugf("Reading %s", versionsURL)
	resp, err := http.Get(versionsURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Don't want to pull in yaml parser, we'll do this manually
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "current: ") {
			log.Debugf("Found %s", line)
			return fmt.Sprintf(isoURL, strings.Split(line, ":")[2]), err
		}
	}

	return "", fmt.Errorf("Failed to find current version")
}

func selectDocker(p Provisioner, baseURL string) error {
	// TODO: detect if its a cloud-init, or a ros setting - and use that..
	if output, err := p.SSHCommand(fmt.Sprintf("wget -O- %s | sh -", baseURL)); err != nil {
		return fmt.Errorf("error selecting docker: (%s) %s", err, output)
	}

	return nil
}
<commit_msg>Skip the default docker provisioning script, RancherOS already has a docker installed<commit_after>package provision

import (
	"bufio"
	"fmt"
	"net/http"
	"strings"

	"github.com/docker/machine/commands/mcndirs"
	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnutils"
	"github.com/docker/machine/libmachine/provision/pkgaction"
	"github.com/docker/machine/libmachine/provision/serviceaction"
	"github.com/docker/machine/libmachine/state"
	"github.com/docker/machine/libmachine/swarm"
)

const (
	versionsURL  = "http://releases.rancher.com/os/versions.yml"
	isoURL       = "https://github.com/rancherio/os/releases/download/%s/machine-rancheros.iso"
	hostnameTmpl = `sudo mkdir -p /var/lib/rancher/conf/cloud-config.d/
sudo tee /var/lib/rancher/conf/cloud-config.d/machine-hostname.yml << EOF
#cloud-config

hostname: %s
EOF
`
)

func init() {
	Register("RancherOS", &RegisteredProvisioner{
		New: NewRancherProvisioner,
	})
}

func NewRancherProvisioner(d drivers.Driver) Provisioner {
	return &RancherProvisioner{
		GenericProvisioner{
			SSHCommander:      GenericSSHCommander{Driver: d},
			DockerOptionsDir:  "/var/lib/rancher/conf",
			DaemonOptionsFile: "/var/lib/rancher/conf/docker",
			OsReleaseID:       "rancheros",
			Driver:            d,
		},
	}
}

type RancherProvisioner struct {
	GenericProvisioner
}

func (provisioner *RancherProvisioner) String() string {
	return "rancheros"
}

func (provisioner *RancherProvisioner) Service(name string, action serviceaction.ServiceAction) error {
	command := fmt.Sprintf("sudo system-docker %s %s", action.String(), name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) Package(name string, action pkgaction.PackageAction) error {
	var packageAction string

	if name == "docker" && action == pkgaction.Upgrade {
		return provisioner.upgrade()
	}

	switch action {
	case pkgaction.Install:
		packageAction = "enabled"
	case pkgaction.Remove:
		packageAction = "disable"
	case pkgaction.Upgrade:
		// TODO: support upgrade
		packageAction = "upgrade"
	}

	command := fmt.Sprintf("sudo rancherctl service %s %s", packageAction, name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) Provision(swarmOptions swarm.Options, authOptions auth.Options, engineOptions engine.Options) error {
	log.Debugf("Running RancherOS provisioner on %s", provisioner.Driver.GetMachineName())

	provisioner.SwarmOptions = swarmOptions
	provisioner.AuthOptions = authOptions
	provisioner.EngineOptions = engineOptions
	swarmOptions.Env = engineOptions.Env

	if provisioner.EngineOptions.StorageDriver == "" {
		provisioner.EngineOptions.StorageDriver = "overlay"
	} else if provisioner.EngineOptions.StorageDriver != "overlay" {
		return fmt.Errorf("Unsupported storage driver: %s", provisioner.EngineOptions.StorageDriver)
	}

	log.Debugf("Setting hostname %s", provisioner.Driver.GetMachineName())
	if err := provisioner.SetHostname(provisioner.Driver.GetMachineName()); err != nil {
		return err
	}

	for _, pkg := range provisioner.Packages {
		log.Debugf("Installing package %s", pkg)
		if err := provisioner.Package(pkg, pkgaction.Install); err != nil {
			return err
		}
	}

	if engineOptions.InstallURL == drivers.DefaultEngineInstallURL {
		log.Debugf("Skipping docker engine default: %s", engineOptions.InstallURL)
	} else {
		log.Debugf("Selecting docker engine: %s", engineOptions.InstallURL)
		if err := selectDocker(provisioner, engineOptions.InstallURL); err != nil {
			return err
		}
	}

	log.Debugf("Preparing certificates")
	provisioner.AuthOptions = setRemoteAuthOptions(provisioner)

	log.Debugf("Setting up certificates")
	if err := ConfigureAuth(provisioner); err != nil {
		return err
	}

	log.Debugf("Configuring swarm")
	if err := configureSwarm(provisioner, swarmOptions, provisioner.AuthOptions); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) SetHostname(hostname string) error {
	// /etc/hosts is bind mounted from Docker, this is hack to that the generic provisioner doesn't try to mv /etc/hosts
	if _, err := provisioner.SSHCommand("sed /127.0.1.1/d /etc/hosts > /tmp/hosts && cat /tmp/hosts | sudo tee /etc/hosts"); err != nil {
		return err
	}

	if err := provisioner.GenericProvisioner.SetHostname(hostname); err != nil {
		return err
	}

	if _, err := provisioner.SSHCommand(fmt.Sprintf(hostnameTmpl, hostname)); err != nil {
		return err
	}

	return nil
}

func (provisioner *RancherProvisioner) upgrade() error {
	switch provisioner.Driver.DriverName() {
	case "virtualbox":
		return provisioner.upgradeIso()
	default:
		log.Infof("Running upgrade")
		if _, err := provisioner.SSHCommand("sudo rancherctl os upgrade -f --no-reboot"); err != nil {
			return err
		}

		log.Infof("Upgrade succeeded, rebooting")
		// ignore errors here because the SSH connection will close
		provisioner.SSHCommand("sudo reboot")

		return nil
	}
}

func (provisioner *RancherProvisioner) upgradeIso() error {
	// Largely copied from Boot2Docker provisioner, we should find a way to share this code
	log.Info("Stopping machine to do the upgrade...")

	if err := provisioner.Driver.Stop(); err != nil {
		return err
	}

	if err := mcnutils.WaitFor(drivers.MachineInState(provisioner.Driver, state.Stopped)); err != nil {
		return err
	}

	machineName := provisioner.GetDriver().GetMachineName()

	log.Infof("Upgrading machine %s...", machineName)

	// TODO: Ideally, we should not read from mcndirs directory at all.
	// The driver should be able to communicate how and where to place the
	// relevant files.
	b2dutils := mcnutils.NewB2dUtils(mcndirs.GetBaseDir())

	url, err := provisioner.getLatestISOURL()
	if err != nil {
		return err
	}

	if err := b2dutils.DownloadISOFromURL(url); err != nil {
		return err
	}

	// Copy the latest version of boot2docker ISO to the machine's directory
	if err := b2dutils.CopyIsoToMachineDir("", machineName); err != nil {
		return err
	}

	log.Infof("Starting machine back up...")

	if err := provisioner.Driver.Start(); err != nil {
		return err
	}

	return mcnutils.WaitFor(drivers.MachineInState(provisioner.Driver, state.Running))
}

func (provisioner *RancherProvisioner) getLatestISOURL() (string, error) {
	log.Debugf("Reading %s", versionsURL)
	resp, err := http.Get(versionsURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Don't want to pull in yaml parser, we'll do this manually
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "current: ") {
			log.Debugf("Found %s", line)
			return fmt.Sprintf(isoURL, strings.Split(line, ":")[2]), err
		}
	}

	return "", fmt.Errorf("Failed to find current version")
}

func selectDocker(p Provisioner, baseURL string) error {
	// TODO: detect if its a cloud-init, or a ros setting - and use that..
	if output, err := p.SSHCommand(fmt.Sprintf("wget -O- %s | sh -", baseURL)); err != nil {
		return fmt.Errorf("error selecting docker: (%s) %s", err, output)
	}

	return nil
}
<|endoftext|>


<commit_before>package maxminddb

import (
	"errors"
	"fmt"
	. "launchpad.net/gocheck"
	"math/big"
	"net"
	"testing"
)

func TestMaxMindDbReader(t *testing.T) { TestingT(t) }

type MySuite struct{}

var _ = Suite(&MySuite{})

func (s *MySuite) TestReader(c *C) {
	for _, recordSize := range []uint{24, 28, 32} {
		for _, ipVersion := range []uint{4, 6} {
			fileName := fmt.Sprintf("test-data/test-data/MaxMind-DB-test-ipv%d-%d.mmdb", ipVersion, recordSize)
			reader, _ := Open(fileName)

			checkMetadata(c, reader, ipVersion, recordSize)

			if ipVersion == 4 {
				checkIpv4(c, reader)
			} else {
				checkIpv6(c, reader)
			}
		}
	}
}

func (s *MySuite) TestDecoder(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")
	recordInterface, _ := reader.Lookup(net.ParseIP("::1.1.1.0"))
	record := recordInterface.(map[string]interface{})

	c.Assert(record["array"], DeepEquals, []interface{}{uint(1), uint(2), uint(3)})
	c.Assert(record["boolean"], Equals, true)
	c.Assert(record["bytes"], DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(record["double"], Equals, 42.123456)
	c.Assert(record["float"], Equals, float32(1.1))
	c.Assert(record["int32"], Equals, -268435456)
	c.Assert(record["map"], DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint(7), uint(8), uint(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(record["uint16"], Equals, uint(100))
	c.Assert(record["uint32"], Equals, uint(268435456))
	c.Assert(record["uint64"], Equals, uint(1152921504606846976))
	c.Assert(record["utf8_string"], Equals, "unicode! ☯ - ♫")
	bigInt := new(big.Int)
	bigInt.SetString("1329227995784915872903807060280344576", 10)
	c.Assert(record["uint128"], DeepEquals, bigInt)

	type TestType struct {
		Array   []uint
		Boolean bool
		Bytes   []byte
		Double  float64
		Float   float32
		Int32   int32
		Map     map[string]interface{}
		Uint16  uint16
		Uint32  uint32
		Uint64  uint64
		// Uint128 big.Int
		Utf8String string `maxminddb:"utf8_string"`
	}

	var result TestType
	err := reader.Unmarshal(net.ParseIP("::1.1.1.0"), &result)
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(result.Array, DeepEquals, []uint{uint(1), uint(2), uint(3)})
	c.Assert(result.Boolean, Equals, true)
	c.Assert(result.Bytes, DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(result.Double, Equals, 42.123456)
	c.Assert(result.Float, Equals, float32(1.1))
	c.Assert(result.Int32, Equals, int32(-268435456))
	c.Assert(result.Map, DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint(7), uint(8), uint(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(result.Uint16, Equals, uint16(100))
	c.Assert(result.Uint32, Equals, uint32(268435456))
	c.Assert(result.Uint64, Equals, uint64(1152921504606846976))
	c.Assert(result.Utf8String, Equals, "unicode! ☯ - ♫")
	// c.Assert(result.Uint128, DeepEquals, bigInt)

	reader.Close()
}

func (s *MySuite) TestIpv6inIpv4(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-ipv4-24.mmdb")
	record, err := reader.Lookup(net.ParseIP("2001::"))
	if record != nil {
		c.Log("nil record from lookup expected")
		c.Fail()
	}
	expected := errors.New("error looking up '2001::': you attempted to look up an IPv6 address in an IPv4-only database")
	c.Assert(err, DeepEquals, expected)
	reader.Close()

}

func (s *MySuite) TestBrokenDatabase(c *C) {
	c.Skip("NO TYPE DECODING VALIDATION DONE")
	reader, _ := Open("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb")
	// // Should return an error like: "The MaxMind DB file's data "
	// //                              "section contains bad data (unknown data "
	// //                              "type or corrupt data)"
	reader.Lookup(net.ParseIP("2001:220::"))
	reader.Close()

}

func (s *MySuite) TestMissingDatabase(c *C) {
	reader, err := Open("file-does-not-exist.mmdb")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "open file-does-not-exist.mmdb: no such file or directory")
}

func (s *MySuite) TestNonDatabase(c *C) {
	reader, err := Open("README.md")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "error opening database file (README.md): invalid MaxMind DB file")
}

func checkMetadata(c *C, reader *Reader, ipVersion uint, recordSize uint) {
	metadata := reader.Metadata

	c.Assert(metadata.BinaryFormatMajorVersion, Equals, uint(2))

	c.Assert(metadata.BinaryFormatMinorVersion, Equals, uint(0))
	c.Assert(metadata.BuildEpoch, FitsTypeOf, uint(0))
	c.Assert(metadata.DatabaseType, Equals, "Test")

	c.Assert(metadata.Description, DeepEquals,
		map[string]string{
			"en": "Test Database",
			"zh": "Test Database Chinese",
		})
	c.Assert(metadata.IPVersion, Equals, ipVersion)
	c.Assert(metadata.Languages, DeepEquals, []string{"en", "zh"})

	if ipVersion == 4 {
		c.Assert(metadata.NodeCount, Equals, uint(37))
	} else {
		c.Assert(metadata.NodeCount, Equals, uint(188))
	}

	c.Assert(metadata.RecordSize, Equals, recordSize)
}

func checkIpv4(c *C, reader *Reader) {

	for i := uint(0); i < 6; i++ {
		address := fmt.Sprintf("1.1.1.%d", uint(1)<<i)
		ip := net.ParseIP(address)

		// XXX - Figure out why ParseIP always returns 16 byte address.
		// Maybe update reader to accommodated
		record, _ := reader.Lookup(ip[12:])
		c.Assert(record, DeepEquals, map[string]interface{}{
			"ip": address})
	}
	pairs := map[string]string{
		"1.1.1.3":  "1.1.1.2",
		"1.1.1.5":  "1.1.1.4",
		"1.1.1.7":  "1.1.1.4",
		"1.1.1.9":  "1.1.1.8",
		"1.1.1.15": "1.1.1.8",
		"1.1.1.17": "1.1.1.16",
		"1.1.1.31": "1.1.1.16",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]interface{}{"ip": valueAddress}

		ip := net.ParseIP(keyAddress)

		record, _ := reader.Lookup(ip[12:])
		c.Assert(record, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123"} {
		ip := net.ParseIP(address)

		record, _ := reader.Lookup(ip[12:])
		c.Assert(record, IsNil)
	}
}

func checkIpv6(c *C, reader *Reader) {

	subnets := []string{"::1:ffff:ffff", "::2:0:0",
		"::2:0:40", "::2:0:50", "::2:0:58"}

	for _, address := range subnets {
		record, _ := reader.Lookup(net.ParseIP(address))
		c.Assert(record, DeepEquals, map[string]interface{}{"ip": address})
	}

	pairs := map[string]string{
		"::2:0:1":  "::2:0:0",
		"::2:0:33": "::2:0:0",
		"::2:0:39": "::2:0:0",
		"::2:0:41": "::2:0:40",
		"::2:0:49": "::2:0:40",
		"::2:0:52": "::2:0:50",
		"::2:0:57": "::2:0:50",
		"::2:0:59": "::2:0:58",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]interface{}{"ip": valueAddress}
		record, _ := reader.Lookup(net.ParseIP(keyAddress))
		c.Assert(record, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123", "89fa::"} {
		record, _ := reader.Lookup(net.ParseIP(address))
		c.Assert(record, IsNil)
	}
}
<commit_msg>Added basic benchmark<commit_after>package maxminddb

import (
	"errors"
	"fmt"
	. "launchpad.net/gocheck"
	"math/big"
	"math/rand"
	"net"
	"testing"
	"time"
)

func TestMaxMindDbReader(t *testing.T) { TestingT(t) }

type MySuite struct{}

var _ = Suite(&MySuite{})

func (s *MySuite) TestReader(c *C) {
	for _, recordSize := range []uint{24, 28, 32} {
		for _, ipVersion := range []uint{4, 6} {
			fileName := fmt.Sprintf("test-data/test-data/MaxMind-DB-test-ipv%d-%d.mmdb", ipVersion, recordSize)
			reader, _ := Open(fileName)

			checkMetadata(c, reader, ipVersion, recordSize)

			if ipVersion == 4 {
				checkIpv4(c, reader)
			} else {
				checkIpv6(c, reader)
			}
		}
	}
}

func (s *MySuite) TestDecoder(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")
	recordInterface, _ := reader.Lookup(net.ParseIP("::1.1.1.0"))
	record := recordInterface.(map[string]interface{})

	c.Assert(record["array"], DeepEquals, []interface{}{uint(1), uint(2), uint(3)})
	c.Assert(record["boolean"], Equals, true)
	c.Assert(record["bytes"], DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(record["double"], Equals, 42.123456)
	c.Assert(record["float"], Equals, float32(1.1))
	c.Assert(record["int32"], Equals, -268435456)
	c.Assert(record["map"], DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint(7), uint(8), uint(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(record["uint16"], Equals, uint(100))
	c.Assert(record["uint32"], Equals, uint(268435456))
	c.Assert(record["uint64"], Equals, uint(1152921504606846976))
	c.Assert(record["utf8_string"], Equals, "unicode! ☯ - ♫")
	bigInt := new(big.Int)
	bigInt.SetString("1329227995784915872903807060280344576", 10)
	c.Assert(record["uint128"], DeepEquals, bigInt)

	type TestType struct {
		Array   []uint
		Boolean bool
		Bytes   []byte
		Double  float64
		Float   float32
		Int32   int32
		Map     map[string]interface{}
		Uint16  uint16
		Uint32  uint32
		Uint64  uint64
		// Uint128 big.Int
		Utf8String string `maxminddb:"utf8_string"`
	}

	var result TestType
	err := reader.Unmarshal(net.ParseIP("::1.1.1.0"), &result)
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(result.Array, DeepEquals, []uint{uint(1), uint(2), uint(3)})
	c.Assert(result.Boolean, Equals, true)
	c.Assert(result.Bytes, DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(result.Double, Equals, 42.123456)
	c.Assert(result.Float, Equals, float32(1.1))
	c.Assert(result.Int32, Equals, int32(-268435456))
	c.Assert(result.Map, DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint(7), uint(8), uint(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(result.Uint16, Equals, uint16(100))
	c.Assert(result.Uint32, Equals, uint32(268435456))
	c.Assert(result.Uint64, Equals, uint64(1152921504606846976))
	c.Assert(result.Utf8String, Equals, "unicode! ☯ - ♫")
	// c.Assert(result.Uint128, DeepEquals, bigInt)

	reader.Close()
}

func (s *MySuite) TestIpv6inIpv4(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-ipv4-24.mmdb")
	record, err := reader.Lookup(net.ParseIP("2001::"))
	if record != nil {
		c.Log("nil record from lookup expected")
		c.Fail()
	}
	expected := errors.New("error looking up '2001::': you attempted to look up an IPv6 address in an IPv4-only database")
	c.Assert(err, DeepEquals, expected)
	reader.Close()

}

func (s *MySuite) TestBrokenDatabase(c *C) {
	c.Skip("NO TYPE DECODING VALIDATION DONE")
	reader, _ := Open("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb")
	// // Should return an error like: "The MaxMind DB file's data "
	// //                              "section contains bad data (unknown data "
	// //                              "type or corrupt data)"
	reader.Lookup(net.ParseIP("2001:220::"))
	reader.Close()

}

func (s *MySuite) TestMissingDatabase(c *C) {
	reader, err := Open("file-does-not-exist.mmdb")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "open file-does-not-exist.mmdb: no such file or directory")
}

func (s *MySuite) TestNonDatabase(c *C) {
	reader, err := Open("README.md")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "error opening database file (README.md): invalid MaxMind DB file")
}

func checkMetadata(c *C, reader *Reader, ipVersion uint, recordSize uint) {
	metadata := reader.Metadata

	c.Assert(metadata.BinaryFormatMajorVersion, Equals, uint(2))

	c.Assert(metadata.BinaryFormatMinorVersion, Equals, uint(0))
	c.Assert(metadata.BuildEpoch, FitsTypeOf, uint(0))
	c.Assert(metadata.DatabaseType, Equals, "Test")

	c.Assert(metadata.Description, DeepEquals,
		map[string]string{
			"en": "Test Database",
			"zh": "Test Database Chinese",
		})
	c.Assert(metadata.IPVersion, Equals, ipVersion)
	c.Assert(metadata.Languages, DeepEquals, []string{"en", "zh"})

	if ipVersion == 4 {
		c.Assert(metadata.NodeCount, Equals, uint(37))
	} else {
		c.Assert(metadata.NodeCount, Equals, uint(188))
	}

	c.Assert(metadata.RecordSize, Equals, recordSize)
}

func checkIpv4(c *C, reader *Reader) {

	for i := uint(0); i < 6; i++ {
		address := fmt.Sprintf("1.1.1.%d", uint(1)<<i)
		ip := net.ParseIP(address)

		// XXX - Figure out why ParseIP always returns 16 byte address.
		// Maybe update reader to accommodated
		record, _ := reader.Lookup(ip[12:])
		c.Assert(record, DeepEquals, map[string]interface{}{
			"ip": address})
	}
	pairs := map[string]string{
		"1.1.1.3":  "1.1.1.2",
		"1.1.1.5":  "1.1.1.4",
		"1.1.1.7":  "1.1.1.4",
		"1.1.1.9":  "1.1.1.8",
		"1.1.1.15": "1.1.1.8",
		"1.1.1.17": "1.1.1.16",
		"1.1.1.31": "1.1.1.16",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]interface{}{"ip": valueAddress}

		ip := net.ParseIP(keyAddress)

		record, _ := reader.Lookup(ip[12:])
		c.Assert(record, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123"} {
		ip := net.ParseIP(address)

		record, _ := reader.Lookup(ip[12:])
		c.Assert(record, IsNil)
	}
}

func checkIpv6(c *C, reader *Reader) {

	subnets := []string{"::1:ffff:ffff", "::2:0:0",
		"::2:0:40", "::2:0:50", "::2:0:58"}

	for _, address := range subnets {
		record, _ := reader.Lookup(net.ParseIP(address))
		c.Assert(record, DeepEquals, map[string]interface{}{"ip": address})
	}

	pairs := map[string]string{
		"::2:0:1":  "::2:0:0",
		"::2:0:33": "::2:0:0",
		"::2:0:39": "::2:0:0",
		"::2:0:41": "::2:0:40",
		"::2:0:49": "::2:0:40",
		"::2:0:52": "::2:0:50",
		"::2:0:57": "::2:0:50",
		"::2:0:59": "::2:0:58",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]interface{}{"ip": valueAddress}
		record, _ := reader.Lookup(net.ParseIP(keyAddress))
		c.Assert(record, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123", "89fa::"} {
		record, _ := reader.Lookup(net.ParseIP(address))
		c.Assert(record, IsNil)
	}
}

func BenchmarkMaxMindDB(b *testing.B) {
	db, err := Open("GeoLite2-City.mmdb")
	if err != nil {
		b.Fatal(err)
	}

	r := rand.New(rand.NewSource(time.Now().UnixNano()))

	for i := 0; i < b.N; i++ {
		num := r.Uint32()
		ip := net.ParseIP(fmt.Sprintf("%d.%d.%d.%d", (0xFF000000&num)>>24,
			(0x00FF0000&num)>>16, (0x0000FF00&num)>>8, 0x000000F&num))
		_, err := db.Lookup(ip)
		if err != nil {
			b.Fatal(err)
		}
	}
	db.Close()
}
<|endoftext|>


<commit_before>// Copyright 2015 LinkedIn Corp. Licensed under the Apache License,
// Version 2.0 (the "License"); you may not use this file except in
// compliance with the License.
 You may obtain a copy of the License
// at http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
//
distributed under the License is distributed on an "AS IS" BASIS,
//
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.Copyright [201X] LinkedIn Corp. Licensed under the Apache
// License, Version 2.0 (the "License"); you may not use this file
// except in compliance with the License.
 You may obtain a copy of
// the License at http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
//
distributed under the License is distributed on an "AS IS" BASIS,
//
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.

package goavro

import (
	"bytes"
	"fmt"
	"testing"
)

func TestRecordRequiresSchema(t *testing.T) {
	_, err := NewRecord()
	checkErrorFatal(t, err, "cannot build record: no schema defined")
}

func TestRecordFieldNames(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"org.foo.Y","fields":[{"type":"int","name":"X"},{"type":"string","name":"W"}]}`
	someRecord, err := NewRecord(RecordSchema(someJSONSchema))
	checkErrorFatal(t, err, nil)
	if someRecord.Name != "org.foo.Y" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecord.Name, "org.foo.Y")
	}
	if someRecord.Fields[0].Name != "org.foo.X" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecord.Fields[0].Name, "org.foo.X")
	}
	if someRecord.Fields[1].Name != "org.foo.W" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecord.Fields[1].Name, "org.foo.W")
	}
}

func TestRecordFieldBailsWithoutName(t *testing.T) {
	schema := make(map[string]interface{})

	schema["type"] = "int"
	_, err := newRecordField(schema)
	checkError(t, err, "ought to have name key")

	schema["name"] = 5
	_, err = newRecordField(schema)
	checkError(t, err, "name ought to be non-empty string")

	schema["name"] = ""
	_, err = newRecordField(schema)
	checkError(t, err, "name ought to be non-empty string")
}

func TestRecordFieldChecksSchema(t *testing.T) {
	var err error
	schema := make(map[string]interface{})

	schema["name"] = ""
	_, err = newRecordField(schema)
	checkError(t, err, "name ought to be non-empty string")

	schema["name"] = "someRecordField"
	_, err = newRecordField(schema)
	checkError(t, err, fmt.Errorf("ought to have type key"))
}

func TestRecordField(t *testing.T) {
	schema := make(map[string]interface{})
	schema["name"] = "someRecordField"
	schema["type"] = "int"
	schema["doc"] = "contans some integer"
	someRecordField, err := newRecordField(schema)
	checkError(t, err, nil)
	if someRecordField.Name != "someRecordField" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecordField.Name, "someRecordField")
	}
}

func TestRecordBailsWithoutName(t *testing.T) {
	recordFields := make([]*recordField, 0)
	{
		schema := make(map[string]interface{})
		schema["name"] = "someRecordField"
		schema["type"] = "int"
		schema["doc"] = "contans some integer"
		someRecordField, err := newRecordField(schema)
		checkErrorFatal(t, err, nil)
		recordFields = append(recordFields, someRecordField)
	}

	schema := make(map[string]interface{})
	schema["fields"] = recordFields

	schema["name"] = 5
	_, err := NewRecord(recordSchemaRaw(schema))
	checkErrorFatal(t, err, "ought to be non-empty string")

	schema["name"] = ""
	_, err = NewRecord(recordSchemaRaw(schema))
	checkError(t, err, "ought to be non-empty string")
}

func TestRecordBailsWithoutFields(t *testing.T) {
	schema := make(map[string]interface{})

	schema["name"] = "someRecord"
	_, err := NewRecord(recordSchemaRaw(schema))
	checkError(t, err, fmt.Errorf("record requires one or more fields"))

	schema["fields"] = 5
	_, err = NewRecord(recordSchemaRaw(schema))
	checkError(t, err, fmt.Errorf("record fields ought to be non-empty array"))

	schema["fields"] = make([]interface{}, 0)
	// Empty unions are only checked if RecordPedantic is set
	_, err = NewRecord(recordSchemaRaw(schema), RecordPedantic())
	checkError(t, err, fmt.Errorf("record fields ought to be non-empty array"))

	fields := make([]interface{}, 0)
	fields = append(fields, "int")
	schema["fields"] = fields
	_, err = NewRecord(recordSchemaRaw(schema))
	checkError(t, err, fmt.Errorf("expected: map[string]interface{}; received: string"))
}

func TestRecordFieldUnion(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"Foo","fields":[{"type":["null","string"],"name":"field1"}]}`
	_, err := NewRecord(RecordSchema(someJSONSchema))
	checkError(t, err, nil)
}

func TestRecordGetFieldSchema(t *testing.T) {
	outerSchema := `
{
  "type": "record",
  "name": "TestRecord",
  "fields": [
    {
      "name": "value",
      "type": "int"
    },
    {
      "name": "rec",
      "type": {
        "type": "array",
        "items": {
          "type": "record",
          "name": "TestRecord2",
          "fields": [
            {
              "name": "stringValue",
              "type": "string"
            },
            {
              "name": "intValue",
              "type": "int"
            }
          ]
        }
      }
    }
  ]
}
`
	outerRecord, err := NewRecord(RecordSchema(outerSchema))
	checkErrorFatal(t, err, nil)
	// make sure it bails when no such schema
	_, err = outerRecord.GetFieldSchema("no_such_field")
	checkError(t, err, "no such field: \"no_such_field\"")
	// get the inner schema
	schema, err := outerRecord.GetFieldSchema("rec")
	checkErrorFatal(t, err, nil)
	_, ok := schema.(map[string]interface{})
	if !ok {
		t.Errorf("Actual: %#v; Expected: %#v", ok, true)
	}
}

func TestNullField(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"Foo","fields":[{"type":"null","name":"field1"}]}`

	codec, err := NewCodec(someJSONSchema)
	checkError(t, err, nil)

	bb := bytes.NewBufferString("")

	rec, err := NewRecord(RecordSchema(someJSONSchema))
	checkError(t, err, nil)

	err = codec.Encode(bb, rec)
	checkError(t, err, nil)
}

func TestNullableStringField(t *testing.T) {
	schema := `
    {
        "name": "record",
        "type": "record",
        "fields": [
            { "type": "string", "name": "string" },
            { "type": ["null", "string"], "name": "nil_or_string" }
        ]
    }
    `

	codec, err := NewCodec(schema)
	checkErrorFatal(t, err, nil)

	record, err := NewRecord(RecordSchema(schema))
	checkErrorFatal(t, err, nil)

	record.Set("nil_or_string", nil) // or a string
	record.Set("string", "can't be empty")
	buf := new(bytes.Buffer)

	err = codec.Encode(buf, record) // "\x1ccan't be empty\x00"
	checkErrorFatal(t, err, nil)

	decode, err := codec.Decode(buf)
	checkErrorFatal(t, err, nil)

	nil_or_string, err := decode.(*Record).Get("nil_or_string")
	checkErrorFatal(t, err, nil)

	if nil_or_string != nil {
		t.Fatalf("Expected nil, got (%T) - (%q)", nil_or_string, nil_or_string)
	}
}
<commit_msg>isolated bug<commit_after>// Copyright 2015 LinkedIn Corp. Licensed under the Apache License,
// Version 2.0 (the "License"); you may not use this file except in
// compliance with the License.
 You may obtain a copy of the License
// at http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
//
distributed under the License is distributed on an "AS IS" BASIS,
//
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.Copyright [201X] LinkedIn Corp. Licensed under the Apache
// License, Version 2.0 (the "License"); you may not use this file
// except in compliance with the License.
 You may obtain a copy of
// the License at http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
//
distributed under the License is distributed on an "AS IS" BASIS,
//
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.

package goavro

import (
	"bytes"
	"fmt"
	"testing"
)

func TestRecordRequiresSchema(t *testing.T) {
	_, err := NewRecord()
	checkErrorFatal(t, err, "cannot build record: no schema defined")
}

func TestRecordFieldNames(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"org.foo.Y","fields":[{"type":"int","name":"X"},{"type":"string","name":"W"}]}`
	someRecord, err := NewRecord(RecordSchema(someJSONSchema))
	checkErrorFatal(t, err, nil)
	if someRecord.Name != "org.foo.Y" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecord.Name, "org.foo.Y")
	}
	if someRecord.Fields[0].Name != "org.foo.X" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecord.Fields[0].Name, "org.foo.X")
	}
	if someRecord.Fields[1].Name != "org.foo.W" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecord.Fields[1].Name, "org.foo.W")
	}
}

func TestRecordFieldBailsWithoutName(t *testing.T) {
	schema := make(map[string]interface{})

	schema["type"] = "int"
	_, err := newRecordField(schema)
	checkError(t, err, "ought to have name key")

	schema["name"] = 5
	_, err = newRecordField(schema)
	checkError(t, err, "name ought to be non-empty string")

	schema["name"] = ""
	_, err = newRecordField(schema)
	checkError(t, err, "name ought to be non-empty string")
}

func TestRecordFieldChecksSchema(t *testing.T) {
	var err error
	schema := make(map[string]interface{})

	schema["name"] = ""
	_, err = newRecordField(schema)
	checkError(t, err, "name ought to be non-empty string")

	schema["name"] = "someRecordField"
	_, err = newRecordField(schema)
	checkError(t, err, fmt.Errorf("ought to have type key"))
}

func TestRecordField(t *testing.T) {
	schema := make(map[string]interface{})
	schema["name"] = "someRecordField"
	schema["type"] = "int"
	schema["doc"] = "contans some integer"
	someRecordField, err := newRecordField(schema)
	checkError(t, err, nil)
	if someRecordField.Name != "someRecordField" {
		t.Errorf("Actual: %#v; Expected: %#v", someRecordField.Name, "someRecordField")
	}
}

func TestRecordBailsWithoutName(t *testing.T) {
	recordFields := make([]*recordField, 0)
	{
		schema := make(map[string]interface{})
		schema["name"] = "someRecordField"
		schema["type"] = "int"
		schema["doc"] = "contans some integer"
		someRecordField, err := newRecordField(schema)
		checkErrorFatal(t, err, nil)
		recordFields = append(recordFields, someRecordField)
	}

	schema := make(map[string]interface{})
	schema["fields"] = recordFields

	schema["name"] = 5
	_, err := NewRecord(recordSchemaRaw(schema))
	checkErrorFatal(t, err, "ought to be non-empty string")

	schema["name"] = ""
	_, err = NewRecord(recordSchemaRaw(schema))
	checkError(t, err, "ought to be non-empty string")
}

func TestRecordBailsWithoutFields(t *testing.T) {
	schema := make(map[string]interface{})

	schema["name"] = "someRecord"
	_, err := NewRecord(recordSchemaRaw(schema))
	checkError(t, err, fmt.Errorf("record requires one or more fields"))

	schema["fields"] = 5
	_, err = NewRecord(recordSchemaRaw(schema))
	checkError(t, err, fmt.Errorf("record fields ought to be non-empty array"))

	schema["fields"] = make([]interface{}, 0)
	// Empty unions are only checked if RecordPedantic is set
	_, err = NewRecord(recordSchemaRaw(schema), RecordPedantic())
	checkError(t, err, fmt.Errorf("record fields ought to be non-empty array"))

	fields := make([]interface{}, 0)
	fields = append(fields, "int")
	schema["fields"] = fields
	_, err = NewRecord(recordSchemaRaw(schema))
	checkError(t, err, fmt.Errorf("expected: map[string]interface{}; received: string"))
}

func TestRecordFieldUnionNullOrStringCanBeNull(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"Foo","fields":[{"type":["null","string"],"name":"field1"}]}`

	codec, err := NewCodec(someJSONSchema)
	checkErrorFatal(t, err, nil)

	record, err := NewRecord(RecordSchema(someJSONSchema))
	checkErrorFatal(t, err, nil)

	bb := new(bytes.Buffer)
	err = codec.Encode(bb, record)
	checkError(t, err, nil)
}

func TestRecordFieldUnionNullOrStringCanBeString(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"Foo","fields":[{"type":["null","string"],"name":"field1"}]}`

	codec, err := NewCodec(someJSONSchema)
	checkErrorFatal(t, err, nil)

	record, err := NewRecord(RecordSchema(someJSONSchema))
	checkErrorFatal(t, err, nil)

	record.Set("field1", "something")

	bb := new(bytes.Buffer)
	err = codec.Encode(bb, record)
	checkError(t, err, nil)
}

func TestRecordGetFieldSchema(t *testing.T) {
	outerSchema := `
{
  "type": "record",
  "name": "TestRecord",
  "fields": [
    {
      "name": "value",
      "type": "int"
    },
    {
      "name": "rec",
      "type": {
        "type": "array",
        "items": {
          "type": "record",
          "name": "TestRecord2",
          "fields": [
            {
              "name": "stringValue",
              "type": "string"
            },
            {
              "name": "intValue",
              "type": "int"
            }
          ]
        }
      }
    }
  ]
}
`
	outerRecord, err := NewRecord(RecordSchema(outerSchema))
	checkErrorFatal(t, err, nil)
	// make sure it bails when no such schema
	_, err = outerRecord.GetFieldSchema("no_such_field")
	checkError(t, err, "no such field: \"no_such_field\"")
	// get the inner schema
	schema, err := outerRecord.GetFieldSchema("rec")
	checkErrorFatal(t, err, nil)
	_, ok := schema.(map[string]interface{})
	if !ok {
		t.Errorf("Actual: %#v; Expected: %#v", ok, true)
	}
}

func TestNullField(t *testing.T) {
	someJSONSchema := `{"type":"record","name":"Foo","fields":[{"type":"null","name":"field1"}]}`

	codec, err := NewCodec(someJSONSchema)
	checkError(t, err, nil)

	bb := bytes.NewBufferString("")

	rec, err := NewRecord(RecordSchema(someJSONSchema))
	checkError(t, err, nil)

	err = codec.Encode(bb, rec)
	checkError(t, err, nil)
}

func TestNullableStringField(t *testing.T) {
	schema := `
    {
        "name": "record",
        "type": "record",
        "fields": [
            { "type": "string", "name": "string" },
            { "type": ["null", "string"], "name": "nil_or_string" }
        ]
    }
    `

	codec, err := NewCodec(schema)
	checkErrorFatal(t, err, nil)

	record, err := NewRecord(RecordSchema(schema))
	checkErrorFatal(t, err, nil)

	record.Set("nil_or_string", nil) // or a string
	record.Set("string", "can't be empty")
	buf := new(bytes.Buffer)

	err = codec.Encode(buf, record) // "\x1ccan't be empty\x00"
	checkErrorFatal(t, err, nil)

	decode, err := codec.Decode(buf)
	checkErrorFatal(t, err, nil)

	nil_or_string, err := decode.(*Record).Get("nil_or_string")
	checkErrorFatal(t, err, nil)

	if nil_or_string != nil {
		t.Fatalf("Expected nil, got (%T) - (%q)", nil_or_string, nil_or_string)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"github.com/nsf/termbox-go"
	"io/ioutil"
	"log"
	"math"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

var columnNames = statusLine{"Job name", "Queued", "Running", "Workers"}

type statusLine struct {
	name    string
	queued  string
	running string
	workers string
}

type fieldWidths struct {
	name    int
	queued  int
	running int
	workers int
}

type gearmanStatus struct {
	statusLines    []statusLine
	statusLineDims fieldWidths
}

func fieldWidthsFactory(line statusLine) fieldWidths {
	return fieldWidths{
		len(line.name),
		len(line.queued),
		len(line.running),
		len(line.workers)}
}

func statusLineFromString(line string) (statusLine, error) {
	parts := strings.Fields(line)
	if len(parts) != 4 {
		return statusLine{}, errors.New("Wrong number of fields")
	}
	return statusLine{parts[0], parts[1], parts[2], parts[3]}, nil
}

func getStatus(c chan gearmanStatus) {
	const waitTime = 1000 * time.Millisecond
	for {
		log.Println("Getting status")
		start := time.Now()
		widths := fieldWidthsFactory(columnNames)
		data, err := exec.Command(gearadmin, "--host="+gearmanHost, "--port="+gearmanPort, "--status").Output()
		if err != nil {
			log.Fatal("Couldn't get status from gearadmin: " + err.Error())
		}
		strData := string(data)
		statusLines := make([]statusLine, 0)
		for _, line := range strings.Split(strData, "\n") {
			statusLine, ok := statusLineFromString(line)
			if ok != nil {
				continue
			}
			widths.name = int(math.Max(float64(len(statusLine.name)), float64(widths.name)))
			widths.queued = int(math.Max(float64(len(statusLine.queued)), float64(widths.queued)))
			widths.running = int(math.Max(float64(len(statusLine.running)), float64(widths.running)))
			widths.workers = int(math.Max(float64(len(statusLine.workers)), float64(widths.workers)))
			statusLines = append(statusLines, statusLine)
		}
		c <- gearmanStatus{statusLines, widths}
		duration := time.Since(start)
		time.Sleep(waitTime - duration)
	}
}

func print_tb(x, y int, fg, bg termbox.Attribute, msg string) {
	for _, c := range msg {
		termbox.SetCell(x, y, c, fg, bg)
		x++
	}
}

func printf_tb(x, y int, fg, bg termbox.Attribute, format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	print_tb(x, y, fg, bg, s)
}

func printField(x, y, fieldWidth int, value string, bold bool) int {
	intValue, ok := strconv.Atoi(value)
	if ok == nil {
		value = fmt.Sprintf("%"+strconv.Itoa(fieldWidth)+"d", intValue) + " "
	}
	fg := termbox.ColorDefault
	if bold {
		fg |= termbox.AttrBold
	}
	print_tb(x, y, fg, termbox.ColorDefault, value)
	return x + fieldWidth
}

func printLine(y int, widths fieldWidths, line statusLine, bold bool) {
	x := 0
	if len(line.name) > widths.name {
		line.name = line.name[:widths.name]
	}
	x = printField(x, y, widths.name+1, line.name, bold)
	x = printField(x, y, widths.queued+1, line.queued, bold)
	x = printField(x, y, widths.running+1, line.running, bold)
	x = printField(x, y, widths.workers, line.workers, bold)
}

func drawStatus(gearmanStatus gearmanStatus, position int, height int, width int) {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	y := 0
	printY := 0

	widths := gearmanStatus.statusLineDims
	totalWidth := widths.name + widths.queued + widths.running + widths.workers + 3
	if totalWidth > width {
		widths.name += width - totalWidth
	}
	printLine(0, widths, columnNames, true)
	for _, line := range gearmanStatus.statusLines {
		printY = y - position
		if printY > height {
			break
		}
		if printY < 1 {
			y++
			continue
		}
		printLine(printY, widths, line, false)
		y++
	}
	termbox.Flush()
}

func handleEvents(direction chan int, resized chan termbox.Event, quit chan bool) {
	for {
		event := termbox.PollEvent()
		log.Println("Recieved event: ", event)
		switch event.Type {
		case termbox.EventKey:
			if event.Ch == 'q' {
				quit <- true
			}
			switch event.Key {
			case termbox.KeyArrowUp:
				direction <- -1
			case termbox.KeyArrowDown:
				direction <- 1
			}
		case termbox.EventResize:
			resized <- event
		}
	}
}

func calculatePosition(currentPosition int, direction int, gearmanStatus gearmanStatus) (int, bool) {
	_, height := termbox.Size()
	scrolledToBottom := len(gearmanStatus.statusLines) <= (currentPosition + height)
	scrolledToTop := currentPosition == 0
	if (direction < 0 && !scrolledToTop) || (direction > 0 && !scrolledToBottom) {
		log.Println("Moving")
		return currentPosition + direction, true
	}
	return currentPosition, false
}

func scrollOutput(direction int, scroll chan int, position int, currentGearmanStatus gearmanStatus) int {
	positionUpdated := false
	log.Println("Scrolling")
	for {
		width, height := termbox.Size()
		position, positionUpdated = calculatePosition(position, direction, currentGearmanStatus)
		if positionUpdated {
			drawStatus(currentGearmanStatus, position, height, width)
		}
		select {
		case direction = <-scroll:
			log.Println("Collating scrolling")
		default:
			return position
		}
	}
	panic("unreachable")
}

func initLogging() *os.File {
	f, err := os.OpenFile("/tmp/gearman_gtop.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		panic(err)
	}
	log.SetOutput(f)
	return f
}

var doLogging bool
var gearadmin string
var gearmanHost string
var gearmanPort string

func init() {
	logDefault := false
	logUsage := "Log debug to /tmp/gearman_gtop.log"
	flag.BoolVar(&doLogging, "log", logDefault, logUsage)
	flag.BoolVar(&doLogging, "l", logDefault, logUsage+" (shorthand)")
	gearadminDefault := "/usr/bin/gearadmin"
	gearadminUsage := "Path to gearadmin, e.g. `which gearadmin`"
	flag.StringVar(&gearadmin, "gearadmin", gearadminDefault, gearadminUsage)
	hostDefault := "localhost"
	hostUsage := "Gearmand host to connect to"
	flag.StringVar(&gearmanHost, "host", hostDefault, hostUsage)
	flag.StringVar(&gearmanHost, "h", hostDefault, hostUsage+" (shorthand)")
	portDefault := "4730"
	portUsage := "Gearmand port to connect to"
	flag.StringVar(&gearmanPort, "port", portDefault, portUsage)
	flag.StringVar(&gearmanPort, "p", portDefault, portUsage+" (shorthand)")
}

func main() {

	flag.Parse()

	if doLogging {
		defer (initLogging()).Close()
	} else {
		log.SetOutput(ioutil.Discard)
	}

	var currentGearmanStatus gearmanStatus
	position := 0
	status := make(chan gearmanStatus)
	quit := make(chan bool)
	resized := make(chan termbox.Event)
	scroll := make(chan int, 3)

	err := termbox.Init()
	if err != nil {
		panic(err)
	}
	defer termbox.Close()
	termbox.SetInputMode(termbox.InputEsc)

	go getStatus(status)
	go handleEvents(scroll, resized, quit)
	for {
		select {
		case currentGearmanStatus = <-status:
			log.Println("Redrawing for updated status")
			width, height := termbox.Size()
			drawStatus(currentGearmanStatus, position, height, width)
		case ev := <-resized:
			log.Println("Redrawing for resize")
			drawStatus(currentGearmanStatus, position, ev.Height, ev.Width)
		case direction := <-scroll:
			position = scrollOutput(direction, scroll, position, currentGearmanStatus)
		case <-quit:
			log.Println("Exiting")
			return
		}
	}
}
<commit_msg>Added handler for Ctrl+C Removed dependency on gearadmin Simplified max calculation<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"github.com/nsf/termbox-go"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"time"
	"net"
	"bufio"
)

var columnNames = statusLine{"Job name", "Queued", "Running", "Workers"}

type statusLine struct {
	name    string
	queued  string
	running string
	workers string
}

type fieldWidths struct {
	name    int
	queued  int
	running int
	workers int
}

type gearmanStatus struct {
	statusLines    []statusLine
	statusLineDims fieldWidths
}

func fieldWidthsFactory(line statusLine) fieldWidths {
	return fieldWidths{
		len(line.name),
		len(line.queued),
		len(line.running),
		len(line.workers)}
}

func fatal(msg string) {
	termbox.Close()
	log.Fatal(msg)
}

func statusLineFromString(line string) (statusLine, error) {
	parts := strings.Fields(line)
	if len(parts) != 4 {
		return statusLine{}, errors.New("Wrong number of fields")
	}
	return statusLine{parts[0], parts[1], parts[2], parts[3]}, nil
}

func max(a, b int) int {
    if a >= b {
        return a
    }
    return b
}

func getStatus(c chan gearmanStatus) {
	log.Println("Connecting to gearman")
	const waitTime = 1000 * time.Millisecond
	gearman, err := net.DialTimeout("tcp", gearmanHost+":"+gearmanPort, 1 * time.Second)
	if err != nil {
		fatal("Couldn't connect to gearman on "+gearmanHost+":"+gearmanPort)
	}
	defer gearman.Close()
	gearmanStream := bufio.NewReader(gearman)
	for {
		log.Println("Getting status")
		start := time.Now()
		widths := fieldWidthsFactory(columnNames)
		gearman.Write([]byte("status\n"))
		statusLines := make([]statusLine, 0)
		for {
			line, err := gearmanStream.ReadString('\n')
			if err != nil {
				break
			}
			if line == ".\n" {
				break
			}
			statusLine, ok := statusLineFromString(line)
			if ok != nil {
				continue
			}
			widths.name = max(len(statusLine.name), widths.name)
			widths.queued = max(len(statusLine.queued), widths.queued)
			widths.running = max(len(statusLine.running), widths.running)
			widths.workers = max(len(statusLine.workers), widths.workers)
			statusLines = append(statusLines, statusLine)
		}
		c <- gearmanStatus{statusLines, widths}
		duration := time.Since(start)
		time.Sleep(waitTime - duration)
	}
}

func print_tb(x, y int, fg, bg termbox.Attribute, msg string) {
	for _, c := range msg {
		termbox.SetCell(x, y, c, fg, bg)
		x++
	}
}

func printf_tb(x, y int, fg, bg termbox.Attribute, format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	print_tb(x, y, fg, bg, s)
}

func printField(x, y, fieldWidth int, value string, bold bool) int {
	intValue, ok := strconv.Atoi(value)
	if ok == nil {
		value = fmt.Sprintf("%"+strconv.Itoa(fieldWidth)+"d", intValue) + " "
	}
	fg := termbox.ColorDefault
	if bold {
		fg |= termbox.AttrBold
	}
	print_tb(x, y, fg, termbox.ColorDefault, value)
	return x + fieldWidth
}

func printLine(y int, widths fieldWidths, line statusLine, bold bool) {
	x := 0
	if len(line.name) > widths.name {
		line.name = line.name[:widths.name]
	}
	x = printField(x, y, widths.name+1, line.name, bold)
	x = printField(x, y, widths.queued+1, line.queued, bold)
	x = printField(x, y, widths.running+1, line.running, bold)
	x = printField(x, y, widths.workers, line.workers, bold)
}

func drawStatus(gearmanStatus gearmanStatus, position int, height int, width int) {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	y := 0
	printY := 0

	widths := gearmanStatus.statusLineDims
	totalWidth := widths.name + widths.queued + widths.running + widths.workers + 3
	if totalWidth > width {
		widths.name += width - totalWidth
	}
	printLine(0, widths, columnNames, true)
	for _, line := range gearmanStatus.statusLines {
		printY = y - position
		if printY > height {
			break
		}
		if printY < 1 {
			y++
			continue
		}
		printLine(printY, widths, line, false)
		y++
	}
	termbox.Flush()
}

func handleEvents(direction chan int, resized chan termbox.Event, quit chan bool) {
	for {
		event := termbox.PollEvent()
		log.Println("Recieved event: ", event)
		switch event.Type {
		case termbox.EventKey:
			if event.Ch == 'q' {
				quit <- true
			}
			switch event.Key {
			case termbox.KeyCtrlC:
				quit <- true
			case termbox.KeyArrowUp:
				direction <- -1
			case termbox.KeyArrowDown:
				direction <- 1
			}
		case termbox.EventResize:
			resized <- event
		}
	}
}

func calculatePosition(currentPosition int, direction int, gearmanStatus gearmanStatus) (int, bool) {
	_, height := termbox.Size()
	scrolledToBottom := len(gearmanStatus.statusLines) <= (currentPosition + height)
	scrolledToTop := currentPosition == 0
	if (direction < 0 && !scrolledToTop) || (direction > 0 && !scrolledToBottom) {
		log.Println("Moving")
		return currentPosition + direction, true
	}
	return currentPosition, false
}

func scrollOutput(direction int, scroll chan int, position int, currentGearmanStatus gearmanStatus) int {
	positionUpdated := false
	log.Println("Scrolling")
	for {
		width, height := termbox.Size()
		position, positionUpdated = calculatePosition(position, direction, currentGearmanStatus)
		if positionUpdated {
			drawStatus(currentGearmanStatus, position, height, width)
		}
		select {
		case direction = <-scroll:
			log.Println("Collating scrolling")
		default:
			return position
		}
	}
	panic("unreachable")
}

func initLogging() *os.File {
	f, err := os.OpenFile("/tmp/gearman_gtop.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		panic(err)
	}
	log.SetOutput(f)
	return f
}

var doLogging bool
var gearmanHost string
var gearmanPort string

func init() {
	logDefault := false
	logUsage := "Log debug to /tmp/gearman_gtop.log"
	flag.BoolVar(&doLogging, "log", logDefault, logUsage)
	flag.BoolVar(&doLogging, "l", logDefault, logUsage+" (shorthand)")
	hostDefault := "localhost"
	hostUsage := "Gearmand host to connect to"
	flag.StringVar(&gearmanHost, "host", hostDefault, hostUsage)
	flag.StringVar(&gearmanHost, "h", hostDefault, hostUsage+" (shorthand)")
	portDefault := "4730"
	portUsage := "Gearmand port to connect to"
	flag.StringVar(&gearmanPort, "port", portDefault, portUsage)
	flag.StringVar(&gearmanPort, "p", portDefault, portUsage+" (shorthand)")
}

func main(){

	flag.Parse()

	if doLogging {
		defer (initLogging()).Close()
	} else {
		log.SetOutput(ioutil.Discard)
	}

	var currentGearmanStatus gearmanStatus
	position := 0
	status := make(chan gearmanStatus)
	quit := make(chan bool)
	resized := make(chan termbox.Event)
	scroll := make(chan int, 3)

	err := termbox.Init()
	if err != nil {
		panic(err)
	}
	defer termbox.Close()
	termbox.SetInputMode(termbox.InputEsc)

	go getStatus(status)
	go handleEvents(scroll, resized, quit)
	for {
		select {
		case currentGearmanStatus = <-status:
			log.Println("Redrawing for updated status")
			width, height := termbox.Size()
			drawStatus(currentGearmanStatus, position, height, width)
		case ev := <-resized:
			log.Println("Redrawing for resize")
			drawStatus(currentGearmanStatus, position, ev.Height, ev.Width)
		case direction := <-scroll:
			position = scrollOutput(direction, scroll, position, currentGearmanStatus)
		case <-quit:
			log.Println("Exiting")
			return
		}
	}
}
<|endoftext|>


<commit_before>package worker_test

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/onsi/gomega/gbytes"

	"github.com/concourse/concourse/worker/workerfakes"

	"code.cloudfoundry.org/lager/lagertest"
	"github.com/concourse/concourse/atc/worker/gclient"
	"github.com/concourse/concourse/worker"
	"github.com/onsi/gomega/ghttp"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Container Sweeper", func() {
	var (
		garden *ghttp.Server

		testLogger = lagertest.NewTestLogger("container-sweeper")

		sweepInterval time.Duration
		maxInFlight   uint16

		fakeTSAClient workerfakes.FakeTSAClient
		gardenClient  gclient.Client

		sweeper *worker.ContainerSweeper

		gardenClientRequestTimeout time.Duration

		osSignal  chan os.Signal
		readyChan chan struct{}
		exited    chan struct{}

		//err error
	)

	BeforeEach(func() {
		sweepInterval = 50 * time.Millisecond
		maxInFlight = 1

		gardenClientRequestTimeout = 5 * time.Millisecond

		garden = ghttp.NewServer()

		osSignal = make(chan os.Signal)
		readyChan = make(chan struct{})
		exited = make(chan struct{})

		gardenAddr := fmt.Sprintf("http://%s", garden.Addr())
		gardenClient = gclient.BasicGardenClientWithRequestTimeout(testLogger, gardenClientRequestTimeout, gardenAddr)

		fakeTSAClient = workerfakes.FakeTSAClient{}
		fakeTSAClient.ReportContainersReturns(nil)

		sweeper = worker.NewContainerSweeper(testLogger, sweepInterval, &fakeTSAClient, gardenClient, maxInFlight)

	})

	AfterEach(func() {
		garden.Close()
	})

	Context("when garden doesn't respond on DELETE", func() {
		var (
			gardenContext context.Context
			gardenCancel  context.CancelFunc
		)
		JustBeforeEach(func() {
			go func() {
				_ = sweeper.Run(osSignal, readyChan)
				close(exited)
			}()
		})

		BeforeEach(func() {
			gardenContext, gardenCancel = context.WithCancel(context.Background())
			garden.AppendHandlers(
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("GET", "/containers"),
					ghttp.RespondWithJSONEncoded(200, []map[string]string{}),
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-1"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-2"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("GET", "/containers"),
					ghttp.RespondWithJSONEncoded(200, []map[string]string{}),
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-3"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-4"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
			)
			// First GC Tick
			fakeTSAClient.ContainersToDestroyReturnsOnCall(0, []string{"some-handle-1", "some-handle-2"}, nil)
			// Second GC Tick
			fakeTSAClient.ContainersToDestroyReturnsOnCall(1, []string{"some-handle-3", "some-handle-4"}, nil)

			garden.AllowUnhandledRequests = true

		})
		AfterEach(func() {
			close(osSignal)
			<-exited
			gardenCancel()
		})

		It("request to garden times out eventually", func() {
			Eventually(testLogger.Buffer()).Should(gbytes.Say("failed-to-destroy-container\".*net/http: request canceled \\(Client.Timeout exceeded while awaiting headers\\)"))
		})
		It("sweeper continues ticking and GC'ing", func() {
			// ensure all 4 DELETEs are issues over 2 successive ticks
			Eventually(func() []string {
				// Gather all containers deleted
				var deleteRequestPaths []string
				for _, req := range garden.ReceivedRequests() {
					if req.Method == http.MethodDelete {
						deleteRequestPaths = append(deleteRequestPaths, req.RequestURI)
					}
				}
				return deleteRequestPaths
			}).Should(ConsistOf(
				"/containers/some-handle-1",
				"/containers/some-handle-2",
				"/containers/some-handle-3",
				"/containers/some-handle-4"))

			// Check calls to TSA for containers to destroy > 1
			Expect(fakeTSAClient.ContainersToDestroyCallCount()).To(BeNumerically(">=", 2))
		})
	})

})
<commit_msg>container sweeper test refactor<commit_after>package worker_test

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/onsi/gomega/gbytes"

	"github.com/concourse/concourse/worker/workerfakes"

	"code.cloudfoundry.org/lager/lagertest"
	"github.com/concourse/concourse/atc/worker/gclient"
	"github.com/concourse/concourse/worker"
	"github.com/onsi/gomega/ghttp"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Container Sweeper", func() {
	var (
		garden *ghttp.Server

		testLogger = lagertest.NewTestLogger("container-sweeper")

		sweepInterval time.Duration
		maxInFlight   uint16

		fakeTSAClient workerfakes.FakeTSAClient
		gardenClient  gclient.Client

		sweeper *worker.ContainerSweeper

		gardenClientRequestTimeout time.Duration

		osSignal  chan os.Signal
		readyChan chan struct{}
		exited    chan struct{}

		//err error
	)

	BeforeEach(func() {
		sweepInterval = 50 * time.Millisecond
		maxInFlight = 1

		gardenClientRequestTimeout = 5 * time.Millisecond

		garden = ghttp.NewServer()

		osSignal = make(chan os.Signal)
		readyChan = make(chan struct{})
		exited = make(chan struct{})

		gardenAddr := fmt.Sprintf("http://%s", garden.Addr())
		gardenClient = gclient.BasicGardenClientWithRequestTimeout(testLogger, gardenClientRequestTimeout, gardenAddr)

		fakeTSAClient = workerfakes.FakeTSAClient{}
		fakeTSAClient.ReportContainersReturns(nil)

		sweeper = worker.NewContainerSweeper(testLogger, sweepInterval, &fakeTSAClient, gardenClient, maxInFlight)

	})

	JustBeforeEach(func() {
		go func() {
			_ = sweeper.Run(osSignal, readyChan)
			close(exited)
		}()
	})

	AfterEach(func() {
		close(osSignal)
		<-exited
		garden.Close()
	})

	Context("when garden doesn't respond on DELETE", func() {
		var (
			gardenContext context.Context
			gardenCancel  context.CancelFunc
		)

		BeforeEach(func() {
			gardenContext, gardenCancel = context.WithCancel(context.Background())
			garden.AppendHandlers(
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("GET", "/containers"),
					ghttp.RespondWithJSONEncoded(200, []map[string]string{}),
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-1"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-2"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("GET", "/containers"),
					ghttp.RespondWithJSONEncoded(200, []map[string]string{}),
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-3"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
				ghttp.CombineHandlers(
					ghttp.VerifyRequest("DELETE", "/containers/some-handle-4"),
					func(w http.ResponseWriter, r *http.Request) {
						<-gardenContext.Done()
					},
				),
			)
			// First GC Tick
			fakeTSAClient.ContainersToDestroyReturnsOnCall(0, []string{"some-handle-1", "some-handle-2"}, nil)
			// Second GC Tick
			fakeTSAClient.ContainersToDestroyReturnsOnCall(1, []string{"some-handle-3", "some-handle-4"}, nil)

			garden.AllowUnhandledRequests = true

		})
		AfterEach(func() {
			gardenCancel()

		})

		It("request to garden times out eventually", func() {
			Eventually(testLogger.Buffer()).Should(gbytes.Say("failed-to-destroy-container\".*net/http: request canceled \\(Client.Timeout exceeded while awaiting headers\\)"))
		})
		It("sweeper continues ticking and GC'ing", func() {
			// ensure all 4 DELETEs are issues over 2 successive ticks
			Eventually(func() []string {
				// Gather all containers deleted
				var deleteRequestPaths []string
				for _, req := range garden.ReceivedRequests() {
					if req.Method == http.MethodDelete {
						deleteRequestPaths = append(deleteRequestPaths, req.RequestURI)
					}
				}
				return deleteRequestPaths
			}).Should(ConsistOf(
				"/containers/some-handle-1",
				"/containers/some-handle-2",
				"/containers/some-handle-3",
				"/containers/some-handle-4"))

			// Check calls to TSA for containers to destroy > 1
			Expect(fakeTSAClient.ContainersToDestroyCallCount()).To(BeNumerically(">=", 2))
		})
	})

})
<|endoftext|>


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

package profiler

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"testing"
	"text/template"
	"time"

	"cloud.google.com/go/profiler/proftest"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"
)

const (
	cloudScope        = "https://www.googleapis.com/auth/cloud-platform"
	benchFinishString = "busybench finished profiling"
	errorString       = "failed to set up or run the benchmark"
)

const startupTemplate = `
#! /bin/bash

# Signal any unexpected error.
trap 'echo "{{.ErrorString}}"' ERR

(
# Shut down the VM in 5 minutes after this script exits
# to stop accounting the VM for billing and cores quota.
trap "sleep 300 && poweroff" EXIT

retry() {
  for i in {1..3}; do
    "${@}" && return 0
  done
  return 1
}

# Fail on any error.
set -eo pipefail

# Display commands being run.
set -x

# Suppress debconf warnings to minimize noise in logs
export DEBIAN_FRONTEND="noninteractive"

# Building go from master will fail without $HOME set.
# Set $HOME becasue it is not automatically set when this script runs.
# If $HOME is unset, $GOCACHE must be set for Go 1.12+
cd /root
export HOME=$PWD

# Install git
retry apt-get update >/dev/null
retry apt-get -y -q install git >/dev/null

# Set $GOPATH
export GOPATH="$HOME/go"

export GOCLOUD_HOME=$GOPATH/src/cloud.google.com/go
mkdir -p $GOCLOUD_HOME

# Install gcc, needed to install go master
if [ "{{.GoVersion}}" = "master" ]
then
retry apt-get -y -q install gcc >/dev/null
fi

# Install desired Go version
mkdir -p /tmp/bin
retry curl -sL -o /tmp/bin/gimme https://raw.githubusercontent.com/travis-ci/gimme/master/gimme
chmod +x /tmp/bin/gimme
export PATH=$PATH:/tmp/bin

retry gimme {{.GoVersion}} > out.gimme
eval "$(cat out.gimme)"

# Install agent
retry git clone https://code.googlesource.com/gocloud $GOCLOUD_HOME >/dev/null
cd $GOCLOUD_HOME
retry git fetch origin {{.Commit}}
git reset --hard {{.Commit}}

cd $GOCLOUD_HOME/profiler/busybench
retry go get >/dev/null

# Run benchmark with agent
go run busybench.go --service="{{.Service}}" --mutex_profiling="{{.MutexProfiling}}"

# Write output to serial port 2 with timestamp.
) 2>&1 | while read line; do echo "$(date): ${line}"; done >/dev/ttyS1
`

type goGCETestCase struct {
	proftest.InstanceConfig
	name             string
	goVersion        string
	mutexProfiling   bool
	wantProfileTypes []string
}

func (tc *goGCETestCase) initializeStartupScript(template *template.Template, commit string) error {
	var buf bytes.Buffer
	err := template.Execute(&buf,
		struct {
			Service        string
			GoVersion      string
			Commit         string
			ErrorString    string
			MutexProfiling bool
		}{
			Service:        tc.name,
			GoVersion:      tc.goVersion,
			Commit:         commit,
			ErrorString:    errorString,
			MutexProfiling: tc.mutexProfiling,
		})
	if err != nil {
		return fmt.Errorf("failed to render startup script for %s: %v", tc.name, err)
	}
	tc.StartupScript = buf.String()
	return nil
}

func TestAgentIntegration(t *testing.T) {
	// Testing against master requires building go code and may take up to 10 minutes.
	// Allow this test to run in parallel with other top level tests to avoid timeouts.
	t.Parallel()
	if testing.Short() {
		t.Skip("skipping profiler integration test in short mode")
	}

	projectID := os.Getenv("GCLOUD_TESTS_GOLANG_PROJECT_ID")
	if projectID == "" {
		t.Skip("skipping profiler integration test when GCLOUD_TESTS_GOLANG_PROJECT_ID variable is not set")
	}

	zone := os.Getenv("GCLOUD_TESTS_GOLANG_PROFILER_ZONE")
	if zone == "" {
		t.Fatalf("GCLOUD_TESTS_GOLANG_PROFILER_ZONE environment variable must be set when integration test is requested")
	}

	// Figure out the Git commit of the current directory. The source checkout in
	// the test VM will run in the same commit. Note that any local changes to
	// the profiler agent won't be tested in the integration test. This flow only
	// works with code that has been committed and pushed to the public repo
	// (either to master or to a branch).
	output, err := exec.Command("git", "rev-parse", "HEAD").CombinedOutput()
	if err != nil {
		t.Fatalf("failed to gather the Git revision of the current source: %v", err)
	}
	commit := strings.Trim(string(output), "\n")
	t.Logf("using Git commit %q for the profiler integration test", commit)

	pst, err := time.LoadLocation("America/Los_Angeles")
	if err != nil {
		t.Fatalf("failed to initiate PST location: %v", err)
	}
	runID := strings.Replace(time.Now().In(pst).Format("2006-01-02-15-04-05.000000-0700"), ".", "-", -1)

	ctx := context.Background()

	client, err := google.DefaultClient(ctx, cloudScope)
	if err != nil {
		t.Fatalf("failed to get default client: %v", err)
	}

	computeService, err := compute.New(client)
	if err != nil {
		t.Fatalf("failed to initialize compute service: %v", err)
	}

	template, err := template.New("startupScript").Parse(startupTemplate)
	if err != nil {
		t.Fatalf("failed to parse startup script template: %v", err)
	}

	tr := proftest.TestRunner{
		Client: client,
	}

	gceTr := proftest.GCETestRunner{
		TestRunner:     tr,
		ComputeService: computeService,
	}

	testcases := []goGCETestCase{
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-gomaster-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-gomaster-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "master",
			mutexProfiling:   true,
		},
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-go111-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-go111-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "1.11",
			mutexProfiling:   true,
		},
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-go110-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-go110-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "1.10",
			mutexProfiling:   true,
		},
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-go19-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-go19-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "1.9",
			mutexProfiling:   true,
		},
	}
	// The number of tests run in parallel is the current value of GOMAXPROCS.
	runtime.GOMAXPROCS(len(testcases))
	for _, tc := range testcases {
		tc := tc // capture range variable
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if err := tc.initializeStartupScript(template, commit); err != nil {
				t.Fatalf("failed to initialize startup script")
			}

			if err := gceTr.StartInstance(ctx, &tc.InstanceConfig); err != nil {
				t.Fatal(err)
			}
			defer func() {
				if gceTr.DeleteInstance(ctx, &tc.InstanceConfig); err != nil {
					t.Fatal(err)
				}
			}()

			timeoutCtx, cancel := context.WithTimeout(ctx, time.Minute*25)
			defer cancel()
			if err := gceTr.PollForSerialOutput(timeoutCtx, &tc.InstanceConfig, benchFinishString, errorString); err != nil {
				t.Fatalf("PollForSerialOutput() got error: %v", err)
			}

			timeNow := time.Now()
			endTime := timeNow.Format(time.RFC3339)
			startTime := timeNow.Add(-1 * time.Hour).Format(time.RFC3339)
			for _, pType := range tc.wantProfileTypes {
				pr, err := tr.QueryProfiles(tc.ProjectID, tc.name, startTime, endTime, pType)
				if err != nil {
					t.Errorf("QueryProfiles(%s, %s, %s, %s, %s) got error: %v", tc.ProjectID, tc.name, startTime, endTime, pType, err)
					continue
				}
				if err := pr.HasFunction("busywork"); err != nil {
					t.Error(err)
				}
			}
		})
	}
}
<commit_msg>Revert "profiler: fix gimme and reenable integration test"<commit_after>// Copyright 2017 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package profiler

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"testing"
	"text/template"
	"time"

	"cloud.google.com/go/profiler/proftest"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"
)

const (
	cloudScope        = "https://www.googleapis.com/auth/cloud-platform"
	benchFinishString = "busybench finished profiling"
	errorString       = "failed to set up or run the benchmark"
)

const startupTemplate = `
#! /bin/bash

# Signal any unexpected error.
trap 'echo "{{.ErrorString}}"' ERR

(
# Shut down the VM in 5 minutes after this script exits
# to stop accounting the VM for billing and cores quota.
trap "sleep 300 && poweroff" EXIT

retry() {
  for i in {1..3}; do
    "${@}" && return 0
  done
  return 1
}

# Fail on any error.
set -eo pipefail

# Display commands being run.
set -x

# Install git
retry apt-get update >/dev/null
retry apt-get -y -q install git >/dev/null

# $GOCACHE is required from Go 1.12. See https://golang.org/doc/go1.11#gocache
# $GOCACHE is explicitly set becasue $HOME is not set when this code runs
mkdir -p /tmp/gocache
export GOCACHE=/tmp/gocache

# Install gcc, needed to install go master
if [ "{{.GoVersion}}" = "master" ]
then
retry apt-get -y -q install gcc >/dev/null
fi

# Install desired Go version
mkdir -p /tmp/bin
retry curl -sL -o /tmp/bin/gimme https://raw.githubusercontent.com/travis-ci/gimme/master/gimme
chmod +x /tmp/bin/gimme
export PATH=$PATH:/tmp/bin

retry eval "$(gimme {{.GoVersion}})"

# Set $GOPATH
export GOPATH="$HOME/go"

export GOCLOUD_HOME=$GOPATH/src/cloud.google.com/go
mkdir -p $GOCLOUD_HOME

# Install agent
retry git clone https://code.googlesource.com/gocloud $GOCLOUD_HOME >/dev/null
cd $GOCLOUD_HOME
retry git fetch origin {{.Commit}}
git reset --hard {{.Commit}}

cd $GOCLOUD_HOME/profiler/busybench
retry go get >/dev/null

# Run benchmark with agent
go run busybench.go --service="{{.Service}}" --mutex_profiling="{{.MutexProfiling}}"

# Write output to serial port 2 with timestamp.
) 2>&1 | while read line; do echo "$(date): ${line}"; done >/dev/ttyS1
`

type goGCETestCase struct {
	proftest.InstanceConfig
	name             string
	goVersion        string
	mutexProfiling   bool
	wantProfileTypes []string
}

func (tc *goGCETestCase) initializeStartupScript(template *template.Template, commit string) error {
	var buf bytes.Buffer
	err := template.Execute(&buf,
		struct {
			Service        string
			GoVersion      string
			Commit         string
			ErrorString    string
			MutexProfiling bool
		}{
			Service:        tc.name,
			GoVersion:      tc.goVersion,
			Commit:         commit,
			ErrorString:    errorString,
			MutexProfiling: tc.mutexProfiling,
		})
	if err != nil {
		return fmt.Errorf("failed to render startup script for %s: %v", tc.name, err)
	}
	tc.StartupScript = buf.String()
	return nil
}

func TestAgentIntegration(t *testing.T) {
	t.Skip("https://github.com/googleapis/google-cloud-go/issues/1366")

	// Testing against master requires building go code and may take up to 10 minutes.
	// Allow this test to run in parallel with other top level tests to avoid timeouts.
	t.Parallel()
	if testing.Short() {
		t.Skip("skipping profiler integration test in short mode")
	}

	projectID := os.Getenv("GCLOUD_TESTS_GOLANG_PROJECT_ID")
	if projectID == "" {
		t.Skip("skipping profiler integration test when GCLOUD_TESTS_GOLANG_PROJECT_ID variable is not set")
	}

	zone := os.Getenv("GCLOUD_TESTS_GOLANG_PROFILER_ZONE")
	if zone == "" {
		t.Fatalf("GCLOUD_TESTS_GOLANG_PROFILER_ZONE environment variable must be set when integration test is requested")
	}

	// Figure out the Git commit of the current directory. The source checkout in
	// the test VM will run in the same commit. Note that any local changes to
	// the profiler agent won't be tested in the integration test. This flow only
	// works with code that has been committed and pushed to the public repo
	// (either to master or to a branch).
	output, err := exec.Command("git", "rev-parse", "HEAD").CombinedOutput()
	if err != nil {
		t.Fatalf("failed to gather the Git revision of the current source: %v", err)
	}
	commit := strings.Trim(string(output), "\n")
	t.Logf("using Git commit %q for the profiler integration test", commit)

	pst, err := time.LoadLocation("America/Los_Angeles")
	if err != nil {
		t.Fatalf("failed to initiate PST location: %v", err)
	}
	runID := strings.Replace(time.Now().In(pst).Format("2006-01-02-15-04-05.000000-0700"), ".", "-", -1)

	ctx := context.Background()

	client, err := google.DefaultClient(ctx, cloudScope)
	if err != nil {
		t.Fatalf("failed to get default client: %v", err)
	}

	computeService, err := compute.New(client)
	if err != nil {
		t.Fatalf("failed to initialize compute service: %v", err)
	}

	template, err := template.New("startupScript").Parse(startupTemplate)
	if err != nil {
		t.Fatalf("failed to parse startup script template: %v", err)
	}

	tr := proftest.TestRunner{
		Client: client,
	}

	gceTr := proftest.GCETestRunner{
		TestRunner:     tr,
		ComputeService: computeService,
	}

	testcases := []goGCETestCase{
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-gomaster-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-gomaster-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "master",
			mutexProfiling:   true,
		},
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-go111-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-go111-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "1.11",
			mutexProfiling:   true,
		},
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-go110-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-go110-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "1.10",
			mutexProfiling:   true,
		},
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:   projectID,
				Zone:        zone,
				Name:        fmt.Sprintf("profiler-test-go19-%s", runID),
				MachineType: "n1-standard-1",
			},
			name:             fmt.Sprintf("profiler-test-go19-%s-gce", runID),
			wantProfileTypes: []string{"CPU", "HEAP", "THREADS", "CONTENTION", "HEAP_ALLOC"},
			goVersion:        "1.9",
			mutexProfiling:   true,
		},
	}
	// The number of tests run in parallel is the current value of GOMAXPROCS.
	runtime.GOMAXPROCS(len(testcases))
	for _, tc := range testcases {
		tc := tc // capture range variable
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if err := tc.initializeStartupScript(template, commit); err != nil {
				t.Fatalf("failed to initialize startup script")
			}

			if err := gceTr.StartInstance(ctx, &tc.InstanceConfig); err != nil {
				t.Fatal(err)
			}
			defer func() {
				if gceTr.DeleteInstance(ctx, &tc.InstanceConfig); err != nil {
					t.Fatal(err)
				}
			}()

			timeoutCtx, cancel := context.WithTimeout(ctx, time.Minute*25)
			defer cancel()
			if err := gceTr.PollForSerialOutput(timeoutCtx, &tc.InstanceConfig, benchFinishString, errorString); err != nil {
				t.Fatalf("PollForSerialOutput() got error: %v", err)
			}

			timeNow := time.Now()
			endTime := timeNow.Format(time.RFC3339)
			startTime := timeNow.Add(-1 * time.Hour).Format(time.RFC3339)
			for _, pType := range tc.wantProfileTypes {
				pr, err := tr.QueryProfiles(tc.ProjectID, tc.name, startTime, endTime, pType)
				if err != nil {
					t.Errorf("QueryProfiles(%s, %s, %s, %s, %s) got error: %v", tc.ProjectID, tc.name, startTime, endTime, pType, err)
					continue
				}
				if err := pr.HasFunction("busywork"); err != nil {
					t.Error(err)
				}
			}
		})
	}
}
<|endoftext|>