<commit_before>package main

import (
	"fmt"

	"flag"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/coredns/coredns/coremain"
	_ "github.com/coredns/coredns/plugin/bind"
	_ "github.com/coredns/coredns/plugin/cache"
	_ "github.com/coredns/coredns/plugin/errors"
	_ "github.com/coredns/coredns/plugin/forward"
	_ "github.com/coredns/coredns/plugin/health"
	_ "github.com/coredns/coredns/plugin/log"
	_ "github.com/coredns/coredns/plugin/loop"
	_ "github.com/coredns/coredns/plugin/metrics"
	clog "github.com/coredns/coredns/plugin/pkg/log"
	_ "github.com/coredns/coredns/plugin/reload"
	"github.com/mholt/caddy"
	"k8s.io/dns/pkg/netif"
	"k8s.io/kubernetes/pkg/util/dbus"
	utilexec "k8s.io/kubernetes/pkg/util/exec"
	utiliptables "k8s.io/kubernetes/pkg/util/iptables"
)

// configParams lists the configuration options that can be provided to dns-cache
type configParams struct {
	localIP              string        // ip address for the local cache agent to listen for dns requests
	localPort            string        // port to listen for dns requests
	metricsListenAddress string        // address to serve metrics on
	interfaceName        string        // Name of the interface to be created
	interval             time.Duration // specifies how often to run iptables rules check
	exitChan             chan bool     // Channel to terminate background goroutines
}

type iptablesRule struct {
	table utiliptables.Table
	chain utiliptables.Chain
	args  []string
}

type cacheApp struct {
	iptables      utiliptables.Interface
	iptablesRules []iptablesRule
	params        configParams
	netifHandle   *netif.NetifManager
}

var cache = cacheApp{params: configParams{localPort: "53"}}

func isLockedErr(err error) bool {
	return strings.Contains(err.Error(), "holding the xtables lock")
}

func (c *cacheApp) Init() {
	err := c.parseAndValidateFlags()
	if err != nil {
		clog.Fatalf("Error parsing flags - %s, Exiting", err)
	}
	c.netifHandle = netif.NewNetifManager(net.ParseIP(c.params.localIP))
	c.initIptables()
	err = c.teardownNetworking()
	if err != nil {
		// It is likely to hit errors here if previous shutdown cleaned up all iptables rules and interface.
		// Logging error at info level
		clog.Infof("Hit error during teardown - %s", err)
	}
	err = c.setupNetworking()
	if err != nil {
		cache.teardownNetworking()
		clog.Fatalf("Failed to setup - %s, Exiting", err)
	}
	initMetrics(c.params.metricsListenAddress)
}

func init() {
	cache.Init()
	caddy.OnProcessExit = append(caddy.OnProcessExit, func() { cache.teardownNetworking() })
}

