<commit_before>// Package build implements Pipe and can build Go projects for
// several platforms, with pre and post hook support.
package build

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

	"github.com/apex/log"
	"github.com/goreleaser/goreleaser/context"
	"github.com/goreleaser/goreleaser/internal/ext"
	"golang.org/x/sync/errgroup"
)

// Pipe for build
type Pipe struct{}

// Description of the pipe
func (Pipe) Description() string {
	return "Building binaries"
}

// Run the pipe
func (Pipe) Run(ctx *context.Context) error {
	if err := runHook(ctx.Config.Build.Env, ctx.Config.Build.Hooks.Pre); err != nil {
		return err
	}
	sem := make(chan bool, 4)
	var g errgroup.Group
	for _, target := range buildTargets(ctx) {
		name, err := nameFor(ctx, target)
		if err != nil {
			return err
		}
		ctx.Archives[target.String()] = name

		sem <- true
		target := target
		g.Go(func() error {
			defer func() {
				<-sem
			}()
			return build(ctx, name, target)
		})
	}
	if err := g.Wait(); err != nil {
		return err
	}
	return runHook(ctx.Config.Build.Env, ctx.Config.Build.Hooks.Post)
}

func runHook(env []string, hook string) error {
	if hook == "" {
		return nil
	}
	log.WithField("hook", hook).Info("running hook")
	cmd := strings.Fields(hook)
	return run(runtimeTarget, cmd, env)
}

func build(ctx *context.Context, name string, target buildTarget) error {
	output := filepath.Join(
		ctx.Config.Dist,
		name,
		ctx.Config.Build.Binary+ext.For(target.goos),
	)
	log.WithField("binary", output).Info("building")
	cmd := []string{"go", "build"}
	if ctx.Config.Build.Flags != "" {
		cmd = append(cmd, strings.Fields(ctx.Config.Build.Flags)...)
	}
	flags, err := ldflags(ctx)
	if err != nil {
		return err
	}
	cmd = append(cmd, "-ldflags="+flags, "-o", output, ctx.Config.Build.Main)
	return run(target, cmd, ctx.Config.Build.Env)
}

func run(target buildTarget, command, env []string) error {
	var cmd = exec.Command(command[0], command[1:]...)
	env = append(env, "GOOS="+target.goos, "GOARCH="+target.goarch, "GOARM="+target.goarm)
	cmd.Env = append(cmd.Env, os.Environ()...)
	cmd.Env = append(cmd.Env, env...)
	log.WithField("target", target.PrettyString()).
		WithField("env", env).
		WithField("args", cmd.Args).
		Debug("running")
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("build failed: %s\n%v", target.PrettyString(), string(out))
	}
	return nil
}
<commit_msg>improved debug logging<commit_after>// Package build implements Pipe and can build Go projects for
// several platforms, with pre and post hook support.
package build

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

	"github.com/apex/log"
	"github.com/goreleaser/goreleaser/context"
	"github.com/goreleaser/goreleaser/internal/ext"
	"golang.org/x/sync/errgroup"
)

// Pipe for build
type Pipe struct{}

// Description of the pipe
func (Pipe) Description() string {
	return "Building binaries"
}

// Run the pipe
func (Pipe) Run(ctx *context.Context) error {
	if err := runHook(ctx.Config.Build.Env, ctx.Config.Build.Hooks.Pre); err != nil {
		return err
	}
	sem := make(chan bool, 4)
	var g errgroup.Group
	for _, target := range buildTargets(ctx) {
		name, err := nameFor(ctx, target)
		if err != nil {
			return err
		}
		ctx.Archives[target.String()] = name

		sem <- true
		target := target
		g.Go(func() error {
			defer func() {
				<-sem
			}()
			return build(ctx, name, target)
		})
	}
	if err := g.Wait(); err != nil {
		return err
	}
	return runHook(ctx.Config.Build.Env, ctx.Config.Build.Hooks.Post)
}

func runHook(env []string, hook string) error {
	if hook == "" {
		return nil
	}
	log.WithField("hook", hook).Info("running hook")
	cmd := strings.Fields(hook)
	return run(runtimeTarget, cmd, env)
}

func build(ctx *context.Context, name string, target buildTarget) error {
	output := filepath.Join(
		ctx.Config.Dist,
		name,
		ctx.Config.Build.Binary+ext.For(target.goos),
	)
	log.WithField("binary", output).Info("building")
	cmd := []string{"go", "build"}
	if ctx.Config.Build.Flags != "" {
		cmd = append(cmd, strings.Fields(ctx.Config.Build.Flags)...)
	}
	flags, err := ldflags(ctx)
	if err != nil {
		return err
	}
	cmd = append(cmd, "-ldflags="+flags, "-o", output, ctx.Config.Build.Main)
	return run(target, cmd, ctx.Config.Build.Env)
}

func run(target buildTarget, command, env []string) error {
	var cmd = exec.Command(command[0], command[1:]...)
	env = append(env, "GOOS="+target.goos, "GOARCH="+target.goarch, "GOARM="+target.goarm)
	var log = log.WithField("target", target.PrettyString()).
		WithField("env", env).
		WithField("cmd", command)
	cmd.Env = append(cmd.Env, os.Environ()...)
	cmd.Env = append(cmd.Env, env...)
	log.Debug("running")
	if out, err := cmd.CombinedOutput(); err != nil {
		log.WithError(err).Debug("failed")
		return fmt.Errorf("build failed for %s:\n%v", target.PrettyString(), string(out))
	}
	return nil
}
<|endoftext|>


<commit_before>package database

import (
	"errors"
	"fmt"

	"github.com/graphql-go/relay"
	"gopkg.in/mgutz/dat.v1"
)

const OrderOnCreatedAt = "created_at"

var (
	DefaultLimit                  = 50
	ErrScopeInvalidBeforeAndAfter = errors.New("You cannot use before and after in the same query")
	ErrScopeInvalidFirstAndLast   = errors.New("You cannot use first and last in the same query")
)

type GraphQLCursor interface {
	Cursor() relay.ConnectionCursor
}

type GraphQLScope struct {
	relay.ConnectionArguments
	relay.ArraySliceMetaInfo
	Limit   int
	Order   string
	OrderBy string
}

func NewGraphQLScope() GraphQLScope {
	filters := map[string]interface{}{}
	return NewGraphQLScopeWithFilters(filters)
}

func NewGraphQLScopeWithFilters(filters map[string]interface{}) GraphQLScope {
	args := relay.NewConnectionArguments(filters)
	scope := GraphQLScope{
		ConnectionArguments: args,
		Limit:               DefaultLimit,
		Order:               "ASC",
		OrderBy:             "id",
	}

	if filters["order"] != nil {
		scope.Order = filters["order"].(string)
	}

	return scope
}

func ApplyGraphQLScope(builder *dat.SelectBuilder, scope GraphQLScope) (*dat.SelectBuilder, error) {
	// Strongly discouraged from using both
	if scope.Before != "" && scope.After != "" {
		return nil, ErrScopeInvalidBeforeAndAfter
	}

	if scope.First != -1 && scope.Last != -1 {
		return nil, ErrScopeInvalidFirstAndLast
	}

	if scope.After != "" {
		after, err := relay.CursorToOffset(scope.After)
		if err != nil {
			return nil, err
		}

		if scope.Order == "DESC" {
			builder = builder.Where("id <= $1", after)
		} else {
			builder = builder.Where("id >= $1", after)
		}
	}

	if scope.First != -1 {
		scope.Limit = scope.First + 1
	}

	if scope.Before != "" {
		before, err := relay.CursorToOffset(scope.Before)
		if err != nil {
			return nil, err
		}

		if scope.Order == "DESC" {
			builder = builder.Where("id >= $1", before)
		} else {
			builder = builder.Where("id <= $1", before)
		}
	}

	if scope.Last != -1 {
		scope.Limit = scope.Last + 1

		// Flip the order of the results as we want to go in the reverse direction
		if scope.Order == "ASC" {
			scope.Order = "DESC"
		} else {
			scope.Order = "ASC"
		}
	}

	if scope.Limit != -1 {
		builder = builder.Limit(uint64(scope.Limit))
	}

	builder = ApplyOrder(builder, scope)

	return builder, nil
}

func ApplyOrder(builder *dat.SelectBuilder, scope GraphQLScope) *dat.SelectBuilder {
	if scope.OrderBy != "" && scope.Order != "" {
		sql := fmt.Sprintf("%s %s", scope.OrderBy, scope.Order)
		builder = builder.OrderBy(sql)
	}

	return builder
}

func GraphQLConnection(arraySlice []GraphQLCursor, scope GraphQLScope) *relay.Connection {
	if scope.Limit == -1 {
		conn := relay.NewConnection()
		conn.PageInfo = relay.PageInfo{
			StartCursor:     "",
			EndCursor:       "",
			HasPreviousPage: false,
			HasNextPage:     false,
		}

		edges := []*relay.Edge{}
		for _, value := range arraySlice {
			edges = append(edges, &relay.Edge{
				Cursor: value.Cursor(),
				Node:   value,
			})
		}

		conn.Edges = edges

		return conn
	}

	var startCursor, endCursor relay.ConnectionCursor
	args := scope.ConnectionArguments
	// Make sure we're within the bounds of
	limit := min(DefaultLimit, len(arraySlice))
	begin, end := 0, limit

	if args.First != -1 {
		// There are more pages
		if len(arraySlice) == args.First+1 {
			// We don't want to grab the last edge that was used for pagination
			end = end - 1
			endCursor = arraySlice[args.First].Cursor()
		}
	}

	if args.Last != -1 {
		// There are more pages
		if len(arraySlice) == args.Last+1 {
			end = args.Last
			startCursor = arraySlice[args.Last].Cursor()
		}
	}

	slice := arraySlice[begin:end]

	edges := []*relay.Edge{}
	for _, value := range slice {
		edges = append(edges, &relay.Edge{
			Cursor: value.Cursor(),
			Node:   value,
		})
	}

	conn := relay.NewConnection()
	conn.Edges = edges
	conn.PageInfo = relay.PageInfo{
		StartCursor:     startCursor,
		EndCursor:       endCursor,
		HasPreviousPage: startCursor != "",
		HasNextPage:     endCursor != "",
	}

	return conn
}

func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}
<commit_msg>Keep original reference to args for other queries<commit_after>package database

import (
	"errors"
	"fmt"

	"github.com/graphql-go/relay"
	"gopkg.in/mgutz/dat.v1"
)

const OrderOnCreatedAt = "created_at"

var (
	DefaultLimit                  = 50
	ErrScopeInvalidBeforeAndAfter = errors.New("You cannot use before and after in the same query")
	ErrScopeInvalidFirstAndLast   = errors.New("You cannot use first and last in the same query")
)

type GraphQLCursor interface {
	Cursor() relay.ConnectionCursor
}

type GraphQLScope struct {
	relay.ConnectionArguments
	relay.ArraySliceMetaInfo
	Args    map[string]interface{}
	Limit   int
	Order   string
	OrderBy string
}

func NewGraphQLScope() GraphQLScope {
	filters := map[string]interface{}{}
	return NewGraphQLScopeWithFilters(filters)
}

func NewGraphQLScopeWithFilters(args map[string]interface{}) GraphQLScope {
	scope := GraphQLScope{
		Args:                args,
		ConnectionArguments: relay.NewConnectionArguments(args),
		Limit:               DefaultLimit,
		Order:               "ASC",
		OrderBy:             "id",
	}

	if scope.Args["order"] != nil {
		scope.Order = scope.Args["order"].(string)
	}

	return scope
}

func ApplyGraphQLScope(builder *dat.SelectBuilder, scope GraphQLScope) (*dat.SelectBuilder, error) {
	// Strongly discouraged from using both
	if scope.Before != "" && scope.After != "" {
		return nil, ErrScopeInvalidBeforeAndAfter
	}

	if scope.First != -1 && scope.Last != -1 {
		return nil, ErrScopeInvalidFirstAndLast
	}

	if scope.After != "" {
		after, err := relay.CursorToOffset(scope.After)
		if err != nil {
			return nil, err
		}

		if scope.Order == "DESC" {
			builder = builder.Where("id <= $1", after)
		} else {
			builder = builder.Where("id >= $1", after)
		}
	}

	if scope.First != -1 {
		scope.Limit = scope.First + 1
	}

	if scope.Before != "" {
		before, err := relay.CursorToOffset(scope.Before)
		if err != nil {
			return nil, err
		}

		if scope.Order == "DESC" {
			builder = builder.Where("id >= $1", before)
		} else {
			builder = builder.Where("id <= $1", before)
		}
	}

	if scope.Last != -1 {
		scope.Limit = scope.Last + 1

		// Flip the order of the results as we want to go in the reverse direction
		if scope.Order == "ASC" {
			scope.Order = "DESC"
		} else {
			scope.Order = "ASC"
		}
	}

	if scope.Limit != -1 {
		builder = builder.Limit(uint64(scope.Limit))
	}

	builder = ApplyOrder(builder, scope)

	return builder, nil
}

