<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package framework

import (
	"fmt"
	"path"
	"strings"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util/wait"
)

// The following upgrade functions are passed into the framework below and used
// to do the actual upgrades.
var MasterUpgrade = func(v string) error {
	switch TestContext.Provider {
	case "gce":
		return masterUpgradeGCE(v)
	case "gke":
		return masterUpgradeGKE(v)
	default:
		return fmt.Errorf("MasterUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
}

func masterUpgradeGCE(rawV string) error {
	v := "v" + rawV
	_, _, err := RunCmd(path.Join(TestContext.RepoRoot, "cluster/gce/upgrade.sh"), "-M", v)
	return err
}

func masterUpgradeGKE(v string) error {
	Logf("Upgrading master to %q", v)
	_, _, err := RunCmd("gcloud", "container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		"--master",
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet")
	return err
}

var NodeUpgrade = func(f *Framework, v string, img string) error {
	// Perform the upgrade.
	var err error
	switch TestContext.Provider {
	case "gce":
		// TODO(maisem): add GCE support for upgrading to different images.
		err = nodeUpgradeGCE(v)
	case "gke":
		err = nodeUpgradeGKE(v, img)
	default:
		err = fmt.Errorf("NodeUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
	if err != nil {
		return err
	}

	// Wait for it to complete and validate nodes are healthy.
	//
	// TODO(ihmccreery) We shouldn't have to wait for nodes to be ready in
	// GKE; the operation shouldn't return until they all are.
	Logf("Waiting up to %v for all nodes to be ready after the upgrade", RestartNodeReadyAgainTimeout)
	if _, err := CheckNodesReady(f.Client, RestartNodeReadyAgainTimeout, TestContext.CloudConfig.NumNodes); err != nil {
		return err
	}
	return nil
}

func nodeUpgradeGCE(rawV string) error {
	// TODO(ihmccreery) This code path should be identical to how a user
	// would trigger a node update; right now it's very different.
	v := "v" + rawV

	Logf("Getting the node template before the upgrade")
	tmplBefore, err := MigTemplate()
	if err != nil {
		return fmt.Errorf("error getting the node template before the upgrade: %v", err)
	}

	Logf("Preparing node upgrade by creating new instance template for %q", v)
	stdout, _, err := RunCmd(path.Join(TestContext.RepoRoot, "cluster/gce/upgrade.sh"), "-P", v)
	if err != nil {
		cleanupNodeUpgradeGCE(tmplBefore)
		return fmt.Errorf("error preparing node upgrade: %v", err)
	}
	tmpl := strings.TrimSpace(stdout)

	Logf("Performing a node upgrade to %q; waiting at most %v per node", tmpl, RestartPerNodeTimeout)
	if err := MigRollingUpdate(tmpl, RestartPerNodeTimeout); err != nil {
		cleanupNodeUpgradeGCE(tmplBefore)
		return fmt.Errorf("error doing node upgrade via a MigRollingUpdate to %s: %v", tmpl, err)
	}
	return nil
}

// MigRollingUpdate starts a MIG rolling update, upgrading the nodes to a new
// instance template named tmpl, and waits up to nt times the number of nodes
// for it to complete.
func MigRollingUpdate(tmpl string, nt time.Duration) error {
	Logf(fmt.Sprintf("starting the MIG rolling update to %s", tmpl))
	id, err := migRollingUpdateStart(tmpl, nt)
	if err != nil {
		return fmt.Errorf("couldn't start the MIG rolling update: %v", err)
	}

	Logf(fmt.Sprintf("polling the MIG rolling update (%s) until it completes", id))
	if err := migRollingUpdatePoll(id, nt); err != nil {
		return fmt.Errorf("err waiting until update completed: %v", err)
	}

	return nil
}

// migRollingUpdateStart (GCE/GKE-only) starts a MIG rolling update using templ
// as the new template, waiting up to nt per node, and returns the ID of that
// update.
func migRollingUpdateStart(templ string, nt time.Duration) (string, error) {
	var errLast error
	var id string
	prefix, suffix := "Started [", "]."
	if err := wait.Poll(Poll, SingleCallTimeout, func() (bool, error) {
		// TODO(mikedanese): make this hit the compute API directly instead of
		//                 shelling out to gcloud.
		// NOTE(mikedanese): If you are changing this gcloud command, update
		//                 cluster/gce/upgrade.sh to match this EXACTLY.
		// A `rolling-updates start` call outputs what we want to stderr.
		_, output, err := retryCmd("gcloud", "alpha", "compute",
			"rolling-updates",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			"start",
			// Required args.
			fmt.Sprintf("--group=%s", TestContext.CloudConfig.NodeInstanceGroup),
			fmt.Sprintf("--template=%s", templ),
			// Optional args to fine-tune behavior.
			fmt.Sprintf("--instance-startup-timeout=%ds", int(nt.Seconds())),
			// NOTE: We can speed up this process by increasing
			//       --max-num-concurrent-instances.
			fmt.Sprintf("--max-num-concurrent-instances=%d", 1),
			fmt.Sprintf("--max-num-failed-instances=%d", 0),
			fmt.Sprintf("--min-instance-update-time=%ds", 0))
		if err != nil {
			errLast = fmt.Errorf("rolling-updates call failed with err: %v", err)
			return false, nil
		}

		// The 'start' call probably succeeded; parse the output and try to find
		// the line that looks like "Started [url/to/<id>]." and return <id>.
		for _, line := range strings.Split(output, "\n") {
			// As a sanity check, ensure the line starts with prefix and ends
			// with suffix.
			if strings.Index(line, prefix) != 0 || strings.Index(line, suffix) != len(line)-len(suffix) {
				continue
			}
			url := strings.Split(strings.TrimSuffix(strings.TrimPrefix(line, prefix), suffix), "/")
			id = url[len(url)-1]
			Logf("Started MIG rolling update; ID: %s", id)
			return true, nil
		}
		errLast = fmt.Errorf("couldn't find line like '%s ... %s' in output to MIG rolling-update start. Output: %s",
			prefix, suffix, output)
		return false, nil
	}); err != nil {
		return "", fmt.Errorf("migRollingUpdateStart() failed with last error: %v", errLast)
	}
	return id, nil
}

// migRollingUpdatePoll (CKE/GKE-only) polls the progress of the MIG rolling
// update with ID id until it is complete. It returns an error if this takes
// longer than nt times the number of nodes.
func migRollingUpdatePoll(id string, nt time.Duration) error {
	// Two keys and a val.
	status, progress, done := "status", "statusMessage", "ROLLED_OUT"
	start, timeout := time.Now(), nt*time.Duration(TestContext.CloudConfig.NumNodes)
	var errLast error
	Logf("Waiting up to %v for MIG rolling update to complete.", timeout)
	if wait.Poll(RestartPoll, timeout, func() (bool, error) {
		// A `rolling-updates describe` call outputs what we want to stdout.
		output, _, err := retryCmd("gcloud", "alpha", "compute",
			"rolling-updates",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			"describe",
			id)
		if err != nil {
			errLast = fmt.Errorf("Error calling rolling-updates describe %s: %v", id, err)
			Logf("%v", errLast)
			return false, nil
		}

		// The 'describe' call probably succeeded; parse the output and try to
		// find the line that looks like "status: <status>" and see whether it's
		// done.
		Logf("Waiting for MIG rolling update: %s (%v elapsed)",
			ParseKVLines(output, progress), time.Since(start))
		if st := ParseKVLines(output, status); st == done {
			return true, nil
		}
		return false, nil
	}) != nil {
		return fmt.Errorf("timeout waiting %v for MIG rolling update to complete. Last error: %v", timeout, errLast)
	}
	Logf("MIG rolling update complete after %v", time.Since(start))
	return nil
}

func cleanupNodeUpgradeGCE(tmplBefore string) {
	Logf("Cleaning up any unused node templates")
	tmplAfter, err := MigTemplate()
	if err != nil {
		Logf("Could not get node template post-upgrade; may have leaked template %s", tmplBefore)
		return
	}
	if tmplBefore == tmplAfter {
		// The node upgrade failed so there's no need to delete
		// anything.
		Logf("Node template %s is still in use; not cleaning up", tmplBefore)
		return
	}
	Logf("Deleting node template %s", tmplBefore)
	if _, _, err := retryCmd("gcloud", "compute", "instance-templates",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		"delete",
		tmplBefore); err != nil {
		Logf("gcloud compute instance-templates delete %s call failed with err: %v", tmplBefore, err)
		Logf("May have leaked instance template %q", tmplBefore)
	}
}

func nodeUpgradeGKE(v string, img string) error {
	Logf("Upgrading nodes to version %q and image %q", v, img)
	args := []string{
		"container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet",
	}
	if len(img) > 0 {
		args = append(args, fmt.Sprintf("--image-type=%s", img))
	}
	_, _, err := RunCmd("gcloud", args...)
	return err
}

// CheckNodesReady waits up to nt for expect nodes accessed by c to be ready,
// returning an error if this doesn't happen in time. It returns the names of
// nodes it finds.
func CheckNodesReady(c *client.Client, nt time.Duration, expect int) ([]string, error) {
	// First, keep getting all of the nodes until we get the number we expect.
	var nodeList *api.NodeList
	var errLast error
	start := time.Now()
	found := wait.Poll(Poll, nt, func() (bool, error) {
		// A rolling-update (GCE/GKE implementation of restart) can complete before the apiserver
		// knows about all of the nodes. Thus, we retry the list nodes call
		// until we get the expected number of nodes.
		nodeList, errLast = c.Nodes().List(api.ListOptions{
			FieldSelector: fields.Set{"spec.unschedulable": "false"}.AsSelector()})
		if errLast != nil {
			return false, nil
		}
		if len(nodeList.Items) != expect {
			errLast = fmt.Errorf("expected to find %d nodes but found only %d (%v elapsed)",
				expect, len(nodeList.Items), time.Since(start))
			Logf("%v", errLast)
			return false, nil
		}
		return true, nil
	}) == nil
	nodeNames := make([]string, len(nodeList.Items))
	for i, n := range nodeList.Items {
		nodeNames[i] = n.ObjectMeta.Name
	}
	if !found {
		return nodeNames, fmt.Errorf("couldn't find %d nodes within %v; last error: %v",
			expect, nt, errLast)
	}
	Logf("Successfully found %d nodes", expect)

	// Next, ensure in parallel that all the nodes are ready. We subtract the
	// time we spent waiting above.
	timeout := nt - time.Since(start)
	result := make(chan bool, len(nodeList.Items))
	for _, n := range nodeNames {
		n := n
		go func() { result <- WaitForNodeToBeReady(c, n, timeout) }()
	}
	failed := false
	// TODO(mbforbes): Change to `for range` syntax once we support only Go
	// >= 1.4.
	for i := range nodeList.Items {
		_ = i
		if !<-result {
			failed = true
		}
	}
	if failed {
		return nodeNames, fmt.Errorf("at least one node failed to be ready")
	}
	return nodeNames, nil
}

// MigTemplate (GCE-only) returns the name of the MIG template that the
// nodes of the cluster use.
func MigTemplate() (string, error) {
	var errLast error
	var templ string
	key := "instanceTemplate"
	if wait.Poll(Poll, SingleCallTimeout, func() (bool, error) {
		// TODO(mikedanese): make this hit the compute API directly instead of
		// shelling out to gcloud.
		// An `instance-groups managed describe` call outputs what we want to stdout.
		output, _, err := retryCmd("gcloud", "compute", "instance-groups", "managed",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			"describe",
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			TestContext.CloudConfig.NodeInstanceGroup)
		if err != nil {
			errLast = fmt.Errorf("gcloud compute instance-groups managed describe call failed with err: %v", err)
			return false, nil
		}

		// The 'describe' call probably succeeded; parse the output and try to
		// find the line that looks like "instanceTemplate: url/to/<templ>" and
		// return <templ>.
		if val := ParseKVLines(output, key); len(val) > 0 {
			url := strings.Split(val, "/")
			templ = url[len(url)-1]
			Logf("MIG group %s using template: %s", TestContext.CloudConfig.NodeInstanceGroup, templ)
			return true, nil
		}
		errLast = fmt.Errorf("couldn't find %s in output to get MIG template. Output: %s", key, output)
		return false, nil
	}) != nil {
		return "", fmt.Errorf("MigTemplate() failed with last error: %v", errLast)
	}
	return templ, nil
}
<commit_msg>Fix node upgrade in e2e tests<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package framework

import (
	"fmt"
	"path"
	"strings"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util/wait"
)

// The following upgrade functions are passed into the framework below and used
// to do the actual upgrades.
var MasterUpgrade = func(v string) error {
	switch TestContext.Provider {
	case "gce":
		return masterUpgradeGCE(v)
	case "gke":
		return masterUpgradeGKE(v)
	default:
		return fmt.Errorf("MasterUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
}

func masterUpgradeGCE(rawV string) error {
	v := "v" + rawV
	_, _, err := RunCmd(path.Join(TestContext.RepoRoot, "cluster/gce/upgrade.sh"), "-M", v)
	return err
}

func masterUpgradeGKE(v string) error {
	Logf("Upgrading master to %q", v)
	_, _, err := RunCmd("gcloud", "container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		"--master",
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet")
	return err
}

var NodeUpgrade = func(f *Framework, v string, img string) error {
	// Perform the upgrade.
	var err error
	switch TestContext.Provider {
	case "gce":
		// TODO(maisem): add GCE support for upgrading to different images.
		err = nodeUpgradeGCE(v)
	case "gke":
		err = nodeUpgradeGKE(v, img)
	default:
		err = fmt.Errorf("NodeUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
	if err != nil {
		return err
	}

	// Wait for it to complete and validate nodes are healthy.
	//
	// TODO(ihmccreery) We shouldn't have to wait for nodes to be ready in
	// GKE; the operation shouldn't return until they all are.
	Logf("Waiting up to %v for all nodes to be ready after the upgrade", RestartNodeReadyAgainTimeout)
	if _, err := CheckNodesReady(f.Client, RestartNodeReadyAgainTimeout, TestContext.CloudConfig.NumNodes); err != nil {
		return err
	}
	return nil
}

func nodeUpgradeGCE(rawV string) error {
	v := "v" + rawV
	_, _, err := RunCmd(path.Join(TestContext.RepoRoot, "cluster/gce/upgrade.sh"), "-N", v)
	return err
}

// MigRollingUpdate starts a MIG rolling update, upgrading the nodes to a new
// instance template named tmpl, and waits up to nt times the number of nodes
// for it to complete.
func MigRollingUpdate(tmpl string, nt time.Duration) error {
	Logf(fmt.Sprintf("starting the MIG rolling update to %s", tmpl))
	id, err := migRollingUpdateStart(tmpl, nt)
	if err != nil {
		return fmt.Errorf("couldn't start the MIG rolling update: %v", err)
	}

	Logf(fmt.Sprintf("polling the MIG rolling update (%s) until it completes", id))
	if err := migRollingUpdatePoll(id, nt); err != nil {
		return fmt.Errorf("err waiting until update completed: %v", err)
	}

	return nil
}

// migRollingUpdateStart (GCE/GKE-only) starts a MIG rolling update using templ
// as the new template, waiting up to nt per node, and returns the ID of that
// update.
func migRollingUpdateStart(templ string, nt time.Duration) (string, error) {
	var errLast error
	var id string
	prefix, suffix := "Started [", "]."
	if err := wait.Poll(Poll, SingleCallTimeout, func() (bool, error) {
		// TODO(mikedanese): make this hit the compute API directly instead of
		//                 shelling out to gcloud.
		// NOTE(mikedanese): If you are changing this gcloud command, update
		//                 cluster/gce/upgrade.sh to match this EXACTLY.
		// A `rolling-updates start` call outputs what we want to stderr.
		_, output, err := retryCmd("gcloud", "alpha", "compute",
			"rolling-updates",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			"start",
			// Required args.
			fmt.Sprintf("--group=%s", TestContext.CloudConfig.NodeInstanceGroup),
			fmt.Sprintf("--template=%s", templ),
			// Optional args to fine-tune behavior.
			fmt.Sprintf("--instance-startup-timeout=%ds", int(nt.Seconds())),
			// NOTE: We can speed up this process by increasing
			//       --max-num-concurrent-instances.
			fmt.Sprintf("--max-num-concurrent-instances=%d", 1),
			fmt.Sprintf("--max-num-failed-instances=%d", 0),
			fmt.Sprintf("--min-instance-update-time=%ds", 0))
		if err != nil {
			errLast = fmt.Errorf("rolling-updates call failed with err: %v", err)
			return false, nil
		}

		// The 'start' call probably succeeded; parse the output and try to find
		// the line that looks like "Started [url/to/<id>]." and return <id>.
		for _, line := range strings.Split(output, "\n") {
			// As a sanity check, ensure the line starts with prefix and ends
			// with suffix.
			if strings.Index(line, prefix) != 0 || strings.Index(line, suffix) != len(line)-len(suffix) {
				continue
			}
			url := strings.Split(strings.TrimSuffix(strings.TrimPrefix(line, prefix), suffix), "/")
			id = url[len(url)-1]
			Logf("Started MIG rolling update; ID: %s", id)
			return true, nil
		}
		errLast = fmt.Errorf("couldn't find line like '%s ... %s' in output to MIG rolling-update start. Output: %s",
			prefix, suffix, output)
		return false, nil
	}); err != nil {
		return "", fmt.Errorf("migRollingUpdateStart() failed with last error: %v", errLast)
	}
	return id, nil
}

// migRollingUpdatePoll (CKE/GKE-only) polls the progress of the MIG rolling
// update with ID id until it is complete. It returns an error if this takes
// longer than nt times the number of nodes.
func migRollingUpdatePoll(id string, nt time.Duration) error {
	// Two keys and a val.
	status, progress, done := "status", "statusMessage", "ROLLED_OUT"
	start, timeout := time.Now(), nt*time.Duration(TestContext.CloudConfig.NumNodes)
	var errLast error
	Logf("Waiting up to %v for MIG rolling update to complete.", timeout)
	if wait.Poll(RestartPoll, timeout, func() (bool, error) {
		// A `rolling-updates describe` call outputs what we want to stdout.
		output, _, err := retryCmd("gcloud", "alpha", "compute",
			"rolling-updates",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			"describe",
			id)
		if err != nil {
			errLast = fmt.Errorf("Error calling rolling-updates describe %s: %v", id, err)
			Logf("%v", errLast)
			return false, nil
		}

		// The 'describe' call probably succeeded; parse the output and try to
		// find the line that looks like "status: <status>" and see whether it's
		// done.
		Logf("Waiting for MIG rolling update: %s (%v elapsed)",
			ParseKVLines(output, progress), time.Since(start))
		if st := ParseKVLines(output, status); st == done {
			return true, nil
		}
		return false, nil
	}) != nil {
		return fmt.Errorf("timeout waiting %v for MIG rolling update to complete. Last error: %v", timeout, errLast)
	}
	Logf("MIG rolling update complete after %v", time.Since(start))
	return nil
}

func cleanupNodeUpgradeGCE(tmplBefore string) {
	Logf("Cleaning up any unused node templates")
	tmplAfter, err := MigTemplate()
	if err != nil {
		Logf("Could not get node template post-upgrade; may have leaked template %s", tmplBefore)
		return
	}
	if tmplBefore == tmplAfter {
		// The node upgrade failed so there's no need to delete
		// anything.
		Logf("Node template %s is still in use; not cleaning up", tmplBefore)
		return
	}
	Logf("Deleting node template %s", tmplBefore)
	if _, _, err := retryCmd("gcloud", "compute", "instance-templates",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		"delete",
		tmplBefore); err != nil {
		Logf("gcloud compute instance-templates delete %s call failed with err: %v", tmplBefore, err)
		Logf("May have leaked instance template %q", tmplBefore)
	}
}

func nodeUpgradeGKE(v string, img string) error {
	Logf("Upgrading nodes to version %q and image %q", v, img)
	args := []string{
		"container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet",
	}
	if len(img) > 0 {
		args = append(args, fmt.Sprintf("--image-type=%s", img))
	}
	_, _, err := RunCmd("gcloud", args...)
	return err
}

// CheckNodesReady waits up to nt for expect nodes accessed by c to be ready,
// returning an error if this doesn't happen in time. It returns the names of
// nodes it finds.
func CheckNodesReady(c *client.Client, nt time.Duration, expect int) ([]string, error) {
	// First, keep getting all of the nodes until we get the number we expect.
	var nodeList *api.NodeList
	var errLast error
	start := time.Now()
	found := wait.Poll(Poll, nt, func() (bool, error) {
		// A rolling-update (GCE/GKE implementation of restart) can complete before the apiserver
		// knows about all of the nodes. Thus, we retry the list nodes call
		// until we get the expected number of nodes.
		nodeList, errLast = c.Nodes().List(api.ListOptions{
			FieldSelector: fields.Set{"spec.unschedulable": "false"}.AsSelector()})
		if errLast != nil {
			return false, nil
		}
		if len(nodeList.Items) != expect {
			errLast = fmt.Errorf("expected to find %d nodes but found only %d (%v elapsed)",
				expect, len(nodeList.Items), time.Since(start))
			Logf("%v", errLast)
			return false, nil
		}
		return true, nil
	}) == nil
	nodeNames := make([]string, len(nodeList.Items))
	for i, n := range nodeList.Items {
		nodeNames[i] = n.ObjectMeta.Name
	}
	if !found {
		return nodeNames, fmt.Errorf("couldn't find %d nodes within %v; last error: %v",
			expect, nt, errLast)
	}
	Logf("Successfully found %d nodes", expect)

	// Next, ensure in parallel that all the nodes are ready. We subtract the
	// time we spent waiting above.
	timeout := nt - time.Since(start)
	result := make(chan bool, len(nodeList.Items))
	for _, n := range nodeNames {
		n := n
		go func() { result <- WaitForNodeToBeReady(c, n, timeout) }()
	}
	failed := false
	// TODO(mbforbes): Change to `for range` syntax once we support only Go
	// >= 1.4.
	for i := range nodeList.Items {
		_ = i
		if !<-result {
			failed = true
		}
	}
	if failed {
		return nodeNames, fmt.Errorf("at least one node failed to be ready")
	}
	return nodeNames, nil
}

// MigTemplate (GCE-only) returns the name of the MIG template that the
// nodes of the cluster use.
func MigTemplate() (string, error) {
	var errLast error
	var templ string
	key := "instanceTemplate"
	if wait.Poll(Poll, SingleCallTimeout, func() (bool, error) {
		// TODO(mikedanese): make this hit the compute API directly instead of
		// shelling out to gcloud.
		// An `instance-groups managed describe` call outputs what we want to stdout.
		output, _, err := retryCmd("gcloud", "compute", "instance-groups", "managed",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			"describe",
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			TestContext.CloudConfig.NodeInstanceGroup)
		if err != nil {
			errLast = fmt.Errorf("gcloud compute instance-groups managed describe call failed with err: %v", err)
			return false, nil
		}

		// The 'describe' call probably succeeded; parse the output and try to
		// find the line that looks like "instanceTemplate: url/to/<templ>" and
		// return <templ>.
		if val := ParseKVLines(output, key); len(val) > 0 {
			url := strings.Split(val, "/")
			templ = url[len(url)-1]
			Logf("MIG group %s using template: %s", TestContext.CloudConfig.NodeInstanceGroup, templ)
			return true, nil
		}
		errLast = fmt.Errorf("couldn't find %s in output to get MIG template. Output: %s", key, output)
		return false, nil
	}) != nil {
		return "", fmt.Errorf("MigTemplate() failed with last error: %v", errLast)
	}
	return templ, nil
}
<|endoftext|>


<commit_before>package conf

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

const (
	srcConf  = "load_conf_test.conf"
	destConf = "load_conf_test2.conf"
)

func doTestTrimComments(t *testing.T) {
	ast := assert.New(t)
	var (
		err         error
		b, b1       []byte
		conf, conf1 interface{}
	)
	err = LoadEx(&conf, srcConf)
	ast.Nil(err)
	tb, err := ioutil.ReadFile(destConf)
	assert.Nil(t, err)
	err = json.Unmarshal(tb, &conf1)
	ast.Nil(err)
	b, err = json.Marshal(conf)
	ast.Nil(err)
	b1, err = json.Marshal(conf1)
	ast.Nil(err)
	ast.Equal(b1, b)
	return
}

func TestDo(t *testing.T) {
	doTestTrimComments(t)
}

func Test_getEnv(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name   []string
		expect string
	}{
		{},
		{
			name:   []string{"name1", "name2", "name3"},
			expect: "name1name2name3",
		},
		{
			name:   []string{"name1"},
			expect: "name1",
		},
	}

	for _, test := range tests {
		for _, n := range test.name {
			os.Setenv(n, n)
		}
		actual := getEnv(test.name)
		assert.EqualValues(t, test.expect, actual)
	}
}

func TestGetConfigDir(t *testing.T) {
	t.Parallel()
	tests := []struct {
		app         string
		homeEnvName string
		expectErr   error
		expectDir   string
	}{
		{
			app:         "",
			homeEnvName: "",
			expectErr:   ErrHomeNotFound,
			expectDir:   "",
		},
		{
			app:         "app",
			homeEnvName: "TEST",
			expectErr:   nil,
			expectDir:   "TEST/.app",
		},
	}

	for _, test := range tests {
		tmp := homeEnvNames
		homeEnvNames = [][]string{{test.homeEnvName}}
		os.Setenv(test.homeEnvName, test.homeEnvName)
		actualDir, actualErr := GetConfigDir(test.app)
		assert.EqualValues(t, test.expectErr, actualErr)
		assert.EqualValues(t, test.expectDir, actualDir)
		homeEnvNames = tmp
	}
}

func TestInit(t *testing.T) {
	t.Parallel()
	tests := []struct {
		cflag       string
		app         string
		defaultConf string
		homeEnvName string
		expect      string
	}{
		{
			expect: "/",
		},
		{
			cflag:       "f",
			app:         "app",
			defaultConf: "defaultConf",
			homeEnvName: "TEST",
			expect:      "TEST/.app/TEST",
		},
	}
	for _, test := range tests {
		tmp := homeEnvNames
		homeEnvNames = [][]string{{test.homeEnvName}}
		os.Setenv(test.homeEnvName, test.homeEnvName)
		Init(test.cflag, test.app, test.homeEnvName)
		assert.NotEmpty(t, *confName)
		homeEnvNames = tmp
	}
}

type Config struct {
	MaxProcs   int      `json:"max_procs"`
	DebugLevel int      `json:"debug_level"`
	ConfsPath  []string `json:"confs_path"`
}

func TestLoad(t *testing.T) {
	t.Parallel()
	tests := []struct {
		content  string
		confName string
		conf     interface{}
		expect   error
	}{
		{
			content: "",
			conf:    &Config{},
			expect:  errors.New("no such file or directory"),
		},
		{
			confName: "./TestLoad",
			content: `{
"max_procs": 2,
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`,
			conf: &Config{
				MaxProcs:   2,
				DebugLevel: 1,
				ConfsPath:  []string{"conf1", "conf2"},
			},
			expect: nil,
		},
	}

	for _, test := range tests {
		confName = &test.confName
		if test.confName != "" {
			err := ioutil.WriteFile(test.confName, []byte(test.content), 0600)
			assert.Nil(t, err)
		}
		err := Load(test.conf)
		if err != nil {
			t.Log("load conf  ", test.confName, " failed: ", err)
			if !strings.Contains(err.Error(), test.expect.Error()) {
				t.Fatalf("expect contains: %s, but got: %s", test.expect.Error(), err.Error())
			}
		}
		os.RemoveAll(test.confName)
	}
}

func TestLoadEx(t *testing.T) {
	t.Parallel()
	tests := []struct {
		content  string
		confName string
		conf     interface{}
		expect   error
	}{
		{
			content: "",
			conf:    &Config{},
			expect:  errors.New("open : no such file or directory"),
		},
		{
			confName: "./TestLoadEx",
			content: `{
"max_procs": 2,
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`,
			conf: &Config{
				MaxProcs:   2,
				DebugLevel: 1,
				ConfsPath:  []string{"conf1", "conf2"},
			},
			expect: nil,
		},
		{
			confName: "./TestLoadEx",
			content: `{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`,
			conf: &Config{
				MaxProcs:   2,
				DebugLevel: 1,
				ConfsPath:  []string{"conf1", "conf2"},
			},
			expect: errors.New("json: cannot unmarshal string into Go struct field Config.max_procs of type int"),
		},
	}

	for _, test := range tests {
		if test.confName != "" {
			err := ioutil.WriteFile(test.confName, []byte(test.content), 0600)
			assert.Nil(t, err)
		}
		err := LoadEx(test.conf, test.confName)
		assert.EqualValues(t, fmt.Sprintf("%v", test.expect), fmt.Sprintf("%v", err))
		os.RemoveAll(test.confName)
	}
}

func Test_trimComments(t *testing.T) {
	t.Parallel()
	tests := []struct {
		data   []byte
		expect []byte
	}{
		{
			data: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
			expect: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
		},
	}

	for _, test := range tests {
		actual := trimComments(test.data)
		assert.EqualValues(t, test.expect, actual)
	}

}

func Test_trimCommentsLine(t *testing.T) {
	tests := []struct {
		line   []byte
		expect []byte
	}{
		{
			line: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
			expect: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
		},
	}

	for _, test := range tests {
		actual := trimCommentsLine(test.line)
		assert.EqualValues(t, test.expect, actual)
	}
}

func TestLoadFile(t *testing.T) {
	t.Parallel()
	confName := "./TestLoadFile"
	conf := &Config{}
	err := ioutil.WriteFile(confName, []byte(`{
"max_procs": 2,
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`), 0755)
	assert.Nil(t, err)
	defer os.RemoveAll(confName)
	err = LoadFile(conf, confName)
	assert.Nil(t, err)
	assert.EqualValues(t, []string{"conf1", "conf2"}, conf.ConfsPath)
	assert.EqualValues(t, 2, conf.MaxProcs)
	assert.EqualValues(t, 1, conf.DebugLevel)
}
<commit_msg>fix utest<commit_after>package conf

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

const (
	srcConf  = "load_conf_test.conf"
	destConf = "load_conf_test2.conf"
)

func doTestTrimComments(t *testing.T) {
	ast := assert.New(t)
	var (
		err         error
		b, b1       []byte
		conf, conf1 interface{}
	)
	err = LoadEx(&conf, srcConf)
	ast.Nil(err)
	tb, err := ioutil.ReadFile(destConf)
	assert.Nil(t, err)
	err = json.Unmarshal(tb, &conf1)
	ast.Nil(err)
	b, err = json.Marshal(conf)
	ast.Nil(err)
	b1, err = json.Marshal(conf1)
	ast.Nil(err)
	ast.Equal(b1, b)
	return
}

func TestDo(t *testing.T) {
	doTestTrimComments(t)
}

func Test_getEnv(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name   []string
		expect string
	}{
		{},
		{
			name:   []string{"name1", "name2", "name3"},
			expect: "name1name2name3",
		},
		{
			name:   []string{"name1"},
			expect: "name1",
		},
	}

	for _, test := range tests {
		for _, n := range test.name {
			os.Setenv(n, n)
		}
		actual := getEnv(test.name)
		assert.EqualValues(t, test.expect, actual)
	}
}

func TestGetConfigDir(t *testing.T) {
	t.Parallel()
	tests := []struct {
		app         string
		homeEnvName string
		expectErr   error
		expectDir   string
	}{
		{
			app:         "",
			homeEnvName: "",
			expectErr:   ErrHomeNotFound,
			expectDir:   "",
		},
		{
			app:         "app",
			homeEnvName: "TEST",
			expectErr:   nil,
			expectDir:   "TEST/.app",
		},
	}

	for _, test := range tests {
		tmp := homeEnvNames
		homeEnvNames = [][]string{{test.homeEnvName}}
		os.Setenv(test.homeEnvName, test.homeEnvName)
		actualDir, actualErr := GetConfigDir(test.app)
		assert.EqualValues(t, test.expectErr, actualErr)
		assert.EqualValues(t, test.expectDir, actualDir)
		homeEnvNames = tmp
	}
}

func TestInit(t *testing.T) {
	t.Parallel()
	tests := []struct {
		cflag       string
		app         string
		defaultConf string
		homeEnvName string
		expect      string
	}{
		{
			expect: "/",
		},
		{
			cflag:       "f",
			app:         "app",
			defaultConf: "defaultConf",
			homeEnvName: "TEST",
			expect:      "TEST/.app/TEST",
		},
	}
	for _, test := range tests {
		tmp := homeEnvNames
		homeEnvNames = [][]string{{test.homeEnvName}}
		os.Setenv(test.homeEnvName, test.homeEnvName)
		Init(test.cflag, test.app, test.homeEnvName)
		assert.NotEmpty(t, *confName)
		homeEnvNames = tmp
	}
}

type Config struct {
	MaxProcs   int      `json:"max_procs"`
	DebugLevel int      `json:"debug_level"`
	ConfsPath  []string `json:"confs_path"`
}

func TestLoad(t *testing.T) {
	t.Parallel()
	tests := []struct {
		content  string
		confName string
		conf     interface{}
		expect   error
	}{
		{
			content: "",
			conf:    &Config{},
			expect:  errors.New("no such file or directory"),
		},
		{
			confName: "./TestLoad",
			content: `{
"max_procs": 2,
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`,
			conf: &Config{
				MaxProcs:   2,
				DebugLevel: 1,
				ConfsPath:  []string{"conf1", "conf2"},
			},
			expect: nil,
		},
	}

	for _, test := range tests {
		confName = &test.confName
		if test.confName != "" {
			err := ioutil.WriteFile(test.confName, []byte(test.content), 0600)
			assert.Nil(t, err)
			time.Sleep(3 * time.Second)
		}
		err := Load(test.conf)
		if err != nil {
			t.Log("load conf  ", test.confName, " failed: ", err)
			if !strings.Contains(err.Error(), test.expect.Error()) {
				t.Fatalf("expect contains: %s, but got: %s", test.expect.Error(), err.Error())
			}
		}
		os.RemoveAll(test.confName)
	}
}

func TestLoadEx(t *testing.T) {
	t.Parallel()
	tests := []struct {
		content  string
		confName string
		conf     interface{}
		expect   error
	}{
		{
			content: "",
			conf:    &Config{},
			expect:  errors.New("open : no such file or directory"),
		},
		{
			confName: "./TestLoadEx",
			content: `{
"max_procs": 2,
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`,
			conf: &Config{
				MaxProcs:   2,
				DebugLevel: 1,
				ConfsPath:  []string{"conf1", "conf2"},
			},
			expect: nil,
		},
		{
			confName: "./TestLoadEx",
			content: `{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`,
			conf: &Config{
				MaxProcs:   2,
				DebugLevel: 1,
				ConfsPath:  []string{"conf1", "conf2"},
			},
			expect: errors.New("json: cannot unmarshal string into Go struct field Config.max_procs of type int"),
		},
	}

	for _, test := range tests {
		if test.confName != "" {
			err := ioutil.WriteFile(test.confName, []byte(test.content), 0600)
			assert.Nil(t, err)
		}
		err := LoadEx(test.conf, test.confName)
		assert.EqualValues(t, fmt.Sprintf("%v", test.expect), fmt.Sprintf("%v", err))
		os.RemoveAll(test.confName)
	}
}

func Test_trimComments(t *testing.T) {
	t.Parallel()
	tests := []struct {
		data   []byte
		expect []byte
	}{
		{
			data: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
			expect: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
		},
	}

	for _, test := range tests {
		actual := trimComments(test.data)
		assert.EqualValues(t, test.expect, actual)
	}

}

func Test_trimCommentsLine(t *testing.T) {
	tests := []struct {
		line   []byte
		expect []byte
	}{
		{
			line: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
			expect: []byte(`{
"max_procs": "2",
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`),
		},
	}

	for _, test := range tests {
		actual := trimCommentsLine(test.line)
		assert.EqualValues(t, test.expect, actual)
	}
}

func TestLoadFile(t *testing.T) {
	t.Parallel()
	confName := "./TestLoadFile"
	conf := &Config{}
	err := ioutil.WriteFile(confName, []byte(`{
"max_procs": 2,
"debug_level": 1,
"confs_path": ["conf1","conf2"]
}`), 0755)
	assert.Nil(t, err)
	defer os.RemoveAll(confName)
	err = LoadFile(conf, confName)
	assert.Nil(t, err)
	assert.EqualValues(t, []string{"conf1", "conf2"}, conf.ConfsPath)
	assert.EqualValues(t, 2, conf.MaxProcs)
	assert.EqualValues(t, 1, conf.DebugLevel)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"fmt"
	"io"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
)

var TAX_CONFIG = "TAX_CONFIG.csv"

func main() {
	if len(os.Args) <= 1 {
		fmt.Println("Usage: > go run PayrollProcessor.go <inputfile>")
		return
	}

	taxBrackets, err := readTaxBracketsConfig(TAX_CONFIG)
	if err != nil {
		// error reading tax bracket config - quit
		fmt.Errorf("Error reading tax brackets config: %v", err)
		return
	}

	payrollRecords, err := readPayrollRecords(os.Args[1])
	if err != nil {
		fmt.Errorf("%v", err)
	}

	for k, v := range payrollRecords {
		fmt.Printf("%2d. ", k)

		if v != nil {
			v.Print()
		} else {
			fmt.Printf("Nil record object\n")
		}
	}
}

func readPayrollRecords(inputFile string) ([]*PayrollRecord, error) {
	fileHandle, err := os.Open(inputFile)
	if err != nil {
		return nil, err
	}
	defer fileHandle.Close()

	csvReader := csv.NewReader(fileHandle)
	records := []*PayrollRecord{}

	for {
		row, err := csvReader.Read()
		if err != nil {
			if err == io.EOF {
				err = nil
			}

			// TODO: returning at the first erroneous row? can we skip over til EOF?
			return records, err
		}

		// TODO: consider sending entire input row to recod creation function
		if len(row) < 5 { // need at least five fields in a valid record, also check for empty fields
			// return or collect error and move to next record
			// return records, fmt.Errorf("Minimum number of fields not met in input <%s>\n", row)
			continue // read next row
		}

		newPayrollRecord, err := createPayrollRecord(row[0], row[1], row[2], row[3], row[4])
		if err != nil {
			// throw error
			// fmt.Errorf("Error creating payroll record with input %s\n", row)
			// return records, fmt.Errorf("Minimum number of fields not met in input <%s>\n", row)
			continue // read next row
		}

		records = append(records, newPayrollRecord)
	}
}

type PayrollRecord struct {
	FirstName    string
	LastName     string
	AnnualSalary float64
	SuperRate    float64
	PaymentDate  string
	Valid        bool   //	indicates if the record object is valid
	ErrorStr     string // if Valid == false, contains the input data from the input file leading to invalid object
}

func (rec *PayrollRecord) fullName() string {
	return rec.FirstName + " " + rec.LastName
}

func (rec *PayrollRecord) payPeriod() string {
	return rec.PaymentDate
}

func (rec *PayrollRecord) grossIncome() float64 {
	return round(rec.AnnualSalary / 12)
}

func (rec *PayrollRecord) incomeTax() {

}

func (rec *PayrollRecord) netIncome() {

}

func (rec *PayrollRecord) superAmount() {

}

func (rec *PayrollRecord) Print() {
	if rec.Valid {
		fmt.Printf("%s, %s, %.2f, %.2f%%, %s, %f\n", rec.FirstName, rec.LastName, rec.AnnualSalary, rec.SuperRate, rec.PaymentDate, rec.grossIncome())
	} else {
		fmt.Printf("Invalid record %s\n", rec.ErrorStr)
	}
}

func createPayrollRecord(firstname_str string, lastname_str string, annualsalary_str string, superrate_str string, paymentdate_str string) (*PayrollRecord, error) {
	var newRecord PayrollRecord

	// prepare input
	FirstName := strings.TrimSpace(firstname_str)
	LastName := strings.TrimSpace(lastname_str)
	AnnualSalary, err_sal := strconv.ParseFloat(annualsalary_str, 64)
	SuperRate := strings.TrimSpace(superrate_str)
	PaymentDate := strings.TrimSpace(paymentdate_str)

	// validate [TODO: may improve by returning specific error, consider ignoring erroneous record and reading ahead]
	// extract numeric super percentage
	rexp, _ := regexp.Compile("^[0-9]+")
	SuperRate_f, err_sr := strconv.ParseFloat(rexp.FindString(SuperRate), 64)

	if FirstName == "" || LastName == "" || (err_sal != nil) || (err_sr != nil) || AnnualSalary <= 0 || SuperRate_f < 0 || SuperRate_f > 50 {
		newRecord.Valid = false
		newRecord.ErrorStr = fmt.Sprintf("Invalid input record: [%s] [%s] [%s] [%s] [%s]", firstname_str, lastname_str, annualsalary_str, superrate_str, paymentdate_str)
		return &newRecord, fmt.Errorf("Invalid data in payroll input record\n")
	}

	newRecord.FirstName = FirstName
	newRecord.LastName = LastName
	newRecord.AnnualSalary = AnnualSalary
	newRecord.SuperRate = SuperRate_f
	newRecord.PaymentDate = PaymentDate
	newRecord.Valid = true

	return &newRecord, nil
}

func round(n float64) float64 {
	n_floor := math.Floor(n)
	n_dec := n - n_floor

	if n_dec >= .50 {
		// round up
		return n_floor + 1.0
	}

	// round down
	return n_floor
}

// takes a salary amount and returns percentage, lump payment and value above which the percebntage should be calculated
// func getTaxPercent(sal float64) (float64, float64, float64) {
// 	taxmap := map[float64]float64{0: 18200, }
// }

type IncomeTaxBracket struct {
	Lower   float64 // lower salary limit of tax bracket
	Upper   float64 // upper salary limit of tax bracket
	Percent float64 // percentage tax to be levied above a certain threshold
	Lump    float64 // lump sum to be paid for this bracket, if any
	Above   float64 // threshold above which percentage tax has to be paid
}

func readTaxBracketsConfig(inputFile string) ([]*IncomeTaxBracket, error) {
	fileHandle, err := os.Open(inputFile)
	if err != nil {
		return nil, err
	}
	defer fileHandle.Close()

	csvReader := csv.NewReader(fileHandle)
	brackets := []*IncomeTaxBracket{}

	i := 0
	prev_upper := 0.0
	for {
		row, err := csvReader.Read()
		if err != nil {
			if err == io.EOF {
				err = nil
			}

			return brackets, err
		}

		// TODO: consider sending entire input row to recod creation function
		if len(row) < 5 { // need at least five fields in a valid record, also check for empty fields
			// return or collect error and move to next record
			return nil, fmt.Errorf("readTaxBrackets(): Minimum number of fields not met in input <%s>\n", row)
		}

		lower, err_low := strconv.ParseFloat(row[0], 64)

		top := false // flag indicating topmost tax bracket
		upp := strings.TrimSpace(row[1])
		upper := 0.0
		var err_upp error
		if upp != "" {
			// upper field could be empty if this is the top bracket
			upper, err_upp = strconv.ParseFloat(row[1], 64)
		} else {
			// set flag indicating uppermost bracket
			top = true
		}

		percent, err_perc := strconv.ParseFloat(row[2], 64)
		lump, err_lump := strconv.ParseFloat(row[3], 64)
		threshold, err_thr := strconv.ParseFloat(row[4], 64)

		if err_low != nil || err_upp != nil || err_perc != nil || err_lump != nil || err_thr != nil {
			return nil, fmt.Errorf("readTaxBrackets(): Error reading tax bracket config record: <%s>\n", row)
		}

		if lower >= upper {
			return nil, fmt.Errorf("readTaxBrackets(): Lower limit >= upper limit in input <%s>\n", row)
		}

		if i == 0 {
			if lower != 0 {
				return nil, fmt.Errorf("readTaxBrackets(): First bracket lower limit != 0 in input <%s>\n", row)
			}
		} else {
			if lower <= prev_upper {
				return nil, fmt.Errorf("readTaxBrackets(): Current bracket's lower limit <= previous upper limit in input <%s>\n", row)
			}
		}

		prev_upper = upper
		newTaxBracket := &IncomeTaxBracket{lower, upper, percent, lump, threshold}
		brackets = append(brackets, newTaxBracket)
		i++

		if top {
			break
		}
	}

	return brackets, nil
}
<commit_msg>Second commit<commit_after>package main

import (
	"encoding/csv"
	"fmt"
	"io"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
)

var TAX_CONFIG_FILE = "TAX_CONFIG.csv"
var TAX_BRACKETS []*IncomeTaxBracket

func main() {
	if len(os.Args) <= 1 {
		fmt.Println("Usage: > go run PayrollProcessor.go <inputfile>")
		return
	}

	TAX_BRACKETS, err := readTaxBracketsConfig(TAX_CONFIG_FILE)
	if err != nil {
		// error reading tax bracket config - quit
		fmt.Printf("Error reading tax brackets config: %v", err)
		return
	}

	for _, v := range TAX_BRACKETS {
		v.Print();
	}

	payrollRecords, err := readPayrollRecords(os.Args[1])
	if err != nil {
		fmt.Printf("Error reading payroll records: %v", err)
	}

	for k, v := range payrollRecords {
		fmt.Printf("%2d. ", k)

		if v != nil {
			v.Print()
		} else {
			fmt.Printf("Nil record object\n")
		}
	}
}

func readPayrollRecords(inputFile string) ([]*PayrollRecord, error) {
	fileHandle, err := os.Open(inputFile)
	if err != nil {
		return nil, err
	}
	defer fileHandle.Close()

	csvReader := csv.NewReader(fileHandle)
	records := []*PayrollRecord{}

	for {
		row, err := csvReader.Read()
		if err != nil {
			if err == io.EOF {
				err = nil
			}

			// TODO: returning at the first erroneous row? can we skip over til EOF?
			return records, err
		}

		// TODO: consider sending entire input row to recod creation function
		if len(row) < 5 { // need at least five fields in a valid record, also check for empty fields
			// return or collect error and move to next record
			// return records, fmt.Errorf("Minimum number of fields not met in input <%s>\n", row)
			continue // read next row
		}

		newPayrollRecord, err := createPayrollRecord(row[0], row[1], row[2], row[3], row[4])
		if err != nil {
			// throw error
			// fmt.Errorf("Error creating payroll record with input %s\n", row)
			// return records, fmt.Errorf("Minimum number of fields not met in input <%s>\n", row)
			continue // read next row
		}

		records = append(records, newPayrollRecord)
	}
}

type PayrollRecord struct {
	FirstName    string
	LastName     string
	AnnualSalary float64
	SuperRate    float64
	PaymentDate  string
	Valid        bool   //	indicates if the record object is valid
	ErrorStr     string // if Valid == false, contains the input data from the input file leading to invalid object
	TaxBrackets  []*IncomeTaxBracket
}

func (rec *PayrollRecord) fullName() string {
	return rec.FirstName + " " + rec.LastName
}

func (rec *PayrollRecord) payPeriod() string {
	return rec.PaymentDate
}

func (rec *PayrollRecord) grossIncome() float64 {
	return round(rec.AnnualSalary / 12)
}

func (rec *PayrollRecord) incomeTax() (float64, error) {
	// iterate through tax brackets and find the right tax percentage for this salary
	perc := 0.0
	lump := 0.0
	abv := 0.0
	set := false
	for _, brac := range TAX_BRACKETS {
		if brac.Upper == 0 {
			// top tax bracket
			if rec.AnnualSalary >= brac.Lower {
				perc = brac.Percent
				lump = brac.Lump
				abv = brac.Above
				set = true
			}

		} else {
			if rec.AnnualSalary >= brac.Lower && rec.AnnualSalary <= brac.Upper {
				perc = brac.Percent
				lump = brac.Lump
				abv = brac.Above
				set = true
			}
		}
	}

	if !set {
		// no fitting bracket found for this salary
		return -1.0, fmt.Errorf("No fittnig tax bracket was found for salary amount %f", rec.AnnualSalary)
	}

	percentageTax := ((rec.AnnualSalary - abv) * perc) / 100
	return round(percentageTax + lump), nil
}

func (rec *PayrollRecord) netIncome() (float64, error) {
	gross := rec.grossIncome()
	tax, err := rec.incomeTax()

	if err != nil {
		return -1.0, err
	}

	if tax > gross {
		return -1.0, fmt.Errorf("Taxed amount (%f) larger than gross income (%f)", tax, gross)
	}

	return round(gross - tax), nil
}

func (rec *PayrollRecord) superAmount() (float64, error) {
	if rec.SuperRate < 0 || rec.SuperRate > 50 {
		return -1.0, fmt.Errorf("Invalid super rate (%f)", rec.SuperRate)
	}

	super := (rec.grossIncome() * rec.SuperRate) / 100
	return round(super), nil
}

func (rec *PayrollRecord) Print() {
	if rec.Valid {
		fmt.Printf("%s, %s, %.2f, %.2f%%, %s, %f\n", rec.FirstName, rec.LastName, rec.AnnualSalary, rec.SuperRate, rec.PaymentDate, rec.grossIncome())
	} else {
		fmt.Printf("Invalid record %s\n", rec.ErrorStr)
	}
}

func createPayrollRecord(firstname_str string, lastname_str string, annualsalary_str string, superrate_str string, paymentdate_str string) (*PayrollRecord, error) {
	var newRecord PayrollRecord

	// prepare input
	FirstName := strings.TrimSpace(firstname_str)
	LastName := strings.TrimSpace(lastname_str)
	AnnualSalary, err_sal := strconv.ParseFloat(annualsalary_str, 64)
	SuperRate := strings.TrimSpace(superrate_str)
	PaymentDate := strings.TrimSpace(paymentdate_str)

	// validate [TODO: may improve by returning specific error, consider ignoring erroneous record and reading ahead]
	// extract numeric super percentage
	rexp, _ := regexp.Compile("^[0-9]+")
	SuperRate_f, err_sr := strconv.ParseFloat(rexp.FindString(SuperRate), 64)

	if FirstName == "" || LastName == "" || (err_sal != nil) || (err_sr != nil) || AnnualSalary <= 0 || SuperRate_f < 0 || SuperRate_f > 50 {
		newRecord.Valid = false
		newRecord.ErrorStr = fmt.Sprintf("Invalid input record: [%s] [%s] [%s] [%s] [%s]", firstname_str, lastname_str, annualsalary_str, superrate_str, paymentdate_str)
		return &newRecord, fmt.Errorf("Invalid data in payroll input record\n")
	}

	newRecord.FirstName = FirstName
	newRecord.LastName = LastName
	newRecord.AnnualSalary = AnnualSalary
	newRecord.SuperRate = SuperRate_f
	newRecord.PaymentDate = PaymentDate
	newRecord.Valid = true

	return &newRecord, nil
}

func round(n float64) float64 {
	n_floor := math.Floor(n)
	n_dec := n - n_floor

	if n_dec >= .50 {
		// round up
		return n_floor + 1.0
	}

	// round down
	return n_floor
}

func createOutputFile(inFileName) error {
	if strings.TrimSpace(inFileName) == "" {
		return fmt.Errorf("Invalid input filename")
	}

	outFileName := strings.Split(strings.TrimSpace(inFileName), ".")[0] + "-out.txt"
	f, err := os.Create(outFileName)

	if err != nil {
		return fmt.Errorf("Error creating outputfile")
	}

	defer f.Close()

	for scanner.Scan() {
		line := scanner.Text()
		date := strings.TrimSpace(r1.FindString(line))
		ext := strings.TrimSpace(r2.FindString(line))

		if len(date) > 0 && len(ext) > 0 {
			// write line to file
			_, err := f.WriteString(r1.FindString(line) + "\t" + r2.FindString(line) + "\n")
			logErr(err)
		}
	}

	f.Sync()

	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}
}

// takes a salary amount and returns percentage, lump payment and value above which the percebntage should be calculated
// func getTaxPercent(sal float64) (float64, float64, float64) {
// 	taxmap := map[float64]float64{0: 18200, }
// }

type IncomeTaxBracket struct {
	Lower   float64 // lower salary limit of tax bracket
	Upper   float64 // upper salary limit of tax bracket
	Percent float64 // percentage tax to be levied above a certain threshold
	Lump    float64 // lump sum to be paid for this bracket, if any
	Above   float64 // threshold above which percentage tax has to be paid
}

func (itb *IncomeTaxBracket) Print()  {
	fmt.Printf("[Lower: $%.2f, Upper: $%.2f, Percent: %.2f%%, Lump Sum: $%.2f, Above: $%.2f]\n", itb.Lower, itb.Upper, itb.Percent, itb.Lump, itb.Above)
}

func readTaxBracketsConfig(inputFile string) ([]*IncomeTaxBracket, error) {
	fileHandle, err := os.Open(inputFile)
	if err != nil {
		return nil, err
	}
	defer fileHandle.Close()

	csvReader := csv.NewReader(fileHandle)
	brackets := []*IncomeTaxBracket{}

	i := 0
	prev_upper := 0.0
	for {
		row, err := csvReader.Read()
		if err != nil {
			if err == io.EOF {
				err = nil
			}

			return brackets, err
		}

		// TODO: consider sending entire input row to recod creation function
		if len(row) < 5 { // need at least five fields in a valid record, also check for empty fields
			// return or collect error and move to next record
			return nil, fmt.Errorf("readTaxBrackets(): Minimum number of fields not met in input <%s>\n", row)
		}

		lower, err_low := strconv.ParseFloat(strings.TrimSpace(row[0]), 64)

		top := false // flag indicating topmost tax bracket
		upp := strings.TrimSpace(row[1])
		upper := 0.0
		var err_upp error
		if upp != "" {
			// upper field could be empty if this is the top bracket
			upper, err_upp = strconv.ParseFloat(strings.TrimSpace(row[1]), 64)
		} else {
			// set flag indicating uppermost bracket
			top = true
		}

		percent, err_perc := strconv.ParseFloat(strings.TrimSpace(row[2]), 64)
		lump, err_lump := strconv.ParseFloat(strings.TrimSpace(row[3]), 64)
		threshold, err_thr := strconv.ParseFloat(strings.TrimSpace(row[4]), 64)

		if err_low != nil || err_upp != nil || err_perc != nil || err_lump != nil || err_thr != nil {
			fmt.Printf("%v\n", err_thr)
			return nil, fmt.Errorf("readTaxBrackets(): Error reading tax bracket config record: <%s>\n", row)
		}

		if lower >= upper && !top {
			return nil, fmt.Errorf("readTaxBrackets(): Lower limit >= upper limit in input <%s>\n", row)
		}

		if i == 0 {
			if lower != 0 {
				return nil, fmt.Errorf("readTaxBrackets(): First bracket lower limit != 0 in input <%s>\n", row)
			}
		} else {
			if lower <= prev_upper {
				return nil, fmt.Errorf("readTaxBrackets(): Current bracket's lower limit <= previous upper limit in input <%s>\n", row)
			}
		}

		prev_upper = upper
		newTaxBracket := &IncomeTaxBracket{lower, upper, percent, lump, threshold}
		brackets = append(brackets, newTaxBracket)
		i++

		if top {
			break
		}
	}

	if i == 0 {
		return nil, fmt.Errorf("readTaxBrackets(): No valid tax brackets found\n")
	}

	return brackets, nil
}
<|endoftext|>


<commit_before>package pinnedpost

import (
	"socialapi/models"
	"time"

	"github.com/koding/logging"
	"github.com/streadway/amqp"
)

type Controller struct{ log logging.Logger }

func (t *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	if delivery.Redelivered {
		t.log.Error("Redelivered message gave error again, putting to maintenance queue", err)
		delivery.Ack(false)
		return true
	}

	t.log.Error("an error occured putting message back to queue", err)
	delivery.Nack(false, true)
	return false
}

func New(log logging.Logger) *Controller {
	return &Controller{log: log}
}

// MessageReplyCreated handles the created replies
func (c *Controller) MessageReplyCreated(messageReply *models.MessageReply) error {
	parent, err := models.ChannelMessageById(messageReply.MessageId)
	if err != nil {
		return err
	}

	// only posts can be marked as pinned
	if parent.TypeConstant != models.ChannelMessage_TYPE_POST {
		return nil
	}

	reply, err := models.ChannelMessageById(messageReply.ReplyId)
	if err != nil {
		return err
	}

	// update all the channels that contain this message
	// this is done for changing the order of the pinned messages
	defer c.updateAllContainingChannels(parent, reply)

	// add parent message to the author's pinned message list
	err = c.addMessage(parent.AccountId, parent.Id, parent.InitialChannelId)
	if err != nil && err != models.ErrMessageAlreadyInTheChannel {
		return err
	}

	// no need to try to add the same message again to the author's pinned
	// message list
	if parent.AccountId == reply.AccountId {
		return nil
	}

	// add parent message to the replier's pinned message list
	err = c.addMessage(reply.AccountId, parent.Id, parent.InitialChannelId)
	if err != nil && err != models.ErrMessageAlreadyInTheChannel {
		return err
	}

	return nil
}

func (c *Controller) addMessage(accountId, messageId, channelId int64) error {
	// fetch the parent channel for group name
	// get it from cache
	channel, err := models.ChannelById(channelId)
	if err != nil {
		return err
	}

	// get pinning channel for current user if it is created,, else create and get
	pinningChannel, err := models.EnsurePinnedActivityChannel(accountId, channel.GroupName)
	if err != nil {
		return err
	}

	// add parent message into pinning channel
	// but do not force message to be put into the channel, because
	// if the user removed the message from channel, we should not add the same message again
	_, err = pinningChannel.EnsureMessage(messageId, false)

	// if message is already in the channel ignore the error, and mark process as successful
	if err == models.ErrMessageAlreadyInTheChannel {
		return nil
	}

	return err
}

// updateAllContainingChannels fetch all channels that parent is in and
// updates those channels except the user who did the action.
func (c *Controller) updateAllContainingChannels(parent, reply *models.ChannelMessage) error {
	cml := models.NewChannelMessageList()
	channels, err := cml.FetchMessageChannels(parent.Id)
	if err != nil {
		return err
	}

	if len(channels) == 0 {
		return nil
	}

	for _, channel := range channels {
		// if channel type is group, we dont need to update group's updatedAt
		if channel.TypeConstant == models.Channel_TYPE_GROUP {
			continue
		}

		// initiatorAccontId refers to users who did the action
		if channel.CreatorId == reply.AccountId {
			cml, err := channel.FetchMessageList(parent.Id)
			if err != nil {
				c.log.Error("error fetching message list for", parent.Id, err)
				continue
			}

			// `Glance` for author, so on next new message, unread count is right
			err = cml.Glance()
			if err != nil {
				c.log.Error("error glancing for messagelist", parent.Id, err)
				continue
			}

			// no need to tell user they did an action
			continue
		}

		// pinned activity channel holds messages one by one
		//
		// this should be removed after private message refactoring
		if channel.TypeConstant != models.Channel_TYPE_PINNED_ACTIVITY {
			channel.UpdatedAt = time.Now().UTC()
			if err := channel.Update(); err != nil {
				c.log.Error("channel update failed", err)
			}
			continue
		}

		cml := models.NewChannelMessageList()
		err := cml.UpdateAddedAt(channel.Id, parent.Id)
		if err != nil {
			c.log.Error("message list update failed", err)
		}

		pclu := models.PinnedChannelListUpdatedEvent{
			Channel: channel,
			Message: *parent,
			Reply:   *reply,
		}

		if err := cml.Emit("pinned_channel_list_updated", pclu); err != nil {
			c.log.Error("err %s", err.Error())
		}
	}

	return nil
}
<commit_msg>Social: fix style issue<commit_after>package pinnedpost

import (
	"socialapi/models"
	"time"

	"github.com/koding/logging"
	"github.com/streadway/amqp"
)

type Controller struct{ log logging.Logger }

func (t *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	if delivery.Redelivered {
		t.log.Error("Redelivered message gave error again, putting to maintenance queue", err)
		delivery.Ack(false)
		return true
	}

	t.log.Error("an error occured putting message back to queue", err)
	delivery.Nack(false, true)
	return false
}

func New(log logging.Logger) *Controller {
	return &Controller{log: log}
}

// MessageReplyCreated handles the created replies
func (c *Controller) MessageReplyCreated(messageReply *models.MessageReply) error {
	parent, err := models.ChannelMessageById(messageReply.MessageId)
	if err != nil {
		return err
	}

	// only posts can be marked as pinned
	if parent.TypeConstant != models.ChannelMessage_TYPE_POST {
		return nil
	}

	reply, err := models.ChannelMessageById(messageReply.ReplyId)
	if err != nil {
		return err
	}

	// update all the channels that contain this message
	// this is done for changing the order of the pinned messages
	defer c.updateAllContainingChannels(parent, reply)

	// add parent message to the author's pinned message list
	err = c.addMessage(parent.AccountId, parent.Id, parent.InitialChannelId)
	if err != nil && err != models.ErrMessageAlreadyInTheChannel {
		return err
	}

	// no need to try to add the same message again to the author's pinned
	// message list
	if parent.AccountId == reply.AccountId {
		return nil
	}

	// add parent message to the replier's pinned message list
	err = c.addMessage(reply.AccountId, parent.Id, parent.InitialChannelId)
	if err != nil && err != models.ErrMessageAlreadyInTheChannel {
		return err
	}

	return nil
}

func (c *Controller) addMessage(accountId, messageId, channelId int64) error {
	// fetch the parent channel for group name
	// get it from cache
	channel, err := models.ChannelById(channelId)
	if err != nil {
		return err
	}

	// get pinning channel for current user if it is created,, else create and get
	pinningChannel, err := models.EnsurePinnedActivityChannel(accountId, channel.GroupName)
	if err != nil {
		return err
	}

	// add parent message into pinning channel
	// but do not force message to be put into the channel, because
	// if the user removed the message from channel, we should not add the same message again
	_, err = pinningChannel.EnsureMessage(messageId, false)

	// if message is already in the channel ignore the error, and mark process as successful
	if err == models.ErrMessageAlreadyInTheChannel {
		return nil
	}

	return err
}

// updateAllContainingChannels fetch all channels that parent is in and
// updates those channels except the user who did the action.
func (c *Controller) updateAllContainingChannels(parent, reply *models.ChannelMessage) error {
	cml := models.NewChannelMessageList()
	channels, err := cml.FetchMessageChannels(parent.Id)
	if err != nil {
		return err
	}

	if len(channels) == 0 {
		return nil
	}

	for _, channel := range channels {
		// if channel type is group, we dont need to update group's updatedAt
		if channel.TypeConstant == models.Channel_TYPE_GROUP {
			continue
		}

		// initiatorAccontId refers to users who did the action
		if channel.CreatorId == reply.AccountId {
			cml, err := channel.FetchMessageList(parent.Id)
			if err != nil {
				c.log.Error("error fetching message list for", parent.Id, err)
				continue
			}

			// `Glance` for author, so on next new message, unread count is right
			err = cml.Glance()
			if err != nil {
				c.log.Error("error glancing for messagelist", parent.Id, err)
				continue
			}

			// no need to tell user they did an action
			continue
		}

		// pinned activity channel holds messages one by one
		//
		// this should be removed after private message refactoring
		if channel.TypeConstant != models.Channel_TYPE_PINNED_ACTIVITY {
			channel.UpdatedAt = time.Now().UTC()
			if err := channel.Update(); err != nil {
				c.log.Error("channel update failed", err)
			}

			continue
		}

		cml := models.NewChannelMessageList()
		err := cml.UpdateAddedAt(channel.Id, parent.Id)
		if err != nil {
			c.log.Error("message list update failed", err)
		}

		pclu := models.PinnedChannelListUpdatedEvent{
			Channel: channel,
			Message: *parent,
			Reply:   *reply,
		}

		if err := cml.Emit("pinned_channel_list_updated", pclu); err != nil {
			c.log.Error("err %s", err.Error())
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"encoding/json"

	"os/exec"

	"github.com/go-errors/errors"
	"github.com/itchio/butler/comm"
	"github.com/itchio/butler/redist"
	"github.com/natefinch/npipe"
	"github.com/olekukonko/tablewriter"
	"golang.org/x/text/encoding/unicode"
	"golang.org/x/text/transform"
)

// PrereqTask describes something the prereq installer has to do
type PrereqTask struct {
	Name    string             `json:"name"`
	WorkDir string             `json:"workDir"`
	Info    redist.RedistEntry `json:"info"`
}

// PrereqPlan contains a list of tasks for the prereq installer
type PrereqPlan struct {
	Tasks []*PrereqTask `json:"tasks"`
}

// PrereqState informs the caller on the current status of a prereq
type PrereqState struct {
	Type   string `json:"type"`
	Name   string `json:"name"`
	Status string `json:"status"`
}

// PrereqLogEntry sends an information to the caller on the progress of the task
type PrereqLogEntry struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

func installPrereqs(planPath string, pipePath string) {
	must(doInstallPrereqs(planPath, pipePath))
}

func doInstallPrereqs(planPath string, pipePath string) error {
	planReader, err := os.Open(planPath)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	decoder := json.NewDecoder(planReader)

	plan := &PrereqPlan{}
	err = decoder.Decode(plan)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	hasConn := true
	conn, err := npipe.Dial(pipePath)
	if err != nil {
		comm.Warnf("Could not dial pipe %s", conn)
		hasConn = false
	}

	writeLine := func(contents []byte) error {
		if !hasConn {
			return nil
		}

		contents = append(contents, '\n')

		_, err = conn.Write(contents)
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return nil
	}

	doWriteState := func(taskName string, status string) error {
		contents, err := json.Marshal(&PrereqState{
			Type:   "state",
			Name:   taskName,
			Status: status,
		})
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return writeLine(contents)
	}

	writeState := func(taskName string, status string) {
		err := doWriteState(taskName, status)
		if err != nil {
			switch err := err.(type) {
			case *errors.Error:
				comm.Warnf("Couldn't write log entry: %s", err.ErrorStack())
			default:
				comm.Warnf("Couldn't write log entry: %s", err.Error())
			}
		}
	}

	doLogf := func(format string, args ...interface{}) error {
		comm.Logf(format, args...)
		message := fmt.Sprintf(format, args...)

		contents, err := json.Marshal(&PrereqLogEntry{
			Type:    "log",
			Message: message,
		})
		if err != nil {
			return errors.Wrap(err, 0)
		}

		err = writeLine([]byte(contents))
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return nil
	}

	logf := func(format string, args ...interface{}) {
		err := doLogf(format, args...)
		if err != nil {
			switch err := err.(type) {
			case *errors.Error:
				comm.Warnf("Couldn't write log entry: %s", err.ErrorStack())
			default:
				comm.Warnf("Couldn't write log entry: %s", err.Error())
			}
		}
	}

	logf("Installing %d prereqs", len(plan.Tasks))
	startTime := time.Now()

	var failed []string

	for _, task := range plan.Tasks {
		taskStartTime := time.Now()
		writeState(task.Name, "installing")

		logf("")
		logf("# Installing %s", task.Name)
		logf("")

		commandPath := filepath.Join(task.WorkDir, task.Info.Command)
		args := task.Info.Args

		// MSI packages get special treatment for reasons.
		if strings.HasSuffix(strings.ToLower(task.Info.Command), ".msi") {
			tempDir, err := ioutil.TempDir("", "butler-msi-logs")
			if err != nil {
				return errors.Wrap(err, 0)
			}

			defer func() {
				os.RemoveAll(tempDir)
			}()

			logPath := filepath.Join(tempDir, "msi-install-log.txt")

			err = doMsiInstall(commandPath, logPath, "")
			if err != nil {
				logf("MSI install failed: %s", err.Error())
				lf, openErr := os.Open(logPath)
				if openErr != nil {
					logf("And what's more, we can't open the log: %s", openErr.Error())
				} else {
					// grok UTF-16
					win16be := unicode.UTF16(unicode.BigEndian, unicode.IgnoreBOM)
					// ...but abide by the BOM if there's one
					utf16bom := unicode.BOMOverride(win16be.NewDecoder())

					unicodeReader := transform.NewReader(lf, utf16bom)

					defer lf.Close()
					logf("Full MSI log follows:")
					s := bufio.NewScanner(unicodeReader)
					for s.Scan() {
						logf("[msi] %s", s.Text())
					}
					if scanErr := s.Err(); scanErr != nil {
						logf("While reading msi log: %s", scanErr.Error())
					}
				}

				failed = append(failed, task.Name)
			}
		} else {
			cmd := exec.Command(commandPath, args...)
			cmd.Dir = task.WorkDir

			logf("Launching %s %s", task.Info.Command, strings.Join(args, " "))

			err = cmd.Run()
			if err != nil {
				if exitError, ok := err.(*exec.ExitError); ok {
					if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
						code := status.ExitStatus()
						known := false
						for _, exitCode := range task.Info.ExitCodes {
							if code == exitCode.Code {
								if exitCode.Success {
									logf("%s, continuing", exitCode.Message)
								} else {
									logf("%s, we'll error out eventually", exitCode.Message)
									failed = append(failed, task.Name)
								}
								known = true
							}
							break
						}

						if !known {
							logf("Got unknown exit code %d, will error out", code)
							failed = append(failed, task.Name)
						}
					} else {
						return errors.Wrap(err, 0)
					}
				} else {
					return errors.Wrap(err, 0)
				}
			}
		}

		writeState(task.Name, "done")
		logf("(Spent %s)", time.Since(taskStartTime))
	}

	if len(failed) > 0 {
		logf("")
		errMsg := fmt.Sprintf("Some prereqs failed to install: %s", strings.Join(failed, ", "))
		logf(errMsg)
		return errors.Wrap(errors.New(errMsg), 0)
	} else {
		logf("All done! Took %s", time.Since(startTime))
	}

	return nil
}

func testPrereqs(prereqs []string) {
	must(doTestPrereqs(prereqs))
}

func doTestPrereqs(prereqs []string) error {
	comm.Opf("Fetching registry...")

	baseURL := "https://dl.itch.ovh/itch-redists"

	infoURL := fmt.Sprintf("%s/info.json?t=%d", baseURL, time.Now().Unix())
	res, err := http.Get(infoURL)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	if res.StatusCode != 200 {
		return errors.Wrap(fmt.Errorf("While getting redist registry, got HTTP %d", res.StatusCode), 0)
	}

	dec := json.NewDecoder(res.Body)

	registry := &redist.RedistRegistry{}
	err = dec.Decode(registry)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	if len(prereqs) == 0 {
		comm.Logf("")
		comm.Statf("No prereqs specified, here are those we know about: \n")

		table := tablewriter.NewWriter(os.Stdout)
		table.SetAutoFormatHeaders(false)
		table.SetColWidth(60)
		table.SetHeader([]string{"Name", "Arch", "Description", "Version"})
		for name, info := range registry.Entries {
			table.Append([]string{name, info.Arch, info.FullName, info.Version})
		}
		table.Render()
		return nil
	}

	if len(prereqs) == 1 && prereqs[0] == "all" {
		prereqs = nil
		for name := range registry.Entries {
			prereqs = append(prereqs, name)
		}
	}

	comm.Logf("Testing out prereqs %s", strings.Join(prereqs, ", "))

	plan := &PrereqPlan{}

	tempDir := filepath.Join(os.TempDir(), "butler-test-prereqs")
	comm.Logf("Working in %s", tempDir)
	comm.Logf("(This helps not having to re-download the prereqs between runs, but feel free to wipe it)")

	err = os.MkdirAll(tempDir, 0755)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	for _, name := range prereqs {
		info, ok := registry.Entries[name]
		if !ok {
			comm.Warnf("Unknown prereq %s, skipping", name)
			continue
		}

		comm.Opf("Downloading prereq %s", name)

		workDir := filepath.Join(tempDir, name)
		err = os.MkdirAll(workDir, 0755)
		if err != nil {
			return errors.Wrap(err, 0)
		}

		task := &PrereqTask{
			Info:    *info,
			Name:    name,
			WorkDir: workDir,
		}

		url := fmt.Sprintf("%s/%s/%s", baseURL, name, info.Command)
		dest := filepath.Join(workDir, info.Command)
		_, err = tryDl(url, dest)
		if err != nil {
			comm.Logf("Could not donwload prereq %s", name)
			return errors.Wrap(err, 0)
		}

		plan.Tasks = append(plan.Tasks, task)
	}

	planPath := filepath.Join(tempDir, "butler_install_plan.json")
	comm.Logf("Writing plan to %s", planPath)

	planContents, err := json.Marshal(plan)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = ioutil.WriteFile(planPath, planContents, 0644)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	comm.Opf("Handing off to install-prereqs...")

	err = doInstallPrereqs(planPath, "")
	if err != nil {
		return errors.Wrap(err, 0)
	}

	return nil
}
<commit_msg>Last minor output nitpicks<commit_after>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"encoding/json"

	"os/exec"

	"github.com/go-errors/errors"
	"github.com/itchio/butler/comm"
	"github.com/itchio/butler/redist"
	"github.com/natefinch/npipe"
	"github.com/olekukonko/tablewriter"
	"golang.org/x/text/encoding/unicode"
	"golang.org/x/text/transform"
)

// PrereqTask describes something the prereq installer has to do
type PrereqTask struct {
	Name    string             `json:"name"`
	WorkDir string             `json:"workDir"`
	Info    redist.RedistEntry `json:"info"`
}

// PrereqPlan contains a list of tasks for the prereq installer
type PrereqPlan struct {
	Tasks []*PrereqTask `json:"tasks"`
}

// PrereqState informs the caller on the current status of a prereq
type PrereqState struct {
	Type   string `json:"type"`
	Name   string `json:"name"`
	Status string `json:"status"`
}

// PrereqLogEntry sends an information to the caller on the progress of the task
type PrereqLogEntry struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

func installPrereqs(planPath string, pipePath string) {
	must(doInstallPrereqs(planPath, pipePath))
}

func doInstallPrereqs(planPath string, pipePath string) error {
	planReader, err := os.Open(planPath)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	decoder := json.NewDecoder(planReader)

	plan := &PrereqPlan{}
	err = decoder.Decode(plan)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	hasConn := true
	conn, err := npipe.Dial(pipePath)
	if err != nil {
		comm.Warnf("Could not dial pipe %s", conn)
		hasConn = false
	}

	writeLine := func(contents []byte) error {
		if !hasConn {
			return nil
		}

		contents = append(contents, '\n')

		_, err = conn.Write(contents)
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return nil
	}

	doWriteState := func(taskName string, status string) error {
		contents, err := json.Marshal(&PrereqState{
			Type:   "state",
			Name:   taskName,
			Status: status,
		})
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return writeLine(contents)
	}

	writeState := func(taskName string, status string) {
		err := doWriteState(taskName, status)
		if err != nil {
			switch err := err.(type) {
			case *errors.Error:
				comm.Warnf("Couldn't write log entry: %s", err.ErrorStack())
			default:
				comm.Warnf("Couldn't write log entry: %s", err.Error())
			}
		}
	}

	doLogf := func(format string, args ...interface{}) error {
		comm.Logf(format, args...)
		message := fmt.Sprintf(format, args...)

		contents, err := json.Marshal(&PrereqLogEntry{
			Type:    "log",
			Message: message,
		})
		if err != nil {
			return errors.Wrap(err, 0)
		}

		err = writeLine([]byte(contents))
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return nil
	}

	logf := func(format string, args ...interface{}) {
		err := doLogf(format, args...)
		if err != nil {
			switch err := err.(type) {
			case *errors.Error:
				comm.Warnf("Couldn't write log entry: %s", err.ErrorStack())
			default:
				comm.Warnf("Couldn't write log entry: %s", err.Error())
			}
		}
	}

	logf("Installing %d prereqs", len(plan.Tasks))
	startTime := time.Now()

	var failed []string

	for _, task := range plan.Tasks {
		taskStartTime := time.Now()
		writeState(task.Name, "installing")

		logf("")
		logf("# Installing %s", task.Name)
		logf("")

		commandPath := filepath.Join(task.WorkDir, task.Info.Command)
		args := task.Info.Args

		// MSI packages get special treatment for reasons.
		if strings.HasSuffix(strings.ToLower(task.Info.Command), ".msi") {
			tempDir, err := ioutil.TempDir("", "butler-msi-logs")
			if err != nil {
				return errors.Wrap(err, 0)
			}

			defer func() {
				os.RemoveAll(tempDir)
			}()

			logPath := filepath.Join(tempDir, "msi-install-log.txt")

			err = doMsiInstall(commandPath, logPath, "")
			if err != nil {
				logf("MSI install failed: %s", err.Error())
				lf, openErr := os.Open(logPath)
				if openErr != nil {
					logf("And what's more, we can't open the log: %s", openErr.Error())
				} else {
					// grok UTF-16
					win16be := unicode.UTF16(unicode.BigEndian, unicode.IgnoreBOM)
					// ...but abide by the BOM if there's one
					utf16bom := unicode.BOMOverride(win16be.NewDecoder())

					unicodeReader := transform.NewReader(lf, utf16bom)

					defer lf.Close()
					logf("Full MSI log follows:")
					s := bufio.NewScanner(unicodeReader)
					for s.Scan() {
						logf("[msi] %s", s.Text())
					}
					if scanErr := s.Err(); scanErr != nil {
						logf("While reading msi log: %s", scanErr.Error())
					}
				}

				failed = append(failed, task.Name)
			}
		} else {
			cmd := exec.Command(commandPath, args...)
			cmd.Dir = task.WorkDir

			logf("Launching %s %s", task.Info.Command, strings.Join(args, " "))

			err = cmd.Run()
			if err != nil {
				if exitError, ok := err.(*exec.ExitError); ok {
					if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
						code := status.ExitStatus()
						known := false
						for _, exitCode := range task.Info.ExitCodes {
							if code == exitCode.Code {
								if exitCode.Success {
									logf("%s, continuing", exitCode.Message)
								} else {
									logf("%s, we'll error out eventually", exitCode.Message)
									failed = append(failed, task.Name)
								}
								known = true
							}
							break
						}

						if !known {
							logf("Got unknown exit code %d, will error out", code)
							failed = append(failed, task.Name)
						}
					} else {
						return errors.Wrap(err, 0)
					}
				} else {
					return errors.Wrap(err, 0)
				}
			}
		}

		writeState(task.Name, "done")
		logf("(Spent %s)", time.Since(taskStartTime))
	}

	logf("")
	if len(failed) > 0 {
		errMsg := fmt.Sprintf("Some prereqs failed to install: %s", strings.Join(failed, ", "))
		logf(errMsg)
		return errors.Wrap(errors.New(errMsg), 0)
	}

	logf("All done! (Spent %s total)", time.Since(startTime))

	return nil
}

func testPrereqs(prereqs []string) {
	must(doTestPrereqs(prereqs))
}

func doTestPrereqs(prereqs []string) error {
	comm.Opf("Fetching registry...")

	baseURL := "https://dl.itch.ovh/itch-redists"

	infoURL := fmt.Sprintf("%s/info.json?t=%d", baseURL, time.Now().Unix())
	res, err := http.Get(infoURL)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	if res.StatusCode != 200 {
		return errors.Wrap(fmt.Errorf("While getting redist registry, got HTTP %d", res.StatusCode), 0)
	}

	dec := json.NewDecoder(res.Body)

	registry := &redist.RedistRegistry{}
	err = dec.Decode(registry)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	if len(prereqs) == 0 {
		comm.Logf("")
		comm.Statf("No prereqs specified, here are those we know about: \n")

		table := tablewriter.NewWriter(os.Stdout)
		table.SetAutoFormatHeaders(false)
		table.SetColWidth(60)
		table.SetHeader([]string{"Name", "Arch", "Description", "Version"})
		for name, info := range registry.Entries {
			table.Append([]string{name, info.Arch, info.FullName, info.Version})
		}
		table.Render()
		return nil
	}

	if len(prereqs) == 1 && prereqs[0] == "all" {
		prereqs = nil
		for name := range registry.Entries {
			prereqs = append(prereqs, name)
		}
	}

	comm.Logf("Testing out prereqs %s", strings.Join(prereqs, ", "))

	plan := &PrereqPlan{}

	tempDir := filepath.Join(os.TempDir(), "butler-test-prereqs")
	comm.Logf("Working in %s", tempDir)
	comm.Logf("(This helps not having to re-download the prereqs between runs, but feel free to wipe it)")

	err = os.MkdirAll(tempDir, 0755)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	for _, name := range prereqs {
		info, ok := registry.Entries[name]
		if !ok {
			comm.Warnf("Unknown prereq %s, skipping", name)
			continue
		}

		comm.Opf("Downloading prereq %s", name)

		workDir := filepath.Join(tempDir, name)
		err = os.MkdirAll(workDir, 0755)
		if err != nil {
			return errors.Wrap(err, 0)
		}

		task := &PrereqTask{
			Info:    *info,
			Name:    name,
			WorkDir: workDir,
		}

		url := fmt.Sprintf("%s/%s/%s", baseURL, name, info.Command)
		dest := filepath.Join(workDir, info.Command)
		_, err = tryDl(url, dest)
		if err != nil {
			comm.Logf("Could not donwload prereq %s", name)
			return errors.Wrap(err, 0)
		}

		plan.Tasks = append(plan.Tasks, task)
	}

	planPath := filepath.Join(tempDir, "butler_install_plan.json")
	comm.Logf("Writing plan to %s", planPath)

	planContents, err := json.Marshal(plan)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = ioutil.WriteFile(planPath, planContents, 0644)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	comm.Opf("Handing off to install-prereqs...")

	err = doInstallPrereqs(planPath, "")
	if err != nil {
		return errors.Wrap(err, 0)
	}

	return nil
}
<|endoftext|>


<commit_before>package tty

import (
	"log"
	"syscall"
	"os"
	"unsafe"
)

type Termios struct {
	Iflag  uint32
	Oflag  uint32
	Cflag  uint32
	Lflag  uint32
	Cc     [20]byte
	Ispeed uint32
	Ospeed uint32
}

type Tty struct {
	original *Termios
}

func tcSetAttr(fd uintptr, termios *Termios) error {
	// TCSETS+1 == TCSETSW, because TCSAFLUSH doesn't exist
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(syscall.TCSETS+1), uintptr(unsafe.Pointer(termios))); err != 0 {
		return err
	}
	return nil
}

func tcGetAttr(fd uintptr) *Termios {
	var termios = &Termios{}
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, syscall.TCGETS, uintptr(unsafe.Pointer(termios))); err != 0 {
		log.Fatalf("Problem getting terminal attributes: %s\n", err)
	}
	return termios
}

func (t *Tty) EnableRawMode() {
	t.original = tcGetAttr(os.Stdin.Fd())
	var raw Termios
	raw = *t.original
	raw.Iflag &^= syscall.BRKINT | syscall.ICRNL | syscall.INPCK | syscall.ISTRIP | syscall.IXON
	raw.Oflag &^= syscall.OPOST
	raw.Cflag |= syscall.CS8
	raw.Lflag &^= syscall.ECHO | syscall.ICANON | syscall.IEXTEN | syscall.ISIG
	raw.Cc[syscall.VMIN+1] = 0
	raw.Cc[syscall.VTIME+1] = 1
	if e := tcSetAttr(os.Stdin.Fd(), &raw); e != nil {
		log.Fatalf("Problem enabling raw mode: %s\n", e)
	}
}

func (t *Tty) DisableRawMode() {
	if e := tcSetAttr(os.Stdin.Fd(), t.original); e != nil {
		log.Fatalf("Problem disabling raw mode: %s\n", e)
	}
}
<commit_msg>Change Termios to termios on suggestion of golint<commit_after>package tty

import (
	"log"
	"syscall"
	"os"
	"unsafe"
)

type termios struct {
	Iflag  uint32
	Oflag  uint32
	Cflag  uint32
	Lflag  uint32
	Cc     [20]byte
	Ispeed uint32
	Ospeed uint32
}

type Tty struct {
	original *termios
}

func tcSetAttr(fd uintptr, termios *termios) error {
	// TCSETS+1 == TCSETSW, because TCSAFLUSH doesn't exist
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(syscall.TCSETS+1), uintptr(unsafe.Pointer(termios))); err != 0 {
		return err
	}
	return nil
}

func tcGetAttr(fd uintptr) *termios {
	var termios = &termios{}
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, syscall.TCGETS, uintptr(unsafe.Pointer(termios))); err != 0 {
		log.Fatalf("Problem getting terminal attributes: %s\n", err)
	}
	return termios
}

func (t *Tty) EnableRawMode() {
	t.original = tcGetAttr(os.Stdin.Fd())
	var raw termios
	raw = *t.original
	raw.Iflag &^= syscall.BRKINT | syscall.ICRNL | syscall.INPCK | syscall.ISTRIP | syscall.IXON
	raw.Oflag &^= syscall.OPOST
	raw.Cflag |= syscall.CS8
	raw.Lflag &^= syscall.ECHO | syscall.ICANON | syscall.IEXTEN | syscall.ISIG
	raw.Cc[syscall.VMIN+1] = 0
	raw.Cc[syscall.VTIME+1] = 1
	if e := tcSetAttr(os.Stdin.Fd(), &raw); e != nil {
		log.Fatalf("Problem enabling raw mode: %s\n", e)
	}
}

func (t *Tty) DisableRawMode() {
	if e := tcSetAttr(os.Stdin.Fd(), t.original); e != nil {
		log.Fatalf("Problem disabling raw mode: %s\n", e)
	}
}
<|endoftext|>


<commit_before>// +build linux

/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 docker

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/events"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/vishvananda/netns"
	"golang.org/x/net/context"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/graph"
	ns "github.com/skydive-project/skydive/topology/probes/netns"
	sversion "github.com/skydive-project/skydive/version"
)

// DockerClientAPIVersion Client API version used
const DockerClientAPIVersion = "1.18"

type containerInfo struct {
	Pid  int
	Node *graph.Node
}

// DockerProbe describes a Docker topology graph that enhance the graph
type DockerProbe struct {
	sync.RWMutex
	*ns.NetNSProbe
	url          string
	client       *client.Client
	cancel       context.CancelFunc
	state        int64
	connected    atomic.Value
	wg           sync.WaitGroup
	hostNs       netns.NsHandle
	containerMap map[string]containerInfo
}

func (probe *DockerProbe) containerNamespace(pid int) string {
	return fmt.Sprintf("/proc/%d/ns/net", pid)
}

func (probe *DockerProbe) registerContainer(id string) {
	probe.Lock()
	defer probe.Unlock()

	if _, ok := probe.containerMap[id]; ok {
		return
	}
	info, err := probe.client.ContainerInspect(context.Background(), id)
	if err != nil {
		logging.GetLogger().Errorf("Failed to inspect Docker container %s: %s", id, err.Error())
		return
	}

	nsHandle, err := netns.GetFromPid(info.State.Pid)
	if err != nil {
		return
	}
	defer nsHandle.Close()

	namespace := probe.containerNamespace(info.State.Pid)
	logging.GetLogger().Debugf("Register docker container %s and PID %d", info.ID, info.State.Pid)

	var n *graph.Node
	if probe.hostNs.Equal(nsHandle) {
		// The container is in net=host mode
		n = probe.Root
	} else {
		n = probe.Register(namespace, info.Name[1:])

		probe.Graph.Lock()
		probe.Graph.AddMetadata(n, "Manager", "docker")
		probe.Graph.Unlock()
	}

	probe.Graph.Lock()
	metadata := graph.Metadata{
		"Type": "container",
		"Name": info.Name[1:],
		"Docker": map[string]interface{}{
			"ContainerID":   info.ID,
			"ContainerName": info.Name,
			"ContainerPID":  int64(info.State.Pid),
		},
	}

	if len(info.Config.Labels) != 0 {
		labels := make(map[string]interface{})
		for k, v := range info.Config.Labels {
			common.SetField(labels, k, v)
		}
		metadata["Docker"].(map[string]interface{})["Labels"] = labels
	}

	containerNode := probe.Graph.NewNode(graph.GenID(), metadata)
	topology.AddOwnershipLink(probe.Graph, n, containerNode, nil)
	probe.Graph.Unlock()

	probe.containerMap[info.ID] = containerInfo{
		Pid:  info.State.Pid,
		Node: containerNode,
	}
}

func (probe *DockerProbe) unregisterContainer(id string) {
	probe.Lock()
	defer probe.Unlock()

	infos, ok := probe.containerMap[id]
	if !ok {
		return
	}

	probe.Graph.Lock()
	probe.Graph.DelNode(infos.Node)
	probe.Graph.Unlock()

	namespace := probe.containerNamespace(infos.Pid)
	logging.GetLogger().Debugf("Stop listening for namespace %s with PID %d", namespace, infos.Pid)
	probe.Unregister(namespace)

	delete(probe.containerMap, id)
}

func (probe *DockerProbe) handleDockerEvent(event *events.Message) {
	if event.Status == "start" {
		probe.registerContainer(event.ID)
	} else if event.Status == "die" {
		probe.unregisterContainer(event.ID)
	}
}

func (probe *DockerProbe) connect() error {
	var err error

	if probe.hostNs, err = netns.Get(); err != nil {
		return err
	}
	defer probe.hostNs.Close()

	logging.GetLogger().Debugf("Connecting to Docker daemon: %s", probe.url)
	defaultHeaders := map[string]string{"User-Agent": fmt.Sprintf("skydive-agent-%s", sversion.Version)}
	probe.client, err = client.NewClient(probe.url, DockerClientAPIVersion, nil, defaultHeaders)
	if err != nil {
		logging.GetLogger().Errorf("Failed to create client to Docker daemon: %s", err.Error())
		return err
	}
	defer probe.client.Close()

	if _, err := probe.client.ServerVersion(context.Background()); err != nil {
		logging.GetLogger().Errorf("Failed to connect to Docker daemon: %s", err.Error())
		return err
	}

	eventsFilter := filters.NewArgs()
	eventsFilter.Add("event", "start")
	eventsFilter.Add("event", "die")

	ctx, cancel := context.WithCancel(context.Background())
	eventChan, errChan := probe.client.Events(ctx, types.EventsOptions{Filters: eventsFilter})

	probe.cancel = cancel
	probe.wg.Add(2)

	probe.connected.Store(true)
	defer probe.connected.Store(false)

	go func() {
		defer probe.wg.Done()

		containers, err := probe.client.ContainerList(context.Background(), types.ContainerListOptions{})
		if err != nil {
			logging.GetLogger().Errorf("Failed to list containers: %s", err.Error())
			return
		}

		for _, c := range containers {
			if atomic.LoadInt64(&probe.state) != common.RunningState {
				break
			}
			probe.registerContainer(c.ID)
		}
	}()

	defer probe.wg.Done()

	for {
		select {
		case err := <-errChan:
			if atomic.LoadInt64(&probe.state) != common.StoppingState {
				logging.GetLogger().Errorf("Got error while waiting for Docker event: %s", err.Error())
			}
			return err
		case event := <-eventChan:
			probe.handleDockerEvent(&event)
		}
	}
}

// Start the probe
func (probe *DockerProbe) Start() {
	if !atomic.CompareAndSwapInt64(&probe.state, common.StoppedState, common.RunningState) {
		return
	}

	go func() {
		for {
			state := atomic.LoadInt64(&probe.state)
			if state == common.StoppingState || state == common.StoppedState {
				break
			}

			if probe.connect() != nil {
				time.Sleep(1 * time.Second)
			}

			probe.wg.Wait()
		}
	}()
}

// Stop the probe
func (probe *DockerProbe) Stop() {
	if !atomic.CompareAndSwapInt64(&probe.state, common.RunningState, common.StoppingState) {
		return
	}

	if probe.connected.Load() == true {
		probe.cancel()
		probe.wg.Wait()
	}

	atomic.StoreInt64(&probe.state, common.StoppedState)
}

// NewDockerProbe creates a new topology Docker probe
func NewDockerProbe(nsProbe *ns.NetNSProbe, dockerURL string) (*DockerProbe, error) {
	dockerProbe := &DockerProbe{
		NetNSProbe:   nsProbe,
		url:          dockerURL,
		containerMap: make(map[string]containerInfo),
		state:        common.StoppedState,
	}

	if path := config.GetConfig().GetString("docker.netns.run_path"); path != "" {
		nsProbe.Watch(path)
	}

	return dockerProbe, nil
}
<commit_msg>docker: use common.NormalizeValue() to set "Labels"<commit_after>// +build linux

/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 docker

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/events"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/vishvananda/netns"
	"golang.org/x/net/context"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/graph"
	ns "github.com/skydive-project/skydive/topology/probes/netns"
	sversion "github.com/skydive-project/skydive/version"
)

