<commit_before>package provision

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/gofn/gofn/iaas"
	uuid "github.com/satori/go.uuid"
)

var (
	// ErrImageNotFound is raised when image is not found
	ErrImageNotFound = errors.New("provision: image not found")

	// ErrContainerNotFound is raised when image is not found
	ErrContainerNotFound = errors.New("provision: container not found")

	// ErrContainerExecutionFailed is raised if container exited with status different of zero
	ErrContainerExecutionFailed = errors.New("provision: container exited with failure")

	// Input receives a string that will be written to the stdin of the container in function FnRun
	Input string
)

// BuildOptions are options used in the image build
type BuildOptions struct {
	ContextDir              string
	Dockerfile              string
	DoNotUsePrefixImageName bool
	ImageName               string
	RemoteURI               string
	StdIN                   string
	Iaas                    iaas.Iaas
}

// ContainerOptions are options used in container
type ContainerOptions struct {
	Cmd     []string
	Volumes []string
	Image   string
	Env     []string
}

// GetImageName sets preffix gofn when needed
func (opts BuildOptions) GetImageName() string {
	if opts.DoNotUsePrefixImageName {
		return opts.ImageName
	}
	return "gofn/" + opts.ImageName
}

// FnClient instantiate a docker client
func FnClient(endPoint, certsDir string) (client *docker.Client, err error) {
	if endPoint == "" {
		endPoint = "unix:///var/run/docker.sock"
	}

	client, err = docker.NewTLSClient(endPoint, certsDir+"/cert.pem", certsDir+"/key.pem", certsDir+"/ca.pem")
	return
}

// FnRemove remove container
func FnRemove(client *docker.Client, containerID string) (err error) {
	err = client.RemoveContainer(docker.RemoveContainerOptions{ID: containerID, Force: true})
	return
}

// FnContainer create container
func FnContainer(client *docker.Client, opts ContainerOptions) (container *docker.Container, err error) {
	config := &docker.Config{
		Image:     opts.Image,
		Cmd:       opts.Cmd,
		Env:       opts.Env,
		StdinOnce: true,
		OpenStdin: true,
	}
	var uid uuid.UUID
	uid, err = uuid.NewV4()
	if err != nil {
		return
	}
	container, err = client.CreateContainer(docker.CreateContainerOptions{
		Name:       fmt.Sprintf("gofn-%s", uid.String()),
		HostConfig: &docker.HostConfig{Binds: opts.Volumes},
		Config:     config,
	})
	return
}

// FnImageBuild builds an image
func FnImageBuild(client *docker.Client, opts *BuildOptions) (Name string, Stdout *bytes.Buffer, err error) {
	if opts.Dockerfile == "" {
		opts.Dockerfile = "Dockerfile"
	}
	stdout := new(bytes.Buffer)
	Name = opts.GetImageName()
	err = client.BuildImage(docker.BuildImageOptions{
		Name:           Name,
		Dockerfile:     opts.Dockerfile,
		SuppressOutput: true,
		OutputStream:   stdout,
		ContextDir:     opts.ContextDir,
		Remote:         opts.RemoteURI,
	})
	if err != nil {
		return
	}
	Stdout = stdout
	return
}

// FnFindImage returns image data by name
func FnFindImage(client *docker.Client, imageName string) (image docker.APIImages, err error) {
	var imgs []docker.APIImages
	imgs, err = client.ListImages(docker.ListImagesOptions{Filter: imageName})
	if err != nil {
		return
	}
	if len(imgs) == 0 {
		err = ErrImageNotFound
		return
	}
	image = imgs[0]
	return
}

// FnFindContainerByID return container by ID
func FnFindContainerByID(client *docker.Client, ID string) (container docker.APIContainers, err error) {
	var containers []docker.APIContainers
	containers, err = client.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		return
	}
	for _, v := range containers {
		if v.ID == ID {
			container = v
			return
		}
	}
	err = ErrContainerNotFound
	return
}

// FnFindContainer return container by image name
func FnFindContainer(client *docker.Client, imageName string) (container docker.APIContainers, err error) {
	var containers []docker.APIContainers
	containers, err = client.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		return
	}

	if !strings.HasPrefix(imageName, "gofn") {
		imageName = "gofn/" + imageName
	}

	for _, v := range containers {
		if v.Image == imageName {
			container = v
			return
		}
	}
	err = ErrContainerNotFound
	return
}

// FnKillContainer kill the container
func FnKillContainer(client *docker.Client, containerID string) (err error) {
	err = client.KillContainer(docker.KillContainerOptions{ID: containerID})
	return
}

//FnAttach attach into a running container
func FnAttach(client *docker.Client, containerID string, stdin io.Reader, stdout io.Writer, stderr io.Writer) (w docker.CloseWaiter, err error) {
	return client.AttachToContainerNonBlocking(docker.AttachToContainerOptions{
		Container:    containerID,
		RawTerminal:  true,
		Stream:       true,
		Stdin:        true,
		Stderr:       true,
		Stdout:       true,
		Logs:         true,
		InputStream:  stdin,
		ErrorStream:  stderr,
		OutputStream: stdout,
	})
}

// FnStart start the container
func FnStart(client *docker.Client, containerID string) error {
	return client.StartContainer(containerID, nil)
}

// FnRun runs the container
func FnRun(client *docker.Client, containerID, input string) (Stdout *bytes.Buffer, Stderr *bytes.Buffer, err error) {
	err = FnStart(client, containerID)
	if err != nil {
		return
	}

	// attach to write input
	_, err = FnAttach(client, containerID, strings.NewReader(input), nil, nil)
	if err != nil {
		return
	}

	e := FnWaitContainer(client, containerID)
	err = <-e

	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	// omit logs because execution error is more important
	_ = FnLogs(client, containerID, stdout, stderr)

	Stdout = stdout
	Stderr = stderr
	return
}

// FnLogs logs all container activity
func FnLogs(client *docker.Client, containerID string, stdout io.Writer, stderr io.Writer) error {
	return client.Logs(docker.LogsOptions{
		Container:    containerID,
		Stdout:       true,
		Stderr:       true,
		ErrorStream:  stderr,
		OutputStream: stdout,
	})
}

// FnWaitContainer wait until container finnish your processing
func FnWaitContainer(client *docker.Client, containerID string) chan error {
	errs := make(chan error)
	go func() {
		code, err := client.WaitContainer(containerID)
		if err != nil {
			errs <- err
		}
		if code != 0 {
			errs <- ErrContainerExecutionFailed
		}
		errs <- nil
	}()
	return errs
}

// FnListContainers lists all the containers created by the gofn.
// It returns the APIContainers from the API, but have to be formatted for pretty printing
func FnListContainers(client *docker.Client) (containers []docker.APIContainers, err error) {
	hostContainers, err := client.ListContainers(docker.ListContainersOptions{
		All: true,
	})
	if err != nil {
		containers = nil
		return
	}
	for _, container := range hostContainers {
		if strings.HasPrefix(container.Image, "gofn/") {
			containers = append(containers, container)
		}
	}
	return
}
<commit_msg>Add Runtime flag on creation containers<commit_after>package provision

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/gofn/gofn/iaas"
	uuid "github.com/satori/go.uuid"
)

var (
	// ErrImageNotFound is raised when image is not found
	ErrImageNotFound = errors.New("provision: image not found")

	// ErrContainerNotFound is raised when image is not found
	ErrContainerNotFound = errors.New("provision: container not found")

	// ErrContainerExecutionFailed is raised if container exited with status different of zero
	ErrContainerExecutionFailed = errors.New("provision: container exited with failure")

	// Input receives a string that will be written to the stdin of the container in function FnRun
	Input string
)

// BuildOptions are options used in the image build
type BuildOptions struct {
	ContextDir              string
	Dockerfile              string
	DoNotUsePrefixImageName bool
	ImageName               string
	RemoteURI               string
	StdIN                   string
	Iaas                    iaas.Iaas
}

// ContainerOptions are options used in container
type ContainerOptions struct {
	Cmd     []string
	Volumes []string
	Image   string
	Env     []string
	Runtime string
}

// GetImageName sets preffix gofn when needed
func (opts BuildOptions) GetImageName() string {
	if opts.DoNotUsePrefixImageName {
		return opts.ImageName
	}
	return "gofn/" + opts.ImageName
}

// FnClient instantiate a docker client
func FnClient(endPoint, certsDir string) (client *docker.Client, err error) {
	if endPoint == "" {
		endPoint = "unix:///var/run/docker.sock"
	}

	client, err = docker.NewTLSClient(endPoint, certsDir+"/cert.pem", certsDir+"/key.pem", certsDir+"/ca.pem")
	return
}

// FnRemove remove container
func FnRemove(client *docker.Client, containerID string) (err error) {
	err = client.RemoveContainer(docker.RemoveContainerOptions{ID: containerID, Force: true})
	return
}

// FnContainer create container
func FnContainer(client *docker.Client, opts ContainerOptions) (container *docker.Container, err error) {
	config := &docker.Config{
		Image:     opts.Image,
		Cmd:       opts.Cmd,
		Env:       opts.Env,
		StdinOnce: true,
		OpenStdin: true,
	}
	var uid uuid.UUID
	uid, err = uuid.NewV4()
	if err != nil {
		return
	}
	container, err = client.CreateContainer(docker.CreateContainerOptions{
		Name:       fmt.Sprintf("gofn-%s", uid.String()),
		HostConfig: &docker.HostConfig{Binds: opts.Volumes, Runtime: opts.Runtime},
		Config:     config,
	})
	return
}

// FnImageBuild builds an image
func FnImageBuild(client *docker.Client, opts *BuildOptions) (Name string, Stdout *bytes.Buffer, err error) {
	if opts.Dockerfile == "" {
		opts.Dockerfile = "Dockerfile"
	}
	stdout := new(bytes.Buffer)
	Name = opts.GetImageName()
	err = client.BuildImage(docker.BuildImageOptions{
		Name:           Name,
		Dockerfile:     opts.Dockerfile,
		SuppressOutput: true,
		OutputStream:   stdout,
		ContextDir:     opts.ContextDir,
		Remote:         opts.RemoteURI,
	})
	if err != nil {
		return
	}
	Stdout = stdout
	return
}

// FnFindImage returns image data by name
func FnFindImage(client *docker.Client, imageName string) (image docker.APIImages, err error) {
	var imgs []docker.APIImages
	imgs, err = client.ListImages(docker.ListImagesOptions{Filter: imageName})
	if err != nil {
		return
	}
	if len(imgs) == 0 {
		err = ErrImageNotFound
		return
	}
	image = imgs[0]
	return
}

// FnFindContainerByID return container by ID
func FnFindContainerByID(client *docker.Client, ID string) (container docker.APIContainers, err error) {
	var containers []docker.APIContainers
	containers, err = client.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		return
	}
	for _, v := range containers {
		if v.ID == ID {
			container = v
			return
		}
	}
	err = ErrContainerNotFound
	return
}

// FnFindContainer return container by image name
func FnFindContainer(client *docker.Client, imageName string) (container docker.APIContainers, err error) {
	var containers []docker.APIContainers
	containers, err = client.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		return
	}

	if !strings.HasPrefix(imageName, "gofn") {
		imageName = "gofn/" + imageName
	}

	for _, v := range containers {
		if v.Image == imageName {
			container = v
			return
		}
	}
	err = ErrContainerNotFound
	return
}

// FnKillContainer kill the container
func FnKillContainer(client *docker.Client, containerID string) (err error) {
	err = client.KillContainer(docker.KillContainerOptions{ID: containerID})
	return
}

