<commit_before>// The package provides methods for converting amounts between currencies. The
// exchange rates are provided by the ECB (http://www.ecb.europa.eu/).
//
// Author: Michael Banzon
package currency

import (
	"fmt"
	"github.com/mbanzon/simplehttp"
	"time"
)

const ecbResourceUrl = "http://www.ecb.europa.eu/stats/eurofxref/eurofxref-daily.xml"

type Envelope struct {
	subject string
	Sender  string `xml:"Sender>name"`
	Cube    []Cube `xml:"Cube>Cube>Cube"`
}

type Cube struct {
	//time     string `xml:"time,attr"`
	currency string `xml:"currency,attr"`
	rate     string `xml:"rate,attr"`
}

type CurrencyConverter struct {
	date       time.Time
	currencies map[string]float64
}

func NewConverter() (*CurrencyConverter, error) {
	var e Envelope
	r := simplehttp.NewGetRequest(ecbResourceUrl)
	r.MakeXMLRequest(&e)

	fmt.Printf("%#v\n", e)

	var foo map[string]string
	re := simplehttp.NewGetRequest(ecbResourceUrl)
	re.MakeXMLRequest(&foo)

	fmt.Printf("%#v\n", foo)

	return &CurrencyConverter{}, nil
}
<commit_msg>Some parsing kinda works.<commit_after>// The package provides methods for converting amounts between currencies. The
// exchange rates are provided by the ECB (http://www.ecb.europa.eu/).
//
// Author: Michael Banzon
package currency

import (
	"fmt"
	"github.com/mbanzon/simplehttp"
	"time"
)

const ecbResourceUrl = "http://www.ecb.europa.eu/stats/eurofxref/eurofxref-daily.xml"

type Envelope struct {
	subject string
	Sender  string `xml:"Sender>name"`
	Cube    []Cube `xml:"Cube>Cube>Cube"`
	//	Time    TimeCube `xml:"Cube>Cube"`
}

type TimeCube struct {
	Time string `xml:"time,attr"`
}

type Cube struct {
	//time     string `xml:"time,attr"`
	Currency string `xml:"currency,attr"`
	Rate     string `xml:"rate,attr"`
}

type CurrencyConverter struct {
	date       time.Time
	currencies map[string]float64
}

func NewConverter() (*CurrencyConverter, error) {
	var e Envelope
	r := simplehttp.NewGetRequest(ecbResourceUrl)
	r.MakeXMLRequest(&e)

	fmt.Printf("%#v\n", e)

	var foo map[string]string
	re := simplehttp.NewGetRequest(ecbResourceUrl)
	re.MakeXMLRequest(&foo)

	fmt.Printf("%#v\n", foo)

	return &CurrencyConverter{}, nil
}
<|endoftext|>


<commit_before>package dag

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/hashicorp/go-multierror"
)

// walker performs a graph walk
type walker struct {
	Callback WalkFunc

	vertices  Set
	edges     Set
	vertexMap map[Vertex]*walkerVertex

	wait       sync.WaitGroup
	changeLock sync.Mutex

	errMap  map[Vertex]error
	errLock sync.Mutex
}

type walkerVertex struct {
	// These should only be set once on initialization and never written again
	DoneCh   chan struct{}
	CancelCh chan struct{}

	// Dependency information. Any changes to any of these fields requires
	// holding DepsLock.
	DepsCh       chan struct{}
	DepsUpdateCh chan struct{}
	DepsLock     sync.Mutex

	// Below is not safe to read/write in parallel. This behavior is
	// enforced by changes only happening in Update.
	deps         map[Vertex]chan struct{}
	depsCancelCh chan struct{}
}

// Wait waits for the completion of the walk and returns any errors (
// in the form of a multierror) that occurred. Update should be called
// to populate the walk with vertices and edges.
func (w *walker) Wait() error {
	// Wait for completion
	w.wait.Wait()

	// Grab the error lock
	w.errLock.Lock()
	defer w.errLock.Unlock()

	// Build the error
	var result error
	for v, err := range w.errMap {
		result = multierror.Append(result, fmt.Errorf(
			"%s: %s", VertexName(v), err))
	}

	return result
}

// Update updates the currently executing walk with the given vertices
// and edges. It does not block until completion.
//
// Update can be called in parallel to Walk.
func (w *walker) Update(v, e *Set) {
	// Grab the change lock so no more updates happen but also so that
	// no new vertices are executed during this time since we may be
	// removing them.
	w.changeLock.Lock()
	defer w.changeLock.Unlock()

	// Initialize fields
	if w.vertexMap == nil {
		w.vertexMap = make(map[Vertex]*walkerVertex)
	}

	// Calculate all our sets
	newEdges := e.Difference(&w.edges)
	oldEdges := w.edges.Difference(e)
	newVerts := v.Difference(&w.vertices)
	oldVerts := w.vertices.Difference(v)

	// Add the new vertices
	for _, raw := range newVerts.List() {
		v := raw.(Vertex)

		// Add to the waitgroup so our walk is not done until everything finishes
		w.wait.Add(1)

		// Add to our own set so we know about it already
		log.Printf("[DEBUG] dag/walk: added new vertex: %q", VertexName(v))
		w.vertices.Add(raw)

		// Initialize the vertex info
		info := &walkerVertex{
			DoneCh:   make(chan struct{}),
			CancelCh: make(chan struct{}),
			DepsCh:   make(chan struct{}),
			deps:     make(map[Vertex]chan struct{}),
		}

		// Close the deps channel immediately so it passes
		close(info.DepsCh)

		// Add it to the map and kick off the walk
		w.vertexMap[v] = info
	}

	// Remove the old vertices
	for _, raw := range oldVerts.List() {
		v := raw.(Vertex)

		// Get the vertex info so we can cancel it
		info, ok := w.vertexMap[v]
		if !ok {
			// This vertex for some reason was never in our map. This
			// shouldn't be possible.
			continue
		}

		// Cancel the vertex
		close(info.CancelCh)

		// Delete it out of the map
		delete(w.vertexMap, v)

		log.Printf("[DEBUG] dag/walk: removed vertex: %q", VertexName(v))
		w.vertices.Delete(raw)
	}

	// Add the new edges
	var changedDeps Set
	for _, raw := range newEdges.List() {
		edge := raw.(Edge)

		// waiter is the vertex that is "waiting" on this edge
		waiter := edge.Target()

		// dep is the dependency we're waiting on
		dep := edge.Source()

		// Get the info for the waiter
		waiterInfo, ok := w.vertexMap[waiter]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Get the info for the dep
		depInfo, ok := w.vertexMap[dep]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Add the dependency to our waiter
		waiterInfo.deps[dep] = depInfo.DoneCh

		// Record that the deps changed for this waiter
		changedDeps.Add(waiter)

		log.Printf(
			"[DEBUG] dag/walk: added edge: %q waiting on %q",
			VertexName(waiter), VertexName(dep))
		w.edges.Add(raw)
	}

	// Process reoved edges
	for _, raw := range oldEdges.List() {
		edge := raw.(Edge)

		// waiter is the vertex that is "waiting" on this edge
		waiter := edge.Target()

		// dep is the dependency we're waiting on
		dep := edge.Source()

		// Get the info for the waiter
		waiterInfo, ok := w.vertexMap[waiter]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Delete the dependency from the waiter
		delete(waiterInfo.deps, dep)

		// Record that the deps changed for this waiter
		changedDeps.Add(waiter)

		log.Printf(
			"[DEBUG] dag/walk: removed edge: %q waiting on %q",
			VertexName(waiter), VertexName(dep))
		w.edges.Delete(raw)
	}

	// For each vertex with changed dependencies, we need to kick off
	// a new waiter and notify the vertex of the changes.
	for _, raw := range changedDeps.List() {
		v := raw.(Vertex)
		info, ok := w.vertexMap[v]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Create a new done channel
		doneCh := make(chan struct{})

		// Create the channel we close for cancellation
		cancelCh := make(chan struct{})

		// Build a new deps copy
		deps := make(map[Vertex]<-chan struct{})
		for k, v := range info.deps {
			deps[k] = v
		}

		// Update the update channel
		info.DepsLock.Lock()
		if info.DepsUpdateCh != nil {
			close(info.DepsUpdateCh)
		}
		info.DepsCh = doneCh
		info.DepsUpdateCh = make(chan struct{})
		info.DepsLock.Unlock()

		// Cancel the older waiter
		if info.depsCancelCh != nil {
			close(info.depsCancelCh)
		}
		info.depsCancelCh = cancelCh

		// Start the waiter
		go w.waitDeps(v, deps, doneCh, cancelCh)
	}

	// Start all the new vertices. We do this at the end so that all
	// the edge waiters and changes are setup above.
	for _, raw := range newVerts.List() {
		v := raw.(Vertex)
		go w.walkVertex(v, w.vertexMap[v])
	}
}

// walkVertex walks a single vertex, waiting for any dependencies before
// executing the callback.
func (w *walker) walkVertex(v Vertex, info *walkerVertex) {
	// When we're done executing, lower the waitgroup count
	defer w.wait.Done()

	// When we're done, always close our done channel
	defer close(info.DoneCh)

	// Wait for our dependencies
	depsCh := info.DepsCh
	for {
		select {
		case <-info.CancelCh:
			// Cancel
			return

		case <-depsCh:
			// Deps complete!
			depsCh = nil

		case <-info.DepsUpdateCh:
			// New deps, reloop
		}

		// Check if we have updated dependencies
		info.DepsLock.Lock()
		if info.DepsCh != nil {
			depsCh = info.DepsCh
			info.DepsCh = nil
		}
		info.DepsLock.Unlock()

		// If we still have no deps channel set, then we're done!
		if depsCh == nil {
			break
		}
	}

	// Call our callback
	log.Printf("[DEBUG] dag/walk: walking %q", VertexName(v))
	if err := w.Callback(v); err != nil {
		w.errLock.Lock()
		defer w.errLock.Unlock()

		if w.errMap == nil {
			w.errMap = make(map[Vertex]error)
		}
		w.errMap[v] = err
	}
}

func (w *walker) waitDeps(
	v Vertex,
	deps map[Vertex]<-chan struct{},
	doneCh chan<- struct{},
	cancelCh <-chan struct{}) {
	// Whenever we return, mark ourselves as complete
	defer close(doneCh)

	// For each dependency given to us, wait for it to complete
	for dep, depCh := range deps {
	DepSatisfied:
		for {
			select {
			case <-depCh:
				// Dependency satisfied!
				break DepSatisfied

			case <-cancelCh:
				// Wait cancelled
				return

			case <-time.After(time.Second * 5):
				log.Printf("[DEBUG] vertex %q, waiting for: %q",
					VertexName(v), VertexName(dep))
			}
		}
	}
}
<commit_msg>dag: improved comments<commit_after>package dag

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/hashicorp/go-multierror"
)

// walker performs a graph walk and supports walk-time changing of vertices
// and edges.
//
// A single walker is only valid for one graph walk. After the walk is complete
// you must construct a new walker to walk again. State for the walk is never
// deleted in case vertices or edges are changed.
type walker struct {
	// Callback is what is called for each vertex
	Callback WalkFunc

	// changeLock must be held to modify any of the fields below. Only Update
	// should modify these fields. Modifying them outside of Update can cause
	// serious problems.
	changeLock sync.Mutex
	vertices   Set
	edges      Set
	vertexMap  map[Vertex]*walkerVertex

	// wait is done when all vertices have executed. It may become "undone"
	// if new vertices are added.
	wait sync.WaitGroup

	// errMap contains the errors recorded so far for execution. Reading
	// and writing should hold errLock.
	errMap  map[Vertex]error
	errLock sync.Mutex
}

type walkerVertex struct {
	// These should only be set once on initialization and never written again
	DoneCh   chan struct{}
	CancelCh chan struct{}

	// Dependency information. Any changes to any of these fields requires
	// holding DepsLock.
	DepsCh       chan struct{}
	DepsUpdateCh chan struct{}
	DepsLock     sync.Mutex

	// Below is not safe to read/write in parallel. This behavior is
	// enforced by changes only happening in Update.
	deps         map[Vertex]chan struct{}
	depsCancelCh chan struct{}
}

// Wait waits for the completion of the walk and returns any errors (
// in the form of a multierror) that occurred. Update should be called
// to populate the walk with vertices and edges prior to calling this.
//
// Wait will return as soon as all currently known vertices are complete.
// If you plan on calling Update with more vertices in the future, you
// should not call Wait until after this is done.
func (w *walker) Wait() error {
	// Wait for completion
	w.wait.Wait()

	// Grab the error lock
	w.errLock.Lock()
	defer w.errLock.Unlock()

	// Build the error
	var result error
	for v, err := range w.errMap {
		result = multierror.Append(result, fmt.Errorf(
			"%s: %s", VertexName(v), err))
	}

	return result
}

