<commit_before>package edit

import (
	"errors"
	"sort"

	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/eval"
	"github.com/elves/elvish/eval/types"
	"github.com/elves/elvish/eval/vartypes"
	"github.com/elves/elvish/parse"
	"github.com/xiaq/persistent/hashmap"
)

var errValueShouldBeFn = errors.New("value should be function")

func getBinding(bindingVar vartypes.Variable, k ui.Key) eval.Callable {
	binding := bindingVar.Get().(BindingTable)
	switch {
	case binding.HasKey(k):
		return binding.get(k)
	case binding.HasKey(ui.Default):
		return binding.get(ui.Default)
	default:
		return nil
	}
}

// BindingTable is a special Map that converts its key to ui.Key and ensures
// that its values satisfy eval.CallableValue.
type BindingTable struct {
	hashmap.Map
}

// Repr returns the representation of the binding table as if it were an
// ordinary map keyed by strings.
func (bt BindingTable) Repr(indent int) string {
	var builder types.MapReprBuilder
	builder.Indent = indent

	var keys ui.Keys
	for it := bt.Map.Iterator(); it.HasElem(); it.Next() {
		k, _ := it.Elem()
		keys = append(keys, k.(ui.Key))
	}
	sort.Sort(keys)

	for _, k := range keys {
		v, _ := bt.Map.Get(k)
		builder.WritePair(parse.Quote(k.String()), indent+2, types.Repr(v, indent+2))
	}

	return builder.String()
}

// Index converts the index to ui.Key and uses the Index of the inner Map.
func (bt BindingTable) Index(index interface{}) (interface{}, error) {
	return types.Index(bt.Map, ui.ToKey(index))
}

func (bt BindingTable) HasKey(k interface{}) bool {
	_, ok := bt.Map.Get(k)
	return ok
}

func (bt BindingTable) get(k ui.Key) eval.Callable {
	v, ok := bt.Map.Get(k)
	if !ok {
		panic("get called when key not present")
	}
	return v.(eval.Callable)
}

// Assoc converts the index to ui.Key, ensures that the value is CallableValue,
// uses the Assoc of the inner Map and converts the result to a BindingTable.
func (bt BindingTable) Assoc(k, v interface{}) (interface{}, error) {
	key := ui.ToKey(k)
	f, ok := v.(eval.Callable)
	if !ok {
		return nil, errValueShouldBeFn
	}
	map2 := bt.Map.Assoc(key, f)
	return BindingTable{map2}, nil
}

func makeBindingTable(f *eval.Frame, args []interface{}, opts map[string]interface{}) {
	var raw hashmap.Map
	eval.ScanArgs(args, &raw)
	eval.TakeNoOpt(opts)

	converted := types.EmptyMap
	for it := raw.Iterator(); it.HasElem(); it.Next() {
		k, v := it.Elem()
		f, ok := v.(eval.Callable)
		if !ok {
			throw(errValueShouldBeFn)
		}
		converted = converted.Assoc(ui.ToKey(k), f)
	}

	f.OutputChan() <- BindingTable{converted}
}
<commit_msg>edit: Implement Dissoc on BindingTable.<commit_after>package edit

import (
	"errors"
	"sort"

	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/eval"
	"github.com/elves/elvish/eval/types"
	"github.com/elves/elvish/eval/vartypes"
	"github.com/elves/elvish/parse"
	"github.com/xiaq/persistent/hashmap"
)

var errValueShouldBeFn = errors.New("value should be function")

func getBinding(bindingVar vartypes.Variable, k ui.Key) eval.Callable {
	binding := bindingVar.Get().(BindingTable)
	switch {
	case binding.HasKey(k):
		return binding.get(k)
	case binding.HasKey(ui.Default):
		return binding.get(ui.Default)
	default:
		return nil
	}
}

// BindingTable is a special Map that converts its key to ui.Key and ensures
// that its values satisfy eval.CallableValue.
type BindingTable struct {
	hashmap.Map
}

// Repr returns the representation of the binding table as if it were an
// ordinary map keyed by strings.
func (bt BindingTable) Repr(indent int) string {
	var builder types.MapReprBuilder
	builder.Indent = indent

	var keys ui.Keys
	for it := bt.Map.Iterator(); it.HasElem(); it.Next() {
		k, _ := it.Elem()
		keys = append(keys, k.(ui.Key))
	}
	sort.Sort(keys)

	for _, k := range keys {
		v, _ := bt.Map.Get(k)
		builder.WritePair(parse.Quote(k.String()), indent+2, types.Repr(v, indent+2))
	}

	return builder.String()
}

// Index converts the index to ui.Key and uses the Index of the inner Map.
func (bt BindingTable) Index(index interface{}) (interface{}, error) {
	return types.Index(bt.Map, ui.ToKey(index))
}

func (bt BindingTable) HasKey(k interface{}) bool {
	_, ok := bt.Map.Get(k)
	return ok
}

func (bt BindingTable) get(k ui.Key) eval.Callable {
	v, ok := bt.Map.Get(k)
	if !ok {
		panic("get called when key not present")
	}
	return v.(eval.Callable)
}

// Assoc converts the index to ui.Key, ensures that the value is CallableValue,
// uses the Assoc of the inner Map and converts the result to a BindingTable.
func (bt BindingTable) Assoc(k, v interface{}) (interface{}, error) {
	key := ui.ToKey(k)
	f, ok := v.(eval.Callable)
	if !ok {
		return nil, errValueShouldBeFn
	}
	map2 := bt.Map.Assoc(key, f)
	return BindingTable{map2}, nil
}

// Dissoc converts the key to ui.Key and calls the Dissoc method of the inner
// map.
func (bt BindingTable) Dissoc(k interface{}) interface{} {
	return BindingTable{bt.Map.Without(ui.ToKey(k))}
}

func makeBindingTable(f *eval.Frame, args []interface{}, opts map[string]interface{}) {
	var raw hashmap.Map
	eval.ScanArgs(args, &raw)
	eval.TakeNoOpt(opts)

	converted := types.EmptyMap
	for it := raw.Iterator(); it.HasElem(); it.Next() {
		k, v := it.Elem()
		f, ok := v.(eval.Callable)
		if !ok {
			throw(errValueShouldBeFn)
		}
		converted = converted.Assoc(ui.ToKey(k), f)
	}

	f.OutputChan() <- BindingTable{converted}
}
<|endoftext|>


<commit_before>package virtualbox

import (
	"errors"
	"fmt"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/executors"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/ssh"
	vbox "gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/virtualbox"
	"time"
)

type VirtualBoxExecutor struct {
	executors.AbstractExecutor
	// cmd             *exec.Cmd
	vmName          string
	sshCommand      ssh.Command
	sshPort         string
	provisioned     bool
	machineVerified bool
}

func (s *VirtualBoxExecutor) verifyMachine(vmName string, sshPort string) error {
	if s.machineVerified {
		return nil
	}

	// Create SSH command
	sshCommand := ssh.Command{
		Config:         *s.Config.SSH,
		Command:        "exit 0",
		Stdout:         s.BuildLog,
		Stderr:         s.BuildLog,
		ConnectRetries: 30,
	}
	host := `localhost`
	sshCommand.Port = &sshPort
	sshCommand.Host = &host

	s.Debugln("Connecting to SSH...")
	err := sshCommand.Connect()
	if err != nil {
		return err
	}
	defer sshCommand.Cleanup()
	err = sshCommand.Run()
	if err != nil {
		return err
	}
	s.machineVerified = true
	return nil
}

func (s *VirtualBoxExecutor) restoreFromSnapshot() error {
	s.Debugln("Reverting VM to current snapshot...")
	err := vbox.RevertToSnapshot(s.vmName)
	if err != nil {
		return err
	}

	return nil
}

// virtualbox doesn't support templates
func (s *VirtualBoxExecutor) createVM(vmName string) error {
	baseImage := s.Config.VirtualBox.BaseName
	if baseImage == "" {
		return errors.New("Missing Image setting from VirtualBox configuration")
	}

	vmStatus, _ := vbox.Status(vmName)
	if vmStatus == vbox.Invalid {
		vbox.Unregister(vmName)
	}

	if !vbox.Exist(vmName) {
		s.Debugln("Creating testing VM from VM", baseImage, "...")
		err := vbox.CreateOsVM(baseImage, vmName)
		if err != nil {
			return err
		}
	}

	s.Debugln("Creating localhost ssh forwarding...")
	vmSSHPort := helpers.StringOrDefault(s.Config.SSH.Port, "22")
	err := vbox.ConfigureSSH(vmName, vmSSHPort)
	if err != nil {
		return err
	}

	s.Debugln("Bootstraping VM...")
	err = vbox.Start(s.vmName)
	if err != nil {
		return err
	}

	s.Debugln("Identify SSH Port...")
	sshPort, err := vbox.FindSSHPort(s.vmName)
	s.sshPort = sshPort
	if err != nil {
		return err
	}

	s.Debugln("Waiting for VM to become responsive...")
	time.Sleep(10)
	err = s.verifyMachine(s.vmName, s.sshPort)
	if err != nil {
		return err
	}

	return nil
}

func (s *VirtualBoxExecutor) Prepare(globalConfig *common.Config, config *common.RunnerConfig, build *common.Build) error {
	err := s.AbstractExecutor.Prepare(globalConfig, config, build)
	if err != nil {
		return err
	}

	if s.BuildScript.PassFile {
		return errors.New("virtualbox doesn't support shells that require script file")
	}

	if s.Config.SSH == nil {
		return errors.New("Missing SSH config")
	}

	if s.Config.VirtualBox == nil {
		return errors.New("Missing VirtualBox configuration")
	}

	if s.Config.VirtualBox.BaseName == "" {
		return errors.New("Missing BaseName setting from VirtualBox configuration")
	}

	version, err := vbox.Version()
	if err != nil {
		panic(err)
	}

	s.Println("Using VirtualBox version", version, "executor...")

	vmStatus, _ := vbox.Status(s.vmName)
	if vmStatus == vbox.Invalid {
		vbox.Unregister(s.vmName)
	}

	if helpers.BoolOrDefault(s.Config.VirtualBox.DisableSnapshots, false) {
		s.vmName = s.Config.VirtualBox.BaseName + "-" + s.Build.ProjectUniqueName()
		if vbox.Exist(s.vmName) {
			s.Debugln("Deleting old VM...")
			vbox.Stop(s.vmName)
			vbox.Delete(s.vmName)
			vbox.Unregister(s.vmName)
		}
	} else {
		s.vmName = fmt.Sprintf("%s-runner-%s-concurrent-%d",
			s.Config.VirtualBox.BaseName,
			s.Build.Runner.ShortDescription(),
			s.Build.RunnerID)
	}

	if vbox.Exist(s.vmName) {
		s.Println("Restoring VM from snapshot...")
		err := s.restoreFromSnapshot()
		if err != nil {
			s.Println("Previous VM failed. Deleting, because", err)
			vbox.Stop(s.vmName)
			vbox.Delete(s.vmName)
			vbox.Unregister(s.vmName)
		}
	}

	if !vbox.Exist(s.vmName) {
		s.Println("Creating new VM...")
		err := s.createVM(s.vmName)
		if err != nil {
			return err
		}

		if !helpers.BoolOrDefault(s.Config.VirtualBox.DisableSnapshots, false) {
			s.Println("Creating default snapshot...")
			err = vbox.CreateSnapshot(s.vmName, "Started")
			if err != nil {
				return err
			}
		}
	}

	s.Debugln("Checking VM status...")
	status, err := vbox.Status(s.vmName)
	if err != nil {
		return err
	}

	if status == vbox.Stopped || status == vbox.Suspended || status == vbox.Saved {
		s.Println("Starting VM...")
		err := vbox.Start(s.vmName)
		if err != nil {
			return err
		}
	}

	if status != vbox.Running {
		s.Debugln("Waiting for VM to run...")
		err = vbox.WaitForStatus(s.vmName, vbox.Running, 60)
		if err != nil {
			return err
		}
	}

	s.Debugln("Identify SSH Port...")
	sshPort, err := vbox.FindSSHPort(s.vmName)
	s.sshPort = sshPort
	if err != nil {
		return err
	}

	s.Println("Waiting VM to become responsive...")
	err = s.verifyMachine(s.vmName, s.sshPort)
	if err != nil {
		return err
	}

	s.provisioned = true

	return nil
}

func (s *VirtualBoxExecutor) Start() error {
	s.Println("Starting SSH command...")
	s.sshCommand = ssh.Command{
		Config:      *s.Config.SSH,
		Environment: s.BuildScript.Environment,
		Command:     s.BuildScript.GetFullCommand(),
		Stdin:       s.BuildScript.GetScriptBytes(),
		Stdout:      s.BuildLog,
		Stderr:      s.BuildLog,
	}
	host := `localhost`
	s.sshCommand.Port = &s.sshPort
	s.sshCommand.Host = &host

	s.Debugln("Connecting to SSH server...")
	err := s.sshCommand.Connect()
	if err != nil {
		return err
	}

	// Wait for process to exit
	go func() {
		s.Debugln("Will run SSH command...")
		err := s.sshCommand.Run()
		s.Debugln("SSH command finished with", err)
		s.BuildFinish <- err
	}()
	return nil
}

func (s *VirtualBoxExecutor) Cleanup() {
	s.sshCommand.Cleanup()

	if s.vmName != "" {
		vbox.Kill(s.vmName)

		if helpers.BoolOrDefault(s.Config.VirtualBox.DisableSnapshots, false) || !s.provisioned {
			vbox.Delete(s.vmName)
		}
	}
}