//FnAttach attach into a running container
func FnAttach(client *docker.Client, containerID string, stdin io.Reader, stdout io.Writer, stderr io.Writer) (w docker.CloseWaiter, err error) {
	return client.AttachToContainerNonBlocking(docker.AttachToContainerOptions{
		Container:    containerID,
		RawTerminal:  true,
		Stream:       true,
		Stdin:        true,
		Stderr:       true,
		Stdout:       true,
		Logs:         true,
		InputStream:  stdin,
		ErrorStream:  stderr,
		OutputStream: stdout,
	})
}

// FnStart start the container
func FnStart(client *docker.Client, containerID string) error {
	return client.StartContainer(containerID, nil)
}

// FnRun runs the container
func FnRun(client *docker.Client, containerID, input string) (Stdout *bytes.Buffer, Stderr *bytes.Buffer, err error) {
	err = FnStart(client, containerID)
	if err != nil {
		return
	}

	// attach to write input
	_, err = FnAttach(client, containerID, strings.NewReader(input), nil, nil)
	if err != nil {
		return
	}

	e := FnWaitContainer(client, containerID)
	err = <-e

	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	// omit logs because execution error is more important
	_ = FnLogs(client, containerID, stdout, stderr)

	Stdout = stdout
	Stderr = stderr
	return
}

// FnLogs logs all container activity
func FnLogs(client *docker.Client, containerID string, stdout io.Writer, stderr io.Writer) error {
	return client.Logs(docker.LogsOptions{
		Container:    containerID,
		Stdout:       true,
		Stderr:       true,
		ErrorStream:  stderr,
		OutputStream: stdout,
	})
}

// FnWaitContainer wait until container finnish your processing
func FnWaitContainer(client *docker.Client, containerID string) chan error {
	errs := make(chan error)
	go func() {
		code, err := client.WaitContainer(containerID)
		if err != nil {
			errs <- err
		}
		if code != 0 {
			errs <- ErrContainerExecutionFailed
		}
		errs <- nil
	}()
	return errs
}

// FnListContainers lists all the containers created by the gofn.
// It returns the APIContainers from the API, but have to be formatted for pretty printing
func FnListContainers(client *docker.Client) (containers []docker.APIContainers, err error) {
	hostContainers, err := client.ListContainers(docker.ListContainersOptions{
		All: true,
	})
	if err != nil {
		containers = nil
		return
	}
	for _, container := range hostContainers {
		if strings.HasPrefix(container.Image, "gofn/") {
			containers = append(containers, container)
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"

	"github.com/pkg/errors"
)

const (
	VulcandKeyBase = "/vulcand"
)

var (
	httpBackendJSON string
)

type Vulcand struct {
	etcd *Etcd
}

type VulcandBackend struct {
	Type string `json:"Type"`
}

type VulcandServer struct {
	URL string `json:"URL"`
}

type VulcandFrontend struct {
	Type      string                  `json:"Type"`
	BackendId string                  `json:"BackendId"`
	Route     string                  `json:"Route"`
	Settings  VulcandFrontendSettings `json:"Settings"`
}

type VulcandFrontendSettings struct {
	TrustForwardHeader bool `json:"TrustForwardHeader"`
}

func NewVulcand(etcd *Etcd) *Vulcand {
	backend := VulcandBackend{
		Type: "http",
	}

	b, _ := json.Marshal(backend)
	httpBackendJSON = string(b)

	return &Vulcand{
		etcd: etcd,
	}
}

// {"Type": "http"}
func (v *Vulcand) SetBackend(application *Application, baseDomain string) error {
	key := fmt.Sprintf("%s/backends/%s/backend", VulcandKeyBase, application.ProjectName)

	if err := v.etcd.Set(key, httpBackendJSON); err != nil {
		return errors.Wrap(err, "Failed to set vulcand backend in etcd.")
	}

	return nil
}

// {"Type": "http", "BackendId": "$identifier", "Route": "Host(`$identifier.$base_domain`) && PathRegexp(`/`)", "Settings": {"TrustForwardHeader": true}}
func (v *Vulcand) SetFrontend(application *Application, identifier, baseDomain string) error {
	key := fmt.Sprintf("%s/frontends/%s/frontend", VulcandKeyBase, identifier)
	frontend := VulcandFrontend{
		Type:      "http",
		BackendId: application.ProjectName,
		Route:     fmt.Sprintf("Host(`%s.%s`) && PathRegexp(`/`)", identifier, baseDomain),
		Settings: VulcandFrontendSettings{
			TrustForwardHeader: true,
		},
	}

	b, err := json.Marshal(frontend)

	if err != nil {
		return errors.Wrap(err, "Failed to generate vulcand frontend JSON.")
	}

	json := string(b)

	if err := v.etcd.Set(key, json); err != nil {
		return errors.Wrap(err, "Failed to set vulcand frontend in etcd.")
	}

	return nil
}

// {"URL": "http://$web_container_host_ip:$web_container_port"}
func (v *Vulcand) SetServer(application *Application, container *Container, baseDomain string) error {
	key := fmt.Sprintf("%s/backends/%s/servers/%s", VulcandKeyBase, application.ProjectName, container.ContainerId)
	server := VulcandServer{
		URL: fmt.Sprintf("http://%s:%s", container.HostIP(), container.HostPort()),
	}

	b, err := json.Marshal(server)

	if err != nil {
		return errors.Wrap(err, "Failed to generate vulcand server JSON.")
	}

	json := string(b)

	if err := v.etcd.Set(key, json); err != nil {
		return errors.Wrap(err, "Failed to set vulcand server in etcd.")
	}

	return nil
}
<commit_msg>Unescape & in vulcand JSON<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/pkg/errors"
)

const (
	VulcandKeyBase = "/vulcand"
)

var (
	httpBackendJSON string
)

type Vulcand struct {
	etcd *Etcd
}

type VulcandBackend struct {
	Type string `json:"Type"`
}

type VulcandServer struct {
	URL string `json:"URL"`
}

type VulcandFrontend struct {
	Type      string                  `json:"Type"`
	BackendId string                  `json:"BackendId"`
	Route     string                  `json:"Route"`
	Settings  VulcandFrontendSettings `json:"Settings"`
}

type VulcandFrontendSettings struct {
	TrustForwardHeader bool `json:"TrustForwardHeader"`
}

func NewVulcand(etcd *Etcd) *Vulcand {
	backend := VulcandBackend{
		Type: "http",
	}

	b, _ := json.Marshal(backend)
	httpBackendJSON = string(b)

	return &Vulcand{
		etcd: etcd,
	}
}

// {"Type": "http"}
func (v *Vulcand) SetBackend(application *Application, baseDomain string) error {
	key := fmt.Sprintf("%s/backends/%s/backend", VulcandKeyBase, application.ProjectName)

	if err := v.etcd.Set(key, httpBackendJSON); err != nil {
		return errors.Wrap(err, "Failed to set vulcand backend in etcd.")
	}

	return nil
}

// {"Type": "http", "BackendId": "$identifier", "Route": "Host(`$identifier.$base_domain`) && PathRegexp(`/`)", "Settings": {"TrustForwardHeader": true}}
func (v *Vulcand) SetFrontend(application *Application, identifier, baseDomain string) error {
	key := fmt.Sprintf("%s/frontends/%s/frontend", VulcandKeyBase, identifier)
	frontend := VulcandFrontend{
		Type:      "http",
		BackendId: application.ProjectName,
		Route:     fmt.Sprintf("Host(`%s.%s`) && PathRegexp(`/`)", identifier, baseDomain),
		Settings: VulcandFrontendSettings{
			TrustForwardHeader: true,
		},
	}

	b, err := json.Marshal(frontend)

	if err != nil {
		return errors.Wrap(err, "Failed to generate vulcand frontend JSON.")
	}

	// json.Marshal generates HTML-escaped JSON string
	b = bytes.Replace(b, []byte("\\u0026"), []byte("&"), -1)
	json := string(b)

	if err := v.etcd.Set(key, json); err != nil {
		return errors.Wrap(err, "Failed to set vulcand frontend in etcd.")
	}

	return nil
}

// {"URL": "http://$web_container_host_ip:$web_container_port"}
func (v *Vulcand) SetServer(application *Application, container *Container, baseDomain string) error {
	key := fmt.Sprintf("%s/backends/%s/servers/%s", VulcandKeyBase, application.ProjectName, container.ContainerId)
	server := VulcandServer{
		URL: fmt.Sprintf("http://%s:%s", container.HostIP(), container.HostPort()),
	}

	b, err := json.Marshal(server)

	if err != nil {
		return errors.Wrap(err, "Failed to generate vulcand server JSON.")
	}

	json := string(b)

	if err := v.etcd.Set(key, json); err != nil {
		return errors.Wrap(err, "Failed to set vulcand server in etcd.")
	}

	return nil
}
<|endoftext|>


<commit_before>package registry

import (
	"github.com/docker/docker/engine"
)

// Service exposes registry capabilities in the standard Engine
// interface. Once installed, it extends the engine with the
// following calls:
//
//  'auth': Authenticate against the public registry
//  'search': Search for images on the public registry
//  'pull': Download images from any registry (TODO)
//  'push': Upload images to any registry (TODO)
type Service struct {
}

// NewService returns a new instance of Service ready to be
// installed no an engine.
func NewService() *Service {
	return &Service{}
}

// Install installs registry capabilities to eng.
func (s *Service) Install(eng *engine.Engine) error {
	eng.Register("auth", s.Auth)
	eng.Register("search", s.Search)
	return nil
}

// Auth contacts the public registry with the provided credentials,
// and returns OK if authentication was sucessful.
// It can be used to verify the validity of a client's credentials.
func (s *Service) Auth(job *engine.Job) engine.Status {
	var (
		err        error
		authConfig = &AuthConfig{}
	)

	job.GetenvJson("authConfig", authConfig)
	// TODO: this is only done here because auth and registry need to be merged into one pkg
	if addr := authConfig.ServerAddress; addr != "" && addr != IndexServerAddress() {
		addr, err = ExpandAndVerifyRegistryUrl(addr)
		if err != nil {
			return job.Error(err)
		}
		authConfig.ServerAddress = addr
	}
	status, err := Login(authConfig, HTTPRequestFactory(nil))
	if err != nil {
		return job.Error(err)
	}
	job.Printf("%s\n", status)
	return engine.StatusOK
}

// Search queries the public registry for images matching the specified
// search terms, and returns the results.
//
// Argument syntax: search TERM
//
// Option environment:
//	'authConfig': json-encoded credentials to authenticate against the registry.
//		The search extends to images only accessible via the credentials.
//
//	'metaHeaders': extra HTTP headers to include in the request to the registry.
//		The headers should be passed as a json-encoded dictionary.
//
// Output:
//	Results are sent as a collection of structured messages (using engine.Table).
//	Each result is sent as a separate message.
//	Results are ordered by number of stars on the public registry.
func (s *Service) Search(job *engine.Job) engine.Status {
	if n := len(job.Args); n != 1 {
		return job.Errorf("Usage: %s TERM", job.Name)
	}
	var (
		term        = job.Args[0]
		metaHeaders = map[string][]string{}
		authConfig  = &AuthConfig{}
	)
	job.GetenvJson("authConfig", authConfig)
	job.GetenvJson("metaHeaders", metaHeaders)

	hostname, term, err := ResolveRepositoryName(term)
	if err != nil {
		return job.Error(err)
	}
	hostname, err = ExpandAndVerifyRegistryUrl(hostname)
	if err != nil {
		return job.Error(err)
	}
	r, err := NewSession(authConfig, HTTPRequestFactory(metaHeaders), IndexServerAddress(), true)
	if err != nil {
		return job.Error(err)
	}
	results, err := r.SearchRepositories(term)
	if err != nil {
		return job.Error(err)
	}
	outs := engine.NewTable("star_count", 0)
	for _, result := range results.Results {
		out := &engine.Env{}
		out.Import(result)
		outs.Add(out)
	}
	outs.ReverseSort()
	if _, err := outs.WriteListTo(job.Stdout); err != nil {
		return job.Error(err)
	}
	return engine.StatusOK
}
<commit_msg>Expand hostname before passing it to NewRegistry()<commit_after>package registry

import (
	"github.com/docker/docker/engine"
)

// Service exposes registry capabilities in the standard Engine
// interface. Once installed, it extends the engine with the
// following calls:
//
//  'auth': Authenticate against the public registry
//  'search': Search for images on the public registry
//  'pull': Download images from any registry (TODO)
//  'push': Upload images to any registry (TODO)
type Service struct {
}

// NewService returns a new instance of Service ready to be
// installed no an engine.
func NewService() *Service {
	return &Service{}
}

// Install installs registry capabilities to eng.
func (s *Service) Install(eng *engine.Engine) error {
	eng.Register("auth", s.Auth)
	eng.Register("search", s.Search)
	return nil
}

// Auth contacts the public registry with the provided credentials,
// and returns OK if authentication was sucessful.
// It can be used to verify the validity of a client's credentials.
func (s *Service) Auth(job *engine.Job) engine.Status {
	var (
		err        error
		authConfig = &AuthConfig{}
	)

	job.GetenvJson("authConfig", authConfig)
	// TODO: this is only done here because auth and registry need to be merged into one pkg
	if addr := authConfig.ServerAddress; addr != "" && addr != IndexServerAddress() {
		addr, err = ExpandAndVerifyRegistryUrl(addr)
		if err != nil {
			return job.Error(err)
		}
		authConfig.ServerAddress = addr
	}
	status, err := Login(authConfig, HTTPRequestFactory(nil))
	if err != nil {
		return job.Error(err)
	}
	job.Printf("%s\n", status)
	return engine.StatusOK
}

// Search queries the public registry for images matching the specified
// search terms, and returns the results.
//
// Argument syntax: search TERM
//
// Option environment:
//	'authConfig': json-encoded credentials to authenticate against the registry.
//		The search extends to images only accessible via the credentials.
//
//	'metaHeaders': extra HTTP headers to include in the request to the registry.
//		The headers should be passed as a json-encoded dictionary.
//
// Output:
//	Results are sent as a collection of structured messages (using engine.Table).
//	Each result is sent as a separate message.
//	Results are ordered by number of stars on the public registry.
func (s *Service) Search(job *engine.Job) engine.Status {
	if n := len(job.Args); n != 1 {
		return job.Errorf("Usage: %s TERM", job.Name)
	}
	var (
		term        = job.Args[0]
		metaHeaders = map[string][]string{}
		authConfig  = &AuthConfig{}
	)
	job.GetenvJson("authConfig", authConfig)
	job.GetenvJson("metaHeaders", metaHeaders)

	hostname, term, err := ResolveRepositoryName(term)
	if err != nil {
		return job.Error(err)
	}
	hostname, err = ExpandAndVerifyRegistryUrl(hostname)
	if err != nil {
		return job.Error(err)
	}
	r, err := NewSession(authConfig, HTTPRequestFactory(metaHeaders), hostname, true)
	if err != nil {
		return job.Error(err)
	}
	results, err := r.SearchRepositories(term)
	if err != nil {
		return job.Error(err)
	}
	outs := engine.NewTable("star_count", 0)
	for _, result := range results.Results {
		out := &engine.Env{}
		out.Import(result)
		outs.Add(out)
	}
	outs.ReverseSort()
	if _, err := outs.WriteListTo(job.Stdout); err != nil {
		return job.Error(err)
	}
	return engine.StatusOK
}
<|endoftext|>


<commit_before>// Runs a daemon to continuously warm the registry cache.
package registry

import (
	"context"
	"encoding/json"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/pkg/errors"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/registry/cache"
)

const refreshWhenExpiryWithin = time.Minute
const askForNewImagesInterval = time.Minute

type Warmer struct {
	Logger        log.Logger
	ClientFactory ClientFactory
	Creds         Credentials
	Expiry        time.Duration
	Writer        cache.Writer
	Reader        cache.Reader
	Burst         int
}

type ImageCreds map[image.Name]Credentials

// Continuously get the images to populate the cache with, and
// populate the cache with them.
func (w *Warmer) Loop(stop <-chan struct{}, wg *sync.WaitGroup, imagesToFetchFunc func() ImageCreds) {
	defer wg.Done()

	if w.Logger == nil || w.ClientFactory == nil || w.Expiry == 0 || w.Writer == nil || w.Reader == nil {
		panic("registry.Warmer fields are nil")
	}

	for k, v := range imagesToFetchFunc() {
		w.warm(k, v)
	}

	newImages := time.Tick(askForNewImagesInterval)
	for {
		select {
		case <-stop:
			w.Logger.Log("stopping", "true")
			return
		case <-newImages:
			for k, v := range imagesToFetchFunc() {
				w.warm(k, v)
			}
		}
	}
}

func (w *Warmer) warm(id image.Name, creds Credentials) {
	client, err := w.ClientFactory.ClientFor(id.Registry(), creds)
	if err != nil {
		w.Logger.Log("err", err.Error())
		return
	}
	defer client.Cancel()

	username := w.Creds.credsFor(id.Registry()).username

	// Refresh tags first
	// Only, for example, "library/alpine" because we have the host information in the client above.
	tags, err := client.Tags(id)
	if err != nil {
		if !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) && !strings.Contains(err.Error(), "net/http: request canceled") {
			w.Logger.Log("err", errors.Wrap(err, "requesting tags"))
		}
		return
	}

	val, err := json.Marshal(tags)
	if err != nil {
		w.Logger.Log("err", errors.Wrap(err, "serializing tags to store in cache"))
		return
	}

	key, err := cache.NewTagKey(username, id.CanonicalName())
	if err != nil {
		w.Logger.Log("err", errors.Wrap(err, "creating key for cache"))
		return
	}

	err = w.Writer.SetKey(key, val)
	if err != nil {
		w.Logger.Log("err", errors.Wrap(err, "storing tags in cache"))
		return
	}

	// Create a list of manifests that need updating
	var toUpdate []image.Ref
	var expired bool
	for _, tag := range tags {
		// See if we have the manifest already cached
		// We don't want to re-download a manifest again.
		newID := id.ToRef(tag)
		key, err := cache.NewManifestKey(username, newID.CanonicalRef())
		if err != nil {
			w.Logger.Log("err", errors.Wrap(err, "creating key for memcache"))
			continue
		}
		expiry, err := w.Reader.GetExpiration(key)
		// If err, then we don't have it yet. Update.
		if err == nil { // If no error, we've already got it
			// If we're outside of the expiry buffer, skip, no need to update.
			if !withinExpiryBuffer(expiry, refreshWhenExpiryWithin) {
				continue
			}
			// If we're within the expiry buffer, we need to update quick!
			expired = true
		}
		toUpdate = append(toUpdate, newID)
	}

	if len(toUpdate) == 0 {
		return
	}
	w.Logger.Log("fetching", id.String(), "to-update", len(toUpdate))

	if expired {
		w.Logger.Log("expiring", id.String())
	}

	// The upper bound for concurrent fetches against a single host is
	// w.Burst, so limit the number of fetching goroutines to that.
	fetchers := make(chan struct{}, w.Burst)
	awaitFetchers := &sync.WaitGroup{}
	for _, imID := range toUpdate {
		awaitFetchers.Add(1)
		fetchers <- struct{}{}
		go func(imageID image.Ref) {
			defer func() { awaitFetchers.Done(); <-fetchers }()
			// Get the image from the remote
			img, err := client.Manifest(imageID)
			if err != nil {
				if err, ok := errors.Cause(err).(net.Error); ok && err.Timeout() {
					// This was due to a context timeout, don't bother logging
					return
				}
				w.Logger.Log("err", errors.Wrap(err, "requesting manifests"))
				return
			}

			key, err := cache.NewManifestKey(username, img.ID.CanonicalRef())
			if err != nil {
				w.Logger.Log("err", errors.Wrap(err, "creating key for memcache"))
				return
			}
			// Write back to memcache
			val, err := json.Marshal(img)
			if err != nil {
				w.Logger.Log("err", errors.Wrap(err, "serializing tag to store in cache"))
				return
			}
			err = w.Writer.SetKey(key, val)
			if err != nil {
				w.Logger.Log("err", errors.Wrap(err, "storing manifests in cache"))
				return
			}
		}(imID)
	}
	awaitFetchers.Wait()
	w.Logger.Log("updated", id.String())
}

