<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/yuin/gopher-lua"
)

func exportBasicSystemFunctions(L *lua.LState) {

	// Return the version string
	L.SetGlobal("version", L.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(versionString))
		return 1 // number of results
	}))

	// Log text with the "Info" log type
	L.SetGlobal("log", L.NewFunction(func(L *lua.LState) int {
		buf := arguments2buffer(L, false)
		// Log the combined text
		log.Info(buf.String())
		return 0 // number of results
	}))

	// Log text with the "Warn" log type
	L.SetGlobal("warn", L.NewFunction(func(L *lua.LState) int {
		buf := arguments2buffer(L, false)
		// Log the combined text
		log.Warn(buf.String())
		return 0 // number of results
	}))

	// Log text with the "Error" log type
	L.SetGlobal("error", L.NewFunction(func(L *lua.LState) int {
		buf := arguments2buffer(L, false)
		// Log the combined text
		log.Error(buf.String())
		return 0 // number of results
	}))

	// Sleep for the given number of seconds (can be a float)
	L.SetGlobal("sleep", L.NewFunction(func(L *lua.LState) int {
		// Extract the correct number of nanoseconds
		duration := time.Duration(float64(L.ToNumber(1)) * 1000000000.0)
		// Wait and block the current thread of execution.
		time.Sleep(duration)
		return 0
	}))

}

// Make functions related to HTTP requests and responses available to Lua scripts.
// Filename can be an empty string.
// realW is only used for flushing the buffer in debug mode, and is the underlying ResponseWriter.
func exportBasicWeb(w http.ResponseWriter, req *http.Request, L *lua.LState, filename string, flushFunc func()) {

	// Print text to the webpage that is being served. Add a newline.
	L.SetGlobal("print", L.NewFunction(func(L *lua.LState) int {
		var buf bytes.Buffer
		top := L.GetTop()
		for i := 1; i <= top; i++ {
			buf.WriteString(L.Get(i).String())
			if i != top {
				buf.WriteString("\t")
			}
		}
		// Final newline
		buf.WriteString("\n")

		// Write the combined text to the http.ResponseWriter
		w.Write(buf.Bytes())

		return 0 // number of results
	}))

	// Flush the ResponseWriter.
	// Needed in debug mode, where ResponseWriter is buffered.
	L.SetGlobal("flush", L.NewFunction(func(L *lua.LState) int {
		if flushFunc != nil {
			flushFunc()
		}
		return 0 // number of results
	}))

	// Set the Content-Type for the page
	L.SetGlobal("content", L.NewFunction(func(L *lua.LState) int {
		lv := L.ToString(1)
		w.Header().Add("Content-Type", lv)
		return 0 // number of results
	}))

	// Return the current URL Path
	L.SetGlobal("urlpath", L.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(req.URL.Path))
		return 1 // number of results
	}))

	// Return the current HTTP method (GET, POST etc)
	L.SetGlobal("method", L.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(req.Method))
		return 1 // number of results
	}))

	// Return the HTTP header in the request, for a given key/string
	L.SetGlobal("header", L.NewFunction(func(L *lua.LState) int {
		key := L.ToString(1)
		value := req.Header.Get(key)
		L.Push(lua.LString(value))
		return 1 // number of results
	}))

	// Set the HTTP header in the request, for a given key and value
	L.SetGlobal("setheader", L.NewFunction(func(L *lua.LState) int {
		key := L.ToString(1)
		value := L.ToString(2)
		w.Header().Set(key, value)
		return 0 // number of results
	}))

	// Return the HTTP body in the request
	L.SetGlobal("body", L.NewFunction(func(L *lua.LState) int {
		body, err := ioutil.ReadAll(req.Body)
		var result lua.LString
		if err != nil {
			result = lua.LString("")
		} else {
			result = lua.LString(string(body))
		}
		L.Push(result)
		return 1 // number of results
	}))

	// Set the HTTP status code (must come before print)
	L.SetGlobal("status", L.NewFunction(func(L *lua.LState) int {
		code := int(L.ToNumber(1))
		w.WriteHeader(code)
		return 0 // number of results
	}))

	// Print a message and set the HTTP status code
	L.SetGlobal("error", L.NewFunction(func(L *lua.LState) int {
		message := L.ToString(1)
		code := int(L.ToNumber(2))
		w.WriteHeader(code)
		fmt.Fprint(w, message)
		return 0 // number of results
	}))

	// Get the full filename of a given file that is in the directory
	// of the script that is about to be run.
	// If no filename is given, the directory where the script lies
	// is returned.
	L.SetGlobal("scriptdir", L.NewFunction(func(L *lua.LState) int {
		scriptdir := path.Dir(filename)
		top := L.GetTop()
		if top == 1 {
			// Also include a separator and a filename
			fn := L.ToString(1)
			scriptdir += pathsep + fn
		}
		if pathsep != "/" {
			// For operating systems that use another path separator
			scriptdir = strings.Replace(scriptdir, pathsep, "/", everyInstance)
		}
		L.Push(lua.LString(scriptdir))
		return 1 // number of results
	}))

	// Get the full filename of a given file that is in the directory
	// where the server is running (root directory for the server).
	// If no filename is given, the directory where the server is
	// currently running is returned.
	L.SetGlobal("serverdir", L.NewFunction(func(L *lua.LState) int {
		serverdir, err := os.Getwd()
		if err != nil {
			// Could not retrieve a directory
			serverdir = ""
		} else if L.GetTop() == 1 {
			// Also include a separator and a filename
			fn := L.ToString(1)
			serverdir += pathsep + fn
		}
		if pathsep != "/" {
			// For operating systems that use another path separator
			serverdir = strings.Replace(serverdir, pathsep, "/", everyInstance)
		}
		L.Push(lua.LString(serverdir))
		return 1 // number of results
	}))

	// Retrieve a table with keys and values from the form in the request
	L.SetGlobal("formdata", L.NewFunction(func(L *lua.LState) int {
		// Place the form data in a map
		m := make(map[string]string)
		req.ParseForm()
		for k, v := range req.Form {
			m[k] = v[0]
		}
		// Convert the map to a table and return it
		L.Push(map2table(L, m))
		return 1 // number of results
	}))

}
<commit_msg>Using path/filepath instead of path<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/yuin/gopher-lua"
)

func exportBasicSystemFunctions(L *lua.LState) {

	// Return the version string
	L.SetGlobal("version", L.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(versionString))
		return 1 // number of results
	}))

	// Log text with the "Info" log type
	L.SetGlobal("log", L.NewFunction(func(L *lua.LState) int {
		buf := arguments2buffer(L, false)
		// Log the combined text
		log.Info(buf.String())
		return 0 // number of results
	}))

	// Log text with the "Warn" log type
	L.SetGlobal("warn", L.NewFunction(func(L *lua.LState) int {
		buf := arguments2buffer(L, false)
		// Log the combined text
		log.Warn(buf.String())
		return 0 // number of results
	}))

	// Log text with the "Error" log type
	L.SetGlobal("error", L.NewFunction(func(L *lua.LState) int {
		buf := arguments2buffer(L, false)
		// Log the combined text
		log.Error(buf.String())
		return 0 // number of results
	}))

	// Sleep for the given number of seconds (can be a float)
	L.SetGlobal("sleep", L.NewFunction(func(L *lua.LState) int {
		// Extract the correct number of nanoseconds
		duration := time.Duration(float64(L.ToNumber(1)) * 1000000000.0)
		// Wait and block the current thread of execution.
		time.Sleep(duration)
		return 0
	}))

}

// Make functions related to HTTP requests and responses available to Lua scripts.
// Filename can be an empty string.
// realW is only used for flushing the buffer in debug mode, and is the underlying ResponseWriter.
func exportBasicWeb(w http.ResponseWriter, req *http.Request, L *lua.LState, filename string, flushFunc func()) {

	// Print text to the webpage that is being served. Add a newline.
	L.SetGlobal("print", L.NewFunction(func(L *lua.LState) int {
		var buf bytes.Buffer
		top := L.GetTop()
		for i := 1; i <= top; i++ {
			buf.WriteString(L.Get(i).String())
			if i != top {
				buf.WriteString("\t")
			}
		}
		// Final newline
		buf.WriteString("\n")

		// Write the combined text to the http.ResponseWriter
		w.Write(buf.Bytes())

		return 0 // number of results
	}))

	// Flush the ResponseWriter.
	// Needed in debug mode, where ResponseWriter is buffered.
	L.SetGlobal("flush", L.NewFunction(func(L *lua.LState) int {
		if flushFunc != nil {
			flushFunc()
		}
		return 0 // number of results
	}))

	// Set the Content-Type for the page
	L.SetGlobal("content", L.NewFunction(func(L *lua.LState) int {
		lv := L.ToString(1)
		w.Header().Add("Content-Type", lv)
		return 0 // number of results
	}))

	// Return the current URL Path
	L.SetGlobal("urlpath", L.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(req.URL.Path))
		return 1 // number of results
	}))

	// Return the current HTTP method (GET, POST etc)
	L.SetGlobal("method", L.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(req.Method))
		return 1 // number of results
	}))

	// Return the HTTP header in the request, for a given key/string
	L.SetGlobal("header", L.NewFunction(func(L *lua.LState) int {
		key := L.ToString(1)
		value := req.Header.Get(key)
		L.Push(lua.LString(value))
		return 1 // number of results
	}))

	// Set the HTTP header in the request, for a given key and value
	L.SetGlobal("setheader", L.NewFunction(func(L *lua.LState) int {
		key := L.ToString(1)
		value := L.ToString(2)
		w.Header().Set(key, value)
		return 0 // number of results
	}))

	// Return the HTTP body in the request
	L.SetGlobal("body", L.NewFunction(func(L *lua.LState) int {
		body, err := ioutil.ReadAll(req.Body)
		var result lua.LString
		if err != nil {
			result = lua.LString("")
		} else {
			result = lua.LString(string(body))
		}
		L.Push(result)
		return 1 // number of results
	}))

	// Set the HTTP status code (must come before print)
	L.SetGlobal("status", L.NewFunction(func(L *lua.LState) int {
		code := int(L.ToNumber(1))
		w.WriteHeader(code)
		return 0 // number of results
	}))

	// Print a message and set the HTTP status code
	L.SetGlobal("error", L.NewFunction(func(L *lua.LState) int {
		message := L.ToString(1)
		code := int(L.ToNumber(2))
		w.WriteHeader(code)
		fmt.Fprint(w, message)
		return 0 // number of results
	}))

	// Get the full filename of a given file that is in the directory
	// of the script that is about to be run.
	// If no filename is given, the directory where the script lies
	// is returned.
	L.SetGlobal("scriptdir", L.NewFunction(func(L *lua.LState) int {
		scriptdir := filepath.Dir(filename)
		top := L.GetTop()
		if top == 1 {
			// Also include a separator and a filename
			fn := L.ToString(1)
			scriptdir += pathsep + fn
		}
		if pathsep != "/" {
			// For operating systems that use another path separator
			scriptdir = strings.Replace(scriptdir, pathsep, "/", everyInstance)
		}
		L.Push(lua.LString(scriptdir))
		return 1 // number of results
	}))

	// Get the full filename of a given file that is in the directory
	// where the server is running (root directory for the server).
	// If no filename is given, the directory where the server is
	// currently running is returned.
	L.SetGlobal("serverdir", L.NewFunction(func(L *lua.LState) int {
		serverdir, err := os.Getwd()
		if err != nil {
			// Could not retrieve a directory
			serverdir = ""
		} else if L.GetTop() == 1 {
			// Also include a separator and a filename
			fn := L.ToString(1)
			serverdir += pathsep + fn
		}
		if pathsep != "/" {
			// For operating systems that use another path separator
			serverdir = strings.Replace(serverdir, pathsep, "/", everyInstance)
		}
		L.Push(lua.LString(serverdir))
		return 1 // number of results
	}))

	// Retrieve a table with keys and values from the form in the request
	L.SetGlobal("formdata", L.NewFunction(func(L *lua.LState) int {
		// Place the form data in a map
		m := make(map[string]string)
		req.ParseForm()
		for k, v := range req.Form {
			m[k] = v[0]
		}
		// Convert the map to a table and return it
		L.Push(map2table(L, m))
		return 1 // number of results
	}))

}
<|endoftext|>