// Update updates the currently executing walk with the given vertices
// and edges. It does not block until completion.
//
// Update can be called in parallel to Walk.
func (w *walker) Update(v, e *Set) {
	// Grab the change lock so no more updates happen but also so that
	// no new vertices are executed during this time since we may be
	// removing them.
	w.changeLock.Lock()
	defer w.changeLock.Unlock()

	// Initialize fields
	if w.vertexMap == nil {
		w.vertexMap = make(map[Vertex]*walkerVertex)
	}

	// Calculate all our sets
	newEdges := e.Difference(&w.edges)
	oldEdges := w.edges.Difference(e)
	newVerts := v.Difference(&w.vertices)
	oldVerts := w.vertices.Difference(v)

	// Add the new vertices
	for _, raw := range newVerts.List() {
		v := raw.(Vertex)

		// Add to the waitgroup so our walk is not done until everything finishes
		w.wait.Add(1)

		// Add to our own set so we know about it already
		log.Printf("[DEBUG] dag/walk: added new vertex: %q", VertexName(v))
		w.vertices.Add(raw)

		// Initialize the vertex info
		info := &walkerVertex{
			DoneCh:   make(chan struct{}),
			CancelCh: make(chan struct{}),
			DepsCh:   make(chan struct{}),
			deps:     make(map[Vertex]chan struct{}),
		}

		// Close the deps channel immediately so it passes
		close(info.DepsCh)

		// Add it to the map and kick off the walk
		w.vertexMap[v] = info
	}

	// Remove the old vertices
	for _, raw := range oldVerts.List() {
		v := raw.(Vertex)

		// Get the vertex info so we can cancel it
		info, ok := w.vertexMap[v]
		if !ok {
			// This vertex for some reason was never in our map. This
			// shouldn't be possible.
			continue
		}

		// Cancel the vertex
		close(info.CancelCh)

		// Delete it out of the map
		delete(w.vertexMap, v)

		log.Printf("[DEBUG] dag/walk: removed vertex: %q", VertexName(v))
		w.vertices.Delete(raw)
	}

	// Add the new edges
	var changedDeps Set
	for _, raw := range newEdges.List() {
		edge := raw.(Edge)

		// waiter is the vertex that is "waiting" on this edge
		waiter := edge.Target()

		// dep is the dependency we're waiting on
		dep := edge.Source()

		// Get the info for the waiter
		waiterInfo, ok := w.vertexMap[waiter]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Get the info for the dep
		depInfo, ok := w.vertexMap[dep]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Add the dependency to our waiter
		waiterInfo.deps[dep] = depInfo.DoneCh

		// Record that the deps changed for this waiter
		changedDeps.Add(waiter)

		log.Printf(
			"[DEBUG] dag/walk: added edge: %q waiting on %q",
			VertexName(waiter), VertexName(dep))
		w.edges.Add(raw)
	}

	// Process reoved edges
	for _, raw := range oldEdges.List() {
		edge := raw.(Edge)

		// waiter is the vertex that is "waiting" on this edge
		waiter := edge.Target()

		// dep is the dependency we're waiting on
		dep := edge.Source()

		// Get the info for the waiter
		waiterInfo, ok := w.vertexMap[waiter]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Delete the dependency from the waiter
		delete(waiterInfo.deps, dep)

		// Record that the deps changed for this waiter
		changedDeps.Add(waiter)

		log.Printf(
			"[DEBUG] dag/walk: removed edge: %q waiting on %q",
			VertexName(waiter), VertexName(dep))
		w.edges.Delete(raw)
	}

	// For each vertex with changed dependencies, we need to kick off
	// a new waiter and notify the vertex of the changes.
	for _, raw := range changedDeps.List() {
		v := raw.(Vertex)
		info, ok := w.vertexMap[v]
		if !ok {
			// Vertex doesn't exist... shouldn't be possible but ignore.
			continue
		}

		// Create a new done channel
		doneCh := make(chan struct{})

		// Create the channel we close for cancellation
		cancelCh := make(chan struct{})

		// Build a new deps copy
		deps := make(map[Vertex]<-chan struct{})
		for k, v := range info.deps {
			deps[k] = v
		}

		// Update the update channel
		info.DepsLock.Lock()
		if info.DepsUpdateCh != nil {
			close(info.DepsUpdateCh)
		}
		info.DepsCh = doneCh
		info.DepsUpdateCh = make(chan struct{})
		info.DepsLock.Unlock()

		// Cancel the older waiter
		if info.depsCancelCh != nil {
			close(info.depsCancelCh)
		}
		info.depsCancelCh = cancelCh

		// Start the waiter
		go w.waitDeps(v, deps, doneCh, cancelCh)
	}

	// Start all the new vertices. We do this at the end so that all
	// the edge waiters and changes are setup above.
	for _, raw := range newVerts.List() {
		v := raw.(Vertex)
		go w.walkVertex(v, w.vertexMap[v])
	}
}

// walkVertex walks a single vertex, waiting for any dependencies before
// executing the callback.
func (w *walker) walkVertex(v Vertex, info *walkerVertex) {
	// When we're done executing, lower the waitgroup count
	defer w.wait.Done()

	// When we're done, always close our done channel
	defer close(info.DoneCh)

	// Wait for our dependencies
	depsCh := info.DepsCh
	for {
		select {
		case <-info.CancelCh:
			// Cancel
			return

		case <-depsCh:
			// Deps complete!
			depsCh = nil

		case <-info.DepsUpdateCh:
			// New deps, reloop
		}

		// Check if we have updated dependencies. This can happen if the
		// dependencies were satisfied exactly prior to an Update occuring.
		// In that case, we'd like to take into account new dependencies
		// if possible.
		info.DepsLock.Lock()
		if info.DepsCh != nil {
			depsCh = info.DepsCh
			info.DepsCh = nil
		}
		info.DepsLock.Unlock()

		// If we still have no deps channel set, then we're done!
		if depsCh == nil {
			break
		}
	}

	// Call our callback
	log.Printf("[DEBUG] dag/walk: walking %q", VertexName(v))
	if err := w.Callback(v); err != nil {
		w.errLock.Lock()
		defer w.errLock.Unlock()

		if w.errMap == nil {
			w.errMap = make(map[Vertex]error)
		}
		w.errMap[v] = err
	}
}

func (w *walker) waitDeps(
	v Vertex,
	deps map[Vertex]<-chan struct{},
	doneCh chan<- struct{},
	cancelCh <-chan struct{}) {
	// Whenever we return, mark ourselves as complete
	defer close(doneCh)

	// For each dependency given to us, wait for it to complete
	for dep, depCh := range deps {
	DepSatisfied:
		for {
			select {
			case <-depCh:
				// Dependency satisfied!
				break DepSatisfied

			case <-cancelCh:
				// Wait cancelled
				return

			case <-time.After(time.Second * 5):
				log.Printf("[DEBUG] vertex %q, waiting for: %q",
					VertexName(v), VertexName(dep))
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/ci-pipeline/cloudformation-resource/utils"
)

func main() {
	utils.GoToBuildDirectory()
	input := utils.GetInput()
	svc := utils.GetCloudformationService(input)
	metadata, success := out(input, svc, &utils.AwsRequestHandler{})
	fmt.Printf("%s", metadata)

	if !success {
		os.Exit(1)
	}
}

type Parameter struct {
	ParameterValue   string
	ParameterKey     string
	UsePreviousValue bool
}

type Tag struct {
	TagKey   string
	TagValue string
}

func stackExists(reqHandler utils.RequestHandler, svc utils.AwsCloudformationSvc, stackName string) bool {
	params := &cloudformation.DescribeStacksInput{
		StackName: aws.String(stackName),
	}
	req, resp := svc.DescribeStacksRequest(params)
	err := reqHandler.HandleRequest(req)
	if err != nil {
		return false
	}
	fmt.Println(resp)

	return resp.Stacks[0] != nil
}

func waitForStack(reqHandler utils.RequestHandler, svc utils.AwsCloudformationSvc, input utils.Input) (success bool, arn, timestamp string) {
	success = true
	params := &cloudformation.DescribeStackEventsInput{
		StackName: aws.String(input.Source.Name),
	}

	pos := 0
	var stackEvents []*cloudformation.StackEvent
	for {
		req, resp := svc.DescribeStackEventsRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			success = input.Params.Delete
			return
		}
		stackEvents = resp.StackEvents
		for j := len(stackEvents) - 1 - pos; j > -1; j-- {
			fmt.Println(stackEvents[j])
		}
		pos = len(stackEvents)

		arn = *stackEvents[0].StackId
		timestamp = stackEvents[0].Timestamp.String()

		if *stackEvents[0].ResourceType == "AWS::CloudFormation::Stack" {
			if *stackEvents[0].ResourceStatus == "CREATE_COMPLETE" {
				return
			} else if *stackEvents[0].ResourceStatus == "UPDATE_COMPLETE" {
				return
			} else if *stackEvents[0].ResourceStatus == "ROLLBACK_COMPLETE" {
				success = false
				return
			} else if *stackEvents[0].ResourceStatus == "UPDATE_ROLLBACK_COMPLETE" {
				success = false
				return
			} else if strings.HasSuffix(*stackEvents[0].ResourceStatus, "_FAILED") {
				success = false
				return
			}
		}

		time.Sleep(500 * time.Millisecond)
	}
}

func out(input utils.Input, svc utils.AwsCloudformationSvc, reqHandler utils.RequestHandler) (metadata string, success bool) {

	var capabilities []*string
	capabilities = nil
	if input.Params.Capabilities != nil {
		for _, c := range input.Params.Capabilities {
			capabilities = append(capabilities, aws.String(c))
		}
	}

	var templateBody string
	if input.Params.Template != "" {
		bytes, err := ioutil.ReadFile(input.Params.Template)
		if err != nil {
			panic(err)
		}
		templateBody = string(bytes)
	}

	parameters := []Parameter{}
	var cloudformationParams []*cloudformation.Parameter
	if input.Params.Parameters != "" {
		bytes, err := ioutil.ReadFile(input.Params.Parameters)
		err = json.Unmarshal(bytes, &parameters)
		if err != nil {
			panic(err)
		}
	}
	for _, p := range parameters {
		param := cloudformation.Parameter{
			ParameterKey:     aws.String(p.ParameterKey),
			ParameterValue:   aws.String(p.ParameterValue),
			UsePreviousValue: aws.Bool(p.UsePreviousValue),
		}
		cloudformationParams = append(cloudformationParams, &param)
	}

	tags := []Tag{}
	var cloudformationTags []*cloudformation.Tag
	if input.Params.Tags != "" {
		bytes, err := ioutil.ReadFile(input.Params.Tags)
		err = json.Unmarshal(bytes, &tags)
		if err != nil {
			panic(err)
		}
	}
	for _, t := range tags {
		tag := cloudformation.Tag{
			Key:   aws.String(t.TagKey),
			Value: aws.String(t.TagValue),
		}
		cloudformationTags = append(cloudformationTags, &tag)
	}

	if input.Params.Delete == true {
		params := &cloudformation.DeleteStackInput{
			StackName: aws.String(input.Source.Name),
		}
		req, resp := svc.DeleteStackRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println(resp)
	} else if !stackExists(reqHandler, svc, input.Source.Name) {
		params := &cloudformation.CreateStackInput{
			StackName:    aws.String(input.Source.Name),
			Capabilities: capabilities,
			Parameters:   cloudformationParams,
			Tags:         cloudformationTags,
			TemplateBody: aws.String(templateBody),
		}
		req, resp := svc.CreateStackRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println(resp)
	} else {
		params := &cloudformation.UpdateStackInput{
			StackName:    aws.String(input.Source.Name),
			Capabilities: capabilities,
			Parameters:   cloudformationParams,
			Tags:         cloudformationTags,
			TemplateBody: aws.String(templateBody),
		}
		req, resp := svc.UpdateStackRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println(resp)
	}

	success, arn, timestamp := waitForStack(reqHandler, svc, input)
	result := make(map[string]string)
	result["arn"] = arn
	result["timestamp"] = timestamp
	bytes, _ := json.Marshal(result)
	return string(bytes), success
}
<commit_msg>Print action<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/ci-pipeline/cloudformation-resource/utils"
)

func main() {
	utils.GoToBuildDirectory()
	input := utils.GetInput()
	svc := utils.GetCloudformationService(input)
	metadata, success := out(input, svc, &utils.AwsRequestHandler{})
	fmt.Printf("%s", metadata)

	if !success {
		os.Exit(1)
	}
}

type Parameter struct {
	ParameterValue   string
	ParameterKey     string
	UsePreviousValue bool
}

type Tag struct {
	TagKey   string
	TagValue string
}

func stackExists(reqHandler utils.RequestHandler, svc utils.AwsCloudformationSvc, stackName string) bool {
	params := &cloudformation.DescribeStacksInput{
		StackName: aws.String(stackName),
	}
	req, resp := svc.DescribeStacksRequest(params)
	err := reqHandler.HandleRequest(req)
	if err != nil {
		return false
	}
	fmt.Println(resp)

	return resp.Stacks[0] != nil
}

func waitForStack(reqHandler utils.RequestHandler, svc utils.AwsCloudformationSvc, input utils.Input) (success bool, arn, timestamp string) {
	success = true
	params := &cloudformation.DescribeStackEventsInput{
		StackName: aws.String(input.Source.Name),
	}

	pos := 0
	var stackEvents []*cloudformation.StackEvent
	for {
		req, resp := svc.DescribeStackEventsRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			success = input.Params.Delete
			return
		}
		stackEvents = resp.StackEvents
		for j := len(stackEvents) - 1 - pos; j > -1; j-- {
			fmt.Println(stackEvents[j])
		}
		pos = len(stackEvents)

		arn = *stackEvents[0].StackId
		timestamp = stackEvents[0].Timestamp.String()

		if *stackEvents[0].ResourceType == "AWS::CloudFormation::Stack" {
			if *stackEvents[0].ResourceStatus == "CREATE_COMPLETE" {
				return
			} else if *stackEvents[0].ResourceStatus == "UPDATE_COMPLETE" {
				return
			} else if *stackEvents[0].ResourceStatus == "ROLLBACK_COMPLETE" {
				success = false
				return
			} else if *stackEvents[0].ResourceStatus == "UPDATE_ROLLBACK_COMPLETE" {
				success = false
				return
			} else if strings.HasSuffix(*stackEvents[0].ResourceStatus, "_FAILED") {
				success = false
				return
			}
		}

		time.Sleep(10 * time.Millisecond)
	}
}