func ApplyOrder(builder *dat.SelectBuilder, scope GraphQLScope) *dat.SelectBuilder {
	if scope.OrderBy != "" && scope.Order != "" {
		sql := fmt.Sprintf("%s %s", scope.OrderBy, scope.Order)
		builder = builder.OrderBy(sql)
	}

	return builder
}

func GraphQLConnection(arraySlice []GraphQLCursor, scope GraphQLScope) *relay.Connection {
	if scope.Limit == -1 {
		conn := relay.NewConnection()
		conn.PageInfo = relay.PageInfo{
			StartCursor:     "",
			EndCursor:       "",
			HasPreviousPage: false,
			HasNextPage:     false,
		}

		edges := []*relay.Edge{}
		for _, value := range arraySlice {
			edges = append(edges, &relay.Edge{
				Cursor: value.Cursor(),
				Node:   value,
			})
		}

		conn.Edges = edges

		return conn
	}

	var startCursor, endCursor relay.ConnectionCursor
	args := scope.ConnectionArguments
	// Make sure we're within the bounds of
	limit := min(DefaultLimit, len(arraySlice))
	begin, end := 0, limit

	if args.First != -1 {
		// There are more pages
		if len(arraySlice) == args.First+1 {
			// We don't want to grab the last edge that was used for pagination
			end = end - 1
			endCursor = arraySlice[args.First].Cursor()
		}
	}

	if args.Last != -1 {
		// There are more pages
		if len(arraySlice) == args.Last+1 {
			end = args.Last
			startCursor = arraySlice[args.Last].Cursor()
		}
	}

	slice := arraySlice[begin:end]

	edges := []*relay.Edge{}
	for _, value := range slice {
		edges = append(edges, &relay.Edge{
			Cursor: value.Cursor(),
			Node:   value,
		})
	}

	conn := relay.NewConnection()
	conn.Edges = edges
	conn.PageInfo = relay.PageInfo{
		StartCursor:     startCursor,
		EndCursor:       endCursor,
		HasPreviousPage: startCursor != "",
		HasNextPage:     endCursor != "",
	}

	return conn
}

func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}
<|endoftext|>


<commit_before>package debuggers

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"path"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/davecgh/go-spew/spew"
	"github.com/solo-io/squash/pkg/client"
	"github.com/solo-io/squash/pkg/client/debugattachment"
	"github.com/solo-io/squash/pkg/models"
	"github.com/solo-io/squash/pkg/platforms"
)

func RunSquashClient(debugger func(string) Debugger, conttopid platforms.Container2Pid) error {
	log.SetLevel(log.DebugLevel)

	customFormatter := new(log.TextFormatter)
	log.SetFormatter(customFormatter)

	log.Info("Squash Client started")

	server := flag.String("server", os.Getenv("SERVERURL"), "")

	flag.Parse()

	log.WithField("server", *server).Info("handleAttachment")
	u, err := url.Parse(*server)
	if err != nil {
		log.WithField("err", err).Error("RunDebugBridge")
		return err

	}
	cfg := &client.TransportConfig{
		BasePath: path.Join(u.Path, client.DefaultBasePath),
		Host:     u.Host,
		Schemes:  []string{u.Scheme},
	}
	log.WithField("cfg", cfg).Debug("creating client")
	client := client.NewHTTPClientWithConfig(nil, cfg)

	return NewDebugHandler(client, debugger, conttopid).handleAttachments()
}

type DebugHandler struct {
	debugger  func(string) Debugger
	conttopid platforms.Container2Pid
	client    *client.Squash
	debugees  map[int]bool
}

func NewDebugHandler(client *client.Squash, debugger func(string) Debugger,
	conttopid platforms.Container2Pid) *DebugHandler {
	return &DebugHandler{
		client:    client,
		debugger:  debugger,
		conttopid: conttopid,
		debugees:  make(map[int]bool),
	}

}

func getNodeName() string {
	return os.Getenv("NODE_NAME")
}

func (d *DebugHandler) handleAttachments() error {
	for {
		err := d.handleAttachment()
		if err != nil {
			log.WithField("err", err).Warn("error watching for attached container")
		}
	}
}
func (d *DebugHandler) handleAttachment() error {
	attachments, err := d.watchForAttached()

	if err != nil {
		log.WithField("err", err).Warn("error watching for attached container")
		return err
	}

	for _, attachment := range attachments {
		// notify the server that we are attaching, so we won't get the same attachment object next time.
		if err := d.notifyState(attachment, models.DebugAttachmentStatusStateAttaching); err != nil {
			log.WithFields(log.Fields{"attachment.Name": attachment.Metadata.Name, "err": err}).Debug("Failed set state to attaching in squash server. aborting.")

			d.notifyError(attachment)
		}
		go d.handleSingleAttachment(attachment)
	}
	return nil
}

func (d *DebugHandler) handleSingleAttachment(attachment *models.DebugAttachment) {

	err := retry(func() error { return d.tryToAttach(attachment) })

	if err != nil {
		log.WithFields(log.Fields{"attachment.Name": attachment.Metadata.Name}).Debug("Failed to attach... signaling server.")
		d.notifyError(attachment)
	}
}

func retry(f func() error) error {
	tries := 3
	for i := 0; i < (tries - 1); i++ {
		if err := f(); err == nil {
			return nil
		}
		time.Sleep(time.Second)
	}
	return f()
}

func (d *DebugHandler) tryToAttach(attachment *models.DebugAttachment) error {

	// make sure this is not a duplicate

	pid, err := d.conttopid.GetPid(context.Background(), attachment.Spec.Attachment)

	if err != nil {
		log.WithField("err", err).Warn("FindFirstProcess error")
		return err
	}

	log.WithField("app", attachment).Info("Attaching to live session")

	p, err := os.FindProcess(pid)
	if err != nil {
		log.WithField("err", err).Error("can't find process")
		return err
	}
	if !d.debugees[pid] {
		log.WithField("pid", pid).Info("starting to debug")
		d.debugees[pid] = true
		err := d.startDebug(attachment, p)
		if err != nil {
			d.notifyError(attachment)
		}
	} else {
		log.WithField("pid", pid).Warn("Already debugging pid. ignoring")
		d.notifyError(attachment)
	}
	return nil
}

func (d *DebugHandler) notifyError(attachment *models.DebugAttachment) {
	d.notifyState(attachment, models.DebugAttachmentStatusStateError)
}

func (d *DebugHandler) notifyState(attachment *models.DebugAttachment, newstate string) error {

	attachmentCopy := *attachment

	params := debugattachment.NewPatchDebugAttachmentParams()
	if attachmentCopy.Status == nil {
		attachmentCopy.Status = &models.DebugAttachmentStatus{}
	}
	attachmentCopy.Status.State = newstate
	params.Body = &attachmentCopy
	params.DebugAttachmentID = attachment.Metadata.Name

	log.WithFields(log.Fields{"patchDebugAttachment": params.Body, "DebugAttachmentID": params.DebugAttachmentID}).Debug("Notifying server of attachment to debug config object")

	_, err := d.client.Debugattachment.PatchDebugAttachment(params)
	if err != nil {
		log.WithField("err", err).Warn("Error notifing debug session attachment - detaching!")
	} else {
		log.Info("debug attachment notified of attachment!")
	}
	return err
}

func (d *DebugHandler) startDebug(attachment *models.DebugAttachment, p *os.Process) error {
	log.Info("start debug called")

	curdebugger := d.debugger(attachment.Spec.Debugger)

	if curdebugger == nil {
		return errors.New("debugger doesn't exist")
	}

	log.WithFields(log.Fields{"curdebugger": attachment.Spec.Debugger}).Info("start debug params")

	log.WithFields(log.Fields{"pid": p.Pid}).Info("starting debug server")
	var err error
	debugServer, err := curdebugger.Attach(p.Pid)

	if err != nil {
		log.WithField("err", err).Error("Starting debug server error")
		return err
	}

	log.WithField("pid", p.Pid).Info("StartDebugServer - posting debug session")

	attachmentPatch := &models.DebugAttachment{
		Metadata: attachment.Metadata,
		Spec:     attachment.Spec,
	}

	podName := ""
	switch debugServer.PodType() {
	case DebugPodTypeTarget:
		att, ok := attachment.Spec.Attachment.(map[string]string)
		if ok {
			podName = att["pod"]
		}
	case DebugPodTypeClient:
		podName = os.Getenv("HOST_ADDR")
	}

	if len(podName) == 0 {
		err = fmt.Errorf("Cannot find POD name for type: %d", debugServer.PodType())
		log.WithField("err", err).Error("Starting debug server error")
		return err
	}

	attachmentPatch.Status = &models.DebugAttachmentStatus{
		DebugServerAddress: fmt.Sprintf("%s:%d", podName, debugServer.Port()),
		State:              models.DebugAttachmentStatusStateAttached,
	}
	params := debugattachment.NewPatchDebugAttachmentParams()
	params.Body = attachmentPatch
	params.DebugAttachmentID = attachment.Metadata.Name

	log.WithFields(log.Fields{"patchDebugAttachment": params.Body, "DebugAttachmentID": params.DebugAttachmentID}).Debug("Notifying server of attachment to debug config object")
	_, err = d.client.Debugattachment.PatchDebugAttachment(params)

	if err != nil {
		log.WithField("err", err).Warn("Error adding debug session - detaching!")
		debugServer.Detach()
	} else {
		log.Info("debug session added!")
	}
	return nil
}

func (d *DebugHandler) watchForAttached() ([]*models.DebugAttachment, error) {
	for {
		params := debugattachment.NewGetDebugAttachmentsParams()
		nodename := getNodeName()
		params.Node = &nodename
		t := true
		params.Wait = &t
		none := models.DebugAttachmentStatusStateNone
		params.State = &none
		log.WithField("params", params).Debug("watchForAttached - calling PopContainerToDebug")

		resp, err := d.client.Debugattachment.GetDebugAttachments(params)

		if _, ok := err.(*debugattachment.GetDebugAttachmentsRequestTimeout); ok {
			continue
		}

		if err != nil {
			log.WithField("err", err).Warn("watchForAttached - error calling function:")
			time.Sleep(time.Second)
			continue
		}

		attachment := resp.Payload

		log.WithField("attachment", spew.Sdump(attachment)).Info("watchForAttached - got debug attachment!")

		return attachment, nil
	}
}
<commit_msg>Fix getting target pod name logic<commit_after>package debuggers

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"path"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/davecgh/go-spew/spew"
	"github.com/solo-io/squash/pkg/client"
	"github.com/solo-io/squash/pkg/client/debugattachment"
	"github.com/solo-io/squash/pkg/models"
	"github.com/solo-io/squash/pkg/platforms"
)

func RunSquashClient(debugger func(string) Debugger, conttopid platforms.Container2Pid) error {
	log.SetLevel(log.DebugLevel)

	customFormatter := new(log.TextFormatter)
	log.SetFormatter(customFormatter)

	log.Info("Squash Client started")

	server := flag.String("server", os.Getenv("SERVERURL"), "")

	flag.Parse()

	log.WithField("server", *server).Info("handleAttachment")
	u, err := url.Parse(*server)
	if err != nil {
		log.WithField("err", err).Error("RunDebugBridge")
		return err

	}
	cfg := &client.TransportConfig{
		BasePath: path.Join(u.Path, client.DefaultBasePath),
		Host:     u.Host,
		Schemes:  []string{u.Scheme},
	}
	log.WithField("cfg", cfg).Debug("creating client")
	client := client.NewHTTPClientWithConfig(nil, cfg)

	return NewDebugHandler(client, debugger, conttopid).handleAttachments()
}

type DebugHandler struct {
	debugger  func(string) Debugger
	conttopid platforms.Container2Pid
	client    *client.Squash
	debugees  map[int]bool
}