<commit_before>package kapacitor

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/gorhill/cronexpr"
	"github.com/influxdb/influxdb/client"
	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/kapacitor/models"
	"github.com/influxdb/kapacitor/pipeline"
)

type SourceBatchNode struct {
	node
	s   *pipeline.SourceBatchNode
	idx int
}

func newSourceBatchNode(et *ExecutingTask, n *pipeline.SourceBatchNode) (*SourceBatchNode, error) {
	sn := &SourceBatchNode{
		node: node{Node: n, et: et},
		s:    n,
	}
	return sn, nil
}

func (s *SourceBatchNode) linkChild(c Node) error {

	// add child
	if s.Provides() != c.Wants() {
		return fmt.Errorf("cannot add child mismatched edges: %s -> %s", s.Provides(), c.Wants())
	}
	s.children = append(s.children, c)

	// add parent
	c.addParent(s)

	return nil
}

func (s *SourceBatchNode) addParentEdge(in *Edge) {
	// Pass edges down to children
	s.children[s.idx].addParentEdge(in)
	s.idx++
}

func (s *SourceBatchNode) start() {
}

func (s *SourceBatchNode) Err() error {
	return nil
}

// Return list of databases and retention policies
// the batcher will query.
func (s *SourceBatchNode) DBRPs() ([]DBRP, error) {
	var dbrps []DBRP
	for _, b := range s.children {
		d, err := b.(*BatchNode).DBRPs()
		if err != nil {
			return nil, err
		}
		dbrps = append(dbrps, d...)
	}
	return dbrps, nil
}

func (s *SourceBatchNode) Count() int {
	return len(s.children)
}

func (s *SourceBatchNode) Start(collectors []*Edge) {
	for i, b := range s.children {
		b.(*BatchNode).Start(collectors[i])
	}
}

func (s *SourceBatchNode) Queries(start, stop time.Time) [][]string {
	queries := make([][]string, len(s.children))
	for i, b := range s.children {
		queries[i] = b.(*BatchNode).Queries(start, stop)
	}
	return queries
}

type BatchNode struct {
	node
	b       *pipeline.BatchNode
	query   *Query
	ticker  ticker
	wg      sync.WaitGroup
	errC    chan error
	closing chan struct{}
}

func newBatchNode(et *ExecutingTask, n *pipeline.BatchNode) (*BatchNode, error) {
	bn := &BatchNode{
		node:    node{Node: n, et: et},
		b:       n,
		errC:    make(chan error, 1),
		closing: make(chan struct{}),
	}
	bn.node.runF = bn.runBatch
	bn.node.stopF = bn.stopBatch

	// Create query
	q, err := NewQuery(n.QueryStr)
	if err != nil {
		return nil, err
	}
	bn.query = q
	// Add in dimensions
	err = bn.query.Dimensions(n.Dimensions)
	if err != nil {
		return nil, err
	}
	// Set fill
	switch fill := n.FillOption.(type) {
	case string:
		switch fill {
		case "null":
			bn.query.Fill(influxql.NullFill, nil)
		case "none":
			bn.query.Fill(influxql.NoFill, nil)
		case "previous":
			bn.query.Fill(influxql.PreviousFill, nil)
		default:
			return nil, fmt.Errorf("unexpected fill option %s", fill)
		}
	case int64, float64:
		bn.query.Fill(influxql.NumberFill, fill)
	}

	// Determine schedule
	if n.Every != 0 && n.Cron != "" {
		return nil, errors.New("must not set both 'every' and 'cron' properties")
	}
	switch {
	case n.Every != 0:
		bn.ticker = newTimeTicker(n.Every)
	case n.Cron != "":
		var err error
		bn.ticker, err = newCronTicker(n.Cron)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("must define one of 'every' or 'cron'")
	}

	if n.Every != 0 && n.Cron != "" {
		return nil, errors.New("must not set both 'every' and 'cron' properties")
	}
	switch {
	case n.Every != 0:
		bn.ticker = newTimeTicker(n.Every)
	case n.Cron != "":
		var err error
		bn.ticker, err = newCronTicker(n.Cron)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("must define one of 'every' or 'cron'")
	}

	return bn, nil
}

// Return list of databases and retention policies
// the batcher will query.
func (b *BatchNode) DBRPs() ([]DBRP, error) {
	return b.query.DBRPs()
}

func (b *BatchNode) Start(batch BatchCollector) {
	b.wg.Add(1)
	go func() {
		b.errC <- b.doQuery(batch)
	}()
}

func (b *BatchNode) Queries(start, stop time.Time) []string {
	now := time.Now()
	if stop.IsZero() {
		stop = now
	}
	// Crons are sensitive to timezones.
	// Make sure we are using local time.
	start = start.Local()
	queries := make([]string, 0)
	for {
		start = b.ticker.Next(start)
		if start.IsZero() || start.After(stop) {
			break
		}
		b.query.Start(start)
		qstop := start.Add(b.b.Period)
		if qstop.After(now) {
			break
		}
		b.query.Stop(qstop)
		queries = append(queries, b.query.String())
	}
	return queries
}

// Query InfluxDB and collect batches on batch collector.
func (b *BatchNode) doQuery(batch BatchCollector) error {
	defer batch.Close()
	defer b.wg.Done()

	if b.et.tm.InfluxDBService == nil {
		return errors.New("InfluxDB not configured, cannot query InfluxDB for batch query")
	}

	tickC := b.ticker.Start()
	for {
		select {
		case <-b.closing:
			return nil
		case now := <-tickC:

			// Update times for query
			stop := now.Add(-1 * b.b.Offset)
			b.query.Start(stop.Add(-1 * b.b.Period))
			b.query.Stop(stop)

			b.logger.Println("D! starting next batch query:", b.query.String())

			// Connect
			con, err := b.et.tm.InfluxDBService.NewClient()
			if err != nil {
				return err
			}
			q := client.Query{
				Command: b.query.String(),
			}

			// Execute query
			resp, err := con.Query(q)
			if err != nil {
				return err
			}

			if resp.Err != nil {
				return resp.Err
			}

			// Collect batches
			for _, res := range resp.Results {
				batches, err := models.ResultToBatches(res)
				if err != nil {
					return err
				}
				for _, bch := range batches {
					batch.CollectBatch(bch)
				}
			}
		}
		return errors.New("batch ticker schedule stopped")
	}
}

func (b *BatchNode) stopBatch() {
	if b.ticker != nil {
		b.ticker.Stop()
	}
	close(b.closing)
	b.wg.Wait()
}

func (b *BatchNode) runBatch() error {
	errC := make(chan error, 1)
	go func() {
		for bt, ok := b.ins[0].NextBatch(); ok; bt, ok = b.ins[0].NextBatch() {
			for _, child := range b.outs {
				err := child.CollectBatch(bt)
				if err != nil {
					errC <- err
					return
				}
			}
		}
		errC <- nil
	}()
	// Wait for errors
	select {
	case err := <-b.errC:
		return err
	case err := <-errC:
		return err
	}
}

type ticker interface {
	Start() <-chan time.Time
	Stop()
	// Return the next time the ticker will tick after now.
	Next(now time.Time) time.Time
}

type timeTicker struct {
	every  time.Duration
	ticker *time.Ticker
}

func newTimeTicker(every time.Duration) *timeTicker {
	return &timeTicker{every: every}
}

func (t *timeTicker) Start() <-chan time.Time {
	t.ticker = time.NewTicker(t.every)
	return t.ticker.C
}

func (t *timeTicker) Stop() {
	if t.ticker != nil {
		t.ticker.Stop()
	}
}

func (t *timeTicker) Next(now time.Time) time.Time {
	return now.Add(t.every)
}

type cronTicker struct {
	expr    *cronexpr.Expression
	ticker  chan time.Time
	closing chan struct{}
	wg      sync.WaitGroup
}

func newCronTicker(cronExpr string) (*cronTicker, error) {
	expr, err := cronexpr.Parse(cronExpr)
	if err != nil {
		return nil, err
	}
	return &cronTicker{
		expr:    expr,
		ticker:  make(chan time.Time),
		closing: make(chan struct{}),
	}, nil
}

func (c *cronTicker) Start() <-chan time.Time {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		for {
			now := time.Now()
			next := c.expr.Next(now)
			diff := next.Sub(now)
			select {
			case <-time.After(diff):
				c.ticker <- next
			case <-c.closing:
				return
			}
		}
	}()
	return c.ticker
}

func (c *cronTicker) Stop() {
	close(c.closing)
	c.wg.Wait()
}

func (c *cronTicker) Next(now time.Time) time.Time {
	return c.expr.Next(now)
}
<commit_msg>fix early termination on batch timers<commit_after>package kapacitor

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/gorhill/cronexpr"
	"github.com/influxdb/influxdb/client"
	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/kapacitor/models"
	"github.com/influxdb/kapacitor/pipeline"
)

type SourceBatchNode struct {
	node
	s   *pipeline.SourceBatchNode
	idx int
}

func newSourceBatchNode(et *ExecutingTask, n *pipeline.SourceBatchNode) (*SourceBatchNode, error) {
	sn := &SourceBatchNode{
		node: node{Node: n, et: et},
		s:    n,
	}
	return sn, nil
}

func (s *SourceBatchNode) linkChild(c Node) error {

	// add child
	if s.Provides() != c.Wants() {
		return fmt.Errorf("cannot add child mismatched edges: %s -> %s", s.Provides(), c.Wants())
	}
	s.children = append(s.children, c)

	// add parent
	c.addParent(s)

	return nil
}

func (s *SourceBatchNode) addParentEdge(in *Edge) {
	// Pass edges down to children
	s.children[s.idx].addParentEdge(in)
	s.idx++
}

func (s *SourceBatchNode) start() {
}

func (s *SourceBatchNode) Err() error {
	return nil
}

// Return list of databases and retention policies
// the batcher will query.
func (s *SourceBatchNode) DBRPs() ([]DBRP, error) {
	var dbrps []DBRP
	for _, b := range s.children {
		d, err := b.(*BatchNode).DBRPs()
		if err != nil {
			return nil, err
		}
		dbrps = append(dbrps, d...)
	}
	return dbrps, nil
}

