<commit_before>package transport

import (
	"compress/gzip"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"strings"
	"time"

	"github.com/algolia/algoliasearch-client-go/v3/algolia/call"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/compression"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/debug"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/errs"
	iopt "github.com/algolia/algoliasearch-client-go/v3/algolia/internal/opt"
)

const version = "3.24.0"

type Transport struct {
	requester     Requester
	retryStrategy *RetryStrategy
	headers       map[string]string
	compression   compression.Compression
}

func New(
	hosts []*StatefulHost,
	requester Requester,
	appID string,
	apiKey string,
	readTimeout time.Duration,
	writeTimeout time.Duration,
	defaultHeaders map[string]string,
	extraUserAgent string,
	compression compression.Compression,
) *Transport {

	if requester == nil {
		requester = newDefaultRequester()
	}

	userAgents := []string{
		fmt.Sprintf("Algolia for Go (%s)", version),
		fmt.Sprintf("Go (%s)", runtime.Version()),
	}

	if extraUserAgent != "" {
		userAgents = append([]string{extraUserAgent}, userAgents...)
	}

	headers := map[string]string{
		"Connection":               "Keep-Alive",
		"Content-Type":             "application/json; charset=utf-8",
		"User-Agent":               strings.Join(userAgents, ";"),
		"X-Algolia-Application-Id": appID,
		"X-Algolia-API-Key":        apiKey,
	}

	for k, v := range defaultHeaders {
		headers[k] = v
	}

	return &Transport{
		requester:     requester,
		retryStrategy: newRetryStrategy(hosts, readTimeout, writeTimeout),
		headers:       headers,
		compression:   compression,
	}
}

func (t *Transport) Request(
	res interface{},
	method string,
	path string,
	body interface{},
	k call.Kind,
	opts ...interface{},
) error {
	var (
		ctx       = iopt.ExtractContext(opts...)
		headers   = t.headers
		urlParams = make(map[string]string)
	)

	if extraHeaders := iopt.ExtractExtraHeaders(opts...); extraHeaders != nil {
		headers = mergeHeaders(headers, extraHeaders.Get())
	}

	if extraURLParams := iopt.ExtractExtraURLParams(opts...); extraURLParams != nil {
		urlParams = extraURLParams.Get()
	}

	if forwardToReplicas := iopt.ExtractForwardToReplicas(opts...); forwardToReplicas != nil {
		urlParams["forwardToReplicas"] = fmt.Sprintf("%t", forwardToReplicas.Get())
	}

	exposeIntermediateNetworkErrors := iopt.ExtractExposeIntermediateNetworkErrors(opts...).Get()
	var intermediateNetworkErrors []error

	for _, h := range t.retryStrategy.GetTryableHosts(k) {
		req, err := buildRequest(t.compression, method, h.host, path, body, headers, urlParams)
		if err != nil {
			return err
		}

		// Handle per-request timeout by using a context with timeout.
		// Note that because we are in a loop, the cancel() callback cannot be
		// deferred. Instead, we call it precisely after the end of each loop or
		// before the early returns, but when we do so, we do it **after**
		// reading the body content of the response. Otherwise, a `context
		// cancelled` error may happen when the body is read.
		perRequestCtx, cancel := context.WithTimeout(ctx, h.timeout)
		req = req.WithContext(perRequestCtx)
		bodyRes, code, err := t.request(req)

		// Context error only returns a non-nil error upon context
		// cancellation, which is a signal we interpret as an early return.
		// Indeed, we do not want to retry on other hosts if the context is
		// already cancelled.
		if ctx.Err() != nil {
			cancel()
			return err
		}

		switch t.retryStrategy.Decide(h, code, err) {
		case Success:
			err = unmarshalTo(bodyRes, &res)
			cancel()
			return err
		case Failure:
			if bodyRes != nil {
				err = unmarshalToError(bodyRes)
			} else if err == nil {
				err = fmt.Errorf("undefined network error with code: %v", code)
			}
			cancel()
			return err
		default:
			if err != nil {
				intermediateNetworkErrors = append(intermediateNetworkErrors, err)
			} else {
				responseErr := unmarshalToError(bodyRes)
				intermediateNetworkErrors = append(intermediateNetworkErrors, responseErr)
			}
			if bodyRes != nil {
				if err = bodyRes.Close(); err != nil {
					cancel()
					return fmt.Errorf("cannot close response's body before retry: %v", err)
				}
			}
		}

		cancel()
	}

	if exposeIntermediateNetworkErrors {
		return errs.NewNoMoreHostToTryError(intermediateNetworkErrors...)
	}

	return errs.ErrNoMoreHostToTry
}

func (t *Transport) request(req *http.Request) (io.ReadCloser, int, error) {
	debug.Display(req)
	res, err := t.requester.Request(req)
	debug.Display(res)

	if err != nil {
		msg := fmt.Sprintf("cannot perform request:\n\terror=%v\n\tmethod=%s\n\turl=%s", err, req.Method, req.URL)
		nerr, ok := err.(net.Error)
		if ok {
			// Because net.Error and error have different meanings for the
			// retry strategy, we cannot simply return a new error, which
			// would make all net.Error simple errors instead. To keep this
			// behaviour, we wrap the message into a custom netError that
			// implements the net.Error interface if the original error was
			// already a net.Error.
			err = errs.NetError(nerr, msg)
		} else {
			err = errors.New(msg)
		}
		return nil, 0, err
	}

	return res.Body, res.StatusCode, nil
}

func mergeHeaders(defaultHeaders, extraHeaders map[string]string) map[string]string {
	headers := make(map[string]string)

	for key, value := range defaultHeaders {
		headers[key] = value
	}
	for key, value := range extraHeaders {
		headers[key] = value
	}

	return headers
}

func wrapJSONEncoder(in interface{}) io.ReadCloser {
	pr, pw := io.Pipe()
	go func() {
		errEncode := json.NewEncoder(pw).Encode(in)
		errClose := pw.Close()
		if errEncode != nil {
			debug.Printf("cannot JSON encode request body: %v", errEncode)
		}
		if errClose != nil {
			debug.Printf("cannot close JSON encoder writer for request body: %v", errClose)
		}
	}()
	return pr
}

func wrapGZIPEncoder(in io.ReadCloser) io.ReadCloser {
	pr, pw := io.Pipe()
	go func() {
		gw := gzip.NewWriter(pw)
		_, errCopy := io.Copy(gw, in)
		errCloseGw := gw.Close()
		errClosePw := pw.Close()
		if errCopy != nil {
			debug.Printf("cannot GZIP request body: %v", errCopy)
		}
		if errCloseGw != nil {
			debug.Printf("cannot close gzip.Writer of request body: %v", errCloseGw)
		}
		if errClosePw != nil {
			debug.Printf("cannot close io.PipeWriter of request body: %v", errClosePw)
		}
	}()
	return pr
}

func buildRequestWithoutBody(method, url string) (*http.Request, error) {
	return http.NewRequest(method, url, nil)
}

func buildRequestWithBody(method, url string, body interface{}, c compression.Compression) (*http.Request, error) {
	var r io.ReadCloser
	jsonEncoder := wrapJSONEncoder(body)
	switch c {
	case compression.GZIP:
		r = wrapGZIPEncoder(jsonEncoder)
	case compression.None:
		r = jsonEncoder
	default:
		r = jsonEncoder
	}
	return http.NewRequest(method, url, r)
}

func buildRequest(
	c compression.Compression,
	method string,
	host string,
	path string,
	body interface{},
	headers map[string]string,
	urlParams map[string]string,
) (req *http.Request, err error) {

	urlStr := "https://" + host + path
	isCompressionEnabled := shouldCompress(c, method, body)

	if body == nil {
		req, err = buildRequestWithoutBody(method, urlStr)
	} else {
		req, err = buildRequestWithBody(method, urlStr, body, c)
	}

	if err != nil {
		return nil, fmt.Errorf("cannot instantiate request:\n\terr=%v\n\tmethod=%s\n\turl=%s\n\tbody=%#v", err, method, urlStr, body)
	}

	// Add headers
	for k, v := range headers {
		req.Header.Add(k, v)
	}

	// Add Content-Encoding header, if needed
	if isCompressionEnabled {
		switch c {
		case compression.GZIP:
			req.Header.Add("Content-Encoding", "gzip")
		default:
			// Do nothing
		}
	}

	// Add URL params
	values := req.URL.Query()
	for k, v := range urlParams {
		values.Set(k, v)
	}
	req.URL.RawQuery = values.Encode()

	return req, nil
}

func unmarshalTo(r io.ReadCloser, v interface{}) error {
	body, err := ioutil.ReadAll(r)
	errClose := r.Close()
	if err != nil {
		return fmt.Errorf("cannot read body: %v", err)
	}
	err = json.Unmarshal(body, &v)
	if err != nil {
		return fmt.Errorf("cannot deserialize response's body: %v: %s", err, string(body))
	}
	if errClose != nil {
		return fmt.Errorf("cannot close response's body: %v", errClose)
	}
	return nil
}

func unmarshalToError(r io.ReadCloser) error {
	var algoliaErr errs.AlgoliaErr
	err := unmarshalTo(r, &algoliaErr)
	if err != nil {
		return err
	}
	return &algoliaErr
}

func shouldCompress(c compression.Compression, method string, body interface{}) bool {
	isValidMethod := method == http.MethodPut || method == http.MethodPost
	isCompressionEnabled := c != compression.None
	isBodyNonEmpty := body != nil
	return isCompressionEnabled && isValidMethod && isBodyNonEmpty
}
<commit_msg>chore: update version to 3.25.0 [skip ci]<commit_after>package transport

import (
	"compress/gzip"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"strings"
	"time"

	"github.com/algolia/algoliasearch-client-go/v3/algolia/call"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/compression"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/debug"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/errs"
	iopt "github.com/algolia/algoliasearch-client-go/v3/algolia/internal/opt"
)

const version = "3.25.0"

type Transport struct {
	requester     Requester
	retryStrategy *RetryStrategy
	headers       map[string]string
	compression   compression.Compression
}

func New(
	hosts []*StatefulHost,
	requester Requester,
	appID string,
	apiKey string,
	readTimeout time.Duration,
	writeTimeout time.Duration,
	defaultHeaders map[string]string,
	extraUserAgent string,
	compression compression.Compression,
) *Transport {
	if requester == nil {
		requester = newDefaultRequester()
	}

	userAgents := []string{
		fmt.Sprintf("Algolia for Go (%s)", version),
		fmt.Sprintf("Go (%s)", runtime.Version()),
	}

	if extraUserAgent != "" {
		userAgents = append([]string{extraUserAgent}, userAgents...)
	}

	headers := map[string]string{
		"Connection":               "Keep-Alive",
		"Content-Type":             "application/json; charset=utf-8",
		"User-Agent":               strings.Join(userAgents, ";"),
		"X-Algolia-Application-Id": appID,
		"X-Algolia-API-Key":        apiKey,
	}

	for k, v := range defaultHeaders {
		headers[k] = v
	}

	return &Transport{
		requester:     requester,
		retryStrategy: newRetryStrategy(hosts, readTimeout, writeTimeout),
		headers:       headers,
		compression:   compression,
	}
}

