<commit_before>package shark

import(
	"strings"
	"os"
	"libxml"
	"fmt"
	xml "libxml/tree"
	tp "tritium/proto"
	"libxml/xpath"
	"rubex"
)

func (ctx *Ctx) runBuiltIn(fun *Function, scope *Scope, ins *tp.Instruction, args []interface{}) (returnValue interface{}) {
	switch fun.Name {
	case "this":
		returnValue = scope.Value
	case "yield":
		myYieldBlock := ctx.yieldBlock()
		ctx.Yields = ctx.Yields[:(len(ctx.Yields)-1)]
		if (ctx.yieldBlock() != nil) {
			returnValue = ctx.runChildren(scope, myYieldBlock.Ins)
		} else {
			panic("yield() failure")
		}
		ctx.Yields = append(ctx.Yields, myYieldBlock)

	case "var.Text":
		val := ctx.Env[args[0].(string)]
		ts := &Scope{Value: val}
		ctx.runChildren(ts, ins)
		returnValue = ts.Value
		ctx.Env[args[0].(string)] = returnValue.(string)
	case "var.Text.Text":
		ctx.Env[args[0].(string)] = args[1].(string)
		returnValue = args[1].(string)
	case "match.Text":
		// Setup stacks
		ctx.MatchStack = append(ctx.MatchStack, args[0].(string))
		ctx.MatchShouldContinue = append(ctx.MatchShouldContinue, true)

		// Run children
		ctx.runChildren(scope, ins)

		if ctx.matchShouldContinue() {
			returnValue = "false"
		} else {
			returnValue = "true"
		}

		// Clear
		ctx.MatchShouldContinue = ctx.MatchShouldContinue[:len(ctx.MatchShouldContinue)-1]
		ctx.MatchStack = ctx.MatchStack[:len(ctx.MatchStack)-1]
	case "with.Text":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			if args[0].(string) == ctx.matchTarget() {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "with.Regexp":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			//println(matcher.MatchAgainst, matchWith)
			if (args[0].(*rubex.Regexp)).Match([]uint8(ctx.matchTarget())) {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "not.Text":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			if args[0].(string) != ctx.matchTarget() {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "not.Regexp":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			//println(matcher.MatchAgainst, matchWith)
			if !(args[0].(*rubex.Regexp)).Match([]uint8(ctx.matchTarget())) {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "regexp.Text.Text":
		mode := rubex.ONIG_OPTION_DEFAULT
		if strings.Index(args[1].(string), "i") >= 0 {
			mode = rubex.ONIG_OPTION_IGNORECASE
		}
		if strings.Index(args[1].(string), "m") >= 0 {
			mode = rubex.ONIG_OPTION_MULTILINE
		}
		var err os.Error
		returnValue, err = rubex.NewRegexp(args[0].(string), mode)
		if err != nil {
			panic("Invalid regexp")
		}
	case "export.Text":
		val := make([]string, 2)
		val[0] = args[0].(string)
		ts := &Scope{Value:""}
		ctx.runChildren(ts, ins)
		val[1] = ts.Value.(string)
		ctx.Exports = append(ctx.Exports, val)
	case "log.Text":
		ctx.Logs = append(ctx.Logs, args[0].(string))

	// ATOMIC FUNCTIONS
	case "concat.Text.Text":
		//println("Concat:", args[0].(string), "+", args[1].(string))
		returnValue = args[0].(string) + args[1].(string)
	case "concat.Text.Text.Text": //REMOVE
		returnValue = args[0].(string) + args[1].(string) + args[2].(string)
	case "downcase.Text":
		returnValue = strings.ToLower(args[0].(string))
		return
	case "upcase.Text":
		returnValue = strings.ToUpper(args[0].(string))
		return
	case "index.XMLNode":
		returnValue = fmt.Sprintf("%d", scope.Index + 1)

	// TEXT FUNCTIONS
	case "set.Text":
		scope.Value = args[0]
	case "append.Text":
		scope.Value = scope.Value.(string) + args[0].(string)
	case "prepend.Text":
		scope.Value = args[0].(string) + scope.Value.(string)
	case "replace.Text":
		ts := &Scope{Value:""}
		ctx.runChildren(ts, ins)
		scope.Value = strings.Replace(scope.Value.(string), args[0].(string), ts.Value.(string), -1)
	case "replace.Regexp":
		regexp := args[0].(*rubex.Regexp)
		scope.Value = regexp.GsubFunc(scope.Value.(string), func(match string, captures map[string]string) string {
			usesGlobal := (ctx.Env["use_global_replace_vars"] == "true")

			for name, capture := range captures {
				if usesGlobal {
					//println("setting $", name, "to", capture)
					ctx.Env[name] = capture
				}
				ctx.vars()[name] = capture
			}

			replacementScope := &Scope{Value:match}
			ctx.runChildren(replacementScope, ins)
			//println(ins.String())

			//println("Replacement:", replacementScope.Value.(string))
			innerReplacer := rubex.MustCompile(`[\\\$]([\d])`)
			return innerReplacer.GsubFunc(replacementScope.Value.(string), func(_ string, numeric_captures map[string]string) string {
				capture := numeric_captures["1"]
				var val string
				if usesGlobal {
					val = ctx.Env[capture]
				} else {
					val = ctx.vars()[capture].(string)
				}
				return val
		    })
		})
		returnValue = scope.Value

	// XML FUNCTIONS
	case "xml":
		doc := libxml.XmlParseString(scope.Value.(string))
		defer doc.Free()
		ns := &Scope{Value:doc}
		ctx.runChildren(ns, ins)
		scope.Value = doc.String()
		returnValue = scope.Value
	case "html":
		doc := libxml.HtmlParseString(scope.Value.(string))
		defer doc.Free()
		ns := &Scope{Value:doc}
		ctx.runChildren(ns, ins)
		scope.Value = doc.DumpHTML()
		returnValue = scope.Value
	case "html_fragment":
		doc := libxml.HtmlParseFragment(scope.Value.(string))
		defer doc.Free()
		ns := &Scope{Value: doc.RootElement()}
		ctx.runChildren(ns, ins)
		scope.Value = ns.Value.(xml.Node).Content()
		returnValue = scope.Value
	case "select.Text":
		// TODO reuse XPath object
		node := scope.Value.(xml.Node)
		xpCtx := xpath.NewXPath(node.Doc())
		xpath := xpath.CompileXPath(args[0].(string))
		//xpath := ctx.XPath(args[0].(string))
		nodeSet := xpCtx.SearchByCompiledXPath(node, xpath).Slice()
		defer xpCtx.Free()
		if len(nodeSet) == 0 {
			returnValue = "false"
		} else {
			returnValue = "true"
		}

		for index, node := range(nodeSet) {
			if (node != nil) && node.IsLinked() && node.IsValid() {
				ns := &Scope{Value: node, Index: index}
				ctx.runChildren(ns, ins)
			}
		}
	case "position.Text":
		returnValue = Positions[args[0].(string)]

	// SHARED NODE FUNCTIONS
	case "remove":
		scope.Value.(xml.Node).Remove()
	case "inner", "value", "inner_text", "text":
		node := scope.Value.(xml.Node)
		ts := &Scope{Value:node.Content()}
		ctx.runChildren(ts, ins)
		val := ts.Value.(string)
		if node.IsLinked() {
			node.SetContent(val)
		}
		returnValue = val
	case "name":
		node := scope.Value.(xml.Node)
		ts := &Scope{Value:node.Name()}
		ctx.runChildren(ts, ins)
		node.SetName(ts.Value.(string))
		returnValue = ts.Value.(string)
	case "dup":
		node := scope.Value.(xml.Node)
		newNode := node.Duplicate()
		MoveFunc(newNode, node, AFTER)
		ns := &Scope{Value:newNode}
		ctx.runChildren(ns, ins)
	case "fetch.Text":
		searchNode := scope.Value.(xml.Node)
		xPathObj := xpath.NewXPath(searchNode.Doc())
		nodeSet := xPathObj.Search(searchNode, args[0].(string))
		if nodeSet.Size() > 0 {
			node := nodeSet.First()
			attr, ok := node.(*xml.Attribute)
			if ok {
				returnValue = attr.Content()
			} else {
				returnValue = node.String()
			}
		}
		xPathObj.Free()

	// LIBXML FUNCTIONS
	case "insert_at.Position.Text":
		node := scope.Value.(xml.Node)
		position := args[0].(Position)
		tagName := args[1].(string)
		element := node.Doc().NewElement(tagName)
		MoveFunc(element, node, position)
		ns := &Scope{Value: element}
		ctx.runChildren(ns, ins)
	case "inject_at.Position.Text":
		node := scope.Value.(xml.Node)
		position := args[0].(Position)
		nodeSet := node.Doc().ParseHtmlFragment(args[1].(string))
		for _, newNode := range(nodeSet) {
			MoveFunc(newNode, node, position)
		}
		if len(nodeSet) > 0 {
			element, ok := nodeSet[0].(*xml.Element)
			if ok {
				// successfully ran scope
				returnValue = "true"
				ns := &Scope{Value: element}
				ctx.runChildren(ns, ins)
			}
		} else {
			returnValue = "false"
		}
	case "move.XMLNode.XMLNode.Position", "move.Node.Node.Position":
		//for name, value := range(ctx.LocalVar) {
		//	println(name, ":", value)
		//}
		MoveFunc(args[0].(xml.Node), args[1].(xml.Node), args[2].(Position))
	case "wrap_text_children.Text":
		returnValue = "false"
		child := scope.Value.(xml.Node).First()
		index := 0
		tagName := args[0].(string)
		for child != nil {
			text, ok := child.(*xml.Text)
			childNext := child.Next()
			if ok {
				returnValue = "true"
				wrap := text.Wrap(tagName)
				ns := &Scope{wrap, index}
				ctx.runChildren(ns, ins)
				index++
			}
			child = childNext
		}

	// ATTRIBUTE FUNCTIONS
	case "attribute.Text":
		node := scope.Value.(xml.Node)
		name := args[0].(string)
		_, ok := node.(*xml.Element)
		if ok == true {
			attr, _ := node.Attribute(name)
			as := &Scope{Value:attr}
			ctx.runChildren(as, ins)
			if attr.IsLinked() && (attr.Content() == "") {
				attr.Remove()
			}
			if !attr.IsLinked() {
				attr.Free()
			}
			returnValue = "true"
		}
	default:
		ctx.Log.Error("Must implement " + fun.Name)
		returnValue = ""
	}
	return
}<commit_msg>make sure a node is an element before we mess with it<commit_after>package shark

import(
	"strings"
	"os"
	"libxml"
	"fmt"
	xml "libxml/tree"
	tp "tritium/proto"
	"libxml/xpath"
	"rubex"
)

func (ctx *Ctx) runBuiltIn(fun *Function, scope *Scope, ins *tp.Instruction, args []interface{}) (returnValue interface{}) {
	switch fun.Name {
	case "this":
		returnValue = scope.Value
	case "yield":
		myYieldBlock := ctx.yieldBlock()
		ctx.Yields = ctx.Yields[:(len(ctx.Yields)-1)]
		if (ctx.yieldBlock() != nil) {
			returnValue = ctx.runChildren(scope, myYieldBlock.Ins)
		} else {
			panic("yield() failure")
		}
		ctx.Yields = append(ctx.Yields, myYieldBlock)

	case "var.Text":
		val := ctx.Env[args[0].(string)]
		ts := &Scope{Value: val}
		ctx.runChildren(ts, ins)
		returnValue = ts.Value
		ctx.Env[args[0].(string)] = returnValue.(string)
	case "var.Text.Text":
		ctx.Env[args[0].(string)] = args[1].(string)
		returnValue = args[1].(string)
	case "match.Text":
		// Setup stacks
		ctx.MatchStack = append(ctx.MatchStack, args[0].(string))
		ctx.MatchShouldContinue = append(ctx.MatchShouldContinue, true)

		// Run children
		ctx.runChildren(scope, ins)

		if ctx.matchShouldContinue() {
			returnValue = "false"
		} else {
			returnValue = "true"
		}

		// Clear
		ctx.MatchShouldContinue = ctx.MatchShouldContinue[:len(ctx.MatchShouldContinue)-1]
		ctx.MatchStack = ctx.MatchStack[:len(ctx.MatchStack)-1]
	case "with.Text":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			if args[0].(string) == ctx.matchTarget() {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "with.Regexp":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			//println(matcher.MatchAgainst, matchWith)
			if (args[0].(*rubex.Regexp)).Match([]uint8(ctx.matchTarget())) {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "not.Text":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			if args[0].(string) != ctx.matchTarget() {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "not.Regexp":
		returnValue = "false"
		if ctx.matchShouldContinue() {
			//println(matcher.MatchAgainst, matchWith)
			if !(args[0].(*rubex.Regexp)).Match([]uint8(ctx.matchTarget())) {
				ctx.MatchShouldContinue[len(ctx.MatchShouldContinue)-1] = false
				ctx.runChildren(scope, ins)
				returnValue = "true"
			}
		}
	case "regexp.Text.Text":
		mode := rubex.ONIG_OPTION_DEFAULT
		if strings.Index(args[1].(string), "i") >= 0 {
			mode = rubex.ONIG_OPTION_IGNORECASE
		}
		if strings.Index(args[1].(string), "m") >= 0 {
			mode = rubex.ONIG_OPTION_MULTILINE
		}
		var err os.Error
		returnValue, err = rubex.NewRegexp(args[0].(string), mode)
		if err != nil {
			panic("Invalid regexp")
		}
	case "export.Text":
		val := make([]string, 2)
		val[0] = args[0].(string)
		ts := &Scope{Value:""}
		ctx.runChildren(ts, ins)
		val[1] = ts.Value.(string)
		ctx.Exports = append(ctx.Exports, val)
	case "log.Text":
		ctx.Logs = append(ctx.Logs, args[0].(string))

	// ATOMIC FUNCTIONS
	case "concat.Text.Text":
		//println("Concat:", args[0].(string), "+", args[1].(string))
		returnValue = args[0].(string) + args[1].(string)
	case "concat.Text.Text.Text": //REMOVE
		returnValue = args[0].(string) + args[1].(string) + args[2].(string)
	case "downcase.Text":
		returnValue = strings.ToLower(args[0].(string))
		return
	case "upcase.Text":
		returnValue = strings.ToUpper(args[0].(string))
		return
	case "index.XMLNode":
		returnValue = fmt.Sprintf("%d", scope.Index + 1)

	// TEXT FUNCTIONS
	case "set.Text":
		scope.Value = args[0]
	case "append.Text":
		scope.Value = scope.Value.(string) + args[0].(string)
	case "prepend.Text":
		scope.Value = args[0].(string) + scope.Value.(string)
	case "replace.Text":
		ts := &Scope{Value:""}
		ctx.runChildren(ts, ins)
		scope.Value = strings.Replace(scope.Value.(string), args[0].(string), ts.Value.(string), -1)
	case "replace.Regexp":
		regexp := args[0].(*rubex.Regexp)
		scope.Value = regexp.GsubFunc(scope.Value.(string), func(match string, captures map[string]string) string {
			usesGlobal := (ctx.Env["use_global_replace_vars"] == "true")

			for name, capture := range captures {
				if usesGlobal {
					//println("setting $", name, "to", capture)
					ctx.Env[name] = capture
				}
				ctx.vars()[name] = capture
			}

			replacementScope := &Scope{Value:match}
			ctx.runChildren(replacementScope, ins)
			//println(ins.String())

			//println("Replacement:", replacementScope.Value.(string))
			innerReplacer := rubex.MustCompile(`[\\\$]([\d])`)
			return innerReplacer.GsubFunc(replacementScope.Value.(string), func(_ string, numeric_captures map[string]string) string {
				capture := numeric_captures["1"]
				var val string
				if usesGlobal {
					val = ctx.Env[capture]
				} else {
					val = ctx.vars()[capture].(string)
				}
				return val
		    })
		})
		returnValue = scope.Value

	// XML FUNCTIONS
	case "xml":
		doc := libxml.XmlParseString(scope.Value.(string))
		defer doc.Free()
		ns := &Scope{Value:doc}
		ctx.runChildren(ns, ins)
		scope.Value = doc.String()
		returnValue = scope.Value
	case "html":
		doc := libxml.HtmlParseString(scope.Value.(string))
		defer doc.Free()
		ns := &Scope{Value:doc}
		ctx.runChildren(ns, ins)
		scope.Value = doc.DumpHTML()
		returnValue = scope.Value
	case "html_fragment":
		doc := libxml.HtmlParseFragment(scope.Value.(string))
		defer doc.Free()
		ns := &Scope{Value: doc.RootElement()}
		ctx.runChildren(ns, ins)
		scope.Value = ns.Value.(xml.Node).Content()
		returnValue = scope.Value
	case "select.Text":
		// TODO reuse XPath object
		node := scope.Value.(xml.Node)
		xpCtx := xpath.NewXPath(node.Doc())
		xpath := xpath.CompileXPath(args[0].(string))
		//xpath := ctx.XPath(args[0].(string))
		nodeSet := xpCtx.SearchByCompiledXPath(node, xpath).Slice()
		defer xpCtx.Free()
		if len(nodeSet) == 0 {
			returnValue = "false"
		} else {
			returnValue = "true"
		}

		for index, node := range(nodeSet) {
			if (node != nil) && node.IsLinked() && node.IsValid() {
				ns := &Scope{Value: node, Index: index}
				ctx.runChildren(ns, ins)
			}
		}
	case "position.Text":
		returnValue = Positions[args[0].(string)]

	// SHARED NODE FUNCTIONS
	case "remove":
		scope.Value.(xml.Node).Remove()
	case "inner", "value", "inner_text", "text":
		node := scope.Value.(xml.Node)
		ts := &Scope{Value:node.Content()}
		ctx.runChildren(ts, ins)
		val := ts.Value.(string)
		_, ok := node.(*xml.Element)
		if ok && node.IsLinked() {
			node.SetContent(val)
		}
		returnValue = val
	case "name":
		node := scope.Value.(xml.Node)
		ts := &Scope{Value:node.Name()}
		ctx.runChildren(ts, ins)
		node.SetName(ts.Value.(string))
		returnValue = ts.Value.(string)
	case "dup":
		node := scope.Value.(xml.Node)
		newNode := node.Duplicate()
		MoveFunc(newNode, node, AFTER)
		ns := &Scope{Value:newNode}
		ctx.runChildren(ns, ins)
	case "fetch.Text":
		searchNode := scope.Value.(xml.Node)
		xPathObj := xpath.NewXPath(searchNode.Doc())
		nodeSet := xPathObj.Search(searchNode, args[0].(string))
		if nodeSet.Size() > 0 {
			node := nodeSet.First()
			attr, ok := node.(*xml.Attribute)
			if ok {
				returnValue = attr.Content()
			} else {
				returnValue = node.String()
			}
		}
		xPathObj.Free()

	// LIBXML FUNCTIONS
	case "insert_at.Position.Text":
		node := scope.Value.(xml.Node)
		position := args[0].(Position)
		tagName := args[1].(string)
		element := node.Doc().NewElement(tagName)
		MoveFunc(element, node, position)
		ns := &Scope{Value: element}
		ctx.runChildren(ns, ins)
	case "inject_at.Position.Text":
		node := scope.Value.(xml.Node)
		position := args[0].(Position)
		nodeSet := node.Doc().ParseHtmlFragment(args[1].(string))
		for _, newNode := range(nodeSet) {
			MoveFunc(newNode, node, position)
		}
		if len(nodeSet) > 0 {
			element, ok := nodeSet[0].(*xml.Element)
			if ok {
				// successfully ran scope
				returnValue = "true"
				ns := &Scope{Value: element}
				ctx.runChildren(ns, ins)
			}
		} else {
			returnValue = "false"
		}
	case "move.XMLNode.XMLNode.Position", "move.Node.Node.Position":
		//for name, value := range(ctx.LocalVar) {
		//	println(name, ":", value)
		//}
		MoveFunc(args[0].(xml.Node), args[1].(xml.Node), args[2].(Position))
	case "wrap_text_children.Text":
		returnValue = "false"
		child := scope.Value.(xml.Node).First()
		index := 0
		tagName := args[0].(string)
		for child != nil {
			text, ok := child.(*xml.Text)
			childNext := child.Next()
			if ok {
				returnValue = "true"
				wrap := text.Wrap(tagName)
				ns := &Scope{wrap, index}
				ctx.runChildren(ns, ins)
				index++
			}
			child = childNext
		}

	// ATTRIBUTE FUNCTIONS
	case "attribute.Text":
		node := scope.Value.(xml.Node)
		name := args[0].(string)
		_, ok := node.(*xml.Element)
		if ok == true {
			attr, _ := node.Attribute(name)
			as := &Scope{Value:attr}
			ctx.runChildren(as, ins)
			if attr.IsLinked() && (attr.Content() == "") {
				attr.Remove()
			}
			if !attr.IsLinked() {
				attr.Free()
			}
			returnValue = "true"
		}
	default:
		ctx.Log.Error("Must implement " + fun.Name)
		returnValue = ""
	}
	return
}<|endoftext|>


<commit_before>package gateway

import (
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/golib/ratelimiter"
	"github.com/funkygao/golib/sync2"
	"github.com/funkygao/golib/timewheel"
	log "github.com/funkygao/log4go"
	"github.com/samuel/go-zookeeper/zk"
)

type subServer struct {
	*webServer

	idleConnsWg   sync2.WaitGroupTimeout // wait for all inflight http connections done
	closedConnCh  chan string            // channel of remote addr
	idleConns     map[net.Conn]struct{}
	idleConnsLock sync.Mutex

	auditor log.Logger

	timer *timewheel.TimeWheel

	// websocket heartbeat configuration
	wsReadLimit int64
	wsPongWait  time.Duration

	shutdownOnce sync.Once
	ackShutdown  int32                                          // sync shutdown with ack handlers goroutines
	ackCh        chan ackOffsets                                // client ack'ed offsets
	ackedOffsets map[string]map[string]map[string]map[int]int64 // [cluster][topic][group][partition]: offset

	subMetrics *subMetrics

	throttleBadGroup *ratelimiter.LeakyBuckets
	goodGroupClients map[string]struct{} // key is remote addr(port inclusive)
	goodGroupLock    sync.RWMutex
}

func newSubServer(httpAddr, httpsAddr string, maxClients int, gw *Gateway) *subServer {
	this := &subServer{
		webServer:        newWebServer("sub_server", httpAddr, httpsAddr, maxClients, gw),
		closedConnCh:     make(chan string, 1<<10),
		idleConns:        make(map[net.Conn]struct{}, 200),
		wsReadLimit:      8 << 10,
		wsPongWait:       time.Minute,
		timer:            timewheel.NewTimeWheel(time.Second, 120),
		throttleBadGroup: ratelimiter.NewLeakyBuckets(3, time.Minute),
		goodGroupClients: make(map[string]struct{}, 100),
		ackShutdown:      0,
		ackCh:            make(chan ackOffsets, 100),
		ackedOffsets:     make(map[string]map[string]map[string]map[int]int64),
	}
	this.subMetrics = NewSubMetrics(this.gw)
	this.waitExitFunc = this.waitExit
	this.connStateFunc = this.connStateHandler

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
	}

	this.auditor = log.NewDefaultLogger(log.TRACE)
	this.auditor.DeleteFilter("stdout")

	_ = os.Mkdir("audit", os.ModePerm)
	rotateEnabled, discardWhenDiskFull := true, false
	filer := log.NewFileLogWriter("audit/sub_audit.log", rotateEnabled, discardWhenDiskFull, 0644)
	if filer == nil {
		panic("failed to open sub audit log")
	}
	filer.SetFormat("[%d %T] [%L] (%S) %M")
	if Options.LogRotateSize > 0 {
		filer.SetRotateSize(Options.LogRotateSize)
	}
	filer.SetRotateLines(0)
	filer.SetRotateDaily(true)
	this.auditor.AddFilter("file", logLevel, filer)

	return this
}

func (this *subServer) Start() {
	this.gw.wg.Add(1)
	go this.ackCommitter()

	this.subMetrics.Load()
	this.webServer.Start()
}

func (this *subServer) connStateHandler(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		// Connections begin at StateNew and then
		// transition to either StateActive or StateClosed
		this.idleConnsWg.Add(1)
		atomic.AddInt32(&this.activeConnN, 1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Inc(1)
		}

	case http.StateActive:
		// StateActive fires before the request has entered a handler
		// and doesn't fire again until the request has been
		// handled.
		// After the request is handled, the state
		// transitions to StateClosed, StateHijacked, or StateIdle.
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

	case http.StateIdle:
		// StateIdle represents a connection that has finished
		// handling a request and is in the keep-alive state, waiting
		// for a new request. Connections transition from StateIdle
		// to either StateActive or StateClosed.
		select {
		case <-this.gw.shutdownCh:
			// actively close the client safely because IO is all done
			c.Close()

			this.idleConnsLock.Lock()
			delete(this.idleConns, c)
			this.idleConnsLock.Unlock()

		default:
			this.idleConnsLock.Lock()
			this.idleConns[c] = struct{}{}
			this.idleConnsLock.Unlock()
		}

	case http.StateHijacked:
		// websocket steals the socket
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

		atomic.AddInt32(&this.activeConnN, -1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)

			this.gw.svrMetrics.ConcurrentSubWs.Inc(1)
		}

	case http.StateClosed:
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)
		}

		this.closedConnCh <- c.RemoteAddr().String()
		this.idleConnsWg.Done()
		atomic.AddInt32(&this.activeConnN, -1)

		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()
	}
}

func (this *subServer) waitExit(exit <-chan struct{}) {
	<-exit

	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpsListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	this.idleConnsLock.Lock()
	t := time.Now().Add(time.Millisecond * 100)
	for c := range this.idleConns {
		c.SetReadDeadline(t)
	}
	this.idleConnsLock.Unlock()

	if this.idleConnsWg.WaitTimeout(Options.SubTimeout) {
		log.Warn("%s waiting for all connected client close timeout: %s",
			this.name, Options.SubTimeout)
	}

	this.subMetrics.Flush()
	this.timer.Stop()

	this.gw.wg.Done()
	close(this.closed)
}

func (this *subServer) ackCommitter() {
	ticker := time.NewTicker(time.Second * 30)
	defer func() {
		ticker.Stop()
		log.Debug("ack committer done")
		this.gw.wg.Done()
	}()

	for {
		select {
		case <-this.gw.shutdownCh:
			this.shutdownOnce.Do(func() {
				atomic.AddInt32(&this.ackShutdown, -1)

				for {
					// waiting for all ack handlers finish
					if atomic.LoadInt32(&this.ackShutdown) <= -1 {
						break
					}

					time.Sleep(time.Millisecond * 10)
				}
				close(this.ackCh)
			})

		case acks, ok := <-this.ackCh:
			if ok {
				for _, ack := range acks {
					if _, present := this.ackedOffsets[ack.cluster]; !present {
						this.ackedOffsets[ack.cluster] = make(map[string]map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.cluster][ack.topic]; !present {
						this.ackedOffsets[ack.cluster][ack.topic] = make(map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.topic][ack.group]; !present {
						this.ackedOffsets[ack.cluster][ack.topic][ack.group] = make(map[int]int64)
					}

					// TODO validation
					this.ackedOffsets[ack.cluster][ack.topic][ack.group][ack.Partition] = ack.Offset
				}
			} else {
				// channel buffer drained, flush all offsets
				// zk is still alive, safe to commit offsets
				this.commitOffsets()
				return
			}

		case <-ticker.C:
			this.commitOffsets()
		}
	}

}

func (this *subServer) commitOffsets() {
	for cluster, clusterTopic := range this.ackedOffsets {
		zkcluster := meta.Default.ZkCluster(cluster)

		for topic, groupPartition := range clusterTopic {
			for group, partitionOffset := range groupPartition {
				for partition, offset := range partitionOffset {
					if offset == -1 {
						// this slot is empty
						continue
					}

					log.Debug("commit offset {C:%s T:%s G:%s P:%d O:%d}", cluster, topic, group, partition, offset)

					if err := zkcluster.ResetConsumerGroupOffset(topic, group,
						strconv.Itoa(partition), offset); err != nil {
						log.Error("commitOffsets: %v", err)
						if err == zk.ErrNoNode {
							// invalid offset commit request, will not retry
							this.ackedOffsets[cluster][topic][group][partition] = -1
						}
					} else {
						// mark this slot empty
						this.ackedOffsets[cluster][topic][group][partition] = -1
					}
				}
			}
		}
	}

}
<commit_msg>besides interval, group commit also use counter<commit_after>package gateway

import (
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/golib/ratelimiter"
	"github.com/funkygao/golib/sync2"
	"github.com/funkygao/golib/timewheel"
	log "github.com/funkygao/log4go"
	"github.com/samuel/go-zookeeper/zk"
)

type subServer struct {
	*webServer

	idleConnsWg   sync2.WaitGroupTimeout // wait for all inflight http connections done
	closedConnCh  chan string            // channel of remote addr
	idleConns     map[net.Conn]struct{}
	idleConnsLock sync.Mutex

	auditor log.Logger

	timer *timewheel.TimeWheel

	// websocket heartbeat configuration
	wsReadLimit int64
	wsPongWait  time.Duration

	shutdownOnce sync.Once
	ackShutdown  int32                                          // sync shutdown with ack handlers goroutines
	ackCh        chan ackOffsets                                // client ack'ed offsets
	ackedOffsets map[string]map[string]map[string]map[int]int64 // [cluster][topic][group][partition]: offset

	subMetrics *subMetrics

	throttleBadGroup *ratelimiter.LeakyBuckets
	goodGroupClients map[string]struct{} // key is remote addr(port inclusive)
	goodGroupLock    sync.RWMutex
}

func newSubServer(httpAddr, httpsAddr string, maxClients int, gw *Gateway) *subServer {
	this := &subServer{
		webServer:        newWebServer("sub_server", httpAddr, httpsAddr, maxClients, gw),
		closedConnCh:     make(chan string, 1<<10),
		idleConns:        make(map[net.Conn]struct{}, 200),
		wsReadLimit:      8 << 10,
		wsPongWait:       time.Minute,
		timer:            timewheel.NewTimeWheel(time.Second, 120),
		throttleBadGroup: ratelimiter.NewLeakyBuckets(3, time.Minute),
		goodGroupClients: make(map[string]struct{}, 100),
		ackShutdown:      0,
		ackCh:            make(chan ackOffsets, 100),
		ackedOffsets:     make(map[string]map[string]map[string]map[int]int64),
	}
	this.subMetrics = NewSubMetrics(this.gw)
	this.waitExitFunc = this.waitExit
	this.connStateFunc = this.connStateHandler

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
	}

	this.auditor = log.NewDefaultLogger(log.TRACE)
	this.auditor.DeleteFilter("stdout")

	_ = os.Mkdir("audit", os.ModePerm)
	rotateEnabled, discardWhenDiskFull := true, false
	filer := log.NewFileLogWriter("audit/sub_audit.log", rotateEnabled, discardWhenDiskFull, 0644)
	if filer == nil {
		panic("failed to open sub audit log")
	}
	filer.SetFormat("[%d %T] [%L] (%S) %M")
	if Options.LogRotateSize > 0 {
		filer.SetRotateSize(Options.LogRotateSize)
	}
	filer.SetRotateLines(0)
	filer.SetRotateDaily(true)
	this.auditor.AddFilter("file", logLevel, filer)

	return this
}

func (this *subServer) Start() {
	this.gw.wg.Add(1)
	go this.ackCommitter()

	this.subMetrics.Load()
	this.webServer.Start()
}

func (this *subServer) connStateHandler(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		// Connections begin at StateNew and then
		// transition to either StateActive or StateClosed
		this.idleConnsWg.Add(1)
		atomic.AddInt32(&this.activeConnN, 1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Inc(1)
		}

	case http.StateActive:
		// StateActive fires before the request has entered a handler
		// and doesn't fire again until the request has been
		// handled.
		// After the request is handled, the state
		// transitions to StateClosed, StateHijacked, or StateIdle.
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

	case http.StateIdle:
		// StateIdle represents a connection that has finished
		// handling a request and is in the keep-alive state, waiting
		// for a new request. Connections transition from StateIdle
		// to either StateActive or StateClosed.
		select {
		case <-this.gw.shutdownCh:
			// actively close the client safely because IO is all done
			c.Close()

			this.idleConnsLock.Lock()
			delete(this.idleConns, c)
			this.idleConnsLock.Unlock()

		default:
			this.idleConnsLock.Lock()
			this.idleConns[c] = struct{}{}
			this.idleConnsLock.Unlock()
		}

	case http.StateHijacked:
		// websocket steals the socket
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

		atomic.AddInt32(&this.activeConnN, -1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)

			this.gw.svrMetrics.ConcurrentSubWs.Inc(1)
		}

	case http.StateClosed:
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)
		}

		this.closedConnCh <- c.RemoteAddr().String()
		this.idleConnsWg.Done()
		atomic.AddInt32(&this.activeConnN, -1)

		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()
	}
}

