<commit_before>// Copyright 2013-present Greg Hurrell. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
//    this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
)

type Options struct {
	Address    string
	Config     string
	Logfile    string
	Executable string
	Flags      string
	Port       int
}

var version = "unknown"

var config Options   // Options read from disk.
var defaults Options // Default options.
var flags Options    // Options set via commandline flags.
var settings Options // Result of merging: flags > config > defaults.
var showHelp bool
var showVersion bool

func printVersion() {
	fmt.Fprintf(os.Stderr, "clipper version: %s (%s)\n", version, runtime.GOOS)
}

func initFlags() {
	const (
		flagsUsage      = "arguments passed to clipboard executable"
		configFileUsage = "path to (JSON) config file"
		executableUsage = "program called to write to clipboard"
		helpUsage       = "show usage information"
		listenAddrUsage = "address to bind to (default loopback interface)"
		listenPortUsage = "port to listen on"
		logFileUsage    = "path to logfile"
		versionUsage    = "show version information"
	)

	flag.BoolVar(&showHelp, "h", false, helpUsage)
	flag.BoolVar(&showHelp, "help", false, helpUsage)
	flag.IntVar(&flags.Port, "p", defaults.Port, listenPortUsage)
	flag.IntVar(&flags.Port, "port", defaults.Port, listenPortUsage)
	flag.StringVar(&flags.Address, "a", defaults.Address, listenAddrUsage)
	flag.StringVar(&flags.Address, "address", defaults.Address, listenAddrUsage)
	flag.StringVar(&flags.Config, "c", defaults.Config, configFileUsage)
	flag.StringVar(&flags.Config, "config", defaults.Config, configFileUsage)
	flag.StringVar(&flags.Executable, "e", defaults.Executable, executableUsage)
	flag.StringVar(&flags.Executable, "executable", defaults.Executable, executableUsage)
	flag.StringVar(&flags.Flags, "f", defaults.Flags, flagsUsage)
	flag.StringVar(&flags.Flags, "flags", defaults.Flags, flagsUsage)
	flag.StringVar(&flags.Logfile, "l", defaults.Logfile, logFileUsage)
	flag.StringVar(&flags.Logfile, "logfile", defaults.Logfile, logFileUsage)
	flag.BoolVar(&showVersion, "v", false, versionUsage)
	flag.BoolVar(&showVersion, "version", false, versionUsage)
}

func setDefaults() {
	defaults.Address = "" // IPv4/IPv6 loopback.
	defaults.Port = 8377

	if runtime.GOOS == "linux" {
		defaults.Config = "~/.config/clipper/clipper.json"
		defaults.Logfile = "~/.config/clipper/logs/clipper.log"
		defaults.Executable = "xclip"
		defaults.Flags = "-selection clipboard"
	} else {
		defaults.Config = "~/.clipper.json"
		defaults.Logfile = "~/Library/Logs/com.wincent.clipper.log"
		defaults.Executable = "pbcopy"
		defaults.Flags = ""
	}
}

func mergeSettings() {
	// Detect which flags were passed in explicitly, and set them immediately.
	// This is used below to determine response to a missing config file.
	visitor := func(f *flag.Flag) {
		if f.Name == "address" || f.Name == "a" {
			settings.Address = flags.Address
		} else if f.Name == "config" || f.Name == "c" {
			settings.Config = flags.Config
		} else if f.Name == "executable" || f.Name == "e" {
			settings.Executable = flags.Executable
		} else if f.Name == "flags" || f.Name == "f" {
			settings.Flags = flags.Flags
		} else if f.Name == "port" || f.Name == "p" {
			settings.Port = flags.Port
		} else if f.Name == "logfile" || f.Name == "l" {
			settings.Logfile = flags.Logfile
		}
	}
	flag.Visit(visitor)

	expandedPath := expandPath(flags.Config)

	if configData, err := ioutil.ReadFile(expandedPath); err != nil {
		if settings.Config != "" {
			// User explicitly asked for a config file and it wasn't there; fail hard.
			log.Fatal(err)
		} else {
			// Default config file missing; just warn.
			log.Print(err)
		}
	} else {
		if err = json.Unmarshal(configData, &config); err != nil {
			log.Fatal(err)
		}
	}

	// Final merge into settings object.
	if settings.Address == "" {
		if config.Address != "" {
			settings.Address = config.Address
		} else {
			settings.Address = defaults.Address
		}
	}
	if settings.Logfile == "" {
		if config.Logfile != "" {
			settings.Logfile = config.Logfile
		} else {
			settings.Logfile = defaults.Logfile
		}
	}
	if settings.Port != 0 || config.Port != 0 {
		if isPath(settings.Address) {
			log.Print("--port option ignored when listening on UNIX domain socket")
		}
	}
	if settings.Port == 0 {
		if config.Port != 0 {
			settings.Port = config.Port
		} else {
			settings.Port = defaults.Port
		}
	}
	if settings.Executable == "" {
		if config.Executable != "" {
			settings.Executable = config.Executable
		} else {
			settings.Executable = defaults.Executable
		}
	}
	if settings.Flags == "" {
		if config.Flags != "" {
			settings.Flags = config.Flags
		} else {
			settings.Flags = defaults.Flags
		}
	}
}

func main() {
	syscall.Umask(0077)
	// Set this up before we even know where our logfile is, in case we have to
	// bail early and print something to stderr.
	log.SetPrefix("clipper: ")
	// Set default values per GOOS.
	setDefaults()
	// Setup flags subsystem.
	initFlags()

	flag.Parse()
	if flag.NArg() != 0 {
		// Additional command-line options not supported.
		flag.Usage()
		os.Exit(1)
	}
	if showHelp {
		printVersion()
		flag.Usage()
		os.Exit(0)
	}
	if showVersion {
		printVersion()
		os.Exit(0)
	}

	// Merge flags -> config -> defaults.
	mergeSettings()

	expandedPath := expandPath(settings.Logfile)
	outfile, err := os.OpenFile(expandedPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		log.Fatal(err)
	}
	defer outfile.Close()
	log.SetOutput(outfile)

	if _, err := exec.LookPath(settings.Executable); err != nil {
		log.Fatal(err)
	}

	var addr string
	var listeners []net.Listener
	if isPath(settings.Address) {
		addr = expandPath(settings.Address)
	} else {
		addr = settings.Address
	}
	if strings.HasPrefix(addr, "/") {
		log.Print("Starting UNIX domain socket server at ", addr)
		listener, err := net.Listen("unix", addr)
		if err != nil {
			log.Fatal(err)
		}
		listeners = append(listeners, listener)
	} else {
		if addr == "" {
			log.Print("Starting TCP server on loopback interface")
			addr = fmt.Sprintf("%s:%d", "127.0.0.1", settings.Port)
			listener, err := net.Listen("tcp4", addr)
			if err != nil {
				log.Fatal(err)
			}
			listeners = append(listeners, listener)
			addr = fmt.Sprintf("%s:%d", "[::1]", settings.Port)
			listener, err = net.Listen("tcp6", addr)
			if err != nil {
				log.Fatal(err)
			}
			listeners = append(listeners, listener)
		} else {
			log.Print("Starting TCP server on ", addr)
			addr = fmt.Sprintf("%s:%d", settings.Address, settings.Port)
			listener, err := net.Listen("tcp", addr)
			if err != nil {
				log.Fatal(err)
			}
			listeners = append(listeners, listener)
		}
	}

	for i := range listeners {
		defer listeners[i].Close()
		go func(listener net.Listener) {
			for {
				conn, err := listener.Accept()
				if err != nil {
					log.Print(err)
					return
				}

				go handleConnection(conn)
			}
		}(listeners[i])
	}

	// Need to catch signals in order for `defer`-ed clean-up items to run.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill, syscall.SIGTERM)
	sig := <-c
	log.Print("Got signal ", sig)
}

// Returns true for things which look like paths (start with "~", "." or "/").
func isPath(path string) bool {
	return strings.HasPrefix(path, "~") ||
		strings.HasPrefix(path, ".") ||
		strings.HasPrefix(path, "/")
}

func expandPath(path string) string {
	expanded := pathByExpandingTildeInPath(path)
	result, err := filepath.Abs(expanded)
	if err != nil {
		log.Fatal(err)
	}
	return result
}

func pathByExpandingTildeInPath(path string) string {
	if strings.HasPrefix(path, "~") {
		user, err := user.Current()
		if err != nil {
			log.Fatal(err)
		}
		path = user.HomeDir + path[1:]
	}
	return path
}

func handleConnection(conn net.Conn) {
	defer log.Print("Connection closed")
	defer conn.Close()

	var args []string
	if settings.Flags != "" {
		whitespace := regexp.MustCompile("\\s+")
		args = whitespace.Split(strings.TrimSpace(settings.Flags), -1)
	}
	cmd := exec.Command(settings.Executable, args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Printf("[ERROR] pipe init: %v\n", err)
		return
	}

	if err = cmd.Start(); err != nil {
		log.Printf("[ERROR] process start: %v\n", err)
		return
	}

	if copied, err := io.Copy(stdin, conn); err != nil {
		log.Printf("[ERROR] pipe copy: %v\n", err)
	} else {
		log.Print("Echoed ", copied, " bytes")
	}
	stdin.Close()

	if err = cmd.Wait(); err != nil {
		log.Printf("[ERROR] wait: %v\n", err)
	}
}
<commit_msg>Reduce duplication of listener set-up<commit_after>// Copyright 2013-present Greg Hurrell. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
//    this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
)

type Options struct {
	Address    string
	Config     string
	Logfile    string
	Executable string
	Flags      string
	Port       int
}

var version = "unknown"

var config Options   // Options read from disk.
var defaults Options // Default options.
var flags Options    // Options set via commandline flags.
var settings Options // Result of merging: flags > config > defaults.
var showHelp bool
var showVersion bool

func printVersion() {
	fmt.Fprintf(os.Stderr, "clipper version: %s (%s)\n", version, runtime.GOOS)
}

func initFlags() {
	const (
		flagsUsage      = "arguments passed to clipboard executable"
		configFileUsage = "path to (JSON) config file"
		executableUsage = "program called to write to clipboard"
		helpUsage       = "show usage information"
		listenAddrUsage = "address to bind to (default loopback interface)"
		listenPortUsage = "port to listen on"
		logFileUsage    = "path to logfile"
		versionUsage    = "show version information"
	)

	flag.BoolVar(&showHelp, "h", false, helpUsage)
	flag.BoolVar(&showHelp, "help", false, helpUsage)
	flag.IntVar(&flags.Port, "p", defaults.Port, listenPortUsage)
	flag.IntVar(&flags.Port, "port", defaults.Port, listenPortUsage)
	flag.StringVar(&flags.Address, "a", defaults.Address, listenAddrUsage)
	flag.StringVar(&flags.Address, "address", defaults.Address, listenAddrUsage)
	flag.StringVar(&flags.Config, "c", defaults.Config, configFileUsage)
	flag.StringVar(&flags.Config, "config", defaults.Config, configFileUsage)
	flag.StringVar(&flags.Executable, "e", defaults.Executable, executableUsage)
	flag.StringVar(&flags.Executable, "executable", defaults.Executable, executableUsage)
	flag.StringVar(&flags.Flags, "f", defaults.Flags, flagsUsage)
	flag.StringVar(&flags.Flags, "flags", defaults.Flags, flagsUsage)
	flag.StringVar(&flags.Logfile, "l", defaults.Logfile, logFileUsage)
	flag.StringVar(&flags.Logfile, "logfile", defaults.Logfile, logFileUsage)
	flag.BoolVar(&showVersion, "v", false, versionUsage)
	flag.BoolVar(&showVersion, "version", false, versionUsage)
}

func setDefaults() {
	defaults.Address = "" // IPv4/IPv6 loopback.
	defaults.Port = 8377

	if runtime.GOOS == "linux" {
		defaults.Config = "~/.config/clipper/clipper.json"
		defaults.Logfile = "~/.config/clipper/logs/clipper.log"
		defaults.Executable = "xclip"
		defaults.Flags = "-selection clipboard"
	} else {
		defaults.Config = "~/.clipper.json"
		defaults.Logfile = "~/Library/Logs/com.wincent.clipper.log"
		defaults.Executable = "pbcopy"
		defaults.Flags = ""
	}
}

