<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * promise/promise.go                                     *
 *                                                        *
 * promise interface for Go.                              *
 *                                                        *
 * LastModified: Aug 13, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package promise

import (
	"reflect"
	"time"
)

// OnFulfilled is a function called when the Promise is fulfilled.
// This function has one argument, the fulfillment value.
//
// The function type can be the following:
//     func() (interface{}, error)
//     func()
//     func(interface{}) (interface{}, error)
//     func(interface{})
type OnFulfilled interface{}

// OnRejected is a function called when the Promise is rejected.
// This function has one argument, the rejection reason.
//
// The function type can be the following:
//     func() (interface{}, error)
//     func()
//     func(interface{}) (interface{}, error)
//     func(interface{})
//     func(error) (interface{}, error)
//     func(error)
type OnRejected interface{}

// OnCompleted is a function called when the Promise is completed.
// This function has one argument,
// the fulfillment value when the Promise is fulfilled,
// or the rejection reason when the Promise is rejected.
// The function type can be the following:
//     func() (interface{}, error)
//     func()
//     func(interface{}) (interface{}, error)
//     func(interface{})
type OnCompleted interface{}

// Promise is an interface of the JS Promise/A+ spec
// (https://promisesaplus.com/).
type Promise interface {
	// Then method returns a Promise. It takes two arguments: callback functions
	// for the success and failure cases of the Promise.
	Then(onFulfilled OnFulfilled, onRejected ...OnRejected) Promise

	// Catch handles errors emitted by this Promise.
	//
	// This is the asynchronous equivalent of a "catch" block.
	//
	// Returns a new Promise that will be completed with either the result of
	// this promise or the result of calling the onRejected callback.
	//
	// If this promise completes with a value, the returned promise completes
	// with the same value.
	//
	// If this promise completes with an error, then test is first called with
	// the error value.
	//
	// If test returns false, the error is not handled by this Catch, and the
	// returned promise completes with the same error and stack trace as this
	// promise.
	//
	// If test returns true, onRejected is called with the error and possibly
	// stack trace, and the returned promise is completed with the result of
	// this call in exactly the same way as for Then's onRejected.
	//
	// If test is omitted, it defaults to a function that always returns true.
	// The test function should not panic, but if it does, it is handled as if
	// the the onRejected function had panic.
	Catch(onRejected OnRejected, test ...func(error) bool) Promise

	// Complete is the same way as Then(onCompleted, onCompleted)
	Complete(onCompleted OnCompleted) Promise

	// WhenComplete register a function to be called when the promise completes.
	//
	// The action function is called when this promise completes, whether it
	// does so with a value or with an error.
	//
	// If this promise completes with a value, the returned promise completes
	// with the same value.
	//
	// If this promise completes with an error, the returned promise completes
	// with the same error.
	//
	// The action function should not panic, but if it does, the returned
	// promise completes with a PanicError.
	WhenComplete(action func()) Promise

	// Done is the same semantics as Then except that it don't return a Promise.
	// If the callback function (onFulfilled or onRejected) returns error or
	// panics, the application will be crashing.
	// The result of the callback function will be ignored.
	Done(onFulfilled OnFulfilled, onRejected ...OnRejected)

	// State return the current state of the Promise
	State() State

	// Resolve method returns a Promise object that is resolved with the given
	// value. If the value is a Promise, the returned promise will "follow" that Promise, adopting its eventual state; otherwise the returned promise
	// will be fulfilled with the value.
	Resolve(value interface{})

	// Reject method returns a Promise object that is rejected with the given
	// reason.
	Reject(reason error)

	// Fill the promise with this promise if the promise is in PENDING state.
	// otherwise nothing to do.
	Fill(promise Promise)

	// Timeout create a new promise that will reject with a TimeoutError or a
	// custom reason after a timeout if promise does not fulfill or reject
	// beforehand.
	Timeout(duration time.Duration, reason ...error) Promise

	// Delay create a new promise that will, after duration delay, fulfill with
	// the same value as this promise. If this promise rejects, delayed promise
	// will be rejected immediately.
	Delay(duration time.Duration) Promise

	// Tap executes a function as a side effect when promise fulfills.
	//
	// It returns a new promise:
	// 1. If promise fulfills, onFulfilledSideEffect is executed:
	//     * If onFulfilledSideEffect returns successfully, the promise
	//       returned by tap fulfills with promise's original fulfillment
	//       value.
	//     * If onFulfilledSideEffect panics, the promise returned by tap
	//       rejects with the panic message as the reason.
	// 2. If promise rejects, onFulfilledSideEffect is not executed, and the
	//    promise returned by tap rejects with promise's rejection reason.
	Tap(onfulfilledSideEffect func(interface{})) Promise

	// Get the value and reason synchronously, if this promise in PENDING state.
	// this method will block the current goroutine.
	Get() (interface{}, error)
}

func catch(promise Promise) {
	if e := recover(); e != nil {
		promise.Reject(NewPanicError(e))
	}
}

func call(promise Promise, computation func() (interface{}, error)) {
	defer catch(promise)
	if result, err := computation(); err != nil {
		promise.Reject(err)
	} else {
		promise.Resolve(result)
	}
}

func call1(promise Promise, computation func()) {
	defer catch(promise)
	computation()
	promise.Resolve(nil)
}

func call2(promise Promise, computation func(interface{}) (interface{}, error), x interface{}) {
	defer catch(promise)
	if result, err := computation(x); err != nil {
		promise.Reject(err)
	} else {
		promise.Resolve(result)
	}
}

func call3(promise Promise, computation func(interface{}), x interface{}) {
	defer catch(promise)
	computation(x)
	promise.Resolve(nil)
}

func call4(promise Promise, computation func(error) (interface{}, error), e error) {
	defer catch(promise)
	if result, err := computation(e); err != nil {
		promise.Reject(err)
	} else {
		promise.Resolve(result)
	}
}

func call5(promise Promise, computation func(error), e error) {
	defer catch(promise)
	computation(e)
	promise.Resolve(nil)
}

func resolve(next Promise, onFulfilled OnFulfilled, x interface{}) {
	switch f := onFulfilled.(type) {
	case nil:
		next.Resolve(x)
	case func() (interface{}, error):
		go call(next, f)
	case func():
		go call1(next, f)
	case func(interface{}) (interface{}, error):
		go call2(next, f, x)
	case func(interface{}):
		go call3(next, f, x)
	default:
		panic("onFulfilled can't support this type: " + reflect.TypeOf(onFulfilled).Name())
	}
}

func reject(next Promise, onRejected OnRejected, e error) {
	switch f := onRejected.(type) {
	case nil:
		next.Reject(e)
	case func() (interface{}, error):
		go call(next, f)
	case func():
		go call1(next, f)
	case func(interface{}) (interface{}, error):
		go call2(next, f, e)
	case func(interface{}):
		go call3(next, f, e)
	case func(error) (interface{}, error):
		go call4(next, f, e)
	case func(error):
		go call5(next, f, e)
	default:
		panic("onRejected can't support this type: " + reflect.TypeOf(onRejected).Name())
	}
}

func timeout(promise Promise, duration time.Duration, reason ...error) Promise {
	next := New()
	timer := time.AfterFunc(duration, func() {
		if len(reason) > 0 {
			next.Reject(reason[0])
		} else {
			next.Reject(TimeoutError{})
		}
	})
	promise.WhenComplete(func() { timer.Stop() }).Fill(next)
	return next
}

func tap(promise Promise, onfulfilledSideEffect func(interface{})) Promise {
	return promise.Then(func(v interface{}) (interface{}, error) {
		onfulfilledSideEffect(v)
		return v, nil
	})
}

// Create creates a Promise object containing the result of asynchronously
// calling computation.
//
// If calling computation returns error, the returned Promise is rejected with
// the error.
//
// If calling computation returns a Promise object, completion of the created
// Promise will wait until the returned Promise completes, and will then
// complete with the same result.
//
// If calling computation returns a non-Promise value, the returned Promise is
// completed with that value.
func Create(computation func() (interface{}, error)) Promise {
	promise := New()
	go call(promise, computation)
	return promise
}

// Sync creates a Promise object containing the result of immediately calling
// computation.
//
// If calling computation returns error, the returned Promise is rejected with
// the error.
//
// If calling computation returns a Promise object, completion of the created
// Promise will wait until the returned Promise completes, and will then
// complete with the same result.
//
// If calling computation returns a non-Promise value, the returned Promise is
// completed with that value.
func Sync(computation func() (interface{}, error)) Promise {
	promise := New()
	call(promise, computation)
	return promise
}

// Delayed creates a Promise object with the given value after a delay.
//
// If the value is a Callable function, it will be executed after the given
// duration has passed, and the Promise object is completed with the result.
func Delayed(duration time.Duration, value interface{}) Promise {
	promise := New()
	go func() {
		time.Sleep(duration)
		switch computation := value.(type) {
		case func() (interface{}, error):
			call(promise, computation)
		case func():
			call1(promise, computation)
		default:
			promise.Resolve(value)
		}
	}()
	return promise
}
<commit_msg>Refactored code<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * promise/promise.go                                     *
 *                                                        *
 * promise interface for Go.                              *
 *                                                        *
 * LastModified: Aug 13, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package promise

import (
	"reflect"
	"time"
)

// OnFulfilled is a function called when the Promise is fulfilled.
// This function has one argument, the fulfillment value.
//
// The function type can be the following:
//     func() (interface{}, error)
//     func()
//     func(interface{}) (interface{}, error)
//     func(interface{})
type OnFulfilled interface{}

// OnRejected is a function called when the Promise is rejected.
// This function has one argument, the rejection reason.
//
// The function type can be the following:
//     func() (interface{}, error)
//     func()
//     func(interface{}) (interface{}, error)
//     func(interface{})
//     func(error) (interface{}, error)
//     func(error)
type OnRejected interface{}

// OnCompleted is a function called when the Promise is completed.
// This function has one argument,
// the fulfillment value when the Promise is fulfilled,
// or the rejection reason when the Promise is rejected.
// The function type can be the following:
//     func() (interface{}, error)
//     func()
//     func(interface{}) (interface{}, error)
//     func(interface{})
type OnCompleted interface{}

// Promise is an interface of the JS Promise/A+ spec
// (https://promisesaplus.com/).
type Promise interface {
	// Then method returns a Promise. It takes two arguments: callback functions
	// for the success and failure cases of the Promise.
	Then(onFulfilled OnFulfilled, onRejected ...OnRejected) Promise

	// Catch handles errors emitted by this Promise.
	//
	// This is the asynchronous equivalent of a "catch" block.
	//
	// Returns a new Promise that will be completed with either the result of
	// this promise or the result of calling the onRejected callback.
	//
	// If this promise completes with a value, the returned promise completes
	// with the same value.
	//
	// If this promise completes with an error, then test is first called with
	// the error value.
	//
	// If test returns false, the error is not handled by this Catch, and the
	// returned promise completes with the same error and stack trace as this
	// promise.
	//
	// If test returns true, onRejected is called with the error and possibly
	// stack trace, and the returned promise is completed with the result of
	// this call in exactly the same way as for Then's onRejected.
	//
	// If test is omitted, it defaults to a function that always returns true.
	// The test function should not panic, but if it does, it is handled as if
	// the the onRejected function had panic.
	Catch(onRejected OnRejected, test ...func(error) bool) Promise

	// Complete is the same way as Then(onCompleted, onCompleted)
	Complete(onCompleted OnCompleted) Promise

	// WhenComplete register a function to be called when the promise completes.
	//
	// The action function is called when this promise completes, whether it
	// does so with a value or with an error.
	//
	// If this promise completes with a value, the returned promise completes
	// with the same value.
	//
	// If this promise completes with an error, the returned promise completes
	// with the same error.
	//
	// The action function should not panic, but if it does, the returned
	// promise completes with a PanicError.
	WhenComplete(action func()) Promise

	// Done is the same semantics as Then except that it don't return a Promise.
	// If the callback function (onFulfilled or onRejected) returns error or
	// panics, the application will be crashing.
	// The result of the callback function will be ignored.
	Done(onFulfilled OnFulfilled, onRejected ...OnRejected)

	// State return the current state of the Promise
	State() State

	// Resolve method returns a Promise object that is resolved with the given
	// value. If the value is a Promise, the returned promise will "follow" that Promise, adopting its eventual state; otherwise the returned promise
	// will be fulfilled with the value.
	Resolve(value interface{})

	// Reject method returns a Promise object that is rejected with the given
	// reason.
	Reject(reason error)

	// Fill the promise with this promise if the promise is in PENDING state.
	// otherwise nothing to do.
	Fill(promise Promise)

	// Timeout create a new promise that will reject with a TimeoutError or a
	// custom reason after a timeout if promise does not fulfill or reject
	// beforehand.
	Timeout(duration time.Duration, reason ...error) Promise

	// Delay create a new promise that will, after duration delay, fulfill with
	// the same value as this promise. If this promise rejects, delayed promise
	// will be rejected immediately.
	Delay(duration time.Duration) Promise

	// Tap executes a function as a side effect when promise fulfills.
	//
	// It returns a new promise:
	// 1. If promise fulfills, onFulfilledSideEffect is executed:
	//     * If onFulfilledSideEffect returns successfully, the promise
	//       returned by tap fulfills with promise's original fulfillment
	//       value.
	//     * If onFulfilledSideEffect panics, the promise returned by tap
	//       rejects with the panic message as the reason.
	// 2. If promise rejects, onFulfilledSideEffect is not executed, and the
	//    promise returned by tap rejects with promise's rejection reason.
	Tap(onfulfilledSideEffect func(interface{})) Promise

	// Get the value and reason synchronously, if this promise in PENDING state.
	// this method will block the current goroutine.
	Get() (interface{}, error)
}