func (this *subServer) waitExit(exit <-chan struct{}) {
	<-exit

	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpsListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	this.idleConnsLock.Lock()
	t := time.Now().Add(time.Millisecond * 100)
	for c := range this.idleConns {
		c.SetReadDeadline(t)
	}
	this.idleConnsLock.Unlock()

	if this.idleConnsWg.WaitTimeout(Options.SubTimeout) {
		log.Warn("%s waiting for all connected client close timeout: %s",
			this.name, Options.SubTimeout)
	}

	this.subMetrics.Flush()
	this.timer.Stop()

	this.gw.wg.Done()
	close(this.closed)
}

func (this *subServer) ackCommitter() {
	ticker := time.NewTicker(time.Second * 30)
	defer func() {
		ticker.Stop()
		log.Debug("ack committer done")
		this.gw.wg.Done()
	}()

	var n int64
	const flushPerN int64 = 4000
	for {
		select {
		case <-this.gw.shutdownCh:
			this.shutdownOnce.Do(func() {
				atomic.AddInt32(&this.ackShutdown, -1)

				for {
					// waiting for all ack handlers finish
					if atomic.LoadInt32(&this.ackShutdown) <= -1 {
						break
					}

					time.Sleep(time.Millisecond * 10)
				}
				close(this.ackCh)
			})

		case acks, ok := <-this.ackCh:
			if ok {
				for _, ack := range acks {
					if _, present := this.ackedOffsets[ack.cluster]; !present {
						this.ackedOffsets[ack.cluster] = make(map[string]map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.cluster][ack.topic]; !present {
						this.ackedOffsets[ack.cluster][ack.topic] = make(map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.topic][ack.group]; !present {
						this.ackedOffsets[ack.cluster][ack.topic][ack.group] = make(map[int]int64)
					}

					// TODO validation
					this.ackedOffsets[ack.cluster][ack.topic][ack.group][ack.Partition] = ack.Offset
				}

				n++
				if n%flushPerN == 0 {
					this.commitOffsets()
				}
			} else {
				// channel buffer drained, flush all offsets
				// zk is still alive, safe to commit offsets
				this.commitOffsets()
				return
			}

		case <-ticker.C:
			this.commitOffsets()
		}
	}

}

func (this *subServer) commitOffsets() {
	for cluster, clusterTopic := range this.ackedOffsets {
		zkcluster := meta.Default.ZkCluster(cluster)

		for topic, groupPartition := range clusterTopic {
			for group, partitionOffset := range groupPartition {
				for partition, offset := range partitionOffset {
					if offset == -1 {
						// this slot is empty
						continue
					}

					log.Debug("cluster[%s] group[%s] commit offset {T:%s/%d O:%d}", cluster, group, topic, partition, offset)

					if err := zkcluster.ResetConsumerGroupOffset(topic, group, strconv.Itoa(partition), offset); err != nil {
						log.Error("cluster[%s] group[%s] commit offset {T:%s/%d O:%d} %v", cluster, group, topic, partition, offset, err)

						if err == zk.ErrNoNode {
							// invalid offset commit request, will not retry
							this.ackedOffsets[cluster][topic][group][partition] = -1
						}
					} else {
						// mark this slot empty
						this.ackedOffsets[cluster][topic][group][partition] = -1
					}
				}
			}
		}
	}

}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"text/template"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/spf13/cobra"
	"github.com/tleyden/deepstyle/deepstylelib"
	"github.com/tleyden/go-couch"
)

const (
	DesignDocName = "unprocessed_jobs"
	ViewName      = "unprocessed_jobs"
)

// publish_cloudwatch_metricsCmd respresents the publish_cloudwatch_metrics command
var publish_cloudwatch_metricsCmd = &cobra.Command{
	Use:   "publish_cloudwatch_metrics",
	Short: "Publish queue metrics to CloudWatch in order to trigger auto-scale alarms",
	Long:  `Publish queue metrics to CloudWatch in order to trigger auto-scale alarms`,
	Run: func(cmd *cobra.Command, args []string) {

		if err := cmd.ParseFlags(args); err != nil {
			log.Printf("err: %v", err)
			return
		}

		awsKey := cmd.Flag("aws_key")

		awsKeyVal := awsKey.Value.String()

		log.Printf("awsKey: %v", awsKeyVal)
		if awsKeyVal == "" {
			log.Printf("ERROR: Missing: --aws_key.\n  %v", cmd.UsageString())
			return
		}

		urlFlag := cmd.Flag("admin_url")

		urlVal := urlFlag.Value.String()
		if urlVal == "" {
			log.Printf("ERROR: Missing: --url.\n  %v", cmd.UsageString())
			return
		}

		err := addCloudWatchMetrics(urlVal)
		if err != nil {
			log.Printf("ERROR: %v", err)
			return
		}

	},
}

func numJobsReadyOrBeingProcessed(syncGwAdminUrl string) (metricValue float64, err error) {

	// try to query view
	//    curl localhost:4985/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs
	// if we get a 404, then install the view and then requery

	// if it has a trailing slash, remove it
	rawUrl := strings.TrimSuffix(syncGwAdminUrl, "/")

	// url validation
	url, err := url.Parse(rawUrl)
	if err != nil {
		return 0.0, err
	}

	db, err := couch.Connect(url.String())
	if err != nil {
		return 0.0, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}
	log.Printf("connected to db: %v", db)

	viewUrl := fmt.Sprintf("_design/%v/_view/%v", DesignDocName, ViewName)
	options := map[string]interface{}{}
	options["stale"] = "false"
	output := map[string]interface{}{}
	err = db.Query(viewUrl, options, &output)
	if err != nil {
		if strings.Contains(err.Error(), "404") || strings.Contains(err.Error(), "not_found") {
			// the view doesn't exist yet, attempt to install view
			if errInstallView := installView(rawUrl); errInstallView != nil {
				// failed to install view, give up
				return 0.0, errInstallView

			}

			// without this workaround, I'm getting:
			// ERROR: HTTP Error 500 Internal Server Error - {"error":"Internal Server Error","reason":"Internal error: error executing view req at http://127.0.0.1:8092/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs?stale=false: 500 Internal Server Error - {\"error\":\"unknown_error\",\"reason\":\"view_undefined\"}\n"}

			log.Printf("Sleeping 10s to wait for view to be ready")
			<-time.After(time.Duration(10) * time.Second)
			log.Printf("Done sleeping 10s to wait for view to be ready")

			// now retry
			errInner := db.Query(viewUrl, options, &output)
			if errInner != nil {
				// failed again, give up
				return 0.0, errInner
			}
		} else {
			return 0.0, err
		}
	}
	log.Printf("output: %+v", output)
	outputRows := output["total_rows"].(float64)
	return float64(outputRows), nil

}

type ViewParams struct {
	JobDocType string
	JobState1  string
	JobState2  string
	JobState3  string
}

func installView(syncGwAdminUrl string) error {

	viewJsonTemplate := `
{
    "views":{
        "unprocessed_jobs":{
            "map":"function (doc, meta) { if (doc.type != '{{.JobDocType}}') { return; } if (doc.state == '{{.JobState1}}' || doc.state == '{{.JobState2}}' || doc.state == '{{.JobState3}}') { emit(doc.state, meta.id); }}"
        }
    }
}
`

	viewParams := ViewParams{
		JobDocType: deepstylelib.Job,
		JobState1:  deepstylelib.StateNotReadyToProcess,
		JobState2:  deepstylelib.StateReadyToProcess,
		JobState3:  deepstylelib.StateBeingProcessed,
	}
	tmpl, err := template.New("UnprocessedJobsView").Parse(viewJsonTemplate)
	if err != nil {
		return err
	}

	var buffer bytes.Buffer // A Buffer needs no initialization.

	err = tmpl.Execute(&buffer, viewParams)
	if err != nil {
		return err
	}

	log.Printf("installView called")

	// curl -X PUT -H "Content-type: application/json" localhost:4985/todolite/_design/all_lists --data @testview
	viewUrl := fmt.Sprintf("%v/_design/%v", syncGwAdminUrl, DesignDocName)

	bufferBytes := buffer.Bytes()
	log.Printf("view: %v", string(bufferBytes))

	req, err := http.NewRequest("PUT", viewUrl, bytes.NewReader(bufferBytes))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}

	log.Printf("put view resp: %v", resp)

	return nil

}