func withinExpiryBuffer(expiry time.Time, buffer time.Duration) bool {
	// if the `time.Now() + buffer  > expiry`,
	// then we're within the expiry buffer
	if time.Now().Add(buffer).After(expiry) {
		return true
	}
	return false
}
<commit_msg>Give the image cache warmer a priority channel<commit_after>// Runs a daemon to continuously warm the registry cache.
package registry

import (
	"context"
	"encoding/json"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/pkg/errors"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/registry/cache"
)

const refreshWhenExpiryWithin = time.Minute
const askForNewImagesInterval = time.Minute

type Warmer struct {
	Logger        log.Logger
	ClientFactory ClientFactory
	Creds         Credentials
	Expiry        time.Duration
	Writer        cache.Writer
	Reader        cache.Reader
	Burst         int
	Priority      chan image.Name
}

type ImageCreds map[image.Name]Credentials

// Continuously get the images to populate the cache with, and
// populate the cache with them.
func (w *Warmer) Loop(stop <-chan struct{}, wg *sync.WaitGroup, imagesToFetchFunc func() ImageCreds) {
	defer wg.Done()

	if w.Logger == nil || w.ClientFactory == nil || w.Expiry == 0 || w.Writer == nil || w.Reader == nil {
		panic("registry.Warmer fields are nil")
	}

	type imageCred struct {
		image.Name
		Credentials
	}

	refresh := time.Tick(askForNewImagesInterval)
	imageCreds := imagesToFetchFunc()
	backlog := []imageCred{}

	// This loop acts keeps a kind of priority queue, whereby image
	// names coming in on the `Priority` channel are looked up first.
	// If there are none, images used in the cluster are refreshed;
	// but no more often than once every `askForNewImagesInterval`,
	// since there is no effective back-pressure on cache refreshes
	// and it would spin freely otherwise).
	for {
		select {
		case <-stop:
			w.Logger.Log("stopping", "true")
			return
		case name := <-w.Priority:
			w.Logger.Log("priority", name.String())
			// NB the implicit contract here is that the prioritised
			// image has to have been running the last time we
			// requested the credentials.
			if creds, ok := imageCreds[name]; ok {
				w.warm(name, creds)
			} else {
				w.Logger.Log("priority", name.String(), "err", "no creds available")
			}
			continue
		default:
		}

		if len(backlog) > 0 {
			im := backlog[0]
			backlog = backlog[1:]
			w.warm(im.Name, im.Credentials)
		} else {
			select {
			case <-refresh:
				imageCreds = imagesToFetchFunc()
				for name, cred := range imageCreds {
					backlog = append(backlog, imageCred{name, cred})
				}
			default:
			}
		}
	}
}