func mergeSettings() {
	// Detect which flags were passed in explicitly, and set them immediately.
	// This is used below to determine response to a missing config file.
	visitor := func(f *flag.Flag) {
		if f.Name == "address" || f.Name == "a" {
			settings.Address = flags.Address
		} else if f.Name == "config" || f.Name == "c" {
			settings.Config = flags.Config
		} else if f.Name == "executable" || f.Name == "e" {
			settings.Executable = flags.Executable
		} else if f.Name == "flags" || f.Name == "f" {
			settings.Flags = flags.Flags
		} else if f.Name == "port" || f.Name == "p" {
			settings.Port = flags.Port
		} else if f.Name == "logfile" || f.Name == "l" {
			settings.Logfile = flags.Logfile
		}
	}
	flag.Visit(visitor)

	expandedPath := expandPath(flags.Config)

	if configData, err := ioutil.ReadFile(expandedPath); err != nil {
		if settings.Config != "" {
			// User explicitly asked for a config file and it wasn't there; fail hard.
			log.Fatal(err)
		} else {
			// Default config file missing; just warn.
			log.Print(err)
		}
	} else {
		if err = json.Unmarshal(configData, &config); err != nil {
			log.Fatal(err)
		}
	}

	// Final merge into settings object.
	if settings.Address == "" {
		if config.Address != "" {
			settings.Address = config.Address
		} else {
			settings.Address = defaults.Address
		}
	}
	if settings.Logfile == "" {
		if config.Logfile != "" {
			settings.Logfile = config.Logfile
		} else {
			settings.Logfile = defaults.Logfile
		}
	}
	if settings.Port != 0 || config.Port != 0 {
		if isPath(settings.Address) {
			log.Print("--port option ignored when listening on UNIX domain socket")
		}
	}
	if settings.Port == 0 {
		if config.Port != 0 {
			settings.Port = config.Port
		} else {
			settings.Port = defaults.Port
		}
	}
	if settings.Executable == "" {
		if config.Executable != "" {
			settings.Executable = config.Executable
		} else {
			settings.Executable = defaults.Executable
		}
	}
	if settings.Flags == "" {
		if config.Flags != "" {
			settings.Flags = config.Flags
		} else {
			settings.Flags = defaults.Flags
		}
	}
}

func main() {
	syscall.Umask(0077)
	// Set this up before we even know where our logfile is, in case we have to
	// bail early and print something to stderr.
	log.SetPrefix("clipper: ")
	// Set default values per GOOS.
	setDefaults()
	// Setup flags subsystem.
	initFlags()

	flag.Parse()
	if flag.NArg() != 0 {
		// Additional command-line options not supported.
		flag.Usage()
		os.Exit(1)
	}
	if showHelp {
		printVersion()
		flag.Usage()
		os.Exit(0)
	}
	if showVersion {
		printVersion()
		os.Exit(0)
	}

	// Merge flags -> config -> defaults.
	mergeSettings()

	expandedPath := expandPath(settings.Logfile)
	outfile, err := os.OpenFile(expandedPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		log.Fatal(err)
	}
	defer outfile.Close()
	log.SetOutput(outfile)

	if _, err := exec.LookPath(settings.Executable); err != nil {
		log.Fatal(err)
	}

	var addr string
	var listeners []net.Listener
	if isPath(settings.Address) {
		addr = expandPath(settings.Address)
	} else {
		addr = settings.Address
	}
	if strings.HasPrefix(addr, "/") {
		log.Print("Starting UNIX domain socket server at ", addr)
		listeners = append(listeners, listen("unix", addr, -1))
	} else {
		if addr == "" {
			log.Print("Starting TCP server on loopback interface")
			listeners = append(listeners, listen("tcp4", "127.0.0.1", settings.Port))
			listeners = append(listeners, listen("tcp6", "[::1]", settings.Port))
		} else {
			log.Print("Starting TCP server on ", addr)
			listeners = append(listeners, listen("tcp", settings.Address, settings.Port))
		}
	}

	for i := range listeners {
		defer listeners[i].Close()
		go func(listener net.Listener) {
			for {
				conn, err := listener.Accept()
				if err != nil {
					log.Print(err)
					return
				}

				go handleConnection(conn)
			}
		}(listeners[i])
	}

	// Need to catch signals in order for `defer`-ed clean-up items to run.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill, syscall.SIGTERM)
	sig := <-c
	log.Print("Got signal ", sig)
}

func listen(listenType string, addr string, port int) net.Listener {
	if port >= 0 {
		addr = fmt.Sprintf("%s:%d", addr, port)
	}
	listener, err := net.Listen(listenType, addr)
	if err != nil {
		log.Fatal(err)
	}
	return listener
}

// Returns true for things which look like paths (start with "~", "." or "/").
func isPath(path string) bool {
	return strings.HasPrefix(path, "~") ||
		strings.HasPrefix(path, ".") ||
		strings.HasPrefix(path, "/")
}

func expandPath(path string) string {
	expanded := pathByExpandingTildeInPath(path)
	result, err := filepath.Abs(expanded)
	if err != nil {
		log.Fatal(err)
	}
	return result
}

func pathByExpandingTildeInPath(path string) string {
	if strings.HasPrefix(path, "~") {
		user, err := user.Current()
		if err != nil {
			log.Fatal(err)
		}
		path = user.HomeDir + path[1:]
	}
	return path
}

func handleConnection(conn net.Conn) {
	defer log.Print("Connection closed")
	defer conn.Close()

	var args []string
	if settings.Flags != "" {
		whitespace := regexp.MustCompile("\\s+")
		args = whitespace.Split(strings.TrimSpace(settings.Flags), -1)
	}
	cmd := exec.Command(settings.Executable, args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Printf("[ERROR] pipe init: %v\n", err)
		return
	}

	if err = cmd.Start(); err != nil {
		log.Printf("[ERROR] process start: %v\n", err)
		return
	}

	if copied, err := io.Copy(stdin, conn); err != nil {
		log.Printf("[ERROR] pipe copy: %v\n", err)
	} else {
		log.Print("Echoed ", copied, " bytes")
	}
	stdin.Close()

	if err = cmd.Wait(); err != nil {
		log.Printf("[ERROR] wait: %v\n", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/coreos/etcd/clientv3"
	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/labels"
)

type clusterEventType string

const (
	eventNewCluster    clusterEventType = "Add"
	eventDeleteCluster clusterEventType = "Delete"
	eventReconcile     clusterEventType = "Reconcile"
)

type clusterEvent struct {
	typ          clusterEventType
	size         int
	antiAffinity bool
	// currently running pods in kubernetes
	running MemberSet
}

type Cluster struct {
	kclient *unversioned.Client

	antiAffinity bool

	name string

	idCounter int
	eventCh   chan *clusterEvent
	stopCh    chan struct{}

	// members repsersents the members in the etcd cluster.
	// the name of the member is the the name of the pod the member
	// process runs in.
	members MemberSet

	backupDir string
}

func newCluster(kclient *unversioned.Client, name string) *Cluster {
	c := &Cluster{
		kclient: kclient,
		name:    name,
		eventCh: make(chan *clusterEvent, 100),
		stopCh:  make(chan struct{}),
	}
	go c.run()

	return c
}

func (c *Cluster) init(spec Spec) {
	c.send(&clusterEvent{
		typ:          eventNewCluster,
		size:         spec.Size,
		antiAffinity: spec.AntiAffinity,
	})
}

func (c *Cluster) Delete() {
	c.send(&clusterEvent{typ: eventDeleteCluster})
}

func (c *Cluster) send(ev *clusterEvent) {
	select {
	case c.eventCh <- ev:
	case <-c.stopCh:
	default:
		panic("TODO: too many events queued...")
	}
}

func (c *Cluster) run() {
	go c.monitorMembers()
	for {
		select {
		case event := <-c.eventCh:
			switch event.typ {
			case eventNewCluster:
				c.create(event.size, event.antiAffinity)
			case eventReconcile:
				if err := c.reconcile(event.running); err != nil {
					panic(err)
				}
			case eventDeleteCluster:
				c.delete()
				close(c.stopCh)
				return
			}
		}
	}
}

func (c *Cluster) create(size int, antiAffinity bool) {
	c.antiAffinity = antiAffinity

	members := MemberSet{}
	// we want to make use of member's utility methods.
	for i := 0; i < size; i++ {
		etcdName := fmt.Sprintf("%s-%04d", c.name, i)
		members.Add(&Member{Name: etcdName})
	}

	// TODO: parallelize it
	for i := 0; i < size; i++ {
		etcdName := fmt.Sprintf("%s-%04d", c.name, i)
		if err := c.createPodAndService(members, members[etcdName], "new"); err != nil {
			panic(fmt.Sprintf("(TODO: we need to clean up already created ones.)\nError: %v", err))
		}
		c.idCounter++
	}

	fmt.Println("created cluster:", members)
}

func (c *Cluster) updateMembers(etcdcli *clientv3.Client) {
	resp, err := etcdcli.MemberList(context.TODO())
	if err != nil {
		panic(err)
	}
	c.members = MemberSet{}
	for _, m := range resp.Members {
		id := findID(m.Name)
		if id+1 > c.idCounter {
			c.idCounter = id + 1
		}

		c.members[m.Name] = &Member{
			Name: m.Name,
			ID:   m.ID,
		}
	}
}
func findID(name string) int {
	var id int
	fmt.Sscanf(name, "etcd-cluster-%d", &id)
	return id
}
func (c *Cluster) delete() {
	option := api.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			"etcd_cluster": c.name,
		}),
	}

	pods, err := c.kclient.Pods("default").List(option)
	if err != nil {
		panic(err)
	}
	for i := range pods.Items {
		c.removePodAndService(pods.Items[i].Name)
	}
}

// todo: use a struct to replace the huge arg list.
func (c *Cluster) createPodAndService(members MemberSet, m *Member, state string) error {
	if err := createEtcdService(c.kclient, m.Name, c.name); err != nil {
		return err
	}
	return createEtcdPod(c.kclient, members.PeerURLPairs(), m, c.name, state, c.antiAffinity)
}

func (c *Cluster) removePodAndService(name string) error {
	err := c.kclient.Pods("default").Delete(name, nil)
	if err != nil {
		return err
	}
	err = c.kclient.Services("default").Delete(name)
	if err != nil {
		log.Printf("failed removing service (%s): %v", name, err)
	}
	return nil
}

func (c *Cluster) backup() error {
	clientAddr := "todo"
	nextSnapshotName := "todo"

	cfg := clientv3.Config{
		Endpoints: []string{clientAddr},
	}
	etcdcli, err := clientv3.New(cfg)
	if err != nil {
		panic(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)

	log.Println("saving snapshot from cluster", c.name)

	rc, err := etcdcli.Maintenance.Snapshot(ctx)
	cancel()
	if err != nil {
		return err
	}

	tmpfile, err := ioutil.TempFile(c.backupDir, "snapshot")
	n, err := io.Copy(tmpfile, rc)
	if err != nil {
		tmpfile.Close()
		os.Remove(tmpfile.Name())
		log.Printf("saving snapshot from cluster %s error: %v\n", c.name, err)
		return err
	}

	err = os.Rename(tmpfile.Name(), nextSnapshotName)
	if err != nil {
		os.Remove(tmpfile.Name())
		log.Printf("renaming snapshot from cluster %s error: %v\n", c.name, err)
		return err
	}

	log.Printf("saved snapshot %v (size: %d) from cluster %s", n, nextSnapshotName, c.name)

	return nil
}

func (c *Cluster) monitorMembers() {
	opts := api.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			"etcd_cluster": c.name,
		}),
	}
	// TODO: Select "etcd_node" to remove left service.
	for {
		select {
		case <-c.stopCh:
			return
		case <-time.After(5 * time.Second):
		}

		podList, err := c.kclient.Pods("default").List(opts)
		if err != nil {
			panic(err)
		}
		running := MemberSet{}
		for i := range podList.Items {
			running.Add(&Member{Name: podList.Items[i].Name})
		}

		c.send(&clusterEvent{
			typ:     eventReconcile,
			running: running,
		})
	}
}
<commit_msg>rename monitorMembers -> monitorPods<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/coreos/etcd/clientv3"
	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/labels"
)