func (s *SourceBatchNode) Count() int {
	return len(s.children)
}

func (s *SourceBatchNode) Start(collectors []*Edge) {
	for i, b := range s.children {
		b.(*BatchNode).Start(collectors[i])
	}
}

func (s *SourceBatchNode) Queries(start, stop time.Time) [][]string {
	queries := make([][]string, len(s.children))
	for i, b := range s.children {
		queries[i] = b.(*BatchNode).Queries(start, stop)
	}
	return queries
}

type BatchNode struct {
	node
	b       *pipeline.BatchNode
	query   *Query
	ticker  ticker
	wg      sync.WaitGroup
	errC    chan error
	closing chan struct{}
}

func newBatchNode(et *ExecutingTask, n *pipeline.BatchNode) (*BatchNode, error) {
	bn := &BatchNode{
		node:    node{Node: n, et: et},
		b:       n,
		errC:    make(chan error, 1),
		closing: make(chan struct{}),
	}
	bn.node.runF = bn.runBatch
	bn.node.stopF = bn.stopBatch

	// Create query
	q, err := NewQuery(n.QueryStr)
	if err != nil {
		return nil, err
	}
	bn.query = q
	// Add in dimensions
	err = bn.query.Dimensions(n.Dimensions)
	if err != nil {
		return nil, err
	}
	// Set fill
	switch fill := n.FillOption.(type) {
	case string:
		switch fill {
		case "null":
			bn.query.Fill(influxql.NullFill, nil)
		case "none":
			bn.query.Fill(influxql.NoFill, nil)
		case "previous":
			bn.query.Fill(influxql.PreviousFill, nil)
		default:
			return nil, fmt.Errorf("unexpected fill option %s", fill)
		}
	case int64, float64:
		bn.query.Fill(influxql.NumberFill, fill)
	}

	// Determine schedule
	if n.Every != 0 && n.Cron != "" {
		return nil, errors.New("must not set both 'every' and 'cron' properties")
	}
	switch {
	case n.Every != 0:
		bn.ticker = newTimeTicker(n.Every)
	case n.Cron != "":
		var err error
		bn.ticker, err = newCronTicker(n.Cron)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("must define one of 'every' or 'cron'")
	}

	if n.Every != 0 && n.Cron != "" {
		return nil, errors.New("must not set both 'every' and 'cron' properties")
	}
	switch {
	case n.Every != 0:
		bn.ticker = newTimeTicker(n.Every)
	case n.Cron != "":
		var err error
		bn.ticker, err = newCronTicker(n.Cron)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("must define one of 'every' or 'cron'")
	}

	return bn, nil
}

// Return list of databases and retention policies
// the batcher will query.
func (b *BatchNode) DBRPs() ([]DBRP, error) {
	return b.query.DBRPs()
}

func (b *BatchNode) Start(batch BatchCollector) {
	b.wg.Add(1)
	go func() {
		b.errC <- b.doQuery(batch)
	}()
}

func (b *BatchNode) Queries(start, stop time.Time) []string {
	now := time.Now()
	if stop.IsZero() {
		stop = now
	}
	// Crons are sensitive to timezones.
	// Make sure we are using local time.
	start = start.Local()
	queries := make([]string, 0)
	for {
		start = b.ticker.Next(start)
		if start.IsZero() || start.After(stop) {
			break
		}
		b.query.Start(start)
		qstop := start.Add(b.b.Period)
		if qstop.After(now) {
			break
		}
		b.query.Stop(qstop)
		queries = append(queries, b.query.String())
	}
	return queries
}

// Query InfluxDB and collect batches on batch collector.
func (b *BatchNode) doQuery(batch BatchCollector) error {
	defer batch.Close()
	defer b.wg.Done()

	if b.et.tm.InfluxDBService == nil {
		return errors.New("InfluxDB not configured, cannot query InfluxDB for batch query")
	}

	tickC := b.ticker.Start()
	for {
		select {
		case <-b.closing:
			return nil
		case now := <-tickC:

			// Update times for query
			stop := now.Add(-1 * b.b.Offset)
			b.query.Start(stop.Add(-1 * b.b.Period))
			b.query.Stop(stop)

			b.logger.Println("D! starting next batch query:", b.query.String())

			// Connect
			con, err := b.et.tm.InfluxDBService.NewClient()
			if err != nil {
				return err
			}
			q := client.Query{
				Command: b.query.String(),
			}

			// Execute query
			resp, err := con.Query(q)
			if err != nil {
				return err
			}

			if resp.Err != nil {
				return resp.Err
			}

			// Collect batches
			for _, res := range resp.Results {
				batches, err := models.ResultToBatches(res)
				if err != nil {
					return err
				}
				for _, bch := range batches {
					batch.CollectBatch(bch)
				}
			}
		}
	}
}

func (b *BatchNode) stopBatch() {
	if b.ticker != nil {
		b.ticker.Stop()
	}
	close(b.closing)
	b.wg.Wait()
}

func (b *BatchNode) runBatch() error {
	errC := make(chan error, 1)
	go func() {
		for bt, ok := b.ins[0].NextBatch(); ok; bt, ok = b.ins[0].NextBatch() {
			for _, child := range b.outs {
				err := child.CollectBatch(bt)
				if err != nil {
					errC <- err
					return
				}
			}
		}
		errC <- nil
	}()
	// Wait for errors
	select {
	case err := <-b.errC:
		return err
	case err := <-errC:
		return err
	}
}

type ticker interface {
	Start() <-chan time.Time
	Stop()
	// Return the next time the ticker will tick after now.
	Next(now time.Time) time.Time
}

type timeTicker struct {
	every  time.Duration
	ticker *time.Ticker
}

func newTimeTicker(every time.Duration) *timeTicker {
	return &timeTicker{every: every}
}

func (t *timeTicker) Start() <-chan time.Time {
	t.ticker = time.NewTicker(t.every)
	return t.ticker.C
}

func (t *timeTicker) Stop() {
	if t.ticker != nil {
		t.ticker.Stop()
	}
}

func (t *timeTicker) Next(now time.Time) time.Time {
	return now.Add(t.every)
}

type cronTicker struct {
	expr    *cronexpr.Expression
	ticker  chan time.Time
	closing chan struct{}
	wg      sync.WaitGroup
}

func newCronTicker(cronExpr string) (*cronTicker, error) {
	expr, err := cronexpr.Parse(cronExpr)
	if err != nil {
		return nil, err
	}
	return &cronTicker{
		expr:    expr,
		ticker:  make(chan time.Time),
		closing: make(chan struct{}),
	}, nil
}

func (c *cronTicker) Start() <-chan time.Time {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		for {
			now := time.Now()
			next := c.expr.Next(now)
			diff := next.Sub(now)
			select {
			case <-time.After(diff):
				c.ticker <- next
			case <-c.closing:
				return
			}
		}
	}()
	return c.ticker
}

func (c *cronTicker) Stop() {
	close(c.closing)
	c.wg.Wait()
}

func (c *cronTicker) Next(now time.Time) time.Time {
	return c.expr.Next(now)
}
<|endoftext|>


<commit_before>package async

import (
	"context"
	"errors"
	"sync"
	"time"
)

type Method func(ctx context.Context, args ...interface{}) (interface{}, error)

func Lambda(method func() (interface{}, error), timeout time.Duration) (interface{}, error) {
	output := make(chan interface{})
	go func() {
		defer close(output)
		res, err := method()
		if err != nil {
			output <- err
		} else {
			output <- res
		}
	}()
	if timeout <= 0 {
		timer := time.NewTimer(timeout)
		defer timer.Stop()
		select {
		case res := <-output:
			{
				switch err := res.(type) {
				case error:
					{
						return nil, err
					}
				default:
					{
						return res, nil
					}
				}
			}
		case <-timer.C:
			{
				return nil, errors.New("Async_Timeout")
			}
		}
	} else {
		res := <-output
		switch err := res.(type) {
		case error:
			{
				return nil, err
			}
		default:
			{
				return res, nil
			}
		}
	}
}

func Call(m Method, ctx context.Context, timeout time.Duration, args ...interface{}) (interface{}, error) {
	return Lambda(func() (interface{}, error) {
		return m(ctx, args...)
	}, timeout)
}

func All(methods []func() (interface{}, error), timeout time.Duration) []interface{} {
	var wg sync.WaitGroup
	result := make([]interface{}, len(methods))
	for i, m := range methods {
		wg.Add(1)
		go func(index int, method func() (interface{}, error)) {
			defer wg.Done()
			res, err := Lambda(method, timeout)
			if err != nil {
				result[index] = err
			} else {
				result[index] = res
			}
		}(i, m)
	}
	wg.Wait()
	return result
}

func Serise(enter Method, ctx context.Context, args []interface{}, methods []Method, timeout time.Duration) (interface{}, error) {
	var (
		res interface{}
		err error
	)
	res, err = Call(enter, ctx, timeout, args...)
	if err != nil {
		return nil, err
	}
	for _, m := range methods {
		res, err = Call(m, ctx, timeout, res)
		if err != nil {
			return nil, err
		}
	}
	return res, nil
}

func Any(methods []func() (interface{}, error), timeout time.Duration) ([]interface{}, error) {
	resChan := make(chan []interface{})
	errChan := make(chan error)
	go func() {
		defer func() {
			close(resChan)
			close(errChan)
		}()
		var wg sync.WaitGroup
		result := make([]interface{}, len(methods))
		for i, m := range methods {
			wg.Add(1)
			go func(index int, method func() (interface{}, error)) {
				defer wg.Done()
				res, err := Lambda(method, timeout)
				if err != nil {
					errChan <- err
				} else {
					result[index] = res
				}
			}(i, m)
		}
		wg.Wait()
		resChan <- result
	}()
	select {
	case err := <-errChan:
		return nil, err
	case res := <-resChan:
		return res, nil
	}
}
<commit_msg>Update async.go<commit_after>package async

import (
	"context"
	"errors"
	"sync"
	"time"
)

type Method func(ctx context.Context, args ...interface{}) (interface{}, error)
type LambdaMethod func() (interface{}, error)

func Lambda(method LambdaMethod, timeout time.Duration) (interface{}, error) {
	output := make(chan interface{})
	go func() {
		defer close(output)
		res, err := method()
		if err != nil {
			output <- err
		} else {
			output <- res
		}
	}()
	if timeout <= 0 {
		timer := time.NewTimer(timeout)
		defer timer.Stop()
		select {
		case res := <-output:
			{
				switch err := res.(type) {
				case error:
					{
						return nil, err
					}
				default:
					{
						return res, nil
					}
				}
			}
		case <-timer.C:
			{
				return nil, errors.New("Async_Timeout")
			}
		}
	} else {
		res := <-output
		switch err := res.(type) {
		case error:
			{
				return nil, err
			}
		default:
			{
				return res, nil
			}
		}
	}
}

func Call(m Method, ctx context.Context, timeout time.Duration, args ...interface{}) (interface{}, error) {
	return Lambda(LambdaMethod {
		return m(ctx, args...)
	}, timeout)
}