func init() {
	options := executors.ExecutorOptions{
		DefaultBuildsDir: "builds",
		SharedBuildsDir:  false,
		Shell: common.ShellScriptInfo{
			Shell: "bash",
			Type:  common.LoginShell,
		},
		ShowHostname: true,
	}

	creator := func() common.Executor {
		return &VirtualBoxExecutor{
			AbstractExecutor: executors.AbstractExecutor{
				ExecutorOptions: options,
			},
		}
	}

	featuresUpdater := func(features *common.FeaturesInfo) {
		features.Variables = true
	}

	common.RegisterExecutor("virtualbox", executors.DefaultExecutorProvider{
		Creator:         creator,
		FeaturesUpdater: featuresUpdater,
	})
}
<commit_msg>Remove legacy code fragment<commit_after>package virtualbox

import (
	"errors"
	"fmt"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/executors"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/ssh"
	vbox "gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/virtualbox"
	"time"
)

type VirtualBoxExecutor struct {
	executors.AbstractExecutor
	vmName          string
	sshCommand      ssh.Command
	sshPort         string
	provisioned     bool
	machineVerified bool
}

func (s *VirtualBoxExecutor) verifyMachine(vmName string, sshPort string) error {
	if s.machineVerified {
		return nil
	}

	// Create SSH command
	sshCommand := ssh.Command{
		Config:         *s.Config.SSH,
		Command:        "exit 0",
		Stdout:         s.BuildLog,
		Stderr:         s.BuildLog,
		ConnectRetries: 30,
	}
	host := `localhost`
	sshCommand.Port = &sshPort
	sshCommand.Host = &host

	s.Debugln("Connecting to SSH...")
	err := sshCommand.Connect()
	if err != nil {
		return err
	}
	defer sshCommand.Cleanup()
	err = sshCommand.Run()
	if err != nil {
		return err
	}
	s.machineVerified = true
	return nil
}

func (s *VirtualBoxExecutor) restoreFromSnapshot() error {
	s.Debugln("Reverting VM to current snapshot...")
	err := vbox.RevertToSnapshot(s.vmName)
	if err != nil {
		return err
	}

	return nil
}

// virtualbox doesn't support templates
func (s *VirtualBoxExecutor) createVM(vmName string) error {
	baseImage := s.Config.VirtualBox.BaseName
	if baseImage == "" {
		return errors.New("Missing Image setting from VirtualBox configuration")
	}

	vmStatus, _ := vbox.Status(vmName)
	if vmStatus == vbox.Invalid {
		vbox.Unregister(vmName)
	}

	if !vbox.Exist(vmName) {
		s.Debugln("Creating testing VM from VM", baseImage, "...")
		err := vbox.CreateOsVM(baseImage, vmName)
		if err != nil {
			return err
		}
	}

	s.Debugln("Creating localhost ssh forwarding...")
	vmSSHPort := helpers.StringOrDefault(s.Config.SSH.Port, "22")
	err := vbox.ConfigureSSH(vmName, vmSSHPort)
	if err != nil {
		return err
	}

	s.Debugln("Bootstraping VM...")
	err = vbox.Start(s.vmName)
	if err != nil {
		return err
	}

	s.Debugln("Identify SSH Port...")
	sshPort, err := vbox.FindSSHPort(s.vmName)
	s.sshPort = sshPort
	if err != nil {
		return err
	}

	s.Debugln("Waiting for VM to become responsive...")
	time.Sleep(10)
	err = s.verifyMachine(s.vmName, s.sshPort)
	if err != nil {
		return err
	}

	return nil
}

func (s *VirtualBoxExecutor) Prepare(globalConfig *common.Config, config *common.RunnerConfig, build *common.Build) error {
	err := s.AbstractExecutor.Prepare(globalConfig, config, build)
	if err != nil {
		return err
	}

	if s.BuildScript.PassFile {
		return errors.New("virtualbox doesn't support shells that require script file")
	}

	if s.Config.SSH == nil {
		return errors.New("Missing SSH config")
	}

	if s.Config.VirtualBox == nil {
		return errors.New("Missing VirtualBox configuration")
	}

	if s.Config.VirtualBox.BaseName == "" {
		return errors.New("Missing BaseName setting from VirtualBox configuration")
	}

	version, err := vbox.Version()
	if err != nil {
		panic(err)
	}

	s.Println("Using VirtualBox version", version, "executor...")

	vmStatus, _ := vbox.Status(s.vmName)
	if vmStatus == vbox.Invalid {
		vbox.Unregister(s.vmName)
	}

	if helpers.BoolOrDefault(s.Config.VirtualBox.DisableSnapshots, false) {
		s.vmName = s.Config.VirtualBox.BaseName + "-" + s.Build.ProjectUniqueName()
		if vbox.Exist(s.vmName) {
			s.Debugln("Deleting old VM...")
			vbox.Stop(s.vmName)
			vbox.Delete(s.vmName)
			vbox.Unregister(s.vmName)
		}
	} else {
		s.vmName = fmt.Sprintf("%s-runner-%s-concurrent-%d",
			s.Config.VirtualBox.BaseName,
			s.Build.Runner.ShortDescription(),
			s.Build.RunnerID)
	}

	if vbox.Exist(s.vmName) {
		s.Println("Restoring VM from snapshot...")
		err := s.restoreFromSnapshot()
		if err != nil {
			s.Println("Previous VM failed. Deleting, because", err)
			vbox.Stop(s.vmName)
			vbox.Delete(s.vmName)
			vbox.Unregister(s.vmName)
		}
	}

	if !vbox.Exist(s.vmName) {
		s.Println("Creating new VM...")
		err := s.createVM(s.vmName)
		if err != nil {
			return err
		}

		if !helpers.BoolOrDefault(s.Config.VirtualBox.DisableSnapshots, false) {
			s.Println("Creating default snapshot...")
			err = vbox.CreateSnapshot(s.vmName, "Started")
			if err != nil {
				return err
			}
		}
	}

	s.Debugln("Checking VM status...")
	status, err := vbox.Status(s.vmName)
	if err != nil {
		return err
	}

	if status == vbox.Stopped || status == vbox.Suspended || status == vbox.Saved {
		s.Println("Starting VM...")
		err := vbox.Start(s.vmName)
		if err != nil {
			return err
		}
	}

	if status != vbox.Running {
		s.Debugln("Waiting for VM to run...")
		err = vbox.WaitForStatus(s.vmName, vbox.Running, 60)
		if err != nil {
			return err
		}
	}

	s.Debugln("Identify SSH Port...")
	sshPort, err := vbox.FindSSHPort(s.vmName)
	s.sshPort = sshPort
	if err != nil {
		return err
	}

	s.Println("Waiting VM to become responsive...")
	err = s.verifyMachine(s.vmName, s.sshPort)
	if err != nil {
		return err
	}

	s.provisioned = true

	return nil
}

func (s *VirtualBoxExecutor) Start() error {
	s.Println("Starting SSH command...")
	s.sshCommand = ssh.Command{
		Config:      *s.Config.SSH,
		Environment: s.BuildScript.Environment,
		Command:     s.BuildScript.GetFullCommand(),
		Stdin:       s.BuildScript.GetScriptBytes(),
		Stdout:      s.BuildLog,
		Stderr:      s.BuildLog,
	}
	host := `localhost`
	s.sshCommand.Port = &s.sshPort
	s.sshCommand.Host = &host

	s.Debugln("Connecting to SSH server...")
	err := s.sshCommand.Connect()
	if err != nil {
		return err
	}

	// Wait for process to exit
	go func() {
		s.Debugln("Will run SSH command...")
		err := s.sshCommand.Run()
		s.Debugln("SSH command finished with", err)
		s.BuildFinish <- err
	}()
	return nil
}

func (s *VirtualBoxExecutor) Cleanup() {
	s.sshCommand.Cleanup()

	if s.vmName != "" {
		vbox.Kill(s.vmName)

		if helpers.BoolOrDefault(s.Config.VirtualBox.DisableSnapshots, false) || !s.provisioned {
			vbox.Delete(s.vmName)
		}
	}
}

func init() {
	options := executors.ExecutorOptions{
		DefaultBuildsDir: "builds",
		SharedBuildsDir:  false,
		Shell: common.ShellScriptInfo{
			Shell: "bash",
			Type:  common.LoginShell,
		},
		ShowHostname: true,
	}

	creator := func() common.Executor {
		return &VirtualBoxExecutor{
			AbstractExecutor: executors.AbstractExecutor{
				ExecutorOptions: options,
			},
		}
	}

	featuresUpdater := func(features *common.FeaturesInfo) {
		features.Variables = true
	}

	common.RegisterExecutor("virtualbox", executors.DefaultExecutorProvider{
		Creator:         creator,
		FeaturesUpdater: featuresUpdater,
	})
}
<|endoftext|>


<commit_before>package main

import (
//    "encoding/json"
	"bytes"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"runtime"
	"strings"
//	"time"

	"github.com/kr/pretty"
	"github.com/docopt/docopt-go"

    "github.com/greggyNapalm/katyusha/katyushalib"
)

const version = "0.0.1"
const workers_num = 1000
const tgt_host, tgt_port = "127.0.0.1", 80

func pp(data interface{}) {
	fmt.Printf("%# v", pretty.Formatter(data))
}

func compose_url(tgt_addr string, tgt_port int) string {
	return fmt.Sprintf("http://%s:%d", tgt_addr, tgt_port)
}