type clusterEventType string

const (
	eventNewCluster    clusterEventType = "Add"
	eventDeleteCluster clusterEventType = "Delete"
	eventReconcile     clusterEventType = "Reconcile"
)

type clusterEvent struct {
	typ          clusterEventType
	size         int
	antiAffinity bool
	// currently running pods in kubernetes
	running MemberSet
}

type Cluster struct {
	kclient *unversioned.Client

	antiAffinity bool

	name string

	idCounter int
	eventCh   chan *clusterEvent
	stopCh    chan struct{}

	// members repsersents the members in the etcd cluster.
	// the name of the member is the the name of the pod the member
	// process runs in.
	members MemberSet

	backupDir string
}

func newCluster(kclient *unversioned.Client, name string) *Cluster {
	c := &Cluster{
		kclient: kclient,
		name:    name,
		eventCh: make(chan *clusterEvent, 100),
		stopCh:  make(chan struct{}),
	}
	go c.run()

	return c
}

func (c *Cluster) init(spec Spec) {
	c.send(&clusterEvent{
		typ:          eventNewCluster,
		size:         spec.Size,
		antiAffinity: spec.AntiAffinity,
	})
}

func (c *Cluster) Delete() {
	c.send(&clusterEvent{typ: eventDeleteCluster})
}

func (c *Cluster) send(ev *clusterEvent) {
	select {
	case c.eventCh <- ev:
	case <-c.stopCh:
	default:
		panic("TODO: too many events queued...")
	}
}

func (c *Cluster) run() {
	go c.monitorPods()
	for {
		select {
		case event := <-c.eventCh:
			switch event.typ {
			case eventNewCluster:
				c.create(event.size, event.antiAffinity)
			case eventReconcile:
				if err := c.reconcile(event.running); err != nil {
					panic(err)
				}
			case eventDeleteCluster:
				c.delete()
				close(c.stopCh)
				return
			}
		}
	}
}

func (c *Cluster) create(size int, antiAffinity bool) {
	c.antiAffinity = antiAffinity

	members := MemberSet{}
	// we want to make use of member's utility methods.
	for i := 0; i < size; i++ {
		etcdName := fmt.Sprintf("%s-%04d", c.name, i)
		members.Add(&Member{Name: etcdName})
	}

	// TODO: parallelize it
	for i := 0; i < size; i++ {
		etcdName := fmt.Sprintf("%s-%04d", c.name, i)
		if err := c.createPodAndService(members, members[etcdName], "new"); err != nil {
			panic(fmt.Sprintf("(TODO: we need to clean up already created ones.)\nError: %v", err))
		}
		c.idCounter++
	}

	fmt.Println("created cluster:", members)
}

func (c *Cluster) updateMembers(etcdcli *clientv3.Client) {
	resp, err := etcdcli.MemberList(context.TODO())
	if err != nil {
		panic(err)
	}
	c.members = MemberSet{}
	for _, m := range resp.Members {
		id := findID(m.Name)
		if id+1 > c.idCounter {
			c.idCounter = id + 1
		}

		c.members[m.Name] = &Member{
			Name: m.Name,
			ID:   m.ID,
		}
	}
}
func findID(name string) int {
	var id int
	fmt.Sscanf(name, "etcd-cluster-%d", &id)
	return id
}
func (c *Cluster) delete() {
	option := api.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			"etcd_cluster": c.name,
		}),
	}

	pods, err := c.kclient.Pods("default").List(option)
	if err != nil {
		panic(err)
	}
	for i := range pods.Items {
		c.removePodAndService(pods.Items[i].Name)
	}
}

// todo: use a struct to replace the huge arg list.
func (c *Cluster) createPodAndService(members MemberSet, m *Member, state string) error {
	if err := createEtcdService(c.kclient, m.Name, c.name); err != nil {
		return err
	}
	return createEtcdPod(c.kclient, members.PeerURLPairs(), m, c.name, state, c.antiAffinity)
}

func (c *Cluster) removePodAndService(name string) error {
	err := c.kclient.Pods("default").Delete(name, nil)
	if err != nil {
		return err
	}
	err = c.kclient.Services("default").Delete(name)
	if err != nil {
		log.Printf("failed removing service (%s): %v", name, err)
	}
	return nil
}

func (c *Cluster) backup() error {
	clientAddr := "todo"
	nextSnapshotName := "todo"

	cfg := clientv3.Config{
		Endpoints: []string{clientAddr},
	}
	etcdcli, err := clientv3.New(cfg)
	if err != nil {
		panic(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)

	log.Println("saving snapshot from cluster", c.name)

	rc, err := etcdcli.Maintenance.Snapshot(ctx)
	cancel()
	if err != nil {
		return err
	}

	tmpfile, err := ioutil.TempFile(c.backupDir, "snapshot")
	n, err := io.Copy(tmpfile, rc)
	if err != nil {
		tmpfile.Close()
		os.Remove(tmpfile.Name())
		log.Printf("saving snapshot from cluster %s error: %v\n", c.name, err)
		return err
	}

	err = os.Rename(tmpfile.Name(), nextSnapshotName)
	if err != nil {
		os.Remove(tmpfile.Name())
		log.Printf("renaming snapshot from cluster %s error: %v\n", c.name, err)
		return err
	}

	log.Printf("saved snapshot %v (size: %d) from cluster %s", n, nextSnapshotName, c.name)

	return nil
}

func (c *Cluster) monitorPods() {
	opts := api.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			"etcd_cluster": c.name,
		}),
	}
	// TODO: Select "etcd_node" to remove left service.
	for {
		select {
		case <-c.stopCh:
			return
		case <-time.After(5 * time.Second):
		}

		podList, err := c.kclient.Pods("default").List(opts)
		if err != nil {
			panic(err)
		}
		running := MemberSet{}
		for i := range podList.Items {
			running.Add(&Member{Name: podList.Items[i].Name})
		}

		c.send(&clusterEvent{
			typ:     eventReconcile,
			running: running,
		})
	}
}
<|endoftext|>


<commit_before>package ec2cluster

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

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// DiscoverClusterMembersByTag returns a list of the private IP addresses of each
// EC2 instance having a tag that matches the tag on the currently running
// instance. The tag is specified by tagName.
//
// For example, if your instances are tagged with a tag named `app` whose
// value is unique per cluster, then you can invoke:
//
//     knownClusterMembers, err := DiscoverClusterMembersByTag(nil, "app")
//
func DiscoverClusterMembersByTag(config *aws.Config, tagName string) ([]string, error) {
	instanceID, err := discoverInstanceID()
	if err != nil {
		return nil, err
	}

	EC2 := ec2.New(config)
	resp, err := EC2.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(instanceID)},
		MaxResults:  aws.Int64(1),
	})
	if err != nil {
		return nil, err
	}
	if len(resp.Reservations) != 1 || len(resp.Reservations[0].Instances) != 1 {
		return nil, fmt.Errorf("Cannot find instance %s", instanceID)
	}

	tagValue := ""
	for _, tag := range resp.Reservations[0].Instances[0].Tags {
		if *tag.Key == tagName {
			tagValue = *tag.Value
		}
	}
	if tagValue == "" {
		return nil, fmt.Errorf("Current instance (%s) does not have a tag %s.",
			instanceID, tagName)
	}

	// List all the instances having that tag
	rv := []string{}

	err = EC2.DescribeInstancesPages(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name:   aws.String(fmt.Sprintf("tag:%s", tagName)),
				Values: []*string{aws.String(tagValue)},
			},
		},
	}, func(resp *ec2.DescribeInstancesOutput, lastPage bool) (shouldContinue bool) {
		for _, reservation := range resp.Reservations {
			for _, instance := range reservation.Instances {
				rv = append(rv, *instance.PrivateIpAddress)
			}
		}
		return true
	})
	if err != nil {
		return nil, err
	}

	return rv, nil
}

// DiscoverAdvertiseAddress returns the address that should be advertised for
// the current node based on the current EC2 instance's private IP address
func DiscoverAdvertiseAddress() (string, error) {
	return readMetadata("local-ipv4")
}

func discoverInstanceID() (string, error) {
	return readMetadata("instance-id")
}

func readMetadata(suffix string) (string, error) {
	// a nice short timeout so we don't hang too much on non-AWS boxes
	client := http.DefaultClient
	client.Timeout = 200 * time.Millisecond

	resp, err := client.Get("http://169.254.169.254/latest/meta-data/" + suffix)
	if err != nil {
		return "", err
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("fetching metadata: %s", resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(body), nil
}
<commit_msg>bugfix: don't call EC2 DescribeInstances with both InstanceId and MaxResults set<commit_after>package ec2cluster

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

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// DiscoverClusterMembersByTag returns a list of the private IP addresses of each
// EC2 instance having a tag that matches the tag on the currently running
// instance. The tag is specified by tagName.
//
// For example, if your instances are tagged with a tag named `app` whose
// value is unique per cluster, then you can invoke:
//
//     knownClusterMembers, err := DiscoverClusterMembersByTag(nil, "app")
//
func DiscoverClusterMembersByTag(config *aws.Config, tagName string) ([]string, error) {
	instanceID, err := discoverInstanceID()
	if err != nil {
		return nil, err
	}

	EC2 := ec2.New(config)
	resp, err := EC2.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(instanceID)},
	})
	if err != nil {
		return nil, err
	}
	if len(resp.Reservations) != 1 || len(resp.Reservations[0].Instances) != 1 {
		return nil, fmt.Errorf("Cannot find instance %s", instanceID)
	}

	tagValue := ""
	for _, tag := range resp.Reservations[0].Instances[0].Tags {
		if *tag.Key == tagName {
			tagValue = *tag.Value
		}
	}
	if tagValue == "" {
		return nil, fmt.Errorf("Current instance (%s) does not have a tag %s.",
			instanceID, tagName)
	}

	// List all the instances having that tag
	rv := []string{}

	err = EC2.DescribeInstancesPages(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name:   aws.String(fmt.Sprintf("tag:%s", tagName)),
				Values: []*string{aws.String(tagValue)},
			},
		},
	}, func(resp *ec2.DescribeInstancesOutput, lastPage bool) (shouldContinue bool) {
		for _, reservation := range resp.Reservations {
			for _, instance := range reservation.Instances {
				rv = append(rv, *instance.PrivateIpAddress)
			}
		}
		return true
	})
	if err != nil {
		return nil, err
	}

	return rv, nil
}

// DiscoverAdvertiseAddress returns the address that should be advertised for
// the current node based on the current EC2 instance's private IP address
func DiscoverAdvertiseAddress() (string, error) {
	return readMetadata("local-ipv4")
}

func discoverInstanceID() (string, error) {
	return readMetadata("instance-id")
}