func catch(promise Promise) {
	if e := recover(); e != nil {
		promise.Reject(NewPanicError(e))
	}
}

type func0 func() (interface{}, error)
type func1 func()
type func2 func(interface{}) (interface{}, error)
type func3 func(interface{})
type func4 func(error) (interface{}, error)
type func5 func(error)

func call(promise Promise, computation func0) {
	defer catch(promise)
	if result, err := computation(); err != nil {
		promise.Reject(err)
	} else {
		promise.Resolve(result)
	}
}

func call1(promise Promise, computation func1) {
	defer catch(promise)
	computation()
	promise.Resolve(nil)
}

func call2(promise Promise, computation func2, x interface{}) {
	defer catch(promise)
	if result, err := computation(x); err != nil {
		promise.Reject(err)
	} else {
		promise.Resolve(result)
	}
}

func call3(promise Promise, computation func3, x interface{}) {
	defer catch(promise)
	computation(x)
	promise.Resolve(nil)
}

func call4(promise Promise, computation func4, e error) {
	defer catch(promise)
	if result, err := computation(e); err != nil {
		promise.Reject(err)
	} else {
		promise.Resolve(result)
	}
}

func call5(promise Promise, computation func5, e error) {
	defer catch(promise)
	computation(e)
	promise.Resolve(nil)
}

func resolve(next Promise, onFulfilled OnFulfilled, x interface{}) {
	switch f := onFulfilled.(type) {
	case nil:
		next.Resolve(x)
	case func() (interface{}, error):
		go call(next, f)
	case func():
		go call1(next, f)
	case func(interface{}) (interface{}, error):
		go call2(next, f, x)
	case func(interface{}):
		go call3(next, f, x)
	default:
		panic("onFulfilled can't support this type: " + reflect.TypeOf(onFulfilled).Name())
	}
}

func reject(next Promise, onRejected OnRejected, e error) {
	switch f := onRejected.(type) {
	case nil:
		next.Reject(e)
	case func() (interface{}, error):
		go call(next, f)
	case func():
		go call1(next, f)
	case func(interface{}) (interface{}, error):
		go call2(next, f, e)
	case func(interface{}):
		go call3(next, f, e)
	case func(error) (interface{}, error):
		go call4(next, f, e)
	case func(error):
		go call5(next, f, e)
	default:
		panic("onRejected can't support this type: " + reflect.TypeOf(onRejected).Name())
	}
}

func timeout(promise Promise, duration time.Duration, reason ...error) Promise {
	next := New()
	timer := time.AfterFunc(duration, func() {
		if len(reason) > 0 {
			next.Reject(reason[0])
		} else {
			next.Reject(TimeoutError{})
		}
	})
	promise.WhenComplete(func() { timer.Stop() }).Fill(next)
	return next
}

func tap(promise Promise, onfulfilledSideEffect func(interface{})) Promise {
	return promise.Then(func(v interface{}) (interface{}, error) {
		onfulfilledSideEffect(v)
		return v, nil
	})
}

// Create creates a Promise object containing the result of asynchronously
// calling computation.
//
// If calling computation returns error, the returned Promise is rejected with
// the error.
//
// If calling computation returns a Promise object, completion of the created
// Promise will wait until the returned Promise completes, and will then
// complete with the same result.
//
// If calling computation returns a non-Promise value, the returned Promise is
// completed with that value.
func Create(computation func() (interface{}, error)) Promise {
	promise := New()
	go call(promise, computation)
	return promise
}

// Sync creates a Promise object containing the result of immediately calling
// computation.
//
// If calling computation returns error, the returned Promise is rejected with
// the error.
//
// If calling computation returns a Promise object, completion of the created
// Promise will wait until the returned Promise completes, and will then
// complete with the same result.
//
// If calling computation returns a non-Promise value, the returned Promise is
// completed with that value.
func Sync(computation func() (interface{}, error)) Promise {
	promise := New()
	call(promise, computation)
	return promise
}

// Delayed creates a Promise object with the given value after a delay.
//
// If the value is a Callable function, it will be executed after the given
// duration has passed, and the Promise object is completed with the result.
func Delayed(duration time.Duration, value interface{}) Promise {
	promise := New()
	go func() {
		time.Sleep(duration)
		switch computation := value.(type) {
		case func() (interface{}, error):
			call(promise, computation)
		case func():
			call1(promise, computation)
		default:
			promise.Resolve(value)
		}
	}()
	return promise
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"testing"
)

func TestIsURLHTTP(t *testing.T) {
	url := "http://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLEmpty(t *testing.T) {
	b := isURL("")
	expect(t, b, false)
}

func TestIsURLHTTPS(t *testing.T) {
	url := "https://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLWrong(t *testing.T) {
	url := "somehost.com"

	b := isURL(url)
	expect(t, b, false)
}

func TestIsURLWrongTLD(t *testing.T) {
	url := "http://somehost."

	b := isURL(url)
	expect(t, b, false)
}

func TestFileExists(t *testing.T) {
	fp := "examples/exports/readthedocs.json"

	ex, err := exists(fp)
	expect(t, ex, true)
	expect(t, err, nil)
}

func TestFileDoesNotExist(t *testing.T) {
	fp := "shouldnotbehere.yaml"

	ex, err := exists(fp)
	expect(t, ex, false)
	expect(t, err, nil)
}

func TestImportFromFile(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.Cache.DeleteData()

	err := dbClient.Import("examples/exports/readthedocs.json")
	expect(t, err, nil)

	recordsCount, err := dbClient.Cache.RecordsCount()
	expect(t, err, nil)
	expect(t, recordsCount, 5)
}

func TestImportFromDiskBlankPath(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.Cache.DeleteData()

	err := dbClient.ImportFromDisk("")
	refute(t, err, nil)
}
<commit_msg>testing import from disk when file is not really json<commit_after>package hoverfly

import (
	"testing"
)

func TestIsURLHTTP(t *testing.T) {
	url := "http://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLEmpty(t *testing.T) {
	b := isURL("")
	expect(t, b, false)
}

func TestIsURLHTTPS(t *testing.T) {
	url := "https://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLWrong(t *testing.T) {
	url := "somehost.com"

	b := isURL(url)
	expect(t, b, false)
}

func TestIsURLWrongTLD(t *testing.T) {
	url := "http://somehost."

	b := isURL(url)
	expect(t, b, false)
}

func TestFileExists(t *testing.T) {
	fp := "examples/exports/readthedocs.json"

	ex, err := exists(fp)
	expect(t, ex, true)
	expect(t, err, nil)
}

func TestFileDoesNotExist(t *testing.T) {
	fp := "shouldnotbehere.yaml"

	ex, err := exists(fp)
	expect(t, ex, false)
	expect(t, err, nil)
}

func TestImportFromFile(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.Cache.DeleteData()

	err := dbClient.Import("examples/exports/readthedocs.json")
	expect(t, err, nil)

	recordsCount, err := dbClient.Cache.RecordsCount()
	expect(t, err, nil)
	expect(t, recordsCount, 5)
}

func TestImportFromDiskBlankPath(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.Cache.DeleteData()

	err := dbClient.ImportFromDisk("")
	refute(t, err, nil)
}

func TestImportFromDiskWrongJson(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.Cache.DeleteData()

	err := dbClient.ImportFromDisk("examples/exports/README.md")
	refute(t, err, nil)
}

<|endoftext|>


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

package imports

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"code.google.com/p/go.tools/astutil"
)

// importToGroup is a list of functions which map from an import path to
// a group number.
var importToGroup = []func(importPath string) (num int, ok bool){
	func(importPath string) (num int, ok bool) {
		if strings.HasPrefix(importPath, "appengine") {
			return 2, true
		}
		return
	},
	func(importPath string) (num int, ok bool) {
		if strings.Contains(importPath, ".") {
			return 1, true
		}
		return
	},
}

func importGroup(importPath string) int {
	for _, fn := range importToGroup {
		if n, ok := fn(importPath); ok {
			return n
		}
	}
	return 0
}

func fixImports(fset *token.FileSet, f *ast.File) (added []string, err error) {
	// refs are a set of possible package references currently unsatisfied by imports.
	// first key: either base package (e.g. "fmt") or renamed package
	// second key: referenced package symbol (e.g. "Println")
	refs := make(map[string]map[string]bool)

	// decls are the current package imports. key is base package or renamed package.
	decls := make(map[string]*ast.ImportSpec)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`))
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]bool{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			unusedImport[strings.Trim(is.Path.Value, `"`)] = true
		}
	}
	for ipath := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteImport(fset, f, ipath)
	}

	return added, nil
}

// importPathToName returns the package name for the given import path.
var importPathToName = importPathToNameGoPath

// importPathToNameBasic assumes the package name is the base of import path.
func importPathToNameBasic(importPath string) (packageName string) {
	return path.Base(importPath)
}

// importPathToNameGoPath finds out the actual package name, as declared in its .go files.
// If there's a problem, it falls back to using importPathToNameBasic.
func importPathToNameGoPath(importPath string) (packageName string) {
	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
		return buildPkg.Name
	} else {
		return importPathToNameBasic(importPath)
	}
}

type pkg struct {
	importpath string // full pkg import path, e.g. "net/http"
	dir        string // absolute file path to pkg directory e.g. "/usr/lib/go/src/fmt"
}

var pkgIndexOnce sync.Once

var pkgIndex struct {
	sync.Mutex
	m map[string][]pkg // shortname => []pkg, e.g "http" => "net/http"
}

// gate is a semaphore for limiting concurrency.
type gate chan struct{}

func (g gate) enter() { g <- struct{}{} }
func (g gate) leave() { <-g }

// fsgate protects the OS & filesystem from too much concurrency.
// Too much disk I/O -> too many threads -> swapping and bad scheduling.
var fsgate = make(gate, 8)

func loadPkgIndex() {
	pkgIndex.Lock()
	pkgIndex.m = make(map[string][]pkg)
	pkgIndex.Unlock()

	var wg sync.WaitGroup
	for _, path := range build.Default.SrcDirs() {
		fsgate.enter()
		f, err := os.Open(path)
		if err != nil {
			fsgate.leave()
			fmt.Fprint(os.Stderr, err)
			continue
		}
		children, err := f.Readdir(-1)
		f.Close()
		fsgate.leave()
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			continue
		}
		for _, child := range children {
			if child.IsDir() {
				wg.Add(1)
				go func(path, name string) {
					defer wg.Done()
					loadPkg(&wg, path, name)
				}(path, child.Name())
			}
		}
	}
	wg.Wait()
}