func out(input utils.Input, svc utils.AwsCloudformationSvc, reqHandler utils.RequestHandler) (metadata string, success bool) {

	var capabilities []*string
	capabilities = nil
	if input.Params.Capabilities != nil {
		for _, c := range input.Params.Capabilities {
			capabilities = append(capabilities, aws.String(c))
		}
	}

	var templateBody string
	if input.Params.Template != "" {
		bytes, err := ioutil.ReadFile(input.Params.Template)
		if err != nil {
			panic(err)
		}
		templateBody = string(bytes)
	}

	parameters := []Parameter{}
	var cloudformationParams []*cloudformation.Parameter
	if input.Params.Parameters != "" {
		bytes, err := ioutil.ReadFile(input.Params.Parameters)
		err = json.Unmarshal(bytes, &parameters)
		if err != nil {
			panic(err)
		}
	}
	for _, p := range parameters {
		param := cloudformation.Parameter{
			ParameterKey:     aws.String(p.ParameterKey),
			ParameterValue:   aws.String(p.ParameterValue),
			UsePreviousValue: aws.Bool(p.UsePreviousValue),
		}
		cloudformationParams = append(cloudformationParams, &param)
	}

	tags := []Tag{}
	var cloudformationTags []*cloudformation.Tag
	if input.Params.Tags != "" {
		bytes, err := ioutil.ReadFile(input.Params.Tags)
		err = json.Unmarshal(bytes, &tags)
		if err != nil {
			panic(err)
		}
	}
	for _, t := range tags {
		tag := cloudformation.Tag{
			Key:   aws.String(t.TagKey),
			Value: aws.String(t.TagValue),
		}
		cloudformationTags = append(cloudformationTags, &tag)
	}

	if input.Params.Delete == true {
		fmt.Println("Deleting stack")
		params := &cloudformation.DeleteStackInput{
			StackName: aws.String(input.Source.Name),
		}
		req, resp := svc.DeleteStackRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println(resp)
	} else if !stackExists(reqHandler, svc, input.Source.Name) {
		fmt.Println("Creating stack")
		params := &cloudformation.CreateStackInput{
			StackName:    aws.String(input.Source.Name),
			Capabilities: capabilities,
			Parameters:   cloudformationParams,
			Tags:         cloudformationTags,
			TemplateBody: aws.String(templateBody),
		}
		req, resp := svc.CreateStackRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println(resp)
	} else {
		fmt.Println("Updating stack")
		params := &cloudformation.UpdateStackInput{
			StackName:    aws.String(input.Source.Name),
			Capabilities: capabilities,
			Parameters:   cloudformationParams,
			Tags:         cloudformationTags,
			TemplateBody: aws.String(templateBody),
		}
		req, resp := svc.UpdateStackRequest(params)
		err := reqHandler.HandleRequest(req)
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println(resp)
	}

	success, arn, timestamp := waitForStack(reqHandler, svc, input)
	result := make(map[string]string)
	result["arn"] = arn
	result["timestamp"] = timestamp
	bytes, _ := json.Marshal(result)
	return string(bytes), success
}
<|endoftext|>


<commit_before>package db

// Level DB cached wrapper to improve write performance using batching.

import (
	"firempq/common"
	"firempq/conf"
	"firempq/log"
	"sync"
	"time"

	. "firempq/api"

	"github.com/jmhodges/levigo"
)

// Default LevelDB read options.
var defaultReadOptions = levigo.NewReadOptions()

// Default LevelDB write options.
var defaultWriteOptions = levigo.NewWriteOptions()

// LevelDBStorage A high level cached structure on top of LevelDB.
// It caches item storing them into database
// as multiple large batches later.
type LevelDBStorage struct {
	db             *levigo.DB        // Pointer the the instance of level db.
	dbName         string            // LevelDB database name.
	itemCache      map[string]string // Active cache for item metadata.
	tmpItemCache   map[string]string // Active cache during flush operation.
	cacheLock      sync.Mutex        // Used for caches access.
	flushLock      sync.Mutex        // Used to prevent double flush.
	closed         bool
	flushSync      *sync.WaitGroup // Use to wait until flush happens.
	forceFlushChan chan bool
}

// NewLevelDBStorage is a constructor of DataStorage.
func NewLevelDBStorage(dbName string) (*LevelDBStorage, error) {
	ds := LevelDBStorage{
		dbName:         dbName,
		itemCache:      make(map[string]string),
		tmpItemCache:   nil,
		closed:         false,
		forceFlushChan: make(chan bool, 1),
		flushSync:      &sync.WaitGroup{},
	}

	// LevelDB write options.
	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	opts.SetWriteBufferSize(10 * 1024 * 1024)
	opts.SetCompression(levigo.SnappyCompression)

	db, err := levigo.Open(dbName, opts)
	if err != nil {
		return nil, err
	}
	ds.db = db
	ds.flushSync.Add(1)
	go ds.periodicCacheFlush()
	return &ds, nil
}

func (ds *LevelDBStorage) periodicCacheFlush() {
	for !ds.closed {
		select {
		case <-ds.forceFlushChan:
			break
		case <-time.After(conf.CFG.DbFlushInterval * time.Millisecond):
			break
		}
		ds.flushLock.Lock()
		oldFlushSync := ds.flushSync
		ds.flushSync = &sync.WaitGroup{}
		ds.flushSync.Add(1)
		if !ds.closed {
			ds.flushCache()
		}
		oldFlushSync.Done()
		ds.flushLock.Unlock()
	}
	ds.flushSync.Done()
}

// WaitFlush waits until all data is flushed on disk.
func (ds *LevelDBStorage) WaitFlush() {
	ds.flushLock.Lock()
	s := ds.flushSync
	ds.flushLock.Unlock()
	s.Wait()
}

// CachedStoreItem stores data into the cache.
func (ds *LevelDBStorage) CachedStore(data ...string) {
	if len(data)%2 != 0 {
		panic("Number of arguments must be even!")
	}
	ds.cacheLock.Lock()
	for i := 0; i < len(data); i += 2 {
		ds.itemCache[data[i]] = data[i+1]
	}
	ds.cacheLock.Unlock()
}

// DeleteDataWithPrefix deletes all service data such as service metadata, items and payloads.
func (ds *LevelDBStorage) DeleteDataWithPrefix(prefix string) int {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	ds.flushCache()

	limitCounter := 0
	total := 0
	iter := ds.IterData(prefix)
	wb := levigo.NewWriteBatch()

	for iter.Valid() {
		total++
		if limitCounter < 1000 {
			wb.Delete(iter.GetKey())
			limitCounter++
		} else {
			limitCounter = 0
			ds.db.Write(defaultWriteOptions, wb)
			wb = levigo.NewWriteBatch()
		}
		iter.Next()
	}

	ds.db.Write(defaultWriteOptions, wb)
	return total
}

// FlushCache flushes all cache into database.
func (ds *LevelDBStorage) flushCache() {
	ds.cacheLock.Lock()
	ds.tmpItemCache = ds.itemCache
	ds.itemCache = make(map[string]string)
	ds.cacheLock.Unlock()

	wb := levigo.NewWriteBatch()
	for k, v := range ds.tmpItemCache {
		key := common.UnsafeStringToBytes(k)
		if v == "" {
			wb.Delete(key)
		} else {
			wb.Put(key, common.UnsafeStringToBytes(v))
		}
	}
	ds.db.Write(defaultWriteOptions, wb)
}

// IterData returns an iterator over all data with prefix.
func (ds *LevelDBStorage) IterData(prefix string) ItemIterator {
	iter := ds.db.NewIterator(defaultReadOptions)
	return makeItemIterator(iter, common.UnsafeStringToBytes(prefix))
}

// StoreData data directly into the database stores service metadata into database.
func (ds *LevelDBStorage) StoreData(data ...string) error {
	if len(data)%2 != 0 {
		panic("Number of arguments must be even!")
	}
	wb := levigo.NewWriteBatch()
	for i := 0; i < len(data); i += 2 {
		wb.Put(common.UnsafeStringToBytes(data[i]),
			common.UnsafeStringToBytes(data[i+1]))
	}
	return ds.db.Write(defaultWriteOptions, wb)
}

func (ds *LevelDBStorage) DeleteData(id ...string) {
	wb := levigo.NewWriteBatch()
	for _, i := range id {
		wb.Delete(common.UnsafeStringToBytes(i))
	}
	ds.db.Write(defaultWriteOptions, wb)
}

// GetData looks data looks for and item going through each layer of cache finally looking into database.
func (ds *LevelDBStorage) GetData(id string) string {
	ds.cacheLock.Lock()
	data, ok := ds.itemCache[id]
	if ok {
		ds.cacheLock.Unlock()
		return data
	}
	data, ok = ds.tmpItemCache[id]
	if ok {
		ds.cacheLock.Unlock()
		return data
	}
	ds.cacheLock.Unlock()
	value, _ := ds.db.Get(defaultReadOptions, common.UnsafeStringToBytes(id))
	return common.UnsafeBytesToString(value)
}

// CachedDeleteData deletes item metadata and payload, affects cache only until flushed.
func (ds *LevelDBStorage) CachedDeleteData(id ...string) {
	ds.cacheLock.Lock()
	for _, i := range id {
		ds.itemCache[i] = ""
	}
	ds.cacheLock.Unlock()
}

func (ds *LevelDBStorage) IsClosed() bool {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	return ds.closed
}

// Close flushes data on disk and closes database.
func (ds *LevelDBStorage) Close() {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	if !ds.closed {
		ds.flushCache()
		ds.closed = true
		ds.db.Close()
	} else {
		log.Error("Attempt to close database more than once!")
	}
}
<commit_msg>Adding task moved to the gorouting level.<commit_after>package db

// Level DB cached wrapper to improve write performance using batching.

import (
	"firempq/common"
	"firempq/conf"
	"firempq/log"
	"sync"
	"time"

	. "firempq/api"

	"github.com/jmhodges/levigo"
)

// Default LevelDB read options.
var defaultReadOptions = levigo.NewReadOptions()

// Default LevelDB write options.
var defaultWriteOptions = levigo.NewWriteOptions()