// DockerClientAPIVersion Client API version used
const DockerClientAPIVersion = "1.18"

type containerInfo struct {
	Pid  int
	Node *graph.Node
}

// DockerProbe describes a Docker topology graph that enhance the graph
type DockerProbe struct {
	sync.RWMutex
	*ns.NetNSProbe
	url          string
	client       *client.Client
	cancel       context.CancelFunc
	state        int64
	connected    atomic.Value
	wg           sync.WaitGroup
	hostNs       netns.NsHandle
	containerMap map[string]containerInfo
}

func (probe *DockerProbe) containerNamespace(pid int) string {
	return fmt.Sprintf("/proc/%d/ns/net", pid)
}

func (probe *DockerProbe) registerContainer(id string) {
	probe.Lock()
	defer probe.Unlock()

	if _, ok := probe.containerMap[id]; ok {
		return
	}
	info, err := probe.client.ContainerInspect(context.Background(), id)
	if err != nil {
		logging.GetLogger().Errorf("Failed to inspect Docker container %s: %s", id, err.Error())
		return
	}

	nsHandle, err := netns.GetFromPid(info.State.Pid)
	if err != nil {
		return
	}
	defer nsHandle.Close()

	namespace := probe.containerNamespace(info.State.Pid)
	logging.GetLogger().Debugf("Register docker container %s and PID %d", info.ID, info.State.Pid)

	var n *graph.Node
	if probe.hostNs.Equal(nsHandle) {
		// The container is in net=host mode
		n = probe.Root
	} else {
		n = probe.Register(namespace, info.Name[1:])

		probe.Graph.Lock()
		probe.Graph.AddMetadata(n, "Manager", "docker")
		probe.Graph.Unlock()
	}

	probe.Graph.Lock()
	metadata := graph.Metadata{
		"Type": "container",
		"Name": info.Name[1:],
		"Docker": map[string]interface{}{
			"ContainerID":   info.ID,
			"ContainerName": info.Name,
			"ContainerPID":  int64(info.State.Pid),
		},
	}

	if len(info.Config.Labels) != 0 {
		metadata["Docker"].(map[string]interface{})["Labels"] = common.NormalizeValue(info.Config.Labels)
	}

	containerNode := probe.Graph.NewNode(graph.GenID(), metadata)
	topology.AddOwnershipLink(probe.Graph, n, containerNode, nil)
	probe.Graph.Unlock()

	probe.containerMap[info.ID] = containerInfo{
		Pid:  info.State.Pid,
		Node: containerNode,
	}
}