func (t *Transport) Request(
	res interface{},
	method string,
	path string,
	body interface{},
	k call.Kind,
	opts ...interface{},
) error {
	var (
		ctx       = iopt.ExtractContext(opts...)
		headers   = t.headers
		urlParams = make(map[string]string)
	)

	if extraHeaders := iopt.ExtractExtraHeaders(opts...); extraHeaders != nil {
		headers = mergeHeaders(headers, extraHeaders.Get())
	}

	if extraURLParams := iopt.ExtractExtraURLParams(opts...); extraURLParams != nil {
		urlParams = extraURLParams.Get()
	}

	if forwardToReplicas := iopt.ExtractForwardToReplicas(opts...); forwardToReplicas != nil {
		urlParams["forwardToReplicas"] = fmt.Sprintf("%t", forwardToReplicas.Get())
	}

	exposeIntermediateNetworkErrors := iopt.ExtractExposeIntermediateNetworkErrors(opts...).Get()
	var intermediateNetworkErrors []error

	for _, h := range t.retryStrategy.GetTryableHosts(k) {
		req, err := buildRequest(t.compression, method, h.host, path, body, headers, urlParams)
		if err != nil {
			return err
		}

		// Handle per-request timeout by using a context with timeout.
		// Note that because we are in a loop, the cancel() callback cannot be
		// deferred. Instead, we call it precisely after the end of each loop or
		// before the early returns, but when we do so, we do it **after**
		// reading the body content of the response. Otherwise, a `context
		// cancelled` error may happen when the body is read.
		perRequestCtx, cancel := context.WithTimeout(ctx, h.timeout)
		req = req.WithContext(perRequestCtx)
		bodyRes, code, err := t.request(req)

		// Context error only returns a non-nil error upon context
		// cancellation, which is a signal we interpret as an early return.
		// Indeed, we do not want to retry on other hosts if the context is
		// already cancelled.
		if ctx.Err() != nil {
			cancel()
			return err
		}

		switch t.retryStrategy.Decide(h, code, err) {
		case Success:
			err = unmarshalTo(bodyRes, &res)
			cancel()
			return err
		case Failure:
			if bodyRes != nil {
				err = unmarshalToError(bodyRes)
			} else if err == nil {
				err = fmt.Errorf("undefined network error with code: %v", code)
			}
			cancel()
			return err
		default:
			if err != nil {
				intermediateNetworkErrors = append(intermediateNetworkErrors, err)
			} else {
				responseErr := unmarshalToError(bodyRes)
				intermediateNetworkErrors = append(intermediateNetworkErrors, responseErr)
			}
			if bodyRes != nil {
				if err = bodyRes.Close(); err != nil {
					cancel()
					return fmt.Errorf("cannot close response's body before retry: %v", err)
				}
			}
		}

		cancel()
	}

	if exposeIntermediateNetworkErrors {
		return errs.NewNoMoreHostToTryError(intermediateNetworkErrors...)
	}

	return errs.ErrNoMoreHostToTry
}

func (t *Transport) request(req *http.Request) (io.ReadCloser, int, error) {
	debug.Display(req)
	res, err := t.requester.Request(req)
	debug.Display(res)

	if err != nil {
		msg := fmt.Sprintf("cannot perform request:\n\terror=%v\n\tmethod=%s\n\turl=%s", err, req.Method, req.URL)
		nerr, ok := err.(net.Error)
		if ok {
			// Because net.Error and error have different meanings for the
			// retry strategy, we cannot simply return a new error, which
			// would make all net.Error simple errors instead. To keep this
			// behaviour, we wrap the message into a custom netError that
			// implements the net.Error interface if the original error was
			// already a net.Error.
			err = errs.NetError(nerr, msg)
		} else {
			err = errors.New(msg)
		}
		return nil, 0, err
	}

	return res.Body, res.StatusCode, nil
}

func mergeHeaders(defaultHeaders, extraHeaders map[string]string) map[string]string {
	headers := make(map[string]string)

	for key, value := range defaultHeaders {
		headers[key] = value
	}
	for key, value := range extraHeaders {
		headers[key] = value
	}

	return headers
}

func wrapJSONEncoder(in interface{}) io.ReadCloser {
	pr, pw := io.Pipe()
	go func() {
		errEncode := json.NewEncoder(pw).Encode(in)
		errClose := pw.Close()
		if errEncode != nil {
			debug.Printf("cannot JSON encode request body: %v", errEncode)
		}
		if errClose != nil {
			debug.Printf("cannot close JSON encoder writer for request body: %v", errClose)
		}
	}()
	return pr
}

func wrapGZIPEncoder(in io.ReadCloser) io.ReadCloser {
	pr, pw := io.Pipe()
	go func() {
		gw := gzip.NewWriter(pw)
		_, errCopy := io.Copy(gw, in)
		errCloseGw := gw.Close()
		errClosePw := pw.Close()
		if errCopy != nil {
			debug.Printf("cannot GZIP request body: %v", errCopy)
		}
		if errCloseGw != nil {
			debug.Printf("cannot close gzip.Writer of request body: %v", errCloseGw)
		}
		if errClosePw != nil {
			debug.Printf("cannot close io.PipeWriter of request body: %v", errClosePw)
		}
	}()
	return pr
}

func buildRequestWithoutBody(method, url string) (*http.Request, error) {
	return http.NewRequest(method, url, nil)
}

func buildRequestWithBody(method, url string, body interface{}, c compression.Compression) (*http.Request, error) {
	var r io.ReadCloser
	jsonEncoder := wrapJSONEncoder(body)
	switch c {
	case compression.GZIP:
		r = wrapGZIPEncoder(jsonEncoder)
	case compression.None:
		r = jsonEncoder
	default:
		r = jsonEncoder
	}
	return http.NewRequest(method, url, r)
}

func buildRequest(
	c compression.Compression,
	method string,
	host string,
	path string,
	body interface{},
	headers map[string]string,
	urlParams map[string]string,
) (req *http.Request, err error) {
	urlStr := "https://" + host + path
	isCompressionEnabled := shouldCompress(c, method, body)

	if body == nil {
		req, err = buildRequestWithoutBody(method, urlStr)
	} else {
		req, err = buildRequestWithBody(method, urlStr, body, c)
	}

	if err != nil {
		return nil, fmt.Errorf("cannot instantiate request:\n\terr=%v\n\tmethod=%s\n\turl=%s\n\tbody=%#v", err, method, urlStr, body)
	}

	// Add headers
	for k, v := range headers {
		req.Header.Add(k, v)
	}

	// Add Content-Encoding header, if needed
	if isCompressionEnabled {
		switch c {
		case compression.GZIP:
			req.Header.Add("Content-Encoding", "gzip")
		default:
			// Do nothing
		}
	}

	// Add URL params
	values := req.URL.Query()
	for k, v := range urlParams {
		values.Set(k, v)
	}
	req.URL.RawQuery = values.Encode()

	return req, nil
}

func unmarshalTo(r io.ReadCloser, v interface{}) error {
	body, err := ioutil.ReadAll(r)
	errClose := r.Close()
	if err != nil {
		return fmt.Errorf("cannot read body: %v", err)
	}
	err = json.Unmarshal(body, &v)
	if err != nil {
		return fmt.Errorf("cannot deserialize response's body: %v: %s", err, string(body))
	}
	if errClose != nil {
		return fmt.Errorf("cannot close response's body: %v", errClose)
	}
	return nil
}

func unmarshalToError(r io.ReadCloser) error {
	var algoliaErr errs.AlgoliaErr
	err := unmarshalTo(r, &algoliaErr)
	if err != nil {
		return err
	}
	return &algoliaErr
}