func All(methods []LambdaMethod, timeout time.Duration) []interface{} {
	var wg sync.WaitGroup
	result := make([]interface{}, len(methods))
	for i, m := range methods {
		wg.Add(1)
		go func(index int, method LambdaMethod) {
			defer wg.Done()
			res, err := Lambda(method, timeout)
			if err != nil {
				result[index] = err
			} else {
				result[index] = res
			}
		}(i, m)
	}
	wg.Wait()
	return result
}

func Serise(enter Method, ctx context.Context, args []interface{}, methods []Method, timeout time.Duration) (interface{}, error) {
	var (
		res interface{}
		err error
	)
	res, err = Call(enter, ctx, timeout, args...)
	if err != nil {
		return nil, err
	}
	for _, m := range methods {
		res, err = Call(m, ctx, timeout, res)
		if err != nil {
			return nil, err
		}
	}
	return res, nil
}

func Any(methods []LambdaMethod, timeout time.Duration) ([]interface{}, error) {
	resChan := make(chan []interface{})
	errChan := make(chan error)
	go func() {
		defer func() {
			close(resChan)
			close(errChan)
		}()
		var wg sync.WaitGroup
		result := make([]interface{}, len(methods))
		for i, m := range methods {
			wg.Add(1)
			go func(index int, method LambdaMethod) {
				defer wg.Done()
				res, err := Lambda(method, timeout)
				if err != nil {
					errChan <- err
				} else {
					result[index] = res
				}
			}(i, m)
		}
		wg.Wait()
		resChan <- result
	}()
	select {
	case err := <-errChan:
		return nil, err
	case res := <-resChan:
		return res, nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package authz

import (
	"net/http"

	"github.com/casbin/casbin/v2"
	"github.com/gin-gonic/gin"
)

// NewAuthorizer returns the authorizer, uses a Casbin enforcer as input
func NewAuthorizer(e *casbin.Enforcer) gin.HandlerFunc {
	a := &BasicAuthorizer{enforcer: e}

	return func(c *gin.Context) {
		if !a.CheckPermission(c.Request) {
			a.RequirePermission(c)
		}
	}
}

// BasicAuthorizer stores the casbin handler
type BasicAuthorizer struct {
	enforcer *casbin.Enforcer
}

// GetUserName gets the user name from the request.
// Currently, only HTTP basic authentication is supported
func (a *BasicAuthorizer) GetUserName(r *http.Request) string {
	username, _, _ := r.BasicAuth()
	return username
}

// CheckPermission checks the user/method/path combination from the request.
// Returns true (permission granted) or false (permission forbidden)
func (a *BasicAuthorizer) CheckPermission(r *http.Request) bool {
	user := a.GetUserName(r)
	method := r.Method
	path := r.URL.Path

	allowed, err := a.enforcer.Enforce(user, path, method)
	if err != nil {
		panic(err)
	}

	return allowed
}

// RequirePermission returns the 403 Forbidden to the client
func (a *BasicAuthorizer) RequirePermission(c *gin.Context) {
	c.AbortWithStatus(403)
}
<commit_msg>chore: replace 403 with StatusForbidden const<commit_after>// Copyright 2014 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package authz

import (
	"net/http"

	"github.com/casbin/casbin/v2"
	"github.com/gin-gonic/gin"
)

// NewAuthorizer returns the authorizer, uses a Casbin enforcer as input
func NewAuthorizer(e *casbin.Enforcer) gin.HandlerFunc {
	a := &BasicAuthorizer{enforcer: e}

	return func(c *gin.Context) {
		if !a.CheckPermission(c.Request) {
			a.RequirePermission(c)
		}
	}
}

// BasicAuthorizer stores the casbin handler
type BasicAuthorizer struct {
	enforcer *casbin.Enforcer
}

// GetUserName gets the user name from the request.
// Currently, only HTTP basic authentication is supported
func (a *BasicAuthorizer) GetUserName(r *http.Request) string {
	username, _, _ := r.BasicAuth()
	return username
}

// CheckPermission checks the user/method/path combination from the request.
// Returns true (permission granted) or false (permission forbidden)
func (a *BasicAuthorizer) CheckPermission(r *http.Request) bool {
	user := a.GetUserName(r)
	method := r.Method
	path := r.URL.Path

	allowed, err := a.enforcer.Enforce(user, path, method)
	if err != nil {
		panic(err)
	}

	return allowed
}

// RequirePermission returns the 403 Forbidden to the client
func (a *BasicAuthorizer) RequirePermission(c *gin.Context) {
	c.AbortWithStatus(http.StatusForbidden)
}
<|endoftext|>


<commit_before>// Copyright 2015 The Chromium 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 tsmon

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"golang.org/x/net/context"

	"github.com/luci/luci-go/common/auth"
	"github.com/luci/luci-go/common/gcloud/pubsub"
	"github.com/luci/luci-go/common/logging"
	"github.com/luci/luci-go/common/tsmon/monitor"
	"github.com/luci/luci-go/common/tsmon/store"
	"github.com/luci/luci-go/common/tsmon/target"
	"github.com/luci/luci-go/common/tsmon/types"
)

// Store returns the global metric store that contains all the metric values for
// this process.  Applications shouldn't need to access this directly - instead
// use the metric objects which provide type-safe accessors.
func Store(c context.Context) store.Store {
	return GetState(c).S
}

// Monitor returns the global monitor that sends metrics to monitoring
// endpoints.  Defaults to a nil monitor, but changed by InitializeFromFlags.
func Monitor(c context.Context) monitor.Monitor {
	return GetState(c).M
}

// Register is called by metric objects to register themselves.  This will panic
// if another metric with the same name is already registered.
func Register(c context.Context, m types.Metric) {
	state := GetState(c)
	state.RegisteredMetricsLock.Lock()
	defer state.RegisteredMetricsLock.Unlock()

	if _, ok := state.RegisteredMetrics[m.Info().Name]; ok {
		panic(fmt.Sprintf("A metric with the name '%s' was already registered", m.Info().Name))
	}

	state.RegisteredMetrics[m.Info().Name] = m

	if Store(c) != nil {
		Store(c).Register(m)
	}
}

// Unregister is called by metric objects to unregister themselves.
func Unregister(c context.Context, m types.Metric) {
	state := GetState(c)
	state.RegisteredMetricsLock.Lock()
	defer state.RegisteredMetricsLock.Unlock()

	delete(state.RegisteredMetrics, m.Info().Name)

	if Store(c) != nil {
		Store(c).Unregister(m)
	}
}

// SetStore changes the global metric store.  All metrics that were registered
// with the old store will be re-registered on the new store.
func SetStore(c context.Context, s store.Store) {
	state := GetState(c)
	oldStore := state.S
	if s == oldStore {
		return
	}

	state.RegisteredMetricsLock.RLock()
	defer state.RegisteredMetricsLock.RUnlock()

	// Register metrics on the new store.
	for _, m := range state.RegisteredMetrics {
		s.Register(m)
	}

	state.S = s

	// Unregister metrics from the old store.
	if oldStore != nil {
		for _, m := range state.RegisteredMetrics {
			oldStore.Unregister(m)
		}
	}
}

// InitializeFromFlags configures the tsmon library from flag values.
//
// This will set a Target (information about what's reporting metrics) and a
// Monitor (where to send the metrics to).
func InitializeFromFlags(c context.Context, fl *Flags) error {
	mon, err := initMonitor(c, fl)
	switch {
	case err != nil:
		return err
	case mon == nil:
		return nil // tsmon is disabled
	}

	// Monitoring is enabled, so get the expensive default values for hostname,
	// etc.
	fl.Target.SetDefaultsFromHostname()
	t, err := target.NewFromFlags(&fl.Target)
	if err != nil {
		return err
	}

	Initialize(c, mon, store.NewInMemory(t))

	state := GetState(c)
	if state.Flusher != nil {
		logging.Infof(c, "Canceling previous tsmon auto flush")
		state.Flusher.stop()
		state.Flusher = nil
	}

	if fl.Flush == "auto" {
		state.Flusher = &autoFlusher{}
		state.Flusher.start(c, fl.FlushInterval)
	}

	return nil
}

// Initialize configures the tsmon library with the given monitor and store.
func Initialize(c context.Context, m monitor.Monitor, s store.Store) {
	state := GetState(c)
	state.M = m
	SetStore(c, s)
}

// Shutdown gracefully terminates the tsmon by doing the final flush and
// disabling auto flush (if it was enabled).
//
// It resets Monitor and Store.
//
// Logs error to standard logger. Does nothing if tsmon wasn't initialized.
func Shutdown(c context.Context) {
	state := GetState(c)
	if store.IsNilStore(state.S) {
		return
	}

	if state.Flusher != nil {
		logging.Debugf(c, "Stopping tsmon auto flush")
		state.Flusher.stop()
		state.Flusher = nil
	}

	logging.Debugf(c, "Doing the final tsmon flush")
	if err := Flush(c); err != nil {
		logging.Errorf(c, "Final tsmon flush failed - %s", err)
	} else {
		logging.Debugf(c, "Final tsmon flush finished")
	}

	// Reset the state as if 'InitializeFromFlags' was never called.
	Initialize(c, monitor.NewNilMonitor(), store.NewNilStore())
}

// ResetCumulativeMetrics resets only cumulative metrics.
func ResetCumulativeMetrics(c context.Context) {
	state := GetState(c)
	for _, m := range state.RegisteredMetrics {
		if m.Info().ValueType.IsCumulative() {
			state.S.Reset(c, m)
		}
	}
}

// initMonitor examines flags and config and initializes a monitor.
//
// It returns (nil, nil) if tsmon should be disabled.
func initMonitor(c context.Context, fl *Flags) (monitor.Monitor, error) {
	// Load the config file, and override its values with flags.
	config, err := loadConfig(fl.ConfigFile)
	if err != nil {
		logging.Infof(c, "tsmon is disabled because the config file (%s) could not be loaded: %s",
			fl.ConfigFile, err)
		return nil, nil
	}

	if fl.Endpoint != "" {
		config.Endpoint = fl.Endpoint
	}
	if fl.Credentials != "" {
		config.Credentials = fl.Credentials
	}

	if config.Endpoint == "" {
		logging.Infof(c, "tsmon is disabled because no endpoint is configured")
		return nil, nil
	}
	if strings.ToLower(config.Endpoint) == "none" {
		logging.Infof(c, "tsmon is explicitly disabled")
		return nil, nil
	}

	endpointURL, err := url.Parse(config.Endpoint)
	if err != nil {
		return nil, err
	}

	switch endpointURL.Scheme {
	case "file":
		return monitor.NewDebugMonitor(endpointURL.Path), nil
	case "pubsub":
		client := clientFactory(config.Credentials)
		return monitor.NewPubsubMonitor(client, endpointURL.Host, strings.TrimPrefix(endpointURL.Path, "/"))
	default:
		return nil, fmt.Errorf("unknown tsmon endpoint url: %s", config.Endpoint)
	}
}

// makeClient returns http.Client that knows how to send authenticated requests
// to PubSub API.
func clientFactory(credentials string) monitor.ClientFactory {
	return func(ctx context.Context) (*http.Client, error) {
		authOpts := auth.Options{
			Context: ctx,
			Scopes:  pubsub.PublisherScopes,
		}
		if credentials == GCECredentials {
			authOpts.Method = auth.GCEMetadataMethod
		} else {
			authOpts.Method = auth.ServiceAccountMethod
			authOpts.ServiceAccountJSONPath = credentials
		}
		return auth.NewAuthenticator(auth.SilentLogin, authOpts).Client()
	}
}
<commit_msg>Don't fail when the tsmon config is not found.<commit_after>// Copyright 2015 The Chromium 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 tsmon

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"golang.org/x/net/context"

	"github.com/luci/luci-go/common/auth"
	"github.com/luci/luci-go/common/gcloud/pubsub"
	"github.com/luci/luci-go/common/logging"
	"github.com/luci/luci-go/common/tsmon/monitor"
	"github.com/luci/luci-go/common/tsmon/store"
	"github.com/luci/luci-go/common/tsmon/target"
	"github.com/luci/luci-go/common/tsmon/types"
)