func (w *Warmer) warm(id image.Name, creds Credentials) {
	client, err := w.ClientFactory.ClientFor(id.Registry(), creds)
	if err != nil {
		w.Logger.Log("err", err.Error())
		return
	}
	defer client.Cancel()

	username := w.Creds.credsFor(id.Registry()).username

	// Refresh tags first
	// Only, for example, "library/alpine" because we have the host information in the client above.
	tags, err := client.Tags(id)
	if err != nil {
		if !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) && !strings.Contains(err.Error(), "net/http: request canceled") {
			w.Logger.Log("err", errors.Wrap(err, "requesting tags"))
		}
		return
	}

	val, err := json.Marshal(tags)
	if err != nil {
		w.Logger.Log("err", errors.Wrap(err, "serializing tags to store in cache"))
		return
	}

	key, err := cache.NewTagKey(username, id.CanonicalName())
	if err != nil {
		w.Logger.Log("err", errors.Wrap(err, "creating key for cache"))
		return
	}

	err = w.Writer.SetKey(key, val)
	if err != nil {
		w.Logger.Log("err", errors.Wrap(err, "storing tags in cache"))
		return
	}

	// Create a list of manifests that need updating
	var toUpdate []image.Ref
	var expired bool
	for _, tag := range tags {
		// See if we have the manifest already cached
		// We don't want to re-download a manifest again.
		newID := id.ToRef(tag)
		key, err := cache.NewManifestKey(username, newID.CanonicalRef())
		if err != nil {
			w.Logger.Log("err", errors.Wrap(err, "creating key for memcache"))
			continue
		}
		expiry, err := w.Reader.GetExpiration(key)
		// If err, then we don't have it yet. Update.
		if err == nil { // If no error, we've already got it
			// If we're outside of the expiry buffer, skip, no need to update.
			if !withinExpiryBuffer(expiry, refreshWhenExpiryWithin) {
				continue
			}
			// If we're within the expiry buffer, we need to update quick!
			expired = true
		}
		toUpdate = append(toUpdate, newID)
	}

	if len(toUpdate) == 0 {
		return
	}
	w.Logger.Log("fetching", id.String(), "to-update", len(toUpdate))

	if expired {
		w.Logger.Log("expiring", id.String())
	}

	// The upper bound for concurrent fetches against a single host is
	// w.Burst, so limit the number of fetching goroutines to that.
	fetchers := make(chan struct{}, w.Burst)
	awaitFetchers := &sync.WaitGroup{}
	for _, imID := range toUpdate {
		awaitFetchers.Add(1)
		fetchers <- struct{}{}
		go func(imageID image.Ref) {
			defer func() { awaitFetchers.Done(); <-fetchers }()
			// Get the image from the remote
			img, err := client.Manifest(imageID)
			if err != nil {
				if err, ok := errors.Cause(err).(net.Error); ok && err.Timeout() {
					// This was due to a context timeout, don't bother logging
					return
				}
				w.Logger.Log("err", errors.Wrap(err, "requesting manifests"))
				return
			}

			key, err := cache.NewManifestKey(username, img.ID.CanonicalRef())
			if err != nil {
				w.Logger.Log("err", errors.Wrap(err, "creating key for memcache"))
				return
			}
			// Write back to memcache
			val, err := json.Marshal(img)
			if err != nil {
				w.Logger.Log("err", errors.Wrap(err, "serializing tag to store in cache"))
				return
			}
			err = w.Writer.SetKey(key, val)
			if err != nil {
				w.Logger.Log("err", errors.Wrap(err, "storing manifests in cache"))
				return
			}
		}(imID)
	}
	awaitFetchers.Wait()
	w.Logger.Log("updated", id.String())
}