func remote_deal(dst_addr string) {
	// Close TCP connection on each request
	for {
		//http.Get("http://127.0.0.1:80")
		_, err := http.Get(dst_addr)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func remote_deal_reuse(dst_addr string) {
	for {
		//http.Get("http://127.0.0.1:80")
		_, err := http.Get(dst_addr)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func get_uname() string {
	cmd := exec.Command("uname", "-sr")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
	return strings.Trim(out.String(), "\n")
}

func main() {
      usage := `katyusha.

Usage:
  katyusha KCFG_PATH
  katyusha -h | --help
  katyusha --version

Arguments:
  KCFG_PATH     Katyushas config file path.

Options:
  -h --help     Show this screen.
  -v --verbose  Give more verbose output.
  --version     Show version.`

    arguments, _ := docopt.Parse(usage, nil, true, "Katyusha load tool v." + version, false)
    fmt.Println(arguments)
    fmt.Println(arguments["KCFG_PATH"])

    //kcfg := katyushalib.ComposeCfg(arguments["KCFG_PATH"])
    kcfg := katyushalib.ComposeCfg("kconfig.json")

	cpu_num := runtime.NumCPU()
	runtime.GOMAXPROCS(cpu_num)

	dst := compose_url(tgt_host, tgt_port)

	log.Printf("Runtime: %s / golang %s / cores count %d", get_uname(), runtime.Version(), cpu_num)
	log.Printf("Target addr: %s", dst)
}
<commit_msg>fixing syntax errors<commit_after>package main

import (
//    "encoding/json"
	"bytes"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"runtime"
	"strings"
//	"time"

	"github.com/kr/pretty"
	"github.com/docopt/docopt-go"

    "github.com/greggyNapalm/katyusha/katyushalib"
)

const version = "0.0.1"
const workers_num = 1000
const tgt_host, tgt_port = "127.0.0.1", 80

func pp(data interface{}) {
	fmt.Printf("%# v", pretty.Formatter(data))
}

func compose_url(tgt_addr string, tgt_port int) string {
	return fmt.Sprintf("http://%s:%d", tgt_addr, tgt_port)
}

func remote_deal(dst_addr string) {
	// Close TCP connection on each request
	for {
		//http.Get("http://127.0.0.1:80")
		_, err := http.Get(dst_addr)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func remote_deal_reuse(dst_addr string) {
	for {
		//http.Get("http://127.0.0.1:80")
		_, err := http.Get(dst_addr)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func get_uname() string {
	cmd := exec.Command("uname", "-sr")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
	return strings.Trim(out.String(), "\n")
}

func main() {
      usage := `katyusha.

Usage:
  katyusha KCFG_PATH
  katyusha -h | --help
  katyusha --version

Arguments:
  KCFG_PATH     Katyushas config file path.

Options:
  -h --help     Show this screen.
  -v --verbose  Give more verbose output.
  --version     Show version.`

    arguments, _ := docopt.Parse(usage, nil, true, "Katyusha load tool v." + version, false)
    fmt.Println(arguments)
    fmt.Println(arguments["KCFG_PATH"])

    //kcfg := katyushalib.ComposeCfg(arguments["KCFG_PATH"])
    kcfg := katyushalib.ComposeCfg("kconfig.json")
    pp(kcfg)

	cpu_num := runtime.NumCPU()
	runtime.GOMAXPROCS(cpu_num)

	dst := compose_url(tgt_host, tgt_port)

	log.Printf("Runtime: %s / golang %s / cores count %d", get_uname(), runtime.Version(), cpu_num)
	log.Printf("Target addr: %s", dst)
}
<|endoftext|>


<commit_before>package gateway

import (
	"strconv"
	"strings"
	"time"

	"github.com/TykTechnologies/tyk/config"
	"github.com/TykTechnologies/tyk/storage"
)

type HealthPrefix string

const (
	Throttle          HealthPrefix = "Throttle"
	QuotaViolation    HealthPrefix = "QuotaViolation"
	KeyFailure        HealthPrefix = "KeyFailure"
	RequestLog        HealthPrefix = "Request"
	BlockedRequestLog HealthPrefix = "BlockedRequest"
)

type HealthChecker interface {
	Init(storage.Handler)
	ApiHealthValues() (HealthCheckValues, error)
	StoreCounterVal(HealthPrefix, string)
}

type HealthCheckValues struct {
	ThrottledRequestsPS float64 `bson:"throttle_reqests_per_second,omitempty" json:"throttle_reqests_per_second"`
	QuotaViolationsPS   float64 `bson:"quota_violations_per_second,omitempty" json:"quota_violations_per_second"`
	KeyFailuresPS       float64 `bson:"key_failures_per_second,omitempty" json:"key_failures_per_second"`
	AvgUpstreamLatency  float64 `bson:"average_upstream_latency,omitempty" json:"average_upstream_latency"`
	AvgRequestsPS       float64 `bson:"average_requests_per_second,omitempty" json:"average_requests_per_second"`
}

type DefaultHealthChecker struct {
	storage storage.Handler
	APIID   string
}

func (h *DefaultHealthChecker) Init(storeType storage.Handler) {
	if !config.Global().HealthCheck.EnableHealthChecks {
		return
	}

	log.Info("Initializing HealthChecker")
	h.storage = storeType
	h.storage.Connect()
}

func (h *DefaultHealthChecker) CreateKeyName(subKey HealthPrefix) string {
	// Key should be API-ID.SubKey.123456789
	return h.APIID + "." + string(subKey)
}

// reportHealthValue is a shortcut we can use throughout the app to push a health check value
func reportHealthValue(spec *APISpec, counter HealthPrefix, value string) {
	if !spec.GlobalConfig.HealthCheck.EnableHealthChecks {
		return
	}

	spec.Health.StoreCounterVal(counter, value)
}

func (h *DefaultHealthChecker) StoreCounterVal(counterType HealthPrefix, value string) {
	searchStr := h.CreateKeyName(counterType)
	log.Debug("Adding Healthcheck to: ", searchStr)
	log.Debug("Val is: ", value)
	//go h.storage.SetKey(searchStr, value, config.Global.HealthCheck.HealthCheckValueTimeout)
	if value != "-1" {
		// need to ensure uniqueness
		now_string := strconv.Itoa(int(time.Now().UnixNano()))
		value = now_string + "." + value
		log.Debug("Set value to: ", value)
	}
	go h.storage.SetRollingWindow(searchStr, config.Global().HealthCheck.HealthCheckValueTimeout, value, false)
}

func (h *DefaultHealthChecker) getAvgCount(prefix HealthPrefix) float64 {
	searchStr := h.CreateKeyName(prefix)
	log.Debug("Searching for: ", searchStr)

	count, _ := h.storage.SetRollingWindow(searchStr, config.Global().HealthCheck.HealthCheckValueTimeout, "-1", false)
	log.Debug("Count is: ", count)
	divisor := float64(config.Global().HealthCheck.HealthCheckValueTimeout)
	if divisor == 0 {
		log.Warning("The Health Check sample timeout is set to 0, samples will never be deleted!!!")
		divisor = 60.0
	}
	if count > 0 {
		return roundValue((float64(count) - 1) / divisor)
	}

	return 0.00
}

func roundValue(untruncated float64) float64 {
	return float64(int(untruncated*100)) / 100
}

func (h *DefaultHealthChecker) ApiHealthValues() (HealthCheckValues, error) {
	values := HealthCheckValues{}

	// Get the counted / average values
	values.ThrottledRequestsPS = h.getAvgCount(Throttle)
	values.QuotaViolationsPS = h.getAvgCount(QuotaViolation)
	values.KeyFailuresPS = h.getAvgCount(KeyFailure)
	values.AvgRequestsPS = h.getAvgCount(RequestLog)

	// Get the micro latency graph, an average upstream latency
	searchStr := h.APIID + "." + string(RequestLog)
	log.Debug("Searching KV for: ", searchStr)
	_, vals := h.storage.SetRollingWindow(searchStr, config.Global().HealthCheck.HealthCheckValueTimeout, "-1", false)
	log.Debug("Found: ", vals)
	if len(vals) == 0 {
		return values, nil
	}
	var runningTotal int
	for _, v := range vals {
		s := string(v.([]byte))
		log.Debug("V is: ", s)
		splitValues := strings.Split(s, ".")
		if len(splitValues) > 1 {
			vInt, err := strconv.Atoi(splitValues[1])
			if err != nil {
				log.Error("Couldn't convert tracked latency value to Int, vl is: ", err)
			} else {
				runningTotal += vInt
			}
		}

	}
	values.AvgUpstreamLatency = roundValue(float64(runningTotal / len(vals)))
	return values, nil
}
<commit_msg>Fix panic on /health endpoint (#3223)<commit_after>package gateway

import (
	"strconv"
	"strings"
	"time"

	"github.com/TykTechnologies/tyk/config"
	"github.com/TykTechnologies/tyk/storage"
)

type HealthPrefix string

const (
	Throttle          HealthPrefix = "Throttle"
	QuotaViolation    HealthPrefix = "QuotaViolation"
	KeyFailure        HealthPrefix = "KeyFailure"
	RequestLog        HealthPrefix = "Request"
	BlockedRequestLog HealthPrefix = "BlockedRequest"
)

type HealthChecker interface {
	Init(storage.Handler)
	ApiHealthValues() (HealthCheckValues, error)
	StoreCounterVal(HealthPrefix, string)
}

type HealthCheckValues struct {
	ThrottledRequestsPS float64 `bson:"throttle_reqests_per_second,omitempty" json:"throttle_reqests_per_second"`
	QuotaViolationsPS   float64 `bson:"quota_violations_per_second,omitempty" json:"quota_violations_per_second"`
	KeyFailuresPS       float64 `bson:"key_failures_per_second,omitempty" json:"key_failures_per_second"`
	AvgUpstreamLatency  float64 `bson:"average_upstream_latency,omitempty" json:"average_upstream_latency"`
	AvgRequestsPS       float64 `bson:"average_requests_per_second,omitempty" json:"average_requests_per_second"`
}

type DefaultHealthChecker struct {
	storage storage.Handler
	APIID   string
}

func (h *DefaultHealthChecker) Init(storeType storage.Handler) {
	if !config.Global().HealthCheck.EnableHealthChecks {
		return
	}

	log.Info("Initializing HealthChecker")
	h.storage = storeType
	h.storage.Connect()
}

func (h *DefaultHealthChecker) CreateKeyName(subKey HealthPrefix) string {
	// Key should be API-ID.SubKey.123456789
	return h.APIID + "." + string(subKey)
}

// reportHealthValue is a shortcut we can use throughout the app to push a health check value
func reportHealthValue(spec *APISpec, counter HealthPrefix, value string) {
	if !spec.GlobalConfig.HealthCheck.EnableHealthChecks {
		return
	}

	spec.Health.StoreCounterVal(counter, value)
}

func (h *DefaultHealthChecker) StoreCounterVal(counterType HealthPrefix, value string) {
	searchStr := h.CreateKeyName(counterType)
	log.Debug("Adding Healthcheck to: ", searchStr)
	log.Debug("Val is: ", value)
	//go h.storage.SetKey(searchStr, value, config.Global.HealthCheck.HealthCheckValueTimeout)
	if value != "-1" {
		// need to ensure uniqueness
		now_string := strconv.Itoa(int(time.Now().UnixNano()))
		value = now_string + "." + value
		log.Debug("Set value to: ", value)
	}
	go h.storage.SetRollingWindow(searchStr, config.Global().HealthCheck.HealthCheckValueTimeout, value, false)
}

func (h *DefaultHealthChecker) getAvgCount(prefix HealthPrefix) float64 {
	searchStr := h.CreateKeyName(prefix)
	log.Debug("Searching for: ", searchStr)

	count, _ := h.storage.SetRollingWindow(searchStr, config.Global().HealthCheck.HealthCheckValueTimeout, "-1", false)
	log.Debug("Count is: ", count)
	divisor := float64(config.Global().HealthCheck.HealthCheckValueTimeout)
	if divisor == 0 {
		log.Warning("The Health Check sample timeout is set to 0, samples will never be deleted!!!")
		divisor = 60.0
	}
	if count > 0 {
		return roundValue((float64(count) - 1) / divisor)
	}

	return 0.00
}

func roundValue(untruncated float64) float64 {
	return float64(int(untruncated*100)) / 100
}

func (h *DefaultHealthChecker) ApiHealthValues() (HealthCheckValues, error) {
	values := HealthCheckValues{}

	// Get the counted / average values
	values.ThrottledRequestsPS = h.getAvgCount(Throttle)
	values.QuotaViolationsPS = h.getAvgCount(QuotaViolation)
	values.KeyFailuresPS = h.getAvgCount(KeyFailure)
	values.AvgRequestsPS = h.getAvgCount(RequestLog)

	// Get the micro latency graph, an average upstream latency
	searchStr := h.APIID + "." + string(RequestLog)
	log.Debug("Searching KV for: ", searchStr)
	_, vals := h.storage.SetRollingWindow(searchStr, config.Global().HealthCheck.HealthCheckValueTimeout, "-1", false)
	log.Debug("Found: ", vals)
	if len(vals) == 0 {
		return values, nil
	}
	var runningTotal int
	for _, v := range vals {
		s := v.(string)
		log.Debug("V is: ", s)
		splitValues := strings.Split(s, ".")
		if len(splitValues) > 1 {
			vInt, err := strconv.Atoi(splitValues[1])
			if err != nil {
				log.Error("Couldn't convert tracked latency value to Int, vl is: ", err)
			} else {
				runningTotal += vInt
			}
		}

	}
	values.AvgUpstreamLatency = roundValue(float64(runningTotal / len(vals)))
	return values, nil
}
<|endoftext|>


<commit_before>package tcp

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/cihub/seelog"
	"github.com/pkg/errors"
	"github.com/schollz/croc/v6/src/comm"
	"github.com/schollz/croc/v6/src/logger"
	"github.com/schollz/croc/v6/src/models"
)

type server struct {
	port       string
	debugLevel string
	banner     string
	rooms      roomMap
}

type roomInfo struct {
	first  *comm.Comm
	second *comm.Comm
	opened time.Time
	full   bool
}

type roomMap struct {
	rooms map[string]roomInfo
	sync.Mutex
}

// Run starts a tcp listener, run async
func Run(debugLevel, port string, banner ...string) (err error) {
	s := new(server)
	s.port = port
	s.debugLevel = debugLevel
	if len(banner) > 0 {
		s.banner = banner[0]
	}
	return s.start()
}

func (s *server) start() (err error) {
	logger.SetLogLevel(s.debugLevel)
	s.rooms.Lock()
	s.rooms.rooms = make(map[string]roomInfo)
	s.rooms.Unlock()

	// delete old rooms
	go func() {
		for {
			time.Sleep(10 * time.Minute)
			s.rooms.Lock()
			for room := range s.rooms.rooms {
				if time.Since(s.rooms.rooms[room].opened) > 3*time.Hour {
					delete(s.rooms.rooms, room)
				}
			}
			s.rooms.Unlock()
		}
	}()

	err = s.run()
	if err != nil {
		log.Error(err)
	}
	return
}

func (s *server) run() (err error) {
	log.Debugf("starting TCP server on " + s.port)
	server, err := net.Listen("tcp", ":"+s.port)
	if err != nil {
		return errors.Wrap(err, "Error listening on :"+s.port)
	}
	defer server.Close()
	// spawn a new goroutine whenever a client connects
	for {
		connection, err := server.Accept()
		if err != nil {
			return errors.Wrap(err, "problem accepting connection")
		}
		log.Debugf("client %s connected", connection.RemoteAddr().String())
		go func(port string, connection net.Conn) {
			errCommunication := s.clientCommuncation(port, comm.New(connection))
			if errCommunication != nil {
				log.Warnf("relay-%s: %s", connection.RemoteAddr().String(), errCommunication.Error())
			}
		}(s.port, connection)
	}
}

func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
	// send ok to tell client they are connected
	log.Debug("sending ok")
	err = c.Send([]byte(s.banner))
	if err != nil {
		return
	}

	// wait for client to tell me which room they want
	log.Debug("waiting for answer")
	roomBytes, err := c.Receive()
	if err != nil {
		return
	}
	room := string(roomBytes)

	s.rooms.Lock()
	// create the room if it is new
	if _, ok := s.rooms.rooms[room]; !ok {
		s.rooms.rooms[room] = roomInfo{
			first:  c,
			opened: time.Now(),
		}
		s.rooms.Unlock()
		// tell the client that they got the room
		err = c.Send([]byte("ok"))
		if err != nil {
			log.Error(err)
			s.deleteRoom(room)
			return
		}
		log.Debugf("room %s has 1", room)
		return nil
	}
	if s.rooms.rooms[room].full {
		s.rooms.Unlock()
		err = c.Send([]byte("room full"))
		if err != nil {
			log.Error(err)
			s.deleteRoom(room)
			return
		}
		return nil
	}
	log.Debugf("room %s has 2", room)
	s.rooms.rooms[room] = roomInfo{
		first:  s.rooms.rooms[room].first,
		second: c,
		opened: s.rooms.rooms[room].opened,
		full:   true,
	}
	otherConnection := s.rooms.rooms[room].first
	s.rooms.Unlock()

	// second connection is the sender, time to staple connections
	var wg sync.WaitGroup
	wg.Add(1)

	// start piping
	go func(com1, com2 *comm.Comm, wg *sync.WaitGroup) {
		log.Debug("starting pipes")
		pipe(com1.Connection(), com2.Connection())
		wg.Done()
		log.Debug("done piping")
	}(otherConnection, c, &wg)

	// tell the sender everything is ready
	err = c.Send([]byte("ok"))
	if err != nil {
		s.deleteRoom(room)
		return
	}
	wg.Wait()

	// delete room
	s.deleteRoom(room)
	return nil
}

func (s *server) deleteRoom(room string) {
	s.rooms.Lock()
	defer s.rooms.Unlock()
	if _, ok := s.rooms.rooms[room]; !ok {
		return
	}
	log.Debugf("deleting room: %s", room)
	s.rooms.rooms[room].first.Close()
	s.rooms.rooms[room].second.Close()
	s.rooms.rooms[room] = roomInfo{first: nil, second: nil}
	delete(s.rooms.rooms, room)

}

// chanFromConn creates a channel from a Conn object, and sends everything it
//  Read()s from the socket to the channel.
func chanFromConn(conn net.Conn) chan []byte {
	c := make(chan []byte, 1)

	go func() {
		b := make([]byte, models.TCP_BUFFER_SIZE)

		for {
			n, err := conn.Read(b)
			if n > 0 {
				res := make([]byte, n)
				// Copy the buffer so it doesn't get changed while read by the recipient.
				copy(res, b[:n])
				c <- res
			}
			if err != nil {
				log.Debug(err)
				c <- nil
				break
			}
		}
		log.Debug("exiting")
	}()

	return c
}

// pipe creates a full-duplex pipe between the two sockets and
// transfers data from one to the other.
func pipe(conn1 net.Conn, conn2 net.Conn) {
	chan1 := chanFromConn(conn1)
	chan2 := chanFromConn(conn2)

	for {
		select {
		case b1 := <-chan1:
			if b1 == nil {
				return
			}
			conn2.Write(b1)

		case b2 := <-chan2:
			if b2 == nil {
				return
			}
			conn1.Write(b2)
		}
	}
}

func ConnectToTCPServer(address, room string) (c *comm.Comm, banner string, err error) {
	c, err = comm.NewConnection(address)
	if err != nil {
		return
	}
	log.Debug("waiting for first ok")
	data, err := c.Receive()
	if err != nil {
		return
	}
	banner = string(data)
	log.Debug("sending room")
	err = c.Send([]byte(room))
	if err != nil {
		return
	}
	log.Debug("waiting for room confirmation")
	data, err = c.Receive()
	if err != nil {
		return
	}
	if !bytes.Equal(data, []byte("ok")) {
		err = fmt.Errorf("got bad response: %s", data)
		return
	}
	log.Debug("all set")
	return
}
<commit_msg>change message<commit_after>package tcp

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/cihub/seelog"
	"github.com/pkg/errors"
	"github.com/schollz/croc/v6/src/comm"
	"github.com/schollz/croc/v6/src/logger"
	"github.com/schollz/croc/v6/src/models"
)

type server struct {
	port       string
	debugLevel string
	banner     string
	rooms      roomMap
}

type roomInfo struct {
	first  *comm.Comm
	second *comm.Comm
	opened time.Time
	full   bool
}

type roomMap struct {
	rooms map[string]roomInfo
	sync.Mutex
}

// Run starts a tcp listener, run async
func Run(debugLevel, port string, banner ...string) (err error) {
	s := new(server)
	s.port = port
	s.debugLevel = debugLevel
	if len(banner) > 0 {
		s.banner = banner[0]
	}
	return s.start()
}

func (s *server) start() (err error) {
	logger.SetLogLevel(s.debugLevel)
	s.rooms.Lock()
	s.rooms.rooms = make(map[string]roomInfo)
	s.rooms.Unlock()

	// delete old rooms
	go func() {
		for {
			time.Sleep(10 * time.Minute)
			s.rooms.Lock()
			for room := range s.rooms.rooms {
				if time.Since(s.rooms.rooms[room].opened) > 3*time.Hour {
					delete(s.rooms.rooms, room)
				}
			}
			s.rooms.Unlock()
		}
	}()

	err = s.run()
	if err != nil {
		log.Error(err)
	}
	return
}

func (s *server) run() (err error) {
	log.Debugf("starting TCP server on " + s.port)
	server, err := net.Listen("tcp", ":"+s.port)
	if err != nil {
		return errors.Wrap(err, "Error listening on :"+s.port)
	}
	defer server.Close()
	// spawn a new goroutine whenever a client connects
	for {
		connection, err := server.Accept()
		if err != nil {
			return errors.Wrap(err, "problem accepting connection")
		}
		log.Debugf("client %s connected", connection.RemoteAddr().String())
		go func(port string, connection net.Conn) {
			errCommunication := s.clientCommuncation(port, comm.New(connection))
			if errCommunication != nil {
				log.Warnf("relay-%s: %s", connection.RemoteAddr().String(), errCommunication.Error())
			}
		}(s.port, connection)
	}
}

func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
	// send ok to tell client they are connected
	log.Debug("sending " + s.banner)
	err = c.Send([]byte(s.banner))
	if err != nil {
		return
	}

	// wait for client to tell me which room they want
	log.Debug("waiting for answer")
	roomBytes, err := c.Receive()
	if err != nil {
		return
	}
	room := string(roomBytes)

	s.rooms.Lock()
	// create the room if it is new
	if _, ok := s.rooms.rooms[room]; !ok {
		s.rooms.rooms[room] = roomInfo{
			first:  c,
			opened: time.Now(),
		}
		s.rooms.Unlock()
		// tell the client that they got the room
		err = c.Send([]byte("ok"))
		if err != nil {
			log.Error(err)
			s.deleteRoom(room)
			return
		}
		log.Debugf("room %s has 1", room)
		return nil
	}
	if s.rooms.rooms[room].full {
		s.rooms.Unlock()
		err = c.Send([]byte("room full"))
		if err != nil {
			log.Error(err)
			s.deleteRoom(room)
			return
		}
		return nil
	}
	log.Debugf("room %s has 2", room)
	s.rooms.rooms[room] = roomInfo{
		first:  s.rooms.rooms[room].first,
		second: c,
		opened: s.rooms.rooms[room].opened,
		full:   true,
	}
	otherConnection := s.rooms.rooms[room].first
	s.rooms.Unlock()

	// second connection is the sender, time to staple connections
	var wg sync.WaitGroup
	wg.Add(1)

	// start piping
	go func(com1, com2 *comm.Comm, wg *sync.WaitGroup) {
		log.Debug("starting pipes")
		pipe(com1.Connection(), com2.Connection())
		wg.Done()
		log.Debug("done piping")
	}(otherConnection, c, &wg)

	// tell the sender everything is ready
	err = c.Send([]byte("ok"))
	if err != nil {
		s.deleteRoom(room)
		return
	}
	wg.Wait()

	// delete room
	s.deleteRoom(room)
	return nil
}

func (s *server) deleteRoom(room string) {
	s.rooms.Lock()
	defer s.rooms.Unlock()
	if _, ok := s.rooms.rooms[room]; !ok {
		return
	}
	log.Debugf("deleting room: %s", room)
	s.rooms.rooms[room].first.Close()
	s.rooms.rooms[room].second.Close()
	s.rooms.rooms[room] = roomInfo{first: nil, second: nil}
	delete(s.rooms.rooms, room)

}

// chanFromConn creates a channel from a Conn object, and sends everything it
//  Read()s from the socket to the channel.
func chanFromConn(conn net.Conn) chan []byte {
	c := make(chan []byte, 1)

	go func() {
		b := make([]byte, models.TCP_BUFFER_SIZE)

		for {
			n, err := conn.Read(b)
			if n > 0 {
				res := make([]byte, n)
				// Copy the buffer so it doesn't get changed while read by the recipient.
				copy(res, b[:n])
				c <- res
			}
			if err != nil {
				log.Debug(err)
				c <- nil
				break
			}
		}
		log.Debug("exiting")
	}()

	return c
}

// pipe creates a full-duplex pipe between the two sockets and
// transfers data from one to the other.
func pipe(conn1 net.Conn, conn2 net.Conn) {
	chan1 := chanFromConn(conn1)
	chan2 := chanFromConn(conn2)

	for {
		select {
		case b1 := <-chan1:
			if b1 == nil {
				return
			}
			conn2.Write(b1)

		case b2 := <-chan2:
			if b2 == nil {
				return
			}
			conn1.Write(b2)
		}
	}
}

func ConnectToTCPServer(address, room string) (c *comm.Comm, banner string, err error) {
	c, err = comm.NewConnection(address)
	if err != nil {
		return
	}
	log.Debug("waiting for first ok")
	data, err := c.Receive()
	if err != nil {
		return
	}
	banner = string(data)
	log.Debug("sending room")
	err = c.Send([]byte(room))
	if err != nil {
		return
	}
	log.Debug("waiting for room confirmation")
	data, err = c.Receive()
	if err != nil {
		return
	}
	if !bytes.Equal(data, []byte("ok")) {
		err = fmt.Errorf("got bad response: %s", data)
		return
	}
	log.Debug("all set")
	return
}
<|endoftext|>


<commit_before>package gtk

// #include <gtk/gtk.h>
// #include "gtk.go.h"
import "C"
import (
"unsafe"

"github.com/gotk3/gotk3/glib"
)

func init() {
	tm := []glib.TypeMarshaler{
		{glib.Type(C.gtk_font_chooser_get_type()), marshalFontChooser},
		{glib.Type(C.gtk_font_button_get_type()), marshalFontButton},
	}

	glib.RegisterGValueMarshalers(tm)

	WrapMap["GtkFontChooser"] = wrapFontChooser
	WrapMap["GtkFontButton"] = wrapFontButton

}

/*
 * GtkFontChooser
 */

// FontChooser is a representation of GTK's GtkFontChooser GInterface.
type FontChooser struct {
	*glib.Object
}

// IFontChooser is an interface type implemented by all structs
// embedding an FontChooser. It is meant to be used as an argument type
// for wrapper functions that wrap around a C GTK function taking a
// GtkFontChooser.
type IFontChooser interface {
	toFontChooser() *C.GtkFontChooser
}

// native returns a pointer to the underlying GtkFontChooser.
func (v *FontChooser) native() *C.GtkFontChooser {
	if v == nil || v.GObject == nil {
		return nil
	}
	p := unsafe.Pointer(v.GObject)
	return C.toGtkFontChooser(p)
}

func marshalFontChooser(p uintptr) (interface{}, error) {
	c := C.g_value_get_object((*C.GValue)(unsafe.Pointer(p)))
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontChooser(obj), nil
}

func wrapFontChooser(obj *glib.Object) *FontChooser {
	return &FontChooser{obj}
}

func (v *FontChooser) toFontChooser() *C.GtkFontChooser {
	if v == nil {
		return nil
	}
	return v.native()
}

// GetFont is a wrapper around gtk_font_chooser_get_font().
func (v *FontChooser) GetFont() string {
	c := C.gtk_font_chooser_get_font(v.native())
	return goString(c)
}

// SetFont is a wrapper around gtk_font_chooser_set_font().
func (v *FontChooser) SetFont(font string) {
	cstr := C.CString(font)
	defer C.free(unsafe.Pointer(cstr))
	C.gtk_font_chooser_set_font(v.native(), (*C.gchar)(cstr))
}

//PangoFontFamily *	gtk_font_chooser_get_font_family ()
//PangoFontFace *	gtk_font_chooser_get_font_face ()
//gint	gtk_font_chooser_get_font_size ()
//PangoFontDescription *	gtk_font_chooser_get_font_desc ()
//void	gtk_font_chooser_set_font_desc ()
//gchar *	gtk_font_chooser_get_preview_text ()
//void	gtk_font_chooser_set_preview_text ()
//gboolean	gtk_font_chooser_get_show_preview_entry ()
//void	gtk_font_chooser_set_show_preview_entry ()
//gboolean	(*GtkFontFilterFunc) ()
//void	gtk_font_chooser_set_filter_func ()
//void	gtk_font_chooser_set_font_map ()
//PangoFontMap *	gtk_font_chooser_get_font_map ()

/*
 * GtkFontButton
 */

// FontButton is a representation of GTK's GtkFontButton.
type FontButton struct {
	Button

	// Interfaces
	FontChooser
}

// native returns a pointer to the underlying GtkFontButton.
func (v *FontButton) native() *C.GtkFontButton {
	if v == nil || v.GObject == nil {
		return nil
	}
	p := unsafe.Pointer(v.GObject)
	return C.toGtkFontButton(p)
}

func marshalFontButton(p uintptr) (interface{}, error) {
	c := C.g_value_get_object((*C.GValue)(unsafe.Pointer(p)))
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontButton(obj), nil
}

func wrapFontButton(obj *glib.Object) *FontButton {
	button := wrapButton(obj)
	fc := wrapFontChooser(obj)
	return &FontButton{*button,*fc}
}

// FontButtonNew is a wrapper around gtk_font_button_new().
func FontButtonNew() (*FontButton, error) {
	c := C.gtk_font_button_new()
	if c == nil {
		return nil, nilPtrErr
	}
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontButton(obj), nil
}

// FontButtonNewWithFont is a wrapper around gtk_font_button_new_with_font().
func FontButtonNewWithFont(fontname string) (*FontButton, error) {
	cstr := C.CString(fontname)
	defer C.free(unsafe.Pointer(cstr))
	c := C.gtk_font_button_new_with_font((*C.gchar)(cstr))
	if c == nil {
		return nil, nilPtrErr
	}
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontButton(obj), nil
}
<commit_msg>go fmt<commit_after>package gtk

// #include <gtk/gtk.h>
// #include "gtk.go.h"
import "C"
import (
	"unsafe"

	"github.com/gotk3/gotk3/glib"
)

func init() {
	tm := []glib.TypeMarshaler{
		{glib.Type(C.gtk_font_chooser_get_type()), marshalFontChooser},
		{glib.Type(C.gtk_font_button_get_type()), marshalFontButton},
	}

	glib.RegisterGValueMarshalers(tm)

	WrapMap["GtkFontChooser"] = wrapFontChooser
	WrapMap["GtkFontButton"] = wrapFontButton

}

/*
 * GtkFontChooser
 */

// FontChooser is a representation of GTK's GtkFontChooser GInterface.
type FontChooser struct {
	*glib.Object
}

// IFontChooser is an interface type implemented by all structs
// embedding an FontChooser. It is meant to be used as an argument type
// for wrapper functions that wrap around a C GTK function taking a
// GtkFontChooser.
type IFontChooser interface {
	toFontChooser() *C.GtkFontChooser
}

// native returns a pointer to the underlying GtkFontChooser.
func (v *FontChooser) native() *C.GtkFontChooser {
	if v == nil || v.GObject == nil {
		return nil
	}
	p := unsafe.Pointer(v.GObject)
	return C.toGtkFontChooser(p)
}

func marshalFontChooser(p uintptr) (interface{}, error) {
	c := C.g_value_get_object((*C.GValue)(unsafe.Pointer(p)))
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontChooser(obj), nil
}

func wrapFontChooser(obj *glib.Object) *FontChooser {
	return &FontChooser{obj}
}

func (v *FontChooser) toFontChooser() *C.GtkFontChooser {
	if v == nil {
		return nil
	}
	return v.native()
}

// GetFont is a wrapper around gtk_font_chooser_get_font().
func (v *FontChooser) GetFont() string {
	c := C.gtk_font_chooser_get_font(v.native())
	return goString(c)
}

// SetFont is a wrapper around gtk_font_chooser_set_font().
func (v *FontChooser) SetFont(font string) {
	cstr := C.CString(font)
	defer C.free(unsafe.Pointer(cstr))
	C.gtk_font_chooser_set_font(v.native(), (*C.gchar)(cstr))
}

//PangoFontFamily *	gtk_font_chooser_get_font_family ()
//PangoFontFace *	gtk_font_chooser_get_font_face ()
//gint	gtk_font_chooser_get_font_size ()
//PangoFontDescription *	gtk_font_chooser_get_font_desc ()
//void	gtk_font_chooser_set_font_desc ()
//gchar *	gtk_font_chooser_get_preview_text ()
//void	gtk_font_chooser_set_preview_text ()
//gboolean	gtk_font_chooser_get_show_preview_entry ()
//void	gtk_font_chooser_set_show_preview_entry ()
//gboolean	(*GtkFontFilterFunc) ()
//void	gtk_font_chooser_set_filter_func ()
//void	gtk_font_chooser_set_font_map ()
//PangoFontMap *	gtk_font_chooser_get_font_map ()

/*
 * GtkFontButton
 */

// FontButton is a representation of GTK's GtkFontButton.
type FontButton struct {
	Button

	// Interfaces
	FontChooser
}

// native returns a pointer to the underlying GtkFontButton.
func (v *FontButton) native() *C.GtkFontButton {
	if v == nil || v.GObject == nil {
		return nil
	}
	p := unsafe.Pointer(v.GObject)
	return C.toGtkFontButton(p)
}

func marshalFontButton(p uintptr) (interface{}, error) {
	c := C.g_value_get_object((*C.GValue)(unsafe.Pointer(p)))
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontButton(obj), nil
}

func wrapFontButton(obj *glib.Object) *FontButton {
	button := wrapButton(obj)
	fc := wrapFontChooser(obj)
	return &FontButton{*button, *fc}
}

// FontButtonNew is a wrapper around gtk_font_button_new().
func FontButtonNew() (*FontButton, error) {
	c := C.gtk_font_button_new()
	if c == nil {
		return nil, nilPtrErr
	}
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontButton(obj), nil
}

// FontButtonNewWithFont is a wrapper around gtk_font_button_new_with_font().
func FontButtonNewWithFont(fontname string) (*FontButton, error) {
	cstr := C.CString(fontname)
	defer C.free(unsafe.Pointer(cstr))
	c := C.gtk_font_button_new_with_font((*C.gchar)(cstr))
	if c == nil {
		return nil, nilPtrErr
	}
	obj := glib.Take(unsafe.Pointer(c))
	return wrapFontButton(obj), nil
}
<|endoftext|>


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

package server

import (
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/tidb"
)

type TidbTestSuite struct {
	tidbdrv *TiDBDriver
	server  *Server
}

var _ = Suite(new(TidbTestSuite))

func (ts *TidbTestSuite) SetUpSuite(c *C) {
	store, err := tidb.NewStore("memory:///tmp/tidb")
	c.Assert(err, IsNil)
	ts.tidbdrv = NewTiDBDriver(store)
	cfg := &Config{
		Addr:     ":4001",
		LogLevel: "debug",
	}
	server, err := NewServer(cfg, ts.tidbdrv)
	c.Assert(err, IsNil)
	ts.server = server
	go ts.server.Run()
	time.Sleep(time.Millisecond * 100)
}

func (ts *TidbTestSuite) TearDownSuite(c *C) {
	ts.server.Close()
}

func (ts *TidbTestSuite) TestRegression(c *C) {
	if regression {
		runTestRegression(c)
	}
}

func (ts *TidbTestSuite) TestUint64(c *C) {
	runTestPrepareResultFieldType(c)
}

func (ts *TidbTestSuite) TestSpecialType(c *C) {
	runTestSpecialType(c)
}

func (ts *TidbTestSuite) TestPreparedString(c *C) {
	runTestPreparedString(c)
}

func (ts *TidbTestSuite) TestConcurrentUpdate(c *C) {
	runTestConcurrentUpdate(c)
}

func (ts *TidbTestSuite) TestErrorCode(c *C) {
	runTestErrorCode(c)
}

func (ts *TidbTestSuite) TestAuth(c *C) {
	runTestAuth(c)
}

func (ts *TidbTestSuite) TestIssues(c *C) {
	runTestIssues(c)
}
<commit_msg>tide-server: avoid panic when failed in SetupSuite.<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/tidb"
)

type TidbTestSuite struct {
	tidbdrv *TiDBDriver
	server  *Server
}

var _ = Suite(new(TidbTestSuite))

func (ts *TidbTestSuite) SetUpSuite(c *C) {
	store, err := tidb.NewStore("memory:///tmp/tidb")
	c.Assert(err, IsNil)
	ts.tidbdrv = NewTiDBDriver(store)
	cfg := &Config{
		Addr:     ":4001",
		LogLevel: "debug",
	}
	server, err := NewServer(cfg, ts.tidbdrv)
	c.Assert(err, IsNil)
	ts.server = server
	go ts.server.Run()
	time.Sleep(time.Millisecond * 100)
}

func (ts *TidbTestSuite) TearDownSuite(c *C) {
	if ts.server != nil {
		ts.server.Close()
	}
}

func (ts *TidbTestSuite) TestRegression(c *C) {
	if regression {
		runTestRegression(c)
	}
}

func (ts *TidbTestSuite) TestUint64(c *C) {
	runTestPrepareResultFieldType(c)
}

func (ts *TidbTestSuite) TestSpecialType(c *C) {
	runTestSpecialType(c)
}

func (ts *TidbTestSuite) TestPreparedString(c *C) {
	runTestPreparedString(c)
}

func (ts *TidbTestSuite) TestConcurrentUpdate(c *C) {
	runTestConcurrentUpdate(c)
}

func (ts *TidbTestSuite) TestErrorCode(c *C) {
	runTestErrorCode(c)
}

func (ts *TidbTestSuite) TestAuth(c *C) {
	runTestAuth(c)
}

func (ts *TidbTestSuite) TestIssues(c *C) {
	runTestIssues(c)
}
<|endoftext|>


<commit_before>package rollbar

import (
	"encoding/json"
	"fmt"
)

type ItemsService struct {
	C *Client
}

type ItemsResponse struct {
	Err    int         `json:"err"`
	Result ItemsResult `json:"result"`
}

type SingleItemResponse struct {
	Err    int  `json:"err"`
	Result Item `json:"result"`
}

type ItemsResult struct {
	Page  int    `json:"page"`
	Items []Item `json:"items"`
}

type Item struct {
	Id                       int    `json:"id"`
	ProjectId                int    `json:"project_id"`
	Title                    string `json:"title"`
	LastOccurrenceId         int    `json:"last_occurrence_id"`
	LastOccurrenceTimestamp  int64  `json:"last_occurrence_timestamp"`
	TotalOccurrences         int    `json:"total_occurrences"`
	FirstOccurrenceId        int    `json:"first_occurrence_id"`
	FirstOccurrenceTimestamp int64  `json:"first_occurrence_timestamp"`
	Status                   string `json:"status"`
	Level                    string `json:"level"`
}

func (i *ItemsService) All() (*ItemsResponse, error) {
	response := &ItemsResponse{}

	body, err := i.C.Request("GET", "items")
	if err != nil {
		return response, err
	}

	defer body.Close()

	err = json.NewDecoder(body).Decode(&response)
	if err != nil {
		return response, err
	}

	return response, nil
}

func (i *ItemsService) GetItem(itemId int) (*SingleItemResponse, error) {
	response := &SingleItemResponse{}

	var url = fmt.Sprintf("items/%v", itemId)
	body, err := i.C.Request("GET", url)
	if err != nil {
		return response, err
	}

	defer body.Close()

	err = json.NewDecoder(body).Decode(&response)
	if err != nil {
		return response, err
	}

	return response, nil
}
<commit_msg>rollbar: fix get items route<commit_after>package rollbar

import (
	"encoding/json"
	"fmt"
)

type ItemsService struct {
	C *Client
}

type ItemsResponse struct {
	Err    int         `json:"err"`
	Result ItemsResult `json:"result"`
}

type SingleItemResponse struct {
	Err    int  `json:"err"`
	Result Item `json:"result"`
}

type ItemsResult struct {
	Page  int    `json:"page"`
	Items []Item `json:"items"`
}

type Item struct {
	Id                       int    `json:"id"`
	ProjectId                int    `json:"project_id"`
	Title                    string `json:"title"`
	LastOccurrenceId         int    `json:"last_occurrence_id"`
	LastOccurrenceTimestamp  int64  `json:"last_occurrence_timestamp"`
	TotalOccurrences         int    `json:"total_occurrences"`
	FirstOccurrenceId        int    `json:"first_occurrence_id"`
	FirstOccurrenceTimestamp int64  `json:"first_occurrence_timestamp"`
	Status                   string `json:"status"`
	Level                    string `json:"level"`
}

func (i *ItemsService) All() (*ItemsResponse, error) {
	response := &ItemsResponse{}

	body, err := i.C.Request("GET", "items")
	if err != nil {
		return response, err
	}

	defer body.Close()

	err = json.NewDecoder(body).Decode(&response)
	if err != nil {
		return response, err
	}

	return response, nil
}

func (i *ItemsService) GetItem(itemId int) (*SingleItemResponse, error) {
	response := &SingleItemResponse{}

	var url = fmt.Sprintf("item/%v", itemId)
	body, err := i.C.Request("GET", url)
	if err != nil {
		return response, err
	}

	defer body.Close()

	err = json.NewDecoder(body).Decode(&response)
	if err != nil {
		return response, err
	}

	return response, nil
}
<|endoftext|>


<commit_before>package models

import (
	"socialapi/workers/common/runner"
	"testing"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
)

func TestInteractiongetAccountId(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while getting account id", t, func() {
		Convey("it should have error if interaction id is not set", func() {
			i := NewInteraction()

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(in, ShouldEqual, 0)
		})

		Convey("it should have error if account is not set in db", func() {
			i := NewInteraction()
			i.Id = 4590

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(in, ShouldEqual, 0)
		})

		Convey("it should return quickly account id if id is set", func() {
			i := NewInteraction()
			i.AccountId = 1020

			in, err := i.getAccountId()
			So(err, ShouldBeNil)
			So(in, ShouldEqual, i.AccountId)
		})
	})
}

func TestInteractionisExempt(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While testing interaction is exempt or not", t, func() {
		Convey("it should have error while getting account id from db when channel id is not set", func() {
			i := NewInteraction()

			ie, err := i.isExempt()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(ie, ShouldEqual, false)
		})

		Convey("it should have error if account id is not set", func() {
			i := NewInteraction()
			i.Id = 1098

			ie, err := i.isExempt()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(ie, ShouldEqual, false)
		})

		Convey("it should return true if account is troll", func() {
			// create troll account
			accTroll := createAccountWithTest()
			err := accTroll.MarkAsTroll()
			So(err, ShouldBeNil)

			i := NewInteraction()
			i.AccountId = accTroll.Id

			ie, err := i.isExempt()
			So(err, ShouldBeNil)
			So(ie, ShouldEqual, true)
		})

		Convey("it should return false if account is not troll", func() {
			// create account
			acc := createAccountWithTest()

			i := NewInteraction()
			i.AccountId = acc.Id

			ie, err := i.isExempt()
			So(err, ShouldBeNil)
			So(ie, ShouldEqual, false)
		})

	})
}

