<commit_before>package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/go-version"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/helper/args"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reRktVersion  = regexp.MustCompile(`rkt version (\d[.\d]+)`)
	reAppcVersion = regexp.MustCompile(`appc version (\d[.\d]+)`)
)

const (
	// minRktVersion is the earliest supported version of rkt. rkt added support
	// for CPU and memory isolators in 0.14.0. We cannot support an earlier
	// version to maintain an uniform interface across all drivers
	minRktVersion = "0.14.0"

	// bytesToMB is the conversion from bytes to megabytes.
	bytesToMB = 1024 * 1024
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type RktDriverConfig struct {
	ImageName string   `mapstructure:"image"`
	Args      []string `mapstructure:"args"`
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	proc   *os.Process
	image  string
	logger *log.Logger
	waitCh chan *cstructs.WaitResult
	doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	Pid   int
	Image string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{DriverContext: *ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reAppcVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "1"
	node.Attributes["driver.rkt.version"] = rktMatches[1]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	minVersion, _ := version.NewVersion(minRktVersion)
	currentVersion, _ := version.NewVersion(node.Attributes["driver.rkt.version"])
	if currentVersion.LessThan(minVersion) {
		// Do not allow rkt < 0.14.0
		d.logger.Printf("[WARN] driver.rkt: please upgrade rkt to a version >= %s", minVersion)
		node.Attributes["driver.rkt"] = "0"
	}
	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig RktDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}
	// Validate that the config is valid.
	img := driverConfig.ImageName
	if img == "" {
		return nil, fmt.Errorf("Missing ACI image for rkt")
	}

	// Get the tasks local directory.
	taskName := d.DriverContext.taskName
	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}
	taskLocal := filepath.Join(taskDir, allocdir.TaskLocal)

	// Add the given trust prefix
	trustPrefix, trustCmd := task.Config["trust_prefix"]
	if trustCmd {
		var outBuf, errBuf bytes.Buffer
		cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trustPrefix))
		cmd.Stdout = &outBuf
		cmd.Stderr = &errBuf
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
				err, outBuf.String(), errBuf.String())
		}
		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trustPrefix)
	}

	// Build the command.
	var cmdArgs []string

	// Inject the environment variables.
	envVars := TaskEnvironmentVariables(ctx, task)

	// Clear the task directories as they are not currently supported.
	envVars.ClearTaskLocalDir()
	envVars.ClearAllocDir()

	for k, v := range envVars.Map() {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--set-env=%v=%v", k, v))
	}

	// Disble signature verification if the trust command was not run.
	if !trustCmd {
		cmdArgs = append(cmdArgs, "--insecure-skip-verify")
	}

	// Append the run command.
	cmdArgs = append(cmdArgs, "run", "--mds-register=false", img)

	// Check if the user has overriden the exec command.
	if execCmd, ok := task.Config["command"]; ok {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--exec=%v", execCmd))
	}

	if task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Memory limit cannot be zero")
	}
	if task.Resources.CPU == 0 {
		return nil, fmt.Errorf("CPU limit cannot be zero")
	}

	// Add memory isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*bytesToMB))

	// Add CPU isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))

	// Add user passed arguments.
	if len(driverConfig.Args) != 0 {
		parsed := args.ParseAndReplace(driverConfig.Args, envVars.Map())

		// Need to start arguments with "--"
		if len(parsed) > 0 {
			cmdArgs = append(cmdArgs, "--")
		}

		for _, arg := range parsed {
			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", arg))
		}
	}

	// Create files to capture stdin and out.
	stdoutFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stdout", taskName))
	stderrFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stderr", taskName))

	stdo, err := os.OpenFile(stdoutFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stdout: %v", err)
	}

	stde, err := os.OpenFile(stderrFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stderr: %v", err)
	}

	cmd := exec.Command("rkt", cmdArgs...)
	cmd.Stdout = stdo
	cmd.Stderr = stde

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("Error running rkt: %v", err)
	}

	d.logger.Printf("[DEBUG] driver.rkt: started ACI %q with: %v", img, cmd.Args)
	h := &rktHandle{
		proc:   cmd.Process,
		image:  img,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	qpid := &rktPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &rktHandle{
		proc:   proc,
		image:  qpid.Image,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		Pid:   h.proc.Pid,
		Image: h.image,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *rktHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	code := 0
	if !ps.Success() {
		// TODO: Better exit code parsing.
		code = 1
	}
	h.waitCh <- cstructs.NewWaitResult(code, 0, err)
	close(h.waitCh)
}
<commit_msg>Consolidate if else conditions<commit_after>package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/go-version"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/helper/args"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reRktVersion  = regexp.MustCompile(`rkt version (\d[.\d]+)`)
	reAppcVersion = regexp.MustCompile(`appc version (\d[.\d]+)`)
)