func addCloudWatchMetrics(syncGwAdminUrl string) error {

	metricValue, err := numJobsReadyOrBeingProcessed(syncGwAdminUrl)
	log.Printf("numJobsReadyOrBeingProcessed: %v", metricValue)
	if err != nil {
		return err
	}

	// TODO: push something to CloudWatch
	// CLI example:
	//   aws cloudwatch put-metric-data
	//     --metric-name PageViewCount
	//      --namespace "MyService"
	//     --value 2
	//     --timestamp 2014-02-14T12:00:00.000Z
	cloudwatchSvc := cloudwatch.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

	log.Printf("cloudwatchSvc: %v", cloudwatchSvc)

	metricName := "NumJobsReadyOrBeingProcessed"
	timestamp := time.Now()

	metricDatum := &cloudwatch.MetricDatum{
		MetricName: &metricName,
		Value:      &metricValue,
		Timestamp:  &timestamp,
	}

	metricDatumSlice := []*cloudwatch.MetricDatum{metricDatum}
	namespace := "DeepStyleQueue"

	putMetricDataInput := &cloudwatch.PutMetricDataInput{
		MetricData: metricDatumSlice,
		Namespace:  &namespace,
	}

	out, err := cloudwatchSvc.PutMetricData(putMetricDataInput)
	if err != nil {
		log.Printf("ERROR adding metric data  %v", err)
		return err
	}
	log.Printf("Metric data output: %v", out)

	return nil

}