// LevelDBStorage A high level cached structure on top of LevelDB.
// It caches item storing them into database
// as multiple large batches later.
type LevelDBStorage struct {
	db             *levigo.DB        // Pointer the the instance of level db.
	dbName         string            // LevelDB database name.
	itemCache      map[string]string // Active cache for item metadata.
	tmpItemCache   map[string]string // Active cache during flush operation.
	cacheLock      sync.Mutex        // Used for caches access.
	flushLock      sync.Mutex        // Used to prevent double flush.
	closed         bool
	flushSync      *sync.WaitGroup // Use to wait until flush happens.
	forceFlushChan chan bool
}

// NewLevelDBStorage is a constructor of DataStorage.
func NewLevelDBStorage(dbName string) (*LevelDBStorage, error) {
	ds := LevelDBStorage{
		dbName:         dbName,
		itemCache:      make(map[string]string),
		tmpItemCache:   nil,
		closed:         false,
		forceFlushChan: make(chan bool, 1),
		flushSync:      &sync.WaitGroup{},
	}

	// LevelDB write options.
	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	opts.SetWriteBufferSize(10 * 1024 * 1024)
	opts.SetCompression(levigo.SnappyCompression)

	db, err := levigo.Open(dbName, opts)
	if err != nil {
		return nil, err
	}
	ds.db = db
	go ds.periodicCacheFlush()
	return &ds, nil
}

func (ds *LevelDBStorage) periodicCacheFlush() {
	ds.flushSync.Add(1)
	for !ds.closed {
		select {
		case <-ds.forceFlushChan:
			break
		case <-time.After(conf.CFG.DbFlushInterval * time.Millisecond):
			break
		}
		ds.flushLock.Lock()
		oldFlushSync := ds.flushSync
		ds.flushSync = &sync.WaitGroup{}
		ds.flushSync.Add(1)
		if !ds.closed {
			ds.flushCache()
		}
		oldFlushSync.Done()
		ds.flushLock.Unlock()
	}
	ds.flushSync.Done()
}

// WaitFlush waits until all data is flushed on disk.
func (ds *LevelDBStorage) WaitFlush() {
	ds.flushLock.Lock()
	s := ds.flushSync
	ds.flushLock.Unlock()
	s.Wait()
}

// CachedStoreItem stores data into the cache.
func (ds *LevelDBStorage) CachedStore(data ...string) {
	if len(data)%2 != 0 {
		panic("Number of arguments must be even!")
	}
	ds.cacheLock.Lock()
	for i := 0; i < len(data); i += 2 {
		ds.itemCache[data[i]] = data[i+1]
	}
	ds.cacheLock.Unlock()
}

// DeleteDataWithPrefix deletes all service data such as service metadata, items and payloads.
func (ds *LevelDBStorage) DeleteDataWithPrefix(prefix string) int {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	ds.flushCache()

	limitCounter := 0
	total := 0
	iter := ds.IterData(prefix)
	wb := levigo.NewWriteBatch()

	for iter.Valid() {
		total++
		if limitCounter < 1000 {
			wb.Delete(iter.GetKey())
			limitCounter++
		} else {
			limitCounter = 0
			ds.db.Write(defaultWriteOptions, wb)
			wb = levigo.NewWriteBatch()
		}
		iter.Next()
	}

	ds.db.Write(defaultWriteOptions, wb)
	return total
}

// FlushCache flushes all cache into database.
func (ds *LevelDBStorage) flushCache() {
	ds.cacheLock.Lock()
	ds.tmpItemCache = ds.itemCache
	ds.itemCache = make(map[string]string)
	ds.cacheLock.Unlock()

	wb := levigo.NewWriteBatch()
	for k, v := range ds.tmpItemCache {
		key := common.UnsafeStringToBytes(k)
		if v == "" {
			wb.Delete(key)
		} else {
			wb.Put(key, common.UnsafeStringToBytes(v))
		}
	}
	ds.db.Write(defaultWriteOptions, wb)
}

// IterData returns an iterator over all data with prefix.
func (ds *LevelDBStorage) IterData(prefix string) ItemIterator {
	iter := ds.db.NewIterator(defaultReadOptions)
	return makeItemIterator(iter, common.UnsafeStringToBytes(prefix))
}

// StoreData data directly into the database stores service metadata into database.
func (ds *LevelDBStorage) StoreData(data ...string) error {
	if len(data)%2 != 0 {
		panic("Number of arguments must be even!")
	}
	wb := levigo.NewWriteBatch()
	for i := 0; i < len(data); i += 2 {
		wb.Put(common.UnsafeStringToBytes(data[i]),
			common.UnsafeStringToBytes(data[i+1]))
	}
	return ds.db.Write(defaultWriteOptions, wb)
}

func (ds *LevelDBStorage) DeleteData(id ...string) {
	wb := levigo.NewWriteBatch()
	for _, i := range id {
		wb.Delete(common.UnsafeStringToBytes(i))
	}
	ds.db.Write(defaultWriteOptions, wb)
}

// GetData looks data looks for and item going through each layer of cache finally looking into database.
func (ds *LevelDBStorage) GetData(id string) string {
	ds.cacheLock.Lock()
	data, ok := ds.itemCache[id]
	if ok {
		ds.cacheLock.Unlock()
		return data
	}
	data, ok = ds.tmpItemCache[id]
	if ok {
		ds.cacheLock.Unlock()
		return data
	}
	ds.cacheLock.Unlock()
	value, _ := ds.db.Get(defaultReadOptions, common.UnsafeStringToBytes(id))
	return common.UnsafeBytesToString(value)
}

// CachedDeleteData deletes item metadata and payload, affects cache only until flushed.
func (ds *LevelDBStorage) CachedDeleteData(id ...string) {
	ds.cacheLock.Lock()
	for _, i := range id {
		ds.itemCache[i] = ""
	}
	ds.cacheLock.Unlock()
}

func (ds *LevelDBStorage) IsClosed() bool {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	return ds.closed
}

// Close flushes data on disk and closes database.
func (ds *LevelDBStorage) Close() {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	if !ds.closed {
		ds.flushCache()
		ds.closed = true
		ds.db.Close()
	} else {
		log.Error("Attempt to close database more than once!")
	}
}
<|endoftext|>


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

// Package db encapsulates tsuru connection with MongoDB.
//
// The function Conn dials to MongoDB using data from the configuration file
// and returns a connection (represented by the storage.Storage type). It
// manages an internal pool of connections, and reconnects in case of failures.
// That means that you should not store references to the connection, but
// always call Open.
package db

import (
	"fmt"

	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/db/storage"
	"github.com/tsuru/tsuru/hc"
	"gopkg.in/mgo.v2"
)

const (
	DefaultDatabaseURL  = "127.0.0.1:27017"
	DefaultDatabaseName = "tsuru"
)

type Storage struct {
	*storage.Storage
}

type LogStorage struct {
	*storage.Storage
}

func init() {
	hc.AddChecker("MongoDB", healthCheck)
}

func healthCheck() error {
	conn, err := Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Apps().Database.Session.Ping()
}

func DbConfig(prefix string) (string, string) {
	url, _ := config.GetString(fmt.Sprintf("database:%surl", prefix))
	if url == "" {
		url, _ = config.GetString("database:url")
		if url == "" {
			url = DefaultDatabaseURL
		}
	}
	dbname, _ := config.GetString(fmt.Sprintf("database:%sname", prefix))
	if dbname == "" {
		dbname, _ = config.GetString("database:name")
		if dbname == "" {
			dbname = DefaultDatabaseName
		}
	}
	return url, dbname
}

// Conn reads the tsuru config and calls storage.Open to get a database connection.
//
// Most tsuru packages should probably use this function. storage.Open is intended for
// use when supporting more than one database.
func Conn() (*Storage, error) {
	var (
		strg Storage
		err  error
	)
	url, dbname := DbConfig("")
	strg.Storage, err = storage.Open(url, dbname)
	return &strg, err
}

func LogConn() (*LogStorage, error) {
	var (
		strg LogStorage
		err  error
	)
	url, dbname := DbConfig("logdb-")
	strg.Storage, err = storage.Open(url, dbname)
	return &strg, err
}

// Apps returns the apps collection from MongoDB.
func (s *Storage) Apps() *storage.Collection {
	nameIndex := mgo.Index{Key: []string{"name"}, Unique: true}
	c := s.Collection("apps")
	c.EnsureIndex(nameIndex)
	return c
}

// Platforms returns the platforms collection from MongoDB.
func (s *Storage) Platforms() *storage.Collection {
	return s.Collection("platforms")
}

// Services returns the services collection from MongoDB.
func (s *Storage) Services() *storage.Collection {
	return s.Collection("services")
}

// ServiceInstances returns the services_instances collection from MongoDB.
func (s *Storage) ServiceInstances() *storage.Collection {
	return s.Collection("service_instances")
}

// Plans returns the plans collection.
func (s *Storage) Plans() *storage.Collection {
	return s.Collection("plans")
}

// Pools returns the pool collection.
func (s *Storage) Pools() *storage.Collection {
	return s.Collection("pool")
}

// PoolsConstraints return the pool constraints collection.
func (s *Storage) PoolsContraints() *storage.Collection {
	poolConstraintIndex := mgo.Index{Key: []string{"poolexpr", "field"}, Unique: true}
	c := s.Collection("pool_constraints")
	c.EnsureIndex(poolConstraintIndex)
	return c
}

// Users returns the users collection from MongoDB.
func (s *Storage) Users() *storage.Collection {
	emailIndex := mgo.Index{Key: []string{"email"}, Unique: true}
	c := s.Collection("users")
	c.EnsureIndex(emailIndex)
	return c
}

func (s *Storage) Tokens() *storage.Collection {
	coll := s.Collection("tokens")
	coll.EnsureIndex(mgo.Index{Key: []string{"token"}})
	return coll
}

func (s *Storage) PasswordTokens() *storage.Collection {
	return s.Collection("password_tokens")
}

func (s *Storage) UserActions() *storage.Collection {
	return s.Collection("user_actions")
}

// Teams returns the teams collection from MongoDB.
func (s *Storage) Teams() *storage.Collection {
	return s.Collection("teams")
}

// Quota returns the quota collection from MongoDB.
func (s *Storage) Quota() *storage.Collection {
	userIndex := mgo.Index{Key: []string{"owner"}, Unique: true}
	c := s.Collection("quota")
	c.EnsureIndex(userIndex)
	return c
}

// SAMLRequests returns the saml_requests from MongoDB.
func (s *Storage) SAMLRequests() *storage.Collection {
	id := mgo.Index{Key: []string{"id"}}
	coll := s.Collection("saml_requests")
	coll.EnsureIndex(id)
	return coll
}

var logCappedInfo = mgo.CollectionInfo{
	Capped:       true,
	MaxBytes:     200 * 5000,
	MaxDocs:      5000,
	ForceIdIndex: true,
}

// Logs returns the logs collection for one app from MongoDB.
func (s *LogStorage) Logs(appName string) *storage.Collection {
	if appName == "" {
		return nil
	}
	c := s.Collection("logs_" + appName)
	c.Create(&logCappedInfo)
	return c
}

// LogsCollections returns logs collections for all apps from MongoDB.
func (s *LogStorage) LogsCollections() ([]*storage.Collection, error) {
	var names []struct {
		Name string
	}
	conn, err := Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	err = conn.Apps().Find(nil).All(&names)
	if err != nil {
		return nil, err
	}
	var colls []*storage.Collection
	for _, name := range names {
		colls = append(colls, s.Collection("logs_"+name.Name))
	}
	return colls, nil
}

func (s *Storage) Roles() *storage.Collection {
	return s.Collection("roles")
}

func (s *Storage) Limiter() *storage.Collection {
	return s.Collection("limiter")
}

func (s *Storage) Events() *storage.Collection {
	ownerIndex := mgo.Index{Key: []string{"owner"}}
	kindIndex := mgo.Index{Key: []string{"kind"}}
	startTimeIndex := mgo.Index{Key: []string{"-starttime"}}
	c := s.Collection("events")
	c.EnsureIndex(ownerIndex)
	c.EnsureIndex(kindIndex)
	c.EnsureIndex(startTimeIndex)
	return c
}

func (s *Storage) InstallHosts() *storage.Collection {
	nameIndex := mgo.Index{Key: []string{"name"}, Unique: true}
	c := s.Collection("install_hosts")
	c.EnsureIndex(nameIndex)
	return c
}
<commit_msg>db: fix typo<commit_after>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package db encapsulates tsuru connection with MongoDB.
//
// The function Conn dials to MongoDB using data from the configuration file
// and returns a connection (represented by the storage.Storage type). It
// manages an internal pool of connections, and reconnects in case of failures.
// That means that you should not store references to the connection, but
// always call Open.
package db