func readMetadata(suffix string) (string, error) {
	// a nice short timeout so we don't hang too much on non-AWS boxes
	client := http.DefaultClient
	client.Timeout = 200 * time.Millisecond

	resp, err := client.Get("http://169.254.169.254/latest/meta-data/" + suffix)
	if err != nil {
		return "", err
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("fetching metadata: %s", resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(body), nil
}
<|endoftext|>


<commit_before>// Package cmd is an interface for parsing the command line
package cmd

import (
	"fmt"
	"io"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/server"

	// brokers
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/broker/http"

	// registries
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/registry/consul"
	"github.com/micro/go-micro/registry/mdns"

	// selectors
	"github.com/micro/go-micro/selector"
	"github.com/micro/go-micro/selector/cache"

	// transports
	"github.com/micro/go-micro/transport"
	thttp "github.com/micro/go-micro/transport/http"
)

type Cmd interface {
	// The cli app within this cmd
	App() *cli.App
	// Adds options, parses flags and initialise
	// exits on error
	Init(opts ...Option) error
	// Options set within this command
	Options() Options
}

type cmd struct {
	opts Options
	app  *cli.App
}

type Option func(o *Options)

var (
	DefaultCmd = newCmd()

	DefaultFlags = []cli.Flag{
		cli.StringFlag{
			Name:   "client",
			EnvVar: "MICRO_CLIENT",
			Usage:  "Client for go-micro; rpc",
		},
		cli.StringFlag{
			Name:   "client_request_timeout",
			EnvVar: "MICRO_CLIENT_REQUEST_TIMEOUT",
			Usage:  "Sets the client request timeout. e.g 500ms, 5s, 1m. Default: 5s",
		},
		cli.IntFlag{
			Name:   "client_retries",
			EnvVar: "MICRO_CLIENT_RETRIES",
			Usage:  "Sets the client retries. Default: 1",
		},
		cli.IntFlag{
			Name:   "client_pool_size",
			EnvVar: "MICRO_CLIENT_POOL_SIZE",
			Usage:  "Sets the client connection pool size. Default: 0",
		},
		cli.StringFlag{
			Name:   "client_pool_ttl",
			EnvVar: "MICRO_CLIENT_POOL_TTL",
			Usage:  "Sets the client connection pool ttl. e.g 500ms, 5s, 1m. Default: 1m",
		},
		cli.StringFlag{
			Name:   "server_name",
			EnvVar: "MICRO_SERVER_NAME",
			Usage:  "Name of the server. go.micro.srv.example",
		},
		cli.StringFlag{
			Name:   "server_version",
			EnvVar: "MICRO_SERVER_VERSION",
			Usage:  "Version of the server. 1.1.0",
		},
		cli.StringFlag{
			Name:   "server_id",
			EnvVar: "MICRO_SERVER_ID",
			Usage:  "Id of the server. Auto-generated if not specified",
		},
		cli.StringFlag{
			Name:   "server_address",
			EnvVar: "MICRO_SERVER_ADDRESS",
			Usage:  "Bind address for the server. 127.0.0.1:8080",
		},
		cli.StringFlag{
			Name:   "server_advertise",
			EnvVar: "MICRO_SERVER_ADVERTISE",
			Usage:  "Used instead of the server_address when registering with discovery. 127.0.0.1:8080",
		},
		cli.StringSliceFlag{
			Name:   "server_metadata",
			EnvVar: "MICRO_SERVER_METADATA",
			Value:  &cli.StringSlice{},
			Usage:  "A list of key-value pairs defining metadata. version=1.0.0",
		},
		cli.StringFlag{
			Name:   "broker",
			EnvVar: "MICRO_BROKER",
			Usage:  "Broker for pub/sub. http, nats, rabbitmq",
		},
		cli.StringFlag{
			Name:   "broker_address",
			EnvVar: "MICRO_BROKER_ADDRESS",
			Usage:  "Comma-separated list of broker addresses",
		},
		cli.StringFlag{
			Name:   "registry",
			EnvVar: "MICRO_REGISTRY",
			Usage:  "Registry for discovery. consul, mdns",
		},
		cli.StringFlag{
			Name:   "registry_address",
			EnvVar: "MICRO_REGISTRY_ADDRESS",
			Usage:  "Comma-separated list of registry addresses",
		},
		cli.StringFlag{
			Name:   "selector",
			EnvVar: "MICRO_SELECTOR",
			Usage:  "Selector used to pick nodes for querying",
		},
		cli.StringFlag{
			Name:   "server",
			EnvVar: "MICRO_SERVER",
			Usage:  "Server for go-micro; rpc",
		},
		cli.StringFlag{
			Name:   "transport",
			EnvVar: "MICRO_TRANSPORT",
			Usage:  "Transport mechanism used; http",
		},
		cli.StringFlag{
			Name:   "transport_address",
			EnvVar: "MICRO_TRANSPORT_ADDRESS",
			Usage:  "Comma-separated list of transport addresses",
		},
	}

	DefaultBrokers = map[string]func(...broker.Option) broker.Broker{
		"http": http.NewBroker,
	}

	DefaultClients = map[string]func(...client.Option) client.Client{
		"rpc": client.NewClient,
	}

	DefaultRegistries = map[string]func(...registry.Option) registry.Registry{
		"consul": consul.NewRegistry,
		"mdns":   mdns.NewRegistry,
	}

	DefaultSelectors = map[string]func(...selector.Option) selector.Selector{
		"default": selector.NewSelector,
		"cache":   cache.NewSelector,
	}

	DefaultServers = map[string]func(...server.Option) server.Server{
		"rpc": server.NewServer,
	}

	DefaultTransports = map[string]func(...transport.Option) transport.Transport{
		"http": thttp.NewTransport,
	}

	// used for default selection as the fall back
	defaultClient    = "rpc"
	defaultServer    = "rpc"
	defaultBroker    = "http"
	defaultRegistry  = "consul"
	defaultSelector  = "default"
	defaultTransport = "http"
)

func init() {
	rand.Seed(time.Now().Unix())
	help := cli.HelpPrinter
	cli.HelpPrinter = func(writer io.Writer, templ string, data interface{}) {
		help(writer, templ, data)
		os.Exit(0)
	}
}

func newCmd(opts ...Option) Cmd {
	options := Options{
		Broker:    &broker.DefaultBroker,
		Client:    &client.DefaultClient,
		Registry:  &registry.DefaultRegistry,
		Server:    &server.DefaultServer,
		Selector:  &selector.DefaultSelector,
		Transport: &transport.DefaultTransport,

		Brokers:    DefaultBrokers,
		Clients:    DefaultClients,
		Registries: DefaultRegistries,
		Selectors:  DefaultSelectors,
		Servers:    DefaultServers,
		Transports: DefaultTransports,
	}

	for _, o := range opts {
		o(&options)
	}

	if len(options.Description) == 0 {
		options.Description = "a go-micro service"
	}

	cmd := new(cmd)
	cmd.opts = options
	cmd.app = cli.NewApp()
	cmd.app.Name = cmd.opts.Name
	cmd.app.Version = cmd.opts.Version
	cmd.app.Usage = cmd.opts.Description
	cmd.app.Before = cmd.Before
	cmd.app.Flags = DefaultFlags
	cmd.app.Action = func(c *cli.Context) {}

	if len(options.Version) == 0 {
		cmd.app.HideVersion = true
	}

	return cmd
}

func (c *cmd) App() *cli.App {
	return c.app
}

func (c *cmd) Options() Options {
	return c.opts
}

func (c *cmd) Before(ctx *cli.Context) error {
	// If flags are set then use them otherwise do nothing
	var serverOpts []server.Option
	var clientOpts []client.Option

	// Set the client
	if name := ctx.String("client"); len(name) > 0 {
		if cl, ok := c.opts.Clients[name]; ok {
			*c.opts.Client = cl()
		}
	}

	// Set the server
	if name := ctx.String("server"); len(name) > 0 {
		if s, ok := c.opts.Servers[name]; ok {
			*c.opts.Server = s()
		}
	}

	// Set the broker
	if name := ctx.String("broker"); len(name) > 0 || len(ctx.String("broker_address")) > 0 {
		if len(name) == 0 {
			name = defaultBroker
		}

		if b, ok := c.opts.Brokers[name]; ok {
			n := b(broker.Addrs(strings.Split(ctx.String("broker_address"), ",")...))
			*c.opts.Broker = n
		} else {
			return fmt.Errorf("Broker %s not found", name)
		}

		serverOpts = append(serverOpts, server.Broker(*c.opts.Broker))
		clientOpts = append(clientOpts, client.Broker(*c.opts.Broker))
	}

	// Set the registry
	if name := ctx.String("registry"); len(name) > 0 || len(ctx.String("registry_address")) > 0 {
		if len(name) == 0 {
			name = defaultRegistry
		}

		if r, ok := c.opts.Registries[name]; ok {
			n := r(registry.Addrs(strings.Split(ctx.String("registry_address"), ",")...))
			*c.opts.Registry = n
		} else {
			return fmt.Errorf("Registry %s not found", name)
		}

		serverOpts = append(serverOpts, server.Registry(*c.opts.Registry))
		clientOpts = append(clientOpts, client.Registry(*c.opts.Registry))

		(*c.opts.Selector).Init(selector.Registry(*c.opts.Registry))
		clientOpts = append(clientOpts, client.Selector(*c.opts.Selector))

		(*c.opts.Broker).Init(broker.Registry(*c.opts.Registry))
	}

	// Set the selector
	if name := ctx.String("selector"); len(name) > 0 {
		if s, ok := c.opts.Selectors[name]; ok {
			n := s(selector.Registry(*c.opts.Registry))
			*c.opts.Selector = n
		} else {
			return fmt.Errorf("Selector %s not found", name)
		}

		// No server option here. Should there be?
		clientOpts = append(clientOpts, client.Selector(*c.opts.Selector))
	}

	// Set the transport
	if name := ctx.String("transport"); len(name) > 0 || len(ctx.String("transport_address")) > 0 {
		if len(name) == 0 {
			name = defaultTransport
		}

		if t, ok := c.opts.Transports[name]; ok {
			n := t(transport.Addrs(strings.Split(ctx.String("transport_address"), ",")...))
			*c.opts.Transport = n
		} else {
			return fmt.Errorf("Transport %s not found", name)
		}

		serverOpts = append(serverOpts, server.Transport(*c.opts.Transport))
		clientOpts = append(clientOpts, client.Transport(*c.opts.Transport))
	}

	// Parse the server options
	metadata := make(map[string]string)
	for _, d := range ctx.StringSlice("server_metadata") {
		var key, val string
		parts := strings.Split(d, "=")
		key = parts[0]
		if len(parts) > 1 {
			val = strings.Join(parts[1:], "=")
		}
		metadata[key] = val
	}

	if len(metadata) > 0 {
		serverOpts = append(serverOpts, server.Metadata(metadata))
	}

	if len(ctx.String("server_name")) > 0 {
		serverOpts = append(serverOpts, server.Name(ctx.String("server_name")))
	}

	if len(ctx.String("server_version")) > 0 {
		serverOpts = append(serverOpts, server.Version(ctx.String("server_version")))
	}

	if len(ctx.String("server_id")) > 0 {
		serverOpts = append(serverOpts, server.Id(ctx.String("server_id")))
	}

	if len(ctx.String("server_address")) > 0 {
		serverOpts = append(serverOpts, server.Address(ctx.String("server_address")))
	}

	if len(ctx.String("server_advertise")) > 0 {
		serverOpts = append(serverOpts, server.Advertise(ctx.String("server_advertise")))
	}

	// client opts
	if r := ctx.Int("client_retries"); r > 0 {
		clientOpts = append(clientOpts, client.Retries(r))
	}

	if t := ctx.String("client_request_timeout"); len(t) > 0 {
		d, err := time.ParseDuration(t)
		if err != nil {
			return fmt.Errorf("failed to parse client_request_timeout: %v", t)
		}
		clientOpts = append(clientOpts, client.RequestTimeout(d))
	}

	if r := ctx.Int("client_pool_size"); r > 0 {
		clientOpts = append(clientOpts, client.PoolSize(r))
	}

	if t := ctx.String("client_pool_ttl"); len(t) > 0 {
		d, err := time.ParseDuration(t)
		if err != nil {
			return fmt.Errorf("failed to parse client_pool_ttl: %v", t)
		}
		clientOpts = append(clientOpts, client.PoolTTL(d))
	}

	// We have some command line opts for the server.
	// Lets set it up
	if len(serverOpts) > 0 {
		(*c.opts.Server).Init(serverOpts...)
	}

	// Use an init option?
	if len(clientOpts) > 0 {
		(*c.opts.Client).Init(clientOpts...)
	}

	return nil
}

func (c *cmd) Init(opts ...Option) error {
	for _, o := range opts {
		o(&c.opts)
	}
	c.app.Name = c.opts.Name
	c.app.Version = c.opts.Version
	c.app.Usage = c.opts.Description
	c.app.RunAndExitOnError()
	return nil
}

func DefaultOptions() Options {
	return DefaultCmd.Options()
}

func App() *cli.App {
	return DefaultCmd.App()
}

func Init(opts ...Option) error {
	return DefaultCmd.Init(opts...)
}

func NewCmd(opts ...Option) Cmd {
	return newCmd(opts...)
}
<commit_msg>fix version option in cmd<commit_after>// Package cmd is an interface for parsing the command line
package cmd

import (
	"fmt"
	"io"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/server"

	// brokers
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/broker/http"

	// registries
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/registry/consul"
	"github.com/micro/go-micro/registry/mdns"

	// selectors
	"github.com/micro/go-micro/selector"
	"github.com/micro/go-micro/selector/cache"

	// transports
	"github.com/micro/go-micro/transport"
	thttp "github.com/micro/go-micro/transport/http"
)

type Cmd interface {
	// The cli app within this cmd
	App() *cli.App
	// Adds options, parses flags and initialise
	// exits on error
	Init(opts ...Option) error
	// Options set within this command
	Options() Options
}

type cmd struct {
	opts Options
	app  *cli.App
}

type Option func(o *Options)

var (
	DefaultCmd = newCmd()

	DefaultFlags = []cli.Flag{
		cli.StringFlag{
			Name:   "client",
			EnvVar: "MICRO_CLIENT",
			Usage:  "Client for go-micro; rpc",
		},
		cli.StringFlag{
			Name:   "client_request_timeout",
			EnvVar: "MICRO_CLIENT_REQUEST_TIMEOUT",
			Usage:  "Sets the client request timeout. e.g 500ms, 5s, 1m. Default: 5s",
		},
		cli.IntFlag{
			Name:   "client_retries",
			EnvVar: "MICRO_CLIENT_RETRIES",
			Usage:  "Sets the client retries. Default: 1",
		},
		cli.IntFlag{
			Name:   "client_pool_size",
			EnvVar: "MICRO_CLIENT_POOL_SIZE",
			Usage:  "Sets the client connection pool size. Default: 0",
		},
		cli.StringFlag{
			Name:   "client_pool_ttl",
			EnvVar: "MICRO_CLIENT_POOL_TTL",
			Usage:  "Sets the client connection pool ttl. e.g 500ms, 5s, 1m. Default: 1m",
		},
		cli.StringFlag{
			Name:   "server_name",
			EnvVar: "MICRO_SERVER_NAME",
			Usage:  "Name of the server. go.micro.srv.example",
		},
		cli.StringFlag{
			Name:   "server_version",
			EnvVar: "MICRO_SERVER_VERSION",
			Usage:  "Version of the server. 1.1.0",
		},
		cli.StringFlag{
			Name:   "server_id",
			EnvVar: "MICRO_SERVER_ID",
			Usage:  "Id of the server. Auto-generated if not specified",
		},
		cli.StringFlag{
			Name:   "server_address",
			EnvVar: "MICRO_SERVER_ADDRESS",
			Usage:  "Bind address for the server. 127.0.0.1:8080",
		},
		cli.StringFlag{
			Name:   "server_advertise",
			EnvVar: "MICRO_SERVER_ADVERTISE",
			Usage:  "Used instead of the server_address when registering with discovery. 127.0.0.1:8080",
		},
		cli.StringSliceFlag{
			Name:   "server_metadata",
			EnvVar: "MICRO_SERVER_METADATA",
			Value:  &cli.StringSlice{},
			Usage:  "A list of key-value pairs defining metadata. version=1.0.0",
		},
		cli.StringFlag{
			Name:   "broker",
			EnvVar: "MICRO_BROKER",
			Usage:  "Broker for pub/sub. http, nats, rabbitmq",
		},
		cli.StringFlag{
			Name:   "broker_address",
			EnvVar: "MICRO_BROKER_ADDRESS",
			Usage:  "Comma-separated list of broker addresses",
		},
		cli.StringFlag{
			Name:   "registry",
			EnvVar: "MICRO_REGISTRY",
			Usage:  "Registry for discovery. consul, mdns",
		},
		cli.StringFlag{
			Name:   "registry_address",
			EnvVar: "MICRO_REGISTRY_ADDRESS",
			Usage:  "Comma-separated list of registry addresses",
		},
		cli.StringFlag{
			Name:   "selector",
			EnvVar: "MICRO_SELECTOR",
			Usage:  "Selector used to pick nodes for querying",
		},
		cli.StringFlag{
			Name:   "server",
			EnvVar: "MICRO_SERVER",
			Usage:  "Server for go-micro; rpc",
		},
		cli.StringFlag{
			Name:   "transport",
			EnvVar: "MICRO_TRANSPORT",
			Usage:  "Transport mechanism used; http",
		},
		cli.StringFlag{
			Name:   "transport_address",
			EnvVar: "MICRO_TRANSPORT_ADDRESS",
			Usage:  "Comma-separated list of transport addresses",
		},
	}

	DefaultBrokers = map[string]func(...broker.Option) broker.Broker{
		"http": http.NewBroker,
	}

	DefaultClients = map[string]func(...client.Option) client.Client{
		"rpc": client.NewClient,
	}

	DefaultRegistries = map[string]func(...registry.Option) registry.Registry{
		"consul": consul.NewRegistry,
		"mdns":   mdns.NewRegistry,
	}

	DefaultSelectors = map[string]func(...selector.Option) selector.Selector{
		"default": selector.NewSelector,
		"cache":   cache.NewSelector,
	}

	DefaultServers = map[string]func(...server.Option) server.Server{
		"rpc": server.NewServer,
	}

	DefaultTransports = map[string]func(...transport.Option) transport.Transport{
		"http": thttp.NewTransport,
	}

	// used for default selection as the fall back
	defaultClient    = "rpc"
	defaultServer    = "rpc"
	defaultBroker    = "http"
	defaultRegistry  = "consul"
	defaultSelector  = "default"
	defaultTransport = "http"
)

func init() {
	rand.Seed(time.Now().Unix())
	help := cli.HelpPrinter
	cli.HelpPrinter = func(writer io.Writer, templ string, data interface{}) {
		help(writer, templ, data)
		os.Exit(0)
	}
}

func newCmd(opts ...Option) Cmd {
	options := Options{
		Broker:    &broker.DefaultBroker,
		Client:    &client.DefaultClient,
		Registry:  &registry.DefaultRegistry,
		Server:    &server.DefaultServer,
		Selector:  &selector.DefaultSelector,
		Transport: &transport.DefaultTransport,

		Brokers:    DefaultBrokers,
		Clients:    DefaultClients,
		Registries: DefaultRegistries,
		Selectors:  DefaultSelectors,
		Servers:    DefaultServers,
		Transports: DefaultTransports,
	}

	for _, o := range opts {
		o(&options)
	}

	if len(options.Description) == 0 {
		options.Description = "a go-micro service"
	}

	cmd := new(cmd)
	cmd.opts = options
	cmd.app = cli.NewApp()
	cmd.app.Name = cmd.opts.Name
	cmd.app.Version = cmd.opts.Version
	cmd.app.Usage = cmd.opts.Description
	cmd.app.Before = cmd.Before
	cmd.app.Flags = DefaultFlags
	cmd.app.Action = func(c *cli.Context) {}

	if len(options.Version) == 0 {
		cmd.app.HideVersion = true
	}

	return cmd
}

func (c *cmd) App() *cli.App {
	return c.app
}

func (c *cmd) Options() Options {
	return c.opts
}

func (c *cmd) Before(ctx *cli.Context) error {
	// If flags are set then use them otherwise do nothing
	var serverOpts []server.Option
	var clientOpts []client.Option

	// Set the client
	if name := ctx.String("client"); len(name) > 0 {
		if cl, ok := c.opts.Clients[name]; ok {
			*c.opts.Client = cl()
		}
	}

	// Set the server
	if name := ctx.String("server"); len(name) > 0 {
		if s, ok := c.opts.Servers[name]; ok {
			*c.opts.Server = s()
		}
	}

	// Set the broker
	if name := ctx.String("broker"); len(name) > 0 || len(ctx.String("broker_address")) > 0 {
		if len(name) == 0 {
			name = defaultBroker
		}

		if b, ok := c.opts.Brokers[name]; ok {
			n := b(broker.Addrs(strings.Split(ctx.String("broker_address"), ",")...))
			*c.opts.Broker = n
		} else {
			return fmt.Errorf("Broker %s not found", name)
		}

		serverOpts = append(serverOpts, server.Broker(*c.opts.Broker))
		clientOpts = append(clientOpts, client.Broker(*c.opts.Broker))
	}

	// Set the registry
	if name := ctx.String("registry"); len(name) > 0 || len(ctx.String("registry_address")) > 0 {
		if len(name) == 0 {
			name = defaultRegistry
		}

		if r, ok := c.opts.Registries[name]; ok {
			n := r(registry.Addrs(strings.Split(ctx.String("registry_address"), ",")...))
			*c.opts.Registry = n
		} else {
			return fmt.Errorf("Registry %s not found", name)
		}

		serverOpts = append(serverOpts, server.Registry(*c.opts.Registry))
		clientOpts = append(clientOpts, client.Registry(*c.opts.Registry))

		(*c.opts.Selector).Init(selector.Registry(*c.opts.Registry))
		clientOpts = append(clientOpts, client.Selector(*c.opts.Selector))

		(*c.opts.Broker).Init(broker.Registry(*c.opts.Registry))
	}

	// Set the selector
	if name := ctx.String("selector"); len(name) > 0 {
		if s, ok := c.opts.Selectors[name]; ok {
			n := s(selector.Registry(*c.opts.Registry))
			*c.opts.Selector = n
		} else {
			return fmt.Errorf("Selector %s not found", name)
		}

		// No server option here. Should there be?
		clientOpts = append(clientOpts, client.Selector(*c.opts.Selector))
	}

	// Set the transport
	if name := ctx.String("transport"); len(name) > 0 || len(ctx.String("transport_address")) > 0 {
		if len(name) == 0 {
			name = defaultTransport
		}

		if t, ok := c.opts.Transports[name]; ok {
			n := t(transport.Addrs(strings.Split(ctx.String("transport_address"), ",")...))
			*c.opts.Transport = n
		} else {
			return fmt.Errorf("Transport %s not found", name)
		}

		serverOpts = append(serverOpts, server.Transport(*c.opts.Transport))
		clientOpts = append(clientOpts, client.Transport(*c.opts.Transport))
	}

	// Parse the server options
	metadata := make(map[string]string)
	for _, d := range ctx.StringSlice("server_metadata") {
		var key, val string
		parts := strings.Split(d, "=")
		key = parts[0]
		if len(parts) > 1 {
			val = strings.Join(parts[1:], "=")
		}
		metadata[key] = val
	}

	if len(metadata) > 0 {
		serverOpts = append(serverOpts, server.Metadata(metadata))
	}

	if len(ctx.String("server_name")) > 0 {
		serverOpts = append(serverOpts, server.Name(ctx.String("server_name")))
	}

	if len(ctx.String("server_version")) > 0 {
		serverOpts = append(serverOpts, server.Version(ctx.String("server_version")))
	}

	if len(ctx.String("server_id")) > 0 {
		serverOpts = append(serverOpts, server.Id(ctx.String("server_id")))
	}

	if len(ctx.String("server_address")) > 0 {
		serverOpts = append(serverOpts, server.Address(ctx.String("server_address")))
	}

	if len(ctx.String("server_advertise")) > 0 {
		serverOpts = append(serverOpts, server.Advertise(ctx.String("server_advertise")))
	}

	// client opts
	if r := ctx.Int("client_retries"); r > 0 {
		clientOpts = append(clientOpts, client.Retries(r))
	}

	if t := ctx.String("client_request_timeout"); len(t) > 0 {
		d, err := time.ParseDuration(t)
		if err != nil {
			return fmt.Errorf("failed to parse client_request_timeout: %v", t)
		}
		clientOpts = append(clientOpts, client.RequestTimeout(d))
	}

	if r := ctx.Int("client_pool_size"); r > 0 {
		clientOpts = append(clientOpts, client.PoolSize(r))
	}

	if t := ctx.String("client_pool_ttl"); len(t) > 0 {
		d, err := time.ParseDuration(t)
		if err != nil {
			return fmt.Errorf("failed to parse client_pool_ttl: %v", t)
		}
		clientOpts = append(clientOpts, client.PoolTTL(d))
	}

	// We have some command line opts for the server.
	// Lets set it up
	if len(serverOpts) > 0 {
		(*c.opts.Server).Init(serverOpts...)
	}

	// Use an init option?
	if len(clientOpts) > 0 {
		(*c.opts.Client).Init(clientOpts...)
	}

	return nil
}

func (c *cmd) Init(opts ...Option) error {
	for _, o := range opts {
		o(&c.opts)
	}
	c.app.Name = c.opts.Name
	c.app.Version = c.opts.Version
	c.app.HideVersion = len(c.opts.Version) == 0
	c.app.Usage = c.opts.Description
	c.app.RunAndExitOnError()
	return nil
}

func DefaultOptions() Options {
	return DefaultCmd.Options()
}

func App() *cli.App {
	return DefaultCmd.App()
}

func Init(opts ...Option) error {
	return DefaultCmd.Init(opts...)
}

func NewCmd(opts ...Option) Cmd {
	return newCmd(opts...)
}
<|endoftext|>


<commit_before>package cmd

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/abiosoft/ishell"
	"github.com/satori/go.uuid"
	"github.com/tywkeene/autobd/api"
	"github.com/tywkeene/autobd/client"
	"github.com/tywkeene/autobd/index"
	"github.com/tywkeene/autobd/options"
	"github.com/tywkeene/autobd/utils"
	"github.com/tywkeene/autobd/version"
	"os"
	"strings"
	"time"
)

type CliConfig struct {
	HeartbeatInterval string `toml:"heartbeat_interval"` // How often to send a heartbeat to the servers
	MaxMissedBeats    int    `toml:"max_missed_beats"`   //How many heartbeats a server is allowed to miss
	OutputJSON        bool   `toml:"output_json"`        //Output json instead of pretty printing
	SyncDirPath       string `toml:"sync_dir"`           //Where to save synced files/directories
	UUIDPath          string `toml:"uuid_path"`          //Where to read/write cli UUID
}

type Cli struct {
	Servers       []*client.Client //Clients for servers
	CurrentServer *client.Client   //Current working server
	Config        CliConfig        //Cli config
	UUID          string           //Cli UUID
}

func cliUserAgent() string { return "Autobd-cli/" + version.GetCliVersion() }

func newCli(configFile string) *Cli {
	var config CliConfig
	if configFile != "" {
		if _, err := toml.DecodeFile(configFile, &config); err != nil {
			fmt.Printf("Error reading config (%s): %s\n", configFile, err.Error())
			return nil
		}
	}

	cli := &Cli{make([]*client.Client, 0), nil, config, ""}
	if _, err := os.Stat(cli.Config.UUIDPath); os.IsNotExist(err) {
		cli.UUID = uuid.NewV4().String()
		if err := utils.WriteJson(cli.Config.UUIDPath, cli.UUID); err != nil {
			fmt.Println(err)
			return cli
		}
		fmt.Printf("Generated and wrote cli-UUID (%s) to (%s)\n", cli.UUID, cli.Config.UUIDPath)
	} else {
		if err := utils.ReadJson(cli.Config.UUIDPath, &cli.UUID); err != nil {
			fmt.Println(err)
			return cli
		}
		fmt.Printf("Read cli-UUID (%s) from (%s)\n", cli.UUID, cli.Config.UUIDPath)
	}
	return cli
}

func (c *Cli) appendServer(server *client.Client) {
	if c.Servers == nil {
		c.Servers = make([]*client.Client, 0)
	}
	c.Servers = append(c.Servers, server)
}

func (c *Cli) addServer(address string) error {
	server := client.NewClient(address)
	_, err := server.IdentifyWithServer(c.UUID, cliUserAgent())
	if err != nil {
		return err
	}
	c.appendServer(server)
	c.CurrentServer = server
	return nil
}

func (c *Cli) findServer(address string) *client.Client {
	for _, server := range c.Servers {
		if server.Address == address {
			return server
		}
	}
	return nil
}

func (c *Cli) StartHeartbeat() {
	go func(c *Cli) {
		interval, _ := time.ParseDuration(c.Config.HeartbeatInterval)
		for {
			time.Sleep(interval)

			for _, server := range c.Servers {
				_, err := server.SendHeartbeat(c.UUID, false, cliUserAgent())
				if err != nil {
					server.MissedBeats++
					if server.MissedBeats == c.Config.MaxMissedBeats {
						server.Online = false
					}
				}
			}
		}
	}(c)
}

func truncateString(str string, length int) string {
	sub := str[:length]
	return sub
}

func (c *Cli) prettyPrintDir(dir *index.Index) {
	fmt.Printf("Name\tModify time\tMode\n")
	fmt.Println(dir.Name, dir.ModTime, dir.Mode)
}

func (c *Cli) prettyPrintFile(file *index.Index) {
	fmt.Printf("Name: %s\tChecksum: %s\tSize: %d\tModify time: %v\tMode: %v\n",
		file.Name, truncateString(file.Checksum, 8), file.Size, file.ModTime, file.Mode)
}

func (c *Cli) prettyPrintIndex(index map[string]*index.Index) {
	for _, object := range index {
		if object.IsDir == true {
			c.prettyPrintIndex(object.Files)
			c.prettyPrintDir(object)
			continue
		} else {
			c.prettyPrintFile(object)
		}
	}
}

func (c *Cli) prettyPrintNodes(nodes map[string]*api.Node) {
	for uuid, node := range nodes {
		fmt.Println("-----------------------------")
		if uuid == c.UUID {
			fmt.Print("* ")
		}
		fmt.Printf("UUID: %s\nAddress: %s\nVersion: %s\nLast online: %s\nCurrently online: %v\nSynced: %v\n",
			uuid, node.Address, node.Version, node.LastOnline, node.IsOnline, node.Synced)
	}
	fmt.Println("-----------------------------")
}

func (c *Cli) printConfig() {
	fmt.Printf("heartbeat_interval = '%s'\nmax_missed_beats = %d\noutput_json = %v\nsync_dir = '%s'\nuuid_path = '%s'\n",
		c.Config.HeartbeatInterval, c.Config.MaxMissedBeats, c.Config.OutputJSON, c.Config.SyncDirPath, c.Config.UUIDPath)
}

func Start() {
	shell := ishell.New()

	c := newCli(options.Config.CliConfigPath)
	c.printConfig()

	shell.Register("print-config", func(args ...string) (string, error) {
		c.printConfig()
		return "OK", nil
	})

	shell.Register("server", func(args ...string) (string, error) {
		if len(args) == 0 {
			return "", errors.New("Server address required")
		}
		address := args[0]
		if strings.Contains(address, "http://") == false &&
			strings.Contains(address, "https://") == false {
			return "", errors.New("Address must be URL")
		}
		if server := c.findServer(address); server != nil {
			c.CurrentServer = server
			shell.Println("Changed current server to: ", server.Address)
			return "OK", nil
		}
		if err := c.addServer(address); err != nil {
			return "", err
		}
		c.StartHeartbeat()
		return "OK", nil
	})

	shell.Register("list-servers", func(args ...string) (string, error) {
		if len(c.Servers) == 0 {
			return "", errors.New("No servers")
		}
		for _, server := range c.Servers {
			if c.CurrentServer == server {
				fmt.Print("* ")
			} else {
				fmt.Print("| ")
			}
			fmt.Printf("Address: %s\t Missed Heartbeats: %d Online: %v\n",
				server.Address, server.MissedBeats, server.Online)
		}
		return "OK", nil
	})

	shell.Register("get-nodes", func(args ...string) (string, error) {
		if len(c.Servers) == 0 {
			return "", errors.New("No servers")
		}
		serial, err := c.CurrentServer.GetNodes(c.UUID, cliUserAgent())
		if err != nil {
			return "", err
		}
		if c.Config.OutputJSON == true {
			fmt.Println(string(serial))
			return "OK", nil
		}
		var nodes map[string]*api.Node
		if err := json.Unmarshal(serial, &nodes); err != nil {
			return "", err
		}
		c.prettyPrintNodes(nodes)
		return "OK", nil
	})
	shell.Register("sync-file", func(args ...string) (string, error) {
		if len(args) == 0 {
			return "", errors.New("Must specify file to sync")
		}
		fmt.Println("(Not yet implemented)")
		return "OK", nil
	})

	shell.Register("sync-dir", func(args ...string) (string, error) {
		if len(args) == 0 {
			return "", errors.New("Must specify dir to sync")
		}
		fmt.Println("(Not yet implemented)")
		return "OK", nil
	})

	shell.Register("get-index", func(args ...string) (string, error) {
		if len(args) == 0 {
			return "", errors.New("Must specify directory to index")
		}
		if len(c.Servers) == 0 {
			return "", errors.New("No servers")
		}
		dir := args[0]
		fmt.Println(dir)
		if dir == "" {
			return "", errors.New("Must specify directory")
		}
		serial, err := c.CurrentServer.RequestIndex(dir, c.UUID, cliUserAgent())
		if err != nil {
			return "", err
		}

		if c.Config.OutputJSON == true {
			fmt.Println(string(serial))
		} else {
			var remoteIndex map[string]*index.Index
			if err := json.Unmarshal(serial, &remoteIndex); err != nil {
				c.prettyPrintIndex(remoteIndex)
			}
		}
		return "", nil
	})
	shell.Start()
}
<commit_msg>Remove sync-file and sync-dir commands from autobd-cli<commit_after>package cmd

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/abiosoft/ishell"
	"github.com/satori/go.uuid"
	"github.com/tywkeene/autobd/api"
	"github.com/tywkeene/autobd/client"
	"github.com/tywkeene/autobd/index"
	"github.com/tywkeene/autobd/options"
	"github.com/tywkeene/autobd/utils"
	"github.com/tywkeene/autobd/version"
	"os"
	"strings"
	"time"
)