func (probe *DockerProbe) unregisterContainer(id string) {
	probe.Lock()
	defer probe.Unlock()

	infos, ok := probe.containerMap[id]
	if !ok {
		return
	}

	probe.Graph.Lock()
	probe.Graph.DelNode(infos.Node)
	probe.Graph.Unlock()

	namespace := probe.containerNamespace(infos.Pid)
	logging.GetLogger().Debugf("Stop listening for namespace %s with PID %d", namespace, infos.Pid)
	probe.Unregister(namespace)

	delete(probe.containerMap, id)
}

func (probe *DockerProbe) handleDockerEvent(event *events.Message) {
	if event.Status == "start" {
		probe.registerContainer(event.ID)
	} else if event.Status == "die" {
		probe.unregisterContainer(event.ID)
	}
}

func (probe *DockerProbe) connect() error {
	var err error

	if probe.hostNs, err = netns.Get(); err != nil {
		return err
	}
	defer probe.hostNs.Close()

	logging.GetLogger().Debugf("Connecting to Docker daemon: %s", probe.url)
	defaultHeaders := map[string]string{"User-Agent": fmt.Sprintf("skydive-agent-%s", sversion.Version)}
	probe.client, err = client.NewClient(probe.url, DockerClientAPIVersion, nil, defaultHeaders)
	if err != nil {
		logging.GetLogger().Errorf("Failed to create client to Docker daemon: %s", err.Error())
		return err
	}
	defer probe.client.Close()

	if _, err := probe.client.ServerVersion(context.Background()); err != nil {
		logging.GetLogger().Errorf("Failed to connect to Docker daemon: %s", err.Error())
		return err
	}

	eventsFilter := filters.NewArgs()
	eventsFilter.Add("event", "start")
	eventsFilter.Add("event", "die")

	ctx, cancel := context.WithCancel(context.Background())
	eventChan, errChan := probe.client.Events(ctx, types.EventsOptions{Filters: eventsFilter})

	probe.cancel = cancel
	probe.wg.Add(2)

	probe.connected.Store(true)
	defer probe.connected.Store(false)

	go func() {
		defer probe.wg.Done()

		containers, err := probe.client.ContainerList(context.Background(), types.ContainerListOptions{})
		if err != nil {
			logging.GetLogger().Errorf("Failed to list containers: %s", err.Error())
			return
		}

		for _, c := range containers {
			if atomic.LoadInt64(&probe.state) != common.RunningState {
				break
			}
			probe.registerContainer(c.ID)
		}
	}()

	defer probe.wg.Done()

	for {
		select {
		case err := <-errChan:
			if atomic.LoadInt64(&probe.state) != common.StoppingState {
				logging.GetLogger().Errorf("Got error while waiting for Docker event: %s", err.Error())
			}
			return err
		case event := <-eventChan:
			probe.handleDockerEvent(&event)
		}
	}
}