import (
	"fmt"

	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/db/storage"
	"github.com/tsuru/tsuru/hc"
	"gopkg.in/mgo.v2"
)

const (
	DefaultDatabaseURL  = "127.0.0.1:27017"
	DefaultDatabaseName = "tsuru"
)

type Storage struct {
	*storage.Storage
}

type LogStorage struct {
	*storage.Storage
}

func init() {
	hc.AddChecker("MongoDB", healthCheck)
}

func healthCheck() error {
	conn, err := Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Apps().Database.Session.Ping()
}

func DbConfig(prefix string) (string, string) {
	url, _ := config.GetString(fmt.Sprintf("database:%surl", prefix))
	if url == "" {
		url, _ = config.GetString("database:url")
		if url == "" {
			url = DefaultDatabaseURL
		}
	}
	dbname, _ := config.GetString(fmt.Sprintf("database:%sname", prefix))
	if dbname == "" {
		dbname, _ = config.GetString("database:name")
		if dbname == "" {
			dbname = DefaultDatabaseName
		}
	}
	return url, dbname
}

// Conn reads the tsuru config and calls storage.Open to get a database connection.
//
// Most tsuru packages should probably use this function. storage.Open is intended for
// use when supporting more than one database.
func Conn() (*Storage, error) {
	var (
		strg Storage
		err  error
	)
	url, dbname := DbConfig("")
	strg.Storage, err = storage.Open(url, dbname)
	return &strg, err
}

func LogConn() (*LogStorage, error) {
	var (
		strg LogStorage
		err  error
	)
	url, dbname := DbConfig("logdb-")
	strg.Storage, err = storage.Open(url, dbname)
	return &strg, err
}

// Apps returns the apps collection from MongoDB.
func (s *Storage) Apps() *storage.Collection {
	nameIndex := mgo.Index{Key: []string{"name"}, Unique: true}
	c := s.Collection("apps")
	c.EnsureIndex(nameIndex)
	return c
}

// Platforms returns the platforms collection from MongoDB.
func (s *Storage) Platforms() *storage.Collection {
	return s.Collection("platforms")
}

// Services returns the services collection from MongoDB.
func (s *Storage) Services() *storage.Collection {
	return s.Collection("services")
}

// ServiceInstances returns the services_instances collection from MongoDB.
func (s *Storage) ServiceInstances() *storage.Collection {
	return s.Collection("service_instances")
}

// Plans returns the plans collection.
func (s *Storage) Plans() *storage.Collection {
	return s.Collection("plans")
}

// Pools returns the pool collection.
func (s *Storage) Pools() *storage.Collection {
	return s.Collection("pool")
}

// PoolsConstraints return the pool constraints collection.
func (s *Storage) PoolsConstraints() *storage.Collection {
	poolConstraintIndex := mgo.Index{Key: []string{"poolexpr", "field"}, Unique: true}
	c := s.Collection("pool_constraints")
	c.EnsureIndex(poolConstraintIndex)
	return c
}

// Users returns the users collection from MongoDB.
func (s *Storage) Users() *storage.Collection {
	emailIndex := mgo.Index{Key: []string{"email"}, Unique: true}
	c := s.Collection("users")
	c.EnsureIndex(emailIndex)
	return c
}

func (s *Storage) Tokens() *storage.Collection {
	coll := s.Collection("tokens")
	coll.EnsureIndex(mgo.Index{Key: []string{"token"}})
	return coll
}

func (s *Storage) PasswordTokens() *storage.Collection {
	return s.Collection("password_tokens")
}

func (s *Storage) UserActions() *storage.Collection {
	return s.Collection("user_actions")
}

// Teams returns the teams collection from MongoDB.
func (s *Storage) Teams() *storage.Collection {
	return s.Collection("teams")
}

// Quota returns the quota collection from MongoDB.
func (s *Storage) Quota() *storage.Collection {
	userIndex := mgo.Index{Key: []string{"owner"}, Unique: true}
	c := s.Collection("quota")
	c.EnsureIndex(userIndex)
	return c
}

// SAMLRequests returns the saml_requests from MongoDB.
func (s *Storage) SAMLRequests() *storage.Collection {
	id := mgo.Index{Key: []string{"id"}}
	coll := s.Collection("saml_requests")
	coll.EnsureIndex(id)
	return coll
}

var logCappedInfo = mgo.CollectionInfo{
	Capped:       true,
	MaxBytes:     200 * 5000,
	MaxDocs:      5000,
	ForceIdIndex: true,
}

// Logs returns the logs collection for one app from MongoDB.
func (s *LogStorage) Logs(appName string) *storage.Collection {
	if appName == "" {
		return nil
	}
	c := s.Collection("logs_" + appName)
	c.Create(&logCappedInfo)
	return c
}

// LogsCollections returns logs collections for all apps from MongoDB.
func (s *LogStorage) LogsCollections() ([]*storage.Collection, error) {
	var names []struct {
		Name string
	}
	conn, err := Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	err = conn.Apps().Find(nil).All(&names)
	if err != nil {
		return nil, err
	}
	var colls []*storage.Collection
	for _, name := range names {
		colls = append(colls, s.Collection("logs_"+name.Name))
	}
	return colls, nil
}

func (s *Storage) Roles() *storage.Collection {
	return s.Collection("roles")
}

func (s *Storage) Limiter() *storage.Collection {
	return s.Collection("limiter")
}

func (s *Storage) Events() *storage.Collection {
	ownerIndex := mgo.Index{Key: []string{"owner"}}
	kindIndex := mgo.Index{Key: []string{"kind"}}
	startTimeIndex := mgo.Index{Key: []string{"-starttime"}}
	c := s.Collection("events")
	c.EnsureIndex(ownerIndex)
	c.EnsureIndex(kindIndex)
	c.EnsureIndex(startTimeIndex)
	return c
}

func (s *Storage) InstallHosts() *storage.Collection {
	nameIndex := mgo.Index{Key: []string{"name"}, Unique: true}
	c := s.Collection("install_hosts")
	c.EnsureIndex(nameIndex)
	return c
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/dustin/seriesly/timelib"
	"github.com/mschoch/gouchstore"
)

type dbOperation uint8

const (
	opStoreItem = dbOperation(iota)
	opDeleteItem
	opCompact
)

const dbExt = ".couch"

type dbqitem struct {
	dbname string
	k      string
	data   []byte
	op     dbOperation
	cherr  chan error
}

type dbWriter struct {
	dbname string
	ch     chan dbqitem
	quit   chan bool
	db     *gouchstore.Gouchstore
}

var errClosed = errors.New("closed")

func (w *dbWriter) Close() error {
	select {
	case <-w.quit:
		return errClosed
	default:
	}
	close(w.quit)
	return nil
}

var dbLock = sync.Mutex{}
var dbConns = map[string]*dbWriter{}

func dbPath(n string) string {
	return filepath.Join(*dbRoot, n) + dbExt
}

func dbBase(n string) string {
	left := 0
	right := len(n)
	if strings.HasPrefix(n, *dbRoot) {
		left = len(*dbRoot)
		if n[left] == '/' {
			left++
		}
	}
	if strings.HasSuffix(n, dbExt) {
		right = len(n) - len(dbExt)
	}
	return n[left:right]
}

func dbopen(name string) (*gouchstore.Gouchstore, error) {
	path := dbPath(name)
	db, err := gouchstore.Open(dbPath(name), 0)
	if err == nil {
		recordDBConn(path, db)
	}
	return db, err
}

func dbcreate(path string) error {
	db, err := gouchstore.Open(path, gouchstore.OPEN_CREATE)
	if err != nil {
		return err
	}
	recordDBConn(path, db)
	closeDBConn(db)
	return nil
}

func dbRemoveConn(dbname string) {
	dbLock.Lock()
	defer dbLock.Unlock()

	writer := dbConns[dbname]
	if writer != nil && writer.quit != nil {
		writer.Close()
	}
	delete(dbConns, dbname)
}

func dbCloseAll() {
	dbLock.Lock()
	defer dbLock.Unlock()

	for n, c := range dbConns {
		log.Printf("Shutting down open conn %v", n)
		c.Close()
	}
}

func dbdelete(dbname string) error {
	return os.Remove(dbPath(dbname))
}

func dblist(root string) []string {
	rv := []string{}
	filepath.Walk(root, func(p string, info os.FileInfo, err error) error {
		if err == nil {
			if !info.IsDir() && strings.HasSuffix(p, dbExt) {
				rv = append(rv, dbBase(p))
			}
		} else {
			log.Printf("Error on %#v: %v", p, err)
		}
		return nil
	})
	return rv
}

func dbCompact(dq *dbWriter, bulk gouchstore.BulkWriter, queued int,
	qi dbqitem) (gouchstore.BulkWriter, error) {
	start := time.Now()
	if queued > 0 {
		bulk.Commit()
		if *verbose {
			log.Printf("Flushed %d items in %v for pre-compact",
				queued, time.Since(start))
		}
		bulk.Close()
	}
	dbn := dbPath(dq.dbname)
	queued = 0
	start = time.Now()
	err := dq.db.Compact(dbn + ".compact")
	if err != nil {
		log.Printf("Error compacting: %v", err)
		return dq.db.Bulk(), err
	}
	log.Printf("Finished compaction of %v in %v", dq.dbname,
		time.Since(start))
	err = os.Rename(dbn+".compact", dbn)
	if err != nil {
		log.Printf("Error putting compacted data back")
		return dq.db.Bulk(), err
	}

	log.Printf("Reopening post-compact")
	closeDBConn(dq.db)

	dq.db, err = dbopen(dq.dbname)
	if err != nil {
		log.Fatalf("Error reopening DB after compaction: %v", err)
	}
	return dq.db.Bulk(), nil
}

var dbWg = sync.WaitGroup{}

func dbWriteLoop(dq *dbWriter) {
	defer dbWg.Done()

	queued := 0
	bulk := dq.db.Bulk()

	t := time.NewTimer(*flushTime)
	defer t.Stop()
	liveTracker := time.NewTicker(*liveTime)
	defer liveTracker.Stop()
	liveOps := 0

	for {
		select {
		case <-dq.quit:
			sdt := time.Now()
			bulk.Close()
			bulk.Commit()
			closeDBConn(dq.db)
			dbRemoveConn(dq.dbname)
			log.Printf("Closed %v in %v", dq.dbname, time.Since(sdt))
			return
		case <-liveTracker.C:
			if queued == 0 && liveOps == 0 {
				log.Printf("Closing idle DB: %v", dq.dbname)
				close(dq.quit)
			}
			liveOps = 0
		case qi := <-dq.ch:
			liveOps++
			switch qi.op {
			case opStoreItem:
				bulk.Set(gouchstore.NewDocumentInfo(qi.k),
					gouchstore.NewDocument(qi.k, qi.data))
				queued++
			case opDeleteItem:
				queued++
				bulk.Delete(gouchstore.NewDocumentInfo(qi.k))
			case opCompact:
				var err error
				bulk, err = dbCompact(dq, bulk, queued, qi)
				qi.cherr <- err
				queued = 0
			default:
				log.Panicf("Unhandled case: %v", qi.op)
			}
			if queued >= *maxOpQueue {
				start := time.Now()
				bulk.Commit()
				if *verbose {
					log.Printf("Flush of %d items took %v",
						queued, time.Since(start))
				}
				queued = 0
				t.Reset(*flushTime)
			}
		case <-t.C:
			if queued > 0 {
				start := time.Now()
				bulk.Commit()
				if *verbose {
					log.Printf("Flush of %d items from timer took %v",
						queued, time.Since(start))
				}
				queued = 0
			}
			t.Reset(*flushTime)
		}
	}
}

func dbWriteFun(dbname string) (*dbWriter, error) {
	db, err := dbopen(dbname)
	if err != nil {
		return nil, err
	}

	writer := &dbWriter{
		dbname,
		make(chan dbqitem, *maxOpQueue),
		make(chan bool),
		db,
	}

	dbWg.Add(1)
	go dbWriteLoop(writer)

	return writer, nil
}

func getOrCreateDB(dbname string) (*dbWriter, bool, error) {
	dbLock.Lock()
	defer dbLock.Unlock()

	writer := dbConns[dbname]
	var err error
	opened := false
	if writer == nil {
		writer, err = dbWriteFun(dbname)
		if err != nil {
			return nil, false, err
		}
		dbConns[dbname] = writer
		opened = true
	}
	return writer, opened, nil
}

func dbstore(dbname string, k string, body []byte) error {
	writer, _, err := getOrCreateDB(dbname)
	if err != nil {
		return err
	}

	writer.ch <- dbqitem{dbname, k, body, opStoreItem, nil}

	return nil
}