func shouldCompress(c compression.Compression, method string, body interface{}) bool {
	isValidMethod := method == http.MethodPut || method == http.MethodPost
	isCompressionEnabled := c != compression.None
	isBodyNonEmpty := body != nil
	return isCompressionEnabled && isValidMethod && isBodyNonEmpty
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 storage

import (
	"fmt"
	"math/rand"
	"path"

	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/storage/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"

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

var (
	BusyBoxImage          = imageutils.GetE2EImage(imageutils.BusyBox)
	durationForStuckMount = 110 * time.Second
)

var _ = utils.SIGDescribe("Detaching volumes", func() {
	f := framework.NewDefaultFramework("flexvolume")

	// note that namespace deletion is handled by delete-namespace flag

	var cs clientset.Interface
	var ns *v1.Namespace
	var node v1.Node
	var suffix string

	BeforeEach(func() {
		framework.SkipUnlessProviderIs("gce", "local")
		framework.SkipUnlessMasterOSDistroIs("debian", "ubuntu", "gci", "custom")
		framework.SkipUnlessNodeOSDistroIs("debian", "ubuntu", "gci", "custom")
		framework.SkipUnlessSSHKeyPresent()

		cs = f.ClientSet
		ns = f.Namespace
		nodes := framework.GetReadySchedulableNodesOrDie(f.ClientSet)
		node = nodes.Items[rand.Intn(len(nodes.Items))]
		suffix = ns.Name
	})

	It("should not work when mount is in progress", func() {
		driver := "attachable-with-long-format"
		driverInstallAs := driver + "-" + suffix

		By(fmt.Sprintf("installing flexvolume %s on node %s as %s", path.Join(driverDir, driver), node.Name, driverInstallAs))
		installFlex(cs, &node, "k8s", driverInstallAs, path.Join(driverDir, driver))
		By(fmt.Sprintf("installing flexvolume %s on master as %s", path.Join(driverDir, driver), driverInstallAs))
		installFlex(cs, nil, "k8s", driverInstallAs, path.Join(driverDir, driver))
		volumeSource := v1.VolumeSource{
			FlexVolume: &v1.FlexVolumeSource{
				Driver: "k8s/" + driverInstallAs,
			},
		}

		clientPod := getFlexVolumePod(volumeSource, node.Name)
		By("Creating pod that uses slow format volume")
		pod, err := cs.CoreV1().Pods(ns.Name).Create(clientPod)
		Expect(err).NotTo(HaveOccurred())

		uniqueVolumeName := getUniqueVolumeName(pod, driverInstallAs)

		By("waiting for volume-in-use on the node after pod creation")
		err = waitForVolumesInUse(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume in use")

		By("Deleting the flexvolume pod")
		err = framework.DeletePodWithWait(f, cs, pod)
		Expect(err).NotTo(HaveOccurred(), "in deleting the pod")

		// Wait a bit for node to sync the volume status
		time.Sleep(30 * time.Second)

		By("waiting for volume-in-use on the node after pod deletion")
		err = waitForVolumesInUse(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume in use")

		// Wait for 110s because mount device operation has a sleep of 120 seconds
		// we previously already waited for 30s.
		time.Sleep(durationForStuckMount)

		By("waiting for volume to disappear from node in-use")
		err = waitForVolumesNotInUse(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume to be removed from in-use")

		By(fmt.Sprintf("uninstalling flexvolume %s from node %s", driverInstallAs, node.Name))
		uninstallFlex(cs, &node, "k8s", driverInstallAs)
		By(fmt.Sprintf("uninstalling flexvolume %s from master", driverInstallAs))
		uninstallFlex(cs, nil, "k8s", driverInstallAs)
	})
})

func getUniqueVolumeName(pod *v1.Pod, driverName string) string {
	return fmt.Sprintf("flexvolume-k8s/%s/%s", driverName, pod.Spec.Volumes[0].Name)
}

func waitForVolumesNotInUse(client clientset.Interface, nodeName, volumeName string) error {
	return wait.PollImmediate(10*time.Second, 60*time.Second, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("error fetching node %s with %v", nodeName, err)
		}
		volumeInUSe := node.Status.VolumesInUse
		for _, volume := range volumeInUSe {
			if string(volume) == volumeName {
				return false, nil
			}
		}
		return true, nil
	})
}

func waitForVolumesInUse(client clientset.Interface, nodeName, volumeName string) error {
	return wait.PollImmediate(10*time.Second, 60*time.Second, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("error fetching node %s with %v", nodeName, err)
		}
		volumeInUSe := node.Status.VolumesInUse
		for _, volume := range volumeInUSe {
			if string(volume) == volumeName {
				return true, nil
			}
		}
		return false, nil
	})
}

func getFlexVolumePod(volumeSource v1.VolumeSource, nodeName string) *v1.Pod {
	var gracePeriod int64
	clientPod := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: "flexvolume-detach-test" + "-client",
			Labels: map[string]string{
				"role": "flexvolume-detach-test" + "-client",
			},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:       "flexvolume-detach-test" + "-client",
					Image:      BusyBoxImage,
					WorkingDir: "/opt",
					// An imperative and easily debuggable container which reads vol contents for
					// us to scan in the tests or by eye.
					// We expect that /opt is empty in the minimal containers which we use in this test.
					Command: []string{
						"/bin/sh",
						"-c",
						"while true ; do cat /opt/foo/index.html ; sleep 2 ; ls -altrh /opt/  ; sleep 2 ; done ",
					},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      "test-long-detach-flex",
							MountPath: "/opt/foo",
						},
					},
				},
			},
			TerminationGracePeriodSeconds: &gracePeriod,
			SecurityContext: &v1.PodSecurityContext{
				SELinuxOptions: &v1.SELinuxOptions{
					Level: "s0:c0,c1",
				},
			},
			Volumes: []v1.Volume{
				{
					Name:         "test-long-detach-flex",
					VolumeSource: volumeSource,
				},
			},
			NodeName: nodeName,
		},
	}
	return clientPod
}
<commit_msg>Fix flake with e2e test that checks detach while mount in progress<commit_after>/*
Copyright 2017 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 storage

import (
	"fmt"
	"math/rand"
	"path"

	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/storage/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"

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

var (
	BusyBoxImage          = imageutils.GetE2EImage(imageutils.BusyBox)
	durationForStuckMount = 110 * time.Second
)

var _ = utils.SIGDescribe("Detaching volumes", func() {
	f := framework.NewDefaultFramework("flexvolume")

	// note that namespace deletion is handled by delete-namespace flag

	var cs clientset.Interface
	var ns *v1.Namespace
	var node v1.Node
	var suffix string

	BeforeEach(func() {
		framework.SkipUnlessProviderIs("gce", "local")
		framework.SkipUnlessMasterOSDistroIs("debian", "ubuntu", "gci", "custom")
		framework.SkipUnlessNodeOSDistroIs("debian", "ubuntu", "gci", "custom")
		framework.SkipUnlessSSHKeyPresent()

		cs = f.ClientSet
		ns = f.Namespace
		nodes := framework.GetReadySchedulableNodesOrDie(f.ClientSet)
		node = nodes.Items[rand.Intn(len(nodes.Items))]
		suffix = ns.Name
	})

	It("should not work when mount is in progress", func() {
		driver := "attachable-with-long-format"
		driverInstallAs := driver + "-" + suffix

		By(fmt.Sprintf("installing flexvolume %s on node %s as %s", path.Join(driverDir, driver), node.Name, driverInstallAs))
		installFlex(cs, &node, "k8s", driverInstallAs, path.Join(driverDir, driver))
		By(fmt.Sprintf("installing flexvolume %s on master as %s", path.Join(driverDir, driver), driverInstallAs))
		installFlex(cs, nil, "k8s", driverInstallAs, path.Join(driverDir, driver))
		volumeSource := v1.VolumeSource{
			FlexVolume: &v1.FlexVolumeSource{
				Driver: "k8s/" + driverInstallAs,
			},
		}

		clientPod := getFlexVolumePod(volumeSource, node.Name)
		By("Creating pod that uses slow format volume")
		pod, err := cs.CoreV1().Pods(ns.Name).Create(clientPod)
		Expect(err).NotTo(HaveOccurred())

		uniqueVolumeName := getUniqueVolumeName(pod, driverInstallAs)

		By("waiting for volumes to be attached to node")
		err = waitForVolumesAttached(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume to attach to %s node", node.Name)

		By("waiting for volume-in-use on the node after pod creation")
		err = waitForVolumesInUse(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume in use")

		By("waiting for kubelet to start mounting the volume")
		time.Sleep(20 * time.Second)

		By("Deleting the flexvolume pod")
		err = framework.DeletePodWithWait(f, cs, pod)
		Expect(err).NotTo(HaveOccurred(), "in deleting the pod")

		// Wait a bit for node to sync the volume status
		time.Sleep(30 * time.Second)

		By("waiting for volume-in-use on the node after pod deletion")
		err = waitForVolumesInUse(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume in use")

		// Wait for 110s because mount device operation has a sleep of 120 seconds
		// we previously already waited for 30s.
		time.Sleep(durationForStuckMount)

		By("waiting for volume to disappear from node in-use")
		err = waitForVolumesNotInUse(cs, node.Name, uniqueVolumeName)
		Expect(err).NotTo(HaveOccurred(), "while waiting for volume to be removed from in-use")

		By(fmt.Sprintf("uninstalling flexvolume %s from node %s", driverInstallAs, node.Name))
		uninstallFlex(cs, &node, "k8s", driverInstallAs)
		By(fmt.Sprintf("uninstalling flexvolume %s from master", driverInstallAs))
		uninstallFlex(cs, nil, "k8s", driverInstallAs)
	})
})

func getUniqueVolumeName(pod *v1.Pod, driverName string) string {
	return fmt.Sprintf("flexvolume-k8s/%s/%s", driverName, pod.Spec.Volumes[0].Name)
}

func waitForVolumesNotInUse(client clientset.Interface, nodeName, volumeName string) error {
	return wait.PollImmediate(10*time.Second, 60*time.Second, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("error fetching node %s with %v", nodeName, err)
		}
		volumeInUSe := node.Status.VolumesInUse
		for _, volume := range volumeInUSe {
			if string(volume) == volumeName {
				return false, nil
			}
		}
		return true, nil
	})
}

func waitForVolumesAttached(client clientset.Interface, nodeName, volumeName string) error {
	return wait.PollImmediate(2*time.Second, 2*time.Minute, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("error fetching node %s with %v", nodeName, err)
		}
		volumeAttached := node.Status.VolumesAttached
		for _, volume := range volumeAttached {
			if string(volume.Name) == volumeName {
				return true, nil
			}
		}
		return false, nil
	})
}

func waitForVolumesInUse(client clientset.Interface, nodeName, volumeName string) error {
	return wait.PollImmediate(10*time.Second, 60*time.Second, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("error fetching node %s with %v", nodeName, err)
		}
		volumeInUSe := node.Status.VolumesInUse
		for _, volume := range volumeInUSe {
			if string(volume) == volumeName {
				return true, nil
			}
		}
		return false, nil
	})
}

func getFlexVolumePod(volumeSource v1.VolumeSource, nodeName string) *v1.Pod {
	var gracePeriod int64
	clientPod := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: "flexvolume-detach-test" + "-client",
			Labels: map[string]string{
				"role": "flexvolume-detach-test" + "-client",
			},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:       "flexvolume-detach-test" + "-client",
					Image:      BusyBoxImage,
					WorkingDir: "/opt",
					// An imperative and easily debuggable container which reads vol contents for
					// us to scan in the tests or by eye.
					// We expect that /opt is empty in the minimal containers which we use in this test.
					Command: []string{
						"/bin/sh",
						"-c",
						"while true ; do cat /opt/foo/index.html ; sleep 2 ; ls -altrh /opt/  ; sleep 2 ; done ",
					},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      "test-long-detach-flex",
							MountPath: "/opt/foo",
						},
					},
				},
			},
			TerminationGracePeriodSeconds: &gracePeriod,
			SecurityContext: &v1.PodSecurityContext{
				SELinuxOptions: &v1.SELinuxOptions{
					Level: "s0:c0,c1",
				},
			},
			Volumes: []v1.Volume{
				{
					Name:         "test-long-detach-flex",
					VolumeSource: volumeSource,
				},
			},
			NodeName: nodeName,
		},
	}
	return clientPod
}
<|endoftext|>


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

package calc

import (
	"reflect"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"

	"github.com/gavv/monotime"

	"github.com/projectcalico/felix/config"
	"github.com/projectcalico/felix/dispatcher"
	"github.com/projectcalico/felix/proto"
	"github.com/projectcalico/libcalico-go/lib/backend/api"
	"github.com/projectcalico/libcalico-go/lib/health"
)

const (
	tickInterval    = 10 * time.Millisecond
	leakyBucketSize = 10
)

var (
	dataplaneStatusGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "felix_resync_state",
		Help: "Current datastore state.",
	})
	resyncsStarted = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "felix_resyncs_started",
		Help: "Current datastore state.",
	})
	statusToGaugeValue = map[api.SyncStatus]float64{
		api.WaitForDatastore: 1,
		api.ResyncInProgress: 2,
		api.InSync:           3,
	}
	countUpdatesProcessed = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "felix_calc_graph_updates_processed",
		Help: "Number of datastore updates processed by the calculation graph.",
	}, []string{"type"})
	countOutputEvents = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "felix_calc_graph_output_events",
		Help: "Number of events emitted by the calculation graph.",
	})
	summaryUpdateTime = prometheus.NewSummary(prometheus.SummaryOpts{
		Name: "felix_calc_graph_update_time_seconds",
		Help: "Seconds to update calculation graph for each datastore OnUpdate call.",
	})
)

func init() {
	prometheus.MustRegister(dataplaneStatusGauge)
	prometheus.MustRegister(resyncsStarted)
	prometheus.MustRegister(countUpdatesProcessed)
	prometheus.MustRegister(countOutputEvents)
	prometheus.MustRegister(summaryUpdateTime)
}

type AsyncCalcGraph struct {
	Dispatcher       *dispatcher.Dispatcher
	inputEvents      chan interface{}
	outputEvents     chan<- interface{}
	eventBuffer      *EventSequencer
	beenInSync       bool
	needToSendInSync bool
	syncStatusNow    api.SyncStatus
	healthAggregator *health.HealthAggregator

	flushTicks       <-chan time.Time
	flushLeakyBucket int
	dirty            bool

	debugHangC <-chan time.Time
}

const (
	healthName     = "async_calc_graph"
	healthInterval = 10 * time.Second
)

func NewAsyncCalcGraph(conf *config.Config, outputEvents chan<- interface{}, healthAggregator *health.HealthAggregator) *AsyncCalcGraph {
	eventBuffer := NewEventBuffer(conf)
	disp := NewCalculationGraph(eventBuffer, conf.FelixHostname)
	g := &AsyncCalcGraph{
		inputEvents:      make(chan interface{}, 10),
		outputEvents:     outputEvents,
		Dispatcher:       disp,
		eventBuffer:      eventBuffer,
		healthAggregator: healthAggregator,
	}
	if conf.DebugSimulateCalcGraphHangAfter != 0 {
		log.WithField("delay", conf.DebugSimulateCalcGraphHangAfter).Warn(
			"Simulating a calculation graph hang.")
		g.debugHangC = time.After(conf.DebugSimulateDataplaneHangAfter)
	}
	eventBuffer.Callback = g.onEvent
	if healthAggregator != nil {
		healthAggregator.RegisterReporter(healthName, &health.HealthReport{Live: true, Ready: true}, healthInterval*2)
	}
	return g
}

func (acg *AsyncCalcGraph) OnUpdates(updates []api.Update) {
	log.Debugf("Got %v updates; queueing", len(updates))
	acg.inputEvents <- updates
}

func (acg *AsyncCalcGraph) OnStatusUpdated(status api.SyncStatus) {
	log.Debugf("Status updated: %v; queueing", status)
	acg.inputEvents <- status
	dataplaneStatusGauge.Set(statusToGaugeValue[status])
	if status == api.ResyncInProgress {
		resyncsStarted.Inc()
	}
}

func (acg *AsyncCalcGraph) loop() {
	log.Info("AsyncCalcGraph running")
	healthTicks := time.NewTicker(healthInterval).C
	acg.reportHealth()
	for {
		select {
		case update := <-acg.inputEvents:
			switch update := update.(type) {
			case []api.Update:
				// Update; send it to the dispatcher.
				log.Debug("Pulled []KVPair off channel")
				updStartTime := monotime.Now()
				acg.Dispatcher.OnUpdates(update)
				updEndTime := monotime.Now()
				summaryUpdateTime.Observe((updEndTime - updStartTime).Seconds())
				// Record stats for the number of messages processed.
				for _, upd := range update {
					typeName := reflect.TypeOf(upd.Key).Name()
					count := countUpdatesProcessed.WithLabelValues(typeName)
					count.Inc()
				}
			case api.SyncStatus:
				// Sync status changed, check if we're now in-sync.
				log.WithField("status", update).Debug(
					"Pulled status update off channel")
				acg.syncStatusNow = update
				acg.Dispatcher.OnStatusUpdated(update)
				if update == api.InSync && !acg.beenInSync {
					log.Info("First time we've been in sync")
					acg.beenInSync = true
					acg.needToSendInSync = true
					acg.dirty = true
					if acg.flushLeakyBucket == 0 {
						// Force a flush.
						acg.flushLeakyBucket++
					}
				}
				acg.reportHealth()
			default:
				log.Panicf("Unexpected update: %#v", update)
			}
			acg.dirty = true
		case <-acg.flushTicks:
			// Timer tick: fill up the leaky bucket.
			if acg.flushLeakyBucket < leakyBucketSize {
				acg.flushLeakyBucket++
			}
		case <-healthTicks:
			acg.reportHealth()
		case <-acg.debugHangC:
			log.Warning("Debug hang simulation timer popped, hanging the calculation graph!!")
			time.Sleep(1 * time.Hour)
			log.Panic("Woke up after 1 hour, something's probably wrong with the test.")
		}
		acg.maybeFlush()
	}
}

func (acg *AsyncCalcGraph) reportHealth() {
	if acg.healthAggregator != nil {
		acg.healthAggregator.Report(healthName, &health.HealthReport{
			Live:  true,
			Ready: acg.syncStatusNow == api.InSync,
		})
	}
}

// maybeFlush flushes the event buffer if: we know it's dirty and we're not throttled.
func (acg *AsyncCalcGraph) maybeFlush() {
	if !acg.dirty {
		return
	}
	if acg.flushLeakyBucket > 0 {
		log.Debug("Not throttled: flushing event buffer")
		acg.flushLeakyBucket--
		acg.eventBuffer.Flush()
		if acg.needToSendInSync {
			log.Info("First flush after becoming in sync, sending InSync message.")
			acg.onEvent(&proto.InSync{})
			acg.needToSendInSync = false
		}
		acg.dirty = false
	} else {
		log.Debug("Throttled: not flushing event buffer")
	}
}

func (acg *AsyncCalcGraph) onEvent(event interface{}) {
	log.Debug("Sending output event on channel")
	acg.outputEvents <- event
	countOutputEvents.Inc()
	log.Debug("Sent output event on channel")
}

func (acg *AsyncCalcGraph) Start() {
	log.Info("Starting AsyncCalcGraph")
	flushTicker := time.NewTicker(tickInterval)
	acg.flushTicks = flushTicker.C
	go acg.loop()
}
<commit_msg>Fix mixup between dataplane and calc graph hang parameters<commit_after>// Copyright (c) 2016-2017 Tigera, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package calc

import (
	"reflect"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"

	"github.com/gavv/monotime"

	"github.com/projectcalico/felix/config"
	"github.com/projectcalico/felix/dispatcher"
	"github.com/projectcalico/felix/proto"
	"github.com/projectcalico/libcalico-go/lib/backend/api"
	"github.com/projectcalico/libcalico-go/lib/health"
)

const (
	tickInterval    = 10 * time.Millisecond
	leakyBucketSize = 10
)

var (
	dataplaneStatusGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "felix_resync_state",
		Help: "Current datastore state.",
	})
	resyncsStarted = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "felix_resyncs_started",
		Help: "Current datastore state.",
	})
	statusToGaugeValue = map[api.SyncStatus]float64{
		api.WaitForDatastore: 1,
		api.ResyncInProgress: 2,
		api.InSync:           3,
	}
	countUpdatesProcessed = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "felix_calc_graph_updates_processed",
		Help: "Number of datastore updates processed by the calculation graph.",
	}, []string{"type"})
	countOutputEvents = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "felix_calc_graph_output_events",
		Help: "Number of events emitted by the calculation graph.",
	})
	summaryUpdateTime = prometheus.NewSummary(prometheus.SummaryOpts{
		Name: "felix_calc_graph_update_time_seconds",
		Help: "Seconds to update calculation graph for each datastore OnUpdate call.",
	})
)

func init() {
	prometheus.MustRegister(dataplaneStatusGauge)
	prometheus.MustRegister(resyncsStarted)
	prometheus.MustRegister(countUpdatesProcessed)
	prometheus.MustRegister(countOutputEvents)
	prometheus.MustRegister(summaryUpdateTime)
}

type AsyncCalcGraph struct {
	Dispatcher       *dispatcher.Dispatcher
	inputEvents      chan interface{}
	outputEvents     chan<- interface{}
	eventBuffer      *EventSequencer
	beenInSync       bool
	needToSendInSync bool
	syncStatusNow    api.SyncStatus
	healthAggregator *health.HealthAggregator

	flushTicks       <-chan time.Time
	flushLeakyBucket int
	dirty            bool

	debugHangC <-chan time.Time
}

const (
	healthName     = "async_calc_graph"
	healthInterval = 10 * time.Second
)

func NewAsyncCalcGraph(conf *config.Config, outputEvents chan<- interface{}, healthAggregator *health.HealthAggregator) *AsyncCalcGraph {
	eventBuffer := NewEventBuffer(conf)
	disp := NewCalculationGraph(eventBuffer, conf.FelixHostname)
	g := &AsyncCalcGraph{
		inputEvents:      make(chan interface{}, 10),
		outputEvents:     outputEvents,
		Dispatcher:       disp,
		eventBuffer:      eventBuffer,
		healthAggregator: healthAggregator,
	}
	if conf.DebugSimulateCalcGraphHangAfter != 0 {
		log.WithField("delay", conf.DebugSimulateCalcGraphHangAfter).Warn(
			"Simulating a calculation graph hang.")
		g.debugHangC = time.After(conf.DebugSimulateCalcGraphHangAfter)
	}
	eventBuffer.Callback = g.onEvent
	if healthAggregator != nil {
		healthAggregator.RegisterReporter(healthName, &health.HealthReport{Live: true, Ready: true}, healthInterval*2)
	}
	return g
}

func (acg *AsyncCalcGraph) OnUpdates(updates []api.Update) {
	log.Debugf("Got %v updates; queueing", len(updates))
	acg.inputEvents <- updates
}

func (acg *AsyncCalcGraph) OnStatusUpdated(status api.SyncStatus) {
	log.Debugf("Status updated: %v; queueing", status)
	acg.inputEvents <- status
	dataplaneStatusGauge.Set(statusToGaugeValue[status])
	if status == api.ResyncInProgress {
		resyncsStarted.Inc()
	}
}

func (acg *AsyncCalcGraph) loop() {
	log.Info("AsyncCalcGraph running")
	healthTicks := time.NewTicker(healthInterval).C
	acg.reportHealth()
	for {
		select {
		case update := <-acg.inputEvents:
			switch update := update.(type) {
			case []api.Update:
				// Update; send it to the dispatcher.
				log.Debug("Pulled []KVPair off channel")
				updStartTime := monotime.Now()
				acg.Dispatcher.OnUpdates(update)
				updEndTime := monotime.Now()
				summaryUpdateTime.Observe((updEndTime - updStartTime).Seconds())
				// Record stats for the number of messages processed.
				for _, upd := range update {
					typeName := reflect.TypeOf(upd.Key).Name()
					count := countUpdatesProcessed.WithLabelValues(typeName)
					count.Inc()
				}
			case api.SyncStatus:
				// Sync status changed, check if we're now in-sync.
				log.WithField("status", update).Debug(
					"Pulled status update off channel")
				acg.syncStatusNow = update
				acg.Dispatcher.OnStatusUpdated(update)
				if update == api.InSync && !acg.beenInSync {
					log.Info("First time we've been in sync")
					acg.beenInSync = true
					acg.needToSendInSync = true
					acg.dirty = true
					if acg.flushLeakyBucket == 0 {
						// Force a flush.
						acg.flushLeakyBucket++
					}
				}
				acg.reportHealth()
			default:
				log.Panicf("Unexpected update: %#v", update)
			}
			acg.dirty = true
		case <-acg.flushTicks:
			// Timer tick: fill up the leaky bucket.
			if acg.flushLeakyBucket < leakyBucketSize {
				acg.flushLeakyBucket++
			}
		case <-healthTicks:
			acg.reportHealth()
		case <-acg.debugHangC:
			log.Warning("Debug hang simulation timer popped, hanging the calculation graph!!")
			time.Sleep(1 * time.Hour)
			log.Panic("Woke up after 1 hour, something's probably wrong with the test.")
		}
		acg.maybeFlush()
	}
}

func (acg *AsyncCalcGraph) reportHealth() {
	if acg.healthAggregator != nil {
		acg.healthAggregator.Report(healthName, &health.HealthReport{
			Live:  true,
			Ready: acg.syncStatusNow == api.InSync,
		})
	}
}

// maybeFlush flushes the event buffer if: we know it's dirty and we're not throttled.
func (acg *AsyncCalcGraph) maybeFlush() {
	if !acg.dirty {
		return
	}
	if acg.flushLeakyBucket > 0 {
		log.Debug("Not throttled: flushing event buffer")
		acg.flushLeakyBucket--
		acg.eventBuffer.Flush()
		if acg.needToSendInSync {
			log.Info("First flush after becoming in sync, sending InSync message.")
			acg.onEvent(&proto.InSync{})
			acg.needToSendInSync = false
		}
		acg.dirty = false
	} else {
		log.Debug("Throttled: not flushing event buffer")
	}
}

func (acg *AsyncCalcGraph) onEvent(event interface{}) {
	log.Debug("Sending output event on channel")
	acg.outputEvents <- event
	countOutputEvents.Inc()
	log.Debug("Sent output event on channel")
}

func (acg *AsyncCalcGraph) Start() {
	log.Info("Starting AsyncCalcGraph")
	flushTicker := time.NewTicker(tickInterval)
	acg.flushTicks = flushTicker.C
	go acg.loop()
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

var updateTopic = &Topic{
	Name:      "update",
	ShortHelp: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:     "update",
	ShortHelp: "updates heroku-cli",
	Args:      []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		fmt.Printf("updating to ", channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku")

func updateIfNeeded() {
	if !updateNeeded() {
		return
	}
	// TODO: update plugins
	manifest := getUpdateManifest()
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		must(err)
	}
	return f.ModTime().Add(20 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest() manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/heroku-cli/" + Channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		panic(err)
	}
}
<commit_msg>reduced update time while this is under development<commit_after>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

var updateTopic = &Topic{
	Name:      "update",
	ShortHelp: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:     "update",
	ShortHelp: "updates heroku-cli",
	Args:      []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		fmt.Printf("updating to ", channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku")

func updateIfNeeded() {
	if !updateNeeded() {
		return
	}
	// TODO: update plugins
	manifest := getUpdateManifest()
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		must(err)
	}
	// TODO: Increase the autoupdate time later
	return f.ModTime().Add(2 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest() manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/heroku-cli/" + Channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// DYNAMIC UPDATES
//
// Dynamic updates reuses the DNS message format, but renames the three of
// the sections. Question is Zone, Answer is Prerequisite, Authority is
// Update, only the Additional is not renamed. See RFC 2136 for the gory details.
//
// You can set a rather complex set of rules for the existence of absence of
// certain resource records or names in a zone to specify if resource records
// should be added or removed. The table from RFC 2136 supplemented with the Go
// DNS function shows which functions exist to specify the prerequisites.
//
// 3.2.4 - Table Of Metavalues Used In Prerequisite Section
//
//   CLASS    TYPE     RDATA    Meaning                          Function
//   ----------------------------------------------------------------------------
//   ANY      ANY      empty    Name is in use                   NameUsed
//   ANY      rrset    empty    RRset exists (value independent) RRsetUsedNoRdata
//   NONE     ANY      empty    Name is not in use               NameNotUsed
//   NONE     rrset    empty    RRset does not exist             RRsetNotUsed
//   zone     rrset    rr       RRset exists (value dependent)   RRsetUsedRdata
//
// The prerequisite section can also be left empty.
// If you have decided an the prerequisites you can tell what RRs should
// be added or deleted. The next table shows the options you have and
// what function to call.
// 3.4.2.6 - Table Of Metavalues Used In Update Section
//
//   CLASS    TYPE     RDATA    Meaning                          Function
//   -------------------------------------------------------------------------
//   ANY      ANY      empty    Delete all RRsets from a name    NameDelete
//   ANY      rrset    empty    Delete an RRset                  RRsetDelete
//   NONE     rrset    rr       Delete an RR from an RRset       RRsetDeleteRR
//   zone     rrset    rr       Add to an RRset                  RRsetAddRdata
//
package dns

type Update struct{ Msg }

func (u *Update) Zone() []Question {
	return u.Msg.Question
}

func (u *Update) Prereq() []RR {
	return u.Msg.Answer
}

func (u *Update) Update() []RR {
	return u.Msg.Ns
}

func (u *Update) Additional() []RR {
	return u.Msg.Extra
}

// NewUpdate creats a new DNS update packet.
func NewUpdate(zone string, class uint16) *Update {
	u := new(Update)
	u.MsgHdr.Opcode = OpcodeUpdate
	u.Question = make([]Question, 1)
	u.Question[0] = Question{zone, TypeSOA, class}
	return u
}

// The table from RFC 2136 supplemented with the Go DNS function.
//
// 3.2.4 - Table Of Metavalues Used In Prerequisite Section
//
//   CLASS    TYPE     RDATA    Meaning                           Function
//   ----------------------------------------------------------------------
//   ANY      ANY      empty    Name is in use                    NameUsed
//   ANY      rrset    empty    RRset exists (value independent)  RRsetUsedNoRdata
//   NONE     ANY      empty    Name is not in use                NameNotUsed
//   NONE     rrset    empty    RRset does not exist              RRsetNotUsed
//   zone     rrset    rr       RRset exists (value dependent)    RRsetUsedRdata

// NameUsed sets the RRs in the prereq section to
// "Name is in use" RRs. RFC 2136 section 2.4.4.
func (u *Update) NameUsed(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = &RR_ANY{Hdr: RR_Header{Name: r.Header().Name, Ttl: 0, Rrtype: TypeANY, Class: ClassANY}}
	}
}

// NameNotUsed sets the RRs in the prereq section to
// "Name is in not use" RRs. RFC 2136 section 2.4.5.
func (u *Update) NameNotUsed(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = &RR_ANY{Hdr: RR_Header{Name: r.Header().Name, Ttl: 0, Rrtype: TypeANY, Class: ClassNONE}}
	}
}

// RRsetUsedRdata sets the RRs in the prereq section to
// "RRset exists (value dependent -- with rdata)" RRs. RFC 2136 section 2.4.2.
func (u *Update) RRsetUsedRdata(rr []RR) {
	if len(u.Msg.Question) == 0 {
		panic("empty question section")
	}
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = r
		u.Answer[i].Header().Class = u.Msg.Question[0].Qclass
	}
}

// RRsetUsedNoRdata sets the RRs in the prereq section to
// "RRset exists (value independent -- no rdata)" RRs. RFC 2136 section 2.4.1.
func (u *Update) RRsetUsedNoRdata(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = r
		u.Answer[i].Header().Class = ClassANY
		u.Answer[i].Header().Ttl = 0
		u.Answer[i].Header().Rdlength = 0
	}
}

// RRsetNotUsed sets the RRs in the prereq section to
// "RRset does not exist" RRs. RFC 2136 section 2.4.3.
func (u *Update) RRsetNotUsed(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = r
		u.Answer[i].Header().Class = ClassNONE
		u.Answer[i].Header().Rdlength = 0
		u.Answer[i].Header().Ttl = 0
	}
}

// The table from RFC 2136 supplemented with the Go DNS function.
//
// 3.4.2.6 - Table Of Metavalues Used In Update Section
//
//   CLASS    TYPE     RDATA    Meaning                         Function
//   --------------------------------------------------------------------------
//   ANY      ANY      empty    Delete all RRsets from a name   NameDelete
//   ANY      rrset    empty    Delete an RRset                 RRsetDelete
//   NONE     rrset    rr       Delete an RR from an RRset      RRsetDeleteRR
//   zone     rrset    rr       Add to an RRset                 RRsetAddRdata

// RRsetAddRdata adds an complete RRset, see RFC 2136 section 2.5.1
func (u *Update) RRsetAddRdata(rr []RR) {
	if len(u.Msg.Question) == 0 {
		panic("empty question section")
	}
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = r
		u.Ns[i].Header().Class = u.Msg.Question[0].Qclass
	}
}

// RRsetDelete deletes an RRset, see RFC 2136 section 2.5.2
func (u *Update) RRsetDelete(rr []RR) {
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = r
		u.Ns[i].Header().Class = ClassANY
		u.Ns[i].Header().Rdlength = 0
		u.Ns[i].Header().Ttl = 0
	}
}

// NameDelete deletes all RRsets of a name, see RFC 2136 section 2.5.3
func (u *Update) NameDelete(rr []RR) {
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = &RR_ANY{Hdr: RR_Header{Name: r.Header().Name, Ttl: 0, Rrtype: TypeANY, Class: ClassANY}}
	}
}

// RRsetDeleteRR deletes RR from the RRSset, see RFC 2136 section 2.5.4
func (u *Update) RRsetDeleteRR(rr []RR) {
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = r
		u.Ns[i].Header().Class = ClassNONE
		u.Ns[i].Header().Ttl = 0
	}
}
<commit_msg>Dont create a new type for updates<commit_after>// DYNAMIC UPDATES
//
// Dynamic updates reuses the DNS message format, but renames the three of
// the sections. Question is Zone, Answer is Prerequisite, Authority is
// Update, only the Additional is not renamed. See RFC 2136 for the gory details.
//
// You can set a rather complex set of rules for the existence of absence of
// certain resource records or names in a zone to specify if resource records
// should be added or removed. The table from RFC 2136 supplemented with the Go
// DNS function shows which functions exist to specify the prerequisites.
//
// 3.2.4 - Table Of Metavalues Used In Prerequisite Section
//
//   CLASS    TYPE     RDATA    Meaning                          Function
//   ----------------------------------------------------------------------------
//   ANY      ANY      empty    Name is in use                   NameUsed
//   ANY      rrset    empty    RRset exists (value independent) RRsetUsedNoRdata
//   NONE     ANY      empty    Name is not in use               NameNotUsed
//   NONE     rrset    empty    RRset does not exist             RRsetNotUsed
//   zone     rrset    rr       RRset exists (value dependent)   RRsetUsedRdata
//
// The prerequisite section can also be left empty.
// If you have decided an the prerequisites you can tell what RRs should
// be added or deleted. The next table shows the options you have and
// what function to call.
// 3.4.2.6 - Table Of Metavalues Used In Update Section
//
//   CLASS    TYPE     RDATA    Meaning                          Function
//   -------------------------------------------------------------------------
//   ANY      ANY      empty    Delete all RRsets from a name    NameDelete
//   ANY      rrset    empty    Delete an RRset                  RRsetDelete
//   NONE     rrset    rr       Delete an RR from an RRset       RRsetDeleteRR
//   zone     rrset    rr       Add to an RRset                  RRsetAddRdata
//
package dns

func (u *Msg) Zone() []Question {
	return u.Question
}

func (u *Msg) Prereq() []RR {
	return u.Answer
}

func (u *Msg) Update() []RR {
	return u.Ns
}

func (u *Msg) Additional() []RR {
	return u.Extra
}

// NewUpdate creats a new DNS update packet.
func NewUpdate(zone string, class uint16) *Msg {
	u := new(Msg)
	u.MsgHdr.Opcode = OpcodeUpdate
	u.Question = make([]Question, 1)
	u.Question[0] = Question{zone, TypeSOA, class}
	return u
}

// The table from RFC 2136 supplemented with the Go DNS function.
//
// 3.2.4 - Table Of Metavalues Used In Prerequisite Section
//
//   CLASS    TYPE     RDATA    Meaning                           Function
//   ----------------------------------------------------------------------
//   ANY      ANY      empty    Name is in use                    NameUsed
//   ANY      rrset    empty    RRset exists (value independent)  RRsetUsedNoRdata
//   NONE     ANY      empty    Name is not in use                NameNotUsed
//   NONE     rrset    empty    RRset does not exist              RRsetNotUsed
//   zone     rrset    rr       RRset exists (value dependent)    RRsetUsedRdata

// NameUsed sets the RRs in the prereq section to
// "Name is in use" RRs. RFC 2136 section 2.4.4.
func (u *Msg) NameUsed(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = &RR_ANY{Hdr: RR_Header{Name: r.Header().Name, Ttl: 0, Rrtype: TypeANY, Class: ClassANY}}
	}
}

// NameNotUsed sets the RRs in the prereq section to
// "Name is in not use" RRs. RFC 2136 section 2.4.5.
func (u *Msg) NameNotUsed(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = &RR_ANY{Hdr: RR_Header{Name: r.Header().Name, Ttl: 0, Rrtype: TypeANY, Class: ClassNONE}}
	}
}

// RRsetUsedRdata sets the RRs in the prereq section to
// "RRset exists (value dependent -- with rdata)" RRs. RFC 2136 section 2.4.2.
func (u *Msg) RRsetUsedRdata(rr []RR) {
	if len(u.Question) == 0 {
		panic("empty question section")
	}
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = r
		u.Answer[i].Header().Class = u.Question[0].Qclass
	}
}

// RRsetUsedNoRdata sets the RRs in the prereq section to
// "RRset exists (value independent -- no rdata)" RRs. RFC 2136 section 2.4.1.
func (u *Msg) RRsetUsedNoRdata(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = r
		u.Answer[i].Header().Class = ClassANY
		u.Answer[i].Header().Ttl = 0
		u.Answer[i].Header().Rdlength = 0
	}
}

// RRsetNotUsed sets the RRs in the prereq section to
// "RRset does not exist" RRs. RFC 2136 section 2.4.3.
func (u *Msg) RRsetNotUsed(rr []RR) {
	u.Answer = make([]RR, len(rr))
	for i, r := range rr {
		u.Answer[i] = r
		u.Answer[i].Header().Class = ClassNONE
		u.Answer[i].Header().Rdlength = 0
		u.Answer[i].Header().Ttl = 0
	}
}

// The table from RFC 2136 supplemented with the Go DNS function.
//
// 3.4.2.6 - Table Of Metavalues Used In Update Section
//
//   CLASS    TYPE     RDATA    Meaning                         Function
//   --------------------------------------------------------------------------
//   ANY      ANY      empty    Delete all RRsets from a name   NameDelete
//   ANY      rrset    empty    Delete an RRset                 RRsetDelete
//   NONE     rrset    rr       Delete an RR from an RRset      RRsetDeleteRR
//   zone     rrset    rr       Add to an RRset                 RRsetAddRdata

// RRsetAddRdata adds an complete RRset, see RFC 2136 section 2.5.1
func (u *Msg) RRsetAddRdata(rr []RR) {
	if len(u.Question) == 0 {
		panic("empty question section")
	}
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = r
		u.Ns[i].Header().Class = u.Question[0].Qclass
	}
}

// RRsetDelete deletes an RRset, see RFC 2136 section 2.5.2
func (u *Msg) RRsetDelete(rr []RR) {
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = r
		u.Ns[i].Header().Class = ClassANY
		u.Ns[i].Header().Rdlength = 0
		u.Ns[i].Header().Ttl = 0
	}
}

// NameDelete deletes all RRsets of a name, see RFC 2136 section 2.5.3
func (u *Msg) NameDelete(rr []RR) {
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = &RR_ANY{Hdr: RR_Header{Name: r.Header().Name, Ttl: 0, Rrtype: TypeANY, Class: ClassANY}}
	}
}

// RRsetDeleteRR deletes RR from the RRSset, see RFC 2136 section 2.5.4
func (u *Msg) RRsetDeleteRR(rr []RR) {
	u.Ns = make([]RR, len(rr))
	for i, r := range rr {
		u.Ns[i] = r
		u.Ns[i].Header().Class = ClassNONE
		u.Ns[i].Header().Ttl = 0
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"github.com/ansel1/merry"
	"github.com/dghubble/sling"
	"github.com/dickeyxxx/golock"
)

func init() {
	Topics = append(Topics, &Topic{
		Name:        "update",
		Description: "update heroku-cli",
		Commands: CommandSet{
			{
				Topic:            "update",
				Hidden:           true,
				Description:      "updates the Heroku CLI",
				DisableAnalytics: true,
				Args:             []Arg{{Name: "channel", Optional: true}},
				Run: func(ctx *Context) {
					channel := ctx.Args.(map[string]string)["channel"]
					if channel == "" {
						channel = Channel
					}
					Update(channel)
				},
			},
		},
	})
}

// Autoupdate is a flag to enable/disable CLI autoupdating
var Autoupdate = "no"

var updateLockPath = filepath.Join(CacheHome, "updating.lock")
var autoupdateFile = filepath.Join(CacheHome, "autoupdate")

// Update updates the CLI and plugins
func Update(channel string) {
	touchAutoupdateFile()
	SubmitAnalytics()
	updateCLI(channel)
	UserPlugins.Update()
	truncate(ErrLogPath, 1000)
	cleanTmpDirs()
}

func updateCLI(channel string) {
	if Autoupdate != "yes" {
		return
	}
	manifest := GetUpdateManifest(channel)
	binExists, _ := FileExists(expectedBinPath())
	if binExists && manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	DownloadCLI(channel, filepath.Join(DataHome, "cli"), manifest)
	loadNewCLI()
}

// DownloadCLI downloads a CLI update to a given path
func DownloadCLI(channel, path string, manifest *Manifest) {
	locked, err := golock.IsLocked(updateLockPath)
	LogIfError(err)
	if locked {
		must(merry.Errorf("Update in progress"))
	}
	LogIfError(golock.Lock(updateLockPath))
	unlock := func() {
		golock.Unlock(updateLockPath)
	}
	defer unlock()
	hideCursor()
	downloadingMessage := fmt.Sprintf("heroku-cli: Updating to %s...", manifest.Version)
	if manifest.Channel != "stable" {
		downloadingMessage = fmt.Sprintf("%s (%s)", downloadingMessage, manifest.Channel)
	}
	Logln(downloadingMessage)
	build := manifest.Builds[runtime.GOOS+"-"+runtime.GOARCH]
	if build == nil {
		must(merry.Errorf("no build for %s", manifest.Channel))
	}
	reader, getSha, err := downloadXZ(build.URL, downloadingMessage)
	must(err)
	tmp := tmpDir(DataHome)
	must(extractTar(reader, tmp))
	sha := getSha()
	if sha != build.Sha256 {
		must(merry.Errorf("SHA mismatch: expected %s to be %s", sha, build.Sha256))
	}
	exists, _ := FileExists(path)
	if exists {
		must(os.Rename(path, filepath.Join(tmpDir(DataHome), "heroku")))
	}
	Errln(path)
	must(os.Rename(filepath.Join(tmp, "heroku"), path))
	Debugf("updated to %s\n", manifest.Version)
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded() bool {
	if exists, _ := FileExists(expectedBinPath()); !exists {
		return true
	}
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		if os.IsNotExist(err) {
			return true
		}
		LogIfError(err)
		return true
	}
	return time.Since(f.ModTime()) > 4*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	must(err)
	_, err = out.WriteString(time.Now().String())
	must(err)
	err = out.Close()
	must(err)
}

// Manifest is the manifest.json for releases
type Manifest struct {
	ReleasedAt string            `json:"released_at"`
	Version    string            `json:"version"`
	Channel    string            `json:"channel"`
	Builds     map[string]*Build `json:"builds"`
}

// Build is a part of a Manifest
type Build struct {
	URL    string `json:"url"`
	Sha256 string `json:"sha256"`
	Bytes  int64  `json:"bytes"`
}

var updateManifestRetrying = false

// GetUpdateManifest loads the manifest.json for a channel
func GetUpdateManifest(channel string) *Manifest {
	var m Manifest
	url := "https://cli-assets.heroku.com/branches/" + channel + "/manifest.json"
	rsp, err := sling.New().Get(url).ReceiveSuccess(&m)
	if err != nil {
		if updateManifestRetrying {
			must(err)
		} else {
			updateManifestRetrying = true
			return GetUpdateManifest(channel)
		}
	}
	must(getHTTPError(rsp))
	return &m
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if IsUpdateNeeded() {
		Debugln("triggering background update")
		touchAutoupdateFile()
		exec.Command(BinPath, "update").Start()
	}
}

func cleanTmpDirs() {
	clean := func(base string) {
		dir := filepath.Join(base, "tmp")
		if exists, _ := FileExists(dir); !exists {
			return
		}
		files, err := ioutil.ReadDir(dir)
		LogIfError(err)
		for _, file := range files {
			if time.Since(file.ModTime()) > 24*time.Hour {
				path := filepath.Join(dir, file.Name())
				Debugf("removing old tmp dir %s", path)
				LogIfError(os.RemoveAll(path))
			}
		}
	}
	clean(DataHome)
	clean(CacheHome)
}

func expectedBinPath() string {
	bin := filepath.Join(DataHome, "cli", "bin", "heroku")
	if runtime.GOOS == WINDOWS {
		bin = bin + ".exe"
	}
	return bin
}

func loadNewCLI() {
	if Autoupdate == "no" {
		return
	}
	expected := expectedBinPath()
	if BinPath == expected {
		return
	}
	if exists, _ := FileExists(expected); !exists {
		if exists, _ = FileExists(npmBinPath()); !exists {
			// uh oh, npm isn't where it should be.
			// The CLI probably isn't installed right so force an update
			Update(Channel)
		}
		return
	}
	execBin(expected, Args...)
}
<commit_msg>small cleanup<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"github.com/ansel1/merry"
	"github.com/dghubble/sling"
	"github.com/dickeyxxx/golock"
)

func init() {
	Topics = append(Topics, &Topic{
		Name:        "update",
		Description: "update heroku-cli",
		Commands: CommandSet{
			{
				Topic:            "update",
				Hidden:           true,
				Description:      "updates the Heroku CLI",
				DisableAnalytics: true,
				Args:             []Arg{{Name: "channel", Optional: true}},
				Run: func(ctx *Context) {
					channel := ctx.Args.(map[string]string)["channel"]
					if channel == "" {
						channel = Channel
					}
					Update(channel)
				},
			},
		},
	})
}

// Autoupdate is a flag to enable/disable CLI autoupdating
var Autoupdate = "no"

var updateLockPath = filepath.Join(CacheHome, "updating.lock")
var autoupdateFile = filepath.Join(CacheHome, "autoupdate")

// Update updates the CLI and plugins
func Update(channel string) {
	touchAutoupdateFile()
	SubmitAnalytics()
	updateCLI(channel)
	UserPlugins.Update()
	truncate(ErrLogPath, 1000)
	cleanTmpDirs()
}

func updateCLI(channel string) {
	if Autoupdate != "yes" {
		return
	}
	manifest := GetUpdateManifest(channel)
	binExists, _ := FileExists(expectedBinPath())
	if binExists && manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	DownloadCLI(channel, filepath.Join(DataHome, "cli"), manifest)
	loadNewCLI()
}

// DownloadCLI downloads a CLI update to a given path
func DownloadCLI(channel, path string, manifest *Manifest) {
	locked, err := golock.IsLocked(updateLockPath)
	LogIfError(err)
	if locked {
		must(merry.Errorf("Update in progress"))
	}
	LogIfError(golock.Lock(updateLockPath))
	unlock := func() {
		golock.Unlock(updateLockPath)
	}
	defer unlock()
	hideCursor()
	downloadingMessage := fmt.Sprintf("heroku-cli: Updating to %s...", manifest.Version)
	if manifest.Channel != "stable" {
		downloadingMessage = fmt.Sprintf("%s (%s)", downloadingMessage, manifest.Channel)
	}
	Logln(downloadingMessage)
	build := manifest.Builds[runtime.GOOS+"-"+runtime.GOARCH]
	if build == nil {
		must(merry.Errorf("no build for %s", manifest.Channel))
	}
	reader, getSha, err := downloadXZ(build.URL, downloadingMessage)
	must(err)
	tmp := tmpDir(DataHome)
	must(extractTar(reader, tmp))
	sha := getSha()
	if sha != build.Sha256 {
		must(merry.Errorf("SHA mismatch: expected %s to be %s", sha, build.Sha256))
	}
	exists, _ := FileExists(path)
	if exists {
		must(os.Rename(path, filepath.Join(tmpDir(DataHome), "heroku")))
	}
	Errln(path)
	must(os.Rename(filepath.Join(tmp, "heroku"), path))
	Debugf("updated to %s\n", manifest.Version)
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded() bool {
	if exists, _ := FileExists(expectedBinPath()); !exists {
		return true
	}
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		if os.IsNotExist(err) {
			return true
		}
		LogIfError(err)
		return true
	}
	return time.Since(f.ModTime()) > 4*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	must(err)
	_, err = out.WriteString(time.Now().String())
	must(err)
	err = out.Close()
	must(err)
}

// Manifest is the manifest.json for releases
type Manifest struct {
	ReleasedAt string            `json:"released_at"`
	Version    string            `json:"version"`
	Channel    string            `json:"channel"`
	Builds     map[string]*Build `json:"builds"`
}

// Build is a part of a Manifest
type Build struct {
	URL    string `json:"url"`
	Sha256 string `json:"sha256"`
	Bytes  int64  `json:"bytes"`
}

var updateManifestRetrying = false

// GetUpdateManifest loads the manifest.json for a channel
func GetUpdateManifest(channel string) *Manifest {
	var m Manifest
	url := "https://cli-assets.heroku.com/branches/" + channel + "/manifest.json"
	rsp, err := sling.New().Get(url).ReceiveSuccess(&m)
	if err != nil && !updateManifestRetrying {
		updateManifestRetrying = true
		return GetUpdateManifest(channel)
	}
	must(err)
	must(getHTTPError(rsp))
	return &m
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if IsUpdateNeeded() {
		Debugln("triggering background update")
		touchAutoupdateFile()
		exec.Command(BinPath, "update").Start()
	}
}

func cleanTmpDirs() {
	clean := func(base string) {
		dir := filepath.Join(base, "tmp")
		if exists, _ := FileExists(dir); !exists {
			return
		}
		files, err := ioutil.ReadDir(dir)
		LogIfError(err)
		for _, file := range files {
			if time.Since(file.ModTime()) > 24*time.Hour {
				path := filepath.Join(dir, file.Name())
				Debugf("removing old tmp dir %s", path)
				LogIfError(os.RemoveAll(path))
			}
		}
	}
	clean(DataHome)
	clean(CacheHome)
}

func expectedBinPath() string {
	bin := filepath.Join(DataHome, "cli", "bin", "heroku")
	if runtime.GOOS == WINDOWS {
		bin = bin + ".exe"
	}
	return bin
}

func loadNewCLI() {
	if Autoupdate == "no" {
		return
	}
	expected := expectedBinPath()
	if BinPath == expected {
		return
	}
	if exists, _ := FileExists(expected); !exists {
		if exists, _ = FileExists(npmBinPath()); !exists {
			// uh oh, npm isn't where it should be.
			// The CLI probably isn't installed right so force an update
			Update(Channel)
		}
		return
	}
	execBin(expected, Args...)
}
<|endoftext|>


<commit_before>package standard

import (
	"net/http"
	"sync"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine"
	"github.com/labstack/echo/log"
	glog "github.com/labstack/gommon/log"
)

type (
	// Server implements `engine.Server`.
	Server struct {
		*http.Server
		config  engine.Config
		handler engine.Handler
		logger  log.Logger
		pool    *pool
	}

	pool struct {
		request         sync.Pool
		response        sync.Pool
		responseAdapter sync.Pool
		header          sync.Pool
		url             sync.Pool
	}
)

// New returns `Server` instance with provided listen address.
func New(addr string) *Server {
	c := engine.Config{Address: addr}
	return WithConfig(c)
}

// WithTLS returns `Server` instance with provided TLS config.
func WithTLS(addr, certFile, keyFile string) *Server {
	c := engine.Config{
		Address:     addr,
		TLSCertFile: certFile,
		TLSKeyFile:  keyFile,
	}
	return WithConfig(c)
}

// WithConfig returns `Server` instance with provided config.
func WithConfig(c engine.Config) (s *Server) {
	s = &Server{
		Server: new(http.Server),
		config: c,
		pool: &pool{
			request: sync.Pool{
				New: func() interface{} {
					return &Request{logger: s.logger}
				},
			},
			response: sync.Pool{
				New: func() interface{} {
					return &Response{logger: s.logger}
				},
			},
			responseAdapter: sync.Pool{
				New: func() interface{} {
					return &responseAdapter{}
				},
			},
			header: sync.Pool{
				New: func() interface{} {
					return &Header{}
				},
			},
			url: sync.Pool{
				New: func() interface{} {
					return &URL{}
				},
			},
		},
		handler: engine.HandlerFunc(func(req engine.Request, res engine.Response) {
			panic("echo: handler not set, use `Server#SetHandler()` to set it.")
		}),
		logger: glog.New("echo"),
	}
	s.ReadTimeout = c.ReadTimeout
	s.WriteTimeout = c.WriteTimeout
	s.Addr = c.Address
	s.Handler = s
	return
}

// SetHandler implements `engine.Server#SetHandler` function.
func (s *Server) SetHandler(h engine.Handler) {
	s.handler = h
}

// SetLogger implements `engine.Server#SetLogger` function.
func (s *Server) SetLogger(l log.Logger) {
	s.logger = l
}

// Start implements `engine.Server#Start` function.
func (s *Server) Start() error {
	if s.config.Listener == nil {
		return s.startDefaultListener()
	}
	return s.startCustomListener()
}

func (s *Server) startDefaultListener() error {
	c := s.config
	if c.TLSCertFile != "" && c.TLSKeyFile != "" {
		return s.ListenAndServeTLS(c.TLSCertFile, c.TLSKeyFile)
	}
	return s.ListenAndServe()
}

func (s *Server) startCustomListener() error {
	return s.Serve(s.config.Listener)
}

// ServeHTTP implements `http.Handler` interface.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Request
	req := s.pool.request.Get().(*Request)
	reqHdr := s.pool.header.Get().(*Header)
	reqURL := s.pool.url.Get().(*URL)
	reqHdr.reset(r.Header)
	reqURL.reset(r.URL)
	req.reset(r, reqHdr, reqURL)

	// Response
	res := s.pool.response.Get().(*Response)
	resAdpt := s.pool.responseAdapter.Get().(*responseAdapter)
	resAdpt.reset(res)
	resHdr := s.pool.header.Get().(*Header)
	resHdr.reset(w.Header())
	res.reset(w, resAdpt, resHdr)

	s.handler.ServeHTTP(req, res)

	// Return to pool
	s.pool.request.Put(req)
	s.pool.header.Put(reqHdr)
	s.pool.url.Put(reqURL)
	s.pool.response.Put(res)
	s.pool.header.Put(resHdr)
}

// WrapHandler wraps `http.Handler` into `echo.HandlerFunc`.
func WrapHandler(h http.Handler) echo.HandlerFunc {
	return func(c echo.Context) error {
		req := c.Request().(*Request)
		res := c.Response().(*Response)
		h.ServeHTTP(res.adapter, req.Request)
		return nil
	}
}

// WrapMiddleware wraps `func(http.Handler) http.Handler` into `echo.MiddlewareFunc`
func WrapMiddleware(m func(http.Handler) http.Handler) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) (err error) {
			req := c.Request().(*Request)
			res := c.Response().(*Response)
			m(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				err = next(c)
			})).ServeHTTP(res.ResponseWriter, req.Request)
			return
		}
	}
}
<commit_msg>Use response adapter in engine/standard<commit_after>package standard

import (
	"net/http"
	"sync"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine"
	"github.com/labstack/echo/log"
	glog "github.com/labstack/gommon/log"
)

type (
	// Server implements `engine.Server`.
	Server struct {
		*http.Server
		config  engine.Config
		handler engine.Handler
		logger  log.Logger
		pool    *pool
	}

	pool struct {
		request         sync.Pool
		response        sync.Pool
		responseAdapter sync.Pool
		header          sync.Pool
		url             sync.Pool
	}
)

// New returns `Server` instance with provided listen address.
func New(addr string) *Server {
	c := engine.Config{Address: addr}
	return WithConfig(c)
}

// WithTLS returns `Server` instance with provided TLS config.
func WithTLS(addr, certFile, keyFile string) *Server {
	c := engine.Config{
		Address:     addr,
		TLSCertFile: certFile,
		TLSKeyFile:  keyFile,
	}
	return WithConfig(c)
}

// WithConfig returns `Server` instance with provided config.
func WithConfig(c engine.Config) (s *Server) {
	s = &Server{
		Server: new(http.Server),
		config: c,
		pool: &pool{
			request: sync.Pool{
				New: func() interface{} {
					return &Request{logger: s.logger}
				},
			},
			response: sync.Pool{
				New: func() interface{} {
					return &Response{logger: s.logger}
				},
			},
			responseAdapter: sync.Pool{
				New: func() interface{} {
					return &responseAdapter{}
				},
			},
			header: sync.Pool{
				New: func() interface{} {
					return &Header{}
				},
			},
			url: sync.Pool{
				New: func() interface{} {
					return &URL{}
				},
			},
		},
		handler: engine.HandlerFunc(func(req engine.Request, res engine.Response) {
			panic("echo: handler not set, use `Server#SetHandler()` to set it.")
		}),
		logger: glog.New("echo"),
	}
	s.ReadTimeout = c.ReadTimeout
	s.WriteTimeout = c.WriteTimeout
	s.Addr = c.Address
	s.Handler = s
	return
}

// SetHandler implements `engine.Server#SetHandler` function.
func (s *Server) SetHandler(h engine.Handler) {
	s.handler = h
}

// SetLogger implements `engine.Server#SetLogger` function.
func (s *Server) SetLogger(l log.Logger) {
	s.logger = l
}

// Start implements `engine.Server#Start` function.
func (s *Server) Start() error {
	if s.config.Listener == nil {
		return s.startDefaultListener()
	}
	return s.startCustomListener()
}

func (s *Server) startDefaultListener() error {
	c := s.config
	if c.TLSCertFile != "" && c.TLSKeyFile != "" {
		return s.ListenAndServeTLS(c.TLSCertFile, c.TLSKeyFile)
	}
	return s.ListenAndServe()
}

func (s *Server) startCustomListener() error {
	return s.Serve(s.config.Listener)
}

// ServeHTTP implements `http.Handler` interface.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Request
	req := s.pool.request.Get().(*Request)
	reqHdr := s.pool.header.Get().(*Header)
	reqURL := s.pool.url.Get().(*URL)
	reqHdr.reset(r.Header)
	reqURL.reset(r.URL)
	req.reset(r, reqHdr, reqURL)

	// Response
	res := s.pool.response.Get().(*Response)
	resAdpt := s.pool.responseAdapter.Get().(*responseAdapter)
	resAdpt.reset(res)
	resHdr := s.pool.header.Get().(*Header)
	resHdr.reset(w.Header())
	res.reset(w, resAdpt, resHdr)

	s.handler.ServeHTTP(req, res)

	// Return to pool
	s.pool.request.Put(req)
	s.pool.header.Put(reqHdr)
	s.pool.url.Put(reqURL)
	s.pool.response.Put(res)
	s.pool.header.Put(resHdr)
}

// WrapHandler wraps `http.Handler` into `echo.HandlerFunc`.
func WrapHandler(h http.Handler) echo.HandlerFunc {
	return func(c echo.Context) error {
		req := c.Request().(*Request)
		res := c.Response().(*Response)
		h.ServeHTTP(res.adapter, req.Request)
		return nil
	}
}

// WrapMiddleware wraps `func(http.Handler) http.Handler` into `echo.MiddlewareFunc`
func WrapMiddleware(m func(http.Handler) http.Handler) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) (err error) {
			req := c.Request().(*Request)
			res := c.Response().(*Response)
			m(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				err = next(c)
			})).ServeHTTP(res.adapter, req.Request)
			return
		}
	}
}
<|endoftext|>


<commit_before>package nativeengine

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

	"github.com/Sirupsen/logrus"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/engines/native/system"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
	"github.com/taskcluster/taskcluster-worker/runtime/ioext"
	"github.com/taskcluster/taskcluster-worker/runtime/util"
)

type sandbox struct {
	engines.SandboxBase
	engine        *engine
	context       *runtime.TaskContext
	log           *logrus.Entry
	workingFolder runtime.TemporaryFolder
	user          *system.User
	process       *system.Process
	env           map[string]string
	resolve       atomics.Once // Guarding resultSet, resultErr and abortErr
	resultSet     *resultSet
	resultErr     error
	abortErr      error
	wg            atomics.WaitGroup
}

func newSandbox(b *sandboxBuilder) (*sandbox, error) {
	// Create temporary home folder for the task
	workingFolder, err := b.engine.environment.TemporaryStorage.NewFolder()
	if err != nil {
		b.log.Error("Failed to create temporary folder: ", err)
		return nil, fmt.Errorf("Failed to temporary folder, error: %s", err)
	}

	if b.payload.Context != "" {
		if err = fetchContext(b.payload.Context, workingFolder.Path()); err != nil {
			b.context.LogError(err)
			return nil, engines.NewMalformedPayloadError(
				fmt.Sprintf("Error downloading %s: %v", b.payload.Context, err),
			)
		}
	}

	var user *system.User

	if b.engine.config.CreateUser {
		// Create temporary user account
		user, err = system.CreateUser(workingFolder.Path(), b.engine.groups)
		if err != nil {
			workingFolder.Remove() // best-effort clean-up this is a fatal error
			return nil, fmt.Errorf("Failed to create temporary system user, error: %s", err)
		}
	} else {
		user, err = system.CurrentUser()
		if err != nil {
			workingFolder.Remove()
			return nil, err
		}
	}

	env := map[string]string{}
	for k, v := range b.env {
		env[k] = v
	}

	env["HOME"] = workingFolder.Path()
	env["USER"] = user.Name()
	env["LOGNAME"] = user.Name()

	// Start process
	debug("StartProcess: %v", b.payload.Command)
	process, err := system.StartProcess(system.ProcessOptions{
		Arguments:     b.payload.Command,
		Environment:   env,
		WorkingFolder: workingFolder.Path(),
		Owner:         user,
		Stdout:        ioext.WriteNopCloser(b.context.LogDrain()),
		// Stderr defaults to Stdout when not specified
	})
	if err != nil {
		// StartProcess provides human-readable error messages (see docs)
		// We'll convert it to a MalformedPayloadError
		return nil, engines.NewMalformedPayloadError(
			"Unable to start specified command: ", b.payload.Command, "error: ", err,
		)
	}

	s := &sandbox{
		engine:        b.engine,
		context:       b.context,
		log:           b.log,
		workingFolder: workingFolder,
		user:          user,
		process:       process,
		env:           b.env,
	}

	go s.waitForTermination()

	return s, nil
}

func fetchContext(context, destdir string) error {
	// TODO: use future cache subsystem, when we have it
	// TODO: use the soon to be merged fetcher subsystem
	filename, err := util.Download(context, destdir)
	if err != nil {
		return fmt.Errorf("Error downloading '%s': %v", context, err)
	}

	// TODO: verify if this will harm Windows
	// TODO: abstract this away in system package
	if err = os.Chmod(filename, 0700); err != nil {
		return fmt.Errorf("Error setting file '%s' permissions: %v", filename, err)
	}

	unpackedFile := ""
	switch filepath.Ext(filename) {
	case ".zip":
		err = runtime.Unzip(filename)
	case ".gz":
		unpackedFile, err = runtime.Gunzip(filename)
	}

	if err != nil {
		return fmt.Errorf("Error unpacking '%s': %v", context, err)
	}

	if filepath.Ext(unpackedFile) == ".tar" {
		err = runtime.Untar(unpackedFile)
		if err != nil {
			return fmt.Errorf("Error unpacking '%s': %v", context, err)
		}
	}

	return nil
}

func (s *sandbox) NewShell(command []string, tty bool) (engines.Shell, error) {
	// Increment shell counter, if draining we don't allow new shells
	if s.wg.Add(1) != nil {
		return nil, engines.ErrSandboxTerminated
	}

	debug("NewShell with: %v", command)
	shell, err := newShell(s, command, tty)
	if err != nil {
		debug("Failed to start shell, error: %s", err)
		s.wg.Done()
		return nil, engines.NewMalformedPayloadError(
			"Unable to spawn command: ", command, " error: ", err,
		)
	}

	// Wait for the shell to be done and decrement WaitGroup
	go func() {
		result, _ := shell.Wait()
		debug("Shell finished with: %v", result)
		s.wg.Done()
	}()

	return shell, nil
}

func (s *sandbox) waitForTermination() {
	// Wait for process to terminate
	success := s.process.Wait()
	debug("Process finished with: %v", success)

	// Wait for all shell to finish and prevent new shells from being created
	s.wg.WaitAndDrain()
	debug("All shells terminated")

	s.resolve.Do(func() {
		// Halt all other sub-processes
		if s.engine.config.CreateUser {
			system.KillByOwner(s.user)
		}

		// Create resultSet
		s.resultSet = &resultSet{
			engine:        s.engine,
			context:       s.context,
			log:           s.log,
			workingFolder: s.workingFolder,
			user:          s.user,
			success:       success,
		}
		s.abortErr = engines.ErrSandboxTerminated
	})
}

func (s *sandbox) WaitForResult() (engines.ResultSet, error) {
	// Wait for result and terminate
	s.resolve.Wait()
	return s.resultSet, s.resultErr
}

func (s *sandbox) Abort() error {
	s.resolve.Do(func() {
		debug("Aborting sandbox")

		// In case we didn't create a new user, killing
		// the children processes is the only safe way
		// to kill processes created by the task.
		system.KillProcessTree(s.process)

		if s.engine.config.CreateUser {
			// When we have a new user created, we can safely
			// kill any process owned by it.
			err := system.KillByOwner(s.user)
			if err != nil {
				s.log.Error("Failed to kill all processes by owner, error: ", err)
			}

			// Remove temporary user (this will panic if unsuccessful)
			s.user.Remove()
		}

		// Remove temporary home folder
		err := s.workingFolder.Remove()
		if err != nil {
			s.log.Error("Failed to remove temporary home directory, error: ", err)
		}

		// Set result
		s.resultErr = engines.ErrSandboxAborted
	})
	s.resolve.Wait()
	return s.abortErr
}
<commit_msg>Fix resource leak upwon process startup failure. r=jonasfj<commit_after>package nativeengine

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

	"github.com/Sirupsen/logrus"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/engines/native/system"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
	"github.com/taskcluster/taskcluster-worker/runtime/ioext"
	"github.com/taskcluster/taskcluster-worker/runtime/util"
)

type sandbox struct {
	engines.SandboxBase
	engine        *engine
	context       *runtime.TaskContext
	log           *logrus.Entry
	workingFolder runtime.TemporaryFolder
	user          *system.User
	process       *system.Process
	env           map[string]string
	resolve       atomics.Once // Guarding resultSet, resultErr and abortErr
	resultSet     *resultSet
	resultErr     error
	abortErr      error
	wg            atomics.WaitGroup
}

func newSandbox(b *sandboxBuilder) (*sandbox, error) {
	// Create temporary home folder for the task
	workingFolder, err := b.engine.environment.TemporaryStorage.NewFolder()
	if err != nil {
		b.log.Error("Failed to create temporary folder: ", err)
		return nil, fmt.Errorf("Failed to temporary folder, error: %s", err)
	}

	if b.payload.Context != "" {
		if err = fetchContext(b.payload.Context, workingFolder.Path()); err != nil {
			b.context.LogError(err)
			return nil, engines.NewMalformedPayloadError(
				fmt.Sprintf("Error downloading %s: %v", b.payload.Context, err),
			)
		}
	}

	var user *system.User

	if b.engine.config.CreateUser {
		// Create temporary user account
		user, err = system.CreateUser(workingFolder.Path(), b.engine.groups)
		if err != nil {
			workingFolder.Remove() // best-effort clean-up this is a fatal error
			return nil, fmt.Errorf("Failed to create temporary system user, error: %s", err)
		}
	} else {
		user, err = system.CurrentUser()
		if err != nil {
			workingFolder.Remove()
			return nil, err
		}
	}

	env := map[string]string{}
	for k, v := range b.env {
		env[k] = v
	}

	env["HOME"] = workingFolder.Path()
	env["USER"] = user.Name()
	env["LOGNAME"] = user.Name()

	// Start process
	debug("StartProcess: %v", b.payload.Command)
	process, err := system.StartProcess(system.ProcessOptions{
		Arguments:     b.payload.Command,
		Environment:   env,
		WorkingFolder: workingFolder.Path(),
		Owner:         user,
		Stdout:        ioext.WriteNopCloser(b.context.LogDrain()),
		// Stderr defaults to Stdout when not specified
	})
	if err != nil {
		if b.engine.config.CreateUser {
			user.Remove()
		}

		workingFolder.Remove()

		// StartProcess provides human-readable error messages (see docs)
		// We'll convert it to a MalformedPayloadError
		return nil, engines.NewMalformedPayloadError(
			"Unable to start specified command: ", b.payload.Command, "error: ", err,
		)
	}

	s := &sandbox{
		engine:        b.engine,
		context:       b.context,
		log:           b.log,
		workingFolder: workingFolder,
		user:          user,
		process:       process,
		env:           b.env,
	}

	go s.waitForTermination()

	return s, nil
}

func fetchContext(context, destdir string) error {
	// TODO: use future cache subsystem, when we have it
	// TODO: use the soon to be merged fetcher subsystem
	filename, err := util.Download(context, destdir)
	if err != nil {
		return fmt.Errorf("Error downloading '%s': %v", context, err)
	}

	// TODO: verify if this will harm Windows
	// TODO: abstract this away in system package
	if err = os.Chmod(filename, 0700); err != nil {
		return fmt.Errorf("Error setting file '%s' permissions: %v", filename, err)
	}

	unpackedFile := ""
	switch filepath.Ext(filename) {
	case ".zip":
		err = runtime.Unzip(filename)
	case ".gz":
		unpackedFile, err = runtime.Gunzip(filename)
	}

	if err != nil {
		return fmt.Errorf("Error unpacking '%s': %v", context, err)
	}

	if filepath.Ext(unpackedFile) == ".tar" {
		err = runtime.Untar(unpackedFile)
		if err != nil {
			return fmt.Errorf("Error unpacking '%s': %v", context, err)
		}
	}

	return nil
}

func (s *sandbox) NewShell(command []string, tty bool) (engines.Shell, error) {
	// Increment shell counter, if draining we don't allow new shells
	if s.wg.Add(1) != nil {
		return nil, engines.ErrSandboxTerminated
	}

	debug("NewShell with: %v", command)
	shell, err := newShell(s, command, tty)
	if err != nil {
		debug("Failed to start shell, error: %s", err)
		s.wg.Done()
		return nil, engines.NewMalformedPayloadError(
			"Unable to spawn command: ", command, " error: ", err,
		)
	}

	// Wait for the shell to be done and decrement WaitGroup
	go func() {
		result, _ := shell.Wait()
		debug("Shell finished with: %v", result)
		s.wg.Done()
	}()

	return shell, nil
}

func (s *sandbox) waitForTermination() {
	// Wait for process to terminate
	success := s.process.Wait()
	debug("Process finished with: %v", success)

	// Wait for all shell to finish and prevent new shells from being created
	s.wg.WaitAndDrain()
	debug("All shells terminated")

	s.resolve.Do(func() {
		// Halt all other sub-processes
		if s.engine.config.CreateUser {
			system.KillByOwner(s.user)
		}

		// Create resultSet
		s.resultSet = &resultSet{
			engine:        s.engine,
			context:       s.context,
			log:           s.log,
			workingFolder: s.workingFolder,
			user:          s.user,
			success:       success,
		}
		s.abortErr = engines.ErrSandboxTerminated
	})
}

func (s *sandbox) WaitForResult() (engines.ResultSet, error) {
	// Wait for result and terminate
	s.resolve.Wait()
	return s.resultSet, s.resultErr
}

func (s *sandbox) Abort() error {
	s.resolve.Do(func() {
		debug("Aborting sandbox")

		// In case we didn't create a new user, killing
		// the children processes is the only safe way
		// to kill processes created by the task.
		system.KillProcessTree(s.process)

		if s.engine.config.CreateUser {
			// When we have a new user created, we can safely
			// kill any process owned by it.
			err := system.KillByOwner(s.user)
			if err != nil {
				s.log.Error("Failed to kill all processes by owner, error: ", err)
			}

			// Remove temporary user (this will panic if unsuccessful)
			s.user.Remove()
		}

		// Remove temporary home folder
		err := s.workingFolder.Remove()
		if err != nil {
			s.log.Error("Failed to remove temporary home directory, error: ", err)
		}

		// Set result
		s.resultErr = engines.ErrSandboxAborted
	})
	s.resolve.Wait()
	return s.abortErr
}
<|endoftext|>