// Start the probe
func (probe *DockerProbe) Start() {
	if !atomic.CompareAndSwapInt64(&probe.state, common.StoppedState, common.RunningState) {
		return
	}

	go func() {
		for {
			state := atomic.LoadInt64(&probe.state)
			if state == common.StoppingState || state == common.StoppedState {
				break
			}

			if probe.connect() != nil {
				time.Sleep(1 * time.Second)
			}

			probe.wg.Wait()
		}
	}()
}

// Stop the probe
func (probe *DockerProbe) Stop() {
	if !atomic.CompareAndSwapInt64(&probe.state, common.RunningState, common.StoppingState) {
		return
	}

	if probe.connected.Load() == true {
		probe.cancel()
		probe.wg.Wait()
	}

	atomic.StoreInt64(&probe.state, common.StoppedState)
}

// NewDockerProbe creates a new topology Docker probe
func NewDockerProbe(nsProbe *ns.NetNSProbe, dockerURL string) (*DockerProbe, error) {
	dockerProbe := &DockerProbe{
		NetNSProbe:   nsProbe,
		url:          dockerURL,
		containerMap: make(map[string]containerInfo),
		state:        common.StoppedState,
	}

	if path := config.GetConfig().GetString("docker.netns.run_path"); path != "" {
		nsProbe.Watch(path)
	}

	return dockerProbe, nil
}
<|endoftext|>