func dbcompact(dbname string) error {
	writer, opened, err := getOrCreateDB(dbname)
	if err != nil {
		return err
	}
	if opened {
		log.Printf("Requesting post-compaction close of %v", dbname)
		defer writer.Close()
	}

	cherr := make(chan error)
	defer close(cherr)
	writer.ch <- dbqitem{dbname: dbname,
		op:    opCompact,
		cherr: cherr,
	}

	return <-cherr
}

func dbGetDoc(dbname, id string) ([]byte, error) {
	db, err := dbopen(dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", dbname, err)
		return nil, err
	}
	defer closeDBConn(db)

	doc, err := db.DocumentById(id)
	if err != nil {
		return nil, err
	}
	return doc.Body, err
}

func dbwalk(dbname, from, to string, f func(k string, v []byte) error) error {
	db, err := dbopen(dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", dbname, err)
		return err
	}
	defer closeDBConn(db)

	return db.WalkDocs(from, to, func(d *gouchstore.Gouchstore,
		di *gouchstore.DocumentInfo, doc *gouchstore.Document) error {
		return f(di.ID, doc.Body)
	})
}

func dbwalkKeys(dbname, from, to string, f func(k string) error) error {
	db, err := dbopen(dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", dbname, err)
		return err
	}
	defer closeDBConn(db)

	return db.AllDocuments(from, to, func(db *gouchstore.Gouchstore, documentInfo *gouchstore.DocumentInfo, userContext interface{}) error {
		return f(documentInfo.ID)
	}, nil)
}

func parseKey(s string) int64 {
	t, err := timelib.ParseCanonicalTime(s)
	if err != nil {
		return -1
	}
	return t.UnixNano()
}
<commit_msg>Log queue size on close<commit_after>package main

import (
	"errors"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/dustin/seriesly/timelib"
	"github.com/mschoch/gouchstore"
)

type dbOperation uint8

const (
	opStoreItem = dbOperation(iota)
	opDeleteItem
	opCompact
)

const dbExt = ".couch"

type dbqitem struct {
	dbname string
	k      string
	data   []byte
	op     dbOperation
	cherr  chan error
}

type dbWriter struct {
	dbname string
	ch     chan dbqitem
	quit   chan bool
	db     *gouchstore.Gouchstore
}

var errClosed = errors.New("closed")

func (w *dbWriter) Close() error {
	select {
	case <-w.quit:
		return errClosed
	default:
	}
	close(w.quit)
	return nil
}

var dbLock = sync.Mutex{}
var dbConns = map[string]*dbWriter{}

func dbPath(n string) string {
	return filepath.Join(*dbRoot, n) + dbExt
}

func dbBase(n string) string {
	left := 0
	right := len(n)
	if strings.HasPrefix(n, *dbRoot) {
		left = len(*dbRoot)
		if n[left] == '/' {
			left++
		}
	}
	if strings.HasSuffix(n, dbExt) {
		right = len(n) - len(dbExt)
	}
	return n[left:right]
}

func dbopen(name string) (*gouchstore.Gouchstore, error) {
	path := dbPath(name)
	db, err := gouchstore.Open(dbPath(name), 0)
	if err == nil {
		recordDBConn(path, db)
	}
	return db, err
}

func dbcreate(path string) error {
	db, err := gouchstore.Open(path, gouchstore.OPEN_CREATE)
	if err != nil {
		return err
	}
	recordDBConn(path, db)
	closeDBConn(db)
	return nil
}

func dbRemoveConn(dbname string) {
	dbLock.Lock()
	defer dbLock.Unlock()

	writer := dbConns[dbname]
	if writer != nil && writer.quit != nil {
		writer.Close()
	}
	delete(dbConns, dbname)
}

func dbCloseAll() {
	dbLock.Lock()
	defer dbLock.Unlock()

	for n, c := range dbConns {
		log.Printf("Shutting down open conn %v", n)
		c.Close()
	}
}

func dbdelete(dbname string) error {
	return os.Remove(dbPath(dbname))
}

func dblist(root string) []string {
	rv := []string{}
	filepath.Walk(root, func(p string, info os.FileInfo, err error) error {
		if err == nil {
			if !info.IsDir() && strings.HasSuffix(p, dbExt) {
				rv = append(rv, dbBase(p))
			}
		} else {
			log.Printf("Error on %#v: %v", p, err)
		}
		return nil
	})
	return rv
}

func dbCompact(dq *dbWriter, bulk gouchstore.BulkWriter, queued int,
	qi dbqitem) (gouchstore.BulkWriter, error) {
	start := time.Now()
	if queued > 0 {
		bulk.Commit()
		if *verbose {
			log.Printf("Flushed %d items in %v for pre-compact",
				queued, time.Since(start))
		}
		bulk.Close()
	}
	dbn := dbPath(dq.dbname)
	queued = 0
	start = time.Now()
	err := dq.db.Compact(dbn + ".compact")
	if err != nil {
		log.Printf("Error compacting: %v", err)
		return dq.db.Bulk(), err
	}
	log.Printf("Finished compaction of %v in %v", dq.dbname,
		time.Since(start))
	err = os.Rename(dbn+".compact", dbn)
	if err != nil {
		log.Printf("Error putting compacted data back")
		return dq.db.Bulk(), err
	}

	log.Printf("Reopening post-compact")
	closeDBConn(dq.db)

	dq.db, err = dbopen(dq.dbname)
	if err != nil {
		log.Fatalf("Error reopening DB after compaction: %v", err)
	}
	return dq.db.Bulk(), nil
}

var dbWg = sync.WaitGroup{}

func dbWriteLoop(dq *dbWriter) {
	defer dbWg.Done()

	queued := 0
	bulk := dq.db.Bulk()

	t := time.NewTimer(*flushTime)
	defer t.Stop()
	liveTracker := time.NewTicker(*liveTime)
	defer liveTracker.Stop()
	liveOps := 0

	for {
		select {
		case <-dq.quit:
			sdt := time.Now()
			bulk.Close()
			bulk.Commit()
			closeDBConn(dq.db)
			dbRemoveConn(dq.dbname)
			log.Printf("Closed %v with %v items in %v",
				dq.dbname, queued, time.Since(sdt))
			return
		case <-liveTracker.C:
			if queued == 0 && liveOps == 0 {
				log.Printf("Closing idle DB: %v", dq.dbname)
				close(dq.quit)
			}
			liveOps = 0
		case qi := <-dq.ch:
			liveOps++
			switch qi.op {
			case opStoreItem:
				bulk.Set(gouchstore.NewDocumentInfo(qi.k),
					gouchstore.NewDocument(qi.k, qi.data))
				queued++
			case opDeleteItem:
				queued++
				bulk.Delete(gouchstore.NewDocumentInfo(qi.k))
			case opCompact:
				var err error
				bulk, err = dbCompact(dq, bulk, queued, qi)
				qi.cherr <- err
				queued = 0
			default:
				log.Panicf("Unhandled case: %v", qi.op)
			}
			if queued >= *maxOpQueue {
				start := time.Now()
				bulk.Commit()
				if *verbose {
					log.Printf("Flush of %d items took %v",
						queued, time.Since(start))
				}
				queued = 0
				t.Reset(*flushTime)
			}
		case <-t.C:
			if queued > 0 {
				start := time.Now()
				bulk.Commit()
				if *verbose {
					log.Printf("Flush of %d items from timer took %v",
						queued, time.Since(start))
				}
				queued = 0
			}
			t.Reset(*flushTime)
		}
	}
}

func dbWriteFun(dbname string) (*dbWriter, error) {
	db, err := dbopen(dbname)
	if err != nil {
		return nil, err
	}

	writer := &dbWriter{
		dbname,
		make(chan dbqitem, *maxOpQueue),
		make(chan bool),
		db,
	}

	dbWg.Add(1)
	go dbWriteLoop(writer)

	return writer, nil
}

func getOrCreateDB(dbname string) (*dbWriter, bool, error) {
	dbLock.Lock()
	defer dbLock.Unlock()

	writer := dbConns[dbname]
	var err error
	opened := false
	if writer == nil {
		writer, err = dbWriteFun(dbname)
		if err != nil {
			return nil, false, err
		}
		dbConns[dbname] = writer
		opened = true
	}
	return writer, opened, nil
}

func dbstore(dbname string, k string, body []byte) error {
	writer, _, err := getOrCreateDB(dbname)
	if err != nil {
		return err
	}

	writer.ch <- dbqitem{dbname, k, body, opStoreItem, nil}

	return nil
}

func dbcompact(dbname string) error {
	writer, opened, err := getOrCreateDB(dbname)
	if err != nil {
		return err
	}
	if opened {
		log.Printf("Requesting post-compaction close of %v", dbname)
		defer writer.Close()
	}

	cherr := make(chan error)
	defer close(cherr)
	writer.ch <- dbqitem{dbname: dbname,
		op:    opCompact,
		cherr: cherr,
	}

	return <-cherr
}

func dbGetDoc(dbname, id string) ([]byte, error) {
	db, err := dbopen(dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", dbname, err)
		return nil, err
	}
	defer closeDBConn(db)

	doc, err := db.DocumentById(id)
	if err != nil {
		return nil, err
	}
	return doc.Body, err
}

func dbwalk(dbname, from, to string, f func(k string, v []byte) error) error {
	db, err := dbopen(dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", dbname, err)
		return err
	}
	defer closeDBConn(db)

	return db.WalkDocs(from, to, func(d *gouchstore.Gouchstore,
		di *gouchstore.DocumentInfo, doc *gouchstore.Document) error {
		return f(di.ID, doc.Body)
	})
}

func dbwalkKeys(dbname, from, to string, f func(k string) error) error {
	db, err := dbopen(dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", dbname, err)
		return err
	}
	defer closeDBConn(db)

	return db.AllDocuments(from, to, func(db *gouchstore.Gouchstore, documentInfo *gouchstore.DocumentInfo, userContext interface{}) error {
		return f(documentInfo.ID)
	}, nil)
}

func parseKey(s string) int64 {
	t, err := timelib.ParseCanonicalTime(s)
	if err != nil {
		return -1
	}
	return t.UnixNano()
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"github.com/jmoiron/sqlx"
	"github.com/pborman/uuid"
	"github.com/pkg/errors"
	"net/http"
	"strings"
	"time"
)

type Database struct {
	db *sqlx.DB
}

type Resolution int

const (
	All Resolution = iota
	Day
	Hour
	Minute
)

func (resolution Resolution) ToString() (string, error) {
	switch resolution {
	case All:
		return "All", nil
	case Minute:
		return "Minute", nil
	case Hour:
		return "Hour", nil
	case Day:
		return "Day", nil
	default:
		return "", errors.New("Unknown resolution")
	}
}

func (resolution Resolution) String() string {
	str, _ := resolution.ToString()
	return str
}

func ResolutionFromString(resolutionString string) (Resolution, error) {
	lowerCase := strings.ToLower(resolutionString)
	if lowerCase == strings.ToLower("All") {
		return All, nil
	} else if lowerCase == strings.ToLower("Minute") {
		return Minute, nil
	} else if lowerCase == strings.ToLower("Hour") {
		return Hour, nil
	} else if lowerCase == strings.ToLower("Day") {
		return Day, nil
	} else {
		return All, errors.New("Unknown resolution from string: " + resolutionString)
	}
}

func (db *Database) readDataFromPlot(user string, plotId int, startTime time.Time, endTime time.Time, resolution Resolution) ([]Measurement, error) {

	if resolution == All {
		return db.readAllDataFromPlot(user, plotId, startTime, endTime)
	} else {
		return db.readAggregatedDataFromPlot(user, plotId, startTime, endTime, resolution)
	}
}

func (db *Database) readAllDataFromPlot(user string, plotId int, startTime time.Time, endTime time.Time) ([]Measurement, error) {
	measurements := []Measurement{}
	var sql = `
        WITH instruments as (
            SELECT key FROM
            instrument
            WHERE plot = $1
        )
        SELECT m.key, m.value, m.timestamp
        FROM measurement m, plot p
        WHERE m.timestamp >= p.start_time
        AND(p.end_time is null OR m.timestamp <= p.end_time)
        AND m.key IN (SELECT key from instruments)
        AND p.id = $1
        AND p.login = $2
        AND m.timestamp >= $3
        AND m.timestamp <= $4
        ORDER BY m.timestamp desc;
    `
	err := db.db.Select(&measurements, sql, plotId, user, startTime, endTime)
	return measurements, err
}