func init() {
	RootCmd.AddCommand(publish_cloudwatch_metricsCmd)

	// Here you will define your flags and configuration settings

	// Cobra supports Persistent Flags which will work for this command and all subcommands

	publish_cloudwatch_metricsCmd.PersistentFlags().String("aws_key", "", "AWS Key")

	publish_cloudwatch_metricsCmd.PersistentFlags().String("admin_url", "", "Sync Gateway Admin URL")

	// Cobra supports local flags which will only run when this command is called directly
	// publish_cloudwatch_metricsCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle" )

}
<commit_msg>Push metrics every 60s in a loop<commit_after>package cmd

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"text/template"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/spf13/cobra"
	"github.com/tleyden/deepstyle/deepstylelib"
	"github.com/tleyden/go-couch"
)

const (
	DesignDocName = "unprocessed_jobs"
	ViewName      = "unprocessed_jobs"
)

// publish_cloudwatch_metricsCmd respresents the publish_cloudwatch_metrics command
var publish_cloudwatch_metricsCmd = &cobra.Command{
	Use:   "publish_cloudwatch_metrics",
	Short: "Publish queue metrics to CloudWatch in order to trigger auto-scale alarms",
	Long:  `Publish queue metrics to CloudWatch in order to trigger auto-scale alarms.  AWS keys will be taken from environment variables or ~/.aws/.  See github.com/aws/aws-sdk-go`,
	Run: func(cmd *cobra.Command, args []string) {

		if err := cmd.ParseFlags(args); err != nil {
			log.Printf("err: %v", err)
			return
		}

		urlFlag := cmd.Flag("admin_url")

		urlVal := urlFlag.Value.String()
		if urlVal == "" {
			log.Printf("ERROR: Missing: --url.\n  %v", cmd.UsageString())
			return
		}

		err := addCloudWatchMetrics(urlVal)
		if err != nil {
			log.Printf("ERROR: %v", err)
			return
		}

	},
}

func numJobsReadyOrBeingProcessed(syncGwAdminUrl string) (metricValue float64, err error) {

	// try to query view
	//    curl localhost:4985/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs
	// if we get a 404, then install the view and then requery

	// if it has a trailing slash, remove it
	rawUrl := strings.TrimSuffix(syncGwAdminUrl, "/")

	// url validation
	url, err := url.Parse(rawUrl)
	if err != nil {
		return 0.0, err
	}

	db, err := couch.Connect(url.String())
	if err != nil {
		return 0.0, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}

	viewUrl := fmt.Sprintf("_design/%v/_view/%v", DesignDocName, ViewName)
	options := map[string]interface{}{}
	options["stale"] = "false"
	output := map[string]interface{}{}
	err = db.Query(viewUrl, options, &output)
	if err != nil {
		if strings.Contains(err.Error(), "404") || strings.Contains(err.Error(), "not_found") {
			// the view doesn't exist yet, attempt to install view
			if errInstallView := installView(rawUrl); errInstallView != nil {
				// failed to install view, give up
				return 0.0, errInstallView

			}

			// without this workaround, I'm getting:
			// ERROR: HTTP Error 500 Internal Server Error - {"error":"Internal Server Error","reason":"Internal error: error executing view req at http://127.0.0.1:8092/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs?stale=false: 500 Internal Server Error - {\"error\":\"unknown_error\",\"reason\":\"view_undefined\"}\n"}

			log.Printf("Sleeping 10s to wait for view to be ready")
			<-time.After(time.Duration(10) * time.Second)
			log.Printf("Done sleeping 10s to wait for view to be ready")

			// now retry
			errInner := db.Query(viewUrl, options, &output)
			if errInner != nil {
				// failed again, give up
				return 0.0, errInner
			}
		} else {
			return 0.0, err
		}
	}
	outputRows := output["total_rows"].(float64)
	return float64(outputRows), nil

}

type ViewParams struct {
	JobDocType string
	JobState1  string
	JobState2  string
	JobState3  string
}

func installView(syncGwAdminUrl string) error {

	viewJsonTemplate := `
{
    "views":{
        "unprocessed_jobs":{
            "map":"function (doc, meta) { if (doc.type != '{{.JobDocType}}') { return; } if (doc.state == '{{.JobState1}}' || doc.state == '{{.JobState2}}' || doc.state == '{{.JobState3}}') { emit(doc.state, meta.id); }}"
        }
    }
}
`

	viewParams := ViewParams{
		JobDocType: deepstylelib.Job,
		JobState1:  deepstylelib.StateNotReadyToProcess,
		JobState2:  deepstylelib.StateReadyToProcess,
		JobState3:  deepstylelib.StateBeingProcessed,
	}
	tmpl, err := template.New("UnprocessedJobsView").Parse(viewJsonTemplate)
	if err != nil {
		return err
	}

	var buffer bytes.Buffer // A Buffer needs no initialization.

	err = tmpl.Execute(&buffer, viewParams)
	if err != nil {
		return err
	}

	log.Printf("installView called")

	// curl -X PUT -H "Content-type: application/json" localhost:4985/todolite/_design/all_lists --data @testview
	viewUrl := fmt.Sprintf("%v/_design/%v", syncGwAdminUrl, DesignDocName)

	bufferBytes := buffer.Bytes()
	log.Printf("view: %v", string(bufferBytes))

	req, err := http.NewRequest("PUT", viewUrl, bytes.NewReader(bufferBytes))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}

	log.Printf("put view resp: %v", resp)

	return nil

}