func withinExpiryBuffer(expiry time.Time, buffer time.Duration) bool {
	// if the `time.Now() + buffer  > expiry`,
	// then we're within the expiry buffer
	if time.Now().Add(buffer).After(expiry) {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strconv"

	"github.com/hashicorp/terraform/flatmap"
	"github.com/hashicorp/terraform/helper/config"
	"github.com/hashicorp/terraform/helper/diff"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/elb"
)

func resource_aws_elb_create(
	s *terraform.InstanceState,
	d *terraform.InstanceDiff,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	// Merge the diff into the state so that we have all the attributes
	// properly.
	rs := s.MergeDiff(d)

	// The name specified for the ELB. This is also our unique ID
	// we save to state if the creation is successful (amazon verifies
	// it is unique)
	elbName := rs.Attributes["name"]

	// Expand the "listener" array to goamz compat []elb.Listener
	v := flatmap.Expand(rs.Attributes, "listener").([]interface{})
	listeners, err := expandListeners(v)
	if err != nil {
		return nil, err
	}

	// Provision the elb
	elbOpts := &elb.CreateLoadBalancer{
		LoadBalancerName: elbName,
		Listeners:        listeners,
	}

	if rs.Attributes["internal"] == "true" {
		elbOpts.Internal = true
	}

	if _, ok := rs.Attributes["availability_zones.#"]; ok {
		v = flatmap.Expand(rs.Attributes, "availability_zones").([]interface{})
		elbOpts.AvailZone = expandStringList(v)
	}

	if _, ok := rs.Attributes["security_groups.#"]; ok {
		v = flatmap.Expand(rs.Attributes, "security_groups").([]interface{})
		elbOpts.SecurityGroups = expandStringList(v)
	}

	if _, ok := rs.Attributes["subnets.#"]; ok {
		v = flatmap.Expand(rs.Attributes, "subnets").([]interface{})
		elbOpts.Subnets = expandStringList(v)
	}

	log.Printf("[DEBUG] ELB create configuration: %#v", elbOpts)

	_, err = elbconn.CreateLoadBalancer(elbOpts)
	if err != nil {
		return nil, fmt.Errorf("Error creating ELB: %s", err)
	}

	// Assign the elb's unique identifier for use later
	rs.ID = elbName
	log.Printf("[INFO] ELB ID: %s", elbName)

	if _, ok := rs.Attributes["instances.#"]; ok {
		// If we have any instances, we need to register them
		v = flatmap.Expand(rs.Attributes, "instances").([]interface{})
		instances := expandStringList(v)

		if len(instances) > 0 {
			registerInstancesOpts := elb.RegisterInstancesWithLoadBalancer{
				LoadBalancerName: elbName,
				Instances:        instances,
			}

			_, err := elbconn.RegisterInstancesWithLoadBalancer(&registerInstancesOpts)

			if err != nil {
				return rs, fmt.Errorf("Failure registering instances: %s", err)
			}
		}
	}

	if _, ok := rs.Attributes["health_check.#"]; ok {
		v := flatmap.Expand(rs.Attributes, "health_check").([]interface{})
		health_check := v[0].(map[string]interface{})
		healthyThreshold, err := strconv.ParseInt(health_check["healthy_threshold"].(string), 0, 0)
		unhealthyThreshold, err := strconv.ParseInt(health_check["unhealthy_threshold"].(string), 0, 0)
		interval, err := strconv.ParseInt(health_check["interval"].(string), 0, 0)
		timeout, err := strconv.ParseInt(health_check["timeout"].(string), 0, 0)

		if err != nil {
			return nil, err
		}

		configureHealthCheckOpts := elb.ConfigureHealthCheck{
			LoadBalancerName: elbName,
			Check: elb.HealthCheck{
				HealthyThreshold:   healthyThreshold,
				UnhealthyThreshold: unhealthyThreshold,
				Interval:           interval,
				Target:             health_check["target"].(string),
				Timeout:            timeout,
			},
		}

		_, err = elbconn.ConfigureHealthCheck(&configureHealthCheckOpts)
		if err != nil {
			return rs, fmt.Errorf("Failure configuring health check: %s", err)
		}
	}

	loadBalancer, err := resource_aws_elb_retrieve_balancer(rs.ID, elbconn)
	if err != nil {
		return rs, err
	}

	return resource_aws_elb_update_state(rs, loadBalancer)
}

func resource_aws_elb_update(
	s *terraform.InstanceState,
	d *terraform.InstanceDiff,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	rs := s.MergeDiff(d)

	// If we currently have instances, or did have instances,
	// we want to figure out what to add and remove from the load
	// balancer
	if attr, ok := d.Attributes["instances.#"]; ok && attr.Old != "" {
		// The new state of instances merged with the diff
		mergedInstances := expandStringList(flatmap.Expand(
			rs.Attributes, "instances").([]interface{}))

		// The state before the diff merge
		previousInstances := expandStringList(flatmap.Expand(
			s.Attributes, "instances").([]interface{}))

		// keep track of what instances we are removing, and which
		// we are adding
		var toRemove []string
		var toAdd []string

		for _, instanceId := range mergedInstances {
			for _, prevId := range previousInstances {
				// If the merged instance ID existed
				// previously, we don't have to do anything
				if instanceId == prevId {
					continue
					// Otherwise, we need to add it to the load balancer
				} else {
					toAdd = append(toAdd, instanceId)
				}
			}
		}

		for i, instanceId := range toAdd {
			for _, prevId := range previousInstances {
				// If the instance ID we are adding existed
				// previously, we want to not add it, but rather remove
				// it
				if instanceId == prevId {
					toRemove = append(toRemove, instanceId)
					toAdd = append(toAdd[:i], toAdd[i+1:]...)
					// Otherwise, we continue adding it to the ELB
				} else {
					continue
				}
			}
		}

		if len(toAdd) > 0 {
			registerInstancesOpts := elb.RegisterInstancesWithLoadBalancer{
				LoadBalancerName: rs.ID,
				Instances:        toAdd,
			}

			_, err := elbconn.RegisterInstancesWithLoadBalancer(&registerInstancesOpts)

			if err != nil {
				return s, fmt.Errorf("Failure registering instances: %s", err)
			}
		}

		if len(toRemove) > 0 {
			deRegisterInstancesOpts := elb.DeregisterInstancesFromLoadBalancer{
				LoadBalancerName: rs.ID,
				Instances:        toRemove,
			}

			_, err := elbconn.DeregisterInstancesFromLoadBalancer(&deRegisterInstancesOpts)

			if err != nil {
				return s, fmt.Errorf("Failure deregistering instances: %s", err)
			}
		}
	}

	loadBalancer, err := resource_aws_elb_retrieve_balancer(rs.ID, elbconn)

	if err != nil {
		return s, err
	}

	return resource_aws_elb_update_state(rs, loadBalancer)
}

func resource_aws_elb_destroy(
	s *terraform.InstanceState,
	meta interface{}) error {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	log.Printf("[INFO] Deleting ELB: %s", s.ID)

	// Destroy the load balancer
	deleteElbOpts := elb.DeleteLoadBalancer{
		LoadBalancerName: s.ID,
	}
	_, err := elbconn.DeleteLoadBalancer(&deleteElbOpts)

	if err != nil {
		return fmt.Errorf("Error deleting ELB: %s", err)
	}

	return nil
}

func resource_aws_elb_refresh(
	s *terraform.InstanceState,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	loadBalancer, err := resource_aws_elb_retrieve_balancer(s.ID, elbconn)
	if err != nil {
		return nil, err
	}

	return resource_aws_elb_update_state(s, loadBalancer)
}

func resource_aws_elb_diff(
	s *terraform.InstanceState,
	c *terraform.ResourceConfig,
	meta interface{}) (*terraform.InstanceDiff, error) {

	b := &diff.ResourceBuilder{
		Attrs: map[string]diff.AttrType{
			"name":              diff.AttrTypeCreate,
			"availability_zone": diff.AttrTypeCreate,
			"security_groups":   diff.AttrTypeCreate, // TODO could be AttrTypeUpdate
			"subnets":           diff.AttrTypeCreate, // TODO could be AttrTypeUpdate
			"listener":          diff.AttrTypeCreate,
			"instances":         diff.AttrTypeUpdate,
			"health_check":      diff.AttrTypeCreate,
			"internal":          diff.AttrTypeCreate,
		},

		ComputedAttrs: []string{
			"dns_name",
		},
	}

	return b.Diff(s, c)
}

func resource_aws_elb_update_state(
	s *terraform.InstanceState,
	balancer *elb.LoadBalancer) (*terraform.InstanceState, error) {

	s.Attributes["name"] = balancer.LoadBalancerName
	s.Attributes["dns_name"] = balancer.DNSName

	// Flatten our group values
	toFlatten := make(map[string]interface{})

	if len(balancer.Instances) > 0 && balancer.Instances[0].InstanceId != "" {
		toFlatten["instances"] = flattenInstances(balancer.Instances)
	}

	if len(balancer.SecurityGroups) > 0 && balancer.SecurityGroups[0] != "" {
		toFlatten["security_groups"] = balancer.SecurityGroups
	}

	if len(balancer.Subnets) > 0 && balancer.Subnets[0] != "" {
		toFlatten["subnets"] = balancer.Subnets
	}

	// There's only one health check, so save that to state as we
	// currently can
	if balancer.HealthCheck.Target != "" {
		toFlatten["health_check"] = flattenHealthCheck(balancer.HealthCheck)
	}

	for k, v := range flatmap.Flatten(toFlatten) {
		s.Attributes[k] = v
	}

	return s, nil
}

// retrieves an ELB by its ID
func resource_aws_elb_retrieve_balancer(id string, elbconn *elb.ELB) (*elb.LoadBalancer, error) {
	describeElbOpts := &elb.DescribeLoadBalancer{
		Names: []string{id},
	}

	// Retrieve the ELB properties for updating the state
	describeResp, err := elbconn.DescribeLoadBalancers(describeElbOpts)
	if err != nil {
		return nil, fmt.Errorf("Error retrieving ELB: %s", err)
	}

	loadBalancer := describeResp.LoadBalancers[0]

	// Verify AWS returned our ELB
	if len(describeResp.LoadBalancers) != 1 ||
		describeResp.LoadBalancers[0].LoadBalancerName != id {
		if err != nil {
			return nil, fmt.Errorf("Unable to find ELB: %#v", describeResp.LoadBalancers)
		}
	}

	return &loadBalancer, nil
}

func resource_aws_elb_validation() *config.Validator {
	return &config.Validator{
		Required: []string{
			"name",
			"listener.*",
			"listener.*.instance_port",
			"listener.*.instance_protocol",
			"listener.*.lb_port",
			"listener.*.lb_protocol",
		},
		Optional: []string{
			"instances.*",
			"internal",
			"availability_zones.*",
			"security_groups.*",
			"subnets.*",
			"health_check.#",
			"health_check.0.healthy_threshold",
			"health_check.0.unhealthy_threshold",
			"health_check.0.interval",
			"health_check.0.target",
			"health_check.0.timeout",
		},
	}
}
<commit_msg>providers/aws: read internal for ELB<commit_after>package aws

import (
	"fmt"
	"log"
	"strconv"

	"github.com/hashicorp/terraform/flatmap"
	"github.com/hashicorp/terraform/helper/config"
	"github.com/hashicorp/terraform/helper/diff"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/elb"
)

func resource_aws_elb_create(
	s *terraform.InstanceState,
	d *terraform.InstanceDiff,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	// Merge the diff into the state so that we have all the attributes
	// properly.
	rs := s.MergeDiff(d)

	// The name specified for the ELB. This is also our unique ID
	// we save to state if the creation is successful (amazon verifies
	// it is unique)
	elbName := rs.Attributes["name"]

	// Expand the "listener" array to goamz compat []elb.Listener
	v := flatmap.Expand(rs.Attributes, "listener").([]interface{})
	listeners, err := expandListeners(v)
	if err != nil {
		return nil, err
	}

	// Provision the elb
	elbOpts := &elb.CreateLoadBalancer{
		LoadBalancerName: elbName,
		Listeners:        listeners,
	}

	if rs.Attributes["internal"] == "true" {
		elbOpts.Internal = true
	}

	if _, ok := rs.Attributes["availability_zones.#"]; ok {
		v = flatmap.Expand(rs.Attributes, "availability_zones").([]interface{})
		elbOpts.AvailZone = expandStringList(v)
	}

	if _, ok := rs.Attributes["security_groups.#"]; ok {
		v = flatmap.Expand(rs.Attributes, "security_groups").([]interface{})
		elbOpts.SecurityGroups = expandStringList(v)
	}

	if _, ok := rs.Attributes["subnets.#"]; ok {
		v = flatmap.Expand(rs.Attributes, "subnets").([]interface{})
		elbOpts.Subnets = expandStringList(v)
	}

	log.Printf("[DEBUG] ELB create configuration: %#v", elbOpts)

	_, err = elbconn.CreateLoadBalancer(elbOpts)
	if err != nil {
		return nil, fmt.Errorf("Error creating ELB: %s", err)
	}

	// Assign the elb's unique identifier for use later
	rs.ID = elbName
	log.Printf("[INFO] ELB ID: %s", elbName)

	if _, ok := rs.Attributes["instances.#"]; ok {
		// If we have any instances, we need to register them
		v = flatmap.Expand(rs.Attributes, "instances").([]interface{})
		instances := expandStringList(v)

		if len(instances) > 0 {
			registerInstancesOpts := elb.RegisterInstancesWithLoadBalancer{
				LoadBalancerName: elbName,
				Instances:        instances,
			}

			_, err := elbconn.RegisterInstancesWithLoadBalancer(&registerInstancesOpts)

			if err != nil {
				return rs, fmt.Errorf("Failure registering instances: %s", err)
			}
		}
	}

	if _, ok := rs.Attributes["health_check.#"]; ok {
		v := flatmap.Expand(rs.Attributes, "health_check").([]interface{})
		health_check := v[0].(map[string]interface{})
		healthyThreshold, err := strconv.ParseInt(health_check["healthy_threshold"].(string), 0, 0)
		unhealthyThreshold, err := strconv.ParseInt(health_check["unhealthy_threshold"].(string), 0, 0)
		interval, err := strconv.ParseInt(health_check["interval"].(string), 0, 0)
		timeout, err := strconv.ParseInt(health_check["timeout"].(string), 0, 0)

		if err != nil {
			return nil, err
		}

		configureHealthCheckOpts := elb.ConfigureHealthCheck{
			LoadBalancerName: elbName,
			Check: elb.HealthCheck{
				HealthyThreshold:   healthyThreshold,
				UnhealthyThreshold: unhealthyThreshold,
				Interval:           interval,
				Target:             health_check["target"].(string),
				Timeout:            timeout,
			},
		}

		_, err = elbconn.ConfigureHealthCheck(&configureHealthCheckOpts)
		if err != nil {
			return rs, fmt.Errorf("Failure configuring health check: %s", err)
		}
	}

	loadBalancer, err := resource_aws_elb_retrieve_balancer(rs.ID, elbconn)
	if err != nil {
		return rs, err
	}

	return resource_aws_elb_update_state(rs, loadBalancer)
}

func resource_aws_elb_update(
	s *terraform.InstanceState,
	d *terraform.InstanceDiff,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	rs := s.MergeDiff(d)

	// If we currently have instances, or did have instances,
	// we want to figure out what to add and remove from the load
	// balancer
	if attr, ok := d.Attributes["instances.#"]; ok && attr.Old != "" {
		// The new state of instances merged with the diff
		mergedInstances := expandStringList(flatmap.Expand(
			rs.Attributes, "instances").([]interface{}))

		// The state before the diff merge
		previousInstances := expandStringList(flatmap.Expand(
			s.Attributes, "instances").([]interface{}))

		// keep track of what instances we are removing, and which
		// we are adding
		var toRemove []string
		var toAdd []string

		for _, instanceId := range mergedInstances {
			for _, prevId := range previousInstances {
				// If the merged instance ID existed
				// previously, we don't have to do anything
				if instanceId == prevId {
					continue
					// Otherwise, we need to add it to the load balancer
				} else {
					toAdd = append(toAdd, instanceId)
				}
			}
		}

		for i, instanceId := range toAdd {
			for _, prevId := range previousInstances {
				// If the instance ID we are adding existed
				// previously, we want to not add it, but rather remove
				// it
				if instanceId == prevId {
					toRemove = append(toRemove, instanceId)
					toAdd = append(toAdd[:i], toAdd[i+1:]...)
					// Otherwise, we continue adding it to the ELB
				} else {
					continue
				}
			}
		}

		if len(toAdd) > 0 {
			registerInstancesOpts := elb.RegisterInstancesWithLoadBalancer{
				LoadBalancerName: rs.ID,
				Instances:        toAdd,
			}

			_, err := elbconn.RegisterInstancesWithLoadBalancer(&registerInstancesOpts)

			if err != nil {
				return s, fmt.Errorf("Failure registering instances: %s", err)
			}
		}

		if len(toRemove) > 0 {
			deRegisterInstancesOpts := elb.DeregisterInstancesFromLoadBalancer{
				LoadBalancerName: rs.ID,
				Instances:        toRemove,
			}

			_, err := elbconn.DeregisterInstancesFromLoadBalancer(&deRegisterInstancesOpts)

			if err != nil {
				return s, fmt.Errorf("Failure deregistering instances: %s", err)
			}
		}
	}

	loadBalancer, err := resource_aws_elb_retrieve_balancer(rs.ID, elbconn)

	if err != nil {
		return s, err
	}

	return resource_aws_elb_update_state(rs, loadBalancer)
}

func resource_aws_elb_destroy(
	s *terraform.InstanceState,
	meta interface{}) error {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	log.Printf("[INFO] Deleting ELB: %s", s.ID)

	// Destroy the load balancer
	deleteElbOpts := elb.DeleteLoadBalancer{
		LoadBalancerName: s.ID,
	}
	_, err := elbconn.DeleteLoadBalancer(&deleteElbOpts)

	if err != nil {
		return fmt.Errorf("Error deleting ELB: %s", err)
	}

	return nil
}

func resource_aws_elb_refresh(
	s *terraform.InstanceState,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	elbconn := p.elbconn

	loadBalancer, err := resource_aws_elb_retrieve_balancer(s.ID, elbconn)
	if err != nil {
		return nil, err
	}

	return resource_aws_elb_update_state(s, loadBalancer)
}

func resource_aws_elb_diff(
	s *terraform.InstanceState,
	c *terraform.ResourceConfig,
	meta interface{}) (*terraform.InstanceDiff, error) {

	b := &diff.ResourceBuilder{
		Attrs: map[string]diff.AttrType{
			"name":              diff.AttrTypeCreate,
			"availability_zone": diff.AttrTypeCreate,
			"security_groups":   diff.AttrTypeCreate, // TODO could be AttrTypeUpdate
			"subnets":           diff.AttrTypeCreate, // TODO could be AttrTypeUpdate
			"listener":          diff.AttrTypeCreate,
			"instances":         diff.AttrTypeUpdate,
			"health_check":      diff.AttrTypeCreate,
			"internal":          diff.AttrTypeCreate,
		},

		ComputedAttrs: []string{
			"dns_name",
		},
	}

	return b.Diff(s, c)
}

func resource_aws_elb_update_state(
	s *terraform.InstanceState,
	balancer *elb.LoadBalancer) (*terraform.InstanceState, error) {

	s.Attributes["name"] = balancer.LoadBalancerName
	s.Attributes["dns_name"] = balancer.DNSName

	if balancer.Scheme == "internal" {
		s.Attributes["internal"] = "true"
	}

	// Flatten our group values
	toFlatten := make(map[string]interface{})

	if len(balancer.Instances) > 0 && balancer.Instances[0].InstanceId != "" {
		toFlatten["instances"] = flattenInstances(balancer.Instances)
	}

	if len(balancer.SecurityGroups) > 0 && balancer.SecurityGroups[0] != "" {
		toFlatten["security_groups"] = balancer.SecurityGroups
	}

	if len(balancer.Subnets) > 0 && balancer.Subnets[0] != "" {
		toFlatten["subnets"] = balancer.Subnets
	}

	// There's only one health check, so save that to state as we
	// currently can
	if balancer.HealthCheck.Target != "" {
		toFlatten["health_check"] = flattenHealthCheck(balancer.HealthCheck)
	}

	for k, v := range flatmap.Flatten(toFlatten) {
		s.Attributes[k] = v
	}

	return s, nil
}

// retrieves an ELB by its ID
func resource_aws_elb_retrieve_balancer(id string, elbconn *elb.ELB) (*elb.LoadBalancer, error) {
	describeElbOpts := &elb.DescribeLoadBalancer{
		Names: []string{id},
	}

	// Retrieve the ELB properties for updating the state
	describeResp, err := elbconn.DescribeLoadBalancers(describeElbOpts)
	if err != nil {
		return nil, fmt.Errorf("Error retrieving ELB: %s", err)
	}

	loadBalancer := describeResp.LoadBalancers[0]

	// Verify AWS returned our ELB
	if len(describeResp.LoadBalancers) != 1 ||
		describeResp.LoadBalancers[0].LoadBalancerName != id {
		if err != nil {
			return nil, fmt.Errorf("Unable to find ELB: %#v", describeResp.LoadBalancers)
		}
	}

	return &loadBalancer, nil
}

func resource_aws_elb_validation() *config.Validator {
	return &config.Validator{
		Required: []string{
			"name",
			"listener.*",
			"listener.*.instance_port",
			"listener.*.instance_protocol",
			"listener.*.lb_port",
			"listener.*.lb_protocol",
		},
		Optional: []string{
			"instances.*",
			"internal",
			"availability_zones.*",
			"security_groups.*",
			"subnets.*",
			"health_check.#",
			"health_check.0.healthy_threshold",
			"health_check.0.unhealthy_threshold",
			"health_check.0.interval",
			"health_check.0.target",
			"health_check.0.timeout",
		},
	}
}
<|endoftext|>


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

// +build acceptance

package acceptance

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/randutil"
	"github.com/cockroachdb/cockroach/util/stop"
)