func TestInteractionMarkIfExempt(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While marking if interaction isexempt ", t, func() {
		Convey("it should return nil if exempt", func() {
			accTroll := createAccountWithTest()
			err := accTroll.MarkAsTroll()
			So(err, ShouldBeNil)

			msg := createMessageWithTest()
			So(msg.Create(), ShouldBeNil)

			i := NewInteraction()
			i.AccountId = accTroll.Id
			i.MessageId = msg.Id

			errs := i.MarkIfExempt()
			So(errs, ShouldBeNil)
		})

	})
}

func TestInteractionIsInteracted(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while testing account is interacted", t, func() {
		Convey("it should have error if message id is not set", func() {
			i := NewInteraction()

			cnt, err := i.IsInteracted(0)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrMessageIdIsNotSet)
			So(cnt, ShouldEqual, false)
		})

		Convey("it should have ", func() {
			i := NewInteraction()
			i.MessageId = 1050

			cnt, err := i.IsInteracted(0)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
			So(cnt, ShouldEqual, false)
		})

	})
}
<commit_msg>social/tests: isInteraction tests are added<commit_after>package models

import (
	"socialapi/workers/common/runner"
	"testing"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
)

func TestInteractiongetAccountId(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while getting account id", t, func() {
		Convey("it should have error if interaction id is not set", func() {
			i := NewInteraction()

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(in, ShouldEqual, 0)
		})

		Convey("it should have error if account is not set in db", func() {
			i := NewInteraction()
			i.Id = 4590

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(in, ShouldEqual, 0)
		})

		Convey("it should return quickly account id if id is set", func() {
			i := NewInteraction()
			i.AccountId = 1020

			in, err := i.getAccountId()
			So(err, ShouldBeNil)
			So(in, ShouldEqual, i.AccountId)
		})
	})
}