// Store returns the global metric store that contains all the metric values for
// this process.  Applications shouldn't need to access this directly - instead
// use the metric objects which provide type-safe accessors.
func Store(c context.Context) store.Store {
	return GetState(c).S
}

// Monitor returns the global monitor that sends metrics to monitoring
// endpoints.  Defaults to a nil monitor, but changed by InitializeFromFlags.
func Monitor(c context.Context) monitor.Monitor {
	return GetState(c).M
}

// Register is called by metric objects to register themselves.  This will panic
// if another metric with the same name is already registered.
func Register(c context.Context, m types.Metric) {
	state := GetState(c)
	state.RegisteredMetricsLock.Lock()
	defer state.RegisteredMetricsLock.Unlock()

	if _, ok := state.RegisteredMetrics[m.Info().Name]; ok {
		panic(fmt.Sprintf("A metric with the name '%s' was already registered", m.Info().Name))
	}

	state.RegisteredMetrics[m.Info().Name] = m

	if Store(c) != nil {
		Store(c).Register(m)
	}
}

// Unregister is called by metric objects to unregister themselves.
func Unregister(c context.Context, m types.Metric) {
	state := GetState(c)
	state.RegisteredMetricsLock.Lock()
	defer state.RegisteredMetricsLock.Unlock()

	delete(state.RegisteredMetrics, m.Info().Name)

	if Store(c) != nil {
		Store(c).Unregister(m)
	}
}

// SetStore changes the global metric store.  All metrics that were registered
// with the old store will be re-registered on the new store.
func SetStore(c context.Context, s store.Store) {
	state := GetState(c)
	oldStore := state.S
	if s == oldStore {
		return
	}

	state.RegisteredMetricsLock.RLock()
	defer state.RegisteredMetricsLock.RUnlock()

	// Register metrics on the new store.
	for _, m := range state.RegisteredMetrics {
		s.Register(m)
	}

	state.S = s

	// Unregister metrics from the old store.
	if oldStore != nil {
		for _, m := range state.RegisteredMetrics {
			oldStore.Unregister(m)
		}
	}
}

// InitializeFromFlags configures the tsmon library from flag values.
//
// This will set a Target (information about what's reporting metrics) and a
// Monitor (where to send the metrics to).
func InitializeFromFlags(c context.Context, fl *Flags) error {
	mon, err := initMonitor(c, fl)
	switch {
	case err != nil:
		return err
	case mon == nil:
		return nil // tsmon is disabled
	}

	// Monitoring is enabled, so get the expensive default values for hostname,
	// etc.
	fl.Target.SetDefaultsFromHostname()
	t, err := target.NewFromFlags(&fl.Target)
	if err != nil {
		return err
	}

	Initialize(c, mon, store.NewInMemory(t))

	state := GetState(c)
	if state.Flusher != nil {
		logging.Infof(c, "Canceling previous tsmon auto flush")
		state.Flusher.stop()
		state.Flusher = nil
	}

	if fl.Flush == "auto" {
		state.Flusher = &autoFlusher{}
		state.Flusher.start(c, fl.FlushInterval)
	}

	return nil
}

// Initialize configures the tsmon library with the given monitor and store.
func Initialize(c context.Context, m monitor.Monitor, s store.Store) {
	state := GetState(c)
	state.M = m
	SetStore(c, s)
}

// Shutdown gracefully terminates the tsmon by doing the final flush and
// disabling auto flush (if it was enabled).
//
// It resets Monitor and Store.
//
// Logs error to standard logger. Does nothing if tsmon wasn't initialized.
func Shutdown(c context.Context) {
	state := GetState(c)
	if store.IsNilStore(state.S) {
		return
	}

	if state.Flusher != nil {
		logging.Debugf(c, "Stopping tsmon auto flush")
		state.Flusher.stop()
		state.Flusher = nil
	}

	logging.Debugf(c, "Doing the final tsmon flush")
	if err := Flush(c); err != nil {
		logging.Errorf(c, "Final tsmon flush failed - %s", err)
	} else {
		logging.Debugf(c, "Final tsmon flush finished")
	}

	// Reset the state as if 'InitializeFromFlags' was never called.
	Initialize(c, monitor.NewNilMonitor(), store.NewNilStore())
}

// ResetCumulativeMetrics resets only cumulative metrics.
func ResetCumulativeMetrics(c context.Context) {
	state := GetState(c)
	for _, m := range state.RegisteredMetrics {
		if m.Info().ValueType.IsCumulative() {
			state.S.Reset(c, m)
		}
	}
}

// initMonitor examines flags and config and initializes a monitor.
//
// It returns (nil, nil) if tsmon should be disabled.
func initMonitor(c context.Context, fl *Flags) (monitor.Monitor, error) {
	// Load the config file, and override its values with flags.
	config, _ := loadConfig(fl.ConfigFile)

	if fl.Endpoint != "" {
		config.Endpoint = fl.Endpoint
	}
	if fl.Credentials != "" {
		config.Credentials = fl.Credentials
	}

	if config.Endpoint == "" {
		logging.Infof(c, "tsmon is disabled because no endpoint is configured")
		return nil, nil
	}
	if strings.ToLower(config.Endpoint) == "none" {
		logging.Infof(c, "tsmon is explicitly disabled")
		return nil, nil
	}

	endpointURL, err := url.Parse(config.Endpoint)
	if err != nil {
		return nil, err
	}

	switch endpointURL.Scheme {
	case "file":
		return monitor.NewDebugMonitor(endpointURL.Path), nil
	case "pubsub":
		client := clientFactory(config.Credentials)
		return monitor.NewPubsubMonitor(client, endpointURL.Host, strings.TrimPrefix(endpointURL.Path, "/"))
	default:
		return nil, fmt.Errorf("unknown tsmon endpoint url: %s", config.Endpoint)
	}
}

// makeClient returns http.Client that knows how to send authenticated requests
// to PubSub API.
func clientFactory(credentials string) monitor.ClientFactory {
	return func(ctx context.Context) (*http.Client, error) {
		authOpts := auth.Options{
			Context: ctx,
			Scopes:  pubsub.PublisherScopes,
		}
		if credentials == GCECredentials {
			authOpts.Method = auth.GCEMetadataMethod
		} else {
			authOpts.Method = auth.ServiceAccountMethod
			authOpts.ServiceAccountJSONPath = credentials
		}
		return auth.NewAuthenticator(auth.SilentLogin, authOpts).Client()
	}
}
<|endoftext|>


<commit_before>package hostsfile

import (
	"encoding/csv"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode/utf8"
)

const entryTemplate string = "\n%s\t%s\t# HOSTSUP %s"
const entryTag string = "HOSTSUP"

type Hostsfile struct {
	Filename string
	File     *os.File
}

func NewHostsfile(filename string, ro ...bool) (*Hostsfile, error) {
	var f *os.File
	var err error

	// Determine if the hosts file should opened as read only or not.
	if len(ro) > 0 && ro[0] == true {
		f, err = os.OpenFile(filename, os.O_RDONLY, 0666)

	} else {
		f, err = os.OpenFile(filename, os.O_RDWR, 0666)
	}

	if err != nil {
		return nil, err
	}

	return &Hostsfile{filename, f}, nil
}

// Close the hostsfile.
func (h *Hostsfile) Close() error {
	return h.File.Close()
}

func (h *Hostsfile) AddEntry(host *Host) error {
	defer h.File.Seek(0, 0)

	// Go the end of the file to append the new host entry.
	h.File.Seek(0, 2)

	entry := fmt.Sprintf(entryTemplate, host.IP, host.Hostname, host.Id)

	if _, err := h.File.WriteString(entry); err != nil {
		return errors.New("Unable to write entry to hosts file.")
	}

	return nil
}

func (h *Hostsfile) RemoveEntry(host *Host) error {
	defer h.File.Seek(0, 0)

	f, err := ioutil.ReadAll(h.File)

	if err != nil {
		return errors.New("Unable to read hosts file.")
	}

	lines := strings.Split(string(f), "\n")
	updatedLines := []string{}

	for _, line := range lines {
		if !strings.Contains(line, host.Id) {
			updatedLines = append(updatedLines, line)
		}
	}

	output := strings.Join(updatedLines, "\n")

	err = ioutil.WriteFile(h.Filename, []byte(output), 0666)

	if err != nil {
		return errors.New("Unable to remove entry from hosts file.")
	}

	return nil
}

func (h *Hostsfile) FindEntry(hostname string) *Host {
	entries := h.GetEntries()

	for _, entry := range entries {
		if entry.Hostname == hostname {
			return entry
		}
	}

	return nil

}

func (h *Hostsfile) GetEntries() []*Host {
	defer h.File.Seek(0, 0)

	const ipIndex = 0
	const hostnameIndex = 1
	const idIndex = 2

	reader := csv.NewReader(h.File)
	reader.Comma, _ = utf8.DecodeRuneInString("\t")
	reader.Comment, _ = utf8.DecodeRuneInString("#")
	reader.FieldsPerRecord = -1

	lines, _ := reader.ReadAll()

	hosts := make([]*Host, 0)

	for _, line := range lines {
		// Verify that the line contains the entryTag. Hostsup entries will
		// always have 3 columns.
		if len(line) >= 3 && strings.Contains(line[idIndex], entryTag) {
			host := NewHost(line[ipIndex], line[hostnameIndex])
			hosts = append(hosts, host)
		}
	}

	return hosts
}

func (h *Hostsfile) Clean() []*Host {
	entries := h.GetEntries()

	for _, entry := range entries {
		h.RemoveEntry(entry)
	}

	return entries
}
<commit_msg>Add comments to hostsfile.<commit_after>package hostsfile

import (
	"encoding/csv"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode/utf8"
)

const entryTemplate string = "\n%s\t%s\t# HOSTSUP %s"
const entryTag string = "HOSTSUP"

type Hostsfile struct {
	Filename string
	File     *os.File
}

// Instantiate a new Hostsfile.
func NewHostsfile(filename string, ro ...bool) (*Hostsfile, error) {
	var f *os.File
	var err error

	// Determine if the hosts file should opened as read only or not.
	if len(ro) > 0 && ro[0] == true {
		f, err = os.OpenFile(filename, os.O_RDONLY, 0666)

	} else {
		f, err = os.OpenFile(filename, os.O_RDWR, 0666)
	}

	if err != nil {
		return nil, err
	}

	return &Hostsfile{filename, f}, nil
}