type CliConfig struct {
	HeartbeatInterval string `toml:"heartbeat_interval"` // How often to send a heartbeat to the servers
	MaxMissedBeats    int    `toml:"max_missed_beats"`   //How many heartbeats a server is allowed to miss
	OutputJSON        bool   `toml:"output_json"`        //Output json instead of pretty printing
	SyncDirPath       string `toml:"sync_dir"`           //Where to save synced files/directories
	UUIDPath          string `toml:"uuid_path"`          //Where to read/write cli UUID
}

type Cli struct {
	Servers       []*client.Client //Clients for servers
	CurrentServer *client.Client   //Current working server
	Config        CliConfig        //Cli config
	UUID          string           //Cli UUID
}

func cliUserAgent() string { return "Autobd-cli/" + version.GetCliVersion() }

func newCli(configFile string) *Cli {
	var config CliConfig
	if configFile != "" {
		if _, err := toml.DecodeFile(configFile, &config); err != nil {
			fmt.Printf("Error reading config (%s): %s\n", configFile, err.Error())
			return nil
		}
	}

	cli := &Cli{make([]*client.Client, 0), nil, config, ""}
	if _, err := os.Stat(cli.Config.UUIDPath); os.IsNotExist(err) {
		cli.UUID = uuid.NewV4().String()
		if err := utils.WriteJson(cli.Config.UUIDPath, cli.UUID); err != nil {
			fmt.Println(err)
			return cli
		}
		fmt.Printf("Generated and wrote cli-UUID (%s) to (%s)\n", cli.UUID, cli.Config.UUIDPath)
	} else {
		if err := utils.ReadJson(cli.Config.UUIDPath, &cli.UUID); err != nil {
			fmt.Println(err)
			return cli
		}
		fmt.Printf("Read cli-UUID (%s) from (%s)\n", cli.UUID, cli.Config.UUIDPath)
	}
	return cli
}