func loadPkg(wg *sync.WaitGroup, root, pkgrelpath string) {
	importpath := filepath.ToSlash(pkgrelpath)
	dir := filepath.Join(root, importpath)

	fsgate.enter()
	defer fsgate.leave()
	pkgDir, err := os.Open(dir)
	if err != nil {
		return
	}
	children, err := pkgDir.Readdir(-1)
	pkgDir.Close()
	if err != nil {
		return
	}
	// hasGo tracks whether a directory actually appears to be a
	// Go source code directory. If $GOPATH == $HOME, and
	// $HOME/src has lots of other large non-Go projects in it,
	// then the calls to importPathToName below can be expensive.
	hasGo := false
	for _, child := range children {
		name := child.Name()
		if name == "" {
			continue
		}
		if c := name[0]; c == '.' || ('0' <= c && c <= '9') {
			continue
		}
		if strings.HasSuffix(name, ".go") {
			hasGo = true
		}
		if child.IsDir() {
			wg.Add(1)
			go func(root, name string) {
				defer wg.Done()
				loadPkg(wg, root, name)
			}(root, filepath.Join(importpath, name))
		}
	}
	if hasGo {
		shortName := importPathToName(importpath)
		pkgIndex.Lock()
		pkgIndex.m[shortName] = append(pkgIndex.m[shortName], pkg{
			importpath: importpath,
			dir:        dir,
		})
		pkgIndex.Unlock()
	}

}

// loadExports returns a list exports for a package.
var loadExports = loadExportsGoPath

func loadExportsGoPath(dir string) map[string]bool {
	exports := make(map[string]bool)
	buildPkg, err := build.ImportDir(dir, 0)
	if err != nil {
		if strings.Contains(err.Error(), "no buildable Go source files in") {
			return nil
		}
		fmt.Fprintf(os.Stderr, "could not import %q: %v\n", dir, err)
		return nil
	}
	fset := token.NewFileSet()
	for _, file := range buildPkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(dir, file), nil, 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "could not parse %q: %v\n", file, err)
			continue
		}
		for name := range f.Scope.Objects {
			if ast.IsExported(name) {
				exports[name] = true
			}
		}
	}
	return exports
}

// findImport searches for a package with the given symbols.
// If no package is found, findImport returns "".
// Declared as a variable rather than a function so goimports can be easily
// extended by adding a file with an init function.
var findImport = findImportGoPath

func findImportGoPath(pkgName string, symbols map[string]bool) (string, bool, error) {
	// Fast path for the standard library.
	// In the common case we hopefully never have to scan the GOPATH, which can
	// be slow with moving disks.
	if pkg, rename, ok := findImportStdlib(pkgName, symbols); ok {
		return pkg, rename, nil
	}

	// TODO(sameer): look at the import lines for other Go files in the
	// local directory, since the user is likely to import the same packages
	// in the current Go file.  Return rename=true when the other Go files
	// use a renamed package that's also used in the current file.

	pkgIndexOnce.Do(loadPkgIndex)

	// Collect exports for packages with matching names.
	var wg sync.WaitGroup
	var pkgsMu sync.Mutex // guards pkgs
	// full importpath => exported symbol => True
	// e.g. "net/http" => "Client" => True
	pkgs := make(map[string]map[string]bool)
	pkgIndex.Lock()
	for _, pkg := range pkgIndex.m[pkgName] {
		wg.Add(1)
		go func(importpath, dir string) {
			defer wg.Done()
			exports := loadExports(dir)
			if exports != nil {
				pkgsMu.Lock()
				pkgs[importpath] = exports
				pkgsMu.Unlock()
			}
		}(pkg.importpath, pkg.dir)
	}
	pkgIndex.Unlock()
	wg.Wait()

	// Filter out packages missing required exported symbols.
	for symbol := range symbols {
		for importpath, exports := range pkgs {
			if !exports[symbol] {
				delete(pkgs, importpath)
			}
		}
	}
	if len(pkgs) == 0 {
		return "", false, nil
	}

	// If there are multiple candidate packages, the shortest one wins.
	// This is a heuristic to prefer the standard library (e.g. "bytes")
	// over e.g. "github.com/foo/bar/bytes".
	shortest := ""
	for importPath := range pkgs {
		if shortest == "" || len(importPath) < len(shortest) {
			shortest = importPath
		}
	}
	return shortest, false, nil
}

type visitFn func(node ast.Node) ast.Visitor

func (fn visitFn) Visit(node ast.Node) ast.Visitor {
	return fn(node)
}

func findImportStdlib(shortPkg string, symbols map[string]bool) (importPath string, rename, ok bool) {
	for symbol := range symbols {
		path := stdlib[shortPkg+"."+symbol]
		if path == "" {
			return "", false, false
		}
		if importPath != "" && importPath != path {
			// Ambiguous. Symbols pointed to different things.
			return "", false, false
		}
		importPath = path
	}
	return importPath, false, importPath != ""
}
<commit_msg>go.tools/imports: recognize exports in cgo files<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package imports

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"code.google.com/p/go.tools/astutil"
)

// importToGroup is a list of functions which map from an import path to
// a group number.
var importToGroup = []func(importPath string) (num int, ok bool){
	func(importPath string) (num int, ok bool) {
		if strings.HasPrefix(importPath, "appengine") {
			return 2, true
		}
		return
	},
	func(importPath string) (num int, ok bool) {
		if strings.Contains(importPath, ".") {
			return 1, true
		}
		return
	},
}

func importGroup(importPath string) int {
	for _, fn := range importToGroup {
		if n, ok := fn(importPath); ok {
			return n
		}
	}
	return 0
}

func fixImports(fset *token.FileSet, f *ast.File) (added []string, err error) {
	// refs are a set of possible package references currently unsatisfied by imports.
	// first key: either base package (e.g. "fmt") or renamed package
	// second key: referenced package symbol (e.g. "Println")
	refs := make(map[string]map[string]bool)

	// decls are the current package imports. key is base package or renamed package.
	decls := make(map[string]*ast.ImportSpec)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`))
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]bool{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			unusedImport[strings.Trim(is.Path.Value, `"`)] = true
		}
	}
	for ipath := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteImport(fset, f, ipath)
	}

	return added, nil
}

// importPathToName returns the package name for the given import path.
var importPathToName = importPathToNameGoPath

// importPathToNameBasic assumes the package name is the base of import path.
func importPathToNameBasic(importPath string) (packageName string) {
	return path.Base(importPath)
}

// importPathToNameGoPath finds out the actual package name, as declared in its .go files.
// If there's a problem, it falls back to using importPathToNameBasic.
func importPathToNameGoPath(importPath string) (packageName string) {
	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
		return buildPkg.Name
	} else {
		return importPathToNameBasic(importPath)
	}
}

type pkg struct {
	importpath string // full pkg import path, e.g. "net/http"
	dir        string // absolute file path to pkg directory e.g. "/usr/lib/go/src/fmt"
}

var pkgIndexOnce sync.Once

var pkgIndex struct {
	sync.Mutex
	m map[string][]pkg // shortname => []pkg, e.g "http" => "net/http"
}

// gate is a semaphore for limiting concurrency.
type gate chan struct{}

func (g gate) enter() { g <- struct{}{} }
func (g gate) leave() { <-g }

// fsgate protects the OS & filesystem from too much concurrency.
// Too much disk I/O -> too many threads -> swapping and bad scheduling.
var fsgate = make(gate, 8)

func loadPkgIndex() {
	pkgIndex.Lock()
	pkgIndex.m = make(map[string][]pkg)
	pkgIndex.Unlock()

	var wg sync.WaitGroup
	for _, path := range build.Default.SrcDirs() {
		fsgate.enter()
		f, err := os.Open(path)
		if err != nil {
			fsgate.leave()
			fmt.Fprint(os.Stderr, err)
			continue
		}
		children, err := f.Readdir(-1)
		f.Close()
		fsgate.leave()
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			continue
		}
		for _, child := range children {
			if child.IsDir() {
				wg.Add(1)
				go func(path, name string) {
					defer wg.Done()
					loadPkg(&wg, path, name)
				}(path, child.Name())
			}
		}
	}
	wg.Wait()
}

func loadPkg(wg *sync.WaitGroup, root, pkgrelpath string) {
	importpath := filepath.ToSlash(pkgrelpath)
	dir := filepath.Join(root, importpath)

	fsgate.enter()
	defer fsgate.leave()
	pkgDir, err := os.Open(dir)
	if err != nil {
		return
	}
	children, err := pkgDir.Readdir(-1)
	pkgDir.Close()
	if err != nil {
		return
	}
	// hasGo tracks whether a directory actually appears to be a
	// Go source code directory. If $GOPATH == $HOME, and
	// $HOME/src has lots of other large non-Go projects in it,
	// then the calls to importPathToName below can be expensive.
	hasGo := false
	for _, child := range children {
		name := child.Name()
		if name == "" {
			continue
		}
		if c := name[0]; c == '.' || ('0' <= c && c <= '9') {
			continue
		}
		if strings.HasSuffix(name, ".go") {
			hasGo = true
		}
		if child.IsDir() {
			wg.Add(1)
			go func(root, name string) {
				defer wg.Done()
				loadPkg(wg, root, name)
			}(root, filepath.Join(importpath, name))
		}
	}
	if hasGo {
		shortName := importPathToName(importpath)
		pkgIndex.Lock()
		pkgIndex.m[shortName] = append(pkgIndex.m[shortName], pkg{
			importpath: importpath,
			dir:        dir,
		})
		pkgIndex.Unlock()
	}

}

// loadExports returns a list exports for a package.
var loadExports = loadExportsGoPath

func loadExportsGoPath(dir string) map[string]bool {
	exports := make(map[string]bool)
	buildPkg, err := build.ImportDir(dir, 0)
	if err != nil {
		if strings.Contains(err.Error(), "no buildable Go source files in") {
			return nil
		}
		fmt.Fprintf(os.Stderr, "could not import %q: %v\n", dir, err)
		return nil
	}
	fset := token.NewFileSet()
	for _, files := range [...][]string{buildPkg.GoFiles, buildPkg.CgoFiles} {
		for _, file := range files {
			f, err := parser.ParseFile(fset, filepath.Join(dir, file), nil, 0)
			if err != nil {
				fmt.Fprintf(os.Stderr, "could not parse %q: %v\n", file, err)
				continue
			}
			for name := range f.Scope.Objects {
				if ast.IsExported(name) {
					exports[name] = true
				}
			}
		}
	}
	return exports
}

// findImport searches for a package with the given symbols.
// If no package is found, findImport returns "".
// Declared as a variable rather than a function so goimports can be easily
// extended by adding a file with an init function.
var findImport = findImportGoPath

func findImportGoPath(pkgName string, symbols map[string]bool) (string, bool, error) {
	// Fast path for the standard library.
	// In the common case we hopefully never have to scan the GOPATH, which can
	// be slow with moving disks.
	if pkg, rename, ok := findImportStdlib(pkgName, symbols); ok {
		return pkg, rename, nil
	}

	// TODO(sameer): look at the import lines for other Go files in the
	// local directory, since the user is likely to import the same packages
	// in the current Go file.  Return rename=true when the other Go files
	// use a renamed package that's also used in the current file.

	pkgIndexOnce.Do(loadPkgIndex)

	// Collect exports for packages with matching names.
	var wg sync.WaitGroup
	var pkgsMu sync.Mutex // guards pkgs
	// full importpath => exported symbol => True
	// e.g. "net/http" => "Client" => True
	pkgs := make(map[string]map[string]bool)
	pkgIndex.Lock()
	for _, pkg := range pkgIndex.m[pkgName] {
		wg.Add(1)
		go func(importpath, dir string) {
			defer wg.Done()
			exports := loadExports(dir)
			if exports != nil {
				pkgsMu.Lock()
				pkgs[importpath] = exports
				pkgsMu.Unlock()
			}
		}(pkg.importpath, pkg.dir)
	}
	pkgIndex.Unlock()
	wg.Wait()

	// Filter out packages missing required exported symbols.
	for symbol := range symbols {
		for importpath, exports := range pkgs {
			if !exports[symbol] {
				delete(pkgs, importpath)
			}
		}
	}
	if len(pkgs) == 0 {
		return "", false, nil
	}

	// If there are multiple candidate packages, the shortest one wins.
	// This is a heuristic to prefer the standard library (e.g. "bytes")
	// over e.g. "github.com/foo/bar/bytes".
	shortest := ""
	for importPath := range pkgs {
		if shortest == "" || len(importPath) < len(shortest) {
			shortest = importPath
		}
	}
	return shortest, false, nil
}