func TestInteractionisExempt(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While testing interaction is exempt or not", t, func() {
		Convey("it should have error while getting account id from db when channel id is not set", func() {
			i := NewInteraction()

			ie, err := i.isExempt()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(ie, ShouldEqual, false)
		})

		Convey("it should have error if account id is not set", func() {
			i := NewInteraction()
			i.Id = 1098

			ie, err := i.isExempt()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(ie, ShouldEqual, false)
		})

		Convey("it should return true if account is troll", func() {
			// create troll account
			accTroll := createAccountWithTest()
			err := accTroll.MarkAsTroll()
			So(err, ShouldBeNil)

			i := NewInteraction()
			i.AccountId = accTroll.Id

			ie, err := i.isExempt()
			So(err, ShouldBeNil)
			So(ie, ShouldEqual, true)
		})

		Convey("it should return false if account is not troll", func() {
			// create account
			acc := createAccountWithTest()

			i := NewInteraction()
			i.AccountId = acc.Id

			ie, err := i.isExempt()
			So(err, ShouldBeNil)
			So(ie, ShouldEqual, false)
		})

	})
}

func TestInteractionMarkIfExempt(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While marking if interaction isexempt ", t, func() {
		Convey("it should return nil if exempt", func() {
			accTroll := createAccountWithTest()
			err := accTroll.MarkAsTroll()
			So(err, ShouldBeNil)

			msg := createMessageWithTest()
			So(msg.Create(), ShouldBeNil)

			i := NewInteraction()
			i.AccountId = accTroll.Id
			i.MessageId = msg.Id

			errs := i.MarkIfExempt()
			So(errs, ShouldBeNil)
		})

	})
}