func NewDebugHandler(client *client.Squash, debugger func(string) Debugger,
	conttopid platforms.Container2Pid) *DebugHandler {
	return &DebugHandler{
		client:    client,
		debugger:  debugger,
		conttopid: conttopid,
		debugees:  make(map[int]bool),
	}

}

func getNodeName() string {
	return os.Getenv("NODE_NAME")
}

func (d *DebugHandler) handleAttachments() error {
	for {
		err := d.handleAttachment()
		if err != nil {
			log.WithField("err", err).Warn("error watching for attached container")
		}
	}
}
func (d *DebugHandler) handleAttachment() error {
	attachments, err := d.watchForAttached()

	if err != nil {
		log.WithField("err", err).Warn("error watching for attached container")
		return err
	}

	for _, attachment := range attachments {
		// notify the server that we are attaching, so we won't get the same attachment object next time.
		if err := d.notifyState(attachment, models.DebugAttachmentStatusStateAttaching); err != nil {
			log.WithFields(log.Fields{"attachment.Name": attachment.Metadata.Name, "err": err}).Debug("Failed set state to attaching in squash server. aborting.")

			d.notifyError(attachment)
		}
		go d.handleSingleAttachment(attachment)
	}
	return nil
}

func (d *DebugHandler) handleSingleAttachment(attachment *models.DebugAttachment) {

	err := retry(func() error { return d.tryToAttach(attachment) })

	if err != nil {
		log.WithFields(log.Fields{"attachment.Name": attachment.Metadata.Name}).Debug("Failed to attach... signaling server.")
		d.notifyError(attachment)
	}
}

func retry(f func() error) error {
	tries := 3
	for i := 0; i < (tries - 1); i++ {
		if err := f(); err == nil {
			return nil
		}
		time.Sleep(time.Second)
	}
	return f()
}

func (d *DebugHandler) tryToAttach(attachment *models.DebugAttachment) error {

	// make sure this is not a duplicate

	pid, err := d.conttopid.GetPid(context.Background(), attachment.Spec.Attachment)

	if err != nil {
		log.WithField("err", err).Warn("FindFirstProcess error")
		return err
	}

	log.WithField("app", attachment).Info("Attaching to live session")

	p, err := os.FindProcess(pid)
	if err != nil {
		log.WithField("err", err).Error("can't find process")
		return err
	}
	if !d.debugees[pid] {
		log.WithField("pid", pid).Info("starting to debug")
		d.debugees[pid] = true
		err := d.startDebug(attachment, p)
		if err != nil {
			d.notifyError(attachment)
		}
	} else {
		log.WithField("pid", pid).Warn("Already debugging pid. ignoring")
		d.notifyError(attachment)
	}
	return nil
}

func (d *DebugHandler) notifyError(attachment *models.DebugAttachment) {
	d.notifyState(attachment, models.DebugAttachmentStatusStateError)
}

func (d *DebugHandler) notifyState(attachment *models.DebugAttachment, newstate string) error {

	attachmentCopy := *attachment

	params := debugattachment.NewPatchDebugAttachmentParams()
	if attachmentCopy.Status == nil {
		attachmentCopy.Status = &models.DebugAttachmentStatus{}
	}
	attachmentCopy.Status.State = newstate
	params.Body = &attachmentCopy
	params.DebugAttachmentID = attachment.Metadata.Name

	log.WithFields(log.Fields{"patchDebugAttachment": params.Body, "DebugAttachmentID": params.DebugAttachmentID}).Debug("Notifying server of attachment to debug config object")

	_, err := d.client.Debugattachment.PatchDebugAttachment(params)
	if err != nil {
		log.WithField("err", err).Warn("Error notifing debug session attachment - detaching!")
	} else {
		log.Info("debug attachment notified of attachment!")
	}
	return err
}

func (d *DebugHandler) startDebug(attachment *models.DebugAttachment, p *os.Process) error {
	log.Info("start debug called")

	curdebugger := d.debugger(attachment.Spec.Debugger)

	if curdebugger == nil {
		return errors.New("debugger doesn't exist")
	}

	log.WithFields(log.Fields{"curdebugger": attachment.Spec.Debugger}).Info("start debug params")

	log.WithFields(log.Fields{"pid": p.Pid}).Info("starting debug server")
	var err error
	debugServer, err := curdebugger.Attach(p.Pid)

	if err != nil {
		log.WithField("err", err).Error("Starting debug server error")
		return err
	}

	log.WithField("pid", p.Pid).Info("StartDebugServer - posting debug session")

	attachmentPatch := &models.DebugAttachment{
		Metadata: attachment.Metadata,
		Spec:     attachment.Spec,
	}

	podName := ""
	switch debugServer.PodType() {
	case DebugPodTypeTarget:
		att, ok := attachment.Spec.Attachment.(map[string]interface{})
		if ok {
			podName, _ = att["pod"].(string)
		}
	case DebugPodTypeClient:
		podName = os.Getenv("HOST_ADDR")
	}

	if len(podName) == 0 {
		err = fmt.Errorf("Cannot find POD name for type: %d", debugServer.PodType())
		log.WithField("err", err).Error("Starting debug server error")
		return err
	}

	attachmentPatch.Status = &models.DebugAttachmentStatus{
		DebugServerAddress: fmt.Sprintf("%s:%d", podName, debugServer.Port()),
		State:              models.DebugAttachmentStatusStateAttached,
	}
	params := debugattachment.NewPatchDebugAttachmentParams()
	params.Body = attachmentPatch
	params.DebugAttachmentID = attachment.Metadata.Name

	log.WithFields(log.Fields{"patchDebugAttachment": params.Body, "DebugAttachmentID": params.DebugAttachmentID}).Debug("Notifying server of attachment to debug config object")
	_, err = d.client.Debugattachment.PatchDebugAttachment(params)

	if err != nil {
		log.WithField("err", err).Warn("Error adding debug session - detaching!")
		debugServer.Detach()
	} else {
		log.Info("debug session added!")
	}
	return nil
}

func (d *DebugHandler) watchForAttached() ([]*models.DebugAttachment, error) {
	for {
		params := debugattachment.NewGetDebugAttachmentsParams()
		nodename := getNodeName()
		params.Node = &nodename
		t := true
		params.Wait = &t
		none := models.DebugAttachmentStatusStateNone
		params.State = &none
		log.WithField("params", params).Debug("watchForAttached - calling PopContainerToDebug")

		resp, err := d.client.Debugattachment.GetDebugAttachments(params)

		if _, ok := err.(*debugattachment.GetDebugAttachmentsRequestTimeout); ok {
			continue
		}

		if err != nil {
			log.WithField("err", err).Warn("watchForAttached - error calling function:")
			time.Sleep(time.Second)
			continue
		}

		attachment := resp.Payload

		log.WithField("attachment", spew.Sdump(attachment)).Info("watchForAttached - got debug attachment!")

		return attachment, nil
	}
}
<|endoftext|>


<commit_before>package install

import (
	"errors"
	"fmt"
	"net"
	"os"
	"time"
)

// ValidatePlan runs validation against the installation plan to ensure
// that the plan contains valid user input. Returns true, nil if the validation
// is successful. Otherwise, returns false and a collection of validation errors.
func ValidatePlan(p *Plan) (bool, []error) {
	v := newValidator()
	v.validate(p)
	return v.valid()
}

type validatable interface {
	validate() (bool, []error)
}

type validator struct {
	errs []error
}

func newValidator() *validator {
	return &validator{
		errs: []error{},
	}
}

func (v *validator) addError(err ...error) {
	v.errs = append(v.errs, err...)
}

func (v *validator) validate(obj validatable) {
	if ok, err := obj.validate(); !ok {
		v.addError(err...)
	}
}

func (v *validator) validateWithErrPrefix(prefix string, obj validatable) {
	if ok, err := obj.validate(); !ok {
		newErrs := make([]error, len(err), len(err))
		for i, err := range err {
			newErrs[i] = fmt.Errorf("%s: %v", prefix, err)
		}
		v.addError(newErrs...)
	}
}

func (v *validator) valid() (bool, []error) {
	if len(v.errs) > 0 {
		return false, v.errs
	}
	return true, nil
}

func (p *Plan) validate() (bool, []error) {
	v := newValidator()

	v.validate(&p.Cluster)
	v.validateWithErrPrefix("Etcd nodes", &p.Etcd)
	v.validateWithErrPrefix("Master nodes", &p.Master)
	v.validateWithErrPrefix("Worker nodes", &p.Worker)

	return v.valid()
}

func (c *Cluster) validate() (bool, []error) {
	v := newValidator()
	if c.Name == "" {
		v.addError(errors.New("Cluster name cannot be empty"))
	}
	if c.AdminPassword == "" {
		v.addError(errors.New("Admin password cannot be empty"))
	}
	v.validate(&c.Networking)
	v.validate(&c.Certificates)
	v.validate(&c.SSH)

	return v.valid()
}

func (n *NetworkConfig) validate() (bool, []error) {
	v := newValidator()
	if n.Type == "" {
		v.addError(errors.New("Networking type cannot be empty"))
	}
	if n.Type != "bridged" && n.Type != "overlay" {
		v.addError(fmt.Errorf("Invalid networking type %q was provided", n.Type))
	}
	if n.PodCIDRBlock == "" {
		v.addError(errors.New("Pod CIDR block cannot be empty"))
	}
	if _, _, err := net.ParseCIDR(n.PodCIDRBlock); n.PodCIDRBlock != "" && err != nil {
		v.addError(fmt.Errorf("Invalid Pod CIDR block provided: %v", err))
	}

	if n.ServiceCIDRBlock == "" {
		v.addError(errors.New("Service CIDR block cannot be empty"))
	}
	if _, _, err := net.ParseCIDR(n.ServiceCIDRBlock); n.ServiceCIDRBlock != "" && err != nil {
		v.addError(fmt.Errorf("Invalid Service CIDR block provided: %v", err))
	}
	return v.valid()
}

func (c *CertsConfig) validate() (bool, []error) {
	v := newValidator()
	if _, err := time.ParseDuration(c.Expiry); err != nil {
		v.addError(fmt.Errorf("Invalid certificate expiry %q provided: %v", c.Expiry, err))
	}

	if c.LocationCity == "" {
		v.addError(errors.New("Certificate location_city field is required"))
	}
	if c.LocationState == "" {
		v.addError(errors.New("Certificate location_state field is required"))
	}
	if c.LocationCountry == "" {
		v.addError(errors.New("Certificate location_country field is required"))
	}
	return v.valid()
}

func (s *SSHConfig) validate() (bool, []error) {
	v := newValidator()
	if s.User == "" {
		v.addError(errors.New("SSH user field is required"))
	}
	if s.Key == "" {
		v.addError(errors.New("SSH key field is required"))
	}
	if _, err := os.Stat(s.Key); os.IsNotExist(err) {
		v.addError(fmt.Errorf("SSH Key file was not found at %q", s.Key))
	}
	if s.Port < 1 || s.Port > 65535 {
		v.addError(fmt.Errorf("SSH port %d is invalid. Port must be in the range 1-65535", s.Port))
	}
	return v.valid()
}

func (ng *NodeGroup) validate() (bool, []error) {
	v := newValidator()
	if len(ng.Nodes) != ng.ExpectedCount {
		v.addError(fmt.Errorf("Expected node count (%d) does not match the number of nodes provided (%d)", ng.ExpectedCount, len(ng.Nodes)))
	}
	for i, n := range ng.Nodes {
		v.validateWithErrPrefix(fmt.Sprintf("Node #%d", i+1), &n)
	}
	return v.valid()
}

func (mng *MasterNodeGroup) validate() (bool, []error) {
	v := newValidator()
	v.validate(&NodeGroup{mng.ExpectedCount, mng.Nodes})

	if mng.LoadBalancedFQDN == "" {
		v.addError(fmt.Errorf("Load balanced FQDN is required"))
	}

	if mng.LoadBalancedShortName == "" {
		v.addError(fmt.Errorf("Load balanced shortname is required"))
	}

	return v.valid()
}

func (n *Node) validate() (bool, []error) {
	v := newValidator()
	if n.Host == "" {
		v.addError(fmt.Errorf("Node host field is required"))
	}
	if n.IP == "" {
		v.addError(fmt.Errorf("Node IP field is required"))
	}
	if ip := net.ParseIP(n.IP); ip == nil {
		v.addError(fmt.Errorf("Invalid IP provided"))
	}
	return v.valid()
}
<commit_msg>KIS-60: More validation around nodeGroup having empty fields<commit_after>package install