// Close the hostsfile.
func (h *Hostsfile) Close() error {
	return h.File.Close()
}

// Add a new entry to the hosts file. The entry resembles:
// <IP Address>	<hostname>	# HOSTSUP <md5 sum>
func (h *Hostsfile) AddEntry(host *Host) error {
	defer h.File.Seek(0, 0)

	// Go the end of the file to append the new host entry.
	h.File.Seek(0, 2)

	entry := fmt.Sprintf(entryTemplate, host.IP, host.Hostname, host.Id)

	if _, err := h.File.WriteString(entry); err != nil {
		return errors.New("Unable to write entry to hosts file.")
	}

	return nil
}

// Remove an entry from the hosts file.
func (h *Hostsfile) RemoveEntry(host *Host) error {
	defer h.File.Seek(0, 0)

	// Read the contents of the hosts file.
	f, err := ioutil.ReadAll(h.File)

	if err != nil {
		return errors.New("Unable to read hosts file.")
	}

	// Explode the contents into a slice by line. Create an empty slice for
	// the modified file contents.
	lines := strings.Split(string(f), "\n")
	updatedLines := []string{}

	// Filter out any lines of the file that contain the Host ID.
	for _, line := range lines {
		if !strings.Contains(line, host.Id) {
			updatedLines = append(updatedLines, line)
		}
	}

	// Implode the lines and write them back to the file.
	output := strings.Join(updatedLines, "\n")
	err = ioutil.WriteFile(h.Filename, []byte(output), 0666)

	if err != nil {
		return errors.New("Unable to remove entry from hosts file.")
	}

	return nil
}

// Find the first host entry that matches a hostname.
func (h *Hostsfile) FindEntry(hostname string) *Host {
	entries := h.GetEntries()

	for _, entry := range entries {
		if entry.Hostname == hostname {
			return entry
		}
	}

	return nil
}

// Get a list of all the host entries that were added. This will exclude any
// entry that was not added by AddEntry().
func (h *Hostsfile) GetEntries() []*Host {
	defer h.File.Seek(0, 0)

	var (
		ipIndex       = 0
		hostnameIndex = 1
		idIndex       = 2
	)

	// Set up the Reader for a tab delimited file.
	reader := csv.NewReader(h.File)
	reader.Comma, _ = utf8.DecodeRuneInString("\t")
	reader.Comment, _ = utf8.DecodeRuneInString("#")
	reader.FieldsPerRecord = -1

	// Get all lines. Sets lines to [][]string.
	lines, _ := reader.ReadAll()
	hosts := make([]*Host, 0)

	for _, line := range lines {
		// Verify that the line contains the entryTag. Hostsup entries will
		// always have 3 columns.
		if len(line) >= 3 && strings.Contains(line[idIndex], entryTag) {
			host := NewHost(line[ipIndex], line[hostnameIndex])
			hosts = append(hosts, host)
		}
	}

	return hosts
}

// Remove all entries added by AddEntries.
func (h *Hostsfile) Clean() []*Host {
	entries := h.GetEntries()

	for _, entry := range entries {
		h.RemoveEntry(entry)
	}

	return entries
}
<|endoftext|>


<commit_before>package backends

import (
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"reflect"
	"time"
)

// MongodbBackend represents a MongoDB result backend
type MongodbBackend struct {
	cnf                  *config.Config
	session              *mgo.Session
	tasksCollection      *mgo.Collection
	groupMetasCollection *mgo.Collection
}

// NewMongodbBackend creates MongodbBackend instance
func NewMongodbBackend(cnf *config.Config) Interface {
	return &MongodbBackend{cnf: cnf}
}

// InitGroup creates and saves a group meta data object
func (b *MongodbBackend) InitGroup(groupUUID string, taskUUIDs []string) error {
	if err := b.connect(); err != nil {
		return err
	}

	groupMeta := &tasks.GroupMeta{
		GroupUUID: groupUUID,
		TaskUUIDs: taskUUIDs,
	}
	return b.groupMetasCollection.Insert(groupMeta)
}

// GroupCompleted returns true if all tasks in a group finished
func (b *MongodbBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	taskStates, err := b.getStates(groupMeta.TaskUUIDs...)
	if err != nil {
		return false, err
	}

	var countSuccessTasks = 0
	for _, taskState := range taskStates {
		if taskState.IsCompleted() {
			countSuccessTasks++
		}
	}

	return countSuccessTasks == groupTaskCount, nil
}

// GroupTaskStates returns states of all tasks in the group
func (b *MongodbBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return []*tasks.TaskState{}, err
	}

	return b.getStates(groupMeta.TaskUUIDs...)
}

// TriggerChord flags chord as triggered in the backend storage to make sure
// chord is never trigerred multiple times. Returns a boolean flag to indicate
// whether the worker should trigger chord (true) or no if it has been triggered
// already (false)
func (b *MongodbBackend) TriggerChord(groupUUID string) (bool, error) {
	if err := b.connect(); err != nil {
		return false, err
	}

	// Get the group meta data
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	// Chord has already been triggered, return false (should not trigger again)
	if groupMeta.ChordTriggered {
		return false, nil
	}

	// If group meta is locked, wait until it's unlocked
	for groupMeta.Lock {
		groupMeta, _ = b.getGroupMeta(groupUUID)
		log.WARNING.Print("Group meta locked, waiting")
		<-time.After(time.Millisecond * 5)
	}

	// Acquire lock
	if err = b.lockGroupMeta(groupUUID); err != nil {
		return false, err
	}
	defer b.unlockGroupMeta(groupUUID)

	// Update the group meta data
	update := bson.M{"$set": bson.M{"chord_triggered": true}}
	_, err = b.groupMetasCollection.UpsertId(groupUUID, update)
	if err != nil {
		return false, err
	}

	return true, nil
}

// SetStatePending updates task state to PENDING
func (b *MongodbBackend) SetStatePending(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StatePending}
	return b.updateState(signature, update)
}

// SetStateReceived updates task state to RECEIVED
func (b *MongodbBackend) SetStateReceived(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateReceived}
	return b.updateState(signature, update)
}

// SetStateStarted updates task state to STARTED
func (b *MongodbBackend) SetStateStarted(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateStarted}
	return b.updateState(signature, update)
}

// SetStateRetry updates task state to RETRY
func (b *MongodbBackend) SetStateRetry(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateRetry}
	return b.updateState(signature, update)
}

// SetStateSuccess updates task state to SUCCESS
func (b *MongodbBackend) SetStateSuccess(signature *tasks.Signature, results []*tasks.TaskResult) error {
	//edited by surendra tiwari
	var err error
	bsonResults := make([]bson.M, len(results))
	for i, result := range results {
		//to hold the json result
		bsonResult_ := bson.M{}
		result_type := reflect.TypeOf(result.Value).Kind()
		if result_type == reflect.String {
			//convert type to json
			err = bson.UnmarshalJSON([]byte(result.Value.(string)), &bsonResult_)
			if err == nil {
				bsonResults[i] = bson.M{
					"type":  "Json",
					"value": bsonResult_,
				}
			} else {
				bsonResults[i] = bson.M{
					"type":  result.Type,
					"value": result.Value,
				}
			}
		} else {
			bsonResults[i] = bson.M{
				"type":  result.Type,
				"value": result.Value,
			}
		}
	}
	update := bson.M{
		"state":   tasks.StateSuccess,
		"results": bsonResults,
	}
	return b.updateState(signature, update)
}

// SetStateFailure updates task state to FAILURE
func (b *MongodbBackend) SetStateFailure(signature *tasks.Signature, err string) error {
	update := bson.M{"state": tasks.StateFailure, "error": err}
	return b.updateState(signature, update)
}

// GetState returns the latest task state
func (b *MongodbBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	state := new(tasks.TaskState)
	if err := b.tasksCollection.FindId(taskUUID).One(state); err != nil {
		return nil, err
	}
	return state, nil
}

// PurgeState deletes stored task state
func (b *MongodbBackend) PurgeState(taskUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.tasksCollection.RemoveId(taskUUID)
}

// PurgeGroupMeta deletes stored group meta data
func (b *MongodbBackend) PurgeGroupMeta(groupUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.groupMetasCollection.RemoveId(groupUUID)
}

// lockGroupMeta acquires lock on groupUUID document
func (b *MongodbBackend) lockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": true}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// unlockGroupMeta releases lock on groupUUID document
func (b *MongodbBackend) unlockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": false}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// getGroupMeta retrieves group meta data, convenience function to avoid repetition
func (b *MongodbBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	query := bson.M{"_id": groupUUID}

	groupMeta := new(tasks.GroupMeta)
	if err := b.groupMetasCollection.Find(query).One(groupMeta); err != nil {
		return nil, err
	}
	return groupMeta, nil
}

// getStates returns multiple task states
func (b *MongodbBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	states := make([]*tasks.TaskState, 0, len(taskUUIDs))

	iter := b.tasksCollection.Find(bson.M{"_id": bson.M{"$in": taskUUIDs}}).Iter()

	state := new(tasks.TaskState)
	for iter.Next(state) {
		states = append(states, state)

		// otherwise we would end up with the last task being every element of the slice
		state = new(tasks.TaskState)
	}

	return states, nil
}

// updateState saves current task state
func (b *MongodbBackend) updateState(signature *tasks.Signature, update bson.M) error {
	if err := b.connect(); err != nil {
		return err
	}

	update = bson.M{"$set": update}
	_, err := b.tasksCollection.UpsertId(signature.UUID, update)
	if err != nil {
		return err
	}
	return nil
}

// connect returns a session if we are already connected to mongo, otherwise
// (when called for the first time) it will open a new session and ensure
// all required indexes for our collections exist
func (b *MongodbBackend) connect() error {
	if b.session != nil {
		return nil
	}

	session, err := mgo.Dial(b.cnf.ResultBackend)
	if err != nil {
		return err
	}
	b.session = session

	b.tasksCollection = b.session.DB("").C("tasks")
	b.groupMetasCollection = b.session.DB("").C("group_metas")

	return b.createMongoIndexes()
}