func TestInteractionIsInteracted(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while testing account is interacted", t, func() {
		Convey("it should have error if message id is not set", func() {
			i := NewInteraction()

			cnt, err := i.IsInteracted(0)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrMessageIdIsNotSet)
			So(cnt, ShouldEqual, false)
		})

		Convey("it should have error if account id is not set", func() {
			i := NewInteraction()
			i.MessageId = 1050

			cnt, err := i.IsInteracted(0)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
			So(cnt, ShouldEqual, false)
		})

	})
}
<|endoftext|>


<commit_before>package yazawa

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/jue58/kana"
	"github.com/shogo82148/go-mecab"
)

// ToYazawa convert text to yazawanized sentence
func ToYazawa(text string, random bool) string {
	tagger, err := mecab.New(map[string]string{})
	if err != nil {
		panic(err)
	}
	defer tagger.Destroy()

	tagger.Parse("")

	node, err := tagger.ParseToNode(text)
	if err != nil {
		panic(err)
	}

	id := findSuitableIndexForReplace(node, random)
	convertedSentence := []string{}
	for ; node != (mecab.Node{}); node = node.Next() {
		if node.Id() == id {
			feature := strings.Split(node.Feature(), ",")
			convertedSentence = append(convertedSentence, "『"+strings.ToUpper(kana.KanaToRomaji(feature[len(feature)-2]))+"』")
		} else {
			convertedSentence = append(convertedSentence, node.Surface())
		}
	}

	return strings.Join(convertedSentence, "")
}

func findSuitableIndexForReplace(node mecab.Node, random bool) int {
	idForReplace := 0
	maxScore := 0
	for ; node != (mecab.Node{}); node = node.Next() {
		score := 0
		word := node.Surface()
		score += examineWord(word)
		partsOfSpeech := strings.Split(node.Feature(), ",")[0]

		if partsOfSpeech == "形容詞" {
			score += 20
		} else if partsOfSpeech == "名詞" {
			score += 10
		} else if partsOfSpeech == "動詞" {
			score += 8
		}

		if random {
			rand.Seed(time.Now().UnixNano())
			score += rand.Intn(20)
		} else {
			score += utf8.RuneCountInString(word)
		}

		if score > maxScore {
			maxScore = score
			idForReplace = node.Id()
		}
	}
	return idForReplace
}

func examineWord(word string) int {
	score := 0
	hasKatakana := false
	hasKanji := false
	for _, r := range word {
		if !hasKatakana && kana.IsKatakana(fmt.Sprintf("%c", r)) {
			hasKatakana = true
		}
		if !hasKanji && kana.IsKanji(fmt.Sprintf("%c", r)) {
			hasKanji = true
		}
	}
	if hasKatakana {
		score += 10
	}
	if hasKanji {
		score += 10
	}
	if hasKatakana || hasKanji {
		score += 100
	}
	return score
}
<commit_msg>WIP Chnage library<commit_after>package yazawa

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/gojp/kana"
	"github.com/shogo82148/go-mecab"
)

// ToYazawa convert text to yazawanized sentence
func ToYazawa(text string, random bool) string {
	tagger, err := mecab.New(map[string]string{})
	if err != nil {
		panic(err)
	}
	defer tagger.Destroy()

	tagger.Parse("")

	node, err := tagger.ParseToNode(text)
	if err != nil {
		panic(err)
	}

	id := findSuitableIndexForReplace(node, random)
	convertedSentence := []string{}
	for ; node != (mecab.Node{}); node = node.Next() {
		if node.Id() == id {
			feature := strings.Split(node.Feature(), ",")
			convertedSentence = append(convertedSentence, "『"+strings.ToUpper(kana.KanaToRomaji(feature[len(feature)-2]))+"』")
		} else {
			convertedSentence = append(convertedSentence, node.Surface())
		}
	}

	return strings.Join(convertedSentence, "")
}

func findSuitableIndexForReplace(node mecab.Node, random bool) int {
	idForReplace := 0
	maxScore := 0
	for ; node != (mecab.Node{}); node = node.Next() {
		score := 0
		word := node.Surface()
		score += examineWord(word)
		partsOfSpeech := strings.Split(node.Feature(), ",")[0]

		if partsOfSpeech == "形容詞" {
			score += 20
		} else if partsOfSpeech == "名詞" {
			score += 10
		} else if partsOfSpeech == "動詞" {
			score += 8
		}

		if random {
			rand.Seed(time.Now().UnixNano())
			score += rand.Intn(20)
		} else {
			score += utf8.RuneCountInString(word)
		}

		if score > maxScore {
			maxScore = score
			idForReplace = node.Id()
		}
	}
	return idForReplace
}

func examineWord(word string) int {
	score := 0
	hasKatakana := false
	hasKanji := false
	for _, r := range word {
		if !hasKatakana && kana.IsKatakana(fmt.Sprintf("%c", r)) {
			hasKatakana = true
		}
		if !hasKanji && kana.IsKanji(fmt.Sprintf("%c", r)) {
			hasKanji = true
		}
	}
	if hasKatakana {
		score += 10
	}
	if hasKanji {
		score += 10
	}
	if hasKatakana || hasKanji {
		score += 100
	}
	return score
}
<|endoftext|>


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

package ec2

import (
	"fmt"
	"launchpad.net/juju-core/environs/imagemetadata"
	"launchpad.net/juju-core/environs/instances"
)

// signedImageDataOnly is defined here to allow tests to override the content.
// If true, only inline PGP signed image metadata will be used.
var signedImageDataOnly = true

// defaultCpuPower is larger the smallest instance's cpuPower, and no larger than
// any other instance type's cpuPower. It is used when no explicit CpuPower
// constraint exists, preventing the smallest instance from being chosen unless
// the user has clearly indicated that they are willing to accept poor performance.
const defaultCpuPower = 100

// findInstanceSpec returns an InstanceSpec satisfying the supplied instanceConstraint.
func findInstanceSpec(baseURLs []string, ic *instances.InstanceConstraint) (*instances.InstanceSpec, error) {
	if ic.Constraints.CpuPower == nil {
		ic.Constraints.CpuPower = instances.CpuPower(defaultCpuPower)
	}
	ec2Region := allRegions[ic.Region]
	imageConstraint := imagemetadata.ImageConstraint{
		CloudSpec: imagemetadata.CloudSpec{ic.Region, ec2Region.EC2Endpoint},
		Series:    ic.Series,
		Arches:    ic.Arches,
	}
	matchingImages, err := imagemetadata.Fetch(
		baseURLs, imagemetadata.DefaultIndexPath, &imageConstraint, signedImageDataOnly)
	if err != nil {
		return nil, err
	}
	var images []instances.Image
	for _, imageMetadata := range matchingImages {
		// For now, we only want images with "ebs" storage.
		if imageMetadata.Storage != ebsStorage {
			continue
		}
		im := *imageMetadata
		images = append(images, instances.Image{
			Id:    im.Id,
			VType: im.VType,
			Arch:  im.Arch,
		})
	}

	// Make a copy of the known EC2 instance types, filling in the cost for the specified region.
	regionCosts := allRegionCosts[ic.Region]
	if len(regionCosts) == 0 && len(allRegionCosts) > 0 {
		return nil, fmt.Errorf("no instance types found in %s", ic.Region)
	}

	var itypesWithCosts []instances.InstanceType
	for _, itype := range allInstanceTypes {
		cost, ok := regionCosts[itype.Name]
		if !ok {
			continue
		}
		itWithCost := itype
		itWithCost.Cost = cost
		itypesWithCosts = append(itypesWithCosts, itWithCost)
	}
	return instances.FindInstanceSpec(images, ic, itypesWithCosts)
}
<commit_msg>[r=dave-cheney] environs/ec2: warn if image meta data missing<commit_after>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package ec2

import (
	"fmt"

	"launchpad.net/juju-core/environs/imagemetadata"
	"launchpad.net/juju-core/environs/instances"
	"launchpad.net/loggo"
)

var logger = loggo.GetLogger("juju.environs.ec2")

// signedImageDataOnly is defined here to allow tests to override the content.
// If true, only inline PGP signed image metadata will be used.
var signedImageDataOnly = true

// defaultCpuPower is larger the smallest instance's cpuPower, and no larger than
// any other instance type's cpuPower. It is used when no explicit CpuPower
// constraint exists, preventing the smallest instance from being chosen unless
// the user has clearly indicated that they are willing to accept poor performance.
const defaultCpuPower = 100

// findInstanceSpec returns an InstanceSpec satisfying the supplied instanceConstraint.
func findInstanceSpec(baseURLs []string, ic *instances.InstanceConstraint) (*instances.InstanceSpec, error) {
	if ic.Constraints.CpuPower == nil {
		ic.Constraints.CpuPower = instances.CpuPower(defaultCpuPower)
	}
	ec2Region := allRegions[ic.Region]
	imageConstraint := imagemetadata.ImageConstraint{
		CloudSpec: imagemetadata.CloudSpec{ic.Region, ec2Region.EC2Endpoint},
		Series:    ic.Series,
		Arches:    ic.Arches,
	}
	matchingImages, err := imagemetadata.Fetch(
		baseURLs, imagemetadata.DefaultIndexPath, &imageConstraint, signedImageDataOnly)
	if err != nil {
		return nil, err
	}
	if len(matchingImages) == 0 {
		logger.Warningf("no matching image meta data for constraints: %v", ic)
	}
	var images []instances.Image
	for _, imageMetadata := range matchingImages {
		// For now, we only want images with "ebs" storage.
		if imageMetadata.Storage != ebsStorage {
			continue
		}
		im := *imageMetadata
		images = append(images, instances.Image{
			Id:    im.Id,
			VType: im.VType,
			Arch:  im.Arch,
		})
	}

	// Make a copy of the known EC2 instance types, filling in the cost for the specified region.
	regionCosts := allRegionCosts[ic.Region]
	if len(regionCosts) == 0 && len(allRegionCosts) > 0 {
		return nil, fmt.Errorf("no instance types found in %s", ic.Region)
	}

	var itypesWithCosts []instances.InstanceType
	for _, itype := range allInstanceTypes {
		cost, ok := regionCosts[itype.Name]
		if !ok {
			continue
		}
		itWithCost := itype
		itWithCost.Cost = cost
		itypesWithCosts = append(itypesWithCosts, itWithCost)
	}
	return instances.FindInstanceSpec(images, ic, itypesWithCosts)
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/raintank/raintank-apps/task-server/model"
)