type visitFn func(node ast.Node) ast.Visitor

func (fn visitFn) Visit(node ast.Node) ast.Visitor {
	return fn(node)
}

func findImportStdlib(shortPkg string, symbols map[string]bool) (importPath string, rename, ok bool) {
	for symbol := range symbols {
		path := stdlib[shortPkg+"."+symbol]
		if path == "" {
			return "", false, false
		}
		if importPath != "" && importPath != path {
			// Ambiguous. Symbols pointed to different things.
			return "", false, false
		}
		importPath = path
	}
	return importPath, false, importPath != ""
}
<|endoftext|>


<commit_before>package providers

import (
	"io"
	"os"
	"path"

	"github.com/forklift/fl/flp"
	"github.com/omeid/semver"
)

func init() {
	//List["local"] = &Local{}
}

//				 Name     Versions
type Local struct {
	location string //file system path.
}

func (p *Local) SetLocation(location string) error {
	p.location = location
	return nil
}

func (p Local) Location() string {
	return p.location
}

func (p *Local) Packages() ([]string, error) {
	return nil, nil
}

func (p *Local) Versions() ([]string, error) {

	versions := []string{}

	return versions, nil
}

func (p *Local) Get(name string, ranges string) (*semver.Version, error) {

	versions, err := p.Versions()
	if err != nil {
		return nil, err
	}

	c, err := semver.NewCollection(versions)
	if err != nil {
		return nil, err
	}

	return c.Latest(ranges)
}

func (p *Local) Fetch(ver *semver.Version) (io.Reader, error) {
	return os.Open(path.Join(p.location, flp.Tag(ver)))
}
<commit_msg>Update local provider.<commit_after>package providers

import (
	"io"
	"os"
	"path"

	"github.com/forklift/fl/flp"
	"github.com/omeid/semver"
)

func init() {
	List["local"] = &Local{}
}

//				 Name     Versions
type Local struct {
	location string //file system path.
}

func (p *Local) SetLocation(location string) error {
	p.location = location
	return nil
}

func (p Local) Location() string {
	return p.location
}

func (p *Local) Update() error {
	return nil //Error provider Local doesn't support update?
}
func (p *Local) Packages(filter string) ([]string, error) {
	return nil, nil
}

func (p *Local) Versions(product string) ([]string, error) {

	versions := []string{}

	return versions, nil
}

func (p *Local) Fetch(ver *semver.Version) (io.Reader, error) {
	return os.Open(path.Join(p.location, flp.Tag(ver)))
}

func (p *Local) Source(ver *semver.Version) (string, error) {
	return "", nil
}
<|endoftext|>


<commit_before>package urknall

import (
	"crypto/sha256"
	"fmt"
	"github.com/dynport/gologger"
	"github.com/dynport/gossh"
	"github.com/dynport/urknall/cmd"
	"path"
	"strings"
)

type ProvisionOptions struct {
	LogStdout bool // log stdout of commands
	DryRun    bool
}

type sshClient struct {
	client           *gossh.Client
	host             *Host
	provisionOptions ProvisionOptions
}

func newSSHClient(host *Host, opts *ProvisionOptions) (client *sshClient) {
	if opts == nil {
		opts = &ProvisionOptions{}
	}
	return &sshClient{host: host, client: gossh.New(host.IP, host.user()), provisionOptions: *opts}
}

func (sc *sshClient) provision() (e error) {
	logger.PushPrefix(sc.host.IP)
	defer logger.PopPrefix()

	if e = sc.host.precompileRunlists(); e != nil {
		return e
	}

	return provisionRunlists(sc.host.runlists(), sc.provisionRunlist)
}

func (sc *sshClient) provisionRunlist(rl *Runlist) (e error) {
	tasks := sc.buildTasksForRunlist(rl)

	checksumDir := fmt.Sprintf("/var/cache/urknall/%s", rl.name)

	checksumHash, e := sc.buildChecksumHash(checksumDir)
	if e != nil {
		return fmt.Errorf("failed to build checksum hash: %s", e.Error())
	}

	if sc.host.isSudoRequired() {
		logger.PushPrefix("SUDO")
		defer logger.PopPrefix()
	}

	for i := range tasks {
		task := tasks[i]
		logMsg := task.command.Logging()
		if _, found := checksumHash[task.checksum]; found { // Task is cached.
			logger.Infof("\b[%s][%.8s]%s", gologger.Colorize(33, "CACHED"), task.checksum, logMsg)
			delete(checksumHash, task.checksum) // Delete checksums of cached tasks from hash.
			continue
		}

		if len(checksumHash) > 0 { // All remaining checksums are invalid, as something changed.
			if e = sc.cleanUpRemainingCachedEntries(checksumDir, checksumHash); e != nil {
				return e
			}
			checksumHash = make(map[string]struct{})
		}

		logger.Infof("\b[%s  ][%.8s]%s", gologger.Colorize(34, "EXEC"), task.checksum, logMsg)
		if e = sc.runTask(task, checksumDir); e != nil {
			return e
		}
	}

	return nil
}

func (sc *sshClient) runTask(task *taskData, checksumDir string) (e error) {
	if sc.provisionOptions.DryRun {
		return nil
	}

	stderr := fmt.Sprintf(`>(while read line; do echo "$(date --iso-8601=ns):stderr:$line"; done | tee /tmp/%s.%s.stderr)`, sc.host.user(), task.checksum)
	stdout := fmt.Sprintf(`>(while read line; do echo "$(date --iso-8601=ns):stdout:$line"; done | tee /tmp/%s.%s.stdout)`, sc.host.user(), task.checksum)

	sc.client.ErrorWriter = logger.Error
	if sc.provisionOptions.LogStdout {
		sc.client.DebugWriter = logger.Info
	}

	sCmd := fmt.Sprintf("bash <<EOF_RUNTASK 1> %s 2> %s\n%s\nEOF_RUNTASK\n", stdout, stderr, task.command.Shell())
	if sc.host.isSudoRequired() {
		sCmd = fmt.Sprintf("sudo %s", sCmd)
	}
	rsp, e := sc.client.Execute(sCmd)

	// Write the checksum file (containing information on the command run).
	sc.writeChecksumFile(checksumDir, task.checksum, e != nil, task.command.Logging(), rsp)

	if e != nil {
		return fmt.Errorf("%s (see %s/%s.failed for more information)", e.Error(), checksumDir, task.checksum)
	}
	return nil
}

func (sc *sshClient) executeCommand(cmdRaw string) *gossh.Result {
	if sc.host.isSudoRequired() {
		cmdRaw = fmt.Sprintf("sudo bash <<EOF_ZWO_SUDO\n%s\nEOF_ZWO_SUDO\n", cmdRaw)
	}
	c := &cmd.ShellCommand{Command: cmdRaw}
	result, e := sc.client.Execute(c.Shell())
	if e != nil {
		stderr := ""
		if result != nil {
			stderr = strings.TrimSpace(result.Stderr())
		}
		panic(fmt.Errorf("internal error: %s (%s)", e.Error(), stderr))
	}
	return result
}

func (sc *sshClient) buildChecksumHash(checksumDir string) (checksumMap map[string]struct{}, e error) {
	// Make sure the directory exists.
	sc.executeCommand(fmt.Sprintf("mkdir -p %s", checksumDir))

	checksums := []string{}
	rsp := sc.executeCommand(fmt.Sprintf("ls %s/*.done | xargs", checksumDir))
	for _, checksumFile := range strings.Fields(rsp.Stdout()) {
		checksum := strings.TrimSuffix(path.Base(checksumFile), ".done")
		checksums = append(checksums, checksum)
	}

	checksumMap = make(map[string]struct{})
	for i := range checksums {
		if len(checksums[i]) != 64 {
			return nil, fmt.Errorf("invalid checksum '%s' found in '%s'", checksums[i], checksumDir)
		}
		checksumMap[checksums[i]] = struct{}{}
	}
	return checksumMap, nil
}

func (sc *sshClient) cleanUpRemainingCachedEntries(checksumDir string, checksumHash map[string]struct{}) (e error) {
	invalidCacheEntries := make([]string, 0, len(checksumHash))
	for k, _ := range checksumHash {
		invalidCacheEntries = append(invalidCacheEntries, fmt.Sprintf("%s.done", k))
	}
	if sc.provisionOptions.DryRun {
		logger.Info("invalidated commands:", invalidCacheEntries)
	} else {
		cmd := fmt.Sprintf("cd %s && rm -f *.failed %s", checksumDir, strings.Join(invalidCacheEntries, " "))
		logger.Debug(cmd)
		sc.executeCommand(cmd)
	}
	return nil
}

func (sc *sshClient) writeChecksumFile(checksumDir, checksum string, failed bool, logMsg string, response *gossh.Result) {
	tmpChecksumFiles := "/tmp/" + sc.host.user() + "." + checksum + ".std*"
	checksumFile := checksumDir + "/" + checksum
	if failed {
		checksumFile += ".failed"
	} else {
		checksumFile += ".done"
	}

	// Whoa, super hacky stuff to get the command to the checksum file. The command might contain a lot of stuff, like
	// apostrophes and the like, that would totally nuke a quoted string. Though there is a here doc.
	c := []string{
		fmt.Sprintf(`cat %s | sort >> %s`, tmpChecksumFiles, checksumFile),
		fmt.Sprintf(`rm -f %s`, tmpChecksumFiles),
	}
	sc.executeCommand(fmt.Sprintf("cat <<EOF_COMMAND > %s && %s\n%s\nEOF_COMMAND\n", checksumFile, strings.Join(c, " && "), logMsg))
}

type taskData struct {
	command  cmd.Command // The command to be executed.
	checksum string      // The checksum of the command.
}

func (sc *sshClient) buildTasksForRunlist(rl *Runlist) (tasks []*taskData) {
	tasks = make([]*taskData, 0, len(rl.commands))

	cmdHash := sha256.New()
	for i := range rl.commands {
		rawCmd := rl.commands[i].Shell()
		cmdHash.Write([]byte(rawCmd))

		task := &taskData{command: rl.commands[i], checksum: fmt.Sprintf("%x", cmdHash.Sum(nil))}
		tasks = append(tasks, task)
	}
	return tasks
}
<commit_msg>fixed problem with ls returning error<commit_after>package urknall

import (
	"crypto/sha256"
	"fmt"
	"github.com/dynport/gologger"
	"github.com/dynport/gossh"
	"github.com/dynport/urknall/cmd"
	"path"
	"strings"
)

type ProvisionOptions struct {
	LogStdout bool // log stdout of commands
	DryRun    bool
}

type sshClient struct {
	client           *gossh.Client
	host             *Host
	provisionOptions ProvisionOptions
}

func newSSHClient(host *Host, opts *ProvisionOptions) (client *sshClient) {
	if opts == nil {
		opts = &ProvisionOptions{}
	}
	return &sshClient{host: host, client: gossh.New(host.IP, host.user()), provisionOptions: *opts}
}

func (sc *sshClient) provision() (e error) {
	logger.PushPrefix(sc.host.IP)
	defer logger.PopPrefix()

	if e = sc.host.precompileRunlists(); e != nil {
		return e
	}

	return provisionRunlists(sc.host.runlists(), sc.provisionRunlist)
}