var maxTransfer = flag.Int("max-transfer", 999, "Maximum amount to transfer in one transaction.")
var numAccounts = flag.Int("num-accounts", 999, "Number of accounts.")

// TestChaos starts up a cluster with an "accounts" table.
// It starts transferring money between accounts, while nodes are
// being killed and restarted continuously.
// The test doesn't measure write performance, but cluster recovery.
// TODO(vivek): Expand this test to check that write performance
// is unaffected by chaos.
func TestChaos(t *testing.T) {
	c := StartCluster(t)
	defer c.AssertAndStop(t)

	num := c.NumNodes()
	if num <= 0 {
		t.Fatalf("%d nodes in cluster", num)
	}

	// One error sent by each client. A successful client sends a nil error.
	errs := make(chan error, num)
	// The number of successful writes (puts) to the database.
	var count int64
	// The number of times chaos monkey has run.
	var round int64
	// Set to 1 if chaos monkey has stalled the writes.
	var stalled int32
	// One client for each node.
	clients := make([]struct {
		sync.RWMutex
		db      *sql.DB
		stopper *stop.Stopper
		count   int64
	}, num)

	// initClient initializes the client talking to node "i".
	// It requires that the caller hold the client's write lock.
	initClient := func(i int) {
		if clients[i].stopper != nil {
			clients[i].stopper.Stop()
		}
		clients[i].db = makePGClient(t, c.PGUrl(i))
		clients[i].stopper = stop.NewStopper()
	}

	// Initialize the "accounts" table.
	db := makePGClient(t, c.PGUrl(0))

	if _, err := db.Exec(`CREATE DATABASE IF NOT EXISTS bank`); err != nil {
		t.Fatal(err)
	}

	// Delete table created by a prior instance of a test.
	if _, err := db.Exec(`DROP TABLE IF EXISTS bank.accounts`); err != nil {
		t.Fatal(err)
	}

	schema := `
CREATE TABLE bank.accounts (
  id INT PRIMARY KEY,
  balance INT NOT NULL
)`
	if _, err := db.Exec(schema); err != nil {
		t.Fatal(err)
	}

	var placeholders bytes.Buffer
	var values []interface{}
	for i := 0; i < *numAccounts; i++ {
		if i > 0 {
			placeholders.WriteString(", ")
		}
		fmt.Fprintf(&placeholders, "($%d, 0)", i+1)
		values = append(values, i)
	}
	stmt := `INSERT INTO bank.accounts (id, balance) VALUES ` + placeholders.String()
	if _, err := db.Exec(stmt, values...); err != nil {
		t.Fatal(err)
	}

	start := time.Now()
	deadline := start.Add(*flagDuration)
	done := func() bool {
		return !time.Now().Before(deadline) || atomic.LoadInt32(&stalled) == 1
	}

	for i := 0; i < num; i++ {
		clients[i].Lock()
		initClient(i)
		clients[i].Unlock()
		go func(i int) {
			for !done() {
				if err := func() error {
					clients[i].RLock()
					defer clients[i].RUnlock()
					from := rand.Intn(*numAccounts)
					to := rand.Intn(*numAccounts - 1)
					if from == to {
						to = *numAccounts - 1
					}
					amount := rand.Intn(*maxTransfer)

					const update = `
									UPDATE bank.accounts
									  SET balance = CASE id WHEN $1 THEN balance-$3 WHEN $2 THEN balance+$3 END
										  WHERE id IN ($1, $2) AND (SELECT balance >= $3 FROM bank.accounts WHERE id = $1)`
					if _, err := clients[i].db.Exec(update, from, to, amount); err != nil {
						// Ignore some errors.
						if testutils.IsError(err, "connection refused") {
							return nil
						}
						return err
					}

					// Only advance the counts on a successful update.
					_ = atomic.AddInt64(&count, 1)
					atomic.AddInt64(&clients[i].count, 1)
					return nil
				}(); err != nil {
					// Report the err and terminate.
					errs <- err
					return
				}
			}
			log.Infof("client %d shutting down", i)
			errs <- nil
		}(i)
	}

	teardown := make(chan struct{})
	defer func() {
		<-teardown
		for i := range clients {
			clients[i].RLock()
			clients[i].stopper.Stop()
			clients[i].stopper = nil
			clients[i].RUnlock()
		}
	}()

	// Chaos monkey.
	go func() {
		defer close(teardown)
		rnd, seed := randutil.NewPseudoRand()
		log.Warningf("monkey starts (seed %d)", seed)
		for atomic.StoreInt64(&round, 1); !done(); atomic.AddInt64(&round, 1) {
			curRound := atomic.LoadInt64(&round)
			select {
			case <-stopper:
				return
			default:
			}
			nodes := rnd.Perm(num)[:rnd.Intn(num)+1]
			// Prevent all clients from writing while nodes are being restarted.
			for i := 0; i < num; i++ {
				clients[i].Lock()
			}
			log.Infof("round %d: restarting nodes %v", curRound, nodes)
			for _, i := range nodes {
				// Two early exit conditions.
				select {
				case <-stopper:
					break
				default:
				}
				if done() {
					break
				}
				log.Infof("round %d: restarting %d", curRound, i)
				c.Kill(i)
				c.Restart(i)
				initClient(i)
			}
			for i := 0; i < num; i++ {
				clients[i].Unlock()
			}
			// Sleep until at least one client is writing successfully.
			first := true
			for cur := atomic.LoadInt64(&count); !done() && atomic.LoadInt64(&count) == cur; time.Sleep(time.Second) {
				c.Assert(t)
				if first {
					first = false
					log.Warningf("round %d: monkey sleeping while cluster recovers...", curRound)
				}
			}
		}
	}()

	prevRound := atomic.LoadInt64(&round)
	stallTime := time.Now().Add(*flagStall)
	var prevOutput string
	// Spin until all clients are shut.
	for numShutClients := 0; numShutClients < num; {
		select {
		case <-teardown:
		case <-stopper:
			t.Fatal("interrupted")

		case err := <-errs:
			if err != nil {
				t.Error(err)
			}
			numShutClients++

		case <-time.After(time.Second):
			var newOutput string
			if time.Now().Before(deadline) {
				curRound := atomic.LoadInt64(&round)
				if curRound == prevRound {
					if time.Now().After(stallTime) {
						atomic.StoreInt32(&stalled, 1)
						t.Fatalf("Stall detected at round %d, no forward progress for %s", curRound, *flagStall)
					}
				} else {
					prevRound = curRound
					stallTime = time.Now().Add(*flagStall)
				}
				// Periodically print out progress so that we know the test is
				// still running and making progress.
				cur := make([]string, num)
				for j := range cur {
					cur[j] = fmt.Sprintf("%d", atomic.LoadInt64(&clients[j].count))
				}
				newOutput = fmt.Sprintf("round %d: %d (%s)", curRound, atomic.LoadInt64(&count), strings.Join(cur, ", "))
			} else {
				newOutput = fmt.Sprintf("test finished, waiting for shutdown of %d clients", num-numShutClients)
			}
			// This just stops the logs from being a bit too spammy.
			if newOutput != prevOutput {
				log.Infof(newOutput)
				prevOutput = newOutput
			}
		}
	}

	// Verify accounts.

	// Hold the read lock on node 0 to prevent it being restarted by
	// chaos monkey.
	var sum int
	clients[0].RLock()
	if err := clients[0].db.QueryRow("SELECT SUM(balance) FROM bank.accounts").Scan(&sum); err != nil {
		t.Fatal(err)
	}
	clients[0].RUnlock()
	if sum != 0 {
		t.Fatalf("The bank is not in good order. Total value: %d", sum)
	}

	elapsed := time.Since(start)
	log.Infof("%d %.1f/sec", count, float64(count)/elapsed.Seconds())
}
<commit_msg>acceptance/chaos: remove count and use uint64<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Tobias Schottdorf (tobias@cockroachlabs.com)

// +build acceptance

package acceptance

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/randutil"
	"github.com/cockroachdb/cockroach/util/stop"
)

var maxTransfer = flag.Int("max-transfer", 999, "Maximum amount to transfer in one transaction.")
var numAccounts = flag.Int("num-accounts", 999, "Number of accounts.")