func addCloudWatchMetrics(syncGwAdminUrl string) error {
	for {

		log.Printf("Adding metrics for queue")
		addCloudWatchMetric(syncGwAdminUrl)

		log.Printf("Sleeping 30s")
		<-time.After(time.Duration(30) * time.Second)

	}
}

func addCloudWatchMetric(syncGwAdminUrl string) error {

	metricValue, err := numJobsReadyOrBeingProcessed(syncGwAdminUrl)
	log.Printf("Adding metric: numJobsReadyOrBeingProcessed = %v", metricValue)
	if err != nil {
		return err
	}

	cloudwatchSvc := cloudwatch.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

	metricName := "NumJobsReadyOrBeingProcessed"
	timestamp := time.Now()

	metricDatum := &cloudwatch.MetricDatum{
		MetricName: &metricName,
		Value:      &metricValue,
		Timestamp:  &timestamp,
	}

	metricDatumSlice := []*cloudwatch.MetricDatum{metricDatum}
	namespace := "DeepStyleQueue"

	putMetricDataInput := &cloudwatch.PutMetricDataInput{
		MetricData: metricDatumSlice,
		Namespace:  &namespace,
	}

	_, err = cloudwatchSvc.PutMetricData(putMetricDataInput)
	if err != nil {
		log.Printf("ERROR adding metric data  %v", err)
		return err
	}

	return nil

}

func init() {
	RootCmd.AddCommand(publish_cloudwatch_metricsCmd)

	// Here you will define your flags and configuration settings

	// Cobra supports Persistent Flags which will work for this command and all subcommands

	publish_cloudwatch_metricsCmd.PersistentFlags().String("admin_url", "", "Sync Gateway Admin URL")

	// Cobra supports local flags which will only run when this command is called directly
	// publish_cloudwatch_metricsCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle" )

}
<|endoftext|>


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

package framework

import (
	"errors"
	"flag"
	"path/filepath"
	"time"

	"istio.io/istio/tests/util"
	"istio.io/pkg/log"
)

var (
	useAutomaticInjection = flag.Bool("use-automatic-injection", false, "Use automatic injection instead of kube-inject for transparent proxy injection")
)

const (
	kubeInjectPrefix = "KubeInject"
)

// App gathers information for Hop app
type App struct {
	AppYamlTemplate string
	AppYaml         string
	KubeInject      bool
	Template        interface{}
	deployedYaml    string
}

// AppManager organize and deploy apps
type AppManager struct {
	Apps             []*App
	tmpDir           string
	namespace        string
	istioctl         *Istioctl
	active           bool
	Kubeconfig       string
	checkDeployments bool
}

// NewAppManager create a new AppManager
func NewAppManager(tmpDir, namespace string, istioctl *Istioctl, kubeconfig string, checkDeployments bool) *AppManager {
	return &AppManager{
		namespace:        namespace,
		tmpDir:           tmpDir,
		istioctl:         istioctl,
		Kubeconfig:       kubeconfig,
		checkDeployments: checkDeployments,
	}
}

// generateAppYaml deploy testing app from tmpl
func (am *AppManager) generateAppYaml(a *App) error {
	if a.AppYamlTemplate == "" {
		return nil
	}
	var err error
	a.AppYaml, err = util.CreateTempfile(am.tmpDir, filepath.Base(a.AppYamlTemplate), yamlSuffix)
	if err != nil {
		log.Errorf("Failed to generate yaml %s: %v", a.AppYamlTemplate, err)
		return err
	}
	if err := util.Fill(a.AppYaml, a.AppYamlTemplate, a.Template); err != nil {
		log.Errorf("Failed to generate yaml for template %s: %v", a.AppYamlTemplate, err)
		return err
	}
	return nil
}

func (am *AppManager) deploy(a *App) error {
	if err := am.generateAppYaml(a); err != nil {
		return err
	}
	finalYaml := a.AppYaml
	if a.KubeInject && !*useAutomaticInjection {
		var err error
		finalYaml, err = util.CreateTempfile(am.tmpDir, kubeInjectPrefix, yamlSuffix)
		if err != nil {
			log.Errorf("CreateTempfile failed %v", err)
			return err
		}
		if err = am.istioctl.KubeInject(a.AppYaml, finalYaml, am.Kubeconfig); err != nil {
			log.Errorf("KubeInject failed for yaml %s: %v", a.AppYaml, err)
			return err
		}
	}
	var err error
	for i := 0; i < 3; i++ {
		if err = util.KubeApply(am.namespace, finalYaml, am.Kubeconfig); err == nil {
			break
		}
		log.Warnf("Kubectl apply %s failed", finalYaml)
		// wait for admission webhook configuration populate.
		time.Sleep(5 * time.Second)
	}
	a.deployedYaml = finalYaml
	return err
}

// Setup deploy apps
func (am *AppManager) Setup() error {
	am.active = true
	log.Info("Setting up apps")
	for _, a := range am.Apps {
		log.Infof("Setup %v", a)
		if err := am.deploy(a); err != nil {
			log.Errorf("error deploying %v: %v", a, err)
			return err
		}
	}
	if am.checkDeployments {
		return am.CheckDeployments()
	}
	return nil
}

// Teardown currently does nothing, only to satisfied cleanable{}
func (am *AppManager) Teardown() error {
	am.active = false
	return nil
}

// AddApp for automated deployment. Must be done before Setup call.
func (am *AppManager) AddApp(a *App) {
	am.Apps = append(am.Apps, a)
}

// DeployApp adds the app and deploys it to the system. Must be called after Setup call.
func (am *AppManager) DeployApp(a *App) error {
	if !am.active {
		return errors.New("function DeployApp must be called after Setup")
	}
	am.AddApp(a)
	return am.deploy(a)
}

// UndeployApp deletes the app from the system. Must be called after Setup call.
func (am *AppManager) UndeployApp(a *App) error {
	if !am.active {
		return errors.New("function UndeployApp must be called after Setup")
	}

	if a.deployedYaml == "" {
		// It wasn't deployed.
		return nil
	}

	if err := util.KubeDelete(am.namespace, a.deployedYaml, am.Kubeconfig); err != nil {
		log.Errorf("Kubectl delete %s failed", a.deployedYaml)
		return err
	}
	return nil
}