func (db *Database) readLatestDataFromPlot(user string, plotId int) ([]Measurement, error) {
	measurements := []Measurement{}
	var sql = `
        WITH
        instruments as (
            SELECT i.key AS keys
            FROM instrument i
            WHERE plot = $1
        ),
        latest_measurement as (

            SELECT max(m.timestamp) as timestamp
            FROM measurement m, plot p
            WHERE m.timestamp >= p.start_time
            AND(p.end_time is null OR m.timestamp <= p.end_time)
            AND p.id = $1
        )
        SELECT m.key, m.value, m.timestamp
        FROM measurement m, plot p
        WHERE m.timestamp >= p.start_time
        and m.timestamp = (select timestamp from latest_measurement)
        AND m.key IN (SELECT i.keys from instruments i)
        AND p.id = $1
        AND p.login = $2
        ORDER BY m.timestamp desc;
    `
	err := db.db.Select(&measurements, sql, plotId, user)
	return measurements, err
}

func (db *Database) getIntervalDefinition(resolution Resolution) (string, string, error) {
	switch resolution {
	case All:
		return "", "", errors.New("No interval definition for All")
	case Minute:
		return "mins", "1 minute", nil
	case Hour:
		return "hours", "1 hour", nil
	case Day:
		return "days", "1 day", nil
	default:
		return "", "", errors.New("Unknown resolution")
	}
}

func (db *Database) readAggregatedDataFromPlot(user string, plotId int, startTime time.Time, endTime time.Time, resolution Resolution) ([]Measurement, error) {
	measurements := []Measurement{}
	var sql = `
        WITH instruments as (
            SELECT key AS keys
            FROM instrument
            WHERE plot = $1
        ),
        intervals AS (
            SELECT start_time FROM
            generate_series(date_trunc($5, GREATEST($3, (select start_time from plot where id = $1))),
                            LEAST($4, NOW()),
                            $6) as start_time
        )
        SELECT
            m.key,
            i.start_time as timestamp,
            AVG(m.value) as value
        FROM measurement m, plot p, intervals i
        WHERE m.timestamp >= p.start_time
        AND(p.end_time is null OR m.timestamp <= p.end_time)
        AND m.key IN (SELECT keys from instruments)
        AND p.id = $1
        AND p.login = $2
        AND m.timestamp > i.start_time
        AND m.timestamp < i.start_time + $6::interval
        AND m.timestamp >= $3
        AND m.timestamp <= $4
        GROUP BY m.key, i.start_time
        ORDER BY i.start_time desc
    `

	trunc, interval, err := db.getIntervalDefinition(resolution)
	if err != nil {
		return measurements, err
	}

	err = db.db.Select(&measurements, sql, plotId, user, startTime, endTime, trunc, interval)
	return measurements, err
}

func (db *Database) readMeasurements(user string, name string) ([]Measurement, error) {
	measurements := []Measurement{}

	var sql = `
        SELECT key, value, timestamp
        FROM measurement
        WHERE name = $1
        AND login = $2
        ORDER BY timestamp
    `

	err := db.db.Select(&measurements, sql, name, user)
	return measurements, err
}

func (db *Database) saveMeasurements(measurements []Measurement, user string) error {
	tx, err := db.db.Beginx()
	if err != nil {
		return errors.Wrap(err, "Unable to save measurement")
	}

	var sql = `
        INSERT INTO measurement (key, value, timestamp, login)
        VALUES (:key, :value, :timestamp, :login)
    `
	for _, measurement := range measurements {
		measurement.Login = user
		tx.NamedExec(sql, &measurement)
	}

	err = tx.Commit()
	if err != nil {
		return errors.Wrap(err, "Unable to save measurement")
	}
	return nil
}

func (db *Database) getPlots(user string) ([]Plot, error) {
	plots := []Plot{}

	var sql = `
        SELECT id, start_time, end_time, name, case when end_time IS null then true else false end as active
        FROM plot
        WHERE login = $1
        ORDER BY start_time DESC
    `

	err := db.db.Select(&plots, sql, user)
	return plots, err
}

func (db *Database) getInstruments(plotId int) ([]Instrument, error) {
	instruments := []Instrument{}

	var sql = `
        SELECT key, id, name, type
        FROM instrument
        WHERE plot = $1
    `

	err := db.db.Select(&instruments, sql, plotId)
	return instruments, err
}

func (db *Database) getPlot(id int, user string) (Plot, error) {
	plot := Plot{}

	var sql = `
        SELECT id, start_time, end_time, name, case when end_time IS null then true else false end as active
        FROM plot
        WHERE login = $1
        AND id = $2
    `

	err := db.db.Get(&plot, sql, user, id)
	return plot, err
}

func (db *Database) savePlot(plot Plot, user string) (Plot, error) {

	plot.Login = user

	var sql = `
        INSERT INTO plot (start_time, end_time, name, login) VALUES (:start_time, :end_time, :name, :login) RETURNING id
    `
	var id int
	rows, err := db.db.NamedQuery(sql, plot)
	if err != nil {
		return plot, errors.Wrap(err, "Unable to save plot")
	}
	if rows.Next() {
		rows.Scan(&id)
	}
	plot.Id = id
	tx, err := db.db.Beginx()
	if err != nil {
		return plot, errors.Wrap(err, "Unable to save instruments for plot")
	}

	var sql2 = `
        INSERT INTO instrument (key, name, type, plot)
        VALUES (:key, :name, :type, :plot)
    `
	for _, instrument := range plot.Instruments {
		instrument.Plot = plot.Id
		tx.NamedExec(sql2, &instrument)
	}
	tx.Commit()
	return plot, err
}

func (db *Database) updatePlot(plot Plot, user string) (Plot, error) {

	plot.Login = user

	var sql = `
        UPDATE plot SET start_time = :start_time, end_time = :end_time, name = :name WHERE id = :id and login = :login
    `
	_, err := db.db.NamedQuery(sql, plot)
	if err != nil {
		return plot, err
	}
	return plot, err
}

func (db *Database) getUser(r *http.Request) (string, error) {
	key := r.Header.Get("X-PYTILT-KEY")
	return db.getUserForKey(key)
}

func (db *Database) getUserForKey(key string) (string, error) {
	var id string
	err := db.db.Get(&id, "SELECT id FROM login WHERE key = $1", key)

	if err == sql.ErrNoRows {
		return "", errors.New("unknown key")
	}
	return id, err
}

func (db *Database) getkeyForUser(user string) (string, error) {
	var key string
	err := db.db.Get(&key, "SELECT key FROM login WHERE id = $1", user)

	if err == sql.ErrNoRows {
		return "", errors.New("unknown user")
	}
	return key, err
}

func (db *Database) userExists(id string) (bool, error) {
	var uid string
	if err := db.db.QueryRow("SELECT id FROM login WHERE id = $1", id).Scan(&uid); err == nil {
		return true, nil
	} else if err == sql.ErrNoRows {
		return false, nil
	} else {
		return false, errors.Wrap(err, "Unable to check if user exists")
	}

}