// TestChaos starts up a cluster with an "accounts" table.
// It starts transferring money between accounts, while nodes are
// being killed and restarted continuously.
// The test doesn't measure write performance, but cluster recovery.
// TODO(vivek): Expand this test to check that write performance
// is unaffected by chaos.
func TestChaos(t *testing.T) {
	c := StartCluster(t)
	defer c.AssertAndStop(t)

	num := c.NumNodes()
	if num <= 0 {
		t.Fatalf("%d nodes in cluster", num)
	}

	// One error sent by each client. A successful client sends a nil error.
	errs := make(chan error, num)
	// The number of times chaos monkey has run.
	var round uint64
	// Set to 1 if chaos monkey has stalled the writes.
	var stalled int32
	// One client for each node.
	clients := make([]struct {
		sync.RWMutex
		db      *sql.DB
		stopper *stop.Stopper
		count   uint64
	}, num)

	// initClient initializes the client talking to node "i".
	// It requires that the caller hold the client's write lock.
	initClient := func(i int) {
		if clients[i].stopper != nil {
			clients[i].stopper.Stop()
		}
		clients[i].db = makePGClient(t, c.PGUrl(i))
		clients[i].stopper = stop.NewStopper()
	}

	// Initialize the "accounts" table.
	db := makePGClient(t, c.PGUrl(0))

	if _, err := db.Exec(`CREATE DATABASE IF NOT EXISTS bank`); err != nil {
		t.Fatal(err)
	}

	// Delete table created by a prior instance of a test.
	if _, err := db.Exec(`DROP TABLE IF EXISTS bank.accounts`); err != nil {
		t.Fatal(err)
	}

	schema := `
CREATE TABLE bank.accounts (
  id INT PRIMARY KEY,
  balance INT NOT NULL
)`
	if _, err := db.Exec(schema); err != nil {
		t.Fatal(err)
	}

	var placeholders bytes.Buffer
	var values []interface{}
	for i := 0; i < *numAccounts; i++ {
		if i > 0 {
			placeholders.WriteString(", ")
		}
		fmt.Fprintf(&placeholders, "($%d, 0)", i+1)
		values = append(values, i)
	}
	stmt := `INSERT INTO bank.accounts (id, balance) VALUES ` + placeholders.String()
	if _, err := db.Exec(stmt, values...); err != nil {
		t.Fatal(err)
	}

	start := time.Now()
	deadline := start.Add(*flagDuration)
	done := func() bool {
		return !time.Now().Before(deadline) || atomic.LoadInt32(&stalled) == 1
	}

	for i := 0; i < num; i++ {
		clients[i].Lock()
		initClient(i)
		clients[i].Unlock()
		go func(i int) {
			for !done() {
				if err := func() error {
					clients[i].RLock()
					defer clients[i].RUnlock()
					from := rand.Intn(*numAccounts)
					to := rand.Intn(*numAccounts - 1)
					if from == to {
						to = *numAccounts - 1
					}
					amount := rand.Intn(*maxTransfer)

					const update = `
									UPDATE bank.accounts
									  SET balance = CASE id WHEN $1 THEN balance-$3 WHEN $2 THEN balance+$3 END
										  WHERE id IN ($1, $2) AND (SELECT balance >= $3 FROM bank.accounts WHERE id = $1)`
					if _, err := clients[i].db.Exec(update, from, to, amount); err != nil {
						// Ignore some errors.
						if testutils.IsError(err, "connection refused") {
							return nil
						}
						return err
					}

					// Only advance the counts on a successful update.
					atomic.AddUint64(&clients[i].count, 1)
					return nil
				}(); err != nil {
					// Report the err and terminate.
					errs <- err
					return
				}
			}
			log.Infof("client %d shutting down", i)
			errs <- nil
		}(i)
	}

	teardown := make(chan struct{})
	defer func() {
		<-teardown
		for i := range clients {
			clients[i].Lock()
			clients[i].stopper.Stop()
			clients[i].stopper = nil
			clients[i].Unlock()
		}
	}()

	// Chaos monkey.
	go func() {
		defer close(teardown)
		rnd, seed := randutil.NewPseudoRand()
		log.Warningf("monkey starts (seed %d)", seed)
		for atomic.StoreUint64(&round, 1); !done(); atomic.AddUint64(&round, 1) {
			curRound := atomic.LoadUint64(&round)
			select {
			case <-stopper:
				return
			default:
			}
			nodes := rnd.Perm(num)[:rnd.Intn(num)+1]
			// Prevent all clients from writing while nodes are being restarted.
			for i := 0; i < num; i++ {
				clients[i].Lock()
			}
			log.Infof("round %d: restarting nodes %v", curRound, nodes)
			for _, i := range nodes {
				// Two early exit conditions.
				select {
				case <-stopper:
					break
				default:
				}
				if done() {
					break
				}
				log.Infof("round %d: restarting %d", curRound, i)
				c.Kill(i)
				c.Restart(i)
				initClient(i)
			}
			for i := 0; i < num; i++ {
				clients[i].Unlock()
			}

			preCount := make([]uint64, len(clients))
			for i, client := range clients {
				preCount[i] = atomic.LoadUint64(&client.count)
			}

			madeProgress := func() bool {
				c.Assert(t)
				for i, client := range clients {
					if atomic.LoadUint64(&client.count) > preCount[i] {
						return true
					}
				}
				return false
			}

			// Sleep until at least one client is writing successfully.
			log.Warningf("round %d: monkey sleeping while cluster recovers...", curRound)
			for !madeProgress() {
				time.Sleep(time.Second)
			}
			log.Warningf("round %d: cluster recovered", curRound)
		}
	}()

	prevRound := atomic.LoadUint64(&round)
	stallTime := time.Now().Add(*flagStall)
	var prevOutput string
	// Spin until all clients are shut.
	for numShutClients := 0; numShutClients < num; {
		select {
		case <-teardown:
		case <-stopper:
			t.Fatal("interrupted")

		case err := <-errs:
			if err != nil {
				t.Error(err)
			}
			numShutClients++

		case <-time.After(time.Second):
			var newOutput string
			if time.Now().Before(deadline) {
				curRound := atomic.LoadUint64(&round)
				if curRound == prevRound {
					if time.Now().After(stallTime) {
						atomic.StoreInt32(&stalled, 1)
						t.Fatalf("Stall detected at round %d, no forward progress for %s", curRound, *flagStall)
					}
				} else {
					prevRound = curRound
					stallTime = time.Now().Add(*flagStall)
				}
				// Periodically print out progress so that we know the test is
				// still running and making progress.
				preCount := make([]string, len(clients))
				for i, client := range clients {
					preCount[i] = strconv.FormatUint(atomic.LoadUint64(&client.count), 10)
				}
				newOutput = fmt.Sprintf("round %d: client counts: (%s)", curRound, strings.Join(preCount, ", "))
			} else {
				newOutput = fmt.Sprintf("test finished, waiting for shutdown of %d clients", num-numShutClients)
			}
			// This just stops the logs from being a bit too spammy.
			if newOutput != prevOutput {
				log.Infof(newOutput)
				prevOutput = newOutput
			}
		}
	}

	// Verify accounts.

	// Hold the read lock on node 0 to prevent it being restarted by
	// chaos monkey.
	var sum int
	clients[0].RLock()
	if err := clients[0].db.QueryRow("SELECT SUM(balance) FROM bank.accounts").Scan(&sum); err != nil {
		t.Fatal(err)
	}
	clients[0].RUnlock()
	if sum != 0 {
		t.Fatalf("The bank is not in good order. Total value: %d", sum)
	}

	elapsed := time.Since(start)
	var count uint64
	for _, client := range clients {
		count += atomic.LoadUint64(&client.count)
	}
	log.Infof("%d %.1f/sec", count, float64(count)/elapsed.Seconds())
}
<|endoftext|>


<commit_before>package aggregator

import (
	"encoding/json"

	"github.com/Sirupsen/logrus"
	"github.com/jacksontj/dnms/graph"
	"github.com/jacksontj/eventsource"
)

// Subscripe to dest, consuming events into PeerGraphMap.
// We'll return a bool channel which can be used to cancel the subscription
func Subscribe(p *PeerGraphMap) chan bool {
	exitChan := make(chan bool)
	go func() {
		stream, err := eventsource.Subscribe("http://"+p.Name+":12345/v1/events/graph", "")
		if err != nil {
			logrus.Fatalf("Error subscribing: %v", err)
		}
		logrus.Infof("connecting to peer: %v", p.Name)
		// defer a removal in case the peer disconnects (or blips)
		defer p.cleanup()

		for {
			select {
			// handle errors-- all of these mean a disconnect/reconnect
			case err, ok := <-stream.Errors:
				logrus.Debugf("stream error, reconnecting: %v %v", err, ok)
				// we need to remove everything we know about this peer-- since
				// the new connection will re-seed on the new connection
				p.cleanup()
			case ev := <-stream.Events:
				//logrus.Infof("Got Event: %v", ev.Event())
				switch ev.Event() {

				// Node events
				case "addNodeEvent":
					n := graph.NetworkNode{}
					err := json.Unmarshal([]byte(ev.Data()), &n)
					if err != nil {
						logrus.Warningf("unable to unmarshal node: %v", err)
					}
					p.AddNode(&n)
				case "updateNodeEvent":
					n := graph.NetworkNode{}
					err := json.Unmarshal([]byte(ev.Data()), &n)
					if err != nil {
						logrus.Warningf("unable to unmarshal node: %v", err)
					}
					node := p.Graph.GetNode(n.Name)
					// TODO: some sort of "merge" method
					if node != nil {
						node.DNSNames = n.DNSNames
					}
				case "removeNodeEvent":
					n := graph.NetworkNode{}
					err := json.Unmarshal([]byte(ev.Data()), &n)
					if err != nil {
						logrus.Warningf("unable to unmarshal node: %v", err)
					}
					p.RemoveNode(&n)

				// Link events
				case "addLinkEvent":
					l := graph.NetworkLink{}
					err := json.Unmarshal([]byte(ev.Data()), &l)
					if err != nil {
						logrus.Warningf("unable to unmarshal link: %v", err)
					}
					p.AddLink(&l)
				// TODO: update event
				case "updateLinkEvent":
					// TODO: implement
					// TODO: some sort of "merge" method
				case "removeLinkEvent":
					l := graph.NetworkLink{}
					err := json.Unmarshal([]byte(ev.Data()), &l)
					if err != nil {
						logrus.Warningf("unable to unmarshal link: %v", err)
					}
					p.RemoveLink(&l)

				// route events
				case "addRouteEvent":
					r := graph.NetworkRoute{}
					err := json.Unmarshal([]byte(ev.Data()), &r)
					if err != nil {
						logrus.Warningf("unable to unmarshal route: %v", err)
					}
					p.AddRoute(&r)
				case "updateRouteEvent":
					r := graph.NetworkRoute{}
					err := json.Unmarshal([]byte(ev.Data()), &r)
					if err != nil {
						logrus.Warningf("unable to unmarshal route: %v", err)
					}
					route := p.Graph.GetRoute(r.Hops())

					// TODO: some sort of "merge" method
					route.State = r.State
				case "removeRouteEvent":
					r := graph.NetworkRoute{}
					err := json.Unmarshal([]byte(ev.Data()), &r)
					if err != nil {
						logrus.Warningf("unable to unmarshal route: %v", err)
					}
					p.RemoveRoute(&r)

				}
			case <-exitChan:
				return
			}
		}
	}()
	return exitChan
}
<commit_msg>Aggregator to retry on failure to connect to peer<commit_after>package aggregator

import (
	"encoding/json"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/jacksontj/dnms/graph"
	"github.com/jacksontj/eventsource"
)