func (sc *sshClient) provisionRunlist(rl *Runlist) (e error) {
	tasks := sc.buildTasksForRunlist(rl)

	checksumDir := fmt.Sprintf("/var/cache/urknall/%s", rl.name)

	checksumHash, e := sc.buildChecksumHash(checksumDir)
	if e != nil {
		return fmt.Errorf("failed to build checksum hash: %s", e.Error())
	}

	if sc.host.isSudoRequired() {
		logger.PushPrefix("SUDO")
		defer logger.PopPrefix()
	}

	for i := range tasks {
		task := tasks[i]
		logMsg := task.command.Logging()
		if _, found := checksumHash[task.checksum]; found { // Task is cached.
			logger.Infof("\b[%s][%.8s]%s", gologger.Colorize(33, "CACHED"), task.checksum, logMsg)
			delete(checksumHash, task.checksum) // Delete checksums of cached tasks from hash.
			continue
		}

		if len(checksumHash) > 0 { // All remaining checksums are invalid, as something changed.
			if e = sc.cleanUpRemainingCachedEntries(checksumDir, checksumHash); e != nil {
				return e
			}
			checksumHash = make(map[string]struct{})
		}

		logger.Infof("\b[%s  ][%.8s]%s", gologger.Colorize(34, "EXEC"), task.checksum, logMsg)
		if e = sc.runTask(task, checksumDir); e != nil {
			return e
		}
	}

	return nil
}

func (sc *sshClient) runTask(task *taskData, checksumDir string) (e error) {
	if sc.provisionOptions.DryRun {
		return nil
	}

	stderr := fmt.Sprintf(`>(while read line; do echo "$(date --iso-8601=ns):stderr:$line"; done | tee /tmp/%s.%s.stderr)`, sc.host.user(), task.checksum)
	stdout := fmt.Sprintf(`>(while read line; do echo "$(date --iso-8601=ns):stdout:$line"; done | tee /tmp/%s.%s.stdout)`, sc.host.user(), task.checksum)

	sc.client.ErrorWriter = logger.Error
	if sc.provisionOptions.LogStdout {
		sc.client.DebugWriter = logger.Info
	}

	sCmd := fmt.Sprintf("bash <<EOF_RUNTASK 1> %s 2> %s\n%s\nEOF_RUNTASK\n", stdout, stderr, task.command.Shell())
	if sc.host.isSudoRequired() {
		sCmd = fmt.Sprintf("sudo %s", sCmd)
	}
	rsp, e := sc.client.Execute(sCmd)

	// Write the checksum file (containing information on the command run).
	sc.writeChecksumFile(checksumDir, task.checksum, e != nil, task.command.Logging(), rsp)

	if e != nil {
		return fmt.Errorf("%s (see %s/%s.failed for more information)", e.Error(), checksumDir, task.checksum)
	}
	return nil
}

func (sc *sshClient) executeCommand(cmdRaw string) *gossh.Result {
	cmdRaw = fmt.Sprintf("bash <<EOF_ZWO_SUDO\n%s\nEOF_ZWO_SUDO\n", cmdRaw)
	if sc.host.isSudoRequired() {
		cmdRaw = "sudo " + cmdRaw
	}
	c := &cmd.ShellCommand{Command: cmdRaw}
	result, e := sc.client.Execute(c.Shell())
	if e != nil {
		stderr := ""
		if result != nil {
			stderr = strings.TrimSpace(result.Stderr())
		}
		panic(fmt.Errorf("internal error: %s (%s)", e.Error(), stderr))
	}
	return result
}

func (sc *sshClient) buildChecksumHash(checksumDir string) (checksumMap map[string]struct{}, e error) {
	// Make sure the directory exists.
	sc.executeCommand(fmt.Sprintf("mkdir -p %s", checksumDir))

	checksums := []string{}
	// The subshell for the if state requires the escaping of the '$' so that the variable is only expanded in the
	// subshell.
	rsp := sc.executeCommand(fmt.Sprintf(`for f in "%s"/*.done; do if [[ -f "\$f" ]]; then echo -n "\$f "; fi; done`, checksumDir))
	for _, checksumFile := range strings.Fields(rsp.Stdout()) {
		checksum := strings.TrimSuffix(path.Base(checksumFile), ".done")
		checksums = append(checksums, checksum)
	}

	checksumMap = make(map[string]struct{})
	for i := range checksums {
		if len(checksums[i]) != 64 {
			return nil, fmt.Errorf("invalid checksum '%s' found in '%s'", checksums[i], checksumDir)
		}
		checksumMap[checksums[i]] = struct{}{}
	}
	return checksumMap, nil
}

func (sc *sshClient) cleanUpRemainingCachedEntries(checksumDir string, checksumHash map[string]struct{}) (e error) {
	invalidCacheEntries := make([]string, 0, len(checksumHash))
	for k, _ := range checksumHash {
		invalidCacheEntries = append(invalidCacheEntries, fmt.Sprintf("%s.done", k))
	}
	if sc.provisionOptions.DryRun {
		logger.Info("invalidated commands:", invalidCacheEntries)
	} else {
		cmd := fmt.Sprintf("cd %s && rm -f *.failed %s", checksumDir, strings.Join(invalidCacheEntries, " "))
		logger.Debug(cmd)
		sc.executeCommand(cmd)
	}
	return nil
}

func (sc *sshClient) writeChecksumFile(checksumDir, checksum string, failed bool, logMsg string, response *gossh.Result) {
	tmpChecksumFiles := "/tmp/" + sc.host.user() + "." + checksum + ".std*"
	checksumFile := checksumDir + "/" + checksum
	if failed {
		checksumFile += ".failed"
	} else {
		checksumFile += ".done"
	}

	// Whoa, super hacky stuff to get the command to the checksum file. The command might contain a lot of stuff, like
	// apostrophes and the like, that would totally nuke a quoted string. Though there is a here doc.
	c := []string{
		fmt.Sprintf(`cat %s | sort >> %s`, tmpChecksumFiles, checksumFile),
		fmt.Sprintf(`rm -f %s`, tmpChecksumFiles),
	}
	sc.executeCommand(fmt.Sprintf("cat <<EOF_COMMAND > %s && %s\n%s\nEOF_COMMAND\n", checksumFile, strings.Join(c, " && "), logMsg))
}

type taskData struct {
	command  cmd.Command // The command to be executed.
	checksum string      // The checksum of the command.
}

func (sc *sshClient) buildTasksForRunlist(rl *Runlist) (tasks []*taskData) {
	tasks = make([]*taskData, 0, len(rl.commands))

	cmdHash := sha256.New()
	for i := range rl.commands {
		rawCmd := rl.commands[i].Shell()
		cmdHash.Write([]byte(rawCmd))

		task := &taskData{command: rl.commands[i], checksum: fmt.Sprintf("%x", cmdHash.Sum(nil))}
		tasks = append(tasks, task)
	}
	return tasks
}
<|endoftext|>


<commit_before>package internal

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"runtime"
	"runtime/pprof"

	profile "github.com/instana/go-sensor/autoprofile/internal/pprof/profile"
)

type blockValues struct {
	delay       float64
	contentions int64
}

type BlockSampler struct {
	top            *CallSite
	prevValues     map[string]*blockValues
	partialProfile *pprof.Profile
}

func NewBlockSampler() *BlockSampler {
	bs := &BlockSampler{
		top:            nil,
		prevValues:     make(map[string]*blockValues),
		partialProfile: nil,
	}

	return bs
}

func (bs *BlockSampler) Reset() {
	bs.top = NewCallSite("", "", 0)
}

func (bs *BlockSampler) Start() error {
	bs.partialProfile = pprof.Lookup("block")
	if bs.partialProfile == nil {
		return errors.New("No block profile found")
	}

	runtime.SetBlockProfileRate(1e6)

	return nil
}

func (bs *BlockSampler) Stop() error {
	p, err := bs.stopBlockSampler()
	if err != nil {
		return err
	}
	if p == nil {
		return errors.New("no profile returned")
	}

	if err := bs.updateBlockProfile(p); err != nil {
		return err
	}

	return nil
}

func (bs *BlockSampler) Profile(duration, timespan int64) (*Profile, error) {
	roots := make([]*CallSite, 0)
	for _, child := range bs.top.children {
		roots = append(roots, child)
	}
	p := NewProfile(CategoryTime, TypeBlockingCalls, UnitMillisecond, roots, duration, timespan)
	return p, nil
}

func (bs *BlockSampler) updateBlockProfile(p *profile.Profile) error {
	contentionIndex := -1
	delayIndex := -1
	for i, s := range p.SampleType {
		if s.Type == "contentions" {
			contentionIndex = i
		} else if s.Type == "delay" {
			delayIndex = i
		}
	}

	if contentionIndex == -1 || delayIndex == -1 {
		return errors.New("Unrecognized profile data")
	}

	for _, s := range p.Sample {
		if shouldSkipStack(s) {
			continue
		}

		delay := float64(s.Value[delayIndex])
		contentions := s.Value[contentionIndex]

		valueKey := generateValueKey(s)
		delay, contentions = bs.getValueChange(valueKey, delay, contentions)

		if contentions == 0 || delay == 0 {
			continue
		}

		// to milliseconds
		delay = delay / 1e6

		current := bs.top
		for i := len(s.Location) - 1; i >= 0; i-- {
			l := s.Location[i]
			funcName, fileName, fileLine := readFuncInfo(l)

			current = current.FindOrAddChild(funcName, fileName, fileLine)
		}
		current.Increment(delay, contentions)
	}

	return nil
}

func generateValueKey(s *profile.Sample) string {
	key := ""
	for _, l := range s.Location {
		key += fmt.Sprintf("%v:", l.Address)
	}

	return key
}

func (bs *BlockSampler) getValueChange(key string, delay float64, contentions int64) (float64, int64) {
	if pv, exists := bs.prevValues[key]; exists {
		delayChange := delay - pv.delay
		contentionsChange := contentions - pv.contentions

		pv.delay = delay
		pv.contentions = contentions

		return delayChange, contentionsChange
	} else {
		bs.prevValues[key] = &blockValues{
			delay:       delay,
			contentions: contentions,
		}

		return delay, contentions
	}
}

func (bs *BlockSampler) stopBlockSampler() (*profile.Profile, error) {
	runtime.SetBlockProfileRate(0)

	var buf bytes.Buffer

	w := bufio.NewWriter(&buf)
	if err := bs.partialProfile.WriteTo(w, 0); err != nil {
		return nil, err
	}

	w.Flush()

	r := bufio.NewReader(&buf)
	p, err := profile.Parse(r)
	if err != nil {
		return nil, err
	}

	if err := symbolizeProfile(p); err != nil {
		return nil, err
	}

	if err := p.CheckValid(); err != nil {
		return nil, err
	}

	return p, nil
}
<commit_msg>Rename (*autoprofile/internal.BlockSampler).stopBlockSampler() -> collectProfile()<commit_after>package internal

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"runtime"
	"runtime/pprof"

	profile "github.com/instana/go-sensor/autoprofile/internal/pprof/profile"
)

type blockValues struct {
	delay       float64
	contentions int64
}

type BlockSampler struct {
	top            *CallSite
	prevValues     map[string]*blockValues
	partialProfile *pprof.Profile
}

func NewBlockSampler() *BlockSampler {
	bs := &BlockSampler{
		top:            nil,
		prevValues:     make(map[string]*blockValues),
		partialProfile: nil,
	}

	return bs
}

func (bs *BlockSampler) Reset() {
	bs.top = NewCallSite("", "", 0)
}

func (bs *BlockSampler) Start() error {
	bs.partialProfile = pprof.Lookup("block")
	if bs.partialProfile == nil {
		return errors.New("No block profile found")
	}

	runtime.SetBlockProfileRate(1e6)

	return nil
}

func (bs *BlockSampler) Stop() error {
	runtime.SetBlockProfileRate(0)

	p, err := bs.collectProfile()
	if err != nil {
		return err
	}

	if p == nil {
		return errors.New("no profile returned")
	}

	if err := bs.updateBlockProfile(p); err != nil {
		return err
	}

	return nil
}