// createMongoIndexes ensures all indexes are in place
func (b *MongodbBackend) createMongoIndexes() error {
	indexes := []mgo.Index{
		{
			Key:         []string{"state"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
		{
			Key:         []string{"lock"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
	}

	for _, index := range indexes {
		// Check if index already exists, if it does, skip
		if err := b.tasksCollection.EnsureIndex(index); err == nil {
			log.INFO.Printf("%s index already exist, skipping create step", index.Key[0])
			continue
		}

		// Create index (keep in mind EnsureIndex is blocking operation)
		log.INFO.Printf("Creating %s index", index.Key[0])
		if err := b.tasksCollection.DropIndex(index.Key[0]); err != nil {
			return err
		}
		if err := b.tasksCollection.EnsureIndex(index); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>goimport to mongodb.go<commit_after>package backends

import (
	"reflect"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// MongodbBackend represents a MongoDB result backend
type MongodbBackend struct {
	cnf                  *config.Config
	session              *mgo.Session
	tasksCollection      *mgo.Collection
	groupMetasCollection *mgo.Collection
}

// NewMongodbBackend creates MongodbBackend instance
func NewMongodbBackend(cnf *config.Config) Interface {
	return &MongodbBackend{cnf: cnf}
}

// InitGroup creates and saves a group meta data object
func (b *MongodbBackend) InitGroup(groupUUID string, taskUUIDs []string) error {
	if err := b.connect(); err != nil {
		return err
	}

	groupMeta := &tasks.GroupMeta{
		GroupUUID: groupUUID,
		TaskUUIDs: taskUUIDs,
	}
	return b.groupMetasCollection.Insert(groupMeta)
}

// GroupCompleted returns true if all tasks in a group finished
func (b *MongodbBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	taskStates, err := b.getStates(groupMeta.TaskUUIDs...)
	if err != nil {
		return false, err
	}

	var countSuccessTasks = 0
	for _, taskState := range taskStates {
		if taskState.IsCompleted() {
			countSuccessTasks++
		}
	}

	return countSuccessTasks == groupTaskCount, nil
}

// GroupTaskStates returns states of all tasks in the group
func (b *MongodbBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return []*tasks.TaskState{}, err
	}

	return b.getStates(groupMeta.TaskUUIDs...)
}

// TriggerChord flags chord as triggered in the backend storage to make sure
// chord is never trigerred multiple times. Returns a boolean flag to indicate
// whether the worker should trigger chord (true) or no if it has been triggered
// already (false)
func (b *MongodbBackend) TriggerChord(groupUUID string) (bool, error) {
	if err := b.connect(); err != nil {
		return false, err
	}

	// Get the group meta data
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	// Chord has already been triggered, return false (should not trigger again)
	if groupMeta.ChordTriggered {
		return false, nil
	}

	// If group meta is locked, wait until it's unlocked
	for groupMeta.Lock {
		groupMeta, _ = b.getGroupMeta(groupUUID)
		log.WARNING.Print("Group meta locked, waiting")
		<-time.After(time.Millisecond * 5)
	}

	// Acquire lock
	if err = b.lockGroupMeta(groupUUID); err != nil {
		return false, err
	}
	defer b.unlockGroupMeta(groupUUID)

	// Update the group meta data
	update := bson.M{"$set": bson.M{"chord_triggered": true}}
	_, err = b.groupMetasCollection.UpsertId(groupUUID, update)
	if err != nil {
		return false, err
	}

	return true, nil
}

// SetStatePending updates task state to PENDING
func (b *MongodbBackend) SetStatePending(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StatePending}
	return b.updateState(signature, update)
}

// SetStateReceived updates task state to RECEIVED
func (b *MongodbBackend) SetStateReceived(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateReceived}
	return b.updateState(signature, update)
}

// SetStateStarted updates task state to STARTED
func (b *MongodbBackend) SetStateStarted(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateStarted}
	return b.updateState(signature, update)
}

// SetStateRetry updates task state to RETRY
func (b *MongodbBackend) SetStateRetry(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateRetry}
	return b.updateState(signature, update)
}

// SetStateSuccess updates task state to SUCCESS
func (b *MongodbBackend) SetStateSuccess(signature *tasks.Signature, results []*tasks.TaskResult) error {
	//edited by surendra tiwari
	var err error
	bsonResults := make([]bson.M, len(results))
	for i, result := range results {
		//to hold the json result
		bsonResult_ := bson.M{}
		result_type := reflect.TypeOf(result.Value).Kind()
		if result_type == reflect.String {
			//convert type to json
			err = bson.UnmarshalJSON([]byte(result.Value.(string)), &bsonResult_)
			if err == nil {
				bsonResults[i] = bson.M{
					"type":  "Json",
					"value": bsonResult_,
				}
			} else {
				bsonResults[i] = bson.M{
					"type":  result.Type,
					"value": result.Value,
				}
			}
		} else {
			bsonResults[i] = bson.M{
				"type":  result.Type,
				"value": result.Value,
			}
		}
	}
	update := bson.M{
		"state":   tasks.StateSuccess,
		"results": bsonResults,
	}
	return b.updateState(signature, update)
}

// SetStateFailure updates task state to FAILURE
func (b *MongodbBackend) SetStateFailure(signature *tasks.Signature, err string) error {
	update := bson.M{"state": tasks.StateFailure, "error": err}
	return b.updateState(signature, update)
}

// GetState returns the latest task state
func (b *MongodbBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	state := new(tasks.TaskState)
	if err := b.tasksCollection.FindId(taskUUID).One(state); err != nil {
		return nil, err
	}
	return state, nil
}

// PurgeState deletes stored task state
func (b *MongodbBackend) PurgeState(taskUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.tasksCollection.RemoveId(taskUUID)
}

// PurgeGroupMeta deletes stored group meta data
func (b *MongodbBackend) PurgeGroupMeta(groupUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.groupMetasCollection.RemoveId(groupUUID)
}

// lockGroupMeta acquires lock on groupUUID document
func (b *MongodbBackend) lockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": true}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// unlockGroupMeta releases lock on groupUUID document
func (b *MongodbBackend) unlockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": false}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// getGroupMeta retrieves group meta data, convenience function to avoid repetition
func (b *MongodbBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	query := bson.M{"_id": groupUUID}

	groupMeta := new(tasks.GroupMeta)
	if err := b.groupMetasCollection.Find(query).One(groupMeta); err != nil {
		return nil, err
	}
	return groupMeta, nil
}

// getStates returns multiple task states
func (b *MongodbBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	states := make([]*tasks.TaskState, 0, len(taskUUIDs))

	iter := b.tasksCollection.Find(bson.M{"_id": bson.M{"$in": taskUUIDs}}).Iter()

	state := new(tasks.TaskState)
	for iter.Next(state) {
		states = append(states, state)

		// otherwise we would end up with the last task being every element of the slice
		state = new(tasks.TaskState)
	}

	return states, nil
}

// updateState saves current task state
func (b *MongodbBackend) updateState(signature *tasks.Signature, update bson.M) error {
	if err := b.connect(); err != nil {
		return err
	}

	update = bson.M{"$set": update}
	_, err := b.tasksCollection.UpsertId(signature.UUID, update)
	if err != nil {
		return err
	}
	return nil
}

// connect returns a session if we are already connected to mongo, otherwise
// (when called for the first time) it will open a new session and ensure
// all required indexes for our collections exist
func (b *MongodbBackend) connect() error {
	if b.session != nil {
		return nil
	}

	session, err := mgo.Dial(b.cnf.ResultBackend)
	if err != nil {
		return err
	}
	b.session = session

	b.tasksCollection = b.session.DB("").C("tasks")
	b.groupMetasCollection = b.session.DB("").C("group_metas")

	return b.createMongoIndexes()
}

// createMongoIndexes ensures all indexes are in place
func (b *MongodbBackend) createMongoIndexes() error {
	indexes := []mgo.Index{
		{
			Key:         []string{"state"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
		{
			Key:         []string{"lock"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
	}

	for _, index := range indexes {
		// Check if index already exists, if it does, skip
		if err := b.tasksCollection.EnsureIndex(index); err == nil {
			log.INFO.Printf("%s index already exist, skipping create step", index.Key[0])
			continue
		}

		// Create index (keep in mind EnsureIndex is blocking operation)
		log.INFO.Printf("Creating %s index", index.Key[0])
		if err := b.tasksCollection.DropIndex(index.Key[0]); err != nil {
			return err
		}
		if err := b.tasksCollection.EnsureIndex(index); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package backends

import (
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
)

// MongodbBackend represents a MongoDB result backend
type MongodbBackend struct {
	cnf                  *config.Config
	session              *mgo.Session
	tasksCollection      *mgo.Collection
	groupMetasCollection *mgo.Collection
}

// NewMongodbBackend creates MongodbBackend instance
func NewMongodbBackend(cnf *config.Config) Interface {
	return &MongodbBackend{cnf: cnf}
}

// InitGroup creates and saves a group meta data object
func (b *MongodbBackend) InitGroup(groupUUID string, taskUUIDs []string) error {
	if err := b.connect(); err != nil {
		return err
	}

	groupMeta := &tasks.GroupMeta{
		GroupUUID: groupUUID,
		TaskUUIDs: taskUUIDs,
	}
	return b.groupMetasCollection.Insert(groupMeta)
}

// GroupCompleted returns true if all tasks in a group finished
func (b *MongodbBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	taskStates, err := b.getStates(groupMeta.TaskUUIDs...)
	if err != nil {
		return false, err
	}

	var countSuccessTasks = 0
	for _, taskState := range taskStates {
		if taskState.IsCompleted() {
			countSuccessTasks++
		}
	}

	return countSuccessTasks == groupTaskCount, nil
}

// GroupTaskStates returns states of all tasks in the group
func (b *MongodbBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return []*tasks.TaskState{}, err
	}

	return b.getStates(groupMeta.TaskUUIDs...)
}

// TriggerChord flags chord as triggered in the backend storage to make sure
// chord is never trigerred multiple times. Returns a boolean flag to indicate
// whether the worker should trigger chord (true) or no if it has been triggered
// already (false)
func (b *MongodbBackend) TriggerChord(groupUUID string) (bool, error) {
	if err := b.connect(); err != nil {
		return false, err
	}

	// Get the group meta data
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	// Chord has already been triggered, return false (should not trigger again)
	if groupMeta.ChordTriggered {
		return false, nil
	}

	// If group meta is locked, wait until it's unlocked
	for groupMeta.Lock {
		groupMeta, _ = b.getGroupMeta(groupUUID)
		log.WARNING.Print("Group meta locked, waiting")
		<-time.After(time.Millisecond * 5)
	}

	// Acquire lock
	if err = b.lockGroupMeta(groupUUID); err != nil {
		return false, err
	}
	defer b.unlockGroupMeta(groupUUID)

	// Update the group meta data
	update := bson.M{"$set": bson.M{"chord_triggered": true}}
	_, err = b.groupMetasCollection.UpsertId(groupUUID, update)
	if err != nil {
		return false, err
	}

	return true, nil
}

// SetStatePending updates task state to PENDING
func (b *MongodbBackend) SetStatePending(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StatePending}
	return b.updateState(signature, update)
}

// SetStateReceived updates task state to RECEIVED
func (b *MongodbBackend) SetStateReceived(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateReceived}
	return b.updateState(signature, update)
}

// SetStateStarted updates task state to STARTED
func (b *MongodbBackend) SetStateStarted(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateStarted}
	return b.updateState(signature, update)
}

// SetStateRetry updates task state to RETRY
func (b *MongodbBackend) SetStateRetry(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateRetry}
	return b.updateState(signature, update)
}

// SetStateSuccess updates task state to SUCCESS
func (b *MongodbBackend) SetStateSuccess(signature *tasks.Signature, results []*tasks.TaskResult) error {
//edited by surendra tiwari
	bsonResults := make([]bson.M, len(results))
	for i, result := range results {
		//to hold the json result
		bsonResult_ := bson.M{}
		//convert type to json
		err := bson.UnmarshalJSON([]byte(result.Value.(string)),&bsonResult_)
        	//if we get error then it means its not convertable to json so push it as it is
		if err != nil {
                	bsonResults[i] = bson.M{
                        "type":  result.Type,
                        "value": result.Value,
                	}
        	}else {
			bsonResults[i] = bson.M{
			"type":  "Json",
			"value": bsonResult_,
			}
		}
	}
	update := bson.M{
		"state":   tasks.StateSuccess,
		"results": bsonResults,
	}
	return b.updateState(signature, update)
}

// SetStateFailure updates task state to FAILURE
func (b *MongodbBackend) SetStateFailure(signature *tasks.Signature, err string) error {
	update := bson.M{"state": tasks.StateFailure, "error": err}
	return b.updateState(signature, update)
}

// GetState returns the latest task state
func (b *MongodbBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	state := new(tasks.TaskState)
	if err := b.tasksCollection.FindId(taskUUID).One(state); err != nil {
		return nil, err
	}
	return state, nil
}

// PurgeState deletes stored task state
func (b *MongodbBackend) PurgeState(taskUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.tasksCollection.RemoveId(taskUUID)
}

// PurgeGroupMeta deletes stored group meta data
func (b *MongodbBackend) PurgeGroupMeta(groupUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.groupMetasCollection.RemoveId(groupUUID)
}

// lockGroupMeta acquires lock on groupUUID document
func (b *MongodbBackend) lockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": true}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// unlockGroupMeta releases lock on groupUUID document
func (b *MongodbBackend) unlockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": false}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// getGroupMeta retrieves group meta data, convenience function to avoid repetition
func (b *MongodbBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	query := bson.M{"_id": groupUUID}

	groupMeta := new(tasks.GroupMeta)
	if err := b.groupMetasCollection.Find(query).One(groupMeta); err != nil {
		return nil, err
	}
	return groupMeta, nil
}

// getStates returns multiple task states
func (b *MongodbBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	states := make([]*tasks.TaskState, 0, len(taskUUIDs))

	iter := b.tasksCollection.Find(bson.M{"_id": bson.M{"$in": taskUUIDs}}).Iter()

	state := new(tasks.TaskState)
	for iter.Next(state) {
		states = append(states, state)

		// otherwise we would end up with the last task being every element of the slice
		state = new(tasks.TaskState)
	}

	return states, nil
}

// updateState saves current task state
func (b *MongodbBackend) updateState(signature *tasks.Signature, update bson.M) error {
	if err := b.connect(); err != nil {
		return err
	}

	update = bson.M{"$set": update}
	_, err := b.tasksCollection.UpsertId(signature.UUID, update)
	if err != nil {
		return err
	}
	return nil
}

// connect returns a session if we are already connected to mongo, otherwise
// (when called for the first time) it will open a new session and ensure
// all required indexes for our collections exist
func (b *MongodbBackend) connect() error {
	if b.session != nil {
		return nil
	}

	session, err := mgo.Dial(b.cnf.ResultBackend)
	if err != nil {
		return err
	}
	b.session = session

	b.tasksCollection = b.session.DB("").C("tasks")
	b.groupMetasCollection = b.session.DB("").C("group_metas")

	return b.createMongoIndexes()
}

// createMongoIndexes ensures all indexes are in place
func (b *MongodbBackend) createMongoIndexes() error {
	indexes := []mgo.Index{
		{
			Key:         []string{"state"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
		{
			Key:         []string{"lock"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
	}

	for _, index := range indexes {
		// Check if index already exists, if it does, skip
		if err := b.tasksCollection.EnsureIndex(index); err == nil {
			log.INFO.Printf("%s index already exist, skipping create step", index.Key[0])
			continue
		}

		// Create index (keep in mind EnsureIndex is blocking operation)
		log.INFO.Printf("Creating %s index", index.Key[0])
		if err := b.tasksCollection.DropIndex(index.Key[0]); err != nil {
			return err
		}
		if err := b.tasksCollection.EnsureIndex(index); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>gofmt to backend mongodb.go<commit_after>package backends

import (
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
)

// MongodbBackend represents a MongoDB result backend
type MongodbBackend struct {
	cnf                  *config.Config
	session              *mgo.Session
	tasksCollection      *mgo.Collection
	groupMetasCollection *mgo.Collection
}

// NewMongodbBackend creates MongodbBackend instance
func NewMongodbBackend(cnf *config.Config) Interface {
	return &MongodbBackend{cnf: cnf}
}

// InitGroup creates and saves a group meta data object
func (b *MongodbBackend) InitGroup(groupUUID string, taskUUIDs []string) error {
	if err := b.connect(); err != nil {
		return err
	}

	groupMeta := &tasks.GroupMeta{
		GroupUUID: groupUUID,
		TaskUUIDs: taskUUIDs,
	}
	return b.groupMetasCollection.Insert(groupMeta)
}

// GroupCompleted returns true if all tasks in a group finished
func (b *MongodbBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	taskStates, err := b.getStates(groupMeta.TaskUUIDs...)
	if err != nil {
		return false, err
	}

	var countSuccessTasks = 0
	for _, taskState := range taskStates {
		if taskState.IsCompleted() {
			countSuccessTasks++
		}
	}

	return countSuccessTasks == groupTaskCount, nil
}

// GroupTaskStates returns states of all tasks in the group
func (b *MongodbBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return []*tasks.TaskState{}, err
	}

	return b.getStates(groupMeta.TaskUUIDs...)
}

// TriggerChord flags chord as triggered in the backend storage to make sure
// chord is never trigerred multiple times. Returns a boolean flag to indicate
// whether the worker should trigger chord (true) or no if it has been triggered
// already (false)
func (b *MongodbBackend) TriggerChord(groupUUID string) (bool, error) {
	if err := b.connect(); err != nil {
		return false, err
	}

	// Get the group meta data
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	// Chord has already been triggered, return false (should not trigger again)
	if groupMeta.ChordTriggered {
		return false, nil
	}

	// If group meta is locked, wait until it's unlocked
	for groupMeta.Lock {
		groupMeta, _ = b.getGroupMeta(groupUUID)
		log.WARNING.Print("Group meta locked, waiting")
		<-time.After(time.Millisecond * 5)
	}

	// Acquire lock
	if err = b.lockGroupMeta(groupUUID); err != nil {
		return false, err
	}
	defer b.unlockGroupMeta(groupUUID)

	// Update the group meta data
	update := bson.M{"$set": bson.M{"chord_triggered": true}}
	_, err = b.groupMetasCollection.UpsertId(groupUUID, update)
	if err != nil {
		return false, err
	}

	return true, nil
}

// SetStatePending updates task state to PENDING
func (b *MongodbBackend) SetStatePending(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StatePending}
	return b.updateState(signature, update)
}

// SetStateReceived updates task state to RECEIVED
func (b *MongodbBackend) SetStateReceived(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateReceived}
	return b.updateState(signature, update)
}

// SetStateStarted updates task state to STARTED
func (b *MongodbBackend) SetStateStarted(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateStarted}
	return b.updateState(signature, update)
}

// SetStateRetry updates task state to RETRY
func (b *MongodbBackend) SetStateRetry(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateRetry}
	return b.updateState(signature, update)
}

// SetStateSuccess updates task state to SUCCESS
func (b *MongodbBackend) SetStateSuccess(signature *tasks.Signature, results []*tasks.TaskResult) error {
	//edited by surendra tiwari
	bsonResults := make([]bson.M, len(results))
	for i, result := range results {
		//to hold the json result
		bsonResult_ := bson.M{}
		//convert type to json
		err := bson.UnmarshalJSON([]byte(result.Value.(string)), &bsonResult_)
		//if we get error then it means its not convertable to json so push it as it is
		if err != nil {
			bsonResults[i] = bson.M{
				"type":  result.Type,
				"value": result.Value,
			}
		} else {
			bsonResults[i] = bson.M{
				"type":  "Json",
				"value": bsonResult_,
			}
		}
	}
	update := bson.M{
		"state":   tasks.StateSuccess,
		"results": bsonResults,
	}
	return b.updateState(signature, update)
}

// SetStateFailure updates task state to FAILURE
func (b *MongodbBackend) SetStateFailure(signature *tasks.Signature, err string) error {
	update := bson.M{"state": tasks.StateFailure, "error": err}
	return b.updateState(signature, update)
}

// GetState returns the latest task state
func (b *MongodbBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	state := new(tasks.TaskState)
	if err := b.tasksCollection.FindId(taskUUID).One(state); err != nil {
		return nil, err
	}
	return state, nil
}

// PurgeState deletes stored task state
func (b *MongodbBackend) PurgeState(taskUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.tasksCollection.RemoveId(taskUUID)
}

// PurgeGroupMeta deletes stored group meta data
func (b *MongodbBackend) PurgeGroupMeta(groupUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.groupMetasCollection.RemoveId(groupUUID)
}

// lockGroupMeta acquires lock on groupUUID document
func (b *MongodbBackend) lockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": true}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// unlockGroupMeta releases lock on groupUUID document
func (b *MongodbBackend) unlockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": false}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// getGroupMeta retrieves group meta data, convenience function to avoid repetition
func (b *MongodbBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	query := bson.M{"_id": groupUUID}

	groupMeta := new(tasks.GroupMeta)
	if err := b.groupMetasCollection.Find(query).One(groupMeta); err != nil {
		return nil, err
	}
	return groupMeta, nil
}

// getStates returns multiple task states
func (b *MongodbBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	states := make([]*tasks.TaskState, 0, len(taskUUIDs))

	iter := b.tasksCollection.Find(bson.M{"_id": bson.M{"$in": taskUUIDs}}).Iter()

	state := new(tasks.TaskState)
	for iter.Next(state) {
		states = append(states, state)

		// otherwise we would end up with the last task being every element of the slice
		state = new(tasks.TaskState)
	}

	return states, nil
}

// updateState saves current task state
func (b *MongodbBackend) updateState(signature *tasks.Signature, update bson.M) error {
	if err := b.connect(); err != nil {
		return err
	}

	update = bson.M{"$set": update}
	_, err := b.tasksCollection.UpsertId(signature.UUID, update)
	if err != nil {
		return err
	}
	return nil
}

// connect returns a session if we are already connected to mongo, otherwise
// (when called for the first time) it will open a new session and ensure
// all required indexes for our collections exist
func (b *MongodbBackend) connect() error {
	if b.session != nil {
		return nil
	}

	session, err := mgo.Dial(b.cnf.ResultBackend)
	if err != nil {
		return err
	}
	b.session = session

	b.tasksCollection = b.session.DB("").C("tasks")
	b.groupMetasCollection = b.session.DB("").C("group_metas")

	return b.createMongoIndexes()
}

// createMongoIndexes ensures all indexes are in place
func (b *MongodbBackend) createMongoIndexes() error {
	indexes := []mgo.Index{
		{
			Key:         []string{"state"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
		{
			Key:         []string{"lock"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
	}

	for _, index := range indexes {
		// Check if index already exists, if it does, skip
		if err := b.tasksCollection.EnsureIndex(index); err == nil {
			log.INFO.Printf("%s index already exist, skipping create step", index.Key[0])
			continue
		}

		// Create index (keep in mind EnsureIndex is blocking operation)
		log.INFO.Printf("Creating %s index", index.Key[0])
		if err := b.tasksCollection.DropIndex(index.Key[0]); err != nil {
			return err
		}
		if err := b.tasksCollection.EnsureIndex(index); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>