func (c *Cli) appendServer(server *client.Client) {
	if c.Servers == nil {
		c.Servers = make([]*client.Client, 0)
	}
	c.Servers = append(c.Servers, server)
}

func (c *Cli) addServer(address string) error {
	server := client.NewClient(address)
	_, err := server.IdentifyWithServer(c.UUID, cliUserAgent())
	if err != nil {
		return err
	}
	c.appendServer(server)
	c.CurrentServer = server
	return nil
}

func (c *Cli) findServer(address string) *client.Client {
	for _, server := range c.Servers {
		if server.Address == address {
			return server
		}
	}
	return nil
}

func (c *Cli) StartHeartbeat() {
	go func(c *Cli) {
		interval, _ := time.ParseDuration(c.Config.HeartbeatInterval)
		for {
			time.Sleep(interval)

			for _, server := range c.Servers {
				_, err := server.SendHeartbeat(c.UUID, false, cliUserAgent())
				if err != nil {
					server.MissedBeats++
					if server.MissedBeats == c.Config.MaxMissedBeats {
						server.Online = false
					}
				}
			}
		}
	}(c)
}

func truncateString(str string, length int) string {
	sub := str[:length]
	return sub
}

func (c *Cli) prettyPrintDir(dir *index.Index) {
	fmt.Printf("Name\tModify time\tMode\n")
	fmt.Println(dir.Name, dir.ModTime, dir.Mode)
}