const (
	// minRktVersion is the earliest supported version of rkt. rkt added support
	// for CPU and memory isolators in 0.14.0. We cannot support an earlier
	// version to maintain an uniform interface across all drivers
	minRktVersion = "0.14.0"

	// bytesToMB is the conversion from bytes to megabytes.
	bytesToMB = 1024 * 1024
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type RktDriverConfig struct {
	ImageName string   `mapstructure:"image"`
	Args      []string `mapstructure:"args"`
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	proc   *os.Process
	image  string
	logger *log.Logger
	waitCh chan *cstructs.WaitResult
	doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	Pid   int
	Image string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{DriverContext: *ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reAppcVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "1"
	node.Attributes["driver.rkt.version"] = rktMatches[1]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	minVersion, _ := version.NewVersion(minRktVersion)
	currentVersion, _ := version.NewVersion(node.Attributes["driver.rkt.version"])
	if currentVersion.LessThan(minVersion) {
		// Do not allow rkt < 0.14.0
		d.logger.Printf("[WARN] driver.rkt: please upgrade rkt to a version >= %s", minVersion)
		node.Attributes["driver.rkt"] = "0"
	}
	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig RktDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}
	// Validate that the config is valid.
	img := driverConfig.ImageName
	if img == "" {
		return nil, fmt.Errorf("Missing ACI image for rkt")
	}

	// Get the tasks local directory.
	taskName := d.DriverContext.taskName
	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}
	taskLocal := filepath.Join(taskDir, allocdir.TaskLocal)

	// Build the command.
	var cmdArgs []string

	// Add the given trust prefix
	trustPrefix, trustCmd := task.Config["trust_prefix"]
	if trustCmd {
		var outBuf, errBuf bytes.Buffer
		cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trustPrefix))
		cmd.Stdout = &outBuf
		cmd.Stderr = &errBuf
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
				err, outBuf.String(), errBuf.String())
		}
		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trustPrefix)
	} else {
		// Disble signature verification if the trust command was not run.
		cmdArgs = append(cmdArgs, "--insecure-skip-verify")
	}

	// Inject the environment variables.
	envVars := TaskEnvironmentVariables(ctx, task)

	// Clear the task directories as they are not currently supported.
	envVars.ClearTaskLocalDir()
	envVars.ClearAllocDir()

	for k, v := range envVars.Map() {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--set-env=%v=%v", k, v))
	}

	// Append the run command.
	cmdArgs = append(cmdArgs, "run", "--mds-register=false", img)

	// Check if the user has overriden the exec command.
	if execCmd, ok := task.Config["command"]; ok {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--exec=%v", execCmd))
	}

	if task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Memory limit cannot be zero")
	}
	if task.Resources.CPU == 0 {
		return nil, fmt.Errorf("CPU limit cannot be zero")
	}

	// Add memory isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*bytesToMB))

	// Add CPU isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))

	// Add user passed arguments.
	if len(driverConfig.Args) != 0 {
		parsed := args.ParseAndReplace(driverConfig.Args, envVars.Map())

		// Need to start arguments with "--"
		if len(parsed) > 0 {
			cmdArgs = append(cmdArgs, "--")
		}

		for _, arg := range parsed {
			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", arg))
		}
	}

	// Create files to capture stdin and out.
	stdoutFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stdout", taskName))
	stderrFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stderr", taskName))

	stdo, err := os.OpenFile(stdoutFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stdout: %v", err)
	}

	stde, err := os.OpenFile(stderrFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stderr: %v", err)
	}

	cmd := exec.Command("rkt", cmdArgs...)
	cmd.Stdout = stdo
	cmd.Stderr = stde

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("Error running rkt: %v", err)
	}

	d.logger.Printf("[DEBUG] driver.rkt: started ACI %q with: %v", img, cmd.Args)
	h := &rktHandle{
		proc:   cmd.Process,
		image:  img,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	qpid := &rktPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &rktHandle{
		proc:   proc,
		image:  qpid.Image,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		Pid:   h.proc.Pid,
		Image: h.image,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *rktHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	code := 0
	if !ps.Success() {
		// TODO: Better exit code parsing.
		code = 1
	}
	h.waitCh <- cstructs.NewWaitResult(code, 0, err)
	close(h.waitCh)
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package client

import (
	"errors"
	"math"
	"sync"
	"time"

	"github.com/m3db/m3db/clock"
	"github.com/m3db/m3db/generated/thrift/rpc"
	"github.com/m3db/m3db/topology"

	"github.com/uber/tchannel-go/thrift"
)

var (
	errQueueNotOpen          = errors.New("host operation queue not open")
	errQueueUnknownOperation = errors.New("host operation queue received unknown operation")
	errQueueFetchNoResponse  = errors.New("host operation queue did not receive response for given fetch")
)

type queue struct {
	sync.RWMutex

	opts                  Options
	nowFn                 clock.NowFn
	host                  topology.Host
	connPool              connectionPool
	writeBatchRequestPool writeBatchRequestPool
	idDatapointArrayPool  idDatapointArrayPool
	size                  int
	ops                   []op
	opsSumSize            int
	opsLastRotatedAt      time.Time
	opsArrayPool          opArrayPool
	drainIn               chan []op
	state                 state
}

func newHostQueue(
	host topology.Host,
	writeBatchRequestPool writeBatchRequestPool,
	idDatapointArrayPool idDatapointArrayPool,
	opts Options,
) hostQueue {
	size := opts.HostQueueOpsFlushSize()

	opArrayPoolCapacity := int(math.Max(float64(size), float64(opts.WriteBatchSize())))
	opArrayPool := newOpArrayPool(opts.HostQueueOpsArrayPoolSize(), opArrayPoolCapacity)
	opArrayPool.Init()

	return &queue{
		opts:                  opts,
		nowFn:                 opts.ClockOptions().NowFn(),
		host:                  host,
		connPool:              newConnectionPool(host, opts),
		writeBatchRequestPool: writeBatchRequestPool,
		idDatapointArrayPool:  idDatapointArrayPool,
		size:                  size,
		ops:                   opArrayPool.Get(),
		opsArrayPool:          opArrayPool,
		// NB(r): specifically use non-buffered queue for single flush at a time
		drainIn: make(chan []op),
	}
}

func (q *queue) Open() {
	q.Lock()
	defer q.Unlock()

	if q.state != stateNotOpen {
		return
	}

	q.state = stateOpen

	// Open the connection pool
	q.connPool.Open()

	// Continually drain the queue until closed
	go q.drain()

	flushInterval := q.opts.HostQueueOpsFlushInterval()
	if flushInterval > 0 {
		// Continually flush the queue at given interval if set
		go q.flushEvery(flushInterval)
	}
}

func (q *queue) flushEvery(interval time.Duration) {
	// sleepForOverride used change the next sleep based on last ops rotation
	var sleepForOverride time.Duration
	for {
		sleepFor := interval
		if sleepForOverride > 0 {
			sleepFor = sleepForOverride
			sleepForOverride = 0
		}

		time.Sleep(sleepFor)

		q.RLock()
		if q.state != stateOpen {
			q.RUnlock()
			return
		}
		lastRotateAt := q.opsLastRotatedAt
		q.RUnlock()

		sinceLastRotate := q.nowFn().Sub(lastRotateAt)
		if sinceLastRotate < interval {
			// Rotated already recently, sleep until we would next consider flushing
			sleepForOverride = interval - sinceLastRotate
			continue
		}

		q.Lock()
		needsDrain := q.rotateOpsWithLock()
		q.Unlock()

		if len(needsDrain) != 0 {
			q.RLock()
			if q.state != stateOpen {
				q.RUnlock()
				return
			}
			// Need to hold lock while writing to the drainIn
			// channel to ensure it has not been closed
			q.drainIn <- needsDrain
			q.RUnlock()
		}
	}
}

func (q *queue) rotateOpsWithLock() []op {
	if q.opsSumSize == 0 {
		// No need to rotate as queue is empty
		return nil
	}

	needsDrain := q.ops

	// Reset ops
	q.ops = q.opsArrayPool.Get()
	q.opsSumSize = 0
	q.opsLastRotatedAt = q.nowFn()

	return needsDrain
}

func (q *queue) drain() {
	var (
		wgAll                       = &sync.WaitGroup{}
		currWriteOpsByNamespace     = make(map[string][]op)
		currIDDatapointsByNamespace = make(map[string][]*rpc.IDDatapoint)
		writeBatchSize              = q.opts.WriteBatchSize()
	)

	for {
		ops, ok := <-q.drainIn
		if !ok {
			break
		}
		var (
			currWriteOps     []op
			currIDDatapoints []*rpc.IDDatapoint
			opsLen           = len(ops)
		)
		for i := 0; i < opsLen; i++ {
			switch v := ops[i].(type) {
			case *writeOp:
				namespace := v.request.NameSpace
				currWriteOps = currWriteOpsByNamespace[namespace]
				currIDDatapoints = currIDDatapointsByNamespace[namespace]
				if currWriteOps == nil {
					currWriteOps = q.opsArrayPool.Get()
					currIDDatapoints = q.idDatapointArrayPool.Get()
				}

				currWriteOps = append(currWriteOps, ops[i])
				currIDDatapoints = append(currIDDatapoints, v.request.IdDatapoint)
				currWriteOpsByNamespace[namespace] = currWriteOps
				currIDDatapointsByNamespace[namespace] = currIDDatapoints

				if len(currWriteOps) == writeBatchSize {
					// Reached write batch limit, write async and reset
					q.asyncWrite(wgAll, namespace, currWriteOps, currIDDatapoints)
					currWriteOpsByNamespace[namespace] = nil
					currIDDatapointsByNamespace[namespace] = nil
				}
			case *fetchBatchOp:
				q.asyncFetch(wgAll, v)
			case *truncateOp:
				q.asyncTruncate(wgAll, v)
			default:
				completionFn := ops[i].CompletionFn()
				completionFn(nil, errQueueUnknownOperation)
			}
		}

		// If any outstanding write ops, async write
		for namespace, writeOps := range currWriteOpsByNamespace {
			if len(writeOps) > 0 {
				q.asyncWrite(wgAll, namespace, writeOps, currIDDatapointsByNamespace[namespace])
				currWriteOpsByNamespace[namespace] = nil
				currIDDatapointsByNamespace[namespace] = nil
			}
		}

		if ops != nil {
			q.opsArrayPool.Put(ops)
		}
	}

	// Close the connection pool after all requests done
	wgAll.Wait()
	q.connPool.Close()
}

func (q *queue) asyncWrite(wg *sync.WaitGroup, namespace string, ops []op, elems []*rpc.IDDatapoint) {
	wg.Add(1)
	// TODO(r): Use a worker pool to avoid creating new go routines for async writes
	go func() {
		req := q.writeBatchRequestPool.Get()
		req.NameSpace = namespace
		req.Elements = elems

		// NB(r): Defer is slow in the hot path unfortunately
		cleanup := func() {
			q.writeBatchRequestPool.Put(req)
			q.idDatapointArrayPool.Put(elems)
			q.opsArrayPool.Put(ops)
			wg.Done()
		}

		client, err := q.connPool.NextClient()
		if err != nil {
			// No client available
			callAllCompletionFns(ops, nil, err)
			cleanup()
			return
		}

		ctx, _ := thrift.NewContext(q.opts.WriteRequestTimeout())
		err = client.WriteBatch(ctx, req)
		if err == nil {
			// All succeeded
			callAllCompletionFns(ops, nil, nil)
			cleanup()
			return
		}

		if batchErrs, ok := err.(*rpc.WriteBatchErrors); ok {
			// Callback all writes with errors
			hasErr := make(map[int]struct{})
			for _, batchErr := range batchErrs.Errors {
				op := ops[batchErr.Index]
				op.CompletionFn()(nil, batchErr.Err)
				hasErr[int(batchErr.Index)] = struct{}{}
			}
			// Callback all writes with no errors
			for i := range ops {
				if _, ok := hasErr[i]; !ok {
					// No error
					ops[i].CompletionFn()(nil, nil)
				}
			}
			cleanup()
			return
		}

		// Entire batch failed
		callAllCompletionFns(ops, nil, err)
		cleanup()
	}()
}

func (q *queue) asyncFetch(wg *sync.WaitGroup, op *fetchBatchOp) {
	wg.Add(1)
	// TODO(r): Use a worker pool to avoid creating new go routines for async fetches
	go func() {
		// NB(r): Defer is slow in the hot path unfortunately
		cleanup := wg.Done

		client, err := q.connPool.NextClient()
		if err != nil {
			// No client available
			op.completeAll(nil, err)
			cleanup()
			return
		}

		ctx, _ := thrift.NewContext(q.opts.FetchRequestTimeout())
		result, err := client.FetchRawBatch(ctx, &op.request)
		if err != nil {
			op.completeAll(nil, err)
			cleanup()
			return
		}

		resultLen := len(result.Elements)
		for i := 0; i < op.Size(); i++ {
			if !(i < resultLen) {
				// No results for this entry, in practice should never occur
				op.complete(i, nil, errQueueFetchNoResponse)
				continue
			}
			if result.Elements[i].Err != nil {
				op.complete(i, nil, result.Elements[i].Err)
				continue
			}
			op.complete(i, result.Elements[i].Segments, nil)
		}
		cleanup()
	}()
}

func (q *queue) asyncTruncate(wg *sync.WaitGroup, op *truncateOp) {
	wg.Add(1)

	go func() {
		client, err := q.connPool.NextClient()
		if err != nil {
			// No client available
			op.completionFn(nil, err)
			wg.Done()
			return
		}

		ctx, _ := thrift.NewContext(q.opts.TruncateRequestTimeout())
		if res, err := client.Truncate(ctx, &op.request); err != nil {
			op.completionFn(nil, err)
		} else {
			op.completionFn(res, nil)
		}

		wg.Done()
	}()
}

func (q *queue) Len() int {
	q.RLock()
	v := q.opsSumSize
	q.RUnlock()
	return v
}

func (q *queue) Enqueue(o op) error {
	var needsDrain []op
	q.Lock()
	if q.state != stateOpen {
		q.Unlock()
		return errQueueNotOpen
	}
	q.ops = append(q.ops, o)
	q.opsSumSize += o.Size()
	// If queue is full flush
	if q.opsSumSize >= q.size {
		needsDrain = q.rotateOpsWithLock()
	}
	// Need to hold lock while writing to the drainIn
	// channel to ensure it has not been closed
	if len(needsDrain) != 0 {
		q.drainIn <- needsDrain
	}
	q.Unlock()
	return nil
}

func (q *queue) Host() topology.Host {
	return q.host
}

func (q *queue) ConnectionCount() int {
	return q.connPool.ConnectionCount()
}

func (q *queue) ConnectionPool() connectionPool {
	return q.connPool
}

func (q *queue) Close() {
	q.Lock()
	if q.state != stateOpen {
		q.Unlock()
		return
	}

	q.state = stateClosed
	// Closed drainIn channel in lock to ensure writers know
	// consistently if channel is open or not by checking state
	close(q.drainIn)
	q.Unlock()
}
<commit_msg>Make host queue drain channel buffered to match size of op arrays pool (#89)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package client

import (
	"errors"
	"math"
	"sync"
	"time"

	"github.com/m3db/m3db/clock"
	"github.com/m3db/m3db/generated/thrift/rpc"
	"github.com/m3db/m3db/topology"

	"github.com/uber/tchannel-go/thrift"
)

var (
	errQueueNotOpen          = errors.New("host operation queue not open")
	errQueueUnknownOperation = errors.New("host operation queue received unknown operation")
	errQueueFetchNoResponse  = errors.New("host operation queue did not receive response for given fetch")
)

type queue struct {
	sync.RWMutex

	opts                  Options
	nowFn                 clock.NowFn
	host                  topology.Host
	connPool              connectionPool
	writeBatchRequestPool writeBatchRequestPool
	idDatapointArrayPool  idDatapointArrayPool
	size                  int
	ops                   []op
	opsSumSize            int
	opsLastRotatedAt      time.Time
	opsArrayPool          opArrayPool
	drainIn               chan []op
	state                 state
}

func newHostQueue(
	host topology.Host,
	writeBatchRequestPool writeBatchRequestPool,
	idDatapointArrayPool idDatapointArrayPool,
	opts Options,
) hostQueue {
	size := opts.HostQueueOpsFlushSize()

	opsArraysLen := opts.HostQueueOpsArrayPoolSize()
	opArrayPoolCapacity := int(math.Max(float64(size), float64(opts.WriteBatchSize())))
	opArrayPool := newOpArrayPool(opsArraysLen, opArrayPoolCapacity)
	opArrayPool.Init()

	return &queue{
		opts:                  opts,
		nowFn:                 opts.ClockOptions().NowFn(),
		host:                  host,
		connPool:              newConnectionPool(host, opts),
		writeBatchRequestPool: writeBatchRequestPool,
		idDatapointArrayPool:  idDatapointArrayPool,
		size:                  size,
		ops:                   opArrayPool.Get(),
		opsArrayPool:          opArrayPool,
		drainIn:               make(chan []op, opsArraysLen),
	}
}

func (q *queue) Open() {
	q.Lock()
	defer q.Unlock()

	if q.state != stateNotOpen {
		return
	}

	q.state = stateOpen

	// Open the connection pool
	q.connPool.Open()

	// Continually drain the queue until closed
	go q.drain()

	flushInterval := q.opts.HostQueueOpsFlushInterval()
	if flushInterval > 0 {
		// Continually flush the queue at given interval if set
		go q.flushEvery(flushInterval)
	}
}

func (q *queue) flushEvery(interval time.Duration) {
	// sleepForOverride used change the next sleep based on last ops rotation
	var sleepForOverride time.Duration
	for {
		sleepFor := interval
		if sleepForOverride > 0 {
			sleepFor = sleepForOverride
			sleepForOverride = 0
		}

		time.Sleep(sleepFor)

		q.RLock()
		if q.state != stateOpen {
			q.RUnlock()
			return
		}
		lastRotateAt := q.opsLastRotatedAt
		q.RUnlock()

		sinceLastRotate := q.nowFn().Sub(lastRotateAt)
		if sinceLastRotate < interval {
			// Rotated already recently, sleep until we would next consider flushing
			sleepForOverride = interval - sinceLastRotate
			continue
		}

		q.Lock()
		needsDrain := q.rotateOpsWithLock()
		q.Unlock()

		if len(needsDrain) != 0 {
			q.RLock()
			if q.state != stateOpen {
				q.RUnlock()
				return
			}
			// Need to hold lock while writing to the drainIn
			// channel to ensure it has not been closed
			q.drainIn <- needsDrain
			q.RUnlock()
		}
	}
}

func (q *queue) rotateOpsWithLock() []op {
	if q.opsSumSize == 0 {
		// No need to rotate as queue is empty
		return nil
	}

	needsDrain := q.ops

	// Reset ops
	q.ops = q.opsArrayPool.Get()
	q.opsSumSize = 0
	q.opsLastRotatedAt = q.nowFn()

	return needsDrain
}

func (q *queue) drain() {
	var (
		wgAll                       = &sync.WaitGroup{}
		currWriteOpsByNamespace     = make(map[string][]op)
		currIDDatapointsByNamespace = make(map[string][]*rpc.IDDatapoint)
		writeBatchSize              = q.opts.WriteBatchSize()
	)

	for {
		ops, ok := <-q.drainIn
		if !ok {
			break
		}
		var (
			currWriteOps     []op
			currIDDatapoints []*rpc.IDDatapoint
			opsLen           = len(ops)
		)
		for i := 0; i < opsLen; i++ {
			switch v := ops[i].(type) {
			case *writeOp:
				namespace := v.request.NameSpace
				currWriteOps = currWriteOpsByNamespace[namespace]
				currIDDatapoints = currIDDatapointsByNamespace[namespace]
				if currWriteOps == nil {
					currWriteOps = q.opsArrayPool.Get()
					currIDDatapoints = q.idDatapointArrayPool.Get()
				}

				currWriteOps = append(currWriteOps, ops[i])
				currIDDatapoints = append(currIDDatapoints, v.request.IdDatapoint)
				currWriteOpsByNamespace[namespace] = currWriteOps
				currIDDatapointsByNamespace[namespace] = currIDDatapoints

				if len(currWriteOps) == writeBatchSize {
					// Reached write batch limit, write async and reset
					q.asyncWrite(wgAll, namespace, currWriteOps, currIDDatapoints)
					currWriteOpsByNamespace[namespace] = nil
					currIDDatapointsByNamespace[namespace] = nil
				}
			case *fetchBatchOp:
				q.asyncFetch(wgAll, v)
			case *truncateOp:
				q.asyncTruncate(wgAll, v)
			default:
				completionFn := ops[i].CompletionFn()
				completionFn(nil, errQueueUnknownOperation)
			}
		}

		// If any outstanding write ops, async write
		for namespace, writeOps := range currWriteOpsByNamespace {
			if len(writeOps) > 0 {
				q.asyncWrite(wgAll, namespace, writeOps, currIDDatapointsByNamespace[namespace])
				currWriteOpsByNamespace[namespace] = nil
				currIDDatapointsByNamespace[namespace] = nil
			}
		}

		if ops != nil {
			q.opsArrayPool.Put(ops)
		}
	}

	// Close the connection pool after all requests done
	wgAll.Wait()
	q.connPool.Close()
}

func (q *queue) asyncWrite(wg *sync.WaitGroup, namespace string, ops []op, elems []*rpc.IDDatapoint) {
	wg.Add(1)
	// TODO(r): Use a worker pool to avoid creating new go routines for async writes
	go func() {
		req := q.writeBatchRequestPool.Get()
		req.NameSpace = namespace
		req.Elements = elems

		// NB(r): Defer is slow in the hot path unfortunately
		cleanup := func() {
			q.writeBatchRequestPool.Put(req)
			q.idDatapointArrayPool.Put(elems)
			q.opsArrayPool.Put(ops)
			wg.Done()
		}

		client, err := q.connPool.NextClient()
		if err != nil {
			// No client available
			callAllCompletionFns(ops, nil, err)
			cleanup()
			return
		}

		ctx, _ := thrift.NewContext(q.opts.WriteRequestTimeout())
		err = client.WriteBatch(ctx, req)
		if err == nil {
			// All succeeded
			callAllCompletionFns(ops, nil, nil)
			cleanup()
			return
		}

		if batchErrs, ok := err.(*rpc.WriteBatchErrors); ok {
			// Callback all writes with errors
			hasErr := make(map[int]struct{})
			for _, batchErr := range batchErrs.Errors {
				op := ops[batchErr.Index]
				op.CompletionFn()(nil, batchErr.Err)
				hasErr[int(batchErr.Index)] = struct{}{}
			}
			// Callback all writes with no errors
			for i := range ops {
				if _, ok := hasErr[i]; !ok {
					// No error
					ops[i].CompletionFn()(nil, nil)
				}
			}
			cleanup()
			return
		}

		// Entire batch failed
		callAllCompletionFns(ops, nil, err)
		cleanup()
	}()
}

func (q *queue) asyncFetch(wg *sync.WaitGroup, op *fetchBatchOp) {
	wg.Add(1)
	// TODO(r): Use a worker pool to avoid creating new go routines for async fetches
	go func() {
		// NB(r): Defer is slow in the hot path unfortunately
		cleanup := wg.Done

		client, err := q.connPool.NextClient()
		if err != nil {
			// No client available
			op.completeAll(nil, err)
			cleanup()
			return
		}

		ctx, _ := thrift.NewContext(q.opts.FetchRequestTimeout())
		result, err := client.FetchRawBatch(ctx, &op.request)
		if err != nil {
			op.completeAll(nil, err)
			cleanup()
			return
		}

		resultLen := len(result.Elements)
		for i := 0; i < op.Size(); i++ {
			if !(i < resultLen) {
				// No results for this entry, in practice should never occur
				op.complete(i, nil, errQueueFetchNoResponse)
				continue
			}
			if result.Elements[i].Err != nil {
				op.complete(i, nil, result.Elements[i].Err)
				continue
			}
			op.complete(i, result.Elements[i].Segments, nil)
		}
		cleanup()
	}()
}

func (q *queue) asyncTruncate(wg *sync.WaitGroup, op *truncateOp) {
	wg.Add(1)

	go func() {
		client, err := q.connPool.NextClient()
		if err != nil {
			// No client available
			op.completionFn(nil, err)
			wg.Done()
			return
		}

		ctx, _ := thrift.NewContext(q.opts.TruncateRequestTimeout())
		if res, err := client.Truncate(ctx, &op.request); err != nil {
			op.completionFn(nil, err)
		} else {
			op.completionFn(res, nil)
		}

		wg.Done()
	}()
}

func (q *queue) Len() int {
	q.RLock()
	v := q.opsSumSize
	q.RUnlock()
	return v
}

func (q *queue) Enqueue(o op) error {
	var needsDrain []op
	q.Lock()
	if q.state != stateOpen {
		q.Unlock()
		return errQueueNotOpen
	}
	q.ops = append(q.ops, o)
	q.opsSumSize += o.Size()
	// If queue is full flush
	if q.opsSumSize >= q.size {
		needsDrain = q.rotateOpsWithLock()
	}
	// Need to hold lock while writing to the drainIn
	// channel to ensure it has not been closed
	if len(needsDrain) != 0 {
		q.drainIn <- needsDrain
	}
	q.Unlock()
	return nil
}

func (q *queue) Host() topology.Host {
	return q.host
}

func (q *queue) ConnectionCount() int {
	return q.connPool.ConnectionCount()
}

func (q *queue) ConnectionPool() connectionPool {
	return q.connPool
}

func (q *queue) Close() {
	q.Lock()
	if q.state != stateOpen {
		q.Unlock()
		return
	}

	q.state = stateClosed
	// Closed drainIn channel in lock to ensure writers know
	// consistently if channel is open or not by checking state
	close(q.drainIn)
	q.Unlock()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"sync/atomic"

	"github.com/kataras/iris"
	"github.com/kataras/iris/mvc"
)

func main() {
	app := iris.New()
	mvc.New(app.Party("/")).Register(&globalVisitorsController{visits: 0})

	// http://localhost:8080
	app.Run(iris.Addr(":8080"))
}

type globalVisitorsController struct {
	// When a singleton controller is used then concurent safe access is up to the developers, because
	// all clients share the same controller instance instead.
	// Note that any controller's methods
	// are per-client, but the struct's field can be shared across multiple clients if the structure
	// does not have any dynamic struct field depenendies that depend on the iris.Context
	// and ALL field's values are NOT zero, at this case we use uint64 which it's no zero (even if we didn't set it
	// manually ease-of-understand reasons) because it's a value of &{0}.
	// All the above declares a Singleton, note that you don't have to write a single line of code to do this, Iris is smart enough.
	//
	// see `Get`.
	visits uint64
}

func (c *globalVisitorsController) Get() string {
	count := atomic.AddUint64(&c.visits, 1)
	return fmt.Sprintf("Total visitors: %d", count)
}
<commit_msg>fix typo<commit_after>package main

import (
	"fmt"
	"sync/atomic"

	"github.com/kataras/iris"
	"github.com/kataras/iris/mvc"
)

func main() {
	app := iris.New()
	mvc.New(app.Party("/")).Register(&globalVisitorsController{visits: 0})

	// http://localhost:8080
	app.Run(iris.Addr(":8080"))
}

type globalVisitorsController struct {
	// When a singleton controller is used then concurent safe access is up to the developers, because
	// all clients share the same controller instance instead.
	// Note that any controller's methods
	// are per-client, but the struct's field can be shared across multiple clients if the structure
	// does not have any dynamic struct field dependencies that depend on the iris.Context
	// and ALL field's values are NOT zero, at this case we use uint64 which it's no zero (even if we didn't set it
	// manually ease-of-understand reasons) because it's a value of &{0}.
	// All the above declares a Singleton, note that you don't have to write a single line of code to do this, Iris is smart enough.
	//
	// see `Get`.
	visits uint64
}

func (c *globalVisitorsController) Get() string {
	count := atomic.AddUint64(&c.visits, 1)
	return fmt.Sprintf("Total visitors: %d", count)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"github.com/larryli/vagrantcloud.v1"
	"log"
	"strings"
)

type Arch struct {
	name, info, box string
}

type Release string

type Box vagrantcloud.Box

type Version vagrantcloud.Version

const (
	url      = "https://cloud-images.ubuntu.com/vagrant/"
	notFound = "404 Not Found"
	see      = "\n\nSee https://github.com/larryli/vagrantcloud.v1/tree/master/update-ubuntu-vagrant-box"
)

var (
	api      *vagrantcloud.Api
	username = flag.String("username", "larryli", "username")
	token    = flag.String("token", "", "access_token")
	test     = flag.Bool("test", false, "test, no effect")
	arches   = []Arch{
		{
			name: "64",
			info: "amd64",
			box:  "-server-cloudimg-amd64-vagrant-disk1.box",
		},
		{
			name: "64juju",
			info: "amd64 with juju",
			box:  "-server-cloudimg-amd64-juju-vagrant-disk1.box",
		},
		{
			name: "32",
			info: "i386",
			box:  "-server-cloudimg-i386-vagrant-disk1.box",
		},
		{
			name: "32juju",
			info: "i386 with juju",
			box:  "-server-cloudimg-i386-juju-vagrant-disk1.box",
		},
	}
)

func Fatal(err error, a ...interface{}) {
	if err != nil {
		a = append(a, err)
		log.Fatalln(a...)
	}
}

func isChildren(s string) bool {
	return !strings.HasPrefix(s, "/") && !strings.HasPrefix(s, "?")
}

func fetchReleases() (releases []Release) {
	doc, err := goquery.NewDocument(url)
	Fatal(err, "fetch "+url)
	doc.Find("a").Each(func(i int, s *goquery.Selection) {
		if link, ok := s.Attr("href"); ok {
			if isChildren(link) {
				releases = append(releases, Release(strings.Trim(link, "/")))
			}
		}
	})
	return
}

func (r Release) name() string {
	return string(r)
}

func (r Release) url() string {
	return url + r.name() + "/"
}

func (r Release) fetch() (versions []string) {
	doc, err := goquery.NewDocument(r.url())
	Fatal(err, "fetch "+r.url())
	doc.Find("a").Each(func(i int, s *goquery.Selection) {
		if link, ok := s.Attr("href"); ok {
			if isChildren(link) && link != "current/" {
				versions = append(versions, strings.Trim(link, "/"))
			}
		}
	})
	return
}

func (r Release) scan() {
	versions := r.fetch()
	for _, t := range arches {
		box := api.Box(*username, r.name()+t.name)
		todo := fmt.Sprintf("fetch \"%s\"", box.Uri())
		err := box.Get()
		if err != nil && strings.EqualFold(err.Error(), notFound) {
			box.ShortDescription = "Ubuntu " + strings.Title(r.name()) + " " + t.info
			box.DescriptionMarkdown = r.url() + see
			todo = fmt.Sprintf("add \"%s\"", box.Uri())
			if !(*test) {
				Fatal(box.New(), todo)
			}
			log.Println(todo)
		} else {
			Fatal(err, todo)
		}
		for _, version := range box.Versions {
			v := Version(version)
			if !v.exists(versions) {
				v.delete()
			}
		}
		for _, version := range versions {
			b := (*Box)(box)
			if !b.find(version) {
				image := r.url() + version + "/" + r.name() + t.box
				b.add(version, image)
			}
		}
	}
}

func (v *Version) exists(versions []string) (found bool) {
	for _, test := range versions {
		if test == v.Version {
			found = true
			return
		}
	}
	return
}

func (v *Version) delete() {
	version := (*vagrantcloud.Version)(v)
	todo := fmt.Sprintf("delete \"%s\" Version: \"%s\"", version.Uri(), version.Version)
	if !(*test) {
		Fatal(version.Delete(), todo)
	}
	log.Println(todo)
}

func (b *Box) find(version string) (found bool) {
	for _, test := range b.Versions {
		if test.Version == version {
			found = true
			return
		}
	}
	return
}

func (b *Box) add(version, image string) {
	box := (*vagrantcloud.Box)(b)
	v := box.Version(0)
	v.Version = version
	v.DescriptionMarkdown = image + see
	todo := fmt.Sprintf("add \"%s\" Version: \"%s\"", v.Uri(), v.Version)
	if !(*test) {
		Fatal(v.New(), todo)
	}
	p := v.Provider(vagrantcloud.ProviderVirtualbox)
	p.OriginalUrl = image
	todo = fmt.Sprintf("add \"%s\" Version: \"%s\"", p.Uri())
	if !(*test) {
		Fatal(p.New(), todo)
	}
	todo = fmt.Sprintf("public \"%s\" Version: \"%s\" Url: \"%s\"", v.Uri(), v.Version, p.OriginalUrl)
	if !(*test) {
		Fatal(v.Release(), todo)
	}
	log.Println(todo)
}

func main() {
	flag.Parse()
	api = vagrantcloud.New(*token)
	log.Println("start")
	for _, release := range fetchReleases() {
		release.scan()
	}
	log.Println("end")
}
<commit_msg>fix juju box name<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"github.com/larryli/vagrantcloud.v1"
	"log"
	"strings"
)

type Arch struct {
	name, info, box string
}

type Release string

type Box vagrantcloud.Box

type Version vagrantcloud.Version

const (
	url      = "https://cloud-images.ubuntu.com/vagrant/"
	notFound = "404 Not Found"
	see      = "\n\nSee https://github.com/larryli/vagrantcloud.v1/tree/master/update-ubuntu-vagrant-box"
)

var (
	api      *vagrantcloud.Api
	username = flag.String("username", "larryli", "username")
	token    = flag.String("token", "", "access_token")
	test     = flag.Bool("test", false, "test, no effect")
	arches   = []Arch{
		{
			name: "64",
			info: "amd64",
			box:  "-server-cloudimg-amd64-vagrant-disk1.box",
		},
		{
			name: "64-juju",
			info: "amd64 with juju",
			box:  "-server-cloudimg-amd64-juju-vagrant-disk1.box",
		},
		{
			name: "32",
			info: "i386",
			box:  "-server-cloudimg-i386-vagrant-disk1.box",
		},
		{
			name: "32-juju",
			info: "i386 with juju",
			box:  "-server-cloudimg-i386-juju-vagrant-disk1.box",
		},
	}
)

func Fatal(err error, a ...interface{}) {
	if err != nil {
		a = append(a, err)
		log.Fatalln(a...)
	}
}

func isChildren(s string) bool {
	return !strings.HasPrefix(s, "/") && !strings.HasPrefix(s, "?")
}

func fetchReleases() (releases []Release) {
	doc, err := goquery.NewDocument(url)
	Fatal(err, "fetch "+url)
	doc.Find("a").Each(func(i int, s *goquery.Selection) {
		if link, ok := s.Attr("href"); ok {
			if isChildren(link) {
				releases = append(releases, Release(strings.Trim(link, "/")))
			}
		}
	})
	return
}

func (r Release) name() string {
	return string(r)
}

func (r Release) url() string {
	return url + r.name() + "/"
}

func (r Release) fetch() (versions []string) {
	doc, err := goquery.NewDocument(r.url())
	Fatal(err, "fetch "+r.url())
	doc.Find("a").Each(func(i int, s *goquery.Selection) {
		if link, ok := s.Attr("href"); ok {
			if isChildren(link) && link != "current/" {
				versions = append(versions, strings.Trim(link, "/"))
			}
		}
	})
	return
}

func (r Release) scan() {
	versions := r.fetch()
	for _, t := range arches {
		box := api.Box(*username, r.name()+t.name)
		todo := fmt.Sprintf("fetch \"%s\"", box.Uri())
		err := box.Get()
		if err != nil && strings.EqualFold(err.Error(), notFound) {
			box.ShortDescription = "Ubuntu " + strings.Title(r.name()) + " " + t.info
			box.DescriptionMarkdown = r.url() + see
			todo = fmt.Sprintf("add \"%s\"", box.Uri())
			if !(*test) {
				Fatal(box.New(), todo)
			}
			log.Println(todo)
		} else {
			Fatal(err, todo)
		}
		for _, version := range box.Versions {
			v := Version(version)
			if !v.exists(versions) {
				v.delete()
			}
		}
		for _, version := range versions {
			b := (*Box)(box)
			if !b.find(version) {
				image := r.url() + version + "/" + r.name() + t.box
				b.add(version, image)
			}
		}
	}
}

func (v *Version) exists(versions []string) (found bool) {
	for _, test := range versions {
		if test == v.Version {
			found = true
			return
		}
	}
	return
}

func (v *Version) delete() {
	version := (*vagrantcloud.Version)(v)
	todo := fmt.Sprintf("delete \"%s\" Version: \"%s\"", version.Uri(), version.Version)
	if !(*test) {
		Fatal(version.Delete(), todo)
	}
	log.Println(todo)
}

func (b *Box) find(version string) (found bool) {
	for _, test := range b.Versions {
		if test.Version == version {
			found = true
			return
		}
	}
	return
}

func (b *Box) add(version, image string) {
	box := (*vagrantcloud.Box)(b)
	v := box.Version(0)
	v.Version = version
	v.DescriptionMarkdown = image + see
	todo := fmt.Sprintf("add \"%s\" Version: \"%s\"", v.Uri(), v.Version)
	if !(*test) {
		Fatal(v.New(), todo)
	}
	p := v.Provider(vagrantcloud.ProviderVirtualbox)
	p.OriginalUrl = image
	todo = fmt.Sprintf("add \"%s\" Version: \"%s\"", p.Uri())
	if !(*test) {
		Fatal(p.New(), todo)
	}
	todo = fmt.Sprintf("public \"%s\" Version: \"%s\" Url: \"%s\"", v.Uri(), v.Version, p.OriginalUrl)
	if !(*test) {
		Fatal(v.Release(), todo)
	}
	log.Println(todo)
}

func main() {
	flag.Parse()
	api = vagrantcloud.New(*token)
	log.Println("start")
	for _, release := range fetchReleases() {
		release.scan()
	}
	log.Println("end")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/ninjasphere/go-ninja/events"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/support"
	"github.com/ninjasphere/go-zigbee"
	"github.com/ninjasphere/go-zigbee/nwkmgr"
)

const (
	ClusterIDBasic    uint32 = 0x00
	ClusterIDOnOff    uint32 = 0x06
	ClusterIDLevel    uint32 = 0x08 // We're always exporting as brightness for now
	ClusterIDColor    uint32 = 0x300
	ClusterIDTemp     uint32 = 0x402
	ClusterIDHumidity uint32 = 0x405
	ClusterIDPower    uint32 = 0x702
)

type ZStackConfig struct {
	Hostname       string
	OtasrvrPort    int
	GatewayPort    int
	NwkmgrPort     int
	StableFlagFile string
}

type Driver struct {
	support.DriverSupport

	devices map[uint64]*Device

	config *ZStackConfig

	nwkmgrConn  *zigbee.ZStackNwkMgr
	gatewayConn *zigbee.ZStackGateway
	otaConn     *zigbee.ZStackOta

	devicesFound int
}

func NewDriver(config *ZStackConfig) (*Driver, error) {
	driver := &Driver{
		config:  config,
		devices: make(map[uint64]*Device),
	}

	err := driver.Init(info)
	if err != nil {
		log.Fatalf("Failed to initialize fake driver: %s", err)
	}

	err = driver.Export(driver)
	if err != nil {
		log.Fatalf("Failed to export fake driver: %s", err)
	}

	userAgent := driver.Conn.GetServiceClient("$device/:deviceId/channel/user-agent")
	userAgent.OnEvent("pairing-requested", func(pairingRequest *events.PairingRequest, values map[string]string) bool {
		log.Infof("Pairing request received from %s for %d seconds", values["deviceId"], pairingRequest.Duration)

		duration := uint32(pairingRequest.Duration)

		if err := driver.EnableJoin(duration); err != nil {
			log.Warningf("Failed to enable joining: %s", err)
		}

		return true
	})

	return driver, nil

}

func (d *Driver) Reset(hard bool) error {
	return d.nwkmgrConn.Reset(hard)
}

func (d *Driver) Start() error {

	waitUntilZStackReady(d.config.StableFlagFile)

	var err error

	d.nwkmgrConn, err = zigbee.ConnectToNwkMgrServer(d.config.Hostname, d.config.NwkmgrPort)
	if err != nil {
		return fmt.Errorf("Error connecting to nwkmgr %s", err)
	}
	d.nwkmgrConn.OnDeviceFound = func(deviceInfo *nwkmgr.NwkDeviceInfoT) {
		d.onDeviceFound(deviceInfo)
	}

	/*done := false
	d.nwkmgrConn.OnNetworkReady = func() {
		if !done {
			done = true
		}
	}*/

	d.otaConn, err = zigbee.ConnectToOtaServer(d.config.Hostname, d.config.OtasrvrPort)
	if err != nil {
		return fmt.Errorf("Error connecting to ota server %s", err)
	}

	d.gatewayConn, err = zigbee.ConnectToGatewayServer(d.config.Hostname, d.config.GatewayPort)
	if err != nil {
		return fmt.Errorf("Error connecting to gateway %s", err)
	}

	/*keyResponse := &nwkmgr.NwkZigbeeGenericCnf{}

	keyRequest := &nwkmgr.NwkChangeNwkKeyReq{
		NewKey: []byte{0x5a, 0x69, 0x67, 0x42, 0x65, 0x65, 0x41, 0x6c, 0x6c, 0x69, 0x61, 0x6e, 0x63, 0x65, 0x30, 0x39},
	}

	err = d.nwkmgrConn.SendCommand(keyRequest, keyResponse)
	if err != nil {
		return fmt.Errorf("Failed setting network key: %s", err)
	}

	spew.Dump(keyResponse)
	log.Println("Sleeping for 10 seconds after setting network key")

	time.Sleep(10 * time.Second)*/

	networkInfo := &nwkmgr.NwkZigbeeNwkInfoCnf{}

	err = d.nwkmgrConn.SendCommand(&nwkmgr.NwkZigbeeNwkInfoReq{}, networkInfo)
	if err != nil {
		if log.IsDebugEnabled() {
			spew.Dump(networkInfo)
		}
		return fmt.Errorf("Failed getting network info: %s", err)
	}

	if *networkInfo.Status == nwkmgr.NwkNetworkStatusT_NWK_DOWN {
		return fmt.Errorf("The ZigBee network is down")
	}

	localDevice := &nwkmgr.NwkGetLocalDeviceInfoCnf{}

	err = d.nwkmgrConn.SendCommand(&nwkmgr.NwkGetLocalDeviceInfoReq{}, localDevice)
	if err != nil {
		return fmt.Errorf("Failed getting local device info: %s", err)
	}

	if log.IsDebugEnabled() {
		spew.Dump("device info", localDevice.String())
	}

	networkKey := &nwkmgr.NwkGetNwkKeyCnf{}

	err = d.nwkmgrConn.SendCommand(&nwkmgr.NwkGetNwkKeyReq{}, networkKey)
	if err != nil {
		log.Fatalf("Failed getting network key: %s", err)
	}

	if log.IsDebugEnabled() {
		spew.Dump(networkKey)
	}

	log.Debugf("Started coordinator. Channel:%d Pan ID:0x%X Key:% X", *networkInfo.NwkChannel, *networkInfo.PanId, networkKey.NewKey)

	return d.FetchDevices()
}

func (d *Driver) EnableJoin(duration uint32) error {

	go func() {
	   	save := d.devicesFound;
		time.Sleep(time.Second * time.Duration(duration))
		d.Log.Infof("Join window closes after %d seconds.", duration);
		d.SendEvent("pairing-ended", &events.PairingEnded{
			DevicesFound: int(d.devicesFound - save),
		})
	}()

	permitJoinRequest := &nwkmgr.NwkSetPermitJoinReq{
		PermitJoinTime: &duration,
		PermitJoin:     nwkmgr.NwkPermitJoinTypeT_PERMIT_ALL.Enum(),
	}

	permitJoinResponse := &nwkmgr.NwkZigbeeGenericCnf{}

	err := d.nwkmgrConn.SendCommand(permitJoinRequest, permitJoinResponse)
	if err != nil {
		return fmt.Errorf("Failed to enable joining: %s", err)
	}
	if permitJoinResponse.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to enable joining: %s", permitJoinResponse.Status)
	}

	d.SendEvent("pairing-started", &events.PairingStarted{
		Duration: int(duration),
	})

	d.Log.Infof("Join window opens for %d seconds", duration);

	return nil
}

func (d *Driver) FetchDevices() error {
	return d.nwkmgrConn.FetchDeviceList()
}

func (d *Driver) onDeviceFound(deviceInfo *nwkmgr.NwkDeviceInfoT) {

	log.Debugf("\n\n")
	log.Debugf("---- Found Device IEEE:%X ----\f", *deviceInfo.IeeeAddress)
	log.Debugf("Device Info: %v", *deviceInfo)

	if d.devices[*deviceInfo.IeeeAddress] != nil {
		// We've seen this already, but it may have been repaired. We *should* just be able to replace
		// the deviceInfo object, which is used for all communication.
		// TODO: Actually verify this. May need to re-run channel init.
		d.devices[*deviceInfo.IeeeAddress].deviceInfo = deviceInfo
	}

	device := &Device{
		driver:     d,
		deviceInfo: deviceInfo,
		info: &model.Device{
			NaturalID:     fmt.Sprintf("%X", *deviceInfo.IeeeAddress),
			NaturalIDType: "zigbee",
			Signatures:    &map[string]string{},
		},
	}

	for _, endpoint := range deviceInfo.SimpleDescList {
		if *endpoint.ProfileId == 0xC05E /*ZLL*/ {

			switch *endpoint.DeviceId {

			case 0x0000: // On/Off Light
				fallthrough
			case 0x0100: // Dimmable Light
				fallthrough
			case 0x0200: // Color Light
				fallthrough
			case 0x210: // Ext Color Light
				(*device.info.Signatures)["ninja:thingType"] = "light"
			}
		}
	}

	err := device.getBasicInfo()
	if err != nil {
		log.Debugf("Failed to get basic info: %s", err)
	}

	if device.ManufacturerName != "" {
		(*device.info.Signatures)["zigbee:ManufacturerName"] = device.ManufacturerName
	} else {
		device.ManufacturerName = "Unknown"
	}
	if device.ModelIdentifier != "" {
		(*device.info.Signatures)["zigbee:ModelIdentifier"] = device.ModelIdentifier
	} else {
		device.ModelIdentifier = fmt.Sprintf("MAC:%X", *deviceInfo.IeeeAddress)
	}

	name := fmt.Sprintf("%s by %s", device.ModelIdentifier, device.ManufacturerName)

	device.info.Name = &name

	if log.IsDebugEnabled() {
		spew.Dump(deviceInfo)
	}

	err = d.Conn.ExportDevice(device)
	if err != nil {
		log.Fatalf("Failed to export zigbee device %s: %s", name, err)
	}
	d.devicesFound++;

	log.Debugf("Got device : %d", *deviceInfo.IeeeAddress)

	for _, endpoint := range deviceInfo.SimpleDescList {
		log.Debugf("Got endpoint : %d", *endpoint.EndpointId)

		if containsUInt32(endpoint.InputClusters, ClusterIDOnOff) {
			log.Debugf("This endpoint has on/off cluster")

			onOff := &OnOffChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDOnOff),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := onOff.init()
			if err != nil {
				log.Debugf("Failed initialising on/off channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDPower) {
			log.Debugf("This endpoint has power cluster")

			power := &PowerChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDPower),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := power.init()
			if err != nil {
				log.Debugf("Failed initialising power channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDTemp) {
			log.Debugf("This endpoint has temperature cluster")

			temp := &TempChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDTemp),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := temp.init()
			if err != nil {
				log.Debugf("Failed initialising temp channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDHumidity) {
			log.Debugf("This endpoint has humidity cluster")

			humidity := &HumidityChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDHumidity),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := humidity.init()
			if err != nil {
				log.Debugf("Failed initialising humidity channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDLevel) {
			log.Debugf("This endpoint has level cluster. Exporting as brightness channel")

			if log.IsDebugEnabled() {
				spew.Dump("brightness cluster", endpoint, ClusterIDLevel)
			}

			brightness := &BrightnessChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDLevel),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := brightness.init()
			if err != nil {
				log.Debugf("Failed initialising brightness channel: %s", err)
			}

		}

	}

	d.devices[*deviceInfo.IeeeAddress] = device

	fmt.Printf("---- Finished Device IEEE:%X ----\n", *deviceInfo.IeeeAddress)

}

func getCurDir() string {
	pwd, _ := os.Getwd()
	return pwd + "/"
}

func containsUInt32(hackstack []uint32, needle uint32) bool {
	for _, cluster := range hackstack {
		if cluster == needle {
			return true
		}
	}
	return false
}

func waitUntilZStackReady(checkFile string) {
	if checkFile == "" {
		return
	}

	// cooperate with zigbeeHAgw so that we don't start the zigbee driver
	// until we look somewhat stable.
	log.Debugf("waiting until zigbeeHAgw writes %s", checkFile)
	for {
		if _, err := os.Stat(checkFile); err == nil {
			break
		} else {
			time.Sleep(time.Second)
		}
	}
	log.Debugf("%s detected. start up continues...", checkFile)
}
<commit_msg>Put an upper bound on the pairing interval.<commit_after>package main

import (
	"fmt"
	"os"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/ninjasphere/go-ninja/events"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/support"
	"github.com/ninjasphere/go-zigbee"
	"github.com/ninjasphere/go-zigbee/nwkmgr"
)

const (
	ClusterIDBasic    uint32 = 0x00
	ClusterIDOnOff    uint32 = 0x06
	ClusterIDLevel    uint32 = 0x08 // We're always exporting as brightness for now
	ClusterIDColor    uint32 = 0x300
	ClusterIDTemp     uint32 = 0x402
	ClusterIDHumidity uint32 = 0x405
	ClusterIDPower    uint32 = 0x702
)

type ZStackConfig struct {
	Hostname       string
	OtasrvrPort    int
	GatewayPort    int
	NwkmgrPort     int
	StableFlagFile string
}

type Driver struct {
	support.DriverSupport

	devices map[uint64]*Device

	config *ZStackConfig

	nwkmgrConn  *zigbee.ZStackNwkMgr
	gatewayConn *zigbee.ZStackGateway
	otaConn     *zigbee.ZStackOta

	devicesFound int
}

func NewDriver(config *ZStackConfig) (*Driver, error) {
	driver := &Driver{
		config:  config,
		devices: make(map[uint64]*Device),
	}

	err := driver.Init(info)
	if err != nil {
		log.Fatalf("Failed to initialize fake driver: %s", err)
	}

	err = driver.Export(driver)
	if err != nil {
		log.Fatalf("Failed to export fake driver: %s", err)
	}

	userAgent := driver.Conn.GetServiceClient("$device/:deviceId/channel/user-agent")
	userAgent.OnEvent("pairing-requested", func(pairingRequest *events.PairingRequest, values map[string]string) bool {

		duration := uint32(pairingRequest.Duration)

		if duration > 254 {
		   duration = 254
		}

		log.Infof("Pairing request received from %s for %d seconds", values["deviceId"], duration)

		if err := driver.EnableJoin(duration); err != nil {
			log.Warningf("Failed to enable joining: %s", err)
		}

		return true
	})

	return driver, nil

}

func (d *Driver) Reset(hard bool) error {
	return d.nwkmgrConn.Reset(hard)
}

func (d *Driver) Start() error {

	waitUntilZStackReady(d.config.StableFlagFile)

	var err error

	d.nwkmgrConn, err = zigbee.ConnectToNwkMgrServer(d.config.Hostname, d.config.NwkmgrPort)
	if err != nil {
		return fmt.Errorf("Error connecting to nwkmgr %s", err)
	}
	d.nwkmgrConn.OnDeviceFound = func(deviceInfo *nwkmgr.NwkDeviceInfoT) {
		d.onDeviceFound(deviceInfo)
	}

	/*done := false
	d.nwkmgrConn.OnNetworkReady = func() {
		if !done {
			done = true
		}
	}*/

	d.otaConn, err = zigbee.ConnectToOtaServer(d.config.Hostname, d.config.OtasrvrPort)
	if err != nil {
		return fmt.Errorf("Error connecting to ota server %s", err)
	}

	d.gatewayConn, err = zigbee.ConnectToGatewayServer(d.config.Hostname, d.config.GatewayPort)
	if err != nil {
		return fmt.Errorf("Error connecting to gateway %s", err)
	}

	/*keyResponse := &nwkmgr.NwkZigbeeGenericCnf{}

	keyRequest := &nwkmgr.NwkChangeNwkKeyReq{
		NewKey: []byte{0x5a, 0x69, 0x67, 0x42, 0x65, 0x65, 0x41, 0x6c, 0x6c, 0x69, 0x61, 0x6e, 0x63, 0x65, 0x30, 0x39},
	}

	err = d.nwkmgrConn.SendCommand(keyRequest, keyResponse)
	if err != nil {
		return fmt.Errorf("Failed setting network key: %s", err)
	}

	spew.Dump(keyResponse)
	log.Println("Sleeping for 10 seconds after setting network key")

	time.Sleep(10 * time.Second)*/

	networkInfo := &nwkmgr.NwkZigbeeNwkInfoCnf{}

	err = d.nwkmgrConn.SendCommand(&nwkmgr.NwkZigbeeNwkInfoReq{}, networkInfo)
	if err != nil {
		if log.IsDebugEnabled() {
			spew.Dump(networkInfo)
		}
		return fmt.Errorf("Failed getting network info: %s", err)
	}

	if *networkInfo.Status == nwkmgr.NwkNetworkStatusT_NWK_DOWN {
		return fmt.Errorf("The ZigBee network is down")
	}

	localDevice := &nwkmgr.NwkGetLocalDeviceInfoCnf{}

	err = d.nwkmgrConn.SendCommand(&nwkmgr.NwkGetLocalDeviceInfoReq{}, localDevice)
	if err != nil {
		return fmt.Errorf("Failed getting local device info: %s", err)
	}

	if log.IsDebugEnabled() {
		spew.Dump("device info", localDevice.String())
	}

	networkKey := &nwkmgr.NwkGetNwkKeyCnf{}

	err = d.nwkmgrConn.SendCommand(&nwkmgr.NwkGetNwkKeyReq{}, networkKey)
	if err != nil {
		log.Fatalf("Failed getting network key: %s", err)
	}

	if log.IsDebugEnabled() {
		spew.Dump(networkKey)
	}

	log.Debugf("Started coordinator. Channel:%d Pan ID:0x%X Key:% X", *networkInfo.NwkChannel, *networkInfo.PanId, networkKey.NewKey)

	return d.FetchDevices()
}

func (d *Driver) EnableJoin(duration uint32) error {

	go func() {
	   	save := d.devicesFound;
		time.Sleep(time.Second * time.Duration(duration))
		d.Log.Infof("Join window closes after %d seconds.", duration);
		d.SendEvent("pairing-ended", &events.PairingEnded{
			DevicesFound: int(d.devicesFound - save),
		})
	}()

	permitJoinRequest := &nwkmgr.NwkSetPermitJoinReq{
		PermitJoinTime: &duration,
		PermitJoin:     nwkmgr.NwkPermitJoinTypeT_PERMIT_ALL.Enum(),
	}

	permitJoinResponse := &nwkmgr.NwkZigbeeGenericCnf{}

	err := d.nwkmgrConn.SendCommand(permitJoinRequest, permitJoinResponse)
	if err != nil {
		return fmt.Errorf("Failed to enable joining: %s", err)
	}
	if permitJoinResponse.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to enable joining: %s", permitJoinResponse.Status)
	}

	d.SendEvent("pairing-started", &events.PairingStarted{
		Duration: int(duration),
	})

	d.Log.Infof("Join window opens for %d seconds", duration);

	return nil
}

func (d *Driver) FetchDevices() error {
	return d.nwkmgrConn.FetchDeviceList()
}

func (d *Driver) onDeviceFound(deviceInfo *nwkmgr.NwkDeviceInfoT) {

	log.Debugf("\n\n")
	log.Debugf("---- Found Device IEEE:%X ----\f", *deviceInfo.IeeeAddress)
	log.Debugf("Device Info: %v", *deviceInfo)

	if d.devices[*deviceInfo.IeeeAddress] != nil {
		// We've seen this already, but it may have been repaired. We *should* just be able to replace
		// the deviceInfo object, which is used for all communication.
		// TODO: Actually verify this. May need to re-run channel init.
		d.devices[*deviceInfo.IeeeAddress].deviceInfo = deviceInfo
	}

	device := &Device{
		driver:     d,
		deviceInfo: deviceInfo,
		info: &model.Device{
			NaturalID:     fmt.Sprintf("%X", *deviceInfo.IeeeAddress),
			NaturalIDType: "zigbee",
			Signatures:    &map[string]string{},
		},
	}

	for _, endpoint := range deviceInfo.SimpleDescList {
		if *endpoint.ProfileId == 0xC05E /*ZLL*/ {

			switch *endpoint.DeviceId {

			case 0x0000: // On/Off Light
				fallthrough
			case 0x0100: // Dimmable Light
				fallthrough
			case 0x0200: // Color Light
				fallthrough
			case 0x210: // Ext Color Light
				(*device.info.Signatures)["ninja:thingType"] = "light"
			}
		}
	}

	err := device.getBasicInfo()
	if err != nil {
		log.Debugf("Failed to get basic info: %s", err)
	}

	if device.ManufacturerName != "" {
		(*device.info.Signatures)["zigbee:ManufacturerName"] = device.ManufacturerName
	} else {
		device.ManufacturerName = "Unknown"
	}
	if device.ModelIdentifier != "" {
		(*device.info.Signatures)["zigbee:ModelIdentifier"] = device.ModelIdentifier
	} else {
		device.ModelIdentifier = fmt.Sprintf("MAC:%X", *deviceInfo.IeeeAddress)
	}

	name := fmt.Sprintf("%s by %s", device.ModelIdentifier, device.ManufacturerName)

	device.info.Name = &name

	if log.IsDebugEnabled() {
		spew.Dump(deviceInfo)
	}

	err = d.Conn.ExportDevice(device)
	if err != nil {
		log.Fatalf("Failed to export zigbee device %s: %s", name, err)
	}
	d.devicesFound++;

	log.Debugf("Got device : %d", *deviceInfo.IeeeAddress)

	for _, endpoint := range deviceInfo.SimpleDescList {
		log.Debugf("Got endpoint : %d", *endpoint.EndpointId)

		if containsUInt32(endpoint.InputClusters, ClusterIDOnOff) {
			log.Debugf("This endpoint has on/off cluster")

			onOff := &OnOffChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDOnOff),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := onOff.init()
			if err != nil {
				log.Debugf("Failed initialising on/off channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDPower) {
			log.Debugf("This endpoint has power cluster")

			power := &PowerChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDPower),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := power.init()
			if err != nil {
				log.Debugf("Failed initialising power channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDTemp) {
			log.Debugf("This endpoint has temperature cluster")

			temp := &TempChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDTemp),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := temp.init()
			if err != nil {
				log.Debugf("Failed initialising temp channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDHumidity) {
			log.Debugf("This endpoint has humidity cluster")

			humidity := &HumidityChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDHumidity),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := humidity.init()
			if err != nil {
				log.Debugf("Failed initialising humidity channel: %s", err)
			}

		}

		if containsUInt32(endpoint.InputClusters, ClusterIDLevel) {
			log.Debugf("This endpoint has level cluster. Exporting as brightness channel")

			if log.IsDebugEnabled() {
				spew.Dump("brightness cluster", endpoint, ClusterIDLevel)
			}

			brightness := &BrightnessChannel{
				Channel: Channel{
					ID:       fmt.Sprintf("%d-%d", *endpoint.EndpointId, ClusterIDLevel),
					device:   device,
					endpoint: endpoint,
				},
			}

			err := brightness.init()
			if err != nil {
				log.Debugf("Failed initialising brightness channel: %s", err)
			}

		}

	}

	d.devices[*deviceInfo.IeeeAddress] = device

	fmt.Printf("---- Finished Device IEEE:%X ----\n", *deviceInfo.IeeeAddress)

}

func getCurDir() string {
	pwd, _ := os.Getwd()
	return pwd + "/"
}

func containsUInt32(hackstack []uint32, needle uint32) bool {
	for _, cluster := range hackstack {
		if cluster == needle {
			return true
		}
	}
	return false
}

func waitUntilZStackReady(checkFile string) {
	if checkFile == "" {
		return
	}

	// cooperate with zigbeeHAgw so that we don't start the zigbee driver
	// until we look somewhat stable.
	log.Debugf("waiting until zigbeeHAgw writes %s", checkFile)
	for {
		if _, err := os.Stat(checkFile); err == nil {
			break
		} else {
			time.Sleep(time.Second)
		}
	}
	log.Debugf("%s detected. start up continues...", checkFile)
}
<|endoftext|>


<commit_before>// Define a type called Client that handles authentication and communication with the sound
// API
package soundclient

import ( "fmt"
          "net/http"
          "io/ioutil"
        )

const apiBaseUrl string = "https://api.soundcloud.com"

type SoundCloud struct {
    ClientId string
    ClientSecret string
}

// Handle get requests to an endpoint
//
func Get(s SoundCloud, apiCall string, kargs string) {
    url := apiBaseUrl+"/"+apiCall+"/"+kargs+"?client_id="+s.ClientId
    resp, err := http.Get(url)
    fmt.Println("[DEBUG] URL-> ", url)
    if (err != nil) {
        fmt.Println("[ERROR] ", err)
    } else {
        body, err := ioutil.ReadAll(resp.Body)
        if (err == nil) {
            fmt.Printf("%s" ,body)
        }
    }
    defer resp.Body.Close()

}

// Return a soundcloud track
func (s SoundCloud) Tracks(trackId string) {
    Get(s, "tracks", trackId)
}

func (s SoundCloud) Playlists(playlistsId string) {
    Get(s, "playlists", playlistsId)
}

// Return a soundcloud user
func (s SoundCloud) Users(userId string) {
    Get(s, "users", userId)
}

// Handle POST requests
func Post() {

}

<commit_msg>return all responses as jason object<commit_after>// Define a type called Client that handles authentication and communication with the sound
// API
package soundclient

import ( "fmt"
          "net/http"
          "io/ioutil"
          "github.com/antonholmquist/jason"
        )

const apiBaseUrl string = "https://api.soundcloud.com"

type SoundCloud struct {
    ClientId string
    ClientSecret string
}

// Handle get requests to an endpoint
// and return a *jason.Object containing data
func Get(s SoundCloud, apiCall string, kargs string) *jason.Object {
    url := apiBaseUrl+"/"+apiCall+"/"+kargs+"?client_id="+s.ClientId
    resp, err := http.Get(url)
    if (err != nil) {
        fmt.Println("[ERROR] ", err)
    } else {
        result, err := ioutil.ReadAll(resp.Body)
        if (err == nil) {
            v, _ := jason.NewObjectFromBytes(result)
            return v
        }
    }
    defer resp.Body.Close()
    return nil

}

// Return a soundcloud track
func (s SoundCloud) Tracks(trackId string) *jason.Object {
    return Get(s, "tracks", trackId)
}

func (s SoundCloud) Playlists(playlistsId string) *jason.Object {
    return Get(s, "playlists", playlistsId)
}

// Return a soundcloud user
func (s SoundCloud) Users(userId string) *jason.Object {
    return Get(s, "users", userId)
}

// Get Group members and contributed tracks
func (s SoundCloud) Groups(groupId string) *jason.Object {
    return Get(s, "groups", groupId)
}

// Handle POST requests
func Post() {

}

<|endoftext|>


<commit_before>package bot

const changeLogPrefix = "\n`->` "

// probably want to move this to the DB, but not bad to have it here
var changeLog = map[string]string{
	"0.2": changeLogPrefix + "Included this command!" +
		changeLogPrefix + "Updated `Rank` command to prevent removal of lowest role." +
		changeLogPrefix + "Added random drops for tickets" +
		changeLogPrefix + "Fixed the cooldown so users wouldn't be spammed due to high luck stat" +
		changeLogPrefix + "Added `Raffle` related commands... For rafflin'" +
		changeLogPrefix + "For future reference, previous versions included help, team, rank, and NSFW commands as well as a welcome message to the server.",

	"0.2.1": changeLogPrefix + "Updated raffle art/relic submissions to post all submissions on command instead of over time.",
}

func commChange(pack *commPackage) {
	pack.session.ChannelMessageSend(pack.channel.ID, "`Moebot update log` (ver "+version+"): \n"+changeLog["0.3"])
}
<commit_msg>fixes typo with changelog<commit_after>package bot

const changeLogPrefix = "\n`->` "

// probably want to move this to the DB, but not bad to have it here
var changeLog = map[string]string{
	"0.2": changeLogPrefix + "Included this command!" +
		changeLogPrefix + "Updated `Rank` command to prevent removal of lowest role." +
		changeLogPrefix + "Added random drops for tickets" +
		changeLogPrefix + "Fixed the cooldown so users wouldn't be spammed due to high luck stat" +
		changeLogPrefix + "Added `Raffle` related commands... For rafflin'" +
		changeLogPrefix + "For future reference, previous versions included help, team, rank, and NSFW commands as well as a welcome message to the server.",

	"0.2.1": changeLogPrefix + "Updated raffle art/relic submissions to post all submissions on command instead of over time.",
}

func commChange(pack *commPackage) {
	pack.session.ChannelMessageSend(pack.channel.ID, "`Moebot update log` (ver "+version+"): \n"+changeLog["0.2.1"])
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"sort"
	"time"

	"owl/common/types"

	"github.com/Knetic/govaluate"
)

const (
	MAX_METHOD    = "max"
	MIN_METHOD    = "min"
	RATIO_METHOD  = "ratio"
	TOP_METHOD    = "top"
	BOTTOM_METHOD = "bottom"
	LAST_METHOD   = "last"
	NODATA_METHOD = "nodata"
)

func maxMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}

	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}
		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}

		sort.Float64s(values)
		parameters := make(map[string]interface{}, 8)
		current_threshold := values[len(values)-1]
		parameters["current_threshold"] = current_threshold
		parameters["threshold"] = trigger.Threshold
		expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
		trigger_result, err := compute(parameters, expression)
		if err != nil {
			lg.Error(err.Error())
			return trigger_result_set, err
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func minMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}

	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}
		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}

		sort.Float64s(values)
		parameters := make(map[string]interface{}, 8)
		current_threshold := values[0]
		parameters["current_threshold"] = current_threshold
		parameters["threshold"] = trigger.Threshold
		expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
		trigger_result, err := compute(parameters, expression)
		if err != nil {
			lg.Error(err.Error())
			return trigger_result_set, err
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func ratioMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	start_time, end_time := unit(cycle, trigger.Number)
	params = NewQueryParams(host_id, start_time, end_time, trigger.Tags, "sum", trigger.Metric)
	results_ago, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for index, result := range results {
		if len(result.Dps) == 0 {
			continue
		}
		data := avg(result.Dps)
		if index > len(results_ago)-1 {
			break
		}
		data_ago := avg(results_ago[index].Dps)

		current_threshold := ((data - data_ago) / data_ago) * 100

		parameters := make(map[string]interface{}, 8)
		parameters["current_threshold"] = current_threshold
		parameters["threshold"] = trigger.Threshold
		expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
		trigger_result, err := compute(parameters, expression)
		if err != nil {
			lg.Error(err.Error())
			return trigger_result_set, err
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func topMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}

		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}
		sort.Sort(sort.Reverse(sort.Float64Slice(values)))

		trigger_result := true
		var count int
		var sum float64
		for _, value := range values {
			count += 1
			if count > trigger.Number {
				break
			}
			sum += value
			parameters := make(map[string]interface{}, 8)
			parameters["current_threshold"] = value
			parameters["threshold"] = trigger.Threshold
			expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
			one_result, err := compute(parameters, expression)
			if err != nil {
				lg.Error(err.Error())
				continue
			}

			if !one_result {
				trigger_result = false
			}
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, sum/float64(trigger.Number), trigger_result))
	}

	return trigger_result_set, nil
}

func bottomMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}

		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}
		sort.Float64s(values)

		trigger_result := true
		var count int
		var sum float64
		for _, value := range values {
			count += 1
			if count > trigger.Number {
				break
			}
			sum += value
			parameters := make(map[string]interface{}, 8)
			parameters["current_threshold"] = value
			parameters["threshold"] = trigger.Threshold
			expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
			one_result, err := compute(parameters, expression)
			if err != nil {
				lg.Error(err.Error())
				continue
			}

			if !one_result {
				trigger_result = false
			}
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, sum/float64(trigger.Number), trigger_result))
	}

	return trigger_result_set, nil
}

func lastMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}

		value_times := make([]string, 0)
		for value_time, _ := range result.Dps {
			value_times = append(value_times, value_time)
		}

		sort.Strings(value_times)

		values := make([]float64, len(value_times))
		for index, value_time := range value_times {
			values[len(value_times)-(index+1)] = result.Dps[value_time]
		}

		trigger_result := true
		var count int
		var sum float64
		for _, value := range values {
			count += 1
			if count > trigger.Number {
				break
			}
			sum += value
			parameters := make(map[string]interface{}, 8)
			parameters["current_threshold"] = value
			parameters["threshold"] = trigger.Threshold
			expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
			one_result, err := compute(parameters, expression)
			if err != nil {
				lg.Error(err.Error())
				continue
			}

			if !one_result {
				trigger_result = false
			}
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, sum/float64(trigger.Number), trigger_result))
	}

	return trigger_result_set, nil
}

func nodataMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}

	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	var current_threshold float64 = 1

	if len(results) != 0 && len(results[0].Dps) != 0 {
		current_threshold = 0
	}

	parameters := make(map[string]interface{}, 8)
	parameters["current_threshold"] = current_threshold
	parameters["threshold"] = trigger.Threshold
	expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
	trigger_result, err := compute(parameters, expression)
	if err != nil {
		return nil, err
	}

	trigger_result_set.Triggered = trigger_result

	trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, make(map[string]string), make([]string, 0), current_threshold, trigger_result))

	return trigger_result_set, nil
}

func compute(params map[string]interface{}, express string) (bool, error) {
	expression, err := govaluate.NewEvaluableExpression(express)
	if err != nil {
		lg.Error("Parse the expression error %s", err.Error())
		return false, err
	}

	result, err := expression.Evaluate(params)
	if err != nil {
		lg.Error("Evaluate the expression error %s", err.Error())
		return false, err
	}
	value, ok := result.(bool)
	if !ok {
		lg.Error("result value is not bool type")
	}

	return value, nil
}

func unit(cycle, number int) (start_time, end_time string) {
	start := time.Now().Add(-time.Duration(time.Duration(number) * 24 * time.Hour)).Add(-time.Duration(time.Duration(cycle) * time.Minute))
	end := time.Now().Add(-time.Duration(time.Duration(number) * 24 * time.Hour))

	start_time = start.Format("2006/01/02 15:04:05")
	end_time = end.Format("2006/01/02 15:04:05")

	return start_time, end_time
}