func (c *cacheApp) initIptables() {

	c.iptablesRules = []iptablesRule{
		// Match traffic destined for localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
		{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "tcp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "NOTRACK"}},
		{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "udp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "NOTRACK"}},
		// There are rules in filter table to allow tracked connections to be accepted. Since we skipped connection tracking,
		// need these additional filter table rules.
		{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "tcp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "ACCEPT"}},
		{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "udp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "ACCEPT"}},
		// Match traffic from localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
		{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "NOTRACK"}},
		{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "udp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "NOTRACK"}},
		// Additional filter table rules for traffic frpm localIp:localPort
		{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "tcp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "ACCEPT"}},
		{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "udp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "ACCEPT"}},
	}
	c.iptables = newIPTables()
}

func newIPTables() utiliptables.Interface {
	execer := utilexec.New()
	dbus := dbus.New()
	return utiliptables.New(execer, dbus, utiliptables.ProtocolIpv4)
}

func (c *cacheApp) setupNetworking() error {
	var err error
	clog.Infof("Setting up networking for node cache")
	err = c.netifHandle.AddDummyDevice(c.params.interfaceName)
	if err != nil {
		return err
	}
	for _, rule := range c.iptablesRules {
		_, err = c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
		if err != nil {
			return err
		}
	}
	return err
}

func (c *cacheApp) teardownNetworking() error {
	clog.Infof("Tearing down")
	if c.params.exitChan != nil {
		// Stop the goroutine that periodically checks for iptables rules/dummy interface
		// exitChan is a buffered channel of size 1, so this will not block
		c.params.exitChan <- true
	}
	err := c.netifHandle.RemoveDummyDevice(c.params.interfaceName)
	for _, rule := range c.iptablesRules {
		exists := true
		for exists == true {
			c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
			exists, _ = c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
		}
		// Delete the rule one last time since EnsureRule creates the rule if it doesn't exist
		c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
	}
	return err
}

func (c *cacheApp) parseAndValidateFlags() error {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Runs coreDNS v1.2.5 as a nodelocal cache listening on the specified ip:port")
		flag.PrintDefaults()
	}

	flag.StringVar(&c.params.localIP, "localip", "", "ip address to bind dnscache to")
	flag.StringVar(&c.params.interfaceName, "interfacename", "nodelocaldns", "name of the interface to be created")
	flag.DurationVar(&c.params.interval, "syncinterval", 60, "interval(in seconds) to check for iptables rules")
	flag.StringVar(&c.params.metricsListenAddress, "metrics-listen-address", "0.0.0.0:9353", "address to serve metrics on")
	flag.Parse()

	if net.ParseIP(c.params.localIP) == nil {
		return fmt.Errorf("Invalid localip specified - %q", c.params.localIP)
	}

	// lookup specified dns port
	if f := flag.Lookup("dns.port"); f == nil {
		return fmt.Errorf("Failed to lookup \"dns.port\" parameter")
	} else {
		c.params.localPort = f.Value.String()
	}
	if _, err := strconv.Atoi(c.params.localPort); err != nil {
		return fmt.Errorf("Invalid port specified - %q", c.params.localPort)
	}
	return nil
}

func (c *cacheApp) runChecks() {
	for _, rule := range c.iptablesRules {
		exists, err := c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
		switch {
		case exists:
			// debug messages can be printed by including "debug" plugin in coreFile.
			clog.Debugf("iptables rule %v for nodelocaldns already exists", rule)
			continue
		case err == nil:
			clog.Infof("Added back nodelocaldns rule - %v", rule)
			continue
		// if we got here, either iptables check failed or adding rule back failed.
		case isLockedErr(err):
			clog.Infof("Error checking/adding iptables rule %v, due to xtables lock in use, retrying in %v", rule, c.params.interval)
			setupErrCount.WithLabelValues("iptables_lock").Inc()
		default:
			clog.Errorf("Error adding iptables rule %v - %s", rule, err)
			setupErrCount.WithLabelValues("iptables").Inc()
		}
	}

	exists, err := c.netifHandle.EnsureDummyDevice(c.params.interfaceName)
	if !exists {
		if err != nil {
			clog.Errorf("Failed to add non-existent interface %s: %s", c.params.interfaceName, err)
			setupErrCount.WithLabelValues("interface_add").Inc()
		}
		clog.Infof("Added back interface - %s", c.params.interfaceName)
	}
	if err != nil {
		clog.Errorf("Error checking dummy device %s - %s", c.params.interfaceName, err)
		setupErrCount.WithLabelValues("interface_check").Inc()
	}
}

func (c *cacheApp) run() {
	c.params.exitChan = make(chan bool, 1)
	tick := time.NewTicker(c.params.interval * time.Second)
	for {
		select {
		case <-tick.C:
			c.runChecks()
		case <-c.params.exitChan:
			clog.Warningf("Exiting iptables check goroutine")
			return
		}
	}
}

func main() {
	// Ensure that the required setup is ready
	// https://github.com/kubernetes/dns/issues/282 sometimes the interface gets the ip and then loses it, if added too soon.
	cache.runChecks()
	go cache.run()
	coremain.Run()
	// Unlikely to reach here, if we did it is because coremain exited and the signal was not trapped.
	clog.Errorf("Untrapped signal, tearing down")
	cache.teardownNetworking()
}
<commit_msg>Add flag to skip iptables rules setup in nodecache<commit_after>package main

import (
	"fmt"

	"flag"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/coredns/coredns/coremain"
	_ "github.com/coredns/coredns/plugin/bind"
	_ "github.com/coredns/coredns/plugin/cache"
	_ "github.com/coredns/coredns/plugin/errors"
	_ "github.com/coredns/coredns/plugin/forward"
	_ "github.com/coredns/coredns/plugin/health"
	_ "github.com/coredns/coredns/plugin/log"
	_ "github.com/coredns/coredns/plugin/loop"
	_ "github.com/coredns/coredns/plugin/metrics"
	clog "github.com/coredns/coredns/plugin/pkg/log"
	_ "github.com/coredns/coredns/plugin/reload"
	"github.com/mholt/caddy"
	"k8s.io/dns/pkg/netif"
	"k8s.io/kubernetes/pkg/util/dbus"
	utilexec "k8s.io/kubernetes/pkg/util/exec"
	utiliptables "k8s.io/kubernetes/pkg/util/iptables"
)

// configParams lists the configuration options that can be provided to dns-cache
type configParams struct {
	localIP              string        // ip address for the local cache agent to listen for dns requests
	localPort            string        // port to listen for dns requests
	metricsListenAddress string        // address to serve metrics on
	interfaceName        string        // Name of the interface to be created
	interval             time.Duration // specifies how often to run iptables rules check
	exitChan             chan bool     // Channel to terminate background goroutines
}

type iptablesRule struct {
	table utiliptables.Table
	chain utiliptables.Chain
	args  []string
}

type cacheApp struct {
	setupIptables bool
	iptables      utiliptables.Interface
	iptablesRules []iptablesRule
	params        configParams
	netifHandle   *netif.NetifManager
}

var cache = cacheApp{params: configParams{localPort: "53"}}

func isLockedErr(err error) bool {
	return strings.Contains(err.Error(), "holding the xtables lock")
}

func (c *cacheApp) Init() {
	err := c.parseAndValidateFlags()
	if err != nil {
		clog.Fatalf("Error parsing flags - %s, Exiting", err)
	}
	c.netifHandle = netif.NewNetifManager(net.ParseIP(c.params.localIP))
	if c.setupIptables {
		c.initIptables()
	}
	err = c.teardownNetworking()
	if err != nil {
		// It is likely to hit errors here if previous shutdown cleaned up all iptables rules and interface.
		// Logging error at info level
		clog.Infof("Hit error during teardown - %s", err)
	}
	err = c.setupNetworking()
	if err != nil {
		cache.teardownNetworking()
		clog.Fatalf("Failed to setup - %s, Exiting", err)
	}
	initMetrics(c.params.metricsListenAddress)
}

func init() {
	cache.Init()
	caddy.OnProcessExit = append(caddy.OnProcessExit, func() { cache.teardownNetworking() })
}

func (c *cacheApp) initIptables() {

	c.iptablesRules = []iptablesRule{
		// Match traffic destined for localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
		{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "tcp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "NOTRACK"}},
		{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "udp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "NOTRACK"}},
		// There are rules in filter table to allow tracked connections to be accepted. Since we skipped connection tracking,
		// need these additional filter table rules.
		{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "tcp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "ACCEPT"}},
		{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "udp", "-d", c.params.localIP,
			"--dport", c.params.localPort, "-j", "ACCEPT"}},
		// Match traffic from localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
		{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "NOTRACK"}},
		{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "udp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "NOTRACK"}},
		// Additional filter table rules for traffic frpm localIp:localPort
		{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "tcp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "ACCEPT"}},
		{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "udp", "-s", c.params.localIP,
			"--sport", c.params.localPort, "-j", "ACCEPT"}},
	}
	c.iptables = newIPTables()
}

func newIPTables() utiliptables.Interface {
	execer := utilexec.New()
	dbus := dbus.New()
	return utiliptables.New(execer, dbus, utiliptables.ProtocolIpv4)
}

func (c *cacheApp) setupNetworking() error {
	var err error
	clog.Infof("Setting up networking for node cache")
	err = c.netifHandle.AddDummyDevice(c.params.interfaceName)
	if err != nil {
		return err
	}
	if c.setupIptables {
		for _, rule := range c.iptablesRules {
			_, err = c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			if err != nil {
				return err
			}
		}
	}
	return err
}

func (c *cacheApp) teardownNetworking() error {
	clog.Infof("Tearing down")
	if c.params.exitChan != nil {
		// Stop the goroutine that periodically checks for iptables rules/dummy interface
		// exitChan is a buffered channel of size 1, so this will not block
		c.params.exitChan <- true
	}
	err := c.netifHandle.RemoveDummyDevice(c.params.interfaceName)
	if c.setupIptables {
		for _, rule := range c.iptablesRules {
			exists := true
			for exists == true {
				c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
				exists, _ = c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			}
			// Delete the rule one last time since EnsureRule creates the rule if it doesn't exist
			c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
		}
	}
	return err
}

func (c *cacheApp) parseAndValidateFlags() error {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Runs coreDNS v1.2.5 as a nodelocal cache listening on the specified ip:port")
		flag.PrintDefaults()
	}

	flag.StringVar(&c.params.localIP, "localip", "", "ip address to bind dnscache to")
	flag.StringVar(&c.params.interfaceName, "interfacename", "nodelocaldns", "name of the interface to be created")
	flag.DurationVar(&c.params.interval, "syncinterval", 60, "interval(in seconds) to check for iptables rules")
	flag.StringVar(&c.params.metricsListenAddress, "metrics-listen-address", "0.0.0.0:9353", "address to serve metrics on")
	flag.BoolVar(&c.setupIptables, "setupiptables", true, "indicates whether iptables rules should be setup")
	flag.Parse()

	if net.ParseIP(c.params.localIP) == nil {
		return fmt.Errorf("Invalid localip specified - %q", c.params.localIP)
	}

	// lookup specified dns port
	if f := flag.Lookup("dns.port"); f == nil {
		return fmt.Errorf("Failed to lookup \"dns.port\" parameter")
	} else {
		c.params.localPort = f.Value.String()
	}
	if _, err := strconv.Atoi(c.params.localPort); err != nil {
		return fmt.Errorf("Invalid port specified - %q", c.params.localPort)
	}
	return nil
}

func (c *cacheApp) runChecks() {
	if c.setupIptables {
		for _, rule := range c.iptablesRules {
			exists, err := c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			switch {
			case exists:
				// debug messages can be printed by including "debug" plugin in coreFile.
				clog.Debugf("iptables rule %v for nodelocaldns already exists", rule)
				continue
			case err == nil:
				clog.Infof("Added back nodelocaldns rule - %v", rule)
				continue
			// if we got here, either iptables check failed or adding rule back failed.
			case isLockedErr(err):
				clog.Infof("Error checking/adding iptables rule %v, due to xtables lock in use, retrying in %v", rule, c.params.interval)
				setupErrCount.WithLabelValues("iptables_lock").Inc()
			default:
				clog.Errorf("Error adding iptables rule %v - %s", rule, err)
				setupErrCount.WithLabelValues("iptables").Inc()
			}
		}
	}

	exists, err := c.netifHandle.EnsureDummyDevice(c.params.interfaceName)
	if !exists {
		if err != nil {
			clog.Errorf("Failed to add non-existent interface %s: %s", c.params.interfaceName, err)
			setupErrCount.WithLabelValues("interface_add").Inc()
		}
		clog.Infof("Added back interface - %s", c.params.interfaceName)
	}
	if err != nil {
		clog.Errorf("Error checking dummy device %s - %s", c.params.interfaceName, err)
		setupErrCount.WithLabelValues("interface_check").Inc()
	}
}

func (c *cacheApp) run() {
	c.params.exitChan = make(chan bool, 1)
	tick := time.NewTicker(c.params.interval * time.Second)
	for {
		select {
		case <-tick.C:
			c.runChecks()
		case <-c.params.exitChan:
			clog.Warningf("Exiting iptables/interface check goroutine")
			return
		}
	}
}

func main() {
	// Ensure that the required setup is ready
	// https://github.com/kubernetes/dns/issues/282 sometimes the interface gets the ip and then loses it, if added too soon.
	cache.runChecks()
	go cache.run()
	coremain.Run()
	// Unlikely to reach here, if we did it is because coremain exited and the signal was not trapped.
	clog.Errorf("Untrapped signal, tearing down")
	cache.teardownNetworking()
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/pkg/transport"
	"github.com/pingcap/kvproto/pkg/metapb"
)

var (
	endpoints   = flag.String("endpoints", "http://127.0.0.1:2379", "endpoints urls")
	allocID     = flag.Uint64("alloc-id", 0, "please make sure alloced ID is safe")
	clusterID   = flag.Uint64("cluster-id", 0, "please make cluster ID match with tikv")
	maxReplicas = flag.Int("max-replicas", 3, "max replicas is the number of replicas for each region")
	caPath      = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs.")
	certPath    = flag.String("cert", "", "path of file that contains X509 certificate in PEM format..")
	keyPath     = flag.String("key", "", "path of file that contains X509 key in PEM format.")
)

const (
	requestTimeout = 10 * time.Second
	etcdTimeout    = 3 * time.Second

	pdRootPath      = "/pd"
	pdClusterIDPath = "/pd/cluster_id"
)

func exitErr(err error) {
	fmt.Println(err.Error())
	os.Exit(1)
}

func main() {
	flag.Parse()
	if *clusterID == 0 {
		fmt.Println("please specify safe cluster-id")
		return
	}
	if *allocID == 0 {
		fmt.Println("please specify safe alloc-id")
		return
	}

	rootPath := path.Join(pdRootPath, strconv.FormatUint(*clusterID, 10))
	clusterRootPath := path.Join(rootPath, "raft")
	raftBootstrapTimeKey := path.Join(clusterRootPath, "status", "raft_bootstrap_time")

	urls := strings.Split(*endpoints, ",")

	tlsInfo := transport.TLSInfo{
		CertFile:      *certPath,
		KeyFile:       *keyPath,
		TrustedCAFile: *caPath,
	}
	tlsConfig, err := tlsInfo.ClientConfig()
	if err != nil {
		fmt.Println("failed to connect: err")
		return
	}

	client, err := clientv3.New(clientv3.Config{
		Endpoints:   urls,
		DialTimeout: etcdTimeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		exitErr(err)
	}
	ctx, cancel := context.WithTimeout(client.Ctx(), requestTimeout)
	defer cancel()

	var ops []clientv3.Op
	// recover cluster_id
	ops = append(ops, clientv3.OpPut(pdClusterIDPath, string(uint64ToBytes(*clusterID))))
	// recover alloc_id
	allocIDPath := path.Join(rootPath, "alloc_id")
	ops = append(ops, clientv3.OpPut(allocIDPath, string(uint64ToBytes(*allocID))))

	// recover bootstrap
	// recover meta of cluster
	clusterMeta := metapb.Cluster{
		Id:           *clusterID,
		MaxPeerCount: uint32(*maxReplicas),
	}
	clusterValue, err := clusterMeta.Marshal()
	if err != nil {
		exitErr(err)
	}
	ops = append(ops, clientv3.OpPut(clusterRootPath, string(clusterValue)))

	// set raft bootstrap time
	nano := time.Now().UnixNano()
	timeData := uint64ToBytes(uint64(nano))
	ops = append(ops, clientv3.OpPut(raftBootstrapTimeKey, string(timeData)))

	// the new pd cluster should not bootstrapped by tikv
	bootstrapCmp := clientv3.Compare(clientv3.CreateRevision(clusterRootPath), "=", 0)
	resp, err := client.Txn(ctx).If(bootstrapCmp).Then(ops...).Commit()
	if err != nil {
		exitErr(err)
	}
	if !resp.Succeeded {
		fmt.Println("failed to recover: the cluster is already bootstrapped")
		return
	}
	fmt.Println("recover success! please restart the PD cluster")
}

func uint64ToBytes(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<commit_msg>pd-recover: delete max replicas (#1076)<commit_after>package main

import (
	"context"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/pkg/transport"
	"github.com/pingcap/kvproto/pkg/metapb"
)

var (
	endpoints = flag.String("endpoints", "http://127.0.0.1:2379", "endpoints urls")
	allocID   = flag.Uint64("alloc-id", 0, "please make sure alloced ID is safe")
	clusterID = flag.Uint64("cluster-id", 0, "please make cluster ID match with tikv")
	caPath    = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs.")
	certPath  = flag.String("cert", "", "path of file that contains X509 certificate in PEM format..")
	keyPath   = flag.String("key", "", "path of file that contains X509 key in PEM format.")
)

const (
	requestTimeout = 10 * time.Second
	etcdTimeout    = 3 * time.Second

	pdRootPath      = "/pd"
	pdClusterIDPath = "/pd/cluster_id"
)

func exitErr(err error) {
	fmt.Println(err.Error())
	os.Exit(1)
}

func main() {
	flag.Parse()
	if *clusterID == 0 {
		fmt.Println("please specify safe cluster-id")
		return
	}
	if *allocID == 0 {
		fmt.Println("please specify safe alloc-id")
		return
	}

	rootPath := path.Join(pdRootPath, strconv.FormatUint(*clusterID, 10))
	clusterRootPath := path.Join(rootPath, "raft")
	raftBootstrapTimeKey := path.Join(clusterRootPath, "status", "raft_bootstrap_time")

	urls := strings.Split(*endpoints, ",")

	tlsInfo := transport.TLSInfo{
		CertFile:      *certPath,
		KeyFile:       *keyPath,
		TrustedCAFile: *caPath,
	}
	tlsConfig, err := tlsInfo.ClientConfig()
	if err != nil {
		fmt.Println("failed to connect: err")
		return
	}

	client, err := clientv3.New(clientv3.Config{
		Endpoints:   urls,
		DialTimeout: etcdTimeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		exitErr(err)
	}
	ctx, cancel := context.WithTimeout(client.Ctx(), requestTimeout)
	defer cancel()

	var ops []clientv3.Op
	// recover cluster_id
	ops = append(ops, clientv3.OpPut(pdClusterIDPath, string(uint64ToBytes(*clusterID))))
	// recover alloc_id
	allocIDPath := path.Join(rootPath, "alloc_id")
	ops = append(ops, clientv3.OpPut(allocIDPath, string(uint64ToBytes(*allocID))))

	// recover bootstrap
	// recover meta of cluster
	clusterMeta := metapb.Cluster{Id: *clusterID}
	clusterValue, err := clusterMeta.Marshal()
	if err != nil {
		exitErr(err)
	}
	ops = append(ops, clientv3.OpPut(clusterRootPath, string(clusterValue)))

	// set raft bootstrap time
	nano := time.Now().UnixNano()
	timeData := uint64ToBytes(uint64(nano))
	ops = append(ops, clientv3.OpPut(raftBootstrapTimeKey, string(timeData)))

	// the new pd cluster should not bootstrapped by tikv
	bootstrapCmp := clientv3.Compare(clientv3.CreateRevision(clusterRootPath), "=", 0)
	resp, err := client.Txn(ctx).If(bootstrapCmp).Then(ops...).Commit()
	if err != nil {
		exitErr(err)
	}
	if !resp.Succeeded {
		fmt.Println("failed to recover: the cluster is already bootstrapped")
		return
	}
	fmt.Println("recover success! please restart the PD cluster")
}

func uint64ToBytes(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/ncw/rclone/cmd"
	"github.com/ncw/rclone/fs"
	"github.com/spf13/cobra"
)

// Globals
var (
	bindAddress = "localhost:8080"
	readWrite   = false
)

func init() {
	Command.Flags().StringVarP(&bindAddress, "addr", "", bindAddress, "IPaddress:Port to bind server to.")
	// Command.Flags().BoolVarP(&readWrite, "rw", "", readWrite, "Serve in read/write mode.")
}

// Command definition for cobra
var Command = &cobra.Command{
	Use:   "http remote:path",
	Short: `Serve the remote over HTTP.`,
	Long: `rclone serve http implements a basic web server to serve the remote
over HTTP.  This can be viewed in a web browser or you can make a
remote of type http read from it.

Use --addr to specify which IP address and port the server should
listen on, eg --addr 1.2.3.4:8000 or --addr :8080 to listen to all
IPs.  By default it only listens on localhost.

You can use the filter flags (eg --include, --exclude) to control what
is served.

The server will log errors.  Use -v to see access logs.

--bwlimit will be respected for file transfers.  Use --stats to
control the stats printing.
`,
	Run: func(command *cobra.Command, args []string) {
		cmd.CheckArgs(1, 1, command, args)
		f := cmd.NewFsSrc(args)
		cmd.Run(false, true, command, func() error {
			s := server{
				f:           f,
				bindAddress: bindAddress,
				readWrite:   readWrite,
			}
			s.serve()
			return nil
		})
	},
}

// server contains everything to run the server
type server struct {
	f           fs.Fs
	bindAddress string
	readWrite   bool
}

// serve creates the http server
func (s *server) serve() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", s.handler)
	// FIXME make a transport?
	httpServer := &http.Server{
		Addr:              s.bindAddress,
		Handler:           mux,
		ReadHeaderTimeout: 10 * time.Second, // time to send the headers
		IdleTimeout:       60 * time.Second, // time to keep idle connections open
		MaxHeaderBytes:    1 << 20,
	}
	fs.Logf(s.f, "Serving on http://%s/", bindAddress)
	log.Fatal(httpServer.ListenAndServe())
}

// handler reads incoming requests and dispatches them
func (s *server) handler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" && r.Method != "HEAD" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	urlPath := r.URL.Path
	isDir := strings.HasSuffix(urlPath, "/")
	remote := strings.Trim(urlPath, "/")
	if isDir {
		s.serveDir(w, r, remote)
	} else {
		s.serveFile(w, r, remote)
	}
}

// entry is a directory entry
type entry struct {
	remote string
	URL    string
	Leaf   string
}

// entries represents a directory
type entries []entry

// indexPage is a directory listing template
var indexPage = `<!DOCTYPE html>

<meta charset="utf-8">
<title>{{ .Title }}</title>

{{ .Title }}

{{ range $i := .Entries }}{{ $i.Leaf }}

{{ end }}

`

// indexTemplate is the instantiated indexPage
var indexTemplate = template.Must(template.New("index").Parse(indexPage))

// indexData is used to fill in the indexTemplate
type indexData struct {
	Title   string
	Entries entries
}

// error returns an http.StatusInternalServerError and logs the error
func internalError(what interface{}, w http.ResponseWriter, text string, err error) {
	fs.Stats.Error()
	fs.Errorf(what, "%s: %v", text, err)
	http.Error(w, text+".", http.StatusInternalServerError)
}

// serveDir serves a directory index at dirRemote
func (s *server) serveDir(w http.ResponseWriter, r *http.Request, dirRemote string) {
	// Check the directory is included in the filters
	if !fs.Config.Filter.IncludeDirectory(dirRemote) {
		fs.Infof(dirRemote, "%s: Directory not found (filtered)", r.RemoteAddr)
		http.Error(w, "Directory not found", http.StatusNotFound)
		return
	}

	// List the directory
	dirEntries, err := fs.ListDirSorted(s.f, false, dirRemote)
	if err == fs.ErrorDirNotFound {
		fs.Infof(dirRemote, "%s: Directory not found", r.RemoteAddr)
		http.Error(w, "Directory not found", http.StatusNotFound)
		return
	} else if err != nil {
		internalError(dirRemote, w, "Failed to list directory", err)
		return
	}

	var out entries
	for _, o := range dirEntries {
		remote := strings.Trim(o.Remote(), "/")
		leaf := path.Base(remote)
		urlRemote := leaf
		if _, ok := o.(*fs.Dir); ok {
			leaf += "/"
			urlRemote += "/"
		}
		out = append(out, entry{remote: remote, URL: urlRemote, Leaf: leaf})
	}

	// Account the transfer
	fs.Stats.Transferring(dirRemote)
	defer fs.Stats.DoneTransferring(dirRemote, true)

	fs.Infof(dirRemote, "%s: Serving directory", r.RemoteAddr)
	err = indexTemplate.Execute(w, indexData{
		Entries: out,
		Title:   fmt.Sprintf("Directory listing of /%s", dirRemote),
	})
	if err != nil {
		internalError(dirRemote, w, "Failed to render template", err)
		return
	}
}

// serveFile serves a file object at remote
func (s *server) serveFile(w http.ResponseWriter, r *http.Request, remote string) {
	// FIXME could cache the directories and objects...
	obj, err := s.f.NewObject(remote)
	if err == fs.ErrorObjectNotFound {
		fs.Infof(remote, "%s: File not found", r.RemoteAddr)
		http.Error(w, "File not found", http.StatusNotFound)
		return
	} else if err != nil {
		internalError(remote, w, "Failed to find file", err)
		return
	}

	// Check the object is included in the filters
	if !fs.Config.Filter.IncludeObject(obj) {
		fs.Infof(remote, "%s: File not found (filtered)", r.RemoteAddr)
		http.Error(w, "File not found", http.StatusNotFound)
		return
	}

	// Set content length since we know how long the object is
	w.Header().Set("Content-Length", strconv.FormatInt(obj.Size(), 10))

	// Set content type
	mimeType := fs.MimeType(obj)
	if mimeType == "application/octet-stream" && path.Ext(remote) == "" {
		// Leave header blank so http server guesses
	} else {
		w.Header().Set("Content-Type", mimeType)
	}

	// If HEAD no need to read the object since we have set the headers
	if r.Method == "HEAD" {
		return
	}

	// open the object
	in, err := obj.Open()
	if err != nil {
		internalError(remote, w, "Failed to open file", err)
		return
	}
	defer func() {
		err := in.Close()
		if err != nil {
			fs.Errorf(remote, "Failed to close file: %v", err)
		}
	}()

	// Account the transfer
	fs.Stats.Transferring(remote)
	defer fs.Stats.DoneTransferring(remote, true)
	in = fs.NewAccount(in, obj).WithBuffer() // account the transfer

	// Copy the contents of the object to the output
	fs.Infof(remote, "%s: Serving file", r.RemoteAddr)
	_, err = io.Copy(w, in)
	if err != nil {
		fs.Errorf(remote, "Failed to write file: %v", err)
	}
}
<commit_msg>serve http: error if Range supplied (not supported yet)<commit_after>package http

import (
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/ncw/rclone/cmd"
	"github.com/ncw/rclone/fs"
	"github.com/spf13/cobra"
)

// Globals
var (
	bindAddress = "localhost:8080"
	readWrite   = false
)

func init() {
	Command.Flags().StringVarP(&bindAddress, "addr", "", bindAddress, "IPaddress:Port to bind server to.")
	// Command.Flags().BoolVarP(&readWrite, "rw", "", readWrite, "Serve in read/write mode.")
}

// Command definition for cobra
var Command = &cobra.Command{
	Use:   "http remote:path",
	Short: `Serve the remote over HTTP.`,
	Long: `rclone serve http implements a basic web server to serve the remote
over HTTP.  This can be viewed in a web browser or you can make a
remote of type http read from it.

Use --addr to specify which IP address and port the server should
listen on, eg --addr 1.2.3.4:8000 or --addr :8080 to listen to all
IPs.  By default it only listens on localhost.

You can use the filter flags (eg --include, --exclude) to control what
is served.

The server will log errors.  Use -v to see access logs.

--bwlimit will be respected for file transfers.  Use --stats to
control the stats printing.

Note the Range header is not supported yet.
`,
	Run: func(command *cobra.Command, args []string) {
		cmd.CheckArgs(1, 1, command, args)
		f := cmd.NewFsSrc(args)
		cmd.Run(false, true, command, func() error {
			s := server{
				f:           f,
				bindAddress: bindAddress,
				readWrite:   readWrite,
			}
			s.serve()
			return nil
		})
	},
}

// server contains everything to run the server
type server struct {
	f           fs.Fs
	bindAddress string
	readWrite   bool
}

// serve creates the http server
func (s *server) serve() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", s.handler)
	// FIXME make a transport?
	httpServer := &http.Server{
		Addr:              s.bindAddress,
		Handler:           mux,
		ReadHeaderTimeout: 10 * time.Second, // time to send the headers
		IdleTimeout:       60 * time.Second, // time to keep idle connections open
		MaxHeaderBytes:    1 << 20,
	}
	fs.Logf(s.f, "Serving on http://%s/", bindAddress)
	log.Fatal(httpServer.ListenAndServe())
}

// handler reads incoming requests and dispatches them
func (s *server) handler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" && r.Method != "HEAD" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	rangeHeader := r.Header.Get("Range")
	if rangeHeader != "" {
		http.Error(w, "Range not supported yet", http.StatusRequestedRangeNotSatisfiable)
		return
	}
	//r.Header().Set("Accept-Ranges", "bytes")
	w.Header().Set("Accept-Ranges", "none") // show we don't support Range yet
	w.Header().Set("Server", "rclone/"+fs.Version)

	urlPath := r.URL.Path
	isDir := strings.HasSuffix(urlPath, "/")
	remote := strings.Trim(urlPath, "/")
	if isDir {
		s.serveDir(w, r, remote)
	} else {
		s.serveFile(w, r, remote)
	}
}

// entry is a directory entry
type entry struct {
	remote string
	URL    string
	Leaf   string
}

// entries represents a directory
type entries []entry

// indexPage is a directory listing template
var indexPage = `<!DOCTYPE html>

<meta charset="utf-8">
<title>{{ .Title }}</title>

{{ .Title }}

{{ range $i := .Entries }}{{ $i.Leaf }}

{{ end }}

`

// indexTemplate is the instantiated indexPage
var indexTemplate = template.Must(template.New("index").Parse(indexPage))

// indexData is used to fill in the indexTemplate
type indexData struct {
	Title   string
	Entries entries
}

// error returns an http.StatusInternalServerError and logs the error
func internalError(what interface{}, w http.ResponseWriter, text string, err error) {
	fs.Stats.Error()
	fs.Errorf(what, "%s: %v", text, err)
	http.Error(w, text+".", http.StatusInternalServerError)
}

// serveDir serves a directory index at dirRemote
func (s *server) serveDir(w http.ResponseWriter, r *http.Request, dirRemote string) {
	// Check the directory is included in the filters
	if !fs.Config.Filter.IncludeDirectory(dirRemote) {
		fs.Infof(dirRemote, "%s: Directory not found (filtered)", r.RemoteAddr)
		http.Error(w, "Directory not found", http.StatusNotFound)
		return
	}

	// List the directory
	dirEntries, err := fs.ListDirSorted(s.f, false, dirRemote)
	if err == fs.ErrorDirNotFound {
		fs.Infof(dirRemote, "%s: Directory not found", r.RemoteAddr)
		http.Error(w, "Directory not found", http.StatusNotFound)
		return
	} else if err != nil {
		internalError(dirRemote, w, "Failed to list directory", err)
		return
	}

	var out entries
	for _, o := range dirEntries {
		remote := strings.Trim(o.Remote(), "/")
		leaf := path.Base(remote)
		urlRemote := leaf
		if _, ok := o.(*fs.Dir); ok {
			leaf += "/"
			urlRemote += "/"
		}
		out = append(out, entry{remote: remote, URL: urlRemote, Leaf: leaf})
	}

	// Account the transfer
	fs.Stats.Transferring(dirRemote)
	defer fs.Stats.DoneTransferring(dirRemote, true)

	fs.Infof(dirRemote, "%s: Serving directory", r.RemoteAddr)
	err = indexTemplate.Execute(w, indexData{
		Entries: out,
		Title:   fmt.Sprintf("Directory listing of /%s", dirRemote),
	})
	if err != nil {
		internalError(dirRemote, w, "Failed to render template", err)
		return
	}
}

// serveFile serves a file object at remote
func (s *server) serveFile(w http.ResponseWriter, r *http.Request, remote string) {
	// FIXME could cache the directories and objects...
	obj, err := s.f.NewObject(remote)
	if err == fs.ErrorObjectNotFound {
		fs.Infof(remote, "%s: File not found", r.RemoteAddr)
		http.Error(w, "File not found", http.StatusNotFound)
		return
	} else if err != nil {
		internalError(remote, w, "Failed to find file", err)
		return
	}

	// Check the object is included in the filters
	if !fs.Config.Filter.IncludeObject(obj) {
		fs.Infof(remote, "%s: File not found (filtered)", r.RemoteAddr)
		http.Error(w, "File not found", http.StatusNotFound)
		return
	}

	// Set content length since we know how long the object is
	w.Header().Set("Content-Length", strconv.FormatInt(obj.Size(), 10))

	// Set content type
	mimeType := fs.MimeType(obj)
	if mimeType == "application/octet-stream" && path.Ext(remote) == "" {
		// Leave header blank so http server guesses
	} else {
		w.Header().Set("Content-Type", mimeType)
	}

	// If HEAD no need to read the object since we have set the headers
	if r.Method == "HEAD" {
		return
	}

	// open the object
	in, err := obj.Open()
	if err != nil {
		internalError(remote, w, "Failed to open file", err)
		return
	}
	defer func() {
		err := in.Close()
		if err != nil {
			fs.Errorf(remote, "Failed to close file: %v", err)
		}
	}()

	// Account the transfer
	fs.Stats.Transferring(remote)
	defer fs.Stats.DoneTransferring(remote, true)
	in = fs.NewAccount(in, obj).WithBuffer() // account the transfer

	// Copy the contents of the object to the output
	fs.Infof(remote, "%s: Serving file", r.RemoteAddr)
	_, err = io.Copy(w, in)
	if err != nil {
		fs.Errorf(remote, "Failed to write file: %v", err)
	}
}
<|endoftext|>


<commit_before>// 每日收盤後產生符合選股條件的報告.
//
/*
Install:

	go install github.com/toomore/gogrs/cmd/twsereport

Usage:

	twsereport [flags]

The flags are:

	-twse
		上市股票代碼，可使用 ',' 分隔多組代碼，例：2618,2329
	-twsecate
		上市股票類別，可使用 ',' 分隔多組代碼，例：11,15
	-ncpu
		指定 CPU 數量，預設為實際 CPU 數量
	-color
		色彩化

*/
package main

import (
	"flag"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"
	"github.com/toomore/gogrs/tradingdays"
	"github.com/toomore/gogrs/twse"
)

type checkGroupList []checkGroup

func (c *checkGroupList) Add(f checkGroup) {
	if (*c)[0] == nil {
		(*c)[0] = f
	} else {
		*c = append(*c, f)
	}
}

var (
	wg         sync.WaitGroup
	twseNo     = flag.String("twse", "", "上市股票代碼，可使用 ',' 分隔多組代碼，例：2618,2329")
	twseCate   = flag.String("twsecate", "", "上市股票類別，可使用 ',' 分隔多組代碼，例：11,15")
	showcolor  = flag.Bool("color", true, "色彩化")
	ncpu       = flag.Int("ncpu", runtime.NumCPU(), "指定 CPU 數量，預設為實際 CPU 數量")
	ckList     = make(checkGroupList, 1)
	white      = color.New(color.FgWhite, color.Bold).SprintfFunc()
	red        = color.New(color.FgRed, color.Bold).SprintfFunc()
	green      = color.New(color.FgGreen, color.Bold).SprintfFunc()
	yellow     = color.New(color.FgYellow).SprintfFunc()
	yellowBold = color.New(color.FgYellow, color.Bold).SprintfFunc()
	blue       = color.New(color.FgBlue).SprintfFunc()
)

func init() {
	runtime.GOMAXPROCS(*ncpu)
}

func prettyprint(stock *twse.Data, check checkGroup) string {
	var (
		Price       = stock.GetPriceList()[len(stock.GetPriceList())-1]
		RangeValue  = stock.GetRangeList()[len(stock.GetRangeList())-1]
		Volume      = stock.GetVolumeList()[len(stock.GetVolumeList())-1] / 1000
		outputcolor func(string, ...interface{}) string
	)

	switch {
	case RangeValue > 0:
		outputcolor = red
	case RangeValue < 0:
		outputcolor = green
	default:
		outputcolor = white
	}

	return fmt.Sprintf("%s %s %s %s%s %s",
		yellow("[%s]", check),
		blue("%s", stock.RawData[stock.Len()-1][0]),
		outputcolor("%s %s", stock.No, stock.Name),
		outputcolor("$%.2f", Price),
		outputcolor("(%.2f)", RangeValue),
		outputcolor("%d", Volume),
	)
}

func main() {
	flag.Parse()
	var (
		datalist   []*twse.Data
		catelist   []twse.StockInfo
		twselist   []string
		catenolist []string
	)

	color.NoColor = !*showcolor

	if *twseCate != "" {
		l := &twse.Lists{Date: tradingdays.FindRecentlyOpened(time.Now())}

		for _, v := range strings.Split(*twseCate, ",") {
			catelist = l.GetCategoryList(v)
			for _, s := range catelist {
				catenolist = append(catenolist, s.No)
			}
		}
	}

	if *twseNo != "" {
		twselist = strings.Split(*twseNo, ",")
	}
	datalist = make([]*twse.Data, len(twselist)+len(catenolist))

	for i, no := range append(twselist, catenolist...) {
		datalist[i] = twse.NewTWSE(no, tradingdays.FindRecentlyOpened(time.Now()))
	}

	if len(datalist) > 0 {
		for _, check := range ckList {
			fmt.Println(yellowBold("----- %v -----", check))
			wg.Add(len(datalist))
			for _, stock := range datalist {
				go func(check checkGroup, stock *twse.Data) {
					defer wg.Done()
					runtime.Gosched()
					if check.CheckFunc(stock) {
						fmt.Println(prettyprint(stock, check))
					}
				}(check, stock)
			}
			wg.Wait()
		}
	} else {
		flag.PrintDefaults()
	}
}
<commit_msg>Add range, precent into twsereport.<commit_after>// 每日收盤後產生符合選股條件的報告.
//
/*
Install:

	go install github.com/toomore/gogrs/cmd/twsereport

Usage:

	twsereport [flags]

The flags are:

	-twse
		上市股票代碼，可使用 ',' 分隔多組代碼，例：2618,2329
	-twsecate
		上市股票類別，可使用 ',' 分隔多組代碼，例：11,15
	-ncpu
		指定 CPU 數量，預設為實際 CPU 數量
	-color
		色彩化

*/
package main

import (
	"flag"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"
	"github.com/toomore/gogrs/tradingdays"
	"github.com/toomore/gogrs/twse"
)

type checkGroupList []checkGroup

func (c *checkGroupList) Add(f checkGroup) {
	if (*c)[0] == nil {
		(*c)[0] = f
	} else {
		*c = append(*c, f)
	}
}

var (
	wg         sync.WaitGroup
	twseNo     = flag.String("twse", "", "上市股票代碼，可使用 ',' 分隔多組代碼，例：2618,2329")
	twseCate   = flag.String("twsecate", "", "上市股票類別，可使用 ',' 分隔多組代碼，例：11,15")
	showcolor  = flag.Bool("color", true, "色彩化")
	ncpu       = flag.Int("ncpu", runtime.NumCPU(), "指定 CPU 數量，預設為實際 CPU 數量")
	ckList     = make(checkGroupList, 1)
	white      = color.New(color.FgWhite, color.Bold).SprintfFunc()
	red        = color.New(color.FgRed, color.Bold).SprintfFunc()
	green      = color.New(color.FgGreen, color.Bold).SprintfFunc()
	yellow     = color.New(color.FgYellow).SprintfFunc()
	yellowBold = color.New(color.FgYellow, color.Bold).SprintfFunc()
	blue       = color.New(color.FgBlue).SprintfFunc()
)

func init() {
	runtime.GOMAXPROCS(*ncpu)
}

func prettyprint(stock *twse.Data, check checkGroup) string {
	var (
		Open        = stock.GetOpenList()[len(stock.GetOpenList())-1]
		Price       = stock.GetPriceList()[len(stock.GetPriceList())-1]
		RangeValue  = stock.GetDailyRangeList()[len(stock.GetDailyRangeList())-1]
		Volume      = stock.GetVolumeList()[len(stock.GetVolumeList())-1] / 1000
		outputcolor func(string, ...interface{}) string
	)

	switch {
	case RangeValue > 0:
		outputcolor = red
	case RangeValue < 0:
		outputcolor = green
	default:
		outputcolor = white
	}

	return fmt.Sprintf("%s %s %s %s%s %s %s",
		yellow("[%s]", check),
		blue("%s", stock.RawData[stock.Len()-1][0]),
		outputcolor("%s %s", stock.No, stock.Name),
		outputcolor("$%.2f", Price),
		outputcolor("(%.2f)", RangeValue),
		outputcolor("%.2f%%", RangeValue/Open*100),
		outputcolor("%d", Volume),
	)
}

func main() {
	flag.Parse()
	var (
		datalist   []*twse.Data
		catelist   []twse.StockInfo
		twselist   []string
		catenolist []string
	)

	color.NoColor = !*showcolor

	if *twseCate != "" {
		l := &twse.Lists{Date: tradingdays.FindRecentlyOpened(time.Now())}

		for _, v := range strings.Split(*twseCate, ",") {
			catelist = l.GetCategoryList(v)
			for _, s := range catelist {
				catenolist = append(catenolist, s.No)
			}
		}
	}

	if *twseNo != "" {
		twselist = strings.Split(*twseNo, ",")
	}
	datalist = make([]*twse.Data, len(twselist)+len(catenolist))

	for i, no := range append(twselist, catenolist...) {
		datalist[i] = twse.NewTWSE(no, tradingdays.FindRecentlyOpened(time.Now()))
	}

	if len(datalist) > 0 {
		for _, check := range ckList {
			fmt.Println(yellowBold("----- %v -----", check))
			wg.Add(len(datalist))
			for _, stock := range datalist {
				go func(check checkGroup, stock *twse.Data) {
					defer wg.Done()
					runtime.Gosched()
					if check.CheckFunc(stock) {
						fmt.Println(prettyprint(stock, check))
					}
				}(check, stock)
			}
			wg.Wait()
		}
	} else {
		flag.PrintDefaults()
	}
}
<|endoftext|>


<commit_before>package watt

import (
	"encoding/json"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/datawire/ambassador/pkg/consulwatch"
	"github.com/datawire/ambassador/pkg/k8s"
	"github.com/datawire/ambassador/pkg/kates"
	"github.com/datawire/ambassador/pkg/limiter"
	"github.com/datawire/ambassador/pkg/supervisor"
	"github.com/datawire/ambassador/pkg/watt"
)

type WatchHook func(p *supervisor.Process, snapshot string) WatchSet

type aggregator struct {
	// Input channel used to tell us about kubernetes state.
	KubernetesEvents chan k8sEvent
	// Input channel used to tell us about consul endpoints.
	ConsulEvents chan consulEvent
	// Output channel used to communicate with the k8s watch manager.
	k8sWatches chan<- []KubernetesWatchSpec
	// Output channel used to communicate with the consul watch manager.
	consulWatches chan<- []ConsulWatchSpec
	// Output channel used to communicate with the invoker.
	snapshots chan<- string
	// We won't consider ourselves "bootstrapped" until we hear
	// about all these kinds.
	requiredKinds       []string
	watchHook           WatchHook
	limiter             limiter.Limiter
	ids                 map[string]bool
	kubernetesResources map[string]map[string][]k8s.Resource
	consulEndpoints     map[string]consulwatch.Endpoints
	bootstrapped        bool
	notifyMux           sync.Mutex
	errors              map[string][]watt.Error
	validator           *kates.Validator
}

func NewAggregator(snapshots chan<- string, k8sWatches chan<- []KubernetesWatchSpec, consulWatches chan<- []ConsulWatchSpec,
	requiredKinds []string, watchHook WatchHook, limiter limiter.Limiter, validator *kates.Validator) *aggregator {
	return &aggregator{
		KubernetesEvents:    make(chan k8sEvent),
		ConsulEvents:        make(chan consulEvent),
		k8sWatches:          k8sWatches,
		consulWatches:       consulWatches,
		snapshots:           snapshots,
		requiredKinds:       requiredKinds,
		watchHook:           watchHook,
		limiter:             limiter,
		ids:                 make(map[string]bool),
		kubernetesResources: make(map[string]map[string][]k8s.Resource),
		consulEndpoints:     make(map[string]consulwatch.Endpoints),
		errors:              make(map[string][]watt.Error),
		validator:           validator,
	}
}

func (a *aggregator) Work(p *supervisor.Process) error {
	// In order to invoke `maybeNotify`, which is a very time consuming
	// operation, we coalesce events:
	//
	// 1. Be continuously reading all available events from
	//    a.KubernetesEvents and a.ConsulEvents and store k8sEvents
	//    in the potentialKubernetesEventSignal variable. This means
	//    at any given point (modulo caveats below), the
	//    potentialKubernetesEventSignal variable will have the
	//    latest Kubernetes event available.
	//
	// 2. At the same time, whenever there is capacity to write
	//    down the kubernetesEventProcessor channel, we send
	//    potentialKubernetesEventSignal to be processed.
	//
	//    The anonymous goroutine below will be constantly reading
	//    from the kubernetesEventProcessor channel and performing
	//    a blocking a.maybeNotify(). This means that we can only
	//    *write* to the kubernetesEventProcessor channel when we are
	//    not currently processing an event, but when that happens, we
	//    will still read from a.KubernetesEvents and a.ConsulEvents
	//    and update potentialKubernetesEventSignal.
	//
	// There are three caveats to the above:
	//
	// 1. At startup, we don't yet have a event to write, but
	//    we're not processing anything, so we will try to write
	//    something down the kubernetesEventProcessor channel.
	//    To cope with this, the invoking goroutine will ignore events
	//    signals that have a event.skip flag.
	//
	// 2. If we process an event quickly, or if there aren't new
	//    events available, then we end up busy looping and
	//    sending the same potentialKubernetesEventSignal value down
	//    the kubernetesEventProcessor channel multiple times. To cope
	//    with this, whenever we have successfully written to the
	//    kubernetesEventProcessor channel, we do a *blocking* read of
	//    the next event from a.KubernetesEvents and a.ConsulEvents.
	//
	// 3. Always be calling a.setKubernetesResources as soon as we
	//    receive an event. This is a fast non-blocking call that
	//    update watches, we can't coalesce this call.

	p.Ready()

	type eventSignal struct {
		kubernetesEvent k8sEvent
		skip            bool
	}

	kubernetesEventProcessor := make(chan eventSignal)
	go func() {
		for event := range kubernetesEventProcessor {
			if event.skip {
				// ignore the initial eventSignal to deal with the
				// corner case where we haven't yet received an event yet.
				continue
			}
			a.maybeNotify(p)
		}
	}()

	potentialKubernetesEventSignal := eventSignal{kubernetesEvent: k8sEvent{}, skip: true}
	for {
		select {
		case potentialKubernetesEvent := <-a.KubernetesEvents:
			// if a new KubernetesEvents is available to be read,
			// and we can't write to the kubernetesEventProcessor channel,
			// then we will overwrite potentialKubernetesEvent
			// with a newer event while still processing a.setKubernetesResources
			a.setKubernetesResources(potentialKubernetesEvent)
			potentialKubernetesEventSignal = eventSignal{kubernetesEvent: potentialKubernetesEvent, skip: false}
		case kubernetesEventProcessor <- potentialKubernetesEventSignal:
			// if we aren't currently blocked in
			// a.maybeNotify() then the above goroutine will be
			// reading from the kubernetesEventProcessor channel and we
			// will send the current potentialKubernetesEventSignal
			// value over the kubernetesEventProcessor channel to be
			// processed
			select {
			case potentialKubernetesEvent := <-a.KubernetesEvents:
				// here we do blocking read of the next event for caveat #2.
				a.setKubernetesResources(potentialKubernetesEvent)
				potentialKubernetesEventSignal = eventSignal{kubernetesEvent: potentialKubernetesEvent, skip: false}
			case event := <-a.ConsulEvents:
				a.updateConsulResources(event)
				a.maybeNotify(p)
			case <-p.Shutdown():
				return nil
			}
		case event := <-a.ConsulEvents:
			// we are always reading and processing ConsulEvents directly,
			// not coalescing them.
			a.updateConsulResources(event)
			a.maybeNotify(p)
		case <-p.Shutdown():
			return nil
		}
	}
}

func (a *aggregator) updateConsulResources(event consulEvent) {
	a.ids[event.WatchId] = true
	a.consulEndpoints[event.Endpoints.Service] = event.Endpoints
}

func (a *aggregator) setKubernetesResources(event k8sEvent) {
	if len(event.errors) > 0 {
		for _, kError := range event.errors {
			a.errors[kError.Source] = append(a.errors[kError.Source], kError)
		}
		return
	}
	a.ids[event.watchId] = true
	submap, ok := a.kubernetesResources[event.watchId]
	if !ok {
		submap = make(map[string][]k8s.Resource)
		a.kubernetesResources[event.watchId] = submap
	}
	submap[event.kind] = event.resources
}

func (a *aggregator) generateSnapshot(p *supervisor.Process) (string, error) {
	k8sResources := make(map[string][]k8s.Resource)
	for _, submap := range a.kubernetesResources {
		for k, v := range submap {
			a.validate(p, v)
			k8sResources[k] = append(k8sResources[k], v...)
		}
	}
	s := watt.Snapshot{
		Consul:     watt.ConsulSnapshot{Endpoints: a.consulEndpoints},
		Kubernetes: k8sResources,
		Errors:     a.errors,
	}

	jsonBytes, err := json.MarshalIndent(s, "", "    ")
	if err != nil {
		return "{}", err
	}

	return string(jsonBytes), nil
}

func (a *aggregator) validate(p *supervisor.Process, resources []k8s.Resource) {
	for _, r := range resources {
		err := a.validator.Validate(p.Context(), map[string]interface{}(r))
		if err == nil {
			delete(r, "errors")
		} else {
			r["errors"] = err.Error()
		}
	}
}

func (a *aggregator) isKubernetesBootstrapped(p *supervisor.Process) bool {
	submap, sok := a.kubernetesResources[""]
	if !sok {
		return false
	}
	for _, k := range a.requiredKinds {
		_, ok := submap[k]
		if !ok {
			return false
		}
	}
	return true
}

// Returns true if the current state of the world is complete. The
// kubernetes state of the world is always complete by definition
// because the kubernetes client provides that guarantee. The
// aggregate state of the world is complete when any consul services
// referenced by kubernetes have populated endpoint information (even
// if the value of the populated info is an empty set of endpoints).
func (a *aggregator) isComplete(p *supervisor.Process, watchset WatchSet) bool {
	complete := true

	for _, w := range watchset.KubernetesWatches {
		if _, ok := a.ids[w.WatchId()]; ok {
			p.Debugf("initialized k8s watch: %s", w.WatchId())
		} else {
			complete = false
			p.Debugf("waiting for k8s watch: %s", w.WatchId())
		}
	}

	for _, w := range watchset.ConsulWatches {
		if _, ok := a.ids[w.WatchId()]; ok {
			p.Debugf("initialized consul watch: %s", w.WatchId())
		} else {
			complete = false
			p.Debugf("waiting for consul watch: %s", w.WatchId())
		}
	}

	return complete
}

func (a *aggregator) maybeNotify(p *supervisor.Process) {
	now := time.Now()
	delay := a.limiter.Limit(now)
	if delay == 0 {
		a.notify(p)
	} else if delay > 0 {
		time.AfterFunc(delay, func() {
			a.notify(p)
		})
	}
}

func (a *aggregator) notify(p *supervisor.Process) {
	a.notifyMux.Lock()
	defer a.notifyMux.Unlock()

	if !a.isKubernetesBootstrapped(p) {
		return
	}

	watchset := a.getWatches(p)

	p.Debugf("found %d kubernetes watches", len(watchset.KubernetesWatches))
	p.Debugf("found %d consul watches", len(watchset.ConsulWatches))
	a.k8sWatches <- watchset.KubernetesWatches
	a.consulWatches <- watchset.ConsulWatches

	if !a.bootstrapped && a.isComplete(p, watchset) {
		p.Logf("bootstrapped!")
		a.bootstrapped = true
	}

	if a.bootstrapped {
		snapshot, err := a.generateSnapshot(p)
		if err != nil {
			p.Logf("generate snapshot failed %v", err)
			return
		}

		a.snapshots <- snapshot
	}
}

func (a *aggregator) getWatches(p *supervisor.Process) WatchSet {
	snapshot, err := a.generateSnapshot(p)
	if err != nil {
		p.Logf("generate snapshot failed %v", err)
		return WatchSet{}
	}
	result := a.watchHook(p, snapshot)
	return result.interpolate()
}

func ExecWatchHook(watchHooks []string) WatchHook {
	return func(p *supervisor.Process, snapshot string) WatchSet {
		result := WatchSet{}

		for _, hook := range watchHooks {
			ws := invokeHook(p, hook, snapshot)
			result.KubernetesWatches = append(result.KubernetesWatches, ws.KubernetesWatches...)
			result.ConsulWatches = append(result.ConsulWatches, ws.ConsulWatches...)
		}

		return result
	}
}

func lines(st string) []string {
	return strings.Split(st, "\n")
}

func invokeHook(p *supervisor.Process, hook, snapshot string) WatchSet {
	cmd := exec.Command("sh", "-c", hook)
	cmd.Stdin = strings.NewReader(snapshot)
	var watches, errors strings.Builder
	cmd.Stdout = &watches
	cmd.Stderr = &errors
	err := cmd.Run()
	stderr := errors.String()
	if stderr != "" {
		for _, line := range lines(stderr) {
			p.Logf("watch hook stderr: %s", line)
		}
	}
	if err != nil {
		p.Logf("watch hook failed: %v", err)
		return WatchSet{}
	}

	encoded := watches.String()

	decoder := json.NewDecoder(strings.NewReader(encoded))
	decoder.DisallowUnknownFields()
	result := WatchSet{}
	err = decoder.Decode(&result)
	if err != nil {
		for _, line := range lines(encoded) {
			p.Debugf("watch hook: %s", line)
		}
		p.Logf("watchset decode failed: %v", err)
		return WatchSet{}
	}

	return result
}
<commit_msg>(from AES) make watt pay attention to AMBASSADOR_FAST_VALIDATION<commit_after>package watt

import (
	"encoding/json"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/datawire/ambassador/pkg/consulwatch"
	"github.com/datawire/ambassador/pkg/k8s"
	"github.com/datawire/ambassador/pkg/kates"
	"github.com/datawire/ambassador/pkg/limiter"
	"github.com/datawire/ambassador/pkg/supervisor"
	"github.com/datawire/ambassador/pkg/watt"
)

type WatchHook func(p *supervisor.Process, snapshot string) WatchSet

type aggregator struct {
	// Input channel used to tell us about kubernetes state.
	KubernetesEvents chan k8sEvent
	// Input channel used to tell us about consul endpoints.
	ConsulEvents chan consulEvent
	// Output channel used to communicate with the k8s watch manager.
	k8sWatches chan<- []KubernetesWatchSpec
	// Output channel used to communicate with the consul watch manager.
	consulWatches chan<- []ConsulWatchSpec
	// Output channel used to communicate with the invoker.
	snapshots chan<- string
	// We won't consider ourselves "bootstrapped" until we hear
	// about all these kinds.
	requiredKinds       []string
	watchHook           WatchHook
	limiter             limiter.Limiter
	ids                 map[string]bool
	kubernetesResources map[string]map[string][]k8s.Resource
	consulEndpoints     map[string]consulwatch.Endpoints
	bootstrapped        bool
	notifyMux           sync.Mutex
	errors              map[string][]watt.Error
	validator           *kates.Validator
}

func NewAggregator(snapshots chan<- string, k8sWatches chan<- []KubernetesWatchSpec, consulWatches chan<- []ConsulWatchSpec,
	requiredKinds []string, watchHook WatchHook, limiter limiter.Limiter, validator *kates.Validator) *aggregator {
	return &aggregator{
		KubernetesEvents:    make(chan k8sEvent),
		ConsulEvents:        make(chan consulEvent),
		k8sWatches:          k8sWatches,
		consulWatches:       consulWatches,
		snapshots:           snapshots,
		requiredKinds:       requiredKinds,
		watchHook:           watchHook,
		limiter:             limiter,
		ids:                 make(map[string]bool),
		kubernetesResources: make(map[string]map[string][]k8s.Resource),
		consulEndpoints:     make(map[string]consulwatch.Endpoints),
		errors:              make(map[string][]watt.Error),
		validator:           validator,
	}
}

func (a *aggregator) Work(p *supervisor.Process) error {
	// In order to invoke `maybeNotify`, which is a very time consuming
	// operation, we coalesce events:
	//
	// 1. Be continuously reading all available events from
	//    a.KubernetesEvents and a.ConsulEvents and store k8sEvents
	//    in the potentialKubernetesEventSignal variable. This means
	//    at any given point (modulo caveats below), the
	//    potentialKubernetesEventSignal variable will have the
	//    latest Kubernetes event available.
	//
	// 2. At the same time, whenever there is capacity to write
	//    down the kubernetesEventProcessor channel, we send
	//    potentialKubernetesEventSignal to be processed.
	//
	//    The anonymous goroutine below will be constantly reading
	//    from the kubernetesEventProcessor channel and performing
	//    a blocking a.maybeNotify(). This means that we can only
	//    *write* to the kubernetesEventProcessor channel when we are
	//    not currently processing an event, but when that happens, we
	//    will still read from a.KubernetesEvents and a.ConsulEvents
	//    and update potentialKubernetesEventSignal.
	//
	// There are three caveats to the above:
	//
	// 1. At startup, we don't yet have a event to write, but
	//    we're not processing anything, so we will try to write
	//    something down the kubernetesEventProcessor channel.
	//    To cope with this, the invoking goroutine will ignore events
	//    signals that have a event.skip flag.
	//
	// 2. If we process an event quickly, or if there aren't new
	//    events available, then we end up busy looping and
	//    sending the same potentialKubernetesEventSignal value down
	//    the kubernetesEventProcessor channel multiple times. To cope
	//    with this, whenever we have successfully written to the
	//    kubernetesEventProcessor channel, we do a *blocking* read of
	//    the next event from a.KubernetesEvents and a.ConsulEvents.
	//
	// 3. Always be calling a.setKubernetesResources as soon as we
	//    receive an event. This is a fast non-blocking call that
	//    update watches, we can't coalesce this call.

	p.Ready()

	type eventSignal struct {
		kubernetesEvent k8sEvent
		skip            bool
	}

	kubernetesEventProcessor := make(chan eventSignal)
	go func() {
		for event := range kubernetesEventProcessor {
			if event.skip {
				// ignore the initial eventSignal to deal with the
				// corner case where we haven't yet received an event yet.
				continue
			}
			a.maybeNotify(p)
		}
	}()

	potentialKubernetesEventSignal := eventSignal{kubernetesEvent: k8sEvent{}, skip: true}
	for {
		select {
		case potentialKubernetesEvent := <-a.KubernetesEvents:
			// if a new KubernetesEvents is available to be read,
			// and we can't write to the kubernetesEventProcessor channel,
			// then we will overwrite potentialKubernetesEvent
			// with a newer event while still processing a.setKubernetesResources
			a.setKubernetesResources(potentialKubernetesEvent)
			potentialKubernetesEventSignal = eventSignal{kubernetesEvent: potentialKubernetesEvent, skip: false}
		case kubernetesEventProcessor <- potentialKubernetesEventSignal:
			// if we aren't currently blocked in
			// a.maybeNotify() then the above goroutine will be
			// reading from the kubernetesEventProcessor channel and we
			// will send the current potentialKubernetesEventSignal
			// value over the kubernetesEventProcessor channel to be
			// processed
			select {
			case potentialKubernetesEvent := <-a.KubernetesEvents:
				// here we do blocking read of the next event for caveat #2.
				a.setKubernetesResources(potentialKubernetesEvent)
				potentialKubernetesEventSignal = eventSignal{kubernetesEvent: potentialKubernetesEvent, skip: false}
			case event := <-a.ConsulEvents:
				a.updateConsulResources(event)
				a.maybeNotify(p)
			case <-p.Shutdown():
				return nil
			}
		case event := <-a.ConsulEvents:
			// we are always reading and processing ConsulEvents directly,
			// not coalescing them.
			a.updateConsulResources(event)
			a.maybeNotify(p)
		case <-p.Shutdown():
			return nil
		}
	}
}

func (a *aggregator) updateConsulResources(event consulEvent) {
	a.ids[event.WatchId] = true
	a.consulEndpoints[event.Endpoints.Service] = event.Endpoints
}

func (a *aggregator) setKubernetesResources(event k8sEvent) {
	if len(event.errors) > 0 {
		for _, kError := range event.errors {
			a.errors[kError.Source] = append(a.errors[kError.Source], kError)
		}
		return
	}
	a.ids[event.watchId] = true
	submap, ok := a.kubernetesResources[event.watchId]
	if !ok {
		submap = make(map[string][]k8s.Resource)
		a.kubernetesResources[event.watchId] = submap
	}
	submap[event.kind] = event.resources
}

func (a *aggregator) generateSnapshot(p *supervisor.Process) (string, error) {
	k8sResources := make(map[string][]k8s.Resource)
	for _, submap := range a.kubernetesResources {
		for k, v := range submap {
			a.validate(p, v)
			k8sResources[k] = append(k8sResources[k], v...)
		}
	}
	s := watt.Snapshot{
		Consul:     watt.ConsulSnapshot{Endpoints: a.consulEndpoints},
		Kubernetes: k8sResources,
		Errors:     a.errors,
	}

	jsonBytes, err := json.MarshalIndent(s, "", "    ")
	if err != nil {
		return "{}", err
	}

	return string(jsonBytes), nil
}

var fastValidation = len(os.Getenv("AMBASSADOR_FAST_VALIDATION")) > 0

func (a *aggregator) validate(p *supervisor.Process, resources []k8s.Resource) {
	if !fastValidation {
		return
	}

	for _, r := range resources {
		err := a.validator.Validate(p.Context(), map[string]interface{}(r))
		if err == nil {
			delete(r, "errors")
		} else {
			r["errors"] = err.Error()
		}
	}
}

func (a *aggregator) isKubernetesBootstrapped(p *supervisor.Process) bool {
	submap, sok := a.kubernetesResources[""]
	if !sok {
		return false
	}
	for _, k := range a.requiredKinds {
		_, ok := submap[k]
		if !ok {
			return false
		}
	}
	return true
}

// Returns true if the current state of the world is complete. The
// kubernetes state of the world is always complete by definition
// because the kubernetes client provides that guarantee. The
// aggregate state of the world is complete when any consul services
// referenced by kubernetes have populated endpoint information (even
// if the value of the populated info is an empty set of endpoints).
func (a *aggregator) isComplete(p *supervisor.Process, watchset WatchSet) bool {
	complete := true

	for _, w := range watchset.KubernetesWatches {
		if _, ok := a.ids[w.WatchId()]; ok {
			p.Debugf("initialized k8s watch: %s", w.WatchId())
		} else {
			complete = false
			p.Debugf("waiting for k8s watch: %s", w.WatchId())
		}
	}

	for _, w := range watchset.ConsulWatches {
		if _, ok := a.ids[w.WatchId()]; ok {
			p.Debugf("initialized consul watch: %s", w.WatchId())
		} else {
			complete = false
			p.Debugf("waiting for consul watch: %s", w.WatchId())
		}
	}

	return complete
}

func (a *aggregator) maybeNotify(p *supervisor.Process) {
	now := time.Now()
	delay := a.limiter.Limit(now)
	if delay == 0 {
		a.notify(p)
	} else if delay > 0 {
		time.AfterFunc(delay, func() {
			a.notify(p)
		})
	}
}

func (a *aggregator) notify(p *supervisor.Process) {
	a.notifyMux.Lock()
	defer a.notifyMux.Unlock()

	if !a.isKubernetesBootstrapped(p) {
		return
	}

	watchset := a.getWatches(p)

	p.Debugf("found %d kubernetes watches", len(watchset.KubernetesWatches))
	p.Debugf("found %d consul watches", len(watchset.ConsulWatches))
	a.k8sWatches <- watchset.KubernetesWatches
	a.consulWatches <- watchset.ConsulWatches

	if !a.bootstrapped && a.isComplete(p, watchset) {
		p.Logf("bootstrapped!")
		a.bootstrapped = true
	}

	if a.bootstrapped {
		snapshot, err := a.generateSnapshot(p)
		if err != nil {
			p.Logf("generate snapshot failed %v", err)
			return
		}

		a.snapshots <- snapshot
	}
}

func (a *aggregator) getWatches(p *supervisor.Process) WatchSet {
	snapshot, err := a.generateSnapshot(p)
	if err != nil {
		p.Logf("generate snapshot failed %v", err)
		return WatchSet{}
	}
	result := a.watchHook(p, snapshot)
	return result.interpolate()
}

func ExecWatchHook(watchHooks []string) WatchHook {
	return func(p *supervisor.Process, snapshot string) WatchSet {
		result := WatchSet{}

		for _, hook := range watchHooks {
			ws := invokeHook(p, hook, snapshot)
			result.KubernetesWatches = append(result.KubernetesWatches, ws.KubernetesWatches...)
			result.ConsulWatches = append(result.ConsulWatches, ws.ConsulWatches...)
		}

		return result
	}
}

func lines(st string) []string {
	return strings.Split(st, "\n")
}

func invokeHook(p *supervisor.Process, hook, snapshot string) WatchSet {
	cmd := exec.Command("sh", "-c", hook)
	cmd.Stdin = strings.NewReader(snapshot)
	var watches, errors strings.Builder
	cmd.Stdout = &watches
	cmd.Stderr = &errors
	err := cmd.Run()
	stderr := errors.String()
	if stderr != "" {
		for _, line := range lines(stderr) {
			p.Logf("watch hook stderr: %s", line)
		}
	}
	if err != nil {
		p.Logf("watch hook failed: %v", err)
		return WatchSet{}
	}

	encoded := watches.String()

	decoder := json.NewDecoder(strings.NewReader(encoded))
	decoder.DisallowUnknownFields()
	result := WatchSet{}
	err = decoder.Decode(&result)
	if err != nil {
		for _, line := range lines(encoded) {
			p.Debugf("watch hook: %s", line)
		}
		p.Logf("watchset decode failed: %v", err)
		return WatchSet{}
	}

	return result
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
)

const version = "0.1.0"

func main() {
	up := flag.Bool("update", true, "update command")
	flag.Parse()

	fmt.Printf("version %s\n", version)

	var code int
	if *up {
		code = update()
	}

	os.Exit(code)
}

const srcPath = "github.com/dvrkps/dojo/cmdupdate"

func update() int {
	cmd := exec.Command("go", "get", "-u", srcPath)
	out, err := cmd.CombinedOutput()
	var code int
	if err != nil {
		fmt.Println("ERROR")
		code = 1
	}
	fmt.Println(string(out))
	return code
}
<commit_msg>cmdupdate: clean<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
)

const version = "0.1.1"

func main() {
	up := flag.Bool("update", false, "update command")
	flag.Parse()

	fmt.Printf("version %s\n", version)

	var code int
	if *up {
		code = update()
	}

	os.Exit(code)
}

const srcPath = "github.com/dvrkps/dojo/cmdupdate"

func update() int {
	cmd := exec.Command("go", "get", "-u", srcPath)
	out, err := cmd.CombinedOutput()
	var code int
	if err != nil {
		fmt.Println("ERROR")
		code = 1
	}
	fmt.Println(string(out))
	if code < 1 {
		fmt.Println("Update done.")
	} else {
		fmt.Println("Update fail.")
	}
	return code
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcsfake

import (
	"errors"
	"io"
	"sort"
	"sync"

	"github.com/jacobsa/gcloud/gcs"
	"golang.org/x/net/context"
	"google.golang.org/cloud/storage"
)

// Create an in-memory bucket with the given name and empty contents.
func NewFakeBucket(name string) gcs.Bucket {
	return &bucket{name: name}
}

type object struct {
	// A storage.Object representing metadata for this object. Never changes.
	metadata *storage.Object

	// The contents of the object. These never change.
	contents []byte
}

// A slice of objects compared by name.
type objectSlice []object

func (s objectSlice) Len() int           { return len(s) }
func (s objectSlice) Less(i, j int) bool { return s[i].metadata.Name < s[j].metadata.Name }
func (s objectSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

type bucket struct {
	name string
	mu   sync.RWMutex

	// The set of extant objects.
	//
	// INVARIANT: Strictly increasing.
	objects objectSlice // GUARDED_BY(mu)
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(
	ctx context.Context,
	query *storage.Query) (*storage.Objects, error) {
	return nil, errors.New("TODO: Implement ListObjects.")
}

func (b *bucket) NewReader(
	ctx context.Context,
	objectName string) (io.ReadCloser, error) {
	return nil, errors.New("TODO: Implement NewReader.")
}

func (b *bucket) NewWriter(
	ctx context.Context,
	attrs *storage.ObjectAttrs) (gcs.ObjectWriter, error) {
	return newObjectWriter(b, attrs), nil
}

func (b *bucket) DeleteObject(
	ctx context.Context,
	name string) error {
	return errors.New("TODO: Implement DeleteObject.")
}

// Create an object struct for the given attributes and contents.
//
// EXCLUSIVE_LOCKS_REQUIRED(b.mu)
func (b *bucket) mintObject(
	attrs *storage.ObjectAttrs,
	contents []byte) object

func (b *bucket) addObject(
	attrs *storage.ObjectAttrs,
	contents []byte) *storage.Object {
	b.mu.Lock()
	defer b.mu.Unlock()

	// Create an object record from the given attributes.
	var o object = b.mintObject(attrs, contents)

	// Add it to our list of object.
	b.objects = append(b.objects, o)
	sort.Sort(b.objects)

	return o.metadata
}
<commit_msg>Implemented some of mintObject.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcsfake

import (
	"errors"
	"io"
	"sort"
	"sync"

	"github.com/jacobsa/gcloud/gcs"
	"golang.org/x/net/context"
	"google.golang.org/cloud/storage"
)

// Create an in-memory bucket with the given name and empty contents.
func NewFakeBucket(name string) gcs.Bucket {
	return &bucket{name: name}
}

type object struct {
	// A storage.Object representing metadata for this object. Never changes.
	metadata *storage.Object

	// The contents of the object. These never change.
	contents []byte
}

// A slice of objects compared by name.
type objectSlice []object

func (s objectSlice) Len() int           { return len(s) }
func (s objectSlice) Less(i, j int) bool { return s[i].metadata.Name < s[j].metadata.Name }
func (s objectSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

type bucket struct {
	name string
	mu   sync.RWMutex

	// The set of extant objects.
	//
	// INVARIANT: Strictly increasing.
	objects objectSlice // GUARDED_BY(mu)
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(
	ctx context.Context,
	query *storage.Query) (*storage.Objects, error) {
	return nil, errors.New("TODO: Implement ListObjects.")
}

func (b *bucket) NewReader(
	ctx context.Context,
	objectName string) (io.ReadCloser, error) {
	return nil, errors.New("TODO: Implement NewReader.")
}

func (b *bucket) NewWriter(
	ctx context.Context,
	attrs *storage.ObjectAttrs) (gcs.ObjectWriter, error) {
	return newObjectWriter(b, attrs), nil
}

func (b *bucket) DeleteObject(
	ctx context.Context,
	name string) error {
	return errors.New("TODO: Implement DeleteObject.")
}

// Create an object struct for the given attributes and contents.
//
// EXCLUSIVE_LOCKS_REQUIRED(b.mu)
func (b *bucket) mintObject(
	attrs *storage.ObjectAttrs,
	contents []byte) (o object) {
	// Set up metadata.
	// TODO(jacobsa): Other fields.
	o.metadata = &storage.Object{
		Bucket: b.Name(),
		Name:   attrs.Name,
	}

	// Set up contents.
	o.contents = contents

	return
}

func (b *bucket) addObject(
	attrs *storage.ObjectAttrs,
	contents []byte) *storage.Object {
	b.mu.Lock()
	defer b.mu.Unlock()

	// Create an object record from the given attributes.
	var o object = b.mintObject(attrs, contents)

	// Add it to our list of object.
	b.objects = append(b.objects, o)
	sort.Sort(b.objects)

	return o.metadata
}
<|endoftext|>


<commit_before>package sphinx

import (
	"testing"
)

func TestNewConfiguration(t *testing.T) {

	// test loading example config
	config, err := NewConfiguration("./example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}

	if config.Forward.Scheme != "http" {
		t.Error("expected http for Forward.Scheme")
	}

	if len(config.Limits) != 4 {
		t.Error("expected 4 bucket definitions")
	}

	for _, limit := range config.Limits {
		if limit.Interval < 1 {
			t.Error("limit interval should be greator than 1")
		}
		if limit.Max < 1 {
			t.Error("limit max should be greator than 1")
		}
		if limit.Keys != nil {
			t.Error("limit was expected to have atleast 1 key")
		}

		if limit.Matches["headers"] == nil && limit.Matches["paths"] == nil {
			t.Error("One of paths or headers was expected to be set for matches")
		}
	}

	// test incorrect config
	invalid_config := []byte(`
forward:
  host: proxy.example.com
`)
	config, err = loadAndValidateConfig(invalid_config)
	if err == nil {
		t.Error("invalid config did not return error")
	}

	invalid_config = []byte(`
forward:
  scheme: http
  host: proxy.example.com

buckets:
  bearer/events:
    keys:
      - 'header:authentication'
`)
	config, err = loadAndValidateConfig(invalid_config)
	if err == nil {
		t.Error("invalid config did not return error")
	}
}
<commit_msg>split up tests for configuration<commit_after>package sphinx

import (
	"testing"
)

// test example config file is loaded correctly
func TestConfigurationFileLoading(t *testing.T) {

	config, err := NewConfiguration("./example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}

	if config.Forward.Scheme != "http" {
		t.Error("expected http for Forward.Scheme")
	}

	if len(config.Limits) != 4 {
		t.Error("expected 4 bucket definitions")
	}

	for _, limit := range config.Limits {
		if limit.Interval < 1 {
			t.Error("limit interval should be greator than 1")
		}
		if limit.Max < 1 {
			t.Error("limit max should be greator than 1")
		}
		if limit.Keys != nil {
			t.Error("limit was expected to have atleast 1 key")
		}

		if limit.Matches["headers"] == nil && limit.Matches["paths"] == nil {
			t.Error("One of paths or headers was expected to be set for matches")
		}
	}
}

// Incorrect configuration file should return errors
func TestConfigurationFileFailures(t *testing.T) {

	invalid_config := []byte(`
forward:
  host: proxy.example.com
`)
	config, err = loadAndValidateConfig(invalid_config)
	if err == nil {
		t.Error("invalid config did not return error")
	}

	invalid_config = []byte(`
forward:
  scheme: http
  host: proxy.example.com

buckets:
  bearer/events:
    keys:
      - 'header:authentication'
`)
	config, err = loadAndValidateConfig(invalid_config)
	if err == nil {
		t.Error("invalid config did not return error")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os/exec"

	"github.com/gorilla/mux"
)

type orchestration struct {
	Scripts []script `json:"scripts"`
}

type script struct {
	Command string   `json:"command"`
	Args    []string `json:"args"`
}

type response struct {
	Results []result `json:"results"`
}

type result struct {
	Stdout string `json:"stdout"`
	Stderr string `json:"stderr"`
}

func hookHandler(w http.ResponseWriter, r *http.Request) {
	params := mux.Vars(r)
	key := params["key"]
	log.Printf("Received Hook for key '%s'\n", key)
	results := processHandler(key)

	data, err := json.MarshalIndent(&response{results}, "", "  ")
	if err != nil {
		log.Println(err.Error())
	}
	w.Write(data)
}

func processHandler(key string) []result {
	results := make([]result, 0)
	script := getScriptFromKey(key)
	for _, x := range script.Scripts {
		r, err := execScript(x)
		if err != nil {
			log.Println("ERROR :" + err.Error())
		}
		results = append(results, r)
	}
	return results
}

func execScript(s script) (result, error) {
	cmd := exec.Command(s.Command, s.Args...)
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	r := result{
		stdout.String(),
		stderr.String(),
	}
	return r, err
}

func getScriptFromKey(key string) orchestration {
	p := fmt.Sprintf("%s/%s.json", configdir, key)
	b, err := ioutil.ReadFile(p)
	if err != nil {
		log.Printf("Error opening %s\n", p)
	}
	var o orchestration
	err = json.Unmarshal(b, &o)
	if err != nil {
		log.Println(err)
	}
	return o
}
<commit_msg>Respect the echo flag<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os/exec"

	"github.com/gorilla/mux"
)

type orchestration struct {
	Scripts []script `json:"scripts"`
}

type script struct {
	Command string   `json:"command"`
	Args    []string `json:"args"`
}

type response struct {
	Results []result `json:"results"`
}

type result struct {
	Stdout string `json:"stdout"`
	Stderr string `json:"stderr"`
}

func hookHandler(w http.ResponseWriter, r *http.Request) {
	params := mux.Vars(r)
	key := params["key"]
	log.Printf("Received Hook for key '%s'\n", key)
	results := processHandler(key)

	if echo {
		data, err := json.MarshalIndent(&response{results}, "", "  ")
		if err != nil {
			log.Println(err.Error())
		}
		w.Write(data)
	}
}

func processHandler(key string) []result {
	results := make([]result, 0)
	script := getScriptFromKey(key)
	for _, x := range script.Scripts {
		r, err := execScript(x)
		if err != nil {
			log.Println("ERROR :" + err.Error())
		}
		results = append(results, r)
	}
	return results
}

func execScript(s script) (result, error) {
	cmd := exec.Command(s.Command, s.Args...)
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	r := result{
		stdout.String(),
		stderr.String(),
	}
	return r, err
}

func getScriptFromKey(key string) orchestration {
	p := fmt.Sprintf("%s/%s.json", configdir, key)
	b, err := ioutil.ReadFile(p)
	if err != nil {
		log.Printf("Error opening %s\n", p)
	}
	var o orchestration
	err = json.Unmarshal(b, &o)
	if err != nil {
		log.Println(err)
	}
	return o
}
<|endoftext|>