func (bs *BlockSampler) Profile(duration, timespan int64) (*Profile, error) {
	roots := make([]*CallSite, 0)
	for _, child := range bs.top.children {
		roots = append(roots, child)
	}
	p := NewProfile(CategoryTime, TypeBlockingCalls, UnitMillisecond, roots, duration, timespan)
	return p, nil
}

func (bs *BlockSampler) updateBlockProfile(p *profile.Profile) error {
	contentionIndex := -1
	delayIndex := -1
	for i, s := range p.SampleType {
		if s.Type == "contentions" {
			contentionIndex = i
		} else if s.Type == "delay" {
			delayIndex = i
		}
	}

	if contentionIndex == -1 || delayIndex == -1 {
		return errors.New("Unrecognized profile data")
	}

	for _, s := range p.Sample {
		if shouldSkipStack(s) {
			continue
		}

		delay := float64(s.Value[delayIndex])
		contentions := s.Value[contentionIndex]

		valueKey := generateValueKey(s)
		delay, contentions = bs.getValueChange(valueKey, delay, contentions)

		if contentions == 0 || delay == 0 {
			continue
		}

		// to milliseconds
		delay = delay / 1e6

		current := bs.top
		for i := len(s.Location) - 1; i >= 0; i-- {
			l := s.Location[i]
			funcName, fileName, fileLine := readFuncInfo(l)

			current = current.FindOrAddChild(funcName, fileName, fileLine)
		}
		current.Increment(delay, contentions)
	}

	return nil
}

func generateValueKey(s *profile.Sample) string {
	key := ""
	for _, l := range s.Location {
		key += fmt.Sprintf("%v:", l.Address)
	}

	return key
}

func (bs *BlockSampler) getValueChange(key string, delay float64, contentions int64) (float64, int64) {
	if pv, exists := bs.prevValues[key]; exists {
		delayChange := delay - pv.delay
		contentionsChange := contentions - pv.contentions

		pv.delay = delay
		pv.contentions = contentions

		return delayChange, contentionsChange
	} else {
		bs.prevValues[key] = &blockValues{
			delay:       delay,
			contentions: contentions,
		}

		return delay, contentions
	}
}

func (bs *BlockSampler) collectProfile() (*profile.Profile, error) {
	var buf bytes.Buffer

	w := bufio.NewWriter(&buf)
	if err := bs.partialProfile.WriteTo(w, 0); err != nil {
		return nil, err
	}

	w.Flush()

	r := bufio.NewReader(&buf)
	p, err := profile.Parse(r)
	if err != nil {
		return nil, err
	}

	if err := symbolizeProfile(p); err != nil {
		return nil, err
	}

	if err := p.CheckValid(); err != nil {
		return nil, err
	}

	return p, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/validation"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/schema"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudhsmv2"
	"github.com/hashicorp/terraform/helper/resource"
)

func resourceAwsCloudHsm2Cluster() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCloudHsm2ClusterCreate,
		Read:   resourceAwsCloudHsm2ClusterRead,
		Update: resourceAwsCloudHsm2ClusterUpdate,
		Delete: resourceAwsCloudHsm2ClusterDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Timeouts: &schema.ResourceTimeout{
			Create: schema.DefaultTimeout(120 * time.Minute),
			Update: schema.DefaultTimeout(120 * time.Minute),
			Delete: schema.DefaultTimeout(120 * time.Minute),
		},

		Schema: map[string]*schema.Schema{
			"source_backup_identifier": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"hsm_type": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringInSlice([]string{"hsm1.medium"}, false),
			},

			"subnet_ids": {
				Type:     schema.TypeSet,
				Required: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
				ForceNew: true,
			},

			"cluster_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"vpc_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"cluster_certificates": {
				Type:     schema.TypeList,
				MaxItems: 1,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"cluster_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"cluster_csr": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"aws_hardware_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"hsm_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"manufacturer_hardware_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},

			"security_group_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"cluster_state": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

func describeCloudHsm2Cluster(conn *cloudhsmv2.CloudHSMV2, clusterId string) (*cloudhsmv2.Cluster, error) {
	filters := []*string{&clusterId}
	result := int64(1)
	out, err := conn.DescribeClusters(&cloudhsmv2.DescribeClustersInput{
		Filters: map[string][]*string{
			"clusterIds": filters,
		},
		MaxResults: &result,
	})
	if err != nil {
		log.Printf("[WARN] Error on retrieving CloudHSMv2 Cluster (%s) when waiting: %s", clusterId, err)
		return nil, err
	}

	var cluster *cloudhsmv2.Cluster

	for _, c := range out.Clusters {
		if aws.StringValue(c.ClusterId) == clusterId {
			cluster = c
			break
		}
	}
	return cluster, nil
}

func resourceAwsCloudHsm2ClusterRefreshFunc(conn *cloudhsmv2.CloudHSMV2, clusterId string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		cluster, err := describeCloudHsm2Cluster(conn, clusterId)

		if cluster == nil {
			return 42, "destroyed", nil
		}

		if cluster.State != nil {
			log.Printf("[DEBUG] CloudHSMv2 Cluster status (%s): %s", clusterId, *cluster.State)
		}

		return cluster, aws.StringValue(cluster.State), err
	}
}