type agentWithTag struct {
	model.Agent    `xorm:"extends"`
	model.AgentTag `xorm:"extends"`
}

type agentWithTags []*agentWithTag

func (agentWithTags) TableName() string {
	return "agent"
}

func (rows agentWithTags) ToAgentDTO() []*model.AgentDTO {
	agentsById := make(map[int64]*model.AgentDTO)
	for _, r := range rows {
		a, ok := agentsById[r.Agent.Id]
		if !ok {
			tags := make([]string, 0)
			if r.AgentTag.Tag != "" {
				tags = append(tags, r.AgentTag.Tag)
			}
			agentsById[r.Agent.Id] = &model.AgentDTO{
				Id:            r.Agent.Id,
				Name:          r.Agent.Name,
				Enabled:       r.Agent.Enabled,
				EnabledChange: r.Agent.EnabledChange,
				OrgId:         r.Agent.OrgId,
				Public:        r.Agent.Public,
				Online:        r.Agent.Online,
				OnlineChange:  r.Agent.OnlineChange,
				Created:       r.Agent.Created,
				Updated:       r.Agent.Updated,
				Tags:          tags,
			}
		} else if r.Tag != "" {
			a.Tags = append(a.Tags, r.Tag)
		}
	}
	agents := make([]*model.AgentDTO, len(agentsById))
	i := 0
	for _, a := range agentsById {
		agents[i] = a
		i++
	}
	return agents
}

func GetAgents(query *model.GetAgentsQuery) ([]*model.AgentDTO, error) {
	sess, err := newSession(false, "agent")
	if err != nil {
		return nil, err
	}
	return getAgents(sess, query)
}

func getAgents(sess *session, query *model.GetAgentsQuery) ([]*model.AgentDTO, error) {
	var a agentWithTags
	var rawSQL bytes.Buffer
	args := make([]interface{}, 0)

	var where bytes.Buffer
	whereArgs := make([]interface{}, 0)
	prefix := "WHERE"

	fmt.Fprint(&rawSQL, "SELECT agent.*, agent_tag.* FROM agent LEFT JOIN agent_tag ON  agent.id = agent_tag.agent_id ")
	if len(query.Tag) > 0 {
		fmt.Fprint(&rawSQL, "INNER JOIN agent_tag as at ON agent.id = at.agent_id ")
		p := make([]string, len(query.Tag))
		for i, tag := range query.Tag {
			p[i] = "?"
			whereArgs = append(whereArgs, tag)
		}
		filter := fmt.Sprintf("at.tag IN (%s)", strings.Join(p, ","))

		fmt.Fprintf(&where, "%s %s ", prefix, filter)
		prefix = "AND"
	}

	if query.Metric != "" {
		sess.Join("INNER", "agent_metric", "agent_metric.agent_id = agent.id").Where("agent_metric.namespace LIKE ?", query.Metric)
		fmt.Fprint(&rawSQL, "INNER JOIN (SELECT DISTINCT(agent_id) FROM agent_metric where namespace like ?) AS am ON am.agent_id = agent.id ")
		ns := strings.Replace(query.Metric, "*", "%", -1)
		args = append(args, ns)
	}

	if query.Name != "" {
		fmt.Fprintf(&where, "%s agent.name=? ", prefix)
		whereArgs = append(whereArgs, query.Name)
		prefix = "AND"
	}
	if query.Enabled != "" {
		enabled, err := strconv.ParseBool(query.Enabled)
		if err != nil {
			return nil, err
		}
		fmt.Fprintf(&where, "%s agent.enabled=? ", prefix)
		whereArgs = append(whereArgs, enabled)
		prefix = "AND"
	}
	if query.Public != "" {
		public, err := strconv.ParseBool(query.Public)
		if err != nil {
			return nil, err
		}
		if public {
			fmt.Fprintf(&where, "%s agent.public=1 ", prefix)
			prefix = "AND"
		} else {
			fmt.Fprintf(&where, "%s (agent.public=0 AND agent.org_id=?) ", prefix)
			whereArgs = append(whereArgs, query.OrgId)
			prefix = "AND"
		}
	} else {
		fmt.Fprintf(&where, "%s (agent.org_id=? OR agent.public=1) ", prefix)
		whereArgs = append(whereArgs, query.OrgId)
		prefix = "AND"
	}

	if query.OrderBy == "" {
		query.OrderBy = "name"
	}
	if query.Limit == 0 {
		query.Limit = 50
	}
	if query.Page == 0 {
		query.Page = 1
	}

	fmt.Fprint(&rawSQL, where.String())
	args = append(args, whereArgs...)
	fmt.Fprintf(&rawSQL, "ORDER BY `%s` ASC LIMIT %d, %d", query.OrderBy, (query.Page-1)*query.Limit, query.Limit)

	err := sess.Sql(rawSQL.String(), args...).Find(&a)

	if err != nil {
		return nil, err
	}
	return a.ToAgentDTO(), nil
}

func onlineAgentsWithNoSession(sess *session) ([]*model.AgentDTO, error) {
	sess.Table("agent")
	sess.Join("LEFT", "agent_tag", "agent.id=agent_tag.agent_id")
	sess.Join("LEFT", "agent_session", "agent.id = agent_session.agent_id")
	sess.Where("agent.online=1").And("agent_session.id is NULL")
	sess.Cols("`agent`.*", "`agent_tag`.*")
	var a agentWithTags
	err := sess.Find(&a)
	if err != nil {
		return nil, err
	}

	return a.ToAgentDTO(), nil
}

func GetAgentById(id int64, orgId int64) (*model.AgentDTO, error) {
	sess, err := newSession(false, "agent")
	if err != nil {
		return nil, err
	}
	return getAgentById(sess, id, orgId)
}

func getAgentById(sess *session, id int64, orgId int64) (*model.AgentDTO, error) {
	var a agentWithTags
	sess.Where("agent.id=?", id)
	if orgId != 0 {
		sess.And("agent.org_id=?", orgId)
	}
	err := sess.Join("LEFT", "agent_tag", "agent.id = agent_tag.agent_id").Find(&a)
	if err != nil {
		return nil, err
	}
	if len(a) == 0 {
		return nil, model.AgentNotFound
	}
	return a.ToAgentDTO()[0], nil
}

func AddAgent(a *model.AgentDTO) error {
	sess, err := newSession(true, "agent")
	if err != nil {
		return err
	}
	defer sess.Cleanup()
	if err = addAgent(sess, a); err != nil {
		return err
	}
	sess.Complete()
	return nil

}

func addAgent(sess *session, a *model.AgentDTO) error {
	agent := &model.Agent{
		Name:          a.Name,
		Enabled:       a.Enabled,
		EnabledChange: time.Now(),
		OrgId:         a.OrgId,
		Public:        a.Public,
		Online:        false,
		OnlineChange:  time.Now(),
		Created:       time.Now(),
		Updated:       time.Now(),
	}

	sess.UseBool("public")
	sess.UseBool("enabled")
	sess.UseBool("online")
	if _, err := sess.Insert(agent); err != nil {
		return err
	}
	a.Id = agent.Id
	a.Created = agent.Created
	a.Updated = agent.Updated

	agentTags := make([]model.AgentTag, 0, len(a.Tags))
	for _, tag := range a.Tags {
		agentTags = append(agentTags, model.AgentTag{
			OrgId:   a.OrgId,
			AgentId: agent.Id,
			Tag:     tag,
			Created: time.Now(),
		})
	}
	if len(agentTags) > 0 {
		sess.Table("agent_tag")
		if _, err := sess.Insert(&agentTags); err != nil {
			return err
		}
	}
	return nil
}

func UpdateAgent(a *model.AgentDTO) error {
	sess, err := newSession(true, "agent")
	if err != nil {
		return err
	}
	defer sess.Cleanup()

	err = updateAgent(sess, a)
	if err != nil {
		return err
	}
	sess.Complete()
	return err
}

func updateAgent(sess *session, a *model.AgentDTO) error {
	existing, err := getAgentById(sess, a.Id, a.OrgId)
	if err != nil {
		return err
	}
	if existing == nil {
		return model.AgentNotFound
	}
	// If the OrgId is different, the only changes that can be made is to Tags.
	if a.OrgId == existing.OrgId {
		enabledChange := existing.EnabledChange
		if existing.Enabled != a.Enabled {
			enabledChange = time.Now()
		}
		agent := &model.Agent{
			Id:            a.Id,
			Name:          a.Name,
			Enabled:       a.Enabled,
			EnabledChange: enabledChange,
			OrgId:         a.OrgId,
			Public:        a.Public,
			Created:       a.Created,
			Updated:       time.Now(),
		}
		sess.UseBool("public")
		sess.UseBool("enabled")
		if _, err := sess.Id(agent.Id).Update(agent); err != nil {
			return err
		}
		a.Updated = agent.Updated
	}

	tagMap := make(map[string]bool)
	tagsToDelete := make([]string, 0)
	tagsToAddMap := make(map[string]bool, 0)
	// create map of current tags
	for _, t := range existing.Tags {
		tagMap[t] = false
	}

	// create map of tags to add. We use a map
	// to ensure that we only add each tag once.
	for _, t := range a.Tags {
		if _, ok := tagMap[t]; !ok {
			tagsToAddMap[t] = true
		}
		// mark that this tag has been seen.
		tagMap[t] = true
	}

	//create list of tags to delete
	for t, seen := range tagMap {
		if !seen {
			tagsToDelete = append(tagsToDelete, t)
		}
	}

	// create list of tags to add.
	tagsToAdd := make([]string, len(tagsToAddMap))
	i := 0
	for t := range tagsToAddMap {
		tagsToAdd[i] = t
		i += 1
	}
	if len(tagsToDelete) > 0 {
		rawParams := make([]interface{}, 0)
		rawParams = append(rawParams, a.Id, a.OrgId)
		p := make([]string, len(tagsToDelete))
		for i, t := range tagsToDelete {
			p[i] = "?"
			rawParams = append(rawParams, t)
		}
		rawSql := fmt.Sprintf("DELETE FROM agent_tag WHERE agent_id=? AND org_id=? AND tag IN (%s)", strings.Join(p, ","))
		if _, err := sess.Exec(rawSql, rawParams...); err != nil {
			return err
		}
	}
	if len(tagsToAdd) > 0 {
		newAgentTags := make([]model.AgentTag, len(tagsToAdd))
		for i, tag := range tagsToAdd {
			newAgentTags[i] = model.AgentTag{
				OrgId:   a.OrgId,
				AgentId: a.Id,
				Tag:     tag,
				Created: time.Now(),
			}
		}
		sess.Table("agent_tag")
		if _, err := sess.Insert(&newAgentTags); err != nil {
			return err
		}
	}

	return nil
}

type AgentId struct {
	Id int64
}

func GetAgentsForTask(task *model.TaskDTO) ([]int64, error) {
	sess, err := newSession(false, "agent")
	if err != nil {
		return nil, err
	}

	return getAgentsForTask(sess, task)
}

func getAgentsForTask(sess *session, t *model.TaskDTO) ([]int64, error) {
	agents := make([]*AgentId, 0)
	switch t.Route.Type {
	case model.RouteAny:
		err := sess.Sql("SELECT agent_id as id FROM route_by_any_index where task_id=?", t.Id).Find(&agents)
		if err != nil {
			return nil, err
		}
	case model.RouteByTags:
		//TODO: this list needs to be filtered by agents that support the metrics listed in the task.
		tags := make([]string, len(t.Route.Config["tags"].([]string)))
		for i, tag := range t.Route.Config["tags"].([]string) {
			tags[i] = tag
		}
		sess.Join("LEFT", "agent_tag", "agent.id = agent_tag.agent_id")
		sess.Where("agent_tag.org_id = ?", t.OrgId)
		sess.In("agent_tag.tag", tags)
		sess.Cols("agent.id")
		err := sess.Find(&agents)
		if err != nil {
			return nil, err
		}
	case model.RouteByIds:
		for _, id := range t.Route.Config["ids"].([]int64) {
			agents = append(agents, &AgentId{Id: id})
		}
	default:
		return nil, fmt.Errorf("unknown routeType")
	}
	agentIds := make([]int64, len(agents))
	for i, a := range agents {
		agentIds[i] = a.Id
	}
	return agentIds, nil
}

func DeleteAgent(id int64, orgId int64) error {
	sess, err := newSession(true, "agent")
	if err != nil {
		return err
	}
	defer sess.Cleanup()
	err = deleteAgent(sess, id, orgId)
	if err != nil {
		return err
	}
	sess.Complete()
	return nil
}

func deleteAgent(sess *session, id int64, orgId int64) error {
	existing, err := getAgentById(sess, id, orgId)
	if err != nil {
		return err
	}
	rawSql := "DELETE FROM agent WHERE id=? and org_id=?"
	if _, err := sess.Exec(rawSql, existing.Id, existing.OrgId); err != nil {
		return err
	}
	rawSql = "DELETE FROM agent_tag WHERE agent_id=? and org_id=?"
	if _, err := sess.Exec(rawSql, existing.Id, existing.OrgId); err != nil {
		return err
	}
	rawSql = "DELETE FROM agent_metric WHERE agent_id=?"
	if _, err := sess.Exec(rawSql, existing.Id); err != nil {
		return err
	}
	rawSql = "DELETE FROM route_by_id_index WHERE agent_id=?"
	if _, err := sess.Exec(rawSql, existing.Id); err != nil {
		return err
	}
	return nil
}
<commit_msg>correctly handle adding tags to public agents<commit_after>package sqlstore

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/raintank/raintank-apps/task-server/model"
)