func (db *Database) createUser(id string, email string, name string) error {
	tx, error := db.db.Beginx()
	if error != nil {
		return errors.New("Unable to connect to database.")
	}

	key := uuid.New()
	var sql = `
        INSERT INTO login (id, name, email, key)
        VALUES ($1, $2, $3, $4)
    `
	_, error = tx.Exec(sql, id, name, email, key)
	if error != nil {
		return errors.New("Unable to create new user")
	}

	tx.Commit()
	return nil
}
<commit_msg>round aggregated values from db (closes #14)<commit_after>package main

import (
	"database/sql"
	"github.com/jmoiron/sqlx"
	"github.com/pborman/uuid"
	"github.com/pkg/errors"
	"net/http"
	"strings"
	"time"
)

type Database struct {
	db *sqlx.DB
}

type Resolution int

const (
	All Resolution = iota
	Day
	Hour
	Minute
)

func (resolution Resolution) ToString() (string, error) {
	switch resolution {
	case All:
		return "All", nil
	case Minute:
		return "Minute", nil
	case Hour:
		return "Hour", nil
	case Day:
		return "Day", nil
	default:
		return "", errors.New("Unknown resolution")
	}
}

func (resolution Resolution) String() string {
	str, _ := resolution.ToString()
	return str
}

func ResolutionFromString(resolutionString string) (Resolution, error) {
	lowerCase := strings.ToLower(resolutionString)
	if lowerCase == strings.ToLower("All") {
		return All, nil
	} else if lowerCase == strings.ToLower("Minute") {
		return Minute, nil
	} else if lowerCase == strings.ToLower("Hour") {
		return Hour, nil
	} else if lowerCase == strings.ToLower("Day") {
		return Day, nil
	} else {
		return All, errors.New("Unknown resolution from string: " + resolutionString)
	}
}

func (db *Database) readDataFromPlot(user string, plotId int, startTime time.Time, endTime time.Time, resolution Resolution) ([]Measurement, error) {

	if resolution == All {
		return db.readAllDataFromPlot(user, plotId, startTime, endTime)
	} else {
		return db.readAggregatedDataFromPlot(user, plotId, startTime, endTime, resolution)
	}
}

func (db *Database) readAllDataFromPlot(user string, plotId int, startTime time.Time, endTime time.Time) ([]Measurement, error) {
	measurements := []Measurement{}
	var sql = `
        WITH instruments as (
            SELECT key FROM
            instrument
            WHERE plot = $1
        )
        SELECT m.key, m.value, m.timestamp
        FROM measurement m, plot p
        WHERE m.timestamp >= p.start_time
        AND(p.end_time is null OR m.timestamp <= p.end_time)
        AND m.key IN (SELECT key from instruments)
        AND p.id = $1
        AND p.login = $2
        AND m.timestamp >= $3
        AND m.timestamp <= $4
        ORDER BY m.timestamp desc;
    `
	err := db.db.Select(&measurements, sql, plotId, user, startTime, endTime)
	return measurements, err
}

func (db *Database) readLatestDataFromPlot(user string, plotId int) ([]Measurement, error) {
	measurements := []Measurement{}
	var sql = `
        WITH
        instruments as (
            SELECT i.key AS keys
            FROM instrument i
            WHERE plot = $1
        ),
        latest_measurement as (

            SELECT max(m.timestamp) as timestamp
            FROM measurement m, plot p
            WHERE m.timestamp >= p.start_time
            AND(p.end_time is null OR m.timestamp <= p.end_time)
            AND p.id = $1
        )
        SELECT m.key, m.value, m.timestamp
        FROM measurement m, plot p
        WHERE m.timestamp >= p.start_time
        and m.timestamp = (select timestamp from latest_measurement)
        AND m.key IN (SELECT i.keys from instruments i)
        AND p.id = $1
        AND p.login = $2
        ORDER BY m.timestamp desc;
    `
	err := db.db.Select(&measurements, sql, plotId, user)
	return measurements, err
}

func (db *Database) getIntervalDefinition(resolution Resolution) (string, string, error) {
	switch resolution {
	case All:
		return "", "", errors.New("No interval definition for All")
	case Minute:
		return "mins", "1 minute", nil
	case Hour:
		return "hours", "1 hour", nil
	case Day:
		return "days", "1 day", nil
	default:
		return "", "", errors.New("Unknown resolution")
	}
}

func (db *Database) readAggregatedDataFromPlot(user string, plotId int, startTime time.Time, endTime time.Time, resolution Resolution) ([]Measurement, error) {
	measurements := []Measurement{}
	var sql = `
        WITH instruments as (
            SELECT key AS keys
            FROM instrument
            WHERE plot = $1
        ),
        intervals AS (
            SELECT start_time FROM
            generate_series(date_trunc($5, GREATEST($3, (select start_time from plot where id = $1))),
                            LEAST($4, NOW()),
                            $6) as start_time
        )
        SELECT
            m.key,
            i.start_time as timestamp,
            round(AVG(m.value)::numeric, 2) as value
        FROM measurement m, plot p, intervals i
        WHERE m.timestamp >= p.start_time
        AND(p.end_time is null OR m.timestamp <= p.end_time)
        AND m.key IN (SELECT keys from instruments)
        AND p.id = $1
        AND p.login = $2
        AND m.timestamp > i.start_time
        AND m.timestamp < i.start_time + $6::interval
        AND m.timestamp >= $3
        AND m.timestamp <= $4
        GROUP BY m.key, i.start_time
        ORDER BY i.start_time desc
    `

	trunc, interval, err := db.getIntervalDefinition(resolution)
	if err != nil {
		return measurements, err
	}

	err = db.db.Select(&measurements, sql, plotId, user, startTime, endTime, trunc, interval)
	return measurements, err
}

func (db *Database) readMeasurements(user string, name string) ([]Measurement, error) {
	measurements := []Measurement{}

	var sql = `
        SELECT key, value, timestamp
        FROM measurement
        WHERE name = $1
        AND login = $2
        ORDER BY timestamp
    `

	err := db.db.Select(&measurements, sql, name, user)
	return measurements, err
}

func (db *Database) saveMeasurements(measurements []Measurement, user string) error {
	tx, err := db.db.Beginx()
	if err != nil {
		return errors.Wrap(err, "Unable to save measurement")
	}

	var sql = `
        INSERT INTO measurement (key, value, timestamp, login)
        VALUES (:key, :value, :timestamp, :login)
    `
	for _, measurement := range measurements {
		measurement.Login = user
		tx.NamedExec(sql, &measurement)
	}

	err = tx.Commit()
	if err != nil {
		return errors.Wrap(err, "Unable to save measurement")
	}
	return nil
}

func (db *Database) getPlots(user string) ([]Plot, error) {
	plots := []Plot{}

	var sql = `
        SELECT id, start_time, end_time, name, case when end_time IS null then true else false end as active
        FROM plot
        WHERE login = $1
        ORDER BY start_time DESC
    `

	err := db.db.Select(&plots, sql, user)
	return plots, err
}

func (db *Database) getInstruments(plotId int) ([]Instrument, error) {
	instruments := []Instrument{}

	var sql = `
        SELECT key, id, name, type
        FROM instrument
        WHERE plot = $1
    `

	err := db.db.Select(&instruments, sql, plotId)
	return instruments, err
}

func (db *Database) getPlot(id int, user string) (Plot, error) {
	plot := Plot{}

	var sql = `
        SELECT id, start_time, end_time, name, case when end_time IS null then true else false end as active
        FROM plot
        WHERE login = $1
        AND id = $2
    `

	err := db.db.Get(&plot, sql, user, id)
	return plot, err
}

func (db *Database) savePlot(plot Plot, user string) (Plot, error) {

	plot.Login = user

	var sql = `
        INSERT INTO plot (start_time, end_time, name, login) VALUES (:start_time, :end_time, :name, :login) RETURNING id
    `
	var id int
	rows, err := db.db.NamedQuery(sql, plot)
	if err != nil {
		return plot, errors.Wrap(err, "Unable to save plot")
	}
	if rows.Next() {
		rows.Scan(&id)
	}
	plot.Id = id
	tx, err := db.db.Beginx()
	if err != nil {
		return plot, errors.Wrap(err, "Unable to save instruments for plot")
	}

	var sql2 = `
        INSERT INTO instrument (key, name, type, plot)
        VALUES (:key, :name, :type, :plot)
    `
	for _, instrument := range plot.Instruments {
		instrument.Plot = plot.Id
		tx.NamedExec(sql2, &instrument)
	}
	tx.Commit()
	return plot, err
}

func (db *Database) updatePlot(plot Plot, user string) (Plot, error) {

	plot.Login = user

	var sql = `
        UPDATE plot SET start_time = :start_time, end_time = :end_time, name = :name WHERE id = :id and login = :login
    `
	_, err := db.db.NamedQuery(sql, plot)
	if err != nil {
		return plot, err
	}
	return plot, err
}

func (db *Database) getUser(r *http.Request) (string, error) {
	key := r.Header.Get("X-PYTILT-KEY")
	return db.getUserForKey(key)
}

func (db *Database) getUserForKey(key string) (string, error) {
	var id string
	err := db.db.Get(&id, "SELECT id FROM login WHERE key = $1", key)

	if err == sql.ErrNoRows {
		return "", errors.New("unknown key")
	}
	return id, err
}

func (db *Database) getkeyForUser(user string) (string, error) {
	var key string
	err := db.db.Get(&key, "SELECT key FROM login WHERE id = $1", user)

	if err == sql.ErrNoRows {
		return "", errors.New("unknown user")
	}
	return key, err
}

func (db *Database) userExists(id string) (bool, error) {
	var uid string
	if err := db.db.QueryRow("SELECT id FROM login WHERE id = $1", id).Scan(&uid); err == nil {
		return true, nil
	} else if err == sql.ErrNoRows {
		return false, nil
	} else {
		return false, errors.Wrap(err, "Unable to check if user exists")
	}

}

func (db *Database) createUser(id string, email string, name string) error {
	tx, error := db.db.Beginx()
	if error != nil {
		return errors.New("Unable to connect to database.")
	}

	key := uuid.New()
	var sql = `
        INSERT INTO login (id, name, email, key)
        VALUES ($1, $2, $3, $4)
    `
	_, error = tx.Exec(sql, id, name, email, key)
	if error != nil {
		return errors.New("Unable to create new user")
	}

	tx.Commit()
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"golang.org/x/net/websocket"
)

func setupConsumers(conf *Config) ([]<-chan *sarama.ConsumerMessage, []io.Closer, error) {
	partitionConsumers := []<-chan *sarama.ConsumerMessage{}
	closeables := []io.Closer{}
	for _, consumerConfig := range conf.Consumers {
		topic, brokerString, partition := consumerConfig.Topic, consumerConfig.Broker, consumerConfig.Partition
		var offset int64 = -1

		if numericOffset, err := strconv.ParseInt(consumerConfig.Offset, 10, 64); err == nil {
			offset = numericOffset
		} else {
			switch consumerConfig.Offset {
			case "oldest":
				offset = -2
			case "newest":
				offset = -1
			default:
				return nil, closeables, fmt.Errorf("Invalid value for consumer offset")
			}
		}

		if topic == "" {
			return nil, closeables, fmt.Errorf("Please define topic name for your consumer")
		}

		if brokerString == "" {
			brokerString = "localhost:9092"
		}

		if offset == 0 {
			offset = sarama.OffsetNewest
		}

		brokers := strings.Split(brokerString, ",")
		consumer, err := sarama.NewConsumer(brokers, nil)
		if err != nil {
			return nil, closeables, fmt.Errorf("Error creating consumer. err=%v", err)
		}

		var partitions []int32
		if partition == -1 {
			partitions, err = consumer.Partitions(topic)
			if err != nil {
				return nil, closeables, fmt.Errorf("Error fetching partitions for topic. err=%v", err)
			}
		} else {
			partitions = append(partitions, int32(partition))
		}

		for _, partition := range partitions {
			partitionConsumer, err := consumer.ConsumePartition(topic, int32(partition), offset)
			if err != nil {
				return nil, closeables, fmt.Errorf("Failed to consume partition %v err=%v\n", partition, err)
			}

			partitionConsumers = append(partitionConsumers, partitionConsumer.Messages())
			closeables = append(closeables, partitionConsumer)
		}
		closeables = append(closeables, consumer)
	}
	return partitionConsumers, closeables, nil
}

func demuxMessages(pc []<-chan *sarama.ConsumerMessage, q chan struct{}) chan *sarama.ConsumerMessage {
	c := make(chan *sarama.ConsumerMessage)
	for _, p := range pc {
		go func(p <-chan *sarama.ConsumerMessage) {
			for {
				select {
				case msg := <-p:
					c <- msg
				case <-q:
					return
				}
			}
		}(p)
	}
	return c
}

func sendMessagesToWsBlocking(ws *websocket.Conn, c chan *sarama.ConsumerMessage, q chan struct{}) {
	for {
		select {
		case cMsg := <-c:
			msg :=
				`{"topic": "` + cMsg.Topic +
					`", "partition": "` + strconv.FormatInt(int64(cMsg.Partition), 10) +
					`", "offset": "` + strconv.FormatInt(cMsg.Offset, 10) +
					`", "key": "` + strings.Replace(string(cMsg.Key), `"`, `\"`, -1) +
					`", "value": "` + strings.Replace(string(cMsg.Value), `"`, `\"`, -1) +
					`", "consumedUnixTimestamp": "` + strconv.FormatInt(time.Now().Unix(), 10) +
					`"}` + "\n"

			log.Println("Sending message to WebSocket: " + msg)
			err := websocket.Message.Send(ws, msg)
			if err != nil {
				log.Printf("Error while trying to send to WebSocket: err=%v\n", err)
				return
			}
		case <-q:
			log.Println("Received quit signal")
			return
		}
	}
}
<commit_msg>Implements start from (newest - n) offset. Abstracts out resolving.<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"golang.org/x/net/websocket"
)

func setupConsumers(conf *Config) ([]<-chan *sarama.ConsumerMessage, []io.Closer, error) {
	partitionConsumers := []<-chan *sarama.ConsumerMessage{}
	closeables := []io.Closer{}
	for _, consumerConfig := range conf.Consumers {
		topic, brokerString, partition := consumerConfig.Topic, consumerConfig.Broker, consumerConfig.Partition

		if topic == "" {
			return nil, closeables, fmt.Errorf("Please define topic name for your consumer")
		}

		if brokerString == "" {
			brokerString = "localhost:9092"
		}

		brokers := strings.Split(brokerString, ",")

		consumer, err := sarama.NewConsumer(brokers, nil)
		if err != nil {
			return nil, closeables, fmt.Errorf("Error creating consumer. err=%v", err)
		}

		var partitions []int32
		if partition == -1 {
			partitions, err = consumer.Partitions(topic)
			if err != nil {
				return nil, closeables, fmt.Errorf("Error fetching partitions for topic. err=%v", err)
			}
		} else {
			partitions = append(partitions, int32(partition))
		}

		for _, partition := range partitions {
			offset, err := resolveOffset(consumerConfig.Offset, brokers, topic, partition)
			if err != nil {
				return nil, closeables, fmt.Errorf("Could not resolve offset for %v, %v, %v. err=%v", brokers, topic, partition, err)
			}

			partitionConsumer, err := consumer.ConsumePartition(topic, int32(partition), offset)
			if err != nil {
				return nil, closeables, fmt.Errorf("Failed to consume partition %v err=%v\n", partition, err)
			}

			partitionConsumers = append(partitionConsumers, partitionConsumer.Messages())
			closeables = append(closeables, partitionConsumer)
		}
		closeables = append(closeables, consumer)
	}
	return partitionConsumers, closeables, nil
}

func resolveOffset(configOffset string, brokers []string, topic string, partition int32) (int64, error) {
	if configOffset == "oldest" {
		return sarama.OffsetOldest, nil
	} else if configOffset == "newest" {
		return sarama.OffsetNewest, nil
	} else if numericOffset, err := strconv.ParseInt(configOffset, 10, 64); err == nil {
		if numericOffset >= -2 {
			return numericOffset, nil
		}

		client, err := sarama.NewClient(brokers, nil)
		if err != nil {
			return 0, fmt.Errorf("Failed to create client for %v, %v, %v", brokers, topic, partition)
		}
		defer client.Close()

		oldest, err := client.GetOffset(topic, partition, sarama.OffsetOldest)
		if err != nil {
			return 0, err
		}

		newest, err := client.GetOffset(topic, partition, sarama.OffsetNewest)
		if err != nil {
			return 0, err
		}

		if newest+numericOffset < oldest {
			return oldest, nil
		}

		return newest + numericOffset, nil
	}

	return 0, fmt.Errorf("Invalid value for consumer offset")
}

func demuxMessages(pc []<-chan *sarama.ConsumerMessage, q chan struct{}) chan *sarama.ConsumerMessage {
	c := make(chan *sarama.ConsumerMessage)
	for _, p := range pc {
		go func(p <-chan *sarama.ConsumerMessage) {
			for {
				select {
				case msg := <-p:
					c <- msg
				case <-q:
					return
				}
			}
		}(p)
	}
	return c
}

func sendMessagesToWsBlocking(ws *websocket.Conn, c chan *sarama.ConsumerMessage, q chan struct{}) {
	for {
		select {
		case cMsg := <-c:
			msg :=
				`{"topic": "` + cMsg.Topic +
					`", "partition": "` + strconv.FormatInt(int64(cMsg.Partition), 10) +
					`", "offset": "` + strconv.FormatInt(cMsg.Offset, 10) +
					`", "key": "` + strings.Replace(string(cMsg.Key), `"`, `\"`, -1) +
					`", "value": "` + strings.Replace(string(cMsg.Value), `"`, `\"`, -1) +
					`", "consumedUnixTimestamp": "` + strconv.FormatInt(time.Now().Unix(), 10) +
					`"}` + "\n"

			log.Println("Sending message to WebSocket: " + msg)
			err := websocket.Message.Send(ws, msg)
			if err != nil {
				log.Printf("Error while trying to send to WebSocket: err=%v\n", err)
				return
			}
		case <-q:
			log.Println("Received quit signal")
			return
		}
	}
}
<|endoftext|>