// CheckDeployments waits for a period for the deployments to be started.
func (am *AppManager) CheckDeployments() error {
	return util.CheckDeployments(am.namespace, maxDeploymentRolloutTime, am.Kubeconfig)
}
<commit_msg>Replace flag --use-automatic-injection as --use_automatic_injection for naming consistent in e2e testing framework. (#19768)<commit_after>// Copyright 2017 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package framework

import (
	"errors"
	"flag"
	"path/filepath"
	"time"

	"istio.io/istio/tests/util"
	"istio.io/pkg/log"
)

var (
	useAutomaticInjection = flag.Bool("use_automatic_injection", false, "Use automatic injection instead of kube-inject for transparent proxy injection")
)

const (
	kubeInjectPrefix = "KubeInject"
)

// App gathers information for Hop app
type App struct {
	AppYamlTemplate string
	AppYaml         string
	KubeInject      bool
	Template        interface{}
	deployedYaml    string
}

// AppManager organize and deploy apps
type AppManager struct {
	Apps             []*App
	tmpDir           string
	namespace        string
	istioctl         *Istioctl
	active           bool
	Kubeconfig       string
	checkDeployments bool
}

// NewAppManager create a new AppManager
func NewAppManager(tmpDir, namespace string, istioctl *Istioctl, kubeconfig string, checkDeployments bool) *AppManager {
	return &AppManager{
		namespace:        namespace,
		tmpDir:           tmpDir,
		istioctl:         istioctl,
		Kubeconfig:       kubeconfig,
		checkDeployments: checkDeployments,
	}
}

// generateAppYaml deploy testing app from tmpl
func (am *AppManager) generateAppYaml(a *App) error {
	if a.AppYamlTemplate == "" {
		return nil
	}
	var err error
	a.AppYaml, err = util.CreateTempfile(am.tmpDir, filepath.Base(a.AppYamlTemplate), yamlSuffix)
	if err != nil {
		log.Errorf("Failed to generate yaml %s: %v", a.AppYamlTemplate, err)
		return err
	}
	if err := util.Fill(a.AppYaml, a.AppYamlTemplate, a.Template); err != nil {
		log.Errorf("Failed to generate yaml for template %s: %v", a.AppYamlTemplate, err)
		return err
	}
	return nil
}

func (am *AppManager) deploy(a *App) error {
	if err := am.generateAppYaml(a); err != nil {
		return err
	}
	finalYaml := a.AppYaml
	if a.KubeInject && !*useAutomaticInjection {
		var err error
		finalYaml, err = util.CreateTempfile(am.tmpDir, kubeInjectPrefix, yamlSuffix)
		if err != nil {
			log.Errorf("CreateTempfile failed %v", err)
			return err
		}
		if err = am.istioctl.KubeInject(a.AppYaml, finalYaml, am.Kubeconfig); err != nil {
			log.Errorf("KubeInject failed for yaml %s: %v", a.AppYaml, err)
			return err
		}
	}
	var err error
	for i := 0; i < 3; i++ {
		if err = util.KubeApply(am.namespace, finalYaml, am.Kubeconfig); err == nil {
			break
		}
		log.Warnf("Kubectl apply %s failed", finalYaml)
		// wait for admission webhook configuration populate.
		time.Sleep(5 * time.Second)
	}
	a.deployedYaml = finalYaml
	return err
}

// Setup deploy apps
func (am *AppManager) Setup() error {
	am.active = true
	log.Info("Setting up apps")
	for _, a := range am.Apps {
		log.Infof("Setup %v", a)
		if err := am.deploy(a); err != nil {
			log.Errorf("error deploying %v: %v", a, err)
			return err
		}
	}
	if am.checkDeployments {
		return am.CheckDeployments()
	}
	return nil
}

// Teardown currently does nothing, only to satisfied cleanable{}
func (am *AppManager) Teardown() error {
	am.active = false
	return nil
}

// AddApp for automated deployment. Must be done before Setup call.
func (am *AppManager) AddApp(a *App) {
	am.Apps = append(am.Apps, a)
}

// DeployApp adds the app and deploys it to the system. Must be called after Setup call.
func (am *AppManager) DeployApp(a *App) error {
	if !am.active {
		return errors.New("function DeployApp must be called after Setup")
	}
	am.AddApp(a)
	return am.deploy(a)
}

// UndeployApp deletes the app from the system. Must be called after Setup call.
func (am *AppManager) UndeployApp(a *App) error {
	if !am.active {
		return errors.New("function UndeployApp must be called after Setup")
	}

	if a.deployedYaml == "" {
		// It wasn't deployed.
		return nil
	}

	if err := util.KubeDelete(am.namespace, a.deployedYaml, am.Kubeconfig); err != nil {
		log.Errorf("Kubectl delete %s failed", a.deployedYaml)
		return err
	}
	return nil
}

// CheckDeployments waits for a period for the deployments to be started.
func (am *AppManager) CheckDeployments() error {
	return util.CheckDeployments(am.namespace, maxDeploymentRolloutTime, am.Kubeconfig)
}
<|endoftext|>


<commit_before>/*

config is a simple library that manages config set-up for boardgame-util and
friends, reading from config.json and config.SECRET.json files. See boardgame-
util/README.md for more on the structure of config.json files.

Although a number of the details are exposed in this package, generally you
just use Get() and then directly read the values of the returned Config's Dev
and Prod properties.

*/
package config

import (
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

const (
	privateConfigFileName = "config.SECRET.json"
	publicConfigFileName  = "config.PUBLIC.json"
)

//DefaultFileNames returns the publicConfig and privateConfig names for the
//given path, even if they don't exist. If dirOrFile ends in ".json" then that
//will be returned, with privateConfig being in the same folder. If it's a
//dir, it will be the default filenames in that folder.
func DefaultFileNames(dirOrFile string) (publicConfig, privateConfig string, err error) {

	if dirOrFile == "" {
		dirOrFile = "."
	}

	if strings.HasSuffix(dirOrFile, ".json") {
		dir := filepath.Dir(dirOrFile)

		if _, err := os.Stat(dir); os.IsNotExist(err) {
			return "", "", errors.New("Dir " + dir + " does not exist.")
		}

		return dirOrFile, filepath.Join(dir, privateConfigFileName), nil
	}

	//OK, we'll interpret as Dir.

	if _, err := os.Stat(dirOrFile); os.IsNotExist(err) {
		return "", "", errors.New(dirOrFile + " is interpreted as a directory but does not exist")
	}
	return filepath.Join(dirOrFile, publicConfigFileName), filepath.Join(dirOrFile, privateConfigFileName), nil
}

//FileNames returns the publicConfig filename and privateConfig filename to
//use given the search path. If dir is a config file itself, loads that (and
//any private component in same directory). Next it interprets dir as a
//directory to search within for any config files. If none are found, walks
//upwards in the directory hierarchy (as long as that's still in $GOPATH)
//until it finds a folder that appears to work. If dir is "", working
//directory is assumed. If skipUpwardSearch is true, then if dir is non-blank
//upward searching in the dir hiearchy will not be done.
func FileNames(dir string, skipUpwardSearch bool) (publicConfig, privateConfig string, err error) {

	if dir == "" {
		dir = "."
		skipUpwardSearch = false
	}

	//Try to interpret it as a file
	if public, private, err := fileNamesToUseWithFile(dir); err == nil {
		return public, private, nil
	}

	//Guess it wasn't a file, try interpreting as a directory.

	goPath, err := filepath.Abs(os.Getenv("GOPATH"))

	if err != nil {
		//Gopath isn't set correctly
		return "", "", errors.New("Couldn't absolutize gopath: " + err.Error())
	}

	for {

		abs, err := filepath.Abs(dir)

		if err != nil {
			//Maybe fell off the end of what is a real file?
			return "", "", errors.New("Got err absolutizing search directory: " + dir + " : " + err.Error())
		}

		if !strings.HasPrefix(abs, goPath) {
			return "", "", errors.New("Fell out of gopath without finding config: " + abs)
		}

		public, private := fileNamesToUseInDir(dir)

		if public != "" || private != "" {
			return public, private, nil
		}

		if skipUpwardSearch {
			break
		}

		dir = filepath.Join("..", dir)
	}

	return "", "", errors.New("Couldn't find a path")

}

//fileNamesToUseWithFile takes a filename of the public component. Returns the
//string to the publicComponent and also the private component if it exists in
//that folder.z
func fileNamesToUseWithFile(filename string) (publicConfig, privateConfig string, err error) {

	if info, err := os.Stat(filename); err != nil {
		return "", "", errors.New("That file does not exist: " + err.Error())
	} else {
		if info.IsDir() {
			return "", "", errors.New(filename + " points to a dir, not a file")
		}
	}

	//Check to see if there's a private config in that folder
	dir := filepath.Dir(filename)

	privatePath := filepath.Join(dir, privateConfigFileName)

	if _, err := os.Stat(privatePath); err != nil {
		// No private path I guess
		return filename, "", nil
	}

	return filename, privatePath, nil

}

//fileNamesToUseInDir looks for public/private values precisely in the given folder.
func fileNamesToUseInDir(dir string) (publicConfig, privateConfig string) {

	possiblePrivateConfig := filepath.Join(dir, privateConfigFileName)

	if _, err := os.Stat(possiblePrivateConfig); err == nil {
		privateConfig = possiblePrivateConfig
	}

	infos, err := ioutil.ReadDir(dir)

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

	foundNames := make(map[string]bool)

	for _, info := range infos {
		if info.Name() == privateConfigFileName {
			continue
		}
		if strings.HasPrefix(info.Name(), "config.") && strings.HasSuffix(info.Name(), ".json") {
			foundNames[info.Name()] = true
		}
	}

	prioritizedNames := []string{
		publicConfigFileName,
		"config.json",
	}

	for _, name := range prioritizedNames {
		if foundNames[name] {
			return filepath.Join(dir, name), privateConfig
		}
	}

	//Whatever, return the first one
	for name := range foundNames {
		return filepath.Join(dir, name), privateConfig
	}

	//None of the preferred names were found, just return whatever is in
	//publicConfig, privateConfig. publicConfig is "", privateConfig already
	//has the dir in it
	return

}

//GetConfig returns a Config for those two named files. publicConfig and
//privateConfig may both be "" without erroring. If createIfNotExist is true,
//then NewRawConfig will be told to create the Configs even if they don't
//exist on disk.
func GetConfig(publicConfigFile, privateConfigFile string, createIfNotExist bool) (*Config, error) {
	publicConfig, err := NewRawConfig(publicConfigFile, createIfNotExist)

	if err != nil {
		return nil, errors.New("Couldn't get public config: " + err.Error())
	}

	privateConfig, err := NewRawConfig(privateConfigFile, createIfNotExist)

	if err != nil {
		return nil, errors.New("Couldn't get private config: " + err.Error())
	}

	return NewConfig(publicConfig, privateConfig), nil
}

//Get fetches a fully realized config. It is a simple convenience wrapper
//around FileNames and GetConfig. If createIfNotExist is true, then if the
//files don't exist on disk we'll generate file names and return raw configs
//for those filenames.
func Get(dir string, createIfNotExist bool) (*Config, error) {
	publicConfigName, privateConfigName, err := FileNames(dir, createIfNotExist)

	if err != nil {
		if createIfNotExist {
			publicConfigName, privateConfigName, err = DefaultFileNames(dir)
			if err != nil {
				return nil, errors.New("Couldn't get default file names to create: " + err.Error())
			}
		} else {
			return nil, errors.New("Couldn't get file names to use: " + err.Error())
		}
	}

	return GetConfig(publicConfigName, privateConfigName, createIfNotExist)

}
<commit_msg>Config searching no longer cares if you fall upward out of gopath, as in many cases in a module world you won't be in gopath. part of #662.<commit_after>/*

config is a simple library that manages config set-up for boardgame-util and
friends, reading from config.json and config.SECRET.json files. See boardgame-
util/README.md for more on the structure of config.json files.

Although a number of the details are exposed in this package, generally you
just use Get() and then directly read the values of the returned Config's Dev
and Prod properties.

*/
package config

import (
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

const (
	privateConfigFileName = "config.SECRET.json"
	publicConfigFileName  = "config.PUBLIC.json"
)

//DefaultFileNames returns the publicConfig and privateConfig names for the
//given path, even if they don't exist. If dirOrFile ends in ".json" then that
//will be returned, with privateConfig being in the same folder. If it's a
//dir, it will be the default filenames in that folder.
func DefaultFileNames(dirOrFile string) (publicConfig, privateConfig string, err error) {

	if dirOrFile == "" {
		dirOrFile = "."
	}

	if strings.HasSuffix(dirOrFile, ".json") {
		dir := filepath.Dir(dirOrFile)

		if _, err := os.Stat(dir); os.IsNotExist(err) {
			return "", "", errors.New("Dir " + dir + " does not exist.")
		}

		return dirOrFile, filepath.Join(dir, privateConfigFileName), nil
	}

	//OK, we'll interpret as Dir.

	if _, err := os.Stat(dirOrFile); os.IsNotExist(err) {
		return "", "", errors.New(dirOrFile + " is interpreted as a directory but does not exist")
	}
	return filepath.Join(dirOrFile, publicConfigFileName), filepath.Join(dirOrFile, privateConfigFileName), nil
}

//FileNames returns the publicConfig filename and privateConfig filename to
//use given the search path. If dir is a config file itself, loads that (and
//any private component in same directory). Next it interprets dir as a
//directory to search within for any config files. If none are found, walks
//upwards in the directory hierarchy (as long as that's still in $GOPATH)
//until it finds a folder that appears to work. If dir is "", working
//directory is assumed. If skipUpwardSearch is true, then if dir is non-blank
//upward searching in the dir hiearchy will not be done.
func FileNames(dir string, skipUpwardSearch bool) (publicConfig, privateConfig string, err error) {

	if dir == "" {
		dir = "."
		skipUpwardSearch = false
	}

	//Try to interpret it as a file
	if public, private, err := fileNamesToUseWithFile(dir); err == nil {
		return public, private, nil
	}

	//Guess it wasn't a file, try interpreting as a directory.

	for {

		public, private := fileNamesToUseInDir(dir)

		if public != "" || private != "" {
			return public, private, nil
		}

		if skipUpwardSearch {
			break
		}

		dir = filepath.Join("..", dir)
	}

	return "", "", errors.New("Couldn't find a path")

}

//fileNamesToUseWithFile takes a filename of the public component. Returns the
//string to the publicComponent and also the private component if it exists in
//that folder.z
func fileNamesToUseWithFile(filename string) (publicConfig, privateConfig string, err error) {

	if info, err := os.Stat(filename); err != nil {
		return "", "", errors.New("That file does not exist: " + err.Error())
	} else {
		if info.IsDir() {
			return "", "", errors.New(filename + " points to a dir, not a file")
		}
	}

	//Check to see if there's a private config in that folder
	dir := filepath.Dir(filename)

	privatePath := filepath.Join(dir, privateConfigFileName)

	if _, err := os.Stat(privatePath); err != nil {
		// No private path I guess
		return filename, "", nil
	}

	return filename, privatePath, nil

}

//fileNamesToUseInDir looks for public/private values precisely in the given folder.
func fileNamesToUseInDir(dir string) (publicConfig, privateConfig string) {

	possiblePrivateConfig := filepath.Join(dir, privateConfigFileName)

	if _, err := os.Stat(possiblePrivateConfig); err == nil {
		privateConfig = possiblePrivateConfig
	}

	infos, err := ioutil.ReadDir(dir)

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

	foundNames := make(map[string]bool)

	for _, info := range infos {
		if info.Name() == privateConfigFileName {
			continue
		}
		if strings.HasPrefix(info.Name(), "config.") && strings.HasSuffix(info.Name(), ".json") {
			foundNames[info.Name()] = true
		}
	}

	prioritizedNames := []string{
		publicConfigFileName,
		"config.json",
	}

	for _, name := range prioritizedNames {
		if foundNames[name] {
			return filepath.Join(dir, name), privateConfig
		}
	}

	//Whatever, return the first one
	for name := range foundNames {
		return filepath.Join(dir, name), privateConfig
	}

	//None of the preferred names were found, just return whatever is in
	//publicConfig, privateConfig. publicConfig is "", privateConfig already
	//has the dir in it
	return

}

//GetConfig returns a Config for those two named files. publicConfig and
//privateConfig may both be "" without erroring. If createIfNotExist is true,
//then NewRawConfig will be told to create the Configs even if they don't
//exist on disk.
func GetConfig(publicConfigFile, privateConfigFile string, createIfNotExist bool) (*Config, error) {
	publicConfig, err := NewRawConfig(publicConfigFile, createIfNotExist)

	if err != nil {
		return nil, errors.New("Couldn't get public config: " + err.Error())
	}

	privateConfig, err := NewRawConfig(privateConfigFile, createIfNotExist)

	if err != nil {
		return nil, errors.New("Couldn't get private config: " + err.Error())
	}

	return NewConfig(publicConfig, privateConfig), nil
}

//Get fetches a fully realized config. It is a simple convenience wrapper
//around FileNames and GetConfig. If createIfNotExist is true, then if the
//files don't exist on disk we'll generate file names and return raw configs
//for those filenames.
func Get(dir string, createIfNotExist bool) (*Config, error) {
	publicConfigName, privateConfigName, err := FileNames(dir, createIfNotExist)

	if err != nil {
		if createIfNotExist {
			publicConfigName, privateConfigName, err = DefaultFileNames(dir)
			if err != nil {
				return nil, errors.New("Couldn't get default file names to create: " + err.Error())
			}
		} else {
			return nil, errors.New("Couldn't get file names to use: " + err.Error())
		}
	}

	return GetConfig(publicConfigName, privateConfigName, createIfNotExist)

}
<|endoftext|>


<commit_before>// Copyright 2013 Daniel Pupius. 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 gifresize resizes animated gifs.
//
// Frames in an animated gif aren't necessarily the same size, subsequent
// frames are overlayed on previous frames. Therefore, resizing the frames
// individually may cause problems due to aliasing of transparent pixels. This
// package tries to avoid this by building frames from all previous frames and
// resizing the frames as RGB.
package gifresize

import (
	"image"
	"image/color/palette"
	"image/draw"
	"image/gif"
	"io"
)

// TransformFunc is a function that transforms an image.
type TransformFunc func(image.Image) image.Image

// Process the GIF read from r, applying transform to each frame, and writing
// the result to w.
func Process(w io.Writer, r io.Reader, transform TransformFunc) error {
	// Decode the original gif.
	im, err := gif.DecodeAll(r)
	if err != nil {
		return err
	}

	// Create a new RGBA image to hold the incremental frames.
	firstFrame := im.Image[0].Bounds()
	b := image.Rect(0, 0, firstFrame.Dx(), firstFrame.Dy())
	img := image.NewRGBA(b)

	// Resize each frame.
	for index, frame := range im.Image {
		bounds := frame.Bounds()
		draw.Draw(img, bounds, frame, bounds.Min, draw.Over)
		im.Image[index] = imageToPaletted(transform(img))
	}

	gif.EncodeAll(w, im)
	return nil
}

func imageToPaletted(img image.Image) *image.Paletted {
	b := img.Bounds()
	pm := image.NewPaletted(b, palette.Plan9)
	draw.FloydSteinberg.Draw(pm, b, img, image.ZP)
	return pm
}
<commit_msg>don't drop error from gif.EncodeAll<commit_after>// Copyright 2013 Daniel Pupius. 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 gifresize resizes animated gifs.
//
// Frames in an animated gif aren't necessarily the same size, subsequent
// frames are overlayed on previous frames. Therefore, resizing the frames
// individually may cause problems due to aliasing of transparent pixels. This
// package tries to avoid this by building frames from all previous frames and
// resizing the frames as RGB.
package gifresize

import (
	"image"
	"image/color/palette"
	"image/draw"
	"image/gif"
	"io"
)

// TransformFunc is a function that transforms an image.
type TransformFunc func(image.Image) image.Image

// Process the GIF read from r, applying transform to each frame, and writing
// the result to w.
func Process(w io.Writer, r io.Reader, transform TransformFunc) error {
	// Decode the original gif.
	im, err := gif.DecodeAll(r)
	if err != nil {
		return err
	}

	// Create a new RGBA image to hold the incremental frames.
	firstFrame := im.Image[0].Bounds()
	b := image.Rect(0, 0, firstFrame.Dx(), firstFrame.Dy())
	img := image.NewRGBA(b)

	// Resize each frame.
	for index, frame := range im.Image {
		bounds := frame.Bounds()
		draw.Draw(img, bounds, frame, bounds.Min, draw.Over)
		im.Image[index] = imageToPaletted(transform(img))
	}

	return gif.EncodeAll(w, im)
}

func imageToPaletted(img image.Image) *image.Paletted {
	b := img.Bounds()
	pm := image.NewPaletted(b, palette.Plan9)
	draw.FloydSteinberg.Draw(pm, b, img, image.ZP)
	return pm
}
<|endoftext|>


<commit_before><commit_msg>Expose `Config` in `SSOConfigureCommand` (#13060)<commit_after><|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package apiserver_test

import (
	stderrors "errors"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	apicommon "launchpad.net/juju-core/state/apiserver/common"
	"launchpad.net/juju-core/testing"
)

type errorsSuite struct {
	testing.LoggingSuite
}

var _ = Suite(&errorsSuite{})

var errorTransformTests = []struct {
	err  error
	code string
}{{
	err:  errors.NotFoundf("hello"),
	code: api.CodeNotFound,
}, {
	err:  errors.Unauthorizedf("hello"),
	code: api.CodeUnauthorized,
}, {
	err:  state.ErrCannotEnterScopeYet,
	code: api.CodeCannotEnterScopeYet,
}, {
	err:  state.ErrCannotEnterScope,
	code: api.CodeCannotEnterScope,
}, {
	err:  state.ErrExcessiveContention,
	code: api.CodeExcessiveContention,
}, {
	err:  state.ErrUnitHasSubordinates,
	code: api.CodeUnitHasSubordinates,
}, {
	err:  apicommon.ErrBadId,
	code: api.CodeNotFound,
}, {
	err:  apicommon.ErrBadCreds,
	code: api.CodeUnauthorized,
}, {
	err:  apicommon.ErrPerm,
	code: api.CodeUnauthorized,
}, {
	err:  apicommon.ErrNotLoggedIn,
	code: api.CodeUnauthorized,
}, {
	err:  apicommon.ErrUnknownWatcher,
	code: api.CodeNotFound,
}, {
	err:  &state.NotAssignedError{&state.Unit{}}, // too sleazy?! nah..
	code: api.CodeNotAssigned,
}, {
	err:  apicommon.ErrStoppedWatcher,
	code: api.CodeStopped,
}, {
	err:  &state.HasAssignedUnitsError{"42", []string{"a"}},
	code: api.CodeHasAssignedUnits,
}, {
	err:  stderrors.New("an error"),
	code: "",
}}

func (s *errorsSuite) TestErrorTransform(c *C) {
	for _, t := range errorTransformTests {
		err1 := apicommon.ServerError(t.err)
		c.Assert(err1.Error(), Equals, t.err.Error())
		if t.code != "" {
			c.Assert(api.ErrCode(err1), Equals, t.code)
		} else {
			c.Assert(err1, Equals, t.err)
		}
	}
}
<commit_msg>Removed api prefix from test package<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package apiserver_test

import (
	stderrors "errors"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/juju-core/state/apiserver/common"
	"launchpad.net/juju-core/testing"
)

type errorsSuite struct {
	testing.LoggingSuite
}

var _ = Suite(&errorsSuite{})

var errorTransformTests = []struct {
	err  error
	code string
}{{
	err:  errors.NotFoundf("hello"),
	code: api.CodeNotFound,
}, {
	err:  errors.Unauthorizedf("hello"),
	code: api.CodeUnauthorized,
}, {
	err:  state.ErrCannotEnterScopeYet,
	code: api.CodeCannotEnterScopeYet,
}, {
	err:  state.ErrCannotEnterScope,
	code: api.CodeCannotEnterScope,
}, {
	err:  state.ErrExcessiveContention,
	code: api.CodeExcessiveContention,
}, {
	err:  state.ErrUnitHasSubordinates,
	code: api.CodeUnitHasSubordinates,
}, {
	err:  common.ErrBadId,
	code: api.CodeNotFound,
}, {
	err:  common.ErrBadCreds,
	code: api.CodeUnauthorized,
}, {
	err:  common.ErrPerm,
	code: api.CodeUnauthorized,
}, {
	err:  common.ErrNotLoggedIn,
	code: api.CodeUnauthorized,
}, {
	err:  common.ErrUnknownWatcher,
	code: api.CodeNotFound,
}, {
	err:  &state.NotAssignedError{&state.Unit{}}, // too sleazy?! nah..
	code: api.CodeNotAssigned,
}, {
	err:  common.ErrStoppedWatcher,
	code: api.CodeStopped,
}, {
	err:  &state.HasAssignedUnitsError{"42", []string{"a"}},
	code: api.CodeHasAssignedUnits,
}, {
	err:  stderrors.New("an error"),
	code: "",
}}

func (s *errorsSuite) TestErrorTransform(c *C) {
	for _, t := range errorTransformTests {
		err1 := common.ServerError(t.err)
		c.Assert(err1.Error(), Equals, t.err.Error())
		if t.code != "" {
			c.Assert(api.ErrCode(err1), Equals, t.code)
		} else {
			c.Assert(err1, Equals, t.err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxdclient

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/juju/errors"
	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"

	"github.com/juju/juju/network"
)

// TODO(ericsnow) We probably need to address some of the things that
// get handled in container/lxc/clonetemplate.go.

type rawInstanceClient interface {
	ListContainers() ([]shared.ContainerInfo, error)
	ContainerInfo(name string) (*shared.ContainerInfo, error)
	Init(name string, imgremote string, image string, profiles *[]string, config map[string]string, ephem bool) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force bool, stateful bool) (*lxd.Response, error)
	Exec(name string, cmd []string, env map[string]string, stdin io.ReadCloser, stdout io.WriteCloser, stderr io.WriteCloser, controlHandler func(*lxd.Client, *websocket.Conn)) (int, error)
	Delete(name string) (*lxd.Response, error)

	WaitForSuccess(waitURL string) error
	ContainerState(name string) (*shared.ContainerState, error)
}

type instanceClient struct {
	raw    rawInstanceClient
	remote string
}

func (client *instanceClient) addInstance(spec InstanceSpec) error {
	imageRemote := spec.ImageRemote
	if imageRemote == "" {
		imageRemote = client.remote
	}

	imageAlias := spec.Image

	var profiles *[]string
	if len(spec.Profiles) > 0 {
		profiles = &spec.Profiles
	}

	// TODO(ericsnow) Copy the image first?

	config := spec.config()
	resp, err := client.raw.Init(spec.Name, imageRemote, imageAlias, profiles, config, spec.Ephemeral)
	if err != nil {
		return errors.Trace(err)
	}

	// Init is an async operation, since the tar -xvf (or whatever) might
	// take a while; the result is an LXD operation id, which we can just
	// wait on until it is finished.
	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

type execFailure struct {
	cmd    string
	code   int
	stderr string
}

// Error returns the string representation of the error.
func (err execFailure) Error() string {
	return fmt.Sprintf("got non-zero code from %q: (%d) %s", err.cmd, err.code, err.stderr)
}

func (client *instanceClient) exec(spec InstanceSpec, cmd []string) error {
	var env map[string]string

	cmdStr := strings.Join(cmd, " ")
	fmt.Println("running", cmdStr)

	var input, output closingBuffer
	stdin, stdout, stderr := &input, &output, &output
	rc, err := client.raw.Exec(spec.Name, cmd, env, stdin, stdout, stderr, nil)
	if err != nil {
		return errors.Trace(err)
	} else if rc != 0 {
		msg := output.String()
		if msg == "" {
			msg = "<reason unknown>"
		}
		err := &execFailure{
			cmd:    cmdStr,
			code:   rc,
			stderr: msg,
		}
		return errors.Trace(err)
	}

	return nil
}

func (client *instanceClient) chmod(spec InstanceSpec, filename string, mode os.FileMode) error {
	cmd := []string{
		"/bin/chmod",
		fmt.Sprintf("%s", mode),
		filename,
	}

	if err := client.exec(spec, cmd); err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (client *instanceClient) startInstance(spec InstanceSpec) error {
	timeout := -1
	force := false
	stateful := false
	resp, err := client.raw.Action(spec.Name, shared.Start, timeout, force, stateful)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// AddInstance creates a new instance based on the spec's data and
// returns it. The instance will be created using the client.
func (client *instanceClient) AddInstance(spec InstanceSpec) (*Instance, error) {
	if err := client.addInstance(spec); err != nil {
		return nil, errors.Trace(err)
	}

	if err := client.startInstance(spec); err != nil {
		if err := client.removeInstance(spec.Name); err != nil {
			logger.Errorf("could not remove container %q after starting it failed", spec.Name)
		}
		return nil, errors.Trace(err)
	}

	inst, err := client.Instance(spec.Name)
	if err != nil {
		return nil, errors.Trace(err)
	}
	inst.spec = &spec

	return inst, nil
}

// Instance gets the up-to-date info about the given instance
// and returns it.
func (client *instanceClient) Instance(name string) (*Instance, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return nil, errors.Trace(err)
	}

	inst := newInstance(info, nil)
	return inst, nil
}

func (client *instanceClient) Status(name string) (string, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return "", errors.Trace(err)
	}

	return info.Status, nil
}

// Instances sends a request to the API for a list of all instances
// (in the Client's namespace) for which the name starts with the
// provided prefix. The result is also limited to those instances with
// one of the specified statuses (if any).
func (client *instanceClient) Instances(prefix string, statuses ...string) ([]Instance, error) {
	infos, err := client.raw.ListContainers()
	if err != nil {
		return nil, errors.Trace(err)
	}

	var insts []Instance
	for _, info := range infos {
		name := info.Name
		if prefix != "" && !strings.HasPrefix(name, prefix) {
			continue
		}
		if len(statuses) > 0 && !checkStatus(info, statuses) {
			continue
		}

		inst := newInstance(&info, nil)
		insts = append(insts, *inst)
	}
	return insts, nil
}

func checkStatus(info shared.ContainerInfo, statuses []string) bool {
	for _, status := range statuses {
		statusCode := allStatuses[status]
		if info.StatusCode == statusCode {
			return true
		}
	}
	return false
}

// removeInstance sends a request to the API to remove the instance
// with the provided ID. The call blocks until the instance is removed
// (or the request fails).
func (client *instanceClient) removeInstance(name string) error {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return errors.Trace(err)
	}

	//if info.Status.StatusCode != 0 && info.Status.StatusCode != shared.Stopped {
	if info.StatusCode != shared.Stopped {
		timeout := -1
		force := true
		stateful := false
		resp, err := client.raw.Action(name, shared.Stop, timeout, force, stateful)
		if err != nil {
			return errors.Trace(err)
		}

		if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
			// TODO(ericsnow) Handle different failures (from the async
			// operation) differently?
			return errors.Trace(err)
		}
	}

	resp, err := client.raw.Delete(name)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// RemoveInstances sends a request to the API to terminate all
// instances (in the Client's namespace) that match one of the
// provided IDs. If a prefix is provided, only IDs that start with the
// prefix will be considered. The call blocks until all the instances
// are removed or the request fails.
func (client *instanceClient) RemoveInstances(prefix string, names ...string) error {
	if len(names) == 0 {
		return nil
	}

	instances, err := client.Instances(prefix)
	if err != nil {
		return errors.Annotatef(err, "while removing instances %v", names)
	}

	var failed []string
	for _, name := range names {
		if !checkInstanceName(name, instances) {
			// We ignore unknown instance names.
			continue
		}

		if err := client.removeInstance(name); err != nil {
			failed = append(failed, name)
			logger.Errorf("while removing instance %q: %v", name, err)
		}
	}
	if len(failed) != 0 {
		return errors.Errorf("some instance removals failed: %v", failed)
	}
	return nil
}

func checkInstanceName(name string, instances []Instance) bool {
	for _, inst := range instances {
		if inst.Name == name {
			return true
		}
	}
	return false
}

// Addresses returns the list of network.Addresses for this instance. It
// converts the information that LXD tracks into the Juju network model.
func (client *instanceClient) Addresses(name string) ([]network.Address, error) {
	state, err := client.raw.ContainerState(name)
	if err != nil {
		return nil, err
	}

	networks := state.Network
	if networks == nil {
		return []network.Address{}, nil
	}

	addrs := []network.Address{}

	for _, net := range networks {
		for _, addr := range net.Addresses {
			if err != nil {
				return nil, err
			}

			addr := network.NewAddress(addr.Address)
			if addr.Scope == network.ScopeLinkLocal || addr.Scope == network.ScopeMachineLocal {
				logger.Tracef("for container %q ignoring address", name, addr)
				continue
			}
			addrs = append(addrs, addr)
		}
	}
	return addrs, nil
}
<commit_msg>address error from go vet 'tools/lxdclient/client_instance.go:329: wrong number of args for format in Tracef call: 1 needed but 2 args'<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxdclient

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/juju/errors"
	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"

	"github.com/juju/juju/network"
)

// TODO(ericsnow) We probably need to address some of the things that
// get handled in container/lxc/clonetemplate.go.

type rawInstanceClient interface {
	ListContainers() ([]shared.ContainerInfo, error)
	ContainerInfo(name string) (*shared.ContainerInfo, error)
	Init(name string, imgremote string, image string, profiles *[]string, config map[string]string, ephem bool) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force bool, stateful bool) (*lxd.Response, error)
	Exec(name string, cmd []string, env map[string]string, stdin io.ReadCloser, stdout io.WriteCloser, stderr io.WriteCloser, controlHandler func(*lxd.Client, *websocket.Conn)) (int, error)
	Delete(name string) (*lxd.Response, error)

	WaitForSuccess(waitURL string) error
	ContainerState(name string) (*shared.ContainerState, error)
}

type instanceClient struct {
	raw    rawInstanceClient
	remote string
}

func (client *instanceClient) addInstance(spec InstanceSpec) error {
	imageRemote := spec.ImageRemote
	if imageRemote == "" {
		imageRemote = client.remote
	}

	imageAlias := spec.Image

	var profiles *[]string
	if len(spec.Profiles) > 0 {
		profiles = &spec.Profiles
	}

	// TODO(ericsnow) Copy the image first?

	config := spec.config()
	resp, err := client.raw.Init(spec.Name, imageRemote, imageAlias, profiles, config, spec.Ephemeral)
	if err != nil {
		return errors.Trace(err)
	}

	// Init is an async operation, since the tar -xvf (or whatever) might
	// take a while; the result is an LXD operation id, which we can just
	// wait on until it is finished.
	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

type execFailure struct {
	cmd    string
	code   int
	stderr string
}

// Error returns the string representation of the error.
func (err execFailure) Error() string {
	return fmt.Sprintf("got non-zero code from %q: (%d) %s", err.cmd, err.code, err.stderr)
}

func (client *instanceClient) exec(spec InstanceSpec, cmd []string) error {
	var env map[string]string

	cmdStr := strings.Join(cmd, " ")
	fmt.Println("running", cmdStr)

	var input, output closingBuffer
	stdin, stdout, stderr := &input, &output, &output
	rc, err := client.raw.Exec(spec.Name, cmd, env, stdin, stdout, stderr, nil)
	if err != nil {
		return errors.Trace(err)
	} else if rc != 0 {
		msg := output.String()
		if msg == "" {
			msg = "<reason unknown>"
		}
		err := &execFailure{
			cmd:    cmdStr,
			code:   rc,
			stderr: msg,
		}
		return errors.Trace(err)
	}

	return nil
}

func (client *instanceClient) chmod(spec InstanceSpec, filename string, mode os.FileMode) error {
	cmd := []string{
		"/bin/chmod",
		fmt.Sprintf("%s", mode),
		filename,
	}

	if err := client.exec(spec, cmd); err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (client *instanceClient) startInstance(spec InstanceSpec) error {
	timeout := -1
	force := false
	stateful := false
	resp, err := client.raw.Action(spec.Name, shared.Start, timeout, force, stateful)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// AddInstance creates a new instance based on the spec's data and
// returns it. The instance will be created using the client.
func (client *instanceClient) AddInstance(spec InstanceSpec) (*Instance, error) {
	if err := client.addInstance(spec); err != nil {
		return nil, errors.Trace(err)
	}

	if err := client.startInstance(spec); err != nil {
		if err := client.removeInstance(spec.Name); err != nil {
			logger.Errorf("could not remove container %q after starting it failed", spec.Name)
		}
		return nil, errors.Trace(err)
	}

	inst, err := client.Instance(spec.Name)
	if err != nil {
		return nil, errors.Trace(err)
	}
	inst.spec = &spec

	return inst, nil
}

// Instance gets the up-to-date info about the given instance
// and returns it.
func (client *instanceClient) Instance(name string) (*Instance, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return nil, errors.Trace(err)
	}

	inst := newInstance(info, nil)
	return inst, nil
}

func (client *instanceClient) Status(name string) (string, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return "", errors.Trace(err)
	}

	return info.Status, nil
}

// Instances sends a request to the API for a list of all instances
// (in the Client's namespace) for which the name starts with the
// provided prefix. The result is also limited to those instances with
// one of the specified statuses (if any).
func (client *instanceClient) Instances(prefix string, statuses ...string) ([]Instance, error) {
	infos, err := client.raw.ListContainers()
	if err != nil {
		return nil, errors.Trace(err)
	}

	var insts []Instance
	for _, info := range infos {
		name := info.Name
		if prefix != "" && !strings.HasPrefix(name, prefix) {
			continue
		}
		if len(statuses) > 0 && !checkStatus(info, statuses) {
			continue
		}

		inst := newInstance(&info, nil)
		insts = append(insts, *inst)
	}
	return insts, nil
}

func checkStatus(info shared.ContainerInfo, statuses []string) bool {
	for _, status := range statuses {
		statusCode := allStatuses[status]
		if info.StatusCode == statusCode {
			return true
		}
	}
	return false
}

// removeInstance sends a request to the API to remove the instance
// with the provided ID. The call blocks until the instance is removed
// (or the request fails).
func (client *instanceClient) removeInstance(name string) error {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return errors.Trace(err)
	}

	//if info.Status.StatusCode != 0 && info.Status.StatusCode != shared.Stopped {
	if info.StatusCode != shared.Stopped {
		timeout := -1
		force := true
		stateful := false
		resp, err := client.raw.Action(name, shared.Stop, timeout, force, stateful)
		if err != nil {
			return errors.Trace(err)
		}

		if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
			// TODO(ericsnow) Handle different failures (from the async
			// operation) differently?
			return errors.Trace(err)
		}
	}

	resp, err := client.raw.Delete(name)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// RemoveInstances sends a request to the API to terminate all
// instances (in the Client's namespace) that match one of the
// provided IDs. If a prefix is provided, only IDs that start with the
// prefix will be considered. The call blocks until all the instances
// are removed or the request fails.
func (client *instanceClient) RemoveInstances(prefix string, names ...string) error {
	if len(names) == 0 {
		return nil
	}

	instances, err := client.Instances(prefix)
	if err != nil {
		return errors.Annotatef(err, "while removing instances %v", names)
	}

	var failed []string
	for _, name := range names {
		if !checkInstanceName(name, instances) {
			// We ignore unknown instance names.
			continue
		}

		if err := client.removeInstance(name); err != nil {
			failed = append(failed, name)
			logger.Errorf("while removing instance %q: %v", name, err)
		}
	}
	if len(failed) != 0 {
		return errors.Errorf("some instance removals failed: %v", failed)
	}
	return nil
}

func checkInstanceName(name string, instances []Instance) bool {
	for _, inst := range instances {
		if inst.Name == name {
			return true
		}
	}
	return false
}

// Addresses returns the list of network.Addresses for this instance. It
// converts the information that LXD tracks into the Juju network model.
func (client *instanceClient) Addresses(name string) ([]network.Address, error) {
	state, err := client.raw.ContainerState(name)
	if err != nil {
		return nil, err
	}

	networks := state.Network
	if networks == nil {
		return []network.Address{}, nil
	}

	addrs := []network.Address{}

	for _, net := range networks {
		for _, addr := range net.Addresses {
			if err != nil {
				return nil, err
			}

			addr := network.NewAddress(addr.Address)
			if addr.Scope == network.ScopeLinkLocal || addr.Scope == network.ScopeMachineLocal {
				logger.Tracef("for container %q ignoring address", name)
				continue
			}
			addrs = append(addrs, addr)
		}
	}
	return addrs, nil
}
<|endoftext|>