import (
	"errors"
	"fmt"
	"net"
	"os"
	"time"
)

// ValidatePlan runs validation against the installation plan to ensure
// that the plan contains valid user input. Returns true, nil if the validation
// is successful. Otherwise, returns false and a collection of validation errors.
func ValidatePlan(p *Plan) (bool, []error) {
	v := newValidator()
	v.validate(p)
	return v.valid()
}

type validatable interface {
	validate() (bool, []error)
}

type validator struct {
	errs []error
}

func newValidator() *validator {
	return &validator{
		errs: []error{},
	}
}

func (v *validator) addError(err ...error) {
	v.errs = append(v.errs, err...)
}

func (v *validator) validate(obj validatable) {
	if ok, err := obj.validate(); !ok {
		v.addError(err...)
	}
}

func (v *validator) validateWithErrPrefix(prefix string, obj validatable) {
	if ok, err := obj.validate(); !ok {
		newErrs := make([]error, len(err), len(err))
		for i, err := range err {
			newErrs[i] = fmt.Errorf("%s: %v", prefix, err)
		}
		v.addError(newErrs...)
	}
}

func (v *validator) valid() (bool, []error) {
	if len(v.errs) > 0 {
		return false, v.errs
	}
	return true, nil
}

func (p *Plan) validate() (bool, []error) {
	v := newValidator()

	v.validate(&p.Cluster)
	v.validateWithErrPrefix("Etcd nodes", &p.Etcd)
	v.validateWithErrPrefix("Master nodes", &p.Master)
	v.validateWithErrPrefix("Worker nodes", &p.Worker)

	return v.valid()
}

func (c *Cluster) validate() (bool, []error) {
	v := newValidator()
	if c.Name == "" {
		v.addError(errors.New("Cluster name cannot be empty"))
	}
	if c.AdminPassword == "" {
		v.addError(errors.New("Admin password cannot be empty"))
	}
	v.validate(&c.Networking)
	v.validate(&c.Certificates)
	v.validate(&c.SSH)

	return v.valid()
}

func (n *NetworkConfig) validate() (bool, []error) {
	v := newValidator()
	if n.Type == "" {
		v.addError(errors.New("Networking type cannot be empty"))
	}
	if n.Type != "bridged" && n.Type != "overlay" {
		v.addError(fmt.Errorf("Invalid networking type %q was provided", n.Type))
	}
	if n.PodCIDRBlock == "" {
		v.addError(errors.New("Pod CIDR block cannot be empty"))
	}
	if _, _, err := net.ParseCIDR(n.PodCIDRBlock); n.PodCIDRBlock != "" && err != nil {
		v.addError(fmt.Errorf("Invalid Pod CIDR block provided: %v", err))
	}

	if n.ServiceCIDRBlock == "" {
		v.addError(errors.New("Service CIDR block cannot be empty"))
	}
	if _, _, err := net.ParseCIDR(n.ServiceCIDRBlock); n.ServiceCIDRBlock != "" && err != nil {
		v.addError(fmt.Errorf("Invalid Service CIDR block provided: %v", err))
	}
	return v.valid()
}

func (c *CertsConfig) validate() (bool, []error) {
	v := newValidator()
	if _, err := time.ParseDuration(c.Expiry); err != nil {
		v.addError(fmt.Errorf("Invalid certificate expiry %q provided: %v", c.Expiry, err))
	}

	if c.LocationCity == "" {
		v.addError(errors.New("Certificate location_city field is required"))
	}
	if c.LocationState == "" {
		v.addError(errors.New("Certificate location_state field is required"))
	}
	if c.LocationCountry == "" {
		v.addError(errors.New("Certificate location_country field is required"))
	}
	return v.valid()
}

func (s *SSHConfig) validate() (bool, []error) {
	v := newValidator()
	if s.User == "" {
		v.addError(errors.New("SSH user field is required"))
	}
	if s.Key == "" {
		v.addError(errors.New("SSH key field is required"))
	}
	if _, err := os.Stat(s.Key); os.IsNotExist(err) {
		v.addError(fmt.Errorf("SSH Key file was not found at %q", s.Key))
	}
	if s.Port < 1 || s.Port > 65535 {
		v.addError(fmt.Errorf("SSH port %d is invalid. Port must be in the range 1-65535", s.Port))
	}
	return v.valid()
}

func (ng *NodeGroup) validate() (bool, []error) {
	v := newValidator()
	if ng == nil || len(ng.Nodes) <= 0 {
		v.addError(fmt.Errorf("Atleast one node is required"))
	}
	if ng.ExpectedCount <= 0 {
		v.addError(fmt.Errorf("Expected node count must be greater than 0"))
	}
	if len(ng.Nodes) != ng.ExpectedCount && (len(ng.Nodes) > 0 && ng.ExpectedCount > 0) {
		v.addError(fmt.Errorf("Expected node count (%d) does not match the number of nodes provided (%d)", ng.ExpectedCount, len(ng.Nodes)))
	}
	for i, n := range ng.Nodes {
		v.validateWithErrPrefix(fmt.Sprintf("Node #%d", i+1), &n)
	}
	return v.valid()
}

func (mng *MasterNodeGroup) validate() (bool, []error) {
	v := newValidator()
	v.validate(&NodeGroup{mng.ExpectedCount, mng.Nodes})

	if mng.LoadBalancedFQDN == "" {
		v.addError(fmt.Errorf("Load balanced FQDN is required"))
	}

	if mng.LoadBalancedShortName == "" {
		v.addError(fmt.Errorf("Load balanced shortname is required"))
	}

	return v.valid()
}