// Subscripe to dest, consuming events into PeerGraphMap.
// We'll return a bool channel which can be used to cancel the subscription
func Subscribe(p *PeerGraphMap) chan bool {
	exitChan := make(chan bool)
	go func() {
		var stream *eventsource.Stream
		for {
			logrus.Infof("connecting to peer: %v", p.Name)
			var err error
			stream, err = eventsource.Subscribe("http://"+p.Name+":12345/v1/events/graph", "")
			if err != nil {
				logrus.Errorf("Error subscribing, retrying: %v", err)
				time.Sleep(time.Second)
			} else {
				break
			}
		}

		// defer a removal in case the peer disconnects (or blips)
		defer p.cleanup()

		for {
			select {
			// handle errors-- all of these mean a disconnect/reconnect
			case err, ok := <-stream.Errors:
				logrus.Debugf("stream error, reconnecting: %v %v", err, ok)
				// we need to remove everything we know about this peer-- since
				// the new connection will re-seed on the new connection
				p.cleanup()
			case ev := <-stream.Events:
				//logrus.Infof("Got Event: %v", ev.Event())
				switch ev.Event() {

				// Node events
				case "addNodeEvent":
					n := graph.NetworkNode{}
					err := json.Unmarshal([]byte(ev.Data()), &n)
					if err != nil {
						logrus.Warningf("unable to unmarshal node: %v", err)
					}
					p.AddNode(&n)
				case "updateNodeEvent":
					n := graph.NetworkNode{}
					err := json.Unmarshal([]byte(ev.Data()), &n)
					if err != nil {
						logrus.Warningf("unable to unmarshal node: %v", err)
					}
					node := p.Graph.GetNode(n.Name)
					// TODO: some sort of "merge" method
					if node != nil {
						node.DNSNames = n.DNSNames
					}
				case "removeNodeEvent":
					n := graph.NetworkNode{}
					err := json.Unmarshal([]byte(ev.Data()), &n)
					if err != nil {
						logrus.Warningf("unable to unmarshal node: %v", err)
					}
					p.RemoveNode(&n)

				// Link events
				case "addLinkEvent":
					l := graph.NetworkLink{}
					err := json.Unmarshal([]byte(ev.Data()), &l)
					if err != nil {
						logrus.Warningf("unable to unmarshal link: %v", err)
					}
					p.AddLink(&l)
				// TODO: update event
				case "updateLinkEvent":
					// TODO: implement
					// TODO: some sort of "merge" method
				case "removeLinkEvent":
					l := graph.NetworkLink{}
					err := json.Unmarshal([]byte(ev.Data()), &l)
					if err != nil {
						logrus.Warningf("unable to unmarshal link: %v", err)
					}
					p.RemoveLink(&l)

				// route events
				case "addRouteEvent":
					r := graph.NetworkRoute{}
					err := json.Unmarshal([]byte(ev.Data()), &r)
					if err != nil {
						logrus.Warningf("unable to unmarshal route: %v", err)
					}
					p.AddRoute(&r)
				case "updateRouteEvent":
					r := graph.NetworkRoute{}
					err := json.Unmarshal([]byte(ev.Data()), &r)
					if err != nil {
						logrus.Warningf("unable to unmarshal route: %v", err)
					}
					route := p.Graph.GetRoute(r.Hops())

					// TODO: some sort of "merge" method
					route.State = r.State
				case "removeRouteEvent":
					r := graph.NetworkRoute{}
					err := json.Unmarshal([]byte(ev.Data()), &r)
					if err != nil {
						logrus.Warningf("unable to unmarshal route: %v", err)
					}
					p.RemoveRoute(&r)

				}
			case <-exitChan:
				return
			}
		}
	}()
	return exitChan
}
<|endoftext|>


<commit_before>package asset

import (
	"github.com/btcsuite/btcutil/hdkeychain"
	"golang.org/x/net/context"

	"chain/api/appdb"
	"chain/errors"
	"chain/fedchain-sandbox/hdkey"
)

var (
	// ErrBadXPub is returned by CreateNode.
	// It may be wrapped using package chain/errors.
	ErrBadXPub = errors.New("bad xpub")

	// ErrTooFewKeys can be returned by CreateNode if not enough keys
	// have been provided or generated.
	ErrTooFewKeys = errors.New("too few keys for signatures required")
)

type nodeType int

// Node types used for CreateNode
const (
	ManagerNode nodeType = iota
	IssuerNode  nodeType = iota
)

// CreateNodeReq is a user filled struct
// passed into CreateManagerNode or CreateIssuerNode
type CreateNodeReq struct {
	Label        string
	Keys         []*XPubInit
	SigsRequired int `json:"signatures_required"`
}

// DeprecatedCreateNodeReq is a user filled struct
// passed into CreateManagerNode or CreateIssuerNode.
// It is deprecated in favor of CreateNodeReq.
type DeprecatedCreateNodeReq struct {
	Label       string
	XPubs       []string
	GenerateKey bool `json:"generate_key"`
}

// XPubInit is a representation of an xpub used when nodes are being created.
// It includes the key itself, as well as two flags:
// Generate specifies whether the key needs to be generated server-side, and
// Variable specifies whether this is a placeholder for an account-specific key.
// If Variable is true, Generate must be false and Key must be empty.
type XPubInit struct {
	Key      string
	Generate bool
	Variable bool
}

// CreateNode is used to create manager and issuer nodes
func CreateNode(ctx context.Context, node nodeType, projID string, req *CreateNodeReq) (interface{}, error) {
	if req.Label == "" {
		return nil, appdb.ErrBadLabel
	}

	var (
		keys       []*hdkey.XKey
		gennedKeys []*hdkey.XKey
	)

	variableKeyCount := 0
	for i, xpub := range req.Keys {
		if xpub.Generate {
			pub, priv, err := newKey()
			if err != nil {
				return nil, err
			}
			keys = append(keys, pub)
			gennedKeys = append(gennedKeys, priv)
		} else if xpub.Variable {
			variableKeyCount++
		} else {
			key, err := hdkey.NewXKey(xpub.Key)
			if err != nil {
				err = errors.Wrap(ErrBadXPub, err.Error())
				return nil, errors.WithDetailf(err, "xpub %d", i)
			}
			keys = append(keys, key)
		}
	}

	if len(keys)+variableKeyCount < req.SigsRequired {
		return nil, ErrTooFewKeys
	}

	for i, key := range keys {
		if key.IsPrivate() {
			return nil, errors.WithDetailf(ErrBadXPub, "key %d is xpriv, not xpub", i)
		}
	}

	if node == ManagerNode {
		return appdb.InsertManagerNode(ctx, projID, req.Label, keys, gennedKeys, variableKeyCount, req.SigsRequired)
	}
	// Do nothing with variable keys for Issuer Nodes since they can't have variable keys yet.
	return appdb.InsertIssuerNode(ctx, projID, req.Label, keys, gennedKeys, req.SigsRequired)
}

func newKey() (pub, priv *hdkey.XKey, err error) {
	seed, err := hdkeychain.GenerateSeed(hdkeychain.RecommendedSeedLen)
	if err != nil {
		return nil, nil, errors.Wrap(err, "generating key seed")
	}
	xprv, err := hdkeychain.NewMaster(seed)
	if err != nil {
		return nil, nil, errors.Wrap(err, "creating root xprv")
	}
	xpub, err := xprv.Neuter()
	if err != nil {
		return nil, nil, errors.Wrap(err, "getting root xpub")
	}
	return &hdkey.XKey{ExtendedKey: *xpub}, &hdkey.XKey{ExtendedKey: *xprv}, nil
}
<commit_msg>bugfix: make sure signatures are required when creating node<commit_after>package asset

import (
	"github.com/btcsuite/btcutil/hdkeychain"
	"golang.org/x/net/context"

	"chain/api/appdb"
	"chain/errors"
	"chain/fedchain-sandbox/hdkey"
)

var (
	// ErrBadXPub is returned by CreateNode.
	// It may be wrapped using package chain/errors.
	ErrBadXPub = errors.New("bad xpub")

	// ErrTooFewKeys can be returned by CreateNode if not enough keys
	// have been provided or generated.
	ErrTooFewKeys = errors.New("too few keys for signatures required")
)

type nodeType int

// Node types used for CreateNode
const (
	ManagerNode nodeType = iota
	IssuerNode  nodeType = iota
)

// CreateNodeReq is a user filled struct
// passed into CreateManagerNode or CreateIssuerNode
type CreateNodeReq struct {
	Label        string
	Keys         []*XPubInit
	SigsRequired int `json:"signatures_required"`
}

// DeprecatedCreateNodeReq is a user filled struct
// passed into CreateManagerNode or CreateIssuerNode.
// It is deprecated in favor of CreateNodeReq.
type DeprecatedCreateNodeReq struct {
	Label       string
	XPubs       []string
	GenerateKey bool `json:"generate_key"`
}

// XPubInit is a representation of an xpub used when nodes are being created.
// It includes the key itself, as well as two flags:
// Generate specifies whether the key needs to be generated server-side, and
// Variable specifies whether this is a placeholder for an account-specific key.
// If Variable is true, Generate must be false and Key must be empty.
type XPubInit struct {
	Key      string
	Generate bool
	Variable bool
}

// CreateNode is used to create manager and issuer nodes
func CreateNode(ctx context.Context, node nodeType, projID string, req *CreateNodeReq) (interface{}, error) {
	if req.Label == "" {
		return nil, appdb.ErrBadLabel
	}

	var (
		keys       []*hdkey.XKey
		gennedKeys []*hdkey.XKey
	)

	variableKeyCount := 0
	for i, xpub := range req.Keys {
		if xpub.Generate {
			pub, priv, err := newKey()
			if err != nil {
				return nil, err
			}
			keys = append(keys, pub)
			gennedKeys = append(gennedKeys, priv)
		} else if xpub.Variable {
			variableKeyCount++
		} else {
			key, err := hdkey.NewXKey(xpub.Key)
			if err != nil {
				err = errors.Wrap(ErrBadXPub, err.Error())
				return nil, errors.WithDetailf(err, "xpub %d", i)
			}
			keys = append(keys, key)
		}
	}

	// If the request is missing a SigsRequired field,
	// it's a deprecated-style request. Default to 1-of-1.
	if req.SigsRequired == 0 {
		req.SigsRequired = 1
	}

	if len(keys)+variableKeyCount < req.SigsRequired {
		return nil, ErrTooFewKeys
	}

	for i, key := range keys {
		if key.IsPrivate() {
			return nil, errors.WithDetailf(ErrBadXPub, "key %d is xpriv, not xpub", i)
		}
	}

	if node == ManagerNode {
		return appdb.InsertManagerNode(ctx, projID, req.Label, keys, gennedKeys, variableKeyCount, req.SigsRequired)
	}
	// Do nothing with variable keys for Issuer Nodes since they can't have variable keys yet.
	return appdb.InsertIssuerNode(ctx, projID, req.Label, keys, gennedKeys, req.SigsRequired)
}

func newKey() (pub, priv *hdkey.XKey, err error) {
	seed, err := hdkeychain.GenerateSeed(hdkeychain.RecommendedSeedLen)
	if err != nil {
		return nil, nil, errors.Wrap(err, "generating key seed")
	}
	xprv, err := hdkeychain.NewMaster(seed)
	if err != nil {
		return nil, nil, errors.Wrap(err, "creating root xprv")
	}
	xpub, err := xprv.Neuter()
	if err != nil {
		return nil, nil, errors.Wrap(err, "getting root xpub")
	}
	return &hdkey.XKey{ExtendedKey: *xpub}, &hdkey.XKey{ExtendedKey: *xprv}, nil
}
<|endoftext|>