func resourceAwsCloudHsm2ClusterCreate(d *schema.ResourceData, meta interface{}) error {
	cloudhsm2 := meta.(*AWSClient).cloudhsmv2conn

	input := &cloudhsmv2.CreateClusterInput{
		HsmType:   aws.String(d.Get("hsm_type").(string)),
		SubnetIds: expandStringSet(d.Get("subnet_ids").(*schema.Set)),
	}

	backupId := d.Get("source_backup_identifier").(string)
	if len(backupId) != 0 {
		input.SourceBackupId = aws.String(backupId)
	}

	log.Printf("[DEBUG] CloudHSMv2 Cluster create %s", input)

	var output *cloudhsmv2.CreateClusterOutput

	err := resource.Retry(180*time.Second, func() *resource.RetryError {
		var err error
		output, err = cloudhsm2.CreateCluster(input)
		if err != nil {
			if isAWSErr(err, cloudhsmv2.ErrCodeCloudHsmInternalFailureException, "request was rejected because of an AWS CloudHSM internal failure") {
				log.Printf("[DEBUG] CloudHSMv2 Cluster re-try creating %s", input)
				return resource.RetryableError(err)
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if isResourceTimeoutError(err) {
		output, err = cloudhsm2.CreateCluster(input)
	}

	if err != nil {
		return fmt.Errorf("error creating CloudHSMv2 Cluster: %s", err)
	}

	d.SetId(aws.StringValue(output.Cluster.ClusterId))
	log.Printf("[INFO] CloudHSMv2 Cluster ID: %s", d.Id())
	log.Println("[INFO] Waiting for CloudHSMv2 Cluster to be available")

	targetState := cloudhsmv2.ClusterStateUninitialized
	if len(backupId) > 0 {
		targetState = cloudhsmv2.ClusterStateActive
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{cloudhsmv2.ClusterStateCreateInProgress, cloudhsmv2.ClusterStateInitializeInProgress},
		Target:     []string{targetState},
		Refresh:    resourceAwsCloudHsm2ClusterRefreshFunc(cloudhsm2, d.Id()),
		Timeout:    d.Timeout(schema.TimeoutCreate),
		MinTimeout: 30 * time.Second,
		Delay:      30 * time.Second,
	}

	// Wait, catching any errors
	_, errWait := stateConf.WaitForState()
	if errWait != nil {
		if len(backupId) == 0 {
			return fmt.Errorf("[WARN] Error waiting for CloudHSMv2 Cluster state to be \"UNINITIALIZED\": %s", errWait)
		} else {
			return fmt.Errorf("[WARN] Error waiting for CloudHSMv2 Cluster state to be \"ACTIVE\": %s", errWait)
		}
	}

	if err := setTagsAwsCloudHsm2Cluster(cloudhsm2, d); err != nil {
		return err
	}

	return resourceAwsCloudHsm2ClusterRead(d, meta)
}

func resourceAwsCloudHsm2ClusterRead(d *schema.ResourceData, meta interface{}) error {

	cluster, err := describeCloudHsm2Cluster(meta.(*AWSClient).cloudhsmv2conn, d.Id())

	if cluster == nil {
		log.Printf("[WARN] CloudHSMv2 Cluster (%s) not found", d.Id())
		d.SetId("")
		return err
	}

	log.Printf("[INFO] Reading CloudHSMv2 Cluster Information: %s", d.Id())

	d.Set("cluster_id", cluster.ClusterId)
	d.Set("cluster_state", cluster.State)
	d.Set("security_group_id", cluster.SecurityGroup)
	d.Set("vpc_id", cluster.VpcId)
	d.Set("source_backup_identifier", cluster.SourceBackupId)
	d.Set("hsm_type", cluster.HsmType)
	if err := d.Set("cluster_certificates", readCloudHsm2ClusterCertificates(cluster)); err != nil {
		return fmt.Errorf("error setting cluster_certificates: %s", err)
	}

	var subnets []string
	for _, sn := range cluster.SubnetMapping {
		subnets = append(subnets, aws.StringValue(sn))
	}
	if err := d.Set("subnet_ids", subnets); err != nil {
		return fmt.Errorf("Error saving Subnet IDs to state for CloudHSMv2 Cluster (%s): %s", d.Id(), err)
	}

	return nil
}

func resourceAwsCloudHsm2ClusterUpdate(d *schema.ResourceData, meta interface{}) error {
	cloudhsm2 := meta.(*AWSClient).cloudhsmv2conn

	if err := setTagsAwsCloudHsm2Cluster(cloudhsm2, d); err != nil {
		return err
	}

	return resourceAwsCloudHsm2ClusterRead(d, meta)
}

func resourceAwsCloudHsm2ClusterDelete(d *schema.ResourceData, meta interface{}) error {
	cloudhsm2 := meta.(*AWSClient).cloudhsmv2conn
	input := &cloudhsmv2.DeleteClusterInput{
		ClusterId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] CloudHSMv2 Delete cluster: %s", d.Id())
	err := resource.Retry(180*time.Second, func() *resource.RetryError {
		var err error
		_, err = cloudhsm2.DeleteCluster(input)
		if err != nil {
			if isAWSErr(err, cloudhsmv2.ErrCodeCloudHsmInternalFailureException, "request was rejected because of an AWS CloudHSM internal failure") {
				log.Printf("[DEBUG] CloudHSMv2 Cluster re-try deleting %s", d.Id())
				return resource.RetryableError(err)
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if isResourceTimeoutError(err) {
		_, err = cloudhsm2.DeleteCluster(input)
	}

	if err != nil {
		return err
	}
	log.Println("[INFO] Waiting for CloudHSMv2 Cluster to be deleted")

	stateConf := &resource.StateChangeConf{
		Pending:    []string{cloudhsmv2.ClusterStateDeleteInProgress},
		Target:     []string{cloudhsmv2.ClusterStateDeleted},
		Refresh:    resourceAwsCloudHsm2ClusterRefreshFunc(cloudhsm2, d.Id()),
		Timeout:    d.Timeout(schema.TimeoutCreate),
		MinTimeout: 30 * time.Second,
		Delay:      30 * time.Second,
	}

	// Wait, catching any errors
	_, errWait := stateConf.WaitForState()
	if errWait != nil {
		return fmt.Errorf("Error waiting for CloudHSMv2 Cluster state to be \"DELETED\": %s", errWait)
	}

	return nil
}

func setTagsAwsCloudHsm2Cluster(conn *cloudhsmv2.CloudHSMV2, d *schema.ResourceData) error {
	if d.HasChange("tags") {
		oraw, nraw := d.GetChange("tags")
		create, remove := diffTagsGeneric(oraw.(map[string]interface{}), nraw.(map[string]interface{}))

		if len(remove) > 0 {
			log.Printf("[DEBUG] Removing tags: %#v", remove)
			keys := make([]*string, 0, len(remove))
			for k := range remove {
				keys = append(keys, aws.String(k))
			}

			_, err := conn.UntagResource(&cloudhsmv2.UntagResourceInput{
				ResourceId: aws.String(d.Id()),
				TagKeyList: keys,
			})
			if err != nil {
				return err
			}
		}
		if len(create) > 0 {
			log.Printf("[DEBUG] Creating tags: %#v", create)
			tagList := make([]*cloudhsmv2.Tag, 0, len(create))
			for k, v := range create {
				tagList = append(tagList, &cloudhsmv2.Tag{
					Key:   &k,
					Value: v,
				})
			}
			_, err := conn.TagResource(&cloudhsmv2.TagResourceInput{
				ResourceId: aws.String(d.Id()),
				TagList:    tagList,
			})
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func readCloudHsm2ClusterCertificates(cluster *cloudhsmv2.Cluster) []map[string]interface{} {
	certs := map[string]interface{}{}
	if cluster.Certificates != nil {
		if aws.StringValue(cluster.State) == "UNINITIALIZED" {
			certs["cluster_csr"] = aws.StringValue(cluster.Certificates.ClusterCsr)
			certs["aws_hardware_certificate"] = aws.StringValue(cluster.Certificates.AwsHardwareCertificate)
			certs["hsm_certificate"] = aws.StringValue(cluster.Certificates.HsmCertificate)
			certs["manufacturer_hardware_certificate"] = aws.StringValue(cluster.Certificates.ManufacturerHardwareCertificate)
		} else if aws.StringValue(cluster.State) == "ACTIVE" {
			certs["cluster_certificate"] = aws.StringValue(cluster.Certificates.ClusterCertificate)
		}
	}
	if len(certs) > 0 {
		return []map[string]interface{}{certs}
	}
	return []map[string]interface{}{}
}
<commit_msg>Small change to fix bug where only 1 tag would be created<commit_after>package aws

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/validation"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/schema"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudhsmv2"
	"github.com/hashicorp/terraform/helper/resource"
)

func resourceAwsCloudHsm2Cluster() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCloudHsm2ClusterCreate,
		Read:   resourceAwsCloudHsm2ClusterRead,
		Update: resourceAwsCloudHsm2ClusterUpdate,
		Delete: resourceAwsCloudHsm2ClusterDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Timeouts: &schema.ResourceTimeout{
			Create: schema.DefaultTimeout(120 * time.Minute),
			Update: schema.DefaultTimeout(120 * time.Minute),
			Delete: schema.DefaultTimeout(120 * time.Minute),
		},

		Schema: map[string]*schema.Schema{
			"source_backup_identifier": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"hsm_type": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringInSlice([]string{"hsm1.medium"}, false),
			},

			"subnet_ids": {
				Type:     schema.TypeSet,
				Required: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
				ForceNew: true,
			},

			"cluster_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"vpc_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"cluster_certificates": {
				Type:     schema.TypeList,
				MaxItems: 1,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"cluster_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"cluster_csr": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"aws_hardware_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"hsm_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"manufacturer_hardware_certificate": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},

			"security_group_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"cluster_state": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

func describeCloudHsm2Cluster(conn *cloudhsmv2.CloudHSMV2, clusterId string) (*cloudhsmv2.Cluster, error) {
	filters := []*string{&clusterId}
	result := int64(1)
	out, err := conn.DescribeClusters(&cloudhsmv2.DescribeClustersInput{
		Filters: map[string][]*string{
			"clusterIds": filters,
		},
		MaxResults: &result,
	})
	if err != nil {
		log.Printf("[WARN] Error on retrieving CloudHSMv2 Cluster (%s) when waiting: %s", clusterId, err)
		return nil, err
	}

	var cluster *cloudhsmv2.Cluster

	for _, c := range out.Clusters {
		if aws.StringValue(c.ClusterId) == clusterId {
			cluster = c
			break
		}
	}
	return cluster, nil
}

func resourceAwsCloudHsm2ClusterRefreshFunc(conn *cloudhsmv2.CloudHSMV2, clusterId string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		cluster, err := describeCloudHsm2Cluster(conn, clusterId)

		if cluster == nil {
			return 42, "destroyed", nil
		}

		if cluster.State != nil {
			log.Printf("[DEBUG] CloudHSMv2 Cluster status (%s): %s", clusterId, *cluster.State)
		}

		return cluster, aws.StringValue(cluster.State), err
	}
}

func resourceAwsCloudHsm2ClusterCreate(d *schema.ResourceData, meta interface{}) error {
	cloudhsm2 := meta.(*AWSClient).cloudhsmv2conn

	input := &cloudhsmv2.CreateClusterInput{
		HsmType:   aws.String(d.Get("hsm_type").(string)),
		SubnetIds: expandStringSet(d.Get("subnet_ids").(*schema.Set)),
	}

	backupId := d.Get("source_backup_identifier").(string)
	if len(backupId) != 0 {
		input.SourceBackupId = aws.String(backupId)
	}

	log.Printf("[DEBUG] CloudHSMv2 Cluster create %s", input)

	var output *cloudhsmv2.CreateClusterOutput

	err := resource.Retry(180*time.Second, func() *resource.RetryError {
		var err error
		output, err = cloudhsm2.CreateCluster(input)
		if err != nil {
			if isAWSErr(err, cloudhsmv2.ErrCodeCloudHsmInternalFailureException, "request was rejected because of an AWS CloudHSM internal failure") {
				log.Printf("[DEBUG] CloudHSMv2 Cluster re-try creating %s", input)
				return resource.RetryableError(err)
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if isResourceTimeoutError(err) {
		output, err = cloudhsm2.CreateCluster(input)
	}

	if err != nil {
		return fmt.Errorf("error creating CloudHSMv2 Cluster: %s", err)
	}

	d.SetId(aws.StringValue(output.Cluster.ClusterId))
	log.Printf("[INFO] CloudHSMv2 Cluster ID: %s", d.Id())
	log.Println("[INFO] Waiting for CloudHSMv2 Cluster to be available")

	targetState := cloudhsmv2.ClusterStateUninitialized
	if len(backupId) > 0 {
		targetState = cloudhsmv2.ClusterStateActive
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{cloudhsmv2.ClusterStateCreateInProgress, cloudhsmv2.ClusterStateInitializeInProgress},
		Target:     []string{targetState},
		Refresh:    resourceAwsCloudHsm2ClusterRefreshFunc(cloudhsm2, d.Id()),
		Timeout:    d.Timeout(schema.TimeoutCreate),
		MinTimeout: 30 * time.Second,
		Delay:      30 * time.Second,
	}

	// Wait, catching any errors
	_, errWait := stateConf.WaitForState()
	if errWait != nil {
		if len(backupId) == 0 {
			return fmt.Errorf("[WARN] Error waiting for CloudHSMv2 Cluster state to be \"UNINITIALIZED\": %s", errWait)
		} else {
			return fmt.Errorf("[WARN] Error waiting for CloudHSMv2 Cluster state to be \"ACTIVE\": %s", errWait)
		}
	}

	if err := setTagsAwsCloudHsm2Cluster(cloudhsm2, d); err != nil {
		return err
	}

	return resourceAwsCloudHsm2ClusterRead(d, meta)
}

func resourceAwsCloudHsm2ClusterRead(d *schema.ResourceData, meta interface{}) error {

	cluster, err := describeCloudHsm2Cluster(meta.(*AWSClient).cloudhsmv2conn, d.Id())

	if cluster == nil {
		log.Printf("[WARN] CloudHSMv2 Cluster (%s) not found", d.Id())
		d.SetId("")
		return err
	}

	log.Printf("[INFO] Reading CloudHSMv2 Cluster Information: %s", d.Id())

	d.Set("cluster_id", cluster.ClusterId)
	d.Set("cluster_state", cluster.State)
	d.Set("security_group_id", cluster.SecurityGroup)
	d.Set("vpc_id", cluster.VpcId)
	d.Set("source_backup_identifier", cluster.SourceBackupId)
	d.Set("hsm_type", cluster.HsmType)
	if err := d.Set("cluster_certificates", readCloudHsm2ClusterCertificates(cluster)); err != nil {
		return fmt.Errorf("error setting cluster_certificates: %s", err)
	}

	var subnets []string
	for _, sn := range cluster.SubnetMapping {
		subnets = append(subnets, aws.StringValue(sn))
	}
	if err := d.Set("subnet_ids", subnets); err != nil {
		return fmt.Errorf("Error saving Subnet IDs to state for CloudHSMv2 Cluster (%s): %s", d.Id(), err)
	}

	return nil
}

func resourceAwsCloudHsm2ClusterUpdate(d *schema.ResourceData, meta interface{}) error {
	cloudhsm2 := meta.(*AWSClient).cloudhsmv2conn

	if err := setTagsAwsCloudHsm2Cluster(cloudhsm2, d); err != nil {
		return err
	}

	return resourceAwsCloudHsm2ClusterRead(d, meta)
}

func resourceAwsCloudHsm2ClusterDelete(d *schema.ResourceData, meta interface{}) error {
	cloudhsm2 := meta.(*AWSClient).cloudhsmv2conn
	input := &cloudhsmv2.DeleteClusterInput{
		ClusterId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] CloudHSMv2 Delete cluster: %s", d.Id())
	err := resource.Retry(180*time.Second, func() *resource.RetryError {
		var err error
		_, err = cloudhsm2.DeleteCluster(input)
		if err != nil {
			if isAWSErr(err, cloudhsmv2.ErrCodeCloudHsmInternalFailureException, "request was rejected because of an AWS CloudHSM internal failure") {
				log.Printf("[DEBUG] CloudHSMv2 Cluster re-try deleting %s", d.Id())
				return resource.RetryableError(err)
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if isResourceTimeoutError(err) {
		_, err = cloudhsm2.DeleteCluster(input)
	}

	if err != nil {
		return err
	}
	log.Println("[INFO] Waiting for CloudHSMv2 Cluster to be deleted")

	stateConf := &resource.StateChangeConf{
		Pending:    []string{cloudhsmv2.ClusterStateDeleteInProgress},
		Target:     []string{cloudhsmv2.ClusterStateDeleted},
		Refresh:    resourceAwsCloudHsm2ClusterRefreshFunc(cloudhsm2, d.Id()),
		Timeout:    d.Timeout(schema.TimeoutCreate),
		MinTimeout: 30 * time.Second,
		Delay:      30 * time.Second,
	}

	// Wait, catching any errors
	_, errWait := stateConf.WaitForState()
	if errWait != nil {
		return fmt.Errorf("Error waiting for CloudHSMv2 Cluster state to be \"DELETED\": %s", errWait)
	}

	return nil
}

func setTagsAwsCloudHsm2Cluster(conn *cloudhsmv2.CloudHSMV2, d *schema.ResourceData) error {
	if d.HasChange("tags") {
		oraw, nraw := d.GetChange("tags")
		create, remove := diffTagsGeneric(oraw.(map[string]interface{}), nraw.(map[string]interface{}))

		if len(remove) > 0 {
			log.Printf("[DEBUG] Removing tags: %#v", remove)
			keys := make([]*string, 0, len(remove))
			for k := range remove {
				keys = append(keys, aws.String(k))
			}

			_, err := conn.UntagResource(&cloudhsmv2.UntagResourceInput{
				ResourceId: aws.String(d.Id()),
				TagKeyList: keys,
			})
			if err != nil {
				return err
			}
		}
		if len(create) > 0 {
			log.Printf("[DEBUG] Creating tags: %#v", create)
			tagList := make([]*cloudhsmv2.Tag, 0, len(create))
			for k, v := range create {
				tagList = append(tagList, &cloudhsmv2.Tag{
					Key:   aws.String(k),
					Value: v,
				})
			}
			_, err := conn.TagResource(&cloudhsmv2.TagResourceInput{
				ResourceId: aws.String(d.Id()),
				TagList:    tagList,
			})
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func readCloudHsm2ClusterCertificates(cluster *cloudhsmv2.Cluster) []map[string]interface{} {
	certs := map[string]interface{}{}
	if cluster.Certificates != nil {
		if aws.StringValue(cluster.State) == "UNINITIALIZED" {
			certs["cluster_csr"] = aws.StringValue(cluster.Certificates.ClusterCsr)
			certs["aws_hardware_certificate"] = aws.StringValue(cluster.Certificates.AwsHardwareCertificate)
			certs["hsm_certificate"] = aws.StringValue(cluster.Certificates.HsmCertificate)
			certs["manufacturer_hardware_certificate"] = aws.StringValue(cluster.Certificates.ManufacturerHardwareCertificate)
		} else if aws.StringValue(cluster.State) == "ACTIVE" {
			certs["cluster_certificate"] = aws.StringValue(cluster.Certificates.ClusterCertificate)
		}
	}
	if len(certs) > 0 {
		return []map[string]interface{}{certs}
	}
	return []map[string]interface{}{}
}
<|endoftext|>


<commit_before>package consul

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	consulapi "github.com/hashicorp/consul/api"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/state"
	"github.com/hashicorp/terraform/state/remote"
)

const (
	lockSuffix     = "/.lock"
	lockInfoSuffix = "/.lockinfo"
)

// RemoteClient is a remote client that stores data in Consul.
type RemoteClient struct {
	Client *consulapi.Client
	Path   string
	GZip   bool

	consulLock *consulapi.Lock
	lockCh     <-chan struct{}
}

func (c *RemoteClient) Get() (*remote.Payload, error) {
	pair, _, err := c.Client.KV().Get(c.Path, nil)
	if err != nil {
		return nil, err
	}
	if pair == nil {
		return nil, nil
	}

	payload := pair.Value
	// If the payload starts with 0x1f, it's gzip, not json
	if len(pair.Value) >= 1 && pair.Value[0] == '\x1f' {
		if data, err := uncompressState(pair.Value); err == nil {
			payload = data
		} else {
			return nil, err
		}
	}

	md5 := md5.Sum(pair.Value)
	return &remote.Payload{
		Data: payload,
		MD5:  md5[:],
	}, nil
}

func (c *RemoteClient) Put(data []byte) error {
	payload := data
	if c.GZip {
		if compressedState, err := compressState(data); err == nil {
			payload = compressedState
		} else {
			return err
		}
	}

	kv := c.Client.KV()
	_, err := kv.Put(&consulapi.KVPair{
		Key:   c.Path,
		Value: payload,
	}, nil)
	return err
}

func (c *RemoteClient) Delete() error {
	kv := c.Client.KV()
	_, err := kv.Delete(c.Path, nil)
	return err
}

func (c *RemoteClient) putLockInfo(info *state.LockInfo) error {
	info.Path = c.Path
	info.Created = time.Now().UTC()

	kv := c.Client.KV()
	_, err := kv.Put(&consulapi.KVPair{
		Key:   c.Path + lockInfoSuffix,
		Value: info.Marshal(),
	}, nil)

	return err
}

func (c *RemoteClient) getLockInfo() (*state.LockInfo, error) {
	path := c.Path + lockInfoSuffix
	pair, _, err := c.Client.KV().Get(path, nil)
	if err != nil {
		return nil, err
	}
	if pair == nil {
		return nil, nil
	}

	li := &state.LockInfo{}
	err = json.Unmarshal(pair.Value, li)
	if err != nil {
		return nil, fmt.Errorf("error unmarshaling lock info: %s", err)
	}

	return li, nil
}

func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
	select {
	case <-c.lockCh:
		// We had a lock, but lost it.
		// Since we typically only call lock once, we shouldn't ever see this.
		return "", errors.New("lost consul lock")
	default:
		if c.lockCh != nil {
			// we have an active lock already
			return "", fmt.Errorf("state %q already locked", c.Path)
		}
	}

	if c.consulLock == nil {
		opts := &consulapi.LockOptions{
			Key: c.Path + lockSuffix,
			// only wait briefly, so terraform has the choice to fail fast or
			// retry as needed.
			LockWaitTime: time.Second,
			LockTryOnce:  true,
		}

		lock, err := c.Client.LockOpts(opts)
		if err != nil {
			return "", err
		}

		c.consulLock = lock
	}

	lockErr := &state.LockError{}

	lockCh, err := c.consulLock.Lock(make(chan struct{}))
	if err != nil {
		lockErr.Err = err
		return "", lockErr
	}

	if lockCh == nil {
		lockInfo, e := c.getLockInfo()
		if e != nil {
			lockErr.Err = e
			return "", lockErr
		}

		lockErr.Info = lockInfo
		return "", lockErr
	}

	c.lockCh = lockCh

	err = c.putLockInfo(info)
	if err != nil {
		if unlockErr := c.Unlock(info.ID); unlockErr != nil {
			err = multierror.Append(err, unlockErr)
		}

		return "", err
	}

	return info.ID, nil
}

func (c *RemoteClient) Unlock(id string) error {
	// this doesn't use the lock id, because the lock is tied to the consul client.
	if c.consulLock == nil || c.lockCh == nil {
		return nil
	}

	select {
	case <-c.lockCh:
		return errors.New("consul lock was lost")
	default:
	}

	err := c.consulLock.Unlock()
	c.lockCh = nil

	// This is only cleanup, and will fail if the lock was immediately taken by
	// another client, so we don't report an error to the user here.
	c.consulLock.Destroy()

	kv := c.Client.KV()
	_, delErr := kv.Delete(c.Path+lockInfoSuffix, nil)
	if delErr != nil {
		err = multierror.Append(err, delErr)
	}

	return err
}

func compressState(data []byte) ([]byte, error) {
	b := new(bytes.Buffer)
	gz := gzip.NewWriter(b)
	if _, err := gz.Write(data); err != nil {
		return nil, err
	}
	if err := gz.Flush(); err != nil {
		return nil, err
	}
	if err := gz.Close(); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func uncompressState(data []byte) ([]byte, error) {
	b := new(bytes.Buffer)
	gz, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	b.ReadFrom(gz)
	if err := gz.Close(); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}
<commit_msg>use CAS for consul state Put<commit_after>package consul

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	consulapi "github.com/hashicorp/consul/api"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/state"
	"github.com/hashicorp/terraform/state/remote"
)

const (
	lockSuffix     = "/.lock"
	lockInfoSuffix = "/.lockinfo"
)

// RemoteClient is a remote client that stores data in Consul.
type RemoteClient struct {
	Client *consulapi.Client
	Path   string
	GZip   bool

	mu sync.Mutex

	// The index of the last state we wrote.
	// If this is > 0, Put will perform a CAS to ensure that the state wasn't
	// changed during the operation. This is important even with locks, because
	// if the client loses the lock for some reason, then reacquires it, we
	// need to make sure that the state was not modified.
	modifyIndex uint64

	consulLock *consulapi.Lock
	lockCh     <-chan struct{}

	info *state.LockInfo
}

func (c *RemoteClient) Get() (*remote.Payload, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	pair, _, err := c.Client.KV().Get(c.Path, nil)
	if err != nil {
		return nil, err
	}
	if pair == nil {
		return nil, nil
	}

	c.modifyIndex = pair.ModifyIndex

	payload := pair.Value
	// If the payload starts with 0x1f, it's gzip, not json
	if len(pair.Value) >= 1 && pair.Value[0] == '\x1f' {
		if data, err := uncompressState(pair.Value); err == nil {
			payload = data
		} else {
			return nil, err
		}
	}

	md5 := md5.Sum(pair.Value)
	return &remote.Payload{
		Data: payload,
		MD5:  md5[:],
	}, nil
}

func (c *RemoteClient) Put(data []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	payload := data
	if c.GZip {
		if compressedState, err := compressState(data); err == nil {
			payload = compressedState
		} else {
			return err
		}
	}

	kv := c.Client.KV()

	verb := consulapi.KVCAS

	// Assume a 0 index doesn't need a CAS for now, since we are either
	// creating a new state or purposely overwriting one.
	if c.modifyIndex == 0 {
		verb = consulapi.KVSet
	}

	// KV.Put doesn't return the new index, so we use a single operation
	// transaction to get the new index with a single request.
	txOps := consulapi.KVTxnOps{
		&consulapi.KVTxnOp{
			Verb:  verb,
			Key:   c.Path,
			Value: payload,
			Index: c.modifyIndex,
		},
	}

	ok, resp, _, err := kv.Txn(txOps, nil)
	if err != nil {
		return err
	}

	// transaction was rolled back
	if !ok {
		return fmt.Errorf("consul CAS failed with transaction errors: %v", resp.Errors)
	}

	if len(resp.Results) != 1 {
		// this probably shouldn't happen
		return fmt.Errorf("expected on 1 response value, got: %d", len(resp.Results))
	}

	c.modifyIndex = resp.Results[0].ModifyIndex

	return nil
}

func (c *RemoteClient) Delete() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	kv := c.Client.KV()
	_, err := kv.Delete(c.Path, nil)
	return err
}

func (c *RemoteClient) putLockInfo(info *state.LockInfo) error {
	info.Path = c.Path
	info.Created = time.Now().UTC()

	kv := c.Client.KV()
	_, err := kv.Put(&consulapi.KVPair{
		Key:   c.Path + lockInfoSuffix,
		Value: info.Marshal(),
	}, nil)

	return err
}

func (c *RemoteClient) getLockInfo() (*state.LockInfo, error) {
	path := c.Path + lockInfoSuffix
	pair, _, err := c.Client.KV().Get(path, nil)
	if err != nil {
		return nil, err
	}
	if pair == nil {
		return nil, nil
	}

	li := &state.LockInfo{}
	err = json.Unmarshal(pair.Value, li)
	if err != nil {
		return nil, fmt.Errorf("error unmarshaling lock info: %s", err)
	}

	return li, nil
}

func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	select {
	case <-c.lockCh:
		// We had a lock, but lost it.
		// Since we typically only call lock once, we shouldn't ever see this.
		return "", errors.New("lost consul lock")
	default:
		if c.lockCh != nil {
			// we have an active lock already
			return "", fmt.Errorf("state %q already locked", c.Path)
		}
	}

	return c.lock(info)
}

func (c *RemoteClient) lock(info *state.LockInfo) (string, error) {
	if c.consulLock == nil {
		opts := &consulapi.LockOptions{
			Key: c.Path + lockSuffix,
			// only wait briefly, so terraform has the choice to fail fast or
			// retry as needed.
			LockWaitTime: time.Second,
			LockTryOnce:  true,
		}

		lock, err := c.Client.LockOpts(opts)
		if err != nil {
			return "", err
		}

		c.consulLock = lock
	}

	lockErr := &state.LockError{}

	lockCh, err := c.consulLock.Lock(make(chan struct{}))
	if err != nil {
		lockErr.Err = err
		return "", lockErr
	}

	if lockCh == nil {
		lockInfo, e := c.getLockInfo()
		if e != nil {
			lockErr.Err = e
			return "", lockErr
		}

		lockErr.Info = lockInfo
		return "", lockErr
	}

	c.lockCh = lockCh

	err = c.putLockInfo(info)
	if err != nil {
		if unlockErr := c.Unlock(info.ID); unlockErr != nil {
			err = multierror.Append(err, unlockErr)
		}

		return "", err
	}

	return info.ID, nil
}

func (c *RemoteClient) Unlock(id string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	// this doesn't use the lock id, because the lock is tied to the consul client.
	if c.consulLock == nil || c.lockCh == nil {
		return nil
	}

	select {
	case <-c.lockCh:
		return errors.New("consul lock was lost")
	default:
	}

	err := c.consulLock.Unlock()
	c.lockCh = nil

	// This is only cleanup, and will fail if the lock was immediately taken by
	// another client, so we don't report an error to the user here.
	c.consulLock.Destroy()

	kv := c.Client.KV()
	_, delErr := kv.Delete(c.Path+lockInfoSuffix, nil)
	if delErr != nil {
		err = multierror.Append(err, delErr)
	}

	return err
}

func compressState(data []byte) ([]byte, error) {
	b := new(bytes.Buffer)
	gz := gzip.NewWriter(b)
	if _, err := gz.Write(data); err != nil {
		return nil, err
	}
	if err := gz.Flush(); err != nil {
		return nil, err
	}
	if err := gz.Close(); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func uncompressState(data []byte) ([]byte, error) {
	b := new(bytes.Buffer)
	gz, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	b.ReadFrom(gz)
	if err := gz.Close(); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}
<|endoftext|>