func (c *Cli) prettyPrintFile(file *index.Index) {
	fmt.Printf("Name: %s\tChecksum: %s\tSize: %d\tModify time: %v\tMode: %v\n",
		file.Name, truncateString(file.Checksum, 8), file.Size, file.ModTime, file.Mode)
}

func (c *Cli) prettyPrintIndex(index map[string]*index.Index) {
	for _, object := range index {
		if object.IsDir == true {
			c.prettyPrintIndex(object.Files)
			c.prettyPrintDir(object)
			continue
		} else {
			c.prettyPrintFile(object)
		}
	}
}

func (c *Cli) prettyPrintNodes(nodes map[string]*api.Node) {
	for uuid, node := range nodes {
		fmt.Println("-----------------------------")
		if uuid == c.UUID {
			fmt.Print("* ")
		}
		fmt.Printf("UUID: %s\nAddress: %s\nVersion: %s\nLast online: %s\nCurrently online: %v\nSynced: %v\n",
			uuid, node.Address, node.Version, node.LastOnline, node.IsOnline, node.Synced)
	}
	fmt.Println("-----------------------------")
}

func (c *Cli) printConfig() {
	fmt.Printf("heartbeat_interval = '%s'\nmax_missed_beats = %d\noutput_json = %v\nsync_dir = '%s'\nuuid_path = '%s'\n",
		c.Config.HeartbeatInterval, c.Config.MaxMissedBeats, c.Config.OutputJSON, c.Config.SyncDirPath, c.Config.UUIDPath)
}

func Start() {
	shell := ishell.New()

	c := newCli(options.Config.CliConfigPath)
	c.printConfig()

	shell.Register("print-config", func(args ...string) (string, error) {
		c.printConfig()
		return "OK", nil
	})

	shell.Register("server", func(args ...string) (string, error) {
		if len(args) == 0 {
			return "", errors.New("Server address required")
		}
		address := args[0]
		if strings.Contains(address, "http://") == false &&
			strings.Contains(address, "https://") == false {
			return "", errors.New("Address must be URL")
		}
		if server := c.findServer(address); server != nil {
			c.CurrentServer = server
			shell.Println("Changed current server to: ", server.Address)
			return "OK", nil
		}
		if err := c.addServer(address); err != nil {
			return "", err
		}
		c.StartHeartbeat()
		return "OK", nil
	})

	shell.Register("list-servers", func(args ...string) (string, error) {
		if len(c.Servers) == 0 {
			return "", errors.New("No servers")
		}
		for _, server := range c.Servers {
			if c.CurrentServer == server {
				fmt.Print("* ")
			} else {
				fmt.Print("| ")
			}
			fmt.Printf("Address: %s\t Missed Heartbeats: %d Online: %v\n",
				server.Address, server.MissedBeats, server.Online)
		}
		return "OK", nil
	})

	shell.Register("get-nodes", func(args ...string) (string, error) {
		if len(c.Servers) == 0 {
			return "", errors.New("No servers")
		}
		serial, err := c.CurrentServer.GetNodes(c.UUID, cliUserAgent())
		if err != nil {
			return "", err
		}
		if c.Config.OutputJSON == true {
			fmt.Println(string(serial))
			return "OK", nil
		}
		var nodes map[string]*api.Node
		if err := json.Unmarshal(serial, &nodes); err != nil {
			return "", err
		}
		c.prettyPrintNodes(nodes)
		return "OK", nil
	})
	shell.Register("get-index", func(args ...string) (string, error) {
		if len(args) == 0 {
			return "", errors.New("Must specify directory to index")
		}
		if len(c.Servers) == 0 {
			return "", errors.New("No servers")
		}
		dir := args[0]
		fmt.Println(dir)
		if dir == "" {
			return "", errors.New("Must specify directory")
		}
		serial, err := c.CurrentServer.RequestIndex(dir, c.UUID, cliUserAgent())
		if err != nil {
			return "", err
		}

		if c.Config.OutputJSON == true {
			fmt.Println(string(serial))
		} else {
			var remoteIndex map[string]*index.Index
			if err := json.Unmarshal(serial, &remoteIndex); err != nil {
				c.prettyPrintIndex(remoteIndex)
			}
		}
		return "", nil
	})
	shell.Start()
}
<|endoftext|>


<commit_before>package cmd

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/minamijoyo/myaws/myaws"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

func init() {
	RootCmd.AddCommand(newECSCmd())
}

func newECSCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "ecs",
		Short: "Manage ECS resources",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	cmd.AddCommand(
		newECSNodeCmd(),
	)

	return cmd
}

func newECSNodeCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "node",
		Short: "Manage ECS node resources",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	cmd.AddCommand(
		newECSNodeLsCmd(),
		newECSNodeUpdateCmd(),
	)

	return cmd
}

func newECSNodeLsCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "ls CLUSTER",
		Short: "List ECS nodes",
		RunE:  runECSNodeLsCmd,
	}

	return cmd
}

func runECSNodeLsCmd(cmd *cobra.Command, args []string) error {
	client, err := newClient()
	if err != nil {
		return errors.Wrap(err, "newClient failed:")
	}

	if len(args) == 0 {
		return errors.New("CLUSTER is required")
	}

	options := myaws.ECSNodeLsOptions{
		Cluster: args[0],
	}
	return client.ECSNodeLs(options)
}

func newECSNodeUpdateCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "update CLUSTER",
		Short: "Update ECS nodes",
		RunE:  runECSNodeUpdateCmd,
	}

	flags := cmd.Flags()
	flags.StringP("container-instances", "i", "", "A list of container instance IDs or full ARN entries separated by space")
	flags.StringP("status", "s", "", "container instance state (ACTIVE | DRAINING)")

	viper.BindPFlag("ecs.node.update.container-instances", flags.Lookup("container-instances"))
	viper.BindPFlag("ecs.node.update.status", flags.Lookup("status"))

	return cmd
}

func runECSNodeUpdateCmd(cmd *cobra.Command, args []string) error {
	client, err := newClient()
	if err != nil {
		return errors.Wrap(err, "newClient failed:")
	}

	if len(args) == 0 {
		return errors.New("CLUSTER is required")
	}

	containerInstances := aws.StringSlice(viper.GetStringSlice("ecs.node.update.container-instances"))
	if len(containerInstances) == 0 {
		return errors.New("container-instances is required")
	}

	status := viper.GetString("ecs.node.update.status")
	if len(status) == 0 {
		return errors.New("status is required")
	}

	options := myaws.ECSNodeUpdateOptions{
		Cluster:            args[0],
		ContainerInstances: containerInstances,
		Status:             status,
	}

	return client.ECSNodeUpdate(options)
}
<commit_msg>ecs/node: Clarify the ecs nodes means container instances in help<commit_after>package cmd

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/minamijoyo/myaws/myaws"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

func init() {
	RootCmd.AddCommand(newECSCmd())
}

func newECSCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "ecs",
		Short: "Manage ECS resources",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	cmd.AddCommand(
		newECSNodeCmd(),
	)

	return cmd
}

func newECSNodeCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "node",
		Short: "Manage ECS node resources (container instances)",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	cmd.AddCommand(
		newECSNodeLsCmd(),
		newECSNodeUpdateCmd(),
	)

	return cmd
}

func newECSNodeLsCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "ls CLUSTER",
		Short: "List ECS nodes (container instances)",
		RunE:  runECSNodeLsCmd,
	}

	return cmd
}

func runECSNodeLsCmd(cmd *cobra.Command, args []string) error {
	client, err := newClient()
	if err != nil {
		return errors.Wrap(err, "newClient failed:")
	}

	if len(args) == 0 {
		return errors.New("CLUSTER is required")
	}

	options := myaws.ECSNodeLsOptions{
		Cluster: args[0],
	}
	return client.ECSNodeLs(options)
}

func newECSNodeUpdateCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "update CLUSTER",
		Short: "Update ECS nodes (container instances)",
		RunE:  runECSNodeUpdateCmd,
	}

	flags := cmd.Flags()
	flags.StringP("container-instances", "i", "", "A list of container instance IDs or full ARN entries separated by space")
	flags.StringP("status", "s", "", "container instance state (ACTIVE | DRAINING)")

	viper.BindPFlag("ecs.node.update.container-instances", flags.Lookup("container-instances"))
	viper.BindPFlag("ecs.node.update.status", flags.Lookup("status"))

	return cmd
}

func runECSNodeUpdateCmd(cmd *cobra.Command, args []string) error {
	client, err := newClient()
	if err != nil {
		return errors.Wrap(err, "newClient failed:")
	}

	if len(args) == 0 {
		return errors.New("CLUSTER is required")
	}

	containerInstances := aws.StringSlice(viper.GetStringSlice("ecs.node.update.container-instances"))
	if len(containerInstances) == 0 {
		return errors.New("container-instances is required")
	}

	status := viper.GetString("ecs.node.update.status")
	if len(status) == 0 {
		return errors.New("status is required")
	}

	options := myaws.ECSNodeUpdateOptions{
		Cluster:            args[0],
		ContainerInstances: containerInstances,
		Status:             status,
	}

	return client.ECSNodeUpdate(options)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"github.com/spf13/cobra"
	"log"
	"os"
	"os/exec"
)

var (
	sshKey  string
	sshPort int
	sshUser string
)

func init() {
	sshCmd.Flags().StringVarP(&sshUser, "user", "u", "root", "--user flag is the ssh user.")
	sshCmd.Flags().IntVarP(&sshPort, "port", "p", 22, "--port to utilize")
	sshCmd.Flags().StringVarP(&sshKey, "key", "i", "", "--key is the ssh key.")

	RootCmd.AddCommand(sshCmd)
}

var sshCmd = &cobra.Command{
	Use:   "ssh",
	Short: "ssh [host]",
	Long:  "ssh [host] will log into specified host using ssh. If the command matches a recipe name first, the recipe will be used in place of the command.",
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) != 1 {
			log.Fatal("A single ssh [host] must be specified")
		}

		sshArgs := []string{
			"-p", fmt.Sprintf("%d", sshPort),
			"-o", "UserKnownHostsFile=/dev/null",
			"-o", "StrictHostKeyChecking=no",
			"-o", "LogLevel=quiet",
			fmt.Sprintf("%s@%s", sshUser, args[0]),
		}

		if sshKey != "" {
			sshArgs = append(sshArgs, "-i")
			sshArgs = append(sshArgs, sshKey)
		}

		currentCmd := exec.Command("ssh", sshArgs...)

		currentCmd.Stdin = os.Stdin
		currentCmd.Stderr = os.Stderr
		currentCmd.Stdout = os.Stdout

		if err := currentCmd.Run(); err != nil {
			log.Fatal(err)
		}
	},
}
<commit_msg>fix usage doc<commit_after>package cmd

import (
	"fmt"
	"github.com/spf13/cobra"
	"log"
	"os"
	"os/exec"
)

var (
	sshKey  string
	sshPort int
	sshUser string
)

func init() {
	sshCmd.Flags().StringVarP(&sshUser, "user", "u", "root", "--user flag is the ssh user.")
	sshCmd.Flags().IntVarP(&sshPort, "port", "p", 22, "--port to utilize")
	sshCmd.Flags().StringVarP(&sshKey, "key", "i", "", "--key is the ssh key.")

	RootCmd.AddCommand(sshCmd)
}