type agentWithTag struct {
	model.Agent    `xorm:"extends"`
	model.AgentTag `xorm:"extends"`
}

type agentWithTags []*agentWithTag

func (agentWithTags) TableName() string {
	return "agent"
}

func (rows agentWithTags) ToAgentDTO() []*model.AgentDTO {
	agentsById := make(map[int64]*model.AgentDTO)
	for _, r := range rows {
		a, ok := agentsById[r.Agent.Id]
		if !ok {
			tags := make([]string, 0)
			if r.AgentTag.Tag != "" {
				tags = append(tags, r.AgentTag.Tag)
			}
			agentsById[r.Agent.Id] = &model.AgentDTO{
				Id:            r.Agent.Id,
				Name:          r.Agent.Name,
				Enabled:       r.Agent.Enabled,
				EnabledChange: r.Agent.EnabledChange,
				OrgId:         r.Agent.OrgId,
				Public:        r.Agent.Public,
				Online:        r.Agent.Online,
				OnlineChange:  r.Agent.OnlineChange,
				Created:       r.Agent.Created,
				Updated:       r.Agent.Updated,
				Tags:          tags,
			}
		} else if r.Tag != "" {
			a.Tags = append(a.Tags, r.Tag)
		}
	}
	agents := make([]*model.AgentDTO, len(agentsById))
	i := 0
	for _, a := range agentsById {
		agents[i] = a
		i++
	}
	return agents
}

func GetAgents(query *model.GetAgentsQuery) ([]*model.AgentDTO, error) {
	sess, err := newSession(false, "agent")
	if err != nil {
		return nil, err
	}
	return getAgents(sess, query)
}

func getAgents(sess *session, query *model.GetAgentsQuery) ([]*model.AgentDTO, error) {
	var a agentWithTags
	var rawSQL bytes.Buffer
	args := make([]interface{}, 0)

	var where bytes.Buffer
	whereArgs := make([]interface{}, 0)
	prefix := "WHERE"

	fmt.Fprint(&rawSQL, "SELECT agent.*, agent_tag.* FROM agent LEFT JOIN agent_tag ON  agent.id = agent_tag.agent_id ")
	if len(query.Tag) > 0 {
		fmt.Fprint(&rawSQL, "INNER JOIN agent_tag as at ON agent.id = at.agent_id ")
		p := make([]string, len(query.Tag))
		for i, tag := range query.Tag {
			p[i] = "?"
			whereArgs = append(whereArgs, tag)
		}
		filter := fmt.Sprintf("at.tag IN (%s)", strings.Join(p, ","))

		fmt.Fprintf(&where, "%s %s ", prefix, filter)
		prefix = "AND"
	}

	if query.Metric != "" {
		sess.Join("INNER", "agent_metric", "agent_metric.agent_id = agent.id").Where("agent_metric.namespace LIKE ?", query.Metric)
		fmt.Fprint(&rawSQL, "INNER JOIN (SELECT DISTINCT(agent_id) FROM agent_metric where namespace like ?) AS am ON am.agent_id = agent.id ")
		ns := strings.Replace(query.Metric, "*", "%", -1)
		args = append(args, ns)
	}

	if query.Name != "" {
		fmt.Fprintf(&where, "%s agent.name=? ", prefix)
		whereArgs = append(whereArgs, query.Name)
		prefix = "AND"
	}
	if query.Enabled != "" {
		enabled, err := strconv.ParseBool(query.Enabled)
		if err != nil {
			return nil, err
		}
		fmt.Fprintf(&where, "%s agent.enabled=? ", prefix)
		whereArgs = append(whereArgs, enabled)
		prefix = "AND"
	}
	if query.Public != "" {
		public, err := strconv.ParseBool(query.Public)
		if err != nil {
			return nil, err
		}
		if public {
			fmt.Fprintf(&where, "%s agent.public=1 ", prefix)
			prefix = "AND"
		} else {
			fmt.Fprintf(&where, "%s (agent.public=0 AND agent.org_id=?) ", prefix)
			whereArgs = append(whereArgs, query.OrgId)
			prefix = "AND"
		}
	} else {
		fmt.Fprintf(&where, "%s (agent.org_id=? OR agent.public=1) ", prefix)
		whereArgs = append(whereArgs, query.OrgId)
		prefix = "AND"
	}

	if query.OrderBy == "" {
		query.OrderBy = "name"
	}
	if query.Limit == 0 {
		query.Limit = 50
	}
	if query.Page == 0 {
		query.Page = 1
	}

	fmt.Fprint(&rawSQL, where.String())
	args = append(args, whereArgs...)
	fmt.Fprintf(&rawSQL, "ORDER BY `%s` ASC LIMIT %d, %d", query.OrderBy, (query.Page-1)*query.Limit, query.Limit)

	err := sess.Sql(rawSQL.String(), args...).Find(&a)

	if err != nil {
		return nil, err
	}
	return a.ToAgentDTO(), nil
}

func onlineAgentsWithNoSession(sess *session) ([]*model.AgentDTO, error) {
	sess.Table("agent")
	sess.Join("LEFT", "agent_tag", "agent.id=agent_tag.agent_id")
	sess.Join("LEFT", "agent_session", "agent.id = agent_session.agent_id")
	sess.Where("agent.online=1").And("agent_session.id is NULL")
	sess.Cols("`agent`.*", "`agent_tag`.*")
	var a agentWithTags
	err := sess.Find(&a)
	if err != nil {
		return nil, err
	}

	return a.ToAgentDTO(), nil
}

func GetAgentById(id int64, orgId int64) (*model.AgentDTO, error) {
	sess, err := newSession(false, "agent")
	if err != nil {
		return nil, err
	}
	return getAgentById(sess, id, orgId)
}

func getAgentById(sess *session, id int64, orgId int64) (*model.AgentDTO, error) {
	var a agentWithTags
	sess.Where("agent.id=?", id)
	if orgId != 0 {
		sess.And("agent.org_id=?", orgId)
	}
	err := sess.Join("LEFT", "agent_tag", "agent.id = agent_tag.agent_id").Find(&a)
	if err != nil {
		return nil, err
	}
	if len(a) == 0 {
		return nil, model.AgentNotFound
	}
	return a.ToAgentDTO()[0], nil
}

func AddAgent(a *model.AgentDTO) error {
	sess, err := newSession(true, "agent")
	if err != nil {
		return err
	}
	defer sess.Cleanup()
	if err = addAgent(sess, a); err != nil {
		return err
	}
	sess.Complete()
	return nil

}

func addAgent(sess *session, a *model.AgentDTO) error {
	agent := &model.Agent{
		Name:          a.Name,
		Enabled:       a.Enabled,
		EnabledChange: time.Now(),
		OrgId:         a.OrgId,
		Public:        a.Public,
		Online:        false,
		OnlineChange:  time.Now(),
		Created:       time.Now(),
		Updated:       time.Now(),
	}

	sess.UseBool("public")
	sess.UseBool("enabled")
	sess.UseBool("online")
	if _, err := sess.Insert(agent); err != nil {
		return err
	}
	a.Id = agent.Id
	a.Created = agent.Created
	a.Updated = agent.Updated

	agentTags := make([]model.AgentTag, 0, len(a.Tags))
	for _, tag := range a.Tags {
		agentTags = append(agentTags, model.AgentTag{
			OrgId:   a.OrgId,
			AgentId: agent.Id,
			Tag:     tag,
			Created: time.Now(),
		})
	}
	if len(agentTags) > 0 {
		sess.Table("agent_tag")
		if _, err := sess.Insert(&agentTags); err != nil {
			return err
		}
	}
	return nil
}

func UpdateAgent(a *model.AgentDTO) error {
	sess, err := newSession(true, "agent")
	if err != nil {
		return err
	}
	defer sess.Cleanup()

	err = updateAgent(sess, a)
	if err != nil {
		return err
	}
	sess.Complete()
	return err
}

func updateAgent(sess *session, a *model.AgentDTO) error {
	existing, err := getAgentById(sess, a.Id, 0)
	if err != nil {
		return err
	}
	if existing == nil {
		return model.AgentNotFound
	}
	// If the OrgId is different, the only changes that can be made is to Tags.
	if a.OrgId == existing.OrgId {
		enabledChange := existing.EnabledChange
		if existing.Enabled != a.Enabled {
			enabledChange = time.Now()
		}
		agent := &model.Agent{
			Id:            a.Id,
			Name:          a.Name,
			Enabled:       a.Enabled,
			EnabledChange: enabledChange,
			OrgId:         a.OrgId,
			Public:        a.Public,
			Created:       a.Created,
			Updated:       time.Now(),
		}
		sess.UseBool("public")
		sess.UseBool("enabled")
		if _, err := sess.Id(agent.Id).Update(agent); err != nil {
			return err
		}
		a.Updated = agent.Updated
	}

	tagMap := make(map[string]bool)
	tagsToDelete := make([]string, 0)
	tagsToAddMap := make(map[string]bool, 0)
	// create map of current tags
	for _, t := range existing.Tags {
		tagMap[t] = false
	}

	// create map of tags to add. We use a map
	// to ensure that we only add each tag once.
	for _, t := range a.Tags {
		if _, ok := tagMap[t]; !ok {
			tagsToAddMap[t] = true
		}
		// mark that this tag has been seen.
		tagMap[t] = true
	}

	//create list of tags to delete
	for t, seen := range tagMap {
		if !seen {
			tagsToDelete = append(tagsToDelete, t)
		}
	}

	// create list of tags to add.
	tagsToAdd := make([]string, len(tagsToAddMap))
	i := 0
	for t := range tagsToAddMap {
		tagsToAdd[i] = t
		i += 1
	}
	if len(tagsToDelete) > 0 {
		rawParams := make([]interface{}, 0)
		rawParams = append(rawParams, a.Id, a.OrgId)
		p := make([]string, len(tagsToDelete))
		for i, t := range tagsToDelete {
			p[i] = "?"
			rawParams = append(rawParams, t)
		}
		rawSql := fmt.Sprintf("DELETE FROM agent_tag WHERE agent_id=? AND org_id=? AND tag IN (%s)", strings.Join(p, ","))
		if _, err := sess.Exec(rawSql, rawParams...); err != nil {
			return err
		}
	}
	if len(tagsToAdd) > 0 {
		newAgentTags := make([]model.AgentTag, len(tagsToAdd))
		for i, tag := range tagsToAdd {
			newAgentTags[i] = model.AgentTag{
				OrgId:   a.OrgId,
				AgentId: a.Id,
				Tag:     tag,
				Created: time.Now(),
			}
		}
		sess.Table("agent_tag")
		if _, err := sess.Insert(&newAgentTags); err != nil {
			return err
		}
	}

	return nil
}

type AgentId struct {
	Id int64
}

func GetAgentsForTask(task *model.TaskDTO) ([]int64, error) {
	sess, err := newSession(false, "agent")
	if err != nil {
		return nil, err
	}

	return getAgentsForTask(sess, task)
}

func getAgentsForTask(sess *session, t *model.TaskDTO) ([]int64, error) {
	agents := make([]*AgentId, 0)
	switch t.Route.Type {
	case model.RouteAny:
		err := sess.Sql("SELECT agent_id as id FROM route_by_any_index where task_id=?", t.Id).Find(&agents)
		if err != nil {
			return nil, err
		}
	case model.RouteByTags:
		//TODO: this list needs to be filtered by agents that support the metrics listed in the task.
		tags := make([]string, len(t.Route.Config["tags"].([]string)))
		for i, tag := range t.Route.Config["tags"].([]string) {
			tags[i] = tag
		}
		sess.Join("LEFT", "agent_tag", "agent.id = agent_tag.agent_id")
		sess.Where("agent_tag.org_id = ?", t.OrgId)
		sess.In("agent_tag.tag", tags)
		sess.Cols("agent.id")
		err := sess.Find(&agents)
		if err != nil {
			return nil, err
		}
	case model.RouteByIds:
		for _, id := range t.Route.Config["ids"].([]int64) {
			agents = append(agents, &AgentId{Id: id})
		}
	default:
		return nil, fmt.Errorf("unknown routeType")
	}
	agentIds := make([]int64, len(agents))
	for i, a := range agents {
		agentIds[i] = a.Id
	}
	return agentIds, nil
}

func DeleteAgent(id int64, orgId int64) error {
	sess, err := newSession(true, "agent")
	if err != nil {
		return err
	}
	defer sess.Cleanup()
	err = deleteAgent(sess, id, orgId)
	if err != nil {
		return err
	}
	sess.Complete()
	return nil
}

func deleteAgent(sess *session, id int64, orgId int64) error {
	existing, err := getAgentById(sess, id, orgId)
	if err != nil {
		return err
	}
	rawSql := "DELETE FROM agent WHERE id=? and org_id=?"
	if _, err := sess.Exec(rawSql, existing.Id, existing.OrgId); err != nil {
		return err
	}
	rawSql = "DELETE FROM agent_tag WHERE agent_id=? and org_id=?"
	if _, err := sess.Exec(rawSql, existing.Id, existing.OrgId); err != nil {
		return err
	}
	rawSql = "DELETE FROM agent_metric WHERE agent_id=?"
	if _, err := sess.Exec(rawSql, existing.Id); err != nil {
		return err
	}
	rawSql = "DELETE FROM route_by_id_index WHERE agent_id=?"
	if _, err := sess.Exec(rawSql, existing.Id); err != nil {
		return err
	}
	return nil
}
<|endoftext|>