func (n *Node) validate() (bool, []error) {
	v := newValidator()
	if n.Host == "" {
		v.addError(fmt.Errorf("Node host field is required"))
	}
	if n.IP == "" {
		v.addError(fmt.Errorf("Node IP field is required"))
	}
	if ip := net.ParseIP(n.IP); ip == nil {
		v.addError(fmt.Errorf("Invalid IP provided"))
	}
	return v.valid()
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Juniper Networks, Inc. All rights reserved.

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

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

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

package network

import (
	"strings"

	"github.com/golang/glog"

	"github.com/Juniper/contrail-go-api"
	"github.com/Juniper/contrail-go-api/types"
)

type InstanceManager interface {
	LocateInstance(namespace, podName string) *types.VirtualMachine
	LocateInterface(network *types.VirtualNetwork, instance *types.VirtualMachine) *types.VirtualMachineInterface
	LocateInstanceIp(network *types.VirtualNetwork, nic *types.VirtualMachineInterface) *types.InstanceIp
}

type InstanceManagerImpl struct {
	client    contrail.ApiClient
	allocator AddressAllocator
}

func NewInstanceManager(client contrail.ApiClient, allocator AddressAllocator) InstanceManager {
	manager := new(InstanceManagerImpl)
	manager.client = client
	manager.allocator = allocator
	return manager
}

func instanceFQName(tenant, podName string) []string {
	fqn := []string{DefaultDomain, tenant, podName}
	return fqn
}

func (m *InstanceManagerImpl) LocateInstance(tenant, podName string) *types.VirtualMachine {
	fqn := instanceFQName(tenant, podName)
	instance, err := types.VirtualMachineByName(m.client, strings.Join(fqn, ":"))
	if err == nil {
		return instance
	}

	instance = new(types.VirtualMachine)
	instance.SetFQName("project", fqn)
	err = m.client.Create(instance)
	if err != nil {
		glog.Errorf("Create %s: %v", podName, err)
		return nil
	}
	return instance
}

func (m *InstanceManagerImpl) DeleteInstance(uid string) error {
	err := m.client.DeleteByUuid("virtual-machine", uid)
	return err
}

func interfaceFQName(namespace, podName string) []string {
	fqn := []string{DefaultDomain, namespace, podName}
	return fqn
}

func (m *InstanceManagerImpl) LookupInterface(namespace, podName string) *types.VirtualMachineInterface {
	fqn := interfaceFQName(namespace, podName)
	obj, err := m.client.FindByName(
		"virtual-machine-interface", strings.Join(fqn, ":"))
	if err != nil {
		glog.Infof("Get vmi %s: %v", podName, err)
		return nil
	}
	return obj.(*types.VirtualMachineInterface)
}

func (m *InstanceManagerImpl) LocateInterface(
	network *types.VirtualNetwork, instance *types.VirtualMachine) *types.VirtualMachineInterface {
	namespace := instance.GetFQName()[len(instance.GetFQName())-2]
	fqn := interfaceFQName(namespace, instance.GetName())

	obj, err := m.client.FindByName(
		"virtual-machine-interface", strings.Join(fqn, ":"))

	if err == nil {
		nic := obj.(*types.VirtualMachineInterface)
		// TODO(prm): ensure network is as expected, else update.
		return nic
	}

	nic := new(types.VirtualMachineInterface)
	nic.SetFQName("project", fqn)
	nic.AddVirtualMachine(instance)
	if network != nil {
		nic.AddVirtualNetwork(network)
	}
	err = m.client.Create(nic)
	if err != nil {
		glog.Errorf("Create interface %s: %v", instance.GetName(), err)
		return nil
	}
	obj, err = m.client.FindByUuid(nic.GetType(), nic.GetUuid())
	if err != nil {
		glog.Errorf("Get vmi %s: %v", nic.GetUuid(), err)
		return nil
	}
	return nic
}

func (m *InstanceManagerImpl) ReleaseInterface(namespace, podName string) {
	fqn := interfaceFQName(namespace, podName)
	obj, err := m.client.FindByName("virtual-machine-interface", strings.Join(fqn, ":"))
	if err != nil {
		glog.Errorf("Get vmi %s: %v", strings.Join(fqn, ":"), err)
		return
	}
	vmi := obj.(*types.VirtualMachineInterface)
	refs, err := vmi.GetFloatingIpBackRefs()
	if err == nil {
		for _, ref := range refs {
			err = m.client.DeleteByUuid("floating-ip", ref.Uuid)
			if err != nil {
				glog.Errorf("Delete floating-ip %s: %v", ref.Uuid, err)
			}
		}
	} else {
		glog.Errorf("Get %s floating-ip back refs: %v", podName, err)
	}
	err = m.client.Delete(obj)
	if err != nil {
		glog.Errorf("Delete vmi %s: %v", obj.GetUuid(), err)
	}
}

func makeInstanceIpFQName(namespace, nicName string) []string {
	return []string{DefaultDomain, namespace, nicName}
}

func (m *InstanceManagerImpl) LocateInstanceIp(
	network *types.VirtualNetwork, nic *types.VirtualMachineInterface) *types.InstanceIp {
	namespace := nic.GetFQName()[len(nic.GetFQName())-2]
	fqn := makeInstanceIpFQName(namespace, nic.GetName())
	obj, err := m.client.FindByName("instance-ip", strings.Join(fqn, ":"))
	if err == nil {
		// TODO(prm): ensure that attributes are as expected
		return obj.(*types.InstanceIp)
	}

	address, err := m.allocator.LocateIpAddress(nic.GetUuid())
	if err != nil {
		return nil
	}

	// Create InstanceIp
	ipObj := new(types.InstanceIp)
	ipObj.SetFQName("", fqn)
	ipObj.AddVirtualNetwork(network)
	ipObj.AddVirtualMachineInterface(nic)
	ipObj.SetInstanceIpAddress(address)
	err = m.client.Create(ipObj)
	if err != nil {
		glog.Errorf("Create instance-ip %s: %v", nic.GetName(), err)
		return nil
	}
	obj, err = m.client.FindByUuid(ipObj.GetType(), ipObj.GetUuid())
	if err != nil {
		glog.Errorf("Get instance-ip %s: %v", ipObj.GetUuid(), err)
		return nil
	}
	return ipObj
}

func (m *InstanceManagerImpl) ReleaseInstanceIp(namespace, nicName, instanceUID string) {
	fqn := makeInstanceIpFQName(namespace, nicName)
	uid, err := m.client.UuidByName("instance-ip", strings.Join(fqn, ":"))
	if err != nil {
		glog.Errorf("Get instance-ip %s: %v", strings.Join(fqn, ":"), err)
		return
	}
	err = m.client.DeleteByUuid("instance-ip", uid)
	if err != nil {
		glog.Errorf("Delete instance-ip %s: %v", uid, err)
	}

	m.allocator.ReleaseIpAddress(instanceUID)
}

func (m *InstanceManagerImpl) AttachFloatingIp(
	podName, projectName string, floatingIp *types.FloatingIp) {

	fqn := append(strings.Split(projectName, ":"), podName)
	obj, err := m.client.FindByName(
		"virtual-machine-interface", strings.Join(fqn, ":"))
	if err != nil {
		glog.Errorf("GET vmi %s: %v", podName, err)
		return
	}

	vmi := obj.(*types.VirtualMachineInterface)

	refs, err := floatingIp.GetVirtualMachineInterfaceRefs()
	if err != nil {
		glog.Errorf("GET floating-ip %s: %v", floatingIp.GetUuid(), err)
		return
	}
	for _, ref := range refs {
		if ref.Uuid == vmi.GetUuid() {
			return
		}
	}

	floatingIp.AddVirtualMachineInterface(vmi)
	err = m.client.Update(floatingIp)
	if err != nil {
		glog.Errorf("Update floating-ip %s: %v", podName, err)
	}
}
<commit_msg>instance-ips have no name hierarchy<commit_after>/*
Copyright 2015 Juniper Networks, Inc. All rights reserved.

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

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

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

package network

import (
	"strings"

	"github.com/golang/glog"

	"github.com/Juniper/contrail-go-api"
	"github.com/Juniper/contrail-go-api/types"
)

type InstanceManager interface {
	LocateInstance(namespace, podName string) *types.VirtualMachine
	LocateInterface(network *types.VirtualNetwork, instance *types.VirtualMachine) *types.VirtualMachineInterface
	LocateInstanceIp(network *types.VirtualNetwork, nic *types.VirtualMachineInterface) *types.InstanceIp
}

type InstanceManagerImpl struct {
	client    contrail.ApiClient
	allocator AddressAllocator
}

func NewInstanceManager(client contrail.ApiClient, allocator AddressAllocator) InstanceManager {
	manager := new(InstanceManagerImpl)
	manager.client = client
	manager.allocator = allocator
	return manager
}

func instanceFQName(tenant, podName string) []string {
	fqn := []string{DefaultDomain, tenant, podName}
	return fqn
}

func (m *InstanceManagerImpl) LocateInstance(tenant, podName string) *types.VirtualMachine {
	fqn := instanceFQName(tenant, podName)
	instance, err := types.VirtualMachineByName(m.client, strings.Join(fqn, ":"))
	if err == nil {
		return instance
	}

	instance = new(types.VirtualMachine)
	instance.SetFQName("project", fqn)
	err = m.client.Create(instance)
	if err != nil {
		glog.Errorf("Create %s: %v", podName, err)
		return nil
	}
	return instance
}

func (m *InstanceManagerImpl) DeleteInstance(uid string) error {
	err := m.client.DeleteByUuid("virtual-machine", uid)
	return err
}

func interfaceFQName(namespace, podName string) []string {
	fqn := []string{DefaultDomain, namespace, podName}
	return fqn
}

func (m *InstanceManagerImpl) LookupInterface(namespace, podName string) *types.VirtualMachineInterface {
	fqn := interfaceFQName(namespace, podName)
	obj, err := m.client.FindByName(
		"virtual-machine-interface", strings.Join(fqn, ":"))
	if err != nil {
		glog.Infof("Get vmi %s: %v", podName, err)
		return nil
	}
	return obj.(*types.VirtualMachineInterface)
}

func (m *InstanceManagerImpl) LocateInterface(
	network *types.VirtualNetwork, instance *types.VirtualMachine) *types.VirtualMachineInterface {
	namespace := instance.GetFQName()[len(instance.GetFQName())-2]
	fqn := interfaceFQName(namespace, instance.GetName())

	obj, err := m.client.FindByName(
		"virtual-machine-interface", strings.Join(fqn, ":"))

	if err == nil {
		nic := obj.(*types.VirtualMachineInterface)
		// TODO(prm): ensure network is as expected, else update.
		return nic
	}

	nic := new(types.VirtualMachineInterface)
	nic.SetFQName("project", fqn)
	nic.AddVirtualMachine(instance)
	if network != nil {
		nic.AddVirtualNetwork(network)
	}
	err = m.client.Create(nic)
	if err != nil {
		glog.Errorf("Create interface %s: %v", instance.GetName(), err)
		return nil
	}
	obj, err = m.client.FindByUuid(nic.GetType(), nic.GetUuid())
	if err != nil {
		glog.Errorf("Get vmi %s: %v", nic.GetUuid(), err)
		return nil
	}
	return nic
}

func (m *InstanceManagerImpl) ReleaseInterface(namespace, podName string) {
	fqn := interfaceFQName(namespace, podName)
	obj, err := m.client.FindByName("virtual-machine-interface", strings.Join(fqn, ":"))
	if err != nil {
		glog.Errorf("Get vmi %s: %v", strings.Join(fqn, ":"), err)
		return
	}
	vmi := obj.(*types.VirtualMachineInterface)
	refs, err := vmi.GetFloatingIpBackRefs()
	if err == nil {
		for _, ref := range refs {
			err = m.client.DeleteByUuid("floating-ip", ref.Uuid)
			if err != nil {
				glog.Errorf("Delete floating-ip %s: %v", ref.Uuid, err)
			}
		}
	} else {
		glog.Errorf("Get %s floating-ip back refs: %v", podName, err)
	}
	err = m.client.Delete(obj)
	if err != nil {
		glog.Errorf("Delete vmi %s: %v", obj.GetUuid(), err)
	}
}

func makeInstanceIpName(tenant, nicName string) string {
	return tenant + "_" + nicName
}

func (m *InstanceManagerImpl) LocateInstanceIp(
	network *types.VirtualNetwork, nic *types.VirtualMachineInterface) *types.InstanceIp {
	tenant := nic.GetFQName()[len(nic.GetFQName())-2]
	ipName := makeInstanceIpName(tenant, nic.GetName())
	obj, err := m.client.FindByName("instance-ip", ipName)
	if err == nil {
		// TODO(prm): ensure that attributes are as expected
		return obj.(*types.InstanceIp)
	}

	address, err := m.allocator.LocateIpAddress(nic.GetUuid())
	if err != nil {
		return nil
	}

	// Create InstanceIp
	ipObj := new(types.InstanceIp)
	ipObj.SetName(ipName)
	ipObj.AddVirtualNetwork(network)
	ipObj.AddVirtualMachineInterface(nic)
	ipObj.SetInstanceIpAddress(address)
	err = m.client.Create(ipObj)
	if err != nil {
		glog.Errorf("Create instance-ip %s: %v", nic.GetName(), err)
		return nil
	}
	obj, err = m.client.FindByUuid(ipObj.GetType(), ipObj.GetUuid())
	if err != nil {
		glog.Errorf("Get instance-ip %s: %v", ipObj.GetUuid(), err)
		return nil
	}
	return ipObj
}

func (m *InstanceManagerImpl) ReleaseInstanceIp(namespace, nicName, instanceUID string) {
	ipName := makeInstanceIpName(namespace, nicName)
	uid, err := m.client.UuidByName("instance-ip", ipName)
	if err != nil {
		glog.Errorf("Get instance-ip %s: %v", ipName, err)
		return
	}
	err = m.client.DeleteByUuid("instance-ip", uid)
	if err != nil {
		glog.Errorf("Delete instance-ip %s: %v", uid, err)
	}

	m.allocator.ReleaseIpAddress(instanceUID)
}

func (m *InstanceManagerImpl) AttachFloatingIp(
	podName, projectName string, floatingIp *types.FloatingIp) {

	fqn := append(strings.Split(projectName, ":"), podName)
	obj, err := m.client.FindByName(
		"virtual-machine-interface", strings.Join(fqn, ":"))
	if err != nil {
		glog.Errorf("GET vmi %s: %v", podName, err)
		return
	}

	vmi := obj.(*types.VirtualMachineInterface)

	refs, err := floatingIp.GetVirtualMachineInterfaceRefs()
	if err != nil {
		glog.Errorf("GET floating-ip %s: %v", floatingIp.GetUuid(), err)
		return
	}
	for _, ref := range refs {
		if ref.Uuid == vmi.GetUuid() {
			return
		}
	}

	floatingIp.AddVirtualMachineInterface(vmi)
	err = m.client.Update(floatingIp)
	if err != nil {
		glog.Errorf("Update floating-ip %s: %v", podName, err)
	}
}
<|endoftext|>


<commit_before>package pubsub // import "github.com/docker/docker/pkg/pubsub"

import "github.com/moby/pubsub"

// NewPublisher creates a new pub/sub publisher to broadcast messages.
// The duration is used as the send timeout as to not block the publisher publishing
// messages to other clients if one client is slow or unresponsive.
// The buffer is used when creating new channels for subscribers.
//
// Deprecated: use github.com/moby/pubsub.NewPublisher
var NewPublisher = pubsub.NewPublisher

// Publisher is basic pub/sub structure. Allows to send events and subscribe
// to them. Can be safely used from multiple goroutines.
//
// Deprecated: use github.com/moby/pubsub.Publisher
type Publisher = pubsub.Publisher
<commit_msg>remove aliases for deprecated pkg/pubsub<commit_after><|endoftext|>


<commit_before>package search

import (
	"github.com/ewhal/nyaa/model"
	"github.com/ewhal/nyaa/service/torrent"
	"github.com/ewhal/nyaa/util/log"
	"html"
	"net/http"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

type SearchParam struct {
	Category string
	Order    string
	Query    string
	Max      int
	Status   string
	Sort     string
}

func SearchByQuery(r *http.Request, pagenum int) (SearchParam, []model.Torrents, int) {
	maxPerPage, errConv := strconv.Atoi(r.URL.Query().Get("max"))
	if errConv != nil {
		maxPerPage = 50 // default Value maxPerPage
	}

	if maxPerPage > 300 {
		maxPerPage = 300
	}
	search_param := SearchParam{}
	search_param.Max = maxPerPage
	search_param.Query = r.URL.Query().Get("q")
	search_param.Category = r.URL.Query().Get("c")
	search_param.Status = r.URL.Query().Get("s")
	search_param.Sort = r.URL.Query().Get("sort")
	search_param.Order = r.URL.Query().Get("order")

	catsSplit := strings.Split(search_param.Category, "_")
	// need this to prevent out of index panics
	var searchCatId, searchSubCatId string
	if len(catsSplit) == 2 {

		searchCatId = html.EscapeString(catsSplit[0])
		searchSubCatId = html.EscapeString(catsSplit[1])
	}

	switch search_param.Sort {
	case "torrent_name":
		search_param.Sort = "torrent_name"
		break
	case "date":
		search_param.Sort = "date"
		break
	case "downloads":
		search_param.Sort = "downloads"
		break
	case "filesize":
		search_param.Sort = "filesize"
	case "torrent_id":
	default:
		search_param.Sort = "torrent_id"
	}

	switch search_param.Order {
	case "asc":
		search_param.Order = "asc"
		break
	case "desc":
	default:
		search_param.Order = "desc"
	}

	order_by := search_param.Sort + " " + search_param.Order

	parameters := torrentService.WhereParams{}
	conditions := []string{}
	if searchCatId != "" {
		conditions = append(conditions, "category_id = ?")
		parameters.Params = append(parameters.Params, searchCatId)
	}
	if searchSubCatId != "" {
		conditions = append(conditions, "sub_category_id = ?")
		parameters.Params = append(parameters.Params, searchSubCatId)
	}
	if search_param.Status != "" {
		if search_param.Status == "2" {
			conditions = append(conditions, "status_id != ?")
		} else {
			conditions = append(conditions, "status_id = ?")
		}
		parameters.Params = append(parameters.Params, search_param.Status)
	}
	searchQuerySplit := strings.Fields(search_param.Query)
	for i, word := range searchQuerySplit {
		firstRune, _ := utf8.DecodeRuneInString(word)
		if len(word) == 1 && unicode.IsPunct(firstRune) {
			// some queries have a single punctuation character
			// which causes a full scan instead of using the index
			// and yields no meaningful results.
			// due to len() == 1 we're just looking at 1-byte/ascii
			// punctuation characters.
			continue
		}
		// TODO: make this faster ?
		conditions = append(conditions, "torrent_name ILIKE ?")
		parameters.Params = append(parameters.Params, "%"+searchQuerySplit[i]+"%")
	}

	parameters.Conditions = strings.Join(conditions[:], " AND ")
	log.Infof("SQL query is :: %s\n", parameters.Conditions)
	torrents, n := torrentService.GetTorrentsOrderBy(&parameters, order_by, maxPerPage, maxPerPage*(pagenum-1))
	return search_param, torrents, n
}
<commit_msg>Added user torrents search<commit_after>package search

import (
	"github.com/ewhal/nyaa/model"
	"github.com/ewhal/nyaa/service/torrent"
	"github.com/ewhal/nyaa/util/log"
	"html"
	"net/http"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

type SearchParam struct {
	Category string
	Order    string
	Query    string
	Max      int
	Status   string
	Sort     string
	UserId   string
}

func SearchByQuery(r *http.Request, pagenum int) (SearchParam, []model.Torrents, int) {
	maxPerPage, errConv := strconv.Atoi(r.URL.Query().Get("max"))
	if errConv != nil {
		maxPerPage = 50 // default Value maxPerPage
	}

	if maxPerPage > 300 {
		maxPerPage = 300
	}
	search_param := SearchParam{}
	search_param.Max = maxPerPage
	search_param.Query = r.URL.Query().Get("q")
	search_param.Category = r.URL.Query().Get("c")
	search_param.Status = r.URL.Query().Get("s")
	search_param.Sort = r.URL.Query().Get("sort")
	search_param.Order = r.URL.Query().Get("order")
	search_param.UserId = r.URL.Query().Get("userId")

	catsSplit := strings.Split(search_param.Category, "_")
	// need this to prevent out of index panics
	var searchCatId, searchSubCatId string
	if len(catsSplit) == 2 {

		searchCatId = html.EscapeString(catsSplit[0])
		searchSubCatId = html.EscapeString(catsSplit[1])
	}

	switch search_param.Sort {
	case "torrent_name":
		search_param.Sort = "torrent_name"
		break
	case "date":
		search_param.Sort = "date"
		break
	case "downloads":
		search_param.Sort = "downloads"
		break
	case "filesize":
		search_param.Sort = "filesize"
	case "torrent_id":
	default:
		search_param.Sort = "torrent_id"
	}

	switch search_param.Order {
	case "asc":
		search_param.Order = "asc"
		break
	case "desc":
	default:
		search_param.Order = "desc"
	}

	order_by := search_param.Sort + " " + search_param.Order

	parameters := torrentService.WhereParams{}
	conditions := []string{}
	if searchCatId != "" {
		conditions = append(conditions, "category_id = ?")
		parameters.Params = append(parameters.Params, searchCatId)
	}
	if searchSubCatId != "" {
		conditions = append(conditions, "sub_category_id = ?")
		parameters.Params = append(parameters.Params, searchSubCatId)
	}
	if search_param.Status != "" {
		if search_param.Status == "2" {
			conditions = append(conditions, "status_id != ?")
		} else {
			conditions = append(conditions, "status_id = ?")
		}
		parameters.Params = append(parameters.Params, search_param.Status)
	}
	if search_param.UserId != "" {

			conditions = append(conditions, "owner_id = ?")
		parameters.Params = append(parameters.Params, search_param.UserId)
	}
	searchQuerySplit := strings.Fields(search_param.Query)
	for i, word := range searchQuerySplit {
		firstRune, _ := utf8.DecodeRuneInString(word)
		if len(word) == 1 && unicode.IsPunct(firstRune) {
			// some queries have a single punctuation character
			// which causes a full scan instead of using the index
			// and yields no meaningful results.
			// due to len() == 1 we're just looking at 1-byte/ascii
			// punctuation characters.
			continue
		}
		// TODO: make this faster ?
		conditions = append(conditions, "torrent_name ILIKE ?")
		parameters.Params = append(parameters.Params, "%"+searchQuerySplit[i]+"%")
	}

	parameters.Conditions = strings.Join(conditions[:], " AND ")
	log.Infof("SQL query is :: %s\n", parameters.Conditions)
	torrents, n := torrentService.GetTorrentsOrderBy(&parameters, order_by, maxPerPage, maxPerPage*(pagenum-1))
	return search_param, torrents, n
}
<|endoftext|>


<commit_before>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
)

// AWSSQSBroker represents a AWS SQS broker
// There are examples on: https://docs.aws.amazon.com/sdk-for-go/v1/developer-guide/sqs-example-create-queue.html
type AWSSQSBroker struct {
	Broker
	processingWG      sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
	receivingWG       sync.WaitGroup
	stopReceivingChan chan int
	sess              *session.Session
	service           sqsiface.SQSAPI
}

// NewAWSSQSBroker creates new Broker instance
func NewAWSSQSBroker(cnf *config.Config) Interface {
	b := &AWSSQSBroker{Broker: New(cnf)}
	if cnf.SQS != nil && cnf.SQS.Client != nil {
		// Use provided *SQS client
		b.service = cnf.SQS.Client
	} else {
		// Initialize a session that the SDK will use to load credentials from the shared credentials file, ~/.aws/credentials.
		// See details on: https://docs.aws.amazon.com/sdk-for-go/v1/developer-guide/configuring-sdk.html
		// Also, env AWS_REGION is also required
		b.sess = session.Must(session.NewSessionWithOptions(session.Options{
			SharedConfigState: session.SharedConfigEnable,
		}))
		b.service = sqs.New(b.sess)
	}

	return b
}

// GetPendingTasks returns a slice of task.Signatures waiting in the queue
func (b *AWSSQSBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
	return nil, errors.New("Not implemented")
}

// StartConsuming enters a loop and waits for incoming messages
func (b *AWSSQSBroker) StartConsuming(consumerTag string, concurrency int, taskProcessor TaskProcessor) (bool, error) {
	b.startConsuming(consumerTag, taskProcessor)
	qURL := b.defaultQueueURL()
	deliveries := make(chan *sqs.ReceiveMessageOutput)

	b.stopReceivingChan = make(chan int)
	b.receivingWG.Add(1)

	go func() {
		defer b.receivingWG.Done()

		log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

		for {
			select {
			// A way to stop this goroutine from b.StopConsuming
			case <-b.stopReceivingChan:
				return
			default:
				output, err := b.receiveMessage(qURL)
				if err != nil {
					log.ERROR.Printf("Queue consume error: %s", err)
					continue
				}
				if len(output.Messages) == 0 {
					continue
				}

				deliveries <- output
			}

			whetherContinue, err := b.continueReceivingMessages(qURL, deliveries)
			if err != nil {
				log.ERROR.Printf("Error when receiving messages. Error: %v", err)
			}
			if whetherContinue == false {
				return
			}
		}
	}()

	if err := b.consume(deliveries, concurrency, taskProcessor); err != nil {
		return b.retry, err
	}

	return b.retry, nil
}

// StopConsuming quits the loop
func (b *AWSSQSBroker) StopConsuming() {
	b.stopConsuming()

	b.stopReceiving()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	// Waiting for the receiving goroutine to have stopped
	b.receivingWG.Wait()
}

// Publish places a new message on the default queue
func (b *AWSSQSBroker) Publish(signature *tasks.Signature) error {

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check that signature.RoutingKey is set, if not switch to DefaultQueue
	AdjustRoutingKey(b, signature)

	MsgInput := &sqs.SendMessageInput{
		MessageBody: aws.String(string(msg)),
		QueueUrl:    aws.String(b.cnf.Broker + "/" + signature.RoutingKey),
	}

	// if this is a fifo queue, there needs to be some additional parameters.
	if strings.HasSuffix(signature.RoutingKey, ".fifo") {
		// Use Machinery's signature Task UUID as SQS Message Group ID.
		MsgDedupID := signature.UUID
		MsgInput.MessageDeduplicationId = aws.String(MsgDedupID)

		// Use Machinery's signature Group UUID as SQS Message Group ID.
		MsgGroupID := signature.GroupUUID
		MsgInput.MessageGroupId = aws.String(MsgGroupID)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// and is not a fifo queue, set a delay in seconds for the task.
	if signature.ETA != nil && !strings.HasSuffix(signature.RoutingKey, ".fifo") {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			MsgInput.DelaySeconds = aws.Int64(signature.ETA.Unix() - now.Unix())
		}
	}

	result, err := b.service.SendMessage(MsgInput)

	if err != nil {
		log.ERROR.Printf("Error when sending a message: %v", err)
		return err

	}
	log.INFO.Printf("Sending a message successfully, the messageId is %v", *result.MessageId)
	return nil

}

// consume is a method which keeps consuming deliveries from a channel, until there is an error or a stop signal
func (b *AWSSQSBroker) consume(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		b.initializePool(pool, concurrency)
	}()

	errorsChan := make(chan error)

	for {
		whetherContinue, err := b.consumeDeliveries(deliveries, concurrency, taskProcessor, pool, errorsChan)
		if err != nil {
			return err
		}
		if whetherContinue == false {
			return nil
		}
	}
}

// consumeOne is a method consumes a delivery. If a delivery was consumed successfully, it will be deleted from AWS SQS
func (b *AWSSQSBroker) consumeOne(delivery *sqs.ReceiveMessageOutput, taskProcessor TaskProcessor) error {
	if len(delivery.Messages) == 0 {
		log.ERROR.Printf("received an empty message, the delivery was %v", delivery)
		return errors.New("received empty message, the delivery is " + delivery.GoString())
	}

	sig := new(tasks.Signature)
	decoder := json.NewDecoder(strings.NewReader(*delivery.Messages[0].Body))
	decoder.UseNumber()
	if err := decoder.Decode(sig); err != nil {
		log.ERROR.Printf("unmarshal error. the delivery is %v", delivery)
		return err
	}

	// If the task is not registered return an error
	// and leave the message in the queue
	if !b.IsTaskRegistered(sig.Name) {
		return fmt.Errorf("task %s is not registered", sig.Name)
	}

	// Delete message after consuming successfully
	err := b.deleteOne(delivery)
	if err != nil {
		log.ERROR.Printf("error when deleting the delivery. the delivery is %v", delivery)
	}

	return taskProcessor.Process(sig)
}

// deleteOne is a method delete a delivery from AWS SQS
func (b *AWSSQSBroker) deleteOne(delivery *sqs.ReceiveMessageOutput) error {
	qURL := b.defaultQueueURL()
	_, err := b.service.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      qURL,
		ReceiptHandle: delivery.Messages[0].ReceiptHandle,
	})

	if err != nil {
		return err
	}
	return nil
}

// defaultQueueURL is a method returns the default queue url
func (b *AWSSQSBroker) defaultQueueURL() *string {
	return aws.String(b.cnf.Broker + "/" + b.cnf.DefaultQueue)
}

// receiveMessage is a method receives a message from specified queue url
func (b *AWSSQSBroker) receiveMessage(qURL *string) (*sqs.ReceiveMessageOutput, error) {
	var waitTimeSeconds int
	if b.cnf.SQS != nil {
		waitTimeSeconds = b.cnf.SQS.WaitTimeSeconds
	} else {
		waitTimeSeconds = 0
	}
	result, err := b.service.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            qURL,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(int64(b.cnf.ResultsExpireIn)), // 10 hours
		WaitTimeSeconds:     aws.Int64(int64(waitTimeSeconds)),
	})
	if err != nil {
		return nil, err
	}
	return result, err
}

// initializePool is a method which initializes concurrency pool
func (b *AWSSQSBroker) initializePool(pool chan struct{}, concurrency int) {
	for i := 0; i < concurrency; i++ {
		pool <- struct{}{}
	}
}

// consumeDeliveries is a method consuming deliveries from deliveries channel
func (b *AWSSQSBroker) consumeDeliveries(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor, pool chan struct{}, errorsChan chan error) (bool, error) {
	select {
	case err := <-errorsChan:
		return false, err
	case d := <-deliveries:
		if concurrency > 0 {
			// get worker from pool (blocks until one is available)
			<-pool
		}

		b.processingWG.Add(1)

		// Consume the task inside a goroutine so multiple tasks
		// can be processed concurrently
		go func() {

			if err := b.consumeOne(d, taskProcessor); err != nil {
				errorsChan <- err
			}

			b.processingWG.Done()

			if concurrency > 0 {
				// give worker back to pool
				pool <- struct{}{}
			}
		}()
	case <-b.stopChan:
		return false, nil
	}
	return true, nil
}

// continueReceivingMessages is a method returns a continue signal
func (b *AWSSQSBroker) continueReceivingMessages(qURL *string, deliveries chan *sqs.ReceiveMessageOutput) (bool, error) {
	select {
	// A way to stop this goroutine from b.StopConsuming
	case <-b.stopReceivingChan:
		return false, nil
	default:
		output, err := b.receiveMessage(qURL)
		if err != nil {
			return true, err
		}
		if len(output.Messages) == 0 {
			return true, nil
		}
		go func() { deliveries <- output }()
	}
	return true, nil
}

// stopReceiving is a method sending a signal to stopReceivingChan
func (b *AWSSQSBroker) stopReceiving() {
	// Stop the receiving goroutine
	b.stopReceivingChan <- 1
}
<commit_msg>Only delete SQS messages after at task successfully processes the message Addresses https://github.com/RichardKnop/machinery/issues/253 for SQS<commit_after>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
)

// AWSSQSBroker represents a AWS SQS broker
// There are examples on: https://docs.aws.amazon.com/sdk-for-go/v1/developer-guide/sqs-example-create-queue.html
type AWSSQSBroker struct {
	Broker
	processingWG      sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
	receivingWG       sync.WaitGroup
	stopReceivingChan chan int
	sess              *session.Session
	service           sqsiface.SQSAPI
}

// NewAWSSQSBroker creates new Broker instance
func NewAWSSQSBroker(cnf *config.Config) Interface {
	b := &AWSSQSBroker{Broker: New(cnf)}
	if cnf.SQS != nil && cnf.SQS.Client != nil {
		// Use provided *SQS client
		b.service = cnf.SQS.Client
	} else {
		// Initialize a session that the SDK will use to load credentials from the shared credentials file, ~/.aws/credentials.
		// See details on: https://docs.aws.amazon.com/sdk-for-go/v1/developer-guide/configuring-sdk.html
		// Also, env AWS_REGION is also required
		b.sess = session.Must(session.NewSessionWithOptions(session.Options{
			SharedConfigState: session.SharedConfigEnable,
		}))
		b.service = sqs.New(b.sess)
	}

	return b
}

// GetPendingTasks returns a slice of task.Signatures waiting in the queue
func (b *AWSSQSBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
	return nil, errors.New("Not implemented")
}

// StartConsuming enters a loop and waits for incoming messages
func (b *AWSSQSBroker) StartConsuming(consumerTag string, concurrency int, taskProcessor TaskProcessor) (bool, error) {
	b.startConsuming(consumerTag, taskProcessor)
	qURL := b.defaultQueueURL()
	deliveries := make(chan *sqs.ReceiveMessageOutput)

	b.stopReceivingChan = make(chan int)
	b.receivingWG.Add(1)

	go func() {
		defer b.receivingWG.Done()

		log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

		for {
			select {
			// A way to stop this goroutine from b.StopConsuming
			case <-b.stopReceivingChan:
				return
			default:
				output, err := b.receiveMessage(qURL)
				if err != nil {
					log.ERROR.Printf("Queue consume error: %s", err)
					continue
				}
				if len(output.Messages) == 0 {
					continue
				}

				deliveries <- output
			}

			whetherContinue, err := b.continueReceivingMessages(qURL, deliveries)
			if err != nil {
				log.ERROR.Printf("Error when receiving messages. Error: %v", err)
			}
			if whetherContinue == false {
				return
			}
		}
	}()

	if err := b.consume(deliveries, concurrency, taskProcessor); err != nil {
		return b.retry, err
	}

	return b.retry, nil
}

// StopConsuming quits the loop
func (b *AWSSQSBroker) StopConsuming() {
	b.stopConsuming()

	b.stopReceiving()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	// Waiting for the receiving goroutine to have stopped
	b.receivingWG.Wait()
}

// Publish places a new message on the default queue
func (b *AWSSQSBroker) Publish(signature *tasks.Signature) error {

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check that signature.RoutingKey is set, if not switch to DefaultQueue
	AdjustRoutingKey(b, signature)

	MsgInput := &sqs.SendMessageInput{
		MessageBody: aws.String(string(msg)),
		QueueUrl:    aws.String(b.cnf.Broker + "/" + signature.RoutingKey),
	}

	// if this is a fifo queue, there needs to be some additional parameters.
	if strings.HasSuffix(signature.RoutingKey, ".fifo") {
		// Use Machinery's signature Task UUID as SQS Message Group ID.
		MsgDedupID := signature.UUID
		MsgInput.MessageDeduplicationId = aws.String(MsgDedupID)

		// Use Machinery's signature Group UUID as SQS Message Group ID.
		MsgGroupID := signature.GroupUUID
		MsgInput.MessageGroupId = aws.String(MsgGroupID)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// and is not a fifo queue, set a delay in seconds for the task.
	if signature.ETA != nil && !strings.HasSuffix(signature.RoutingKey, ".fifo") {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			MsgInput.DelaySeconds = aws.Int64(signature.ETA.Unix() - now.Unix())
		}
	}

	result, err := b.service.SendMessage(MsgInput)

	if err != nil {
		log.ERROR.Printf("Error when sending a message: %v", err)
		return err

	}
	log.INFO.Printf("Sending a message successfully, the messageId is %v", *result.MessageId)
	return nil

}

// consume is a method which keeps consuming deliveries from a channel, until there is an error or a stop signal
func (b *AWSSQSBroker) consume(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		b.initializePool(pool, concurrency)
	}()

	errorsChan := make(chan error)

	for {
		whetherContinue, err := b.consumeDeliveries(deliveries, concurrency, taskProcessor, pool, errorsChan)
		if err != nil {
			return err
		}
		if whetherContinue == false {
			return nil
		}
	}
}

// consumeOne is a method consumes a delivery. If a delivery was consumed successfully, it will be deleted from AWS SQS
func (b *AWSSQSBroker) consumeOne(delivery *sqs.ReceiveMessageOutput, taskProcessor TaskProcessor) error {
	if len(delivery.Messages) == 0 {
		log.ERROR.Printf("received an empty message, the delivery was %v", delivery)
		return errors.New("received empty message, the delivery is " + delivery.GoString())
	}

	sig := new(tasks.Signature)
	decoder := json.NewDecoder(strings.NewReader(*delivery.Messages[0].Body))
	decoder.UseNumber()
	if err := decoder.Decode(sig); err != nil {
		log.ERROR.Printf("unmarshal error. the delivery is %v", delivery)
		return err
	}

	// If the task is not registered return an error
	// and leave the message in the queue
	if !b.IsTaskRegistered(sig.Name) {
		return fmt.Errorf("task %s is not registered", sig.Name)
	}

	err := taskProcessor.Process(sig)
	if err != nil {
		// Delete message after successfully consuming and processing the message
		if err := b.deleteOne(delivery); err != nil {
			log.ERROR.Printf("error when deleting the delivery. the delivery is %v", delivery)
		}
	}
	return err
}

// deleteOne is a method delete a delivery from AWS SQS
func (b *AWSSQSBroker) deleteOne(delivery *sqs.ReceiveMessageOutput) error {
	qURL := b.defaultQueueURL()
	_, err := b.service.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      qURL,
		ReceiptHandle: delivery.Messages[0].ReceiptHandle,
	})

	if err != nil {
		return err
	}
	return nil
}

// defaultQueueURL is a method returns the default queue url
func (b *AWSSQSBroker) defaultQueueURL() *string {
	return aws.String(b.cnf.Broker + "/" + b.cnf.DefaultQueue)
}

// receiveMessage is a method receives a message from specified queue url
func (b *AWSSQSBroker) receiveMessage(qURL *string) (*sqs.ReceiveMessageOutput, error) {
	var waitTimeSeconds int
	if b.cnf.SQS != nil {
		waitTimeSeconds = b.cnf.SQS.WaitTimeSeconds
	} else {
		waitTimeSeconds = 0
	}
	result, err := b.service.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            qURL,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(int64(b.cnf.ResultsExpireIn)), // 10 hours
		WaitTimeSeconds:     aws.Int64(int64(waitTimeSeconds)),
	})
	if err != nil {
		return nil, err
	}
	return result, err
}

// initializePool is a method which initializes concurrency pool
func (b *AWSSQSBroker) initializePool(pool chan struct{}, concurrency int) {
	for i := 0; i < concurrency; i++ {
		pool <- struct{}{}
	}
}

// consumeDeliveries is a method consuming deliveries from deliveries channel
func (b *AWSSQSBroker) consumeDeliveries(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor, pool chan struct{}, errorsChan chan error) (bool, error) {
	select {
	case err := <-errorsChan:
		return false, err
	case d := <-deliveries:
		if concurrency > 0 {
			// get worker from pool (blocks until one is available)
			<-pool
		}

		b.processingWG.Add(1)

		// Consume the task inside a goroutine so multiple tasks
		// can be processed concurrently
		go func() {

			if err := b.consumeOne(d, taskProcessor); err != nil {
				errorsChan <- err
			}

			b.processingWG.Done()

			if concurrency > 0 {
				// give worker back to pool
				pool <- struct{}{}
			}
		}()
	case <-b.stopChan:
		return false, nil
	}
	return true, nil
}

// continueReceivingMessages is a method returns a continue signal
func (b *AWSSQSBroker) continueReceivingMessages(qURL *string, deliveries chan *sqs.ReceiveMessageOutput) (bool, error) {
	select {
	// A way to stop this goroutine from b.StopConsuming
	case <-b.stopReceivingChan:
		return false, nil
	default:
		output, err := b.receiveMessage(qURL)
		if err != nil {
			return true, err
		}
		if len(output.Messages) == 0 {
			return true, nil
		}
		go func() { deliveries <- output }()
	}
	return true, nil
}

// stopReceiving is a method sending a signal to stopReceivingChan
func (b *AWSSQSBroker) stopReceiving() {
	// Stop the receiving goroutine
	b.stopReceivingChan <- 1
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package api

import (
	"strings"
)

func IsPullAlways(p PullPolicy) bool {
	// Default to pull always
	if len(p) == 0 {
		return true
	}
	return pullPoliciesEqual(p, PullAlways)
}

func IsPullNever(p PullPolicy) bool {
	return pullPoliciesEqual(p, PullNever)
}

func IsPullIfNotPresent(p PullPolicy) bool {
	return pullPoliciesEqual(p, PullIfNotPresent)
}

func pullPoliciesEqual(p1, p2 PullPolicy) bool {
	return strings.ToLower(string(p1)) == strings.ToLower(string(p2))
}
<commit_msg>Change the default Pull policy to PullIfNotPresent.<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package api

import (
	"strings"
)

func IsPullAlways(p PullPolicy) bool {
	return pullPoliciesEqual(p, PullAlways)
}

func IsPullNever(p PullPolicy) bool {
	return pullPoliciesEqual(p, PullNever)
}

func IsPullIfNotPresent(p PullPolicy) bool {
	// Default to pull if not present
	if len(p) == 0 {
		return true
	}
	return pullPoliciesEqual(p, PullIfNotPresent)
}

func pullPoliciesEqual(p1, p2 PullPolicy) bool {
	return strings.ToLower(string(p1)) == strings.ToLower(string(p2))
}
<|endoftext|>


<commit_before>package doozer

import (
	"doozer/client"
	"github.com/bmizerany/assert"
	"net"
	"testing"
)

// import "runtime"

func mustListen() net.Listener {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	return l
}

func mustListenPacket(addr string) net.PacketConn {
	c, err := net.ListenPacket("udp", addr)
	if err != nil {
		panic(err)
	}
	return c
}

func TestDoozerSimple(t *testing.T) {
	l := mustListen()
	defer l.Close()
	u := mustListenPacket(l.Addr().String())
	defer u.Close()

	go Main("a", "", u, l, nil)

	cl, err := client.Dial(l.Addr().String())
	assert.Equal(t, nil, err)
	assert.Equal(t, nil, cl.Noop())
}

// func TestWatch(t *testing.T) {
// 	for { runtime.Gosched() }
// }
<commit_msg>use goroutine count for testing<commit_after>package doozer

import (
	"doozer/client"
	"github.com/bmizerany/assert"
	"net"
	"runtime"
	"testing"
)

func mustListen() net.Listener {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	return l
}

func mustListenPacket(addr string) net.PacketConn {
	c, err := net.ListenPacket("udp", addr)
	if err != nil {
		panic(err)
	}
	return c
}

func TestDoozerSimple(t *testing.T) {
	gs := runtime.Goroutines()
	gs = 27 // TODO delete this line

	l := mustListen()
	defer l.Close()
	u := mustListenPacket(l.Addr().String())
	defer u.Close()

	go Main("a", "", u, l, nil)

	cl, err := client.Dial(l.Addr().String())
	assert.Equal(t, nil, err)
	assert.Equal(t, nil, cl.Noop())

	assert.Equal(t, gs, runtime.Goroutines())
}
<|endoftext|>


<commit_before>package ec2

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strconv"
	"strings"

	// Community:
	log "github.com/Sirupsen/logrus"
	"github.com/katosys/kato/pkg/tools"
)

//-----------------------------------------------------------------------------
// func: Add
//-----------------------------------------------------------------------------

// Add a new instance to the cluster.
func (d *Data) Add() {

	// Set current command:
	d.command = "add"

	// Load state from state file:
	if err := d.loadState(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Retrieve the CoreOS AMI ID:
	var err error
	if d.AmiID, err = d.retrieveCoreOSAmiID(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Execute the udata|run pipeline:
	out, err := tools.ExecutePipeline(d.forgeUdataCommand(), d.forgeRunCommand())
	if err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Publish DNS records:
	if err := d.publishDNSRecords(d.Roles, out); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Warning(err)
	}
}

//-----------------------------------------------------------------------------
// func: retrieveCoreOSAmiID
//-----------------------------------------------------------------------------

func (d *Data) retrieveCoreOSAmiID() (string, error) {

	// Send the request:
	res, err := http.Get("https://coreos.com/dist/aws/aws-" +
		d.CoreOSChannel + ".json")
	if err != nil {
		return "", err
	}

	// Retrieve the data:
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}

	// Close the handler:
	if err = res.Body.Close(); err != nil {
		return "", err
	}

	// Decode JSON into Go values:
	var jsonData map[string]interface{}
	if err := json.Unmarshal(data, &jsonData); err != nil {
		return "", err
	}

	// Store the AMI ID:
	amis := jsonData[d.Region].(map[string]interface{})
	amiID := amis["hvm"].(string)

	// Log this action:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": amiID}).
		Info("Latest CoreOS " + d.CoreOSChannel + " AMI located")

	// Return the AMI ID:
	return amiID, nil
}

//-----------------------------------------------------------------------------
// func: forgeUdataCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeUdataCommand() *exec.Cmd {

	// Udata arguments bundle:
	args := []string{"udata",
		"--roles", d.Roles,
		"--cluster-id", d.ClusterID,
		"--cluster-state", d.ClusterState,
		"--quorum-count", strconv.Itoa(d.QuorumCount),
		"--master-count", strconv.Itoa(d.MasterCount),
		"--host-name", d.HostName,
		"--host-id", d.HostID,
		"--domain", d.Domain,
		"--ec2-region", d.Region,
		"--dns-provider", d.DNSProvider,
		"--dns-api-key", d.DNSApiKey,
		"--etcd-token", d.EtcdToken,
		"--calico-ip-pool", d.CalicoIPPool,
		"--rexray-storage-driver", "ebs",
		"--iaas-provider", "ec2",
		"--prometheus",
		"--gzip-udata",
	}

	// Append flags if present:
	if d.SysdigAccessKey != "" {
		args = append(args, "--sysdig-access-key", d.SysdigAccessKey)
	}
	if d.DatadogAPIKey != "" {
		args = append(args, "--datadog-api-key", d.DatadogAPIKey)
	}
	if d.SlackWebhook != "" {
		args = append(args, "--slack-webhook", d.SlackWebhook)
	}
	if d.CaCertPath != "" {
		args = append(args, "--ca-cert-path", d.CaCertPath)
	}
	for _, z := range d.StubZones {
		args = append(args, "--stub-zone", z)
	}
	if d.SMTPURL != "" {
		args = append(args, "--smtp-url", d.SMTPURL)
	}
	if d.AdminEmail != "" {
		args = append(args, "--admin-email", d.AdminEmail)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: forgeRunCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeRunCommand() *exec.Cmd {

	// Ec2 run arguments bundle:
	args := []string{"ec2", "run",
		"--tag-name", d.HostName + "-" + d.HostID + "." + d.Domain,
		"--region", d.Region,
		"--zone", d.Zone,
		"--ami-id", d.AmiID,
		"--instance-type", d.InstanceType,
		"--key-pair", d.KeyPair,
		"--subnet-id", d.ExtSubnetID,
		"--security-group-ids", strings.Join(d.securityGroupIDs(d.Roles), ","),
		"--iam-role", "kato",
		"--source-dest-check", "false",
		"--public-ip", "true",
	}

	// Append flags if present:
	if strings.Contains(d.Roles, "master") {
		i, _ := strconv.Atoi(d.HostID)
		args = append(args, "--private-ip", tools.OffsetIP(d.ExtSubnetCidr, 10+i))
	}
	if strings.Contains(d.Roles, "worker") {
		args = append(args, "--elb-name", d.ClusterID)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: securityGroupIDs
//-----------------------------------------------------------------------------

func (d *Data) securityGroupIDs(roles string) (list []string) {
	for _, role := range strings.Split(roles, ",") {
		switch role {
		case "quorum":
			list = append(list, d.QuorumSecGrp)
		case "master":
			list = append(list, d.MasterSecGrp)
		case "worker":
			list = append(list, d.WorkerSecGrp)
		case "border":
			list = append(list, d.BorderSecGrp)
		}
	}
	return
}

//-----------------------------------------------------------------------------
// func: publishDNSRecords
//-----------------------------------------------------------------------------

func (d *Data) publishDNSRecords(roles string, out []byte) error {

	// Retrieve the instance IPs:
	var dat map[string]string
	if err := json.Unmarshal(out, &dat); err != nil {
		return err
	}

	// For every role in this instance:
	for _, role := range strings.Split(roles, ",") {

		log.Info("katoctl " + d.DNSProvider +
			"--api-key " + d.DNSApiKey +
			"record add --zone " + "int." +
			d.Domain + " " + dat["internal"] +
			":A:" + role + "-" + d.HostID)

		log.Info("katoctl " + d.DNSProvider +
			"--api-key " + d.DNSApiKey +
			"record add --zone " + "ext." +
			d.Domain + " " + dat["external"] +
			":A:" + role + "-" + d.HostID)
	}

	return nil
}
<commit_msg>Publish DNS records<commit_after>package ec2

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"

	// Community:
	log "github.com/Sirupsen/logrus"
	"github.com/katosys/kato/pkg/tools"
)

//-----------------------------------------------------------------------------
// func: Add
//-----------------------------------------------------------------------------

// Add a new instance to the cluster.
func (d *Data) Add() {

	// Set current command:
	d.command = "add"

	// Load state from state file:
	if err := d.loadState(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Retrieve the CoreOS AMI ID:
	var err error
	if d.AmiID, err = d.retrieveCoreOSAmiID(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Execute the udata|run pipeline:
	out, err := tools.ExecutePipeline(d.forgeUdataCommand(), d.forgeRunCommand())
	if err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Publish DNS records:
	if err := d.publishDNSRecords(d.Roles, out); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Warning(err)
	}
}

//-----------------------------------------------------------------------------
// func: retrieveCoreOSAmiID
//-----------------------------------------------------------------------------

func (d *Data) retrieveCoreOSAmiID() (string, error) {

	// Send the request:
	res, err := http.Get("https://coreos.com/dist/aws/aws-" +
		d.CoreOSChannel + ".json")
	if err != nil {
		return "", err
	}

	// Retrieve the data:
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}

	// Close the handler:
	if err = res.Body.Close(); err != nil {
		return "", err
	}

	// Decode JSON into Go values:
	var jsonData map[string]interface{}
	if err := json.Unmarshal(data, &jsonData); err != nil {
		return "", err
	}

	// Store the AMI ID:
	amis := jsonData[d.Region].(map[string]interface{})
	amiID := amis["hvm"].(string)

	// Log this action:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": amiID}).
		Info("Latest CoreOS " + d.CoreOSChannel + " AMI located")

	// Return the AMI ID:
	return amiID, nil
}

//-----------------------------------------------------------------------------
// func: forgeUdataCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeUdataCommand() *exec.Cmd {

	// Udata arguments bundle:
	args := []string{"udata",
		"--roles", d.Roles,
		"--cluster-id", d.ClusterID,
		"--cluster-state", d.ClusterState,
		"--quorum-count", strconv.Itoa(d.QuorumCount),
		"--master-count", strconv.Itoa(d.MasterCount),
		"--host-name", d.HostName,
		"--host-id", d.HostID,
		"--domain", d.Domain,
		"--ec2-region", d.Region,
		"--dns-provider", d.DNSProvider,
		"--dns-api-key", d.DNSApiKey,
		"--etcd-token", d.EtcdToken,
		"--calico-ip-pool", d.CalicoIPPool,
		"--rexray-storage-driver", "ebs",
		"--iaas-provider", "ec2",
		"--prometheus",
		"--gzip-udata",
	}

	// Append flags if present:
	if d.SysdigAccessKey != "" {
		args = append(args, "--sysdig-access-key", d.SysdigAccessKey)
	}
	if d.DatadogAPIKey != "" {
		args = append(args, "--datadog-api-key", d.DatadogAPIKey)
	}
	if d.SlackWebhook != "" {
		args = append(args, "--slack-webhook", d.SlackWebhook)
	}
	if d.CaCertPath != "" {
		args = append(args, "--ca-cert-path", d.CaCertPath)
	}
	for _, z := range d.StubZones {
		args = append(args, "--stub-zone", z)
	}
	if d.SMTPURL != "" {
		args = append(args, "--smtp-url", d.SMTPURL)
	}
	if d.AdminEmail != "" {
		args = append(args, "--admin-email", d.AdminEmail)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: forgeRunCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeRunCommand() *exec.Cmd {

	// Ec2 run arguments bundle:
	args := []string{"ec2", "run",
		"--tag-name", d.HostName + "-" + d.HostID + "." + d.Domain,
		"--region", d.Region,
		"--zone", d.Zone,
		"--ami-id", d.AmiID,
		"--instance-type", d.InstanceType,
		"--key-pair", d.KeyPair,
		"--subnet-id", d.ExtSubnetID,
		"--security-group-ids", strings.Join(d.securityGroupIDs(d.Roles), ","),
		"--iam-role", "kato",
		"--source-dest-check", "false",
		"--public-ip", "true",
	}

	// Append flags if present:
	if strings.Contains(d.Roles, "master") {
		i, _ := strconv.Atoi(d.HostID)
		args = append(args, "--private-ip", tools.OffsetIP(d.ExtSubnetCidr, 10+i))
	}
	if strings.Contains(d.Roles, "worker") {
		args = append(args, "--elb-name", d.ClusterID)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: securityGroupIDs
//-----------------------------------------------------------------------------

func (d *Data) securityGroupIDs(roles string) (list []string) {
	for _, role := range strings.Split(roles, ",") {
		switch role {
		case "quorum":
			list = append(list, d.QuorumSecGrp)
		case "master":
			list = append(list, d.MasterSecGrp)
		case "worker":
			list = append(list, d.WorkerSecGrp)
		case "border":
			list = append(list, d.BorderSecGrp)
		}
	}
	return
}

//-----------------------------------------------------------------------------
// func: publishDNSRecords
//-----------------------------------------------------------------------------

func (d *Data) publishDNSRecords(roles string, out []byte) error {

	// Retrieve the instance IPs:
	var dat map[string]string
	if err := json.Unmarshal(out, &dat); err != nil {
		return err
	}

	// For every role in this instance:
	for _, role := range strings.Split(roles, ",") {

		// Forge the internal record command:
		cmdInt := exec.Command("katoctl", d.DNSProvider,
			"--api-key", d.DNSApiKey,
			"record", "add",
			"--zone", "int."+d.Domain,
			dat["internal"]+":A:"+role+"-"+d.HostID)

		// Execute the internal record command:
		cmdInt.Stderr = os.Stderr
		if err := cmdInt.Run(); err != nil {
			return err
		}

		// Forge the external record command:
		cmdExt := exec.Command("katoctl", d.DNSProvider,
			"--api-key", d.DNSApiKey,
			"record", "add",
			"--zone", "ext."+d.Domain,
			dat["external"]+":A:"+role+"-"+d.HostID)

		// Execute the external record command:
		cmdExt.Stderr = os.Stderr
		if err := cmdExt.Run(); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>