func avg(values map[string]float64) float64 {
	var sum float64
	count := 0
	for _, value := range values {
		sum += value
		count += 1
	}

	return sum / float64(count)
}
<commit_msg>fix bug for restore alarm wrong current threhold<commit_after>package main

import (
	"errors"
	"fmt"
	"sort"
	"time"

	"owl/common/types"

	"github.com/Knetic/govaluate"
)

const (
	MAX_METHOD    = "max"
	MIN_METHOD    = "min"
	RATIO_METHOD  = "ratio"
	TOP_METHOD    = "top"
	BOTTOM_METHOD = "bottom"
	LAST_METHOD   = "last"
	NODATA_METHOD = "nodata"
)

func maxMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}

	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}
		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}

		sort.Float64s(values)
		parameters := make(map[string]interface{}, 8)
		current_threshold := values[len(values)-1]
		parameters["current_threshold"] = current_threshold
		parameters["threshold"] = trigger.Threshold
		expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
		trigger_result, err := compute(parameters, expression)
		if err != nil {
			lg.Error(err.Error())
			return trigger_result_set, err
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func minMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}

	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}
		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}

		sort.Float64s(values)
		parameters := make(map[string]interface{}, 8)
		current_threshold := values[0]
		parameters["current_threshold"] = current_threshold
		parameters["threshold"] = trigger.Threshold
		expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
		trigger_result, err := compute(parameters, expression)
		if err != nil {
			lg.Error(err.Error())
			return trigger_result_set, err
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func ratioMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	start_time, end_time := unit(cycle, trigger.Number)
	params = NewQueryParams(host_id, start_time, end_time, trigger.Tags, "sum", trigger.Metric)
	results_ago, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for index, result := range results {
		if len(result.Dps) == 0 {
			continue
		}
		data := avg(result.Dps)
		if index > len(results_ago)-1 {
			break
		}
		data_ago := avg(results_ago[index].Dps)

		current_threshold := ((data - data_ago) / data_ago) * 100

		parameters := make(map[string]interface{}, 8)
		parameters["current_threshold"] = current_threshold
		parameters["threshold"] = trigger.Threshold
		expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
		trigger_result, err := compute(parameters, expression)
		if err != nil {
			lg.Error(err.Error())
			return trigger_result_set, err
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func topMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}

		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}
		sort.Sort(sort.Reverse(sort.Float64Slice(values)))

		trigger_result := true
		var count int
		var sum float64
		var current_threshold float64
		for _, value := range values {
			count += 1
			if count > trigger.Number {
				break
			}
			sum += value
			parameters := make(map[string]interface{}, 8)
			parameters["current_threshold"] = value
			parameters["threshold"] = trigger.Threshold
			expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
			one_result, err := compute(parameters, expression)
			if err != nil {
				lg.Error(err.Error())
				continue
			}

			if !one_result {
				trigger_result = false
				current_threshold = value
			}
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
			current_threshold = sum / float64(trigger.Number)
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func bottomMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}

		values := make([]float64, 0)
		for _, value := range result.Dps {
			values = append(values, value)
		}
		sort.Float64s(values)

		trigger_result := true
		var count int
		var sum float64
		var current_threshold float64
		for _, value := range values {
			count += 1
			if count > trigger.Number {
				break
			}
			sum += value
			parameters := make(map[string]interface{}, 8)
			parameters["current_threshold"] = value
			parameters["threshold"] = trigger.Threshold
			expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
			one_result, err := compute(parameters, expression)
			if err != nil {
				lg.Error(err.Error())
				continue
			}

			if !one_result {
				trigger_result = false
				current_threshold = value
			}
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
			current_threshold = sum / float64(trigger.Number)
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func lastMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	if trigger.Number == 0 {
		err := errors.New("number can not be 0")
		lg.Error(err.Error())
		return nil, err
	}

	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}
	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	for _, result := range results {
		if len(result.Dps) == 0 {
			continue
		}

		value_times := make([]string, 0)
		for value_time, _ := range result.Dps {
			value_times = append(value_times, value_time)
		}

		sort.Strings(value_times)

		values := make([]float64, len(value_times))
		for index, value_time := range value_times {
			values[len(value_times)-(index+1)] = result.Dps[value_time]
		}

		trigger_result := true
		var count int
		var sum float64
		var current_threshold float64
		for _, value := range values {
			count += 1
			if count > trigger.Number {
				break
			}
			sum += value
			parameters := make(map[string]interface{}, 8)
			parameters["current_threshold"] = value
			parameters["threshold"] = trigger.Threshold
			expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
			one_result, err := compute(parameters, expression)
			if err != nil {
				lg.Error(err.Error())
				continue
			}

			if !one_result {
				trigger_result = false
				current_threshold = value
			}
		}

		if !trigger_result_set.Triggered && trigger_result {
			trigger_result_set.Triggered = trigger_result
			current_threshold = sum / float64(trigger.Number)
		}

		trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, result.Tags, result.AggregateTags, current_threshold, trigger_result))
	}

	return trigger_result_set, nil
}