var sshCmd = &cobra.Command{
	Use:   "ssh [host]",
	Short: "ssh [host]",
	Long:  "ssh [host] will log into specified host using ssh. If the command matches a recipe name first, the recipe will be used in place of the command.",
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) != 1 {
			log.Fatal("A single ssh [host] must be specified")
		}

		sshArgs := []string{
			"-p", fmt.Sprintf("%d", sshPort),
			"-o", "UserKnownHostsFile=/dev/null",
			"-o", "StrictHostKeyChecking=no",
			"-o", "LogLevel=quiet",
			fmt.Sprintf("%s@%s", sshUser, args[0]),
		}

		if sshKey != "" {
			sshArgs = append(sshArgs, "-i")
			sshArgs = append(sshArgs, sshKey)
		}

		currentCmd := exec.Command("ssh", sshArgs...)

		currentCmd.Stdin = os.Stdin
		currentCmd.Stderr = os.Stderr
		currentCmd.Stdout = os.Stdout

		if err := currentCmd.Run(); err != nil {
			log.Fatal(err)
		}
	},
}
<|endoftext|>


<commit_before>package kdtree

import (
	"errors"
	"github.com/gonum/matrix/mat64"
	"github.com/sjwhitworth/golearn/metrics/pairwise"
	"sort"
)

type node struct {
	feature  int
	value    []float64
	srcRowNo int
	left     *node
	right    *node
}

// Tree is a kdtree.
type Tree struct {
	firstDiv *node
	data     [][]float64
}

type SortData struct {
	RowData [][]float64
	Data    []int
	Feature int
}

func (d SortData) Len() int { return len(d.Data) }
func (d SortData) Less(i, j int) bool {
	return d.RowData[d.Data[i]][d.Feature] < d.RowData[d.Data[j]][d.Feature]
}
func (d SortData) Swap(i, j int) { d.Data[i], d.Data[j] = d.Data[j], d.Data[i] }

// New return a Tree pointer.
func New() *Tree {
	return &Tree{}
}

// Build builds the kdtree with specific data.
func (t *Tree) Build(data [][]float64) error {
	if len(data) == 0 {
		return errors.New("no input data")
	}
	size := len(data[0])
	for _, v := range data {
		if len(v) != size {
			return errors.New("amounts of features are not the same")
		}
	}

	t.data = data

	newData := make([]int, len(data))
	for k, _ := range newData {
		newData[k] = k
	}

	if len(data) == 1 {
		t.firstDiv = &node{feature: -1, srcRowNo: 0}
		t.firstDiv.value = make([]float64, len(data[0]))
		copy(t.firstDiv.value, data[0])
	} else {
		t.firstDiv = t.buildHandle(newData, 0)
	}

	return nil
}

// buildHandle builds the kdtree recursively.
func (t *Tree) buildHandle(data []int, featureIndex int) *node {
	n := &node{feature: featureIndex}

	tmp := SortData{RowData: t.data, Data: data, Feature: featureIndex}
	sort.Sort(tmp)
	middle := len(data) / 2

	n.srcRowNo = data[middle]
	n.value = make([]float64, len(t.data[data[middle]]))
	copy(n.value, t.data[data[middle]])

	divPoint := middle
	for i := middle + 1; i < len(data); i++ {
		if t.data[data[i]][featureIndex] == t.data[data[middle]][featureIndex] {
			divPoint = i
		} else {
			break
		}
	}

	if divPoint == 1 {
		n.left = &node{feature: -1}
		n.left.value = make([]float64, len(t.data[data[0]]))
		copy(n.left.value, t.data[data[0]])
		n.left.srcRowNo = data[0]
	} else {
		n.left = t.buildHandle(data[:divPoint], (featureIndex+1)%len(t.data[data[0]]))
	}

	if divPoint == (len(data) - 2) {
		n.right = &node{feature: -1}
		n.right.value = make([]float64, len(t.data[data[divPoint+1]]))
		copy(n.right.value, t.data[data[divPoint+1]])
		n.left.srcRowNo = data[divPoint+1]
	} else if divPoint != (len(data) - 1) {
		n.right = t.buildHandle(data[divPoint+1:], (featureIndex+1)%len(t.data[data[0]]))
	}

	return n
}

// Search return []int contained k nearest neighbor from
// specific distance function.
func (t *Tree) Search(k int, disType pairwise.PairwiseDistanceFunc, target []float64) ([]int, error) {
	if k > len(t.data) {
		return []int{}, errors.New("k is largerer than amount of trainData")
	}

	if len(target) != len(t.data[0]) {
		return []int{}, errors.New("amount of features is not equal")
	}

	h := newHeap()
	t.searchHandle(k, disType, target, h, t.firstDiv)

	out := make([]int, k)
	i := k - 1
	for h.size() != 0 {
		out[i] = h.maximum().srcRowNo
		i--
		h.extractMax()
	}

	return out, nil
}

func (t *Tree) searchHandle(k int, disType pairwise.PairwiseDistanceFunc, target []float64, h *heap, n *node) {
	if n.feature == -1 {
		vectorX := mat64.NewDense(len(target), 1, target)
		vectorY := mat64.NewDense(len(target), 1, n.value)
		length := disType.Distance(vectorX, vectorY)
		h.insert(n.value, length, n.srcRowNo)
		return
	}

	dir := true
	if target[n.feature] <= n.value[n.feature] {
		t.searchHandle(k, disType, target, h, n.left)
	} else {
		dir = false
		t.searchHandle(k, disType, target, h, n.right)
	}

	vectorX := mat64.NewDense(len(target), 1, target)
	vectorY := mat64.NewDense(len(target), 1, n.value)
	length := disType.Distance(vectorX, vectorY)

	if k > h.size() {
		h.insert(n.value, length, n.srcRowNo)
		if dir {
			t.searchAllNode(k, disType, target, h, n.right)
		} else {
			t.searchAllNode(k, disType, target, h, n.left)
		}
	} else if h.maximum().length > length {
		h.extractMax()
		h.insert(n.value, length, n.srcRowNo)
		if dir {
			t.searchAllNode(k, disType, target, h, n.right)
		} else {
			t.searchAllNode(k, disType, target, h, n.left)
		}
	}
}

func (t *Tree) searchAllNode(k int, disType pairwise.PairwiseDistanceFunc, target []float64, h *heap, n *node) {
	vectorX := mat64.NewDense(len(target), 1, target)
	vectorY := mat64.NewDense(len(target), 1, n.value)
	length := disType.Distance(vectorX, vectorY)

	if k > h.size() {
		h.insert(n.value, length, n.srcRowNo)
	} else if h.maximum().length > length {
		h.extractMax()
		h.insert(n.value, length, n.srcRowNo)
	}

	if n.left != nil {
		t.searchAllNode(k, disType, target, h, n.left)
	}
	if n.right != nil {
		t.searchAllNode(k, disType, target, h, n.right)
	}
}
<commit_msg>update to latest and change searchAllNode to searchAllNodes<commit_after>package kdtree

import (
	"errors"
	"github.com/gonum/matrix/mat64"
	"github.com/sjwhitworth/golearn/metrics/pairwise"
	"sort"
)

type node struct {
	feature  int
	value    []float64
	srcRowNo int
	left     *node
	right    *node
}

// Tree is a kdtree.
type Tree struct {
	firstDiv *node
	data     [][]float64
}

type SortData struct {
	RowData [][]float64
	Data    []int
	Feature int
}

func (d SortData) Len() int { return len(d.Data) }
func (d SortData) Less(i, j int) bool {
	return d.RowData[d.Data[i]][d.Feature] < d.RowData[d.Data[j]][d.Feature]
}
func (d SortData) Swap(i, j int) { d.Data[i], d.Data[j] = d.Data[j], d.Data[i] }

// New return a Tree pointer.
func New() *Tree {
	return &Tree{}
}

// Build builds the kdtree with specific data.
func (t *Tree) Build(data [][]float64) error {
	if len(data) == 0 {
		return errors.New("no input data")
	}
	size := len(data[0])
	for _, v := range data {
		if len(v) != size {
			return errors.New("amounts of features are not the same")
		}
	}

	t.data = data

	newData := make([]int, len(data))
	for k, _ := range newData {
		newData[k] = k
	}

	if len(data) == 1 {
		t.firstDiv = &node{feature: -1, srcRowNo: 0}
		t.firstDiv.value = make([]float64, len(data[0]))
		copy(t.firstDiv.value, data[0])
	} else {
		t.firstDiv = t.buildHandle(newData, 0)
	}

	return nil
}

// buildHandle builds the kdtree recursively.
func (t *Tree) buildHandle(data []int, featureIndex int) *node {
	n := &node{feature: featureIndex}

	tmp := SortData{RowData: t.data, Data: data, Feature: featureIndex}
	sort.Sort(tmp)
	middle := len(data) / 2

	n.srcRowNo = data[middle]
	n.value = make([]float64, len(t.data[data[middle]]))
	copy(n.value, t.data[data[middle]])

	divPoint := middle
	for i := middle + 1; i < len(data); i++ {
		if t.data[data[i]][featureIndex] == t.data[data[middle]][featureIndex] {
			divPoint = i
		} else {
			break
		}
	}

	if divPoint == 1 {
		n.left = &node{feature: -1}
		n.left.value = make([]float64, len(t.data[data[0]]))
		copy(n.left.value, t.data[data[0]])
		n.left.srcRowNo = data[0]
	} else {
		n.left = t.buildHandle(data[:divPoint], (featureIndex+1)%len(t.data[data[0]]))
	}

	if divPoint == (len(data) - 2) {
		n.right = &node{feature: -1}
		n.right.value = make([]float64, len(t.data[data[divPoint+1]]))
		copy(n.right.value, t.data[data[divPoint+1]])
		n.left.srcRowNo = data[divPoint+1]
	} else if divPoint != (len(data) - 1) {
		n.right = t.buildHandle(data[divPoint+1:], (featureIndex+1)%len(t.data[data[0]]))
	}

	return n
}

// Search return []int contained k nearest neighbor from
// specific distance function.
func (t *Tree) Search(k int, disType pairwise.PairwiseDistanceFunc, target []float64) ([]int, error) {
	if k > len(t.data) {
		return []int{}, errors.New("k is largerer than amount of trainData")
	}

	if len(target) != len(t.data[0]) {
		return []int{}, errors.New("amount of features is not equal")
	}

	h := newHeap()
	t.searchHandle(k, disType, target, h, t.firstDiv)

	out := make([]int, k)
	i := k - 1
	for h.size() != 0 {
		out[i] = h.maximum().srcRowNo
		i--
		h.extractMax()
	}

	return out, nil
}

func (t *Tree) searchHandle(k int, disType pairwise.PairwiseDistanceFunc, target []float64, h *heap, n *node) {
	if n.feature == -1 {
		vectorX := mat64.NewDense(len(target), 1, target)
		vectorY := mat64.NewDense(len(target), 1, n.value)
		length := disType.Distance(vectorX, vectorY)
		h.insert(n.value, length, n.srcRowNo)
		return
	}

	dir := true
	if target[n.feature] <= n.value[n.feature] {
		t.searchHandle(k, disType, target, h, n.left)
	} else {
		dir = false
		t.searchHandle(k, disType, target, h, n.right)
	}

	vectorX := mat64.NewDense(len(target), 1, target)
	vectorY := mat64.NewDense(len(target), 1, n.value)
	length := disType.Distance(vectorX, vectorY)

	if k > h.size() {
		h.insert(n.value, length, n.srcRowNo)
		if dir {
			t.searchAllNodes(k, disType, target, h, n.right)
		} else {
			t.searchAllNodes(k, disType, target, h, n.left)
		}
	} else if h.maximum().length > length {
		h.extractMax()
		h.insert(n.value, length, n.srcRowNo)
		if dir {
			t.searchAllNodes(k, disType, target, h, n.right)
		} else {
			t.searchAllNodes(k, disType, target, h, n.left)
		}
	}
}

func (t *Tree) searchAllNodes(k int, disType pairwise.PairwiseDistanceFunc, target []float64, h *heap, n *node) {
	vectorX := mat64.NewDense(len(target), 1, target)
	vectorY := mat64.NewDense(len(target), 1, n.value)
	length := disType.Distance(vectorX, vectorY)

	if k > h.size() {
		h.insert(n.value, length, n.srcRowNo)
	} else if h.maximum().length > length {
		h.extractMax()
		h.insert(n.value, length, n.srcRowNo)
	}

	if n.left != nil {
		t.searchAllNodes(k, disType, target, h, n.left)
	}
	if n.right != nil {
		t.searchAllNodes(k, disType, target, h, n.right)
	}
}
<|endoftext|>