func nodataMethod(host_id string, cycle int, trigger *types.Trigger) (*types.TriggerResultSet, error) {
	trigger_result_set := &types.TriggerResultSet{make([]*types.TriggerResult, 0), false}

	params := NewQueryParams(host_id, fmt.Sprintf("%dm-ago", cycle), "", trigger.Tags, "sum", trigger.Metric)
	results, err := tsdbClient.Query(params)
	if err != nil {
		lg.Error(err.Error())
		return nil, err
	}

	var current_threshold float64 = 1

	if len(results) != 0 && len(results[0].Dps) != 0 {
		current_threshold = 0
	}

	parameters := make(map[string]interface{}, 8)
	parameters["current_threshold"] = current_threshold
	parameters["threshold"] = trigger.Threshold
	expression := fmt.Sprintf("current_threshold %s threshold", trigger.Symbol)
	trigger_result, err := compute(parameters, expression)
	if err != nil {
		return nil, err
	}

	trigger_result_set.Triggered = trigger_result

	trigger_result_set.TriggerResults = append(trigger_result_set.TriggerResults, types.NewTriggerResult(trigger.Index, make(map[string]string), make([]string, 0), current_threshold, trigger_result))

	return trigger_result_set, nil
}

func compute(params map[string]interface{}, express string) (bool, error) {
	expression, err := govaluate.NewEvaluableExpression(express)
	if err != nil {
		lg.Error("Parse the expression error %s", err.Error())
		return false, err
	}

	result, err := expression.Evaluate(params)
	if err != nil {
		lg.Error("Evaluate the expression error %s", err.Error())
		return false, err
	}
	value, ok := result.(bool)
	if !ok {
		lg.Error("result value is not bool type")
	}

	return value, nil
}

func unit(cycle, number int) (start_time, end_time string) {
	start := time.Now().Add(-time.Duration(time.Duration(number) * 24 * time.Hour)).Add(-time.Duration(time.Duration(cycle) * time.Minute))
	end := time.Now().Add(-time.Duration(time.Duration(number) * 24 * time.Hour))

	start_time = start.Format("2006/01/02 15:04:05")
	end_time = end.Format("2006/01/02 15:04:05")

	return start_time, end_time
}

func avg(values map[string]float64) float64 {
	var sum float64
	count := 0
	for _, value := range values {
		sum += value
		count += 1
	}

	return sum / float64(count)
}
<|endoftext|>