<commit_before>package xd

import (
	"bufio"
	"io"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"strings"
	"time"
	"xd/lib/bittorrent/swarm"
	"xd/lib/config"
	"xd/lib/log"
	"xd/lib/rpc"
	"xd/lib/sync"
	t "xd/lib/translate"
	"xd/lib/util"
	"xd/lib/version"
)

type httpRPC struct {
	w http.ResponseWriter
	r *http.Request
}

func printHelp(cmd string) {
	log.Infof("usage: %s [config.ini] | --genconf config.ini\n", cmd)
}

func NewContext() *Context {
	pr, pw := io.Pipe()
	return &Context{
		pr:      pr,
		pw:      pw,
		sigchnl: make(chan os.Signal),
	}
}

type Context struct {
	pr         io.ReadCloser
	pw         io.WriteCloser
	closers    sync.Map
	numClosers int
	quit       bool
	swarms     []*swarm.Swarm
	sigchnl    chan os.Signal
	netlost    bool
}

func (c *Context) Run() {
	go func() {
		io.Copy(c.pw, os.Stdin)
	}()
	r := bufio.NewReader(c.pr)
	for {
		line, err := r.ReadString(10)
		if err == nil {
			if strings.ToLower(line) == "f\n" {
				if c.netlost {
					log.Debug("respec paid")
				}
			} else if line == "\n" && c.quit {
				break
			} else if line == "" {
				break
			}
		} else {
			break
		}
	}
	c.pr.Close()
}

func (c *Context) Running() bool {
	return !c.quit
}

func (c *Context) RunSignals() {
	signal.Notify(c.sigchnl, os.Interrupt)
	for {
		sig := <-c.sigchnl
		if sig == os.Interrupt {
			log.Info("Interrupted")
			c.Close()
			return
		} else {
			log.Warnf("got wierd signal wtf: %s", sig)
			continue
		}
	}
}

func (c *Context) AddCloser(cl io.Closer) int {
	c.numClosers++
	c.closers.Store(c.numClosers, cl)
	return c.numClosers
}

func (c *Context) ReplaceCloser(id int, cl io.Closer) {
	c.closers.Store(id, cl)
}

func (c *Context) AddSwarm(sw *swarm.Swarm) {
	c.swarms = append(c.swarms, sw)
}

func (c *Context) Close() error {
	c.quit = true
	c.pw.Close()
	// close swarms first
	for _, sw := range c.swarms {
		sw.Close()
	}
	c.closers.Range(func(k, v interface{}) bool {
		cl := v.(io.Closer)
		cl.Close()
		return true
	})
	return nil
}

// Run runs XD main function
func Run() {

	ctx := NewContext()

	v := version.Version()
	conf := new(config.Config)
	fname := "torrents.ini"
	if len(os.Args) > 1 {
		fname = os.Args[1]
	}
	if fname == "-h" || fname == "--help" {
		printHelp(os.Args[0])
		return
	}
	var err error
	if fname == "--genconf" {
		if len(os.Args) == 3 {
			conf.Load("")
			err = conf.Save(os.Args[2])
			if err != nil {
				log.Errorf("failed to save config: %s", err)
			}
		} else {
			printHelp(os.Args[0])
		}
		return
	}

	log.Info(t.T("starting %s", v))
	if !util.CheckFile(fname) {
		conf.Load(fname)
		err = conf.Save(fname)
		if err != nil {
			log.Errorf("failed to save initial config: %s", err)
			return
		}
		log.Info(t.T("auto-generated new config at %s", fname))
	}
	err = conf.Load(fname)
	if err != nil {
		log.Errorf("failed to config %s", err)
		return
	}
	log.Info(t.T("loaded config %s", fname))
	log.SetLevel(conf.Log.Level)

	if conf.Log.Pprof {
		go func() {
			pprofaddr := "127.0.0.1:6060"
			l, err := net.Listen("tcp", pprofaddr)
			if err == nil {
				ctx.AddCloser(l)
				log.Infof("spawning pprof at %s", pprofaddr)
				log.Warnf("pprof exited: %s", http.Serve(l, nil))
			}
		}()
	}

	st := conf.Storage.CreateStorage()
	err = st.Init()
	if err != nil {
		log.Errorf("error initializing storage: %s", err)
		return
	}
	// start io thread
	go st.Run()
	count := 0
	for count < conf.Bittorrent.Swarms {
		gnutella := conf.Gnutella.CreateSwarm()
		sw := conf.Bittorrent.CreateSwarm(st, gnutella)
		if gnutella != nil {
			ctx.AddCloser(gnutella)
		}
		ctx.AddSwarm(sw)
		count++
	}

	ts, err := st.OpenAllTorrents()
	if err != nil {
		log.Errorf("error opening all torrents: %s", err)
		return
	}
	for _, t := range ts {
		for _, sw := range ctx.swarms {
			err = sw.AddTorrent(t)
			if err != nil {
				log.Errorf("error adding torrent: %s", err)
			}
		}
	}

	// torrent auto adder
	go func() {
		for ctx.Running() {
			nt := st.PollNewTorrents()
			for _, t := range nt {
				e := t.VerifyAll()
				if e != nil {
					log.Errorf("failed to add %s: %s", t.Name(), e.Error())
					continue
				}
				for _, sw := range ctx.swarms {
					sw.AddTorrent(t)
				}
			}
			time.Sleep(time.Second)
		}
	}()

	// start rpc server
	if conf.RPC.Enabled {
		log.Infof("RPC enabled")
		var host string
		var l net.Listener
		var e error
		var cleanSock func()
		if strings.HasPrefix(conf.RPC.Bind, "unix:") {
			sock := conf.RPC.Bind[5:]
			cleanSock = func() {
				os.Remove(sock)
			}
			l, e = net.Listen("unix", sock)
			if e == nil {
				e = os.Chmod(sock, 0640)
			}
		} else {
			l, e = net.Listen("tcp", conf.RPC.Bind)
			cleanSock = func() {
			}
			host = conf.RPC.ExpectedHost
		}
		if e == nil {
			ctx.AddCloser(l)
			s := &http.Server{
				Handler: rpc.NewServer(ctx.swarms, host),
			}
			go func(serv *http.Server) {
				log.Errorf("rpc died: %s", serv.Serve(l))
				cleanSock()
			}(s)
		} else {
			log.Errorf("failed to bind rpc: %s", e)
		}
	}

	runFunc := func(netConf config.I2PConfig, sw *swarm.Swarm) {
		n := netConf.CreateSession()
		id := ctx.AddCloser(n)
		for sw.Running() {
			log.Info("opening i2p session")
			err := n.Open()
			if err == nil {
				log.Infof("i2p session made, we are %s", n.B32Addr())
				sw.ObtainedNetwork(n)
				ctx.netlost = false
				err = sw.Run()
				if err != nil {
					ctx.netlost = true
					log.Errorf("lost i2p session: %s", err)
					sw.LostNetwork()
					n = netConf.CreateSession()
					ctx.ReplaceCloser(id, n)
				}
			} else {
				ctx.netlost = true
				n = netConf.CreateSession()
				ctx.ReplaceCloser(id, n)
				log.Errorf("failed to create i2p session: %s", err)
				time.Sleep(time.Second)
			}
		}
	}

	for idx := range ctx.swarms {
		go runFunc(conf.I2P, ctx.swarms[idx])
	}
	ctx.AddCloser(st)
	go ctx.RunSignals()
	ctx.Run()
}
<commit_msg>start out as net lost<commit_after>package xd

import (
	"bufio"
	"io"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"strings"
	"time"
	"xd/lib/bittorrent/swarm"
	"xd/lib/config"
	"xd/lib/log"
	"xd/lib/rpc"
	"xd/lib/sync"
	t "xd/lib/translate"
	"xd/lib/util"
	"xd/lib/version"
)

type httpRPC struct {
	w http.ResponseWriter
	r *http.Request
}

func printHelp(cmd string) {
	log.Infof("usage: %s [config.ini] | --genconf config.ini\n", cmd)
}

func NewContext() *Context {
	pr, pw := io.Pipe()
	return &Context{
		pr:      pr,
		pw:      pw,
		sigchnl: make(chan os.Signal),
		netlost: true,
	}
}

type Context struct {
	pr         io.ReadCloser
	pw         io.WriteCloser
	closers    sync.Map
	numClosers int
	quit       bool
	swarms     []*swarm.Swarm
	sigchnl    chan os.Signal
	netlost    bool
}

func (c *Context) Run() {
	go func() {
		io.Copy(c.pw, os.Stdin)
	}()
	r := bufio.NewReader(c.pr)
	for {
		line, err := r.ReadString(10)
		if err == nil {
			if strings.ToLower(line) == "f\n" {
				if c.netlost {
					log.Debug("respec paid")
				}
			} else if line == "\n" && c.quit {
				break
			} else if line == "" {
				break
			}
		} else {
			break
		}
	}
	c.pr.Close()
}

func (c *Context) Running() bool {
	return !c.quit
}

func (c *Context) RunSignals() {
	signal.Notify(c.sigchnl, os.Interrupt)
	for {
		sig := <-c.sigchnl
		if sig == os.Interrupt {
			log.Info("Interrupted")
			c.Close()
			return
		} else {
			log.Warnf("got wierd signal wtf: %s", sig)
			continue
		}
	}
}

func (c *Context) AddCloser(cl io.Closer) int {
	c.numClosers++
	c.closers.Store(c.numClosers, cl)
	return c.numClosers
}

func (c *Context) ReplaceCloser(id int, cl io.Closer) {
	c.closers.Store(id, cl)
}

func (c *Context) AddSwarm(sw *swarm.Swarm) {
	c.swarms = append(c.swarms, sw)
}

func (c *Context) Close() error {
	c.quit = true
	c.pw.Close()
	// close swarms first
	for _, sw := range c.swarms {
		sw.Close()
	}
	c.closers.Range(func(k, v interface{}) bool {
		cl := v.(io.Closer)
		cl.Close()
		return true
	})
	return nil
}

// Run runs XD main function
func Run() {

	ctx := NewContext()

	v := version.Version()
	conf := new(config.Config)
	fname := "torrents.ini"
	if len(os.Args) > 1 {
		fname = os.Args[1]
	}
	if fname == "-h" || fname == "--help" {
		printHelp(os.Args[0])
		return
	}
	var err error
	if fname == "--genconf" {
		if len(os.Args) == 3 {
			conf.Load("")
			err = conf.Save(os.Args[2])
			if err != nil {
				log.Errorf("failed to save config: %s", err)
			}
		} else {
			printHelp(os.Args[0])
		}
		return
	}

	log.Info(t.T("starting %s", v))
	if !util.CheckFile(fname) {
		conf.Load(fname)
		err = conf.Save(fname)
		if err != nil {
			log.Errorf("failed to save initial config: %s", err)
			return
		}
		log.Info(t.T("auto-generated new config at %s", fname))
	}
	err = conf.Load(fname)
	if err != nil {
		log.Errorf("failed to config %s", err)
		return
	}
	log.Info(t.T("loaded config %s", fname))
	log.SetLevel(conf.Log.Level)

	if conf.Log.Pprof {
		go func() {
			pprofaddr := "127.0.0.1:6060"
			l, err := net.Listen("tcp", pprofaddr)
			if err == nil {
				ctx.AddCloser(l)
				log.Infof("spawning pprof at %s", pprofaddr)
				log.Warnf("pprof exited: %s", http.Serve(l, nil))
			}
		}()
	}

	st := conf.Storage.CreateStorage()
	err = st.Init()
	if err != nil {
		log.Errorf("error initializing storage: %s", err)
		return
	}
	// start io thread
	go st.Run()
	count := 0
	for count < conf.Bittorrent.Swarms {
		gnutella := conf.Gnutella.CreateSwarm()
		sw := conf.Bittorrent.CreateSwarm(st, gnutella)
		if gnutella != nil {
			ctx.AddCloser(gnutella)
		}
		ctx.AddSwarm(sw)
		count++
	}

	ts, err := st.OpenAllTorrents()
	if err != nil {
		log.Errorf("error opening all torrents: %s", err)
		return
	}
	for _, t := range ts {
		for _, sw := range ctx.swarms {
			err = sw.AddTorrent(t)
			if err != nil {
				log.Errorf("error adding torrent: %s", err)
			}
		}
	}

	// torrent auto adder
	go func() {
		for ctx.Running() {
			nt := st.PollNewTorrents()
			for _, t := range nt {
				e := t.VerifyAll()
				if e != nil {
					log.Errorf("failed to add %s: %s", t.Name(), e.Error())
					continue
				}
				for _, sw := range ctx.swarms {
					sw.AddTorrent(t)
				}
			}
			time.Sleep(time.Second)
		}
	}()

	// start rpc server
	if conf.RPC.Enabled {
		log.Infof("RPC enabled")
		var host string
		var l net.Listener
		var e error
		var cleanSock func()
		if strings.HasPrefix(conf.RPC.Bind, "unix:") {
			sock := conf.RPC.Bind[5:]
			cleanSock = func() {
				os.Remove(sock)
			}
			l, e = net.Listen("unix", sock)
			if e == nil {
				e = os.Chmod(sock, 0640)
			}
		} else {
			l, e = net.Listen("tcp", conf.RPC.Bind)
			cleanSock = func() {
			}
			host = conf.RPC.ExpectedHost
		}
		if e == nil {
			ctx.AddCloser(l)
			s := &http.Server{
				Handler: rpc.NewServer(ctx.swarms, host),
			}
			go func(serv *http.Server) {
				log.Errorf("rpc died: %s", serv.Serve(l))
				cleanSock()
			}(s)
		} else {
			log.Errorf("failed to bind rpc: %s", e)
		}
	}

	runFunc := func(netConf config.I2PConfig, sw *swarm.Swarm) {
		n := netConf.CreateSession()
		id := ctx.AddCloser(n)
		for sw.Running() {
			log.Info("opening i2p session")
			err := n.Open()
			if err == nil {
				log.Infof("i2p session made, we are %s", n.B32Addr())
				sw.ObtainedNetwork(n)
				ctx.netlost = false
				err = sw.Run()
				if err != nil {
					ctx.netlost = true
					log.Errorf("lost i2p session: %s", err)
					sw.LostNetwork()
					n = netConf.CreateSession()
					ctx.ReplaceCloser(id, n)
				}
			} else {
				ctx.netlost = true
				n = netConf.CreateSession()
				ctx.ReplaceCloser(id, n)
				log.Errorf("failed to create i2p session: %s", err)
				time.Sleep(time.Second)
			}
		}
	}

	for idx := range ctx.swarms {
		go runFunc(conf.I2P, ctx.swarms[idx])
	}
	ctx.AddCloser(st)
	go ctx.RunSignals()
	ctx.Run()
}
<|endoftext|>


<commit_before>package config

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"gopkg.in/yaml.v2"
)

// Define Application Settings Structure
type CassabonConfig struct {
	Logging struct {
		Logdir   string // Log Directory
		Loglevel string // Level to log at.
	}
	Statsd   StatsdSettings
	Channels struct {
		DataStoreChanLen  int // Length of the DataStore channel
		IndexStoreChanLen int // Length of the Indexer channel
		GopherChanLen     int // Length of the StatGopher channel
	}
	Carbon struct {
		Listen     string // ip:port on which to listen for Carbon stats
		Protocol   string // "tcp", "udp" or "both" are acceptable
		Parameters struct {
			TCPTimeout int
			UDPTimeout int
		}
		Peers []string // All servers in the Cassabon array, as "ip:port"
	}
	API struct {
		Listen          string // HTTP API listens on this address:port
		HealthCheckFile string // Location of healthcheck file.
	}
	Cassandra struct {
		Hosts []string // List of hostnames or IP addresses of Cassandra ring
		Port  string   // Cassandra port
	}
	Redis struct {
		Index RedisSettings // Settings for Redis Index
		Queue RedisSettings // Settings for Redis Queue
	}
	Rollups map[string]RollupSettings // Map of regex and rollups
}

// Definition of each rollup
type RollupSettings struct {
	Retention   []string
	Aggregation string
}

// Redis struct for redis connection information
type RedisSettings struct {
	Sentinel bool     // True if sentinel, false if standalone.
	Addr     []string // List of addresses in host:port format
	DB       int64    // Redis DB number for the index.
	Pwd      string   // Password for Redis.
	Master   string   // Master config name for sentinel settings.
}

type StatsdSettings struct {
	Host   string // Host or IP address of statsd server
	Port   string // Port that statsd server listens on
	Events struct {
		ReceiveOK struct {
			Key        string
			SampleRate float32
		}
		ReceiveFail struct {
			Key        string
			SampleRate float32
		}
	}
}

// rawCassabonConfig is the decoded YAML from the configuration file.
var rawCassabonConfig *CassabonConfig

// ReadConfigurationFile reads the contents of the specified file from disk, and unmarshals it.
// First time through we have no logger, so we can't log warnings or errors. Panic instead.
func ReadConfigurationFile(configFile string, haveLogger bool) error {

	// Read the configuration file.
	yamlConfig, err := ioutil.ReadFile(configFile)
	if err != nil {
		if haveLogger {
			return err
		} else {
			panic(err)
		}
	}

	// Unmarshal config file contents into raw config struct.
	err = yaml.Unmarshal(yamlConfig, &rawCassabonConfig)
	if err != nil {
		if haveLogger {
			return err
		} else {
			panic(err)
		}
	}
	return nil
}

// LoadStartupValues populates the global config object with values that are used only once.
func LoadStartupValues() {

	// Copy in the logging configuration.
	G.Log.Logdir = rawCassabonConfig.Logging.Logdir
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the statsd configuration.
	G.Statsd = rawCassabonConfig.Statsd

	// Copy in and sanitize the channel lengths.
	G.Channels.DataStoreChanLen = rawCassabonConfig.Channels.DataStoreChanLen
	if G.Channels.DataStoreChanLen < 10 {
		G.Channels.DataStoreChanLen = 10
	}
	if G.Channels.DataStoreChanLen > 1000 {
		G.Channels.DataStoreChanLen = 1000
	}
	G.Channels.IndexStoreChanLen = rawCassabonConfig.Channels.IndexStoreChanLen
	if G.Channels.IndexStoreChanLen < 10 {
		G.Channels.IndexStoreChanLen = 10
	}
	if G.Channels.IndexStoreChanLen > 1000 {
		G.Channels.IndexStoreChanLen = 1000
	}
	G.Channels.GopherChanLen = rawCassabonConfig.Channels.GopherChanLen
	if G.Channels.GopherChanLen < 10 {
		G.Channels.GopherChanLen = 10
	}
	if G.Channels.GopherChanLen > 1000 {
		G.Channels.GopherChanLen = 1000
	}
}

// ValidatePeerList ensures addresses are valid, and that the local address is in the peer list.
func ValidatePeerList(localHostPort string, peers []string) error {

	for _, v := range peers {
		if _, err := net.ResolveTCPAddr("tcp4", v); err != nil {
			return fmt.Errorf("Invalid host:port \"%s\" in peer list: %v", v, err)
		}
		if v == localHostPort {
			localHostPort = ""
		}
	}
	if localHostPort != "" {
		return fmt.Errorf("Local host:port %s is not in peer list: %v", localHostPort, peers)
	}

	return nil
}

// LoadRefreshableValues populates the global config objectwith values that
// take effect again on receipt of a SIGHUP.
func LoadRefreshableValues() {

	// Copy in the logging level (can be changed while running).
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the Carbon listener configuration.
	// NOTE: If any of these change, all rollup accumulators must be flushed.
	G.Carbon.Listen = rawCassabonConfig.Carbon.Listen
	G.Carbon.Protocol = rawCassabonConfig.Carbon.Protocol
	G.Carbon.Peers = rawCassabonConfig.Carbon.Peers

	// Ensure a canonical order for Cassabon peers, as we will allocate
	// metrics paths to a peer by indexing into this array.
	sort.Strings(G.Carbon.Peers)

	// Ensure addresses are valid, and that the local address:port is in the peer list.
	if err := ValidatePeerList(G.Carbon.Listen, G.Carbon.Peers); err != nil {
		G.Log.System.LogFatal(err.Error())
		os.Exit(1)
	}

	// Copy in and sanitize the Carbon TCP listener timeout.
	G.Carbon.Parameters.TCPTimeout = rawCassabonConfig.Carbon.Parameters.TCPTimeout
	if G.Carbon.Parameters.TCPTimeout < 1 {
		G.Carbon.Parameters.TCPTimeout = 1
	}
	if G.Carbon.Parameters.TCPTimeout > 30 {
		G.Carbon.Parameters.TCPTimeout = 30
	}

	// Copy in and sanitize the Carbon UDP listener timeout.
	G.Carbon.Parameters.UDPTimeout = rawCassabonConfig.Carbon.Parameters.UDPTimeout
	if G.Carbon.Parameters.UDPTimeout < 1 {
		G.Carbon.Parameters.UDPTimeout = 1
	}
	if G.Carbon.Parameters.UDPTimeout > 30 {
		G.Carbon.Parameters.UDPTimeout = 30
	}

	// Copy in the API configuration values.
	G.API.Listen = rawCassabonConfig.API.Listen
	G.API.HealthCheckFile = rawCassabonConfig.API.HealthCheckFile

	// Copy in the Cassandra database connection values.
	G.Cassandra = rawCassabonConfig.Cassandra

	// Copy in the Redis database connection values.
	G.Redis.Index = rawCassabonConfig.Redis.Index
	G.Redis.Queue = rawCassabonConfig.Redis.Queue
}

// LoadRollups populates the global config object with the rollup definitions,
// which should only happen when there are no accumulated stats.
func LoadRollups() {

	// Validate, copy in and normalize the rollup definitions.
	G.RollupPriority = make([]string, 0, len(rawCassabonConfig.Rollups))
	G.Rollup = make(map[string]RollupDef)

	var method RollupMethod
	var window, retention time.Duration
	var err error
	var intRetention int64
	var rd *RollupDef
	reDuration := regexp.MustCompile("([0-9]+)([a-z])")

	// Inspect each rollup found in the configuration.
	// Note: YAML decode has already folded duplicate path expressions.
	for expression, v := range rawCassabonConfig.Rollups {

		// Validate and decode the aggregation method.
		switch strings.ToLower(v.Aggregation) {
		case "average":
			method = AVERAGE
		case "max":
			method = MAX
		case "min":
			method = MIN
		case "sum":
			method = SUM
		default:
			G.Log.System.LogWarn("Invalid aggregation method for \"%s\": %s", expression, v.Aggregation)
			continue
		}

		// Build up the rollup definitions for this regular expression.
		rd = new(RollupDef)
		rd.Method = method
		rd.Windows = make([]RollupWindow, 0)
		if expression != ROLLUP_CATCHALL {
			if re, err := regexp.Compile(expression); err == nil {
				rd.Expression = re
			} else {
				G.Log.System.LogWarn("Malformed regular expression for \"%s\": %v", expression, err)
				continue
			}
		}

		// Parse and validate each window:retention pair.
		for _, s := range v.Retention {

			// Split the value on the colon between the parts.
			couplet := strings.Split(s, ":")
			if len(couplet) != 2 {
				G.Log.System.LogWarn("Malformed definition for \"%s\": %s", expression, s)
				continue
			}

			// Convert the window to a time.Duration.
			if window, err = time.ParseDuration(couplet[0]); err != nil {
				G.Log.System.LogWarn("Malformed window for \"%s\": %s %s", expression, s, couplet[0])
				continue
			}

			// Don't permit windows shorter than 1 second.
			if window < time.Second {
				G.Log.System.LogWarn("Duration less than minimum 1 second for \"%s\": %v", expression, window)
				continue
			}

			// Convert the retention to a time.Duration (max: 720 years).
			// ParseDuration doesn't handle anything longer than hours, so do it manually.
			matches := reDuration.FindStringSubmatch(couplet[1]) // "1d" -> [ 1d 1 d ]
			if len(matches) != 3 {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			if intRetention, err = strconv.ParseInt(matches[1], 10, 64); err != nil {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			switch matches[2] {
			case "m":
				retention = time.Minute * time.Duration(intRetention)
			case "h":
				retention = time.Hour * time.Duration(intRetention)
			case "d":
				retention = time.Hour * 24 * time.Duration(intRetention)
			case "w":
				retention = time.Hour * 24 * 7 * time.Duration(intRetention)
			case "y":
				retention = time.Hour * 24 * 365 * time.Duration(intRetention)
			default:
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}

			// Append to the rollups for this expression.
			rd.Windows = append(rd.Windows, RollupWindow{window, retention})
		}

		// If any of the rollup window definitions were valid, add expression to the list.
		if len(rd.Windows) > 0 {

			// Sort windows into ascending duration order, and validate.
			sort.Sort(ByWindow(rd.Windows))
			shortestDuration := rd.Windows[0].Window
			allExactMultiples := true
			for i, v := range rd.Windows {
				if i > 0 {
					remainder := v.Window % shortestDuration
					if remainder != 0 {
						G.Log.System.LogWarn(
							"Next duration is not a multiple for \"%s\": %v %% %v remainder is %v",
							expression, v.Window, shortestDuration, remainder)
						allExactMultiples = false
					}
				}
			}

			// If all durations are exact multiples of the shortest duration, save this expression.
			if allExactMultiples {
				G.Rollup[expression] = *rd
				G.RollupPriority = append(G.RollupPriority, expression)
			} else {
				G.Log.System.LogWarn("Rollup expression rejected due to previous errors: \"%s\"", expression)
			}
		}
	}

	// Sort the path expressions into priority order.
	sort.Sort(ByPriority(G.RollupPriority))
}
<commit_msg>Permit use of "0.0.0.0" for Carbon listen address<commit_after>package config

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"gopkg.in/yaml.v2"
)

// Define Application Settings Structure
type CassabonConfig struct {
	Logging struct {
		Logdir   string // Log Directory
		Loglevel string // Level to log at.
	}
	Statsd   StatsdSettings
	Channels struct {
		DataStoreChanLen  int // Length of the DataStore channel
		IndexStoreChanLen int // Length of the Indexer channel
		GopherChanLen     int // Length of the StatGopher channel
	}
	Carbon struct {
		Listen     string // ip:port on which to listen for Carbon stats
		Protocol   string // "tcp", "udp" or "both" are acceptable
		Parameters struct {
			TCPTimeout int
			UDPTimeout int
		}
		Peers []string // All servers in the Cassabon array, as "ip:port"
	}
	API struct {
		Listen          string // HTTP API listens on this address:port
		HealthCheckFile string // Location of healthcheck file.
	}
	Cassandra struct {
		Hosts []string // List of hostnames or IP addresses of Cassandra ring
		Port  string   // Cassandra port
	}
	Redis struct {
		Index RedisSettings // Settings for Redis Index
		Queue RedisSettings // Settings for Redis Queue
	}
	Rollups map[string]RollupSettings // Map of regex and rollups
}

// Definition of each rollup
type RollupSettings struct {
	Retention   []string
	Aggregation string
}

// Redis struct for redis connection information
type RedisSettings struct {
	Sentinel bool     // True if sentinel, false if standalone.
	Addr     []string // List of addresses in host:port format
	DB       int64    // Redis DB number for the index.
	Pwd      string   // Password for Redis.
	Master   string   // Master config name for sentinel settings.
}

type StatsdSettings struct {
	Host   string // Host or IP address of statsd server
	Port   string // Port that statsd server listens on
	Events struct {
		ReceiveOK struct {
			Key        string
			SampleRate float32
		}
		ReceiveFail struct {
			Key        string
			SampleRate float32
		}
	}
}

// rawCassabonConfig is the decoded YAML from the configuration file.
var rawCassabonConfig *CassabonConfig

// ReadConfigurationFile reads the contents of the specified file from disk, and unmarshals it.
// First time through we have no logger, so we can't log warnings or errors. Panic instead.
func ReadConfigurationFile(configFile string, haveLogger bool) error {

	// Read the configuration file.
	yamlConfig, err := ioutil.ReadFile(configFile)
	if err != nil {
		if haveLogger {
			return err
		} else {
			panic(err)
		}
	}

	// Unmarshal config file contents into raw config struct.
	err = yaml.Unmarshal(yamlConfig, &rawCassabonConfig)
	if err != nil {
		if haveLogger {
			return err
		} else {
			panic(err)
		}
	}
	return nil
}

// LoadStartupValues populates the global config object with values that are used only once.
func LoadStartupValues() {

	// Copy in the logging configuration.
	G.Log.Logdir = rawCassabonConfig.Logging.Logdir
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the statsd configuration.
	G.Statsd = rawCassabonConfig.Statsd

	// Copy in and sanitize the channel lengths.
	G.Channels.DataStoreChanLen = rawCassabonConfig.Channels.DataStoreChanLen
	if G.Channels.DataStoreChanLen < 10 {
		G.Channels.DataStoreChanLen = 10
	}
	if G.Channels.DataStoreChanLen > 1000 {
		G.Channels.DataStoreChanLen = 1000
	}
	G.Channels.IndexStoreChanLen = rawCassabonConfig.Channels.IndexStoreChanLen
	if G.Channels.IndexStoreChanLen < 10 {
		G.Channels.IndexStoreChanLen = 10
	}
	if G.Channels.IndexStoreChanLen > 1000 {
		G.Channels.IndexStoreChanLen = 1000
	}
	G.Channels.GopherChanLen = rawCassabonConfig.Channels.GopherChanLen
	if G.Channels.GopherChanLen < 10 {
		G.Channels.GopherChanLen = 10
	}
	if G.Channels.GopherChanLen > 1000 {
		G.Channels.GopherChanLen = 1000
	}
}

// ValidatePeerList ensures addresses are valid, and that the local address is in the peer list.
func ValidatePeerList(localHostPort string, peers []string) error {

	for _, v := range peers {
		if _, err := net.ResolveTCPAddr("tcp4", v); err != nil {
			return fmt.Errorf("Invalid host:port \"%s\" in peer list: %v", v, err)
		}
		if v == localHostPort {
			localHostPort = ""
		}
	}
	if localHostPort != "" {
		return fmt.Errorf("Local host:port %s is not in peer list: %v", localHostPort, peers)
	}

	return nil
}

// LoadRefreshableValues populates the global config objectwith values that
// take effect again on receipt of a SIGHUP.
func LoadRefreshableValues() {

	// Copy in the logging level (can be changed while running).
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// If the listen address is "0.0.0.0", replace it with the address of
	// the first non-localhost, non-IPv6 address found for this machine.
	if lh, lp, err := net.SplitHostPort(rawCassabonConfig.Carbon.Listen); err == nil {
		// Is this "0.0.0.0"?
		if net.ParseIP(lh).IsUnspecified() {
			// Enumerate all local addresses in CIDR format.
			if addr, err := net.InterfaceAddrs(); err == nil {
				for _, v := range addr {
					if ip, _, err := net.ParseCIDR(v.String()); err == nil {
						// Skip loopback and IPv6 addresses.
						if !ip.IsLoopback() && !strings.Contains(ip.String(), ":") {
							rawCassabonConfig.Carbon.Listen = net.JoinHostPort(ip.String(), lp)
							// That's the one we want, we're done.
							break
						}
					}
				}
			}
		}
	}

	// Copy in the Carbon listener configuration.
	// NOTE: If any of these change, all rollup accumulators must be flushed.
	G.Carbon.Listen = rawCassabonConfig.Carbon.Listen
	G.Carbon.Protocol = rawCassabonConfig.Carbon.Protocol
	G.Carbon.Peers = rawCassabonConfig.Carbon.Peers

	// Ensure a canonical order for Cassabon peers, as we will allocate
	// metrics paths to a peer by indexing into this array.
	sort.Strings(G.Carbon.Peers)

	// Ensure addresses are valid, and that the local address:port is in the peer list.
	if err := ValidatePeerList(G.Carbon.Listen, G.Carbon.Peers); err != nil {
		G.Log.System.LogFatal(err.Error())
		os.Exit(1)
	}

	// Copy in and sanitize the Carbon TCP listener timeout.
	G.Carbon.Parameters.TCPTimeout = rawCassabonConfig.Carbon.Parameters.TCPTimeout
	if G.Carbon.Parameters.TCPTimeout < 1 {
		G.Carbon.Parameters.TCPTimeout = 1
	}
	if G.Carbon.Parameters.TCPTimeout > 30 {
		G.Carbon.Parameters.TCPTimeout = 30
	}

	// Copy in and sanitize the Carbon UDP listener timeout.
	G.Carbon.Parameters.UDPTimeout = rawCassabonConfig.Carbon.Parameters.UDPTimeout
	if G.Carbon.Parameters.UDPTimeout < 1 {
		G.Carbon.Parameters.UDPTimeout = 1
	}
	if G.Carbon.Parameters.UDPTimeout > 30 {
		G.Carbon.Parameters.UDPTimeout = 30
	}

	// Copy in the API configuration values.
	G.API.Listen = rawCassabonConfig.API.Listen
	G.API.HealthCheckFile = rawCassabonConfig.API.HealthCheckFile

	// Copy in the Cassandra database connection values.
	G.Cassandra = rawCassabonConfig.Cassandra

	// Copy in the Redis database connection values.
	G.Redis.Index = rawCassabonConfig.Redis.Index
	G.Redis.Queue = rawCassabonConfig.Redis.Queue
}

// LoadRollups populates the global config object with the rollup definitions,
// which should only happen when there are no accumulated stats.
func LoadRollups() {

	// Validate, copy in and normalize the rollup definitions.
	G.RollupPriority = make([]string, 0, len(rawCassabonConfig.Rollups))
	G.Rollup = make(map[string]RollupDef)

	var method RollupMethod
	var window, retention time.Duration
	var err error
	var intRetention int64
	var rd *RollupDef
	reDuration := regexp.MustCompile("([0-9]+)([a-z])")

	// Inspect each rollup found in the configuration.
	// Note: YAML decode has already folded duplicate path expressions.
	for expression, v := range rawCassabonConfig.Rollups {

		// Validate and decode the aggregation method.
		switch strings.ToLower(v.Aggregation) {
		case "average":
			method = AVERAGE
		case "max":
			method = MAX
		case "min":
			method = MIN
		case "sum":
			method = SUM
		default:
			G.Log.System.LogWarn("Invalid aggregation method for \"%s\": %s", expression, v.Aggregation)
			continue
		}

		// Build up the rollup definitions for this regular expression.
		rd = new(RollupDef)
		rd.Method = method
		rd.Windows = make([]RollupWindow, 0)
		if expression != ROLLUP_CATCHALL {
			if re, err := regexp.Compile(expression); err == nil {
				rd.Expression = re
			} else {
				G.Log.System.LogWarn("Malformed regular expression for \"%s\": %v", expression, err)
				continue
			}
		}

		// Parse and validate each window:retention pair.
		for _, s := range v.Retention {

			// Split the value on the colon between the parts.
			couplet := strings.Split(s, ":")
			if len(couplet) != 2 {
				G.Log.System.LogWarn("Malformed definition for \"%s\": %s", expression, s)
				continue
			}

			// Convert the window to a time.Duration.
			if window, err = time.ParseDuration(couplet[0]); err != nil {
				G.Log.System.LogWarn("Malformed window for \"%s\": %s %s", expression, s, couplet[0])
				continue
			}

			// Don't permit windows shorter than 1 second.
			if window < time.Second {
				G.Log.System.LogWarn("Duration less than minimum 1 second for \"%s\": %v", expression, window)
				continue
			}

			// Convert the retention to a time.Duration (max: 720 years).
			// ParseDuration doesn't handle anything longer than hours, so do it manually.
			matches := reDuration.FindStringSubmatch(couplet[1]) // "1d" -> [ 1d 1 d ]
			if len(matches) != 3 {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			if intRetention, err = strconv.ParseInt(matches[1], 10, 64); err != nil {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			switch matches[2] {
			case "m":
				retention = time.Minute * time.Duration(intRetention)
			case "h":
				retention = time.Hour * time.Duration(intRetention)
			case "d":
				retention = time.Hour * 24 * time.Duration(intRetention)
			case "w":
				retention = time.Hour * 24 * 7 * time.Duration(intRetention)
			case "y":
				retention = time.Hour * 24 * 365 * time.Duration(intRetention)
			default:
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}

			// Append to the rollups for this expression.
			rd.Windows = append(rd.Windows, RollupWindow{window, retention})
		}

		// If any of the rollup window definitions were valid, add expression to the list.
		if len(rd.Windows) > 0 {

			// Sort windows into ascending duration order, and validate.
			sort.Sort(ByWindow(rd.Windows))
			shortestDuration := rd.Windows[0].Window
			allExactMultiples := true
			for i, v := range rd.Windows {
				if i > 0 {
					remainder := v.Window % shortestDuration
					if remainder != 0 {
						G.Log.System.LogWarn(
							"Next duration is not a multiple for \"%s\": %v %% %v remainder is %v",
							expression, v.Window, shortestDuration, remainder)
						allExactMultiples = false
					}
				}
			}

			// If all durations are exact multiples of the shortest duration, save this expression.
			if allExactMultiples {
				G.Rollup[expression] = *rd
				G.RollupPriority = append(G.RollupPriority, expression)
			} else {
				G.Log.System.LogWarn("Rollup expression rejected due to previous errors: \"%s\"", expression)
			}
		}
	}

	// Sort the path expressions into priority order.
	sort.Sort(ByPriority(G.RollupPriority))
}
<|endoftext|>


<commit_before>package blocks

import(
	"time"
    "github.com/nytlabs/streamtools/st/loghub"
    "fmt"
)

type Msg struct {
	Msg   interface{}
	Route string
}

type AddChanMsg struct {
	Route   string
	Channel chan *Msg
}

type QueryMsg struct {
	Route    string
	RespChan chan interface{}
}

type BlockChans struct {
	InChan    chan *Msg
	QueryChan chan *QueryMsg
	AddChan   chan *AddChanMsg
	DelChan   chan *Msg
	ErrChan   chan error
	QuitChan  chan bool
}

type LogStreams struct {
	log			chan interface{}
	ui 			chan interface{}
}

type Block struct {
	Id        	string // the name of the block specifed by the user (like MyBlock)
	Kind        string // the kind of block this is (like count, toFile, fromSQS)
	inRoutes    map[string]chan interface{}
	queryRoutes map[string]chan chan interface{}
	broadcast   chan interface{}
	quit 		chan interface{}
	doesBroadcast bool
	BlockChans
	LogStreams
}

type BlockDef struct {
	Type string
	InRoutes []string
	QueryRoutes [] string
	OutRoutes []string
}

type BlockInterface interface {
	Setup()
	Run()
	CleanUp()
	Build(BlockChans)
	Quit() chan interface{}
	Broadcast() chan interface{}
	InRoute(string) chan interface{}
	QueryRoute(string) chan chan interface{}
	GetBlock() *Block
	GetDef() *BlockDef
	Log(interface{})
	Error(interface{})
	SetId(string)
}

func (b *Block) Build(c BlockChans) {
	// fuck can I do this all in one?
	b.InChan = c.InChan
	b.QueryChan = c.QueryChan
	b.AddChan = c.AddChan
	b.DelChan = c.DelChan
	b.ErrChan = c.ErrChan
	b.QuitChan = c.QuitChan

	// route maps
	b.inRoutes = make(map[string]chan interface{}) // necessary to stop locking...
	b.queryRoutes = make(map[string]chan chan interface{})

	// broadcast channel
	b.broadcast = make(chan interface{}, 10) // necessary to stop locking...

	// quit chan
	b.quit = make(chan interface{})

	b.ui = make(chan interface{})
	b.log = make(chan interface{})
}

func (b *Block) SetId(Id string){
	b.Id = Id
}

func (b *Block) InRoute(routeName string) chan interface{} {
	route := make(chan interface{}, 1000)
	b.inRoutes[routeName] = route
	return route
}

func (b *Block) QueryRoute(routeName string) chan chan interface{} {
	route := make(chan chan interface{}, 1000)
	b.queryRoutes[routeName] = route
	return route
}

func (b *Block) Broadcast() chan interface{} {
	b.doesBroadcast = true
	return b.broadcast
}

func (b *Block) Quit() chan interface{} {
	return b.quit
}

func (b *Block) GetBlock() *Block {
	return b
}

func (b *Block) GetDef() *BlockDef {
	inRoutes := []string{}
	queryRoutes := []string{}
	outRoutes := []string{}

	for k, _ := range b.inRoutes {
		inRoutes = append(inRoutes, k)
	}

	for k, _ := range b.queryRoutes {
		queryRoutes = append(queryRoutes, k)
	}

	if b.doesBroadcast {
		outRoutes = []string{"out"}
	}

	return &BlockDef{
		Type: b.Kind,
		InRoutes: inRoutes,
		QueryRoutes: queryRoutes,
		OutRoutes: outRoutes,
	}
}

func (b *Block) CleanUp() {
	for route := range b.inRoutes {
		defer close(b.inRoutes[route])
	}
	for route := range b.queryRoutes {
		defer close(b.queryRoutes[route])
	}
	defer close(b.InChan)
	defer close(b.QueryChan)
	defer close(b.AddChan)
	defer close(b.DelChan)
	defer close(b.ErrChan)
	defer close(b.QuitChan)
	defer close(b.broadcast)

	go func(id string){
	    loghub.Log <- &loghub.LogMsg{
	        Type: loghub.INFO,
	        Data: fmt.Sprintf("Block %s Quitting...", b.Id),
	        Id: id,
	    }
	}(b.Id)
}

func (b *Block) Error(msg interface{}) {
	go func(id string){
	    loghub.Log <- &loghub.LogMsg{
	        Type: loghub.ERROR,
	        Data: msg,
	        Id: id,
	    }
	}(b.Id)
}

func (b *Block) Log(msg interface{}){
	go func(id string) {
	    loghub.Log <- &loghub.LogMsg{
	        Type: loghub.INFO,
	        Data: msg,
	        Id: id,
	    }
	}(b.Id)
}

func BlockRoutine(bi BlockInterface) {
	outChans := make(map[string]chan *Msg)

	b := bi.GetBlock()
	bi.Setup()
	go bi.Run()

	for {
		select {
		case msg := <-b.InChan:
			_, ok := b.inRoutes[msg.Route]
			if !ok {
				break
			}

			// every in channel is buffered a 1000 messages.
			// if we cannot immediately send to that in channel we place the msg
			// in a go routine and notify the user that the block routine's
			// buffer has overflowed. this still allows for unrecoverable
			// overflows (for example: a stuck run() function), but at least it
			// offloads to memory/cpu instead of blocking.
			select {
				case b.inRoutes[msg.Route] <- msg.Msg:
				default:
					go func(id string, msgOut interface{}){
						b.inRoutes[msg.Route] <- msgOut
						loghub.Log <- &loghub.LogMsg{
							Type: loghub.ERROR,
							Data: "Dropping messages!",
							Id:   id,
						}
					}(b.Id, msg.Msg)
			}

			if msg.Route == "rule" {
				go func(id string){
					loghub.UI <- &loghub.LogMsg{
						Type: loghub.RULE_UPDATED,
						Data: map[string]interface{}{},
						Id:   id,
					}
				}(b.Id)
			}

		case msg := <-b.QueryChan:
			_, ok := b.queryRoutes[msg.Route]
			if !ok {
				break
			}

			select {
				case b.queryRoutes[msg.Route] <- msg.RespChan:
				default:
					go func(){
						b.queryRoutes[msg.Route] <- msg.RespChan
					}()
			}

		case msg := <-b.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-b.DelChan:
			delete(outChans, msg.Route)
		case msg := <-b.broadcast:
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg,
					Route: "",
				}
			}
		case <-b.QuitChan:
			b.quit <- true
			b.CleanUp()
			return
		}
	}
}

type Connection struct {
	Id 	string
	ToRoute string
	BlockChans
	LogStreams
}

func (c *Connection) SetId(Id string){
	c.Id = Id
}

func (c *Connection) Build(chans BlockChans){
	c.InChan = chans.InChan
	c.QueryChan = chans.QueryChan
	c.AddChan = chans.AddChan
	c.DelChan = chans.DelChan
	c.QuitChan = chans.QuitChan
}

func (c *Connection) CleanUp(){
	defer close(c.InChan)
	defer close(c.QueryChan)
	defer close(c.AddChan)
	defer close(c.DelChan)
	defer close(c.QuitChan)

    loghub.Log <- &loghub.LogMsg{
        Type: loghub.INFO,
        Data: fmt.Sprintf("Connection %s Quitting...", c.Id),
        Id:   c.Id,
    }
}

func ConnectionRoutine(c *Connection){
	var last interface{}
	var rate float64

	outChans := make(map[string]chan *Msg)
	times := make([]int64,100,100)
	timesIdx := len(times)
	rateReport := time.NewTicker(200 * time.Millisecond)

	for{
		select{
		case <- rateReport.C:
			if timesIdx == len(times) {
				rate = 0
			} else {
				rate = 1000000000.0 * float64(len(times) - timesIdx)/float64(time.Now().UnixNano() - times[timesIdx])
			}

			go func(id string, r float64){
			    loghub.UI <- &loghub.LogMsg{
			        Type: loghub.UPDATE_RATE,
			        Data: map[string]interface{}{
			        	"Rate": r,
			        },
			        Id:   id,
			    }
		    }(c.Id, rate)

		case msg := <- c.InChan:
			last = msg.Msg
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg.Msg,
					Route: c.ToRoute,
				}
			}

			times = times[1:]
			times = append(times, time.Now().UnixNano())

			if timesIdx > 0 {
				timesIdx--
			}

		case msg := <- c.QueryChan:
			switch msg.Route {
			case "rate":
				msg.RespChan <- map[string]interface{}{
					"Rate" : rate,
				}
			case "last":
				msg.RespChan <- map[string]interface{}{
					"Last" : last,
				}
			}
		case msg := <- c.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <- c.DelChan:
			delete(outChans, msg.Route)
		case <- c.QuitChan:
			c.CleanUp()
			return
		}
	}
}
<commit_msg>indenting commit, ws only<commit_after>package blocks

import (
	"fmt"
	"github.com/nytlabs/streamtools/st/loghub"
	"time"
)

type Msg struct {
	Msg   interface{}
	Route string
}

type AddChanMsg struct {
	Route   string
	Channel chan *Msg
}

type QueryMsg struct {
	Route    string
	RespChan chan interface{}
}

type BlockChans struct {
	InChan    chan *Msg
	QueryChan chan *QueryMsg
	AddChan   chan *AddChanMsg
	DelChan   chan *Msg
	ErrChan   chan error
	QuitChan  chan bool
}

type LogStreams struct {
	log chan interface{}
	ui  chan interface{}
}

type Block struct {
	Id            string // the name of the block specifed by the user (like MyBlock)
	Kind          string // the kind of block this is (like count, toFile, fromSQS)
	inRoutes      map[string]chan interface{}
	queryRoutes   map[string]chan chan interface{}
	broadcast     chan interface{}
	quit          chan interface{}
	doesBroadcast bool
	BlockChans
	LogStreams
}

type BlockDef struct {
	Type        string
	InRoutes    []string
	QueryRoutes []string
	OutRoutes   []string
}

type BlockInterface interface {
	Setup()
	Run()
	CleanUp()
	Build(BlockChans)
	Quit() chan interface{}
	Broadcast() chan interface{}
	InRoute(string) chan interface{}
	QueryRoute(string) chan chan interface{}
	GetBlock() *Block
	GetDef() *BlockDef
	Log(interface{})
	Error(interface{})
	SetId(string)
}

func (b *Block) Build(c BlockChans) {
	// fuck can I do this all in one?
	b.InChan = c.InChan
	b.QueryChan = c.QueryChan
	b.AddChan = c.AddChan
	b.DelChan = c.DelChan
	b.ErrChan = c.ErrChan
	b.QuitChan = c.QuitChan

	// route maps
	b.inRoutes = make(map[string]chan interface{}) // necessary to stop locking...
	b.queryRoutes = make(map[string]chan chan interface{})

	// broadcast channel
	b.broadcast = make(chan interface{}, 10) // necessary to stop locking...

	// quit chan
	b.quit = make(chan interface{})

	b.ui = make(chan interface{})
	b.log = make(chan interface{})
}

func (b *Block) SetId(Id string) {
	b.Id = Id
}

func (b *Block) InRoute(routeName string) chan interface{} {
	route := make(chan interface{}, 1000)
	b.inRoutes[routeName] = route
	return route
}

func (b *Block) QueryRoute(routeName string) chan chan interface{} {
	route := make(chan chan interface{}, 1000)
	b.queryRoutes[routeName] = route
	return route
}

func (b *Block) Broadcast() chan interface{} {
	b.doesBroadcast = true
	return b.broadcast
}

func (b *Block) Quit() chan interface{} {
	return b.quit
}

func (b *Block) GetBlock() *Block {
	return b
}

func (b *Block) GetDef() *BlockDef {
	inRoutes := []string{}
	queryRoutes := []string{}
	outRoutes := []string{}

	for k, _ := range b.inRoutes {
		inRoutes = append(inRoutes, k)
	}

	for k, _ := range b.queryRoutes {
		queryRoutes = append(queryRoutes, k)
	}

	if b.doesBroadcast {
		outRoutes = []string{"out"}
	}

	return &BlockDef{
		Type:        b.Kind,
		InRoutes:    inRoutes,
		QueryRoutes: queryRoutes,
		OutRoutes:   outRoutes,
	}
}

func (b *Block) CleanUp() {
	for route := range b.inRoutes {
		defer close(b.inRoutes[route])
	}
	for route := range b.queryRoutes {
		defer close(b.queryRoutes[route])
	}
	defer close(b.InChan)
	defer close(b.QueryChan)
	defer close(b.AddChan)
	defer close(b.DelChan)
	defer close(b.ErrChan)
	defer close(b.QuitChan)
	defer close(b.broadcast)

	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.INFO,
			Data: fmt.Sprintf("Block %s Quitting...", b.Id),
			Id:   id,
		}
	}(b.Id)
}

func (b *Block) Error(msg interface{}) {
	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.ERROR,
			Data: msg,
			Id:   id,
		}
	}(b.Id)
}

func (b *Block) Log(msg interface{}) {
	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.INFO,
			Data: msg,
			Id:   id,
		}
	}(b.Id)
}

func BlockRoutine(bi BlockInterface) {
	outChans := make(map[string]chan *Msg)

	b := bi.GetBlock()
	bi.Setup()
	go bi.Run()

	for {
		select {
		case msg := <-b.InChan:
			_, ok := b.inRoutes[msg.Route]
			if !ok {
				break
			}

			// every in channel is buffered a 1000 messages.
			// if we cannot immediately send to that in channel we place the msg
			// in a go routine and notify the user that the block routine's
			// buffer has overflowed. this still allows for unrecoverable
			// overflows (for example: a stuck run() function), but at least it
			// offloads to memory/cpu instead of blocking.
			select {
			case b.inRoutes[msg.Route] <- msg.Msg:
			default:
				go func(id string, msgOut interface{}) {
					b.inRoutes[msg.Route] <- msgOut
					loghub.Log <- &loghub.LogMsg{
						Type: loghub.ERROR,
						Data: "Dropping messages!",
						Id:   id,
					}
				}(b.Id, msg.Msg)
			}

			if msg.Route == "rule" {
				go func(id string) {
					loghub.UI <- &loghub.LogMsg{
						Type: loghub.RULE_UPDATED,
						Data: map[string]interface{}{},
						Id:   id,
					}
				}(b.Id)
			}

		case msg := <-b.QueryChan:
			_, ok := b.queryRoutes[msg.Route]
			if !ok {
				break
			}

			select {
			case b.queryRoutes[msg.Route] <- msg.RespChan:
			default:
				go func() {
					b.queryRoutes[msg.Route] <- msg.RespChan
				}()
			}

		case msg := <-b.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-b.DelChan:
			delete(outChans, msg.Route)
		case msg := <-b.broadcast:
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg,
					Route: "",
				}
			}
		case <-b.QuitChan:
			b.quit <- true
			b.CleanUp()
			return
		}
	}
}

type Connection struct {
	Id      string
	ToRoute string
	BlockChans
	LogStreams
}

func (c *Connection) SetId(Id string) {
	c.Id = Id
}

func (c *Connection) Build(chans BlockChans) {
	c.InChan = chans.InChan
	c.QueryChan = chans.QueryChan
	c.AddChan = chans.AddChan
	c.DelChan = chans.DelChan
	c.QuitChan = chans.QuitChan
}

func (c *Connection) CleanUp() {
	defer close(c.InChan)
	defer close(c.QueryChan)
	defer close(c.AddChan)
	defer close(c.DelChan)
	defer close(c.QuitChan)

	loghub.Log <- &loghub.LogMsg{
		Type: loghub.INFO,
		Data: fmt.Sprintf("Connection %s Quitting...", c.Id),
		Id:   c.Id,
	}
}

func ConnectionRoutine(c *Connection) {
	var last interface{}
	var rate float64

	outChans := make(map[string]chan *Msg)
	times := make([]int64, 100, 100)
	timesIdx := len(times)
	rateReport := time.NewTicker(200 * time.Millisecond)

	for {
		select {
		case <-rateReport.C:
			if timesIdx == len(times) {
				rate = 0
			} else {
				rate = 1000000000.0 * float64(len(times)-timesIdx) / float64(time.Now().UnixNano()-times[timesIdx])
			}

			go func(id string, r float64) {
				loghub.UI <- &loghub.LogMsg{
					Type: loghub.UPDATE_RATE,
					Data: map[string]interface{}{
						"Rate": r,
					},
					Id: id,
				}
			}(c.Id, rate)

		case msg := <-c.InChan:
			last = msg.Msg
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg.Msg,
					Route: c.ToRoute,
				}
			}

			times = times[1:]
			times = append(times, time.Now().UnixNano())

			if timesIdx > 0 {
				timesIdx--
			}

		case msg := <-c.QueryChan:
			switch msg.Route {
			case "rate":
				msg.RespChan <- map[string]interface{}{
					"Rate": rate,
				}
			case "last":
				msg.RespChan <- map[string]interface{}{
					"Last": last,
				}
			}
		case msg := <-c.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-c.DelChan:
			delete(outChans, msg.Route)
		case <-c.QuitChan:
			c.CleanUp()
			return
		}
	}
}
<|endoftext|>


<commit_before>package state

//
// DefinitionSelect postgres specific query for definitions
//
const DefinitionSelect = `
select td.definition_id                    as definitionid,
       td.adaptive_resource_allocation     as adaptiveresourceallocation,
       td.image                            as image,
       td.group_name                       as groupname,
       td.alias                            as alias,
       td.memory                           as memory,
       coalesce(td.command, '')            as command,
       coalesce(td.task_type, '')          as tasktype,
       env::TEXT                           as env,
       td.cpu                              as cpu,
       td.gpu                              as gpu,
       array_to_json('{""}'::TEXT[])::TEXT as tags,
       array_to_json('{}'::INT[])::TEXT    as ports
from (select * from task_def) td
`

//
// ListDefinitionsSQL postgres specific query for listing definitions
//
const ListDefinitionsSQL = DefinitionSelect + "\n%s %s limit $1 offset $2"

//
// GetDefinitionSQL postgres specific query for getting a single definition
//
const GetDefinitionSQL = DefinitionSelect + "\nwhere definition_id = $1"

//
// GetDefinitionByAliasSQL get definition by alias
//
const GetDefinitionByAliasSQL = DefinitionSelect + "\nwhere alias = $1"

const TaskResourcesSelectCommandSQL = `
SELECT cast((percentile_disc(0.99) within GROUP (ORDER BY A.max_memory_used)) * 1.75 as int) as memory,
       cast((percentile_disc(0.99) within GROUP (ORDER BY A.max_cpu_used)) * 1.25  as int)  as cpu
FROM (SELECT memory as max_memory_used, cpu as max_cpu_used
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '3 days'
           AND (exit_code = 137 or exit_reason = 'OOMKilled')
           AND engine = 'eks'
           AND definition_id = $1
           AND command_hash = (SELECT command_hash FROM task WHERE run_id = $2)
      LIMIT 30) A
`

const TaskResourcesExecutorCountSQL = `
SELECT least(coalesce(cast((percentile_disc(0.99) within GROUP (ORDER BY A.executor_count)) as int), 25), 100) as executor_count
FROM (SELECT CASE
                 WHEN (exit_reason like '%Exception%')
                     THEN (spark_extension -> 'spark_submit_job_driver' -> 'num_executors')::int * 1.75
                 ELSE (spark_extension -> 'spark_submit_job_driver' -> 'num_executors')::int * 1
                 END as executor_count
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '30 days'
           AND engine = 'eks-spark'
           AND definition_id = $1
           AND command_hash = $2
      LIMIT 30) A
`
const TaskResourcesDriverOOMSQL = `
SELECT (spark_extension -> 'driver_oom')::boolean AS driver_oom
FROM TASK
WHERE queued_at >= CURRENT_TIMESTAMP - INTERVAL '7 days'
  AND engine = 'eks-spark'
  AND definition_id = $1
  AND command_hash = $2
  AND exit_code = 137
  AND spark_extension ? 'driver_oom'
GROUP BY 1
`

const TaskResourcesExecutorOOMSQL = `
SELECT CASE WHEN A.c >= 1 THEN true::boolean ELSE false::boolean END
FROM (SELECT count(*) as c
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '7 days'
           AND definition_id = $1
           AND command_hash = $2
		   AND engine = 'eks-spark'
           AND exit_code !=0
      LIMIT 30) A
`

const TaskResourcesExecutorNodeLifecycleSQL = `
SELECT CASE WHEN A.c >= 1 THEN 'ondemand' ELSE 'spot' END
FROM (SELECT count(*) as c
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '12 hour'
           AND definition_id = $1
           AND command_hash = $2
           AND exit_code !=0
      LIMIT 30) A
`

const TaskExecutionRuntimeCommandSQL = `
SELECT percentile_disc(0.95) within GROUP (ORDER BY A.minutes) as minutes
FROM (SELECT EXTRACT(epoch from finished_at - started_at) / 60 as minutes
      FROM TASK
      WHERE definition_id = $1
        AND exit_code = 0
        AND engine = 'eks'
        AND queued_at >= CURRENT_TIMESTAMP - INTERVAL '7 days'
        AND command_hash = (SELECT command_hash FROM task WHERE run_id = $2)
      LIMIT 30) A
`

const ListFailingNodesSQL = `
SELECT instance_dns_name
FROM (
         SELECT instance_dns_name, count(*) as c
         FROM TASK
         WHERE (exit_code = 128 OR
                pod_events @> '[{"reason": "Failed"}]' OR
                pod_events @> '[{"reason": "FailedSync"}]' OR
                pod_events @> '[{"reason": "FailedCreatePodSandBox"}]' OR
                pod_events @> '[{"reason": "OutOfmemory"}]')
           AND engine = 'eks'
           AND queued_at >= NOW() - INTERVAL '1 HOURS'
           AND instance_dns_name like 'ip-%'
         GROUP BY 1
         order by 2 desc) AS all_nodes
WHERE c >= 5
`

const PodReAttemptRate = `
SELECT (multiple_attempts / (CASE WHEN single_attempts = 0 THEN 1 ELSE single_attempts END)) AS attempts
FROM (
      SELECT COUNT(CASE WHEN attempt_count <= 1 THEN 1 END) * 1.0 AS single_attempts,
             COUNT(CASE WHEN attempt_count > 1 THEN 1 END) * 1.0 AS multiple_attempts
      FROM task
      WHERE engine = 'eks' AND
            queued_at >= NOW() - INTERVAL '18 MINUTES' AND
            node_lifecycle = 'spot') A
`

//
// RunSelect postgres specific query for runs
//
const RunSelect = `
select t.run_id                          as runid,
       coalesce(t.definition_id, '')     as definitionid,
       coalesce(t.alias, '')             as alias,
       coalesce(t.image, '')             as image,
       coalesce(t.cluster_name, '')      as clustername,
       t.exit_code                       as exitcode,
       t.exit_reason                     as exitreason,
       coalesce(t.status, '')            as status,
       queued_at                         as queuedat,
       started_at                        as startedat,
       finished_at                       as finishedat,
       coalesce(t.instance_id, '')       as instanceid,
       coalesce(t.instance_dns_name, '') as instancednsname,
       coalesce(t.group_name, '')        as groupname,
       coalesce(t.task_type, '')         as tasktype,
       env::TEXT                         as env,
       command,
       memory,
       cpu,
       gpu,
       engine,
       ephemeral_storage                 as ephemeralstorage,
       node_lifecycle                    as nodelifecycle,
       pod_name                          as podname,
       namespace,
       max_cpu_used                      as maxcpuused,
       max_memory_used                   as maxmemoryused,
       pod_events::TEXT                  as podevents,
       command_hash                      as commandhash,
       cloudtrail_notifications::TEXT    as cloudtrailnotifications,
       coalesce(executable_id, '')       as executableid,
       coalesce(executable_type, '')     as executabletype,
       execution_request_custom::TEXT    as executionrequestcustom,
       cpu_limit                         as cpulimit,
       memory_limit                      as memorylimit,
       attempt_count                     as attemptcount,
       spawned_runs::TEXT                as spawnedruns,
       run_exceptions::TEXT              as runexceptions,
       active_deadline_seconds           as activedeadlineseconds,
       spark_extension::TEXT             as sparkextension,
       metrics_uri                       as metricsuri,
       description                       as description
from task t
`

//
// ListRunsSQL postgres specific query for listing runs
//
const ListRunsSQL = RunSelect + "\n%s %s limit $1 offset $2"

//
// GetRunSQL postgres specific query for getting a single run
//
const GetRunSQL = RunSelect + "\nwhere run_id = $1"

const GetRunSQLByEMRJobId = RunSelect + "\nwhere spark_extension->>'emr_job_id' = $1"

//
// GetRunSQLForUpdate postgres specific query for getting a single run
// for update
//
const GetRunSQLForUpdate = GetRunSQL + " for update"

//
// GroupsSelect postgres specific query for getting existing definition
// group_names
//
const GroupsSelect = `
select distinct group_name from task_def
`

//
// TagsSelect postgres specific query for getting existing definition tags
//
const TagsSelect = `
select distinct text from tags
`

//
// ListGroupsSQL postgres specific query for listing definition group_names
//
const ListGroupsSQL = GroupsSelect + "\n%s order by group_name asc limit $1 offset $2"

//
// ListTagsSQL postgres specific query for listing definition tags
//
const ListTagsSQL = TagsSelect + "\n%s order by text asc limit $1 offset $2"

//
// WorkerSelect postgres specific query for workers
//
const WorkerSelect = `
  select
    worker_type        as workertype,
    count_per_instance as countperinstance,
    engine
  from worker
`

//
// ListWorkersSQL postgres specific query for listing workers
//
const ListWorkersSQL = WorkerSelect

const GetWorkerEngine = WorkerSelect + "\nwhere engine = $1"

//
// GetWorkerSQL postgres specific query for retrieving data for a specific
// worker type.
//
const GetWorkerSQL = WorkerSelect + "\nwhere worker_type = $1 and engine = $2"

//
// GetWorkerSQLForUpdate postgres specific query for retrieving data for a specific
// worker type; locks the row.
//
const GetWorkerSQLForUpdate = GetWorkerSQL + " for update"

// TemplateSelect selects a template
const TemplateSelect = `
SELECT
  template_id as templateid,
  template_name as templatename,
  version,
  schema,
  command_template as commandtemplate,
  adaptive_resource_allocation as adaptiveresourceallocation,
  image,
  memory,
  env::TEXT as env,
  privileged,
  cpu,
  gpu,
  defaults,
  coalesce(avatar_uri, '') as avataruri
FROM template
`

// ListTemplatesSQL postgres specific query for listing templates
const ListTemplatesSQL = TemplateSelect + "\n%s limit $1 offset $2"

// GetTemplateByIDSQL postgres specific query for getting a single template
const GetTemplateByIDSQL = TemplateSelect + "\nwhere template_id = $1"

// ListTemplatesLatestOnlySQL lists the latest version of each distinct
// template name.
const ListTemplatesLatestOnlySQL = `
  SELECT DISTINCT ON (template_name)
    template_id as templateid,
    template_name as templatename,
    version,
    schema,
    command_template as commandtemplate,
    adaptive_resource_allocation as adaptiveresourceallocation,
    image,
    memory,
    env::TEXT as env,
    privileged,
    cpu,
    gpu,
    defaults,
    coalesce(avatar_uri, '') as avataruri
  FROM template
  ORDER BY template_name, version DESC, template_id
  LIMIT $1 OFFSET $2
`

// GetTemplateLatestOnlySQL get the latest version of a specific template name.
const GetTemplateLatestOnlySQL = TemplateSelect + "\nWHERE template_name = $1 ORDER BY version DESC LIMIT 1;"
const GetTemplateByVersionSQL = TemplateSelect + "\nWHERE template_name = $1 AND version = $2 ORDER BY version DESC LIMIT 1;"
<commit_msg>reduce lookback window<commit_after>package state

//
// DefinitionSelect postgres specific query for definitions
//
const DefinitionSelect = `
select td.definition_id                    as definitionid,
       td.adaptive_resource_allocation     as adaptiveresourceallocation,
       td.image                            as image,
       td.group_name                       as groupname,
       td.alias                            as alias,
       td.memory                           as memory,
       coalesce(td.command, '')            as command,
       coalesce(td.task_type, '')          as tasktype,
       env::TEXT                           as env,
       td.cpu                              as cpu,
       td.gpu                              as gpu,
       array_to_json('{""}'::TEXT[])::TEXT as tags,
       array_to_json('{}'::INT[])::TEXT    as ports
from (select * from task_def) td
`

//
// ListDefinitionsSQL postgres specific query for listing definitions
//
const ListDefinitionsSQL = DefinitionSelect + "\n%s %s limit $1 offset $2"

//
// GetDefinitionSQL postgres specific query for getting a single definition
//
const GetDefinitionSQL = DefinitionSelect + "\nwhere definition_id = $1"

//
// GetDefinitionByAliasSQL get definition by alias
//
const GetDefinitionByAliasSQL = DefinitionSelect + "\nwhere alias = $1"

const TaskResourcesSelectCommandSQL = `
SELECT cast((percentile_disc(0.99) within GROUP (ORDER BY A.max_memory_used)) * 1.75 as int) as memory,
       cast((percentile_disc(0.99) within GROUP (ORDER BY A.max_cpu_used)) * 1.25  as int)  as cpu
FROM (SELECT memory as max_memory_used, cpu as max_cpu_used
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '3 days'
           AND (exit_code = 137 or exit_reason = 'OOMKilled')
           AND engine = 'eks'
           AND definition_id = $1
           AND command_hash = (SELECT command_hash FROM task WHERE run_id = $2)
      LIMIT 30) A
`

const TaskResourcesExecutorCountSQL = `
SELECT least(coalesce(cast((percentile_disc(0.99) within GROUP (ORDER BY A.executor_count)) as int), 25), 100) as executor_count
FROM (SELECT CASE
                 WHEN (exit_reason like '%Exception%')
                     THEN (spark_extension -> 'spark_submit_job_driver' -> 'num_executors')::int * 1.75
                 ELSE (spark_extension -> 'spark_submit_job_driver' -> 'num_executors')::int * 1
                 END as executor_count
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '24 hours'
           AND engine = 'eks-spark'
           AND definition_id = $1
           AND command_hash = $2
      LIMIT 30) A
`
const TaskResourcesDriverOOMSQL = `
SELECT (spark_extension -> 'driver_oom')::boolean AS driver_oom
FROM TASK
WHERE queued_at >= CURRENT_TIMESTAMP - INTERVAL '7 days'
  AND engine = 'eks-spark'
  AND definition_id = $1
  AND command_hash = $2
  AND exit_code = 137
  AND spark_extension ? 'driver_oom'
GROUP BY 1
`

const TaskResourcesExecutorOOMSQL = `
SELECT CASE WHEN A.c >= 1 THEN true::boolean ELSE false::boolean END
FROM (SELECT count(*) as c
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '7 days'
           AND definition_id = $1
           AND command_hash = $2
		   AND engine = 'eks-spark'
           AND exit_code !=0
      LIMIT 30) A
`

const TaskResourcesExecutorNodeLifecycleSQL = `
SELECT CASE WHEN A.c >= 1 THEN 'ondemand' ELSE 'spot' END
FROM (SELECT count(*) as c
      FROM TASK
      WHERE
           queued_at >= CURRENT_TIMESTAMP - INTERVAL '12 hour'
           AND definition_id = $1
           AND command_hash = $2
           AND exit_code !=0
      LIMIT 30) A
`

const TaskExecutionRuntimeCommandSQL = `
SELECT percentile_disc(0.95) within GROUP (ORDER BY A.minutes) as minutes
FROM (SELECT EXTRACT(epoch from finished_at - started_at) / 60 as minutes
      FROM TASK
      WHERE definition_id = $1
        AND exit_code = 0
        AND engine = 'eks'
        AND queued_at >= CURRENT_TIMESTAMP - INTERVAL '7 days'
        AND command_hash = (SELECT command_hash FROM task WHERE run_id = $2)
      LIMIT 30) A
`

const ListFailingNodesSQL = `
SELECT instance_dns_name
FROM (
         SELECT instance_dns_name, count(*) as c
         FROM TASK
         WHERE (exit_code = 128 OR
                pod_events @> '[{"reason": "Failed"}]' OR
                pod_events @> '[{"reason": "FailedSync"}]' OR
                pod_events @> '[{"reason": "FailedCreatePodSandBox"}]' OR
                pod_events @> '[{"reason": "OutOfmemory"}]')
           AND engine = 'eks'
           AND queued_at >= NOW() - INTERVAL '1 HOURS'
           AND instance_dns_name like 'ip-%'
         GROUP BY 1
         order by 2 desc) AS all_nodes
WHERE c >= 5
`

const PodReAttemptRate = `
SELECT (multiple_attempts / (CASE WHEN single_attempts = 0 THEN 1 ELSE single_attempts END)) AS attempts
FROM (
      SELECT COUNT(CASE WHEN attempt_count <= 1 THEN 1 END) * 1.0 AS single_attempts,
             COUNT(CASE WHEN attempt_count > 1 THEN 1 END) * 1.0 AS multiple_attempts
      FROM task
      WHERE engine = 'eks' AND
            queued_at >= NOW() - INTERVAL '18 MINUTES' AND
            node_lifecycle = 'spot') A
`

//
// RunSelect postgres specific query for runs
//
const RunSelect = `
select t.run_id                          as runid,
       coalesce(t.definition_id, '')     as definitionid,
       coalesce(t.alias, '')             as alias,
       coalesce(t.image, '')             as image,
       coalesce(t.cluster_name, '')      as clustername,
       t.exit_code                       as exitcode,
       t.exit_reason                     as exitreason,
       coalesce(t.status, '')            as status,
       queued_at                         as queuedat,
       started_at                        as startedat,
       finished_at                       as finishedat,
       coalesce(t.instance_id, '')       as instanceid,
       coalesce(t.instance_dns_name, '') as instancednsname,
       coalesce(t.group_name, '')        as groupname,
       coalesce(t.task_type, '')         as tasktype,
       env::TEXT                         as env,
       command,
       memory,
       cpu,
       gpu,
       engine,
       ephemeral_storage                 as ephemeralstorage,
       node_lifecycle                    as nodelifecycle,
       pod_name                          as podname,
       namespace,
       max_cpu_used                      as maxcpuused,
       max_memory_used                   as maxmemoryused,
       pod_events::TEXT                  as podevents,
       command_hash                      as commandhash,
       cloudtrail_notifications::TEXT    as cloudtrailnotifications,
       coalesce(executable_id, '')       as executableid,
       coalesce(executable_type, '')     as executabletype,
       execution_request_custom::TEXT    as executionrequestcustom,
       cpu_limit                         as cpulimit,
       memory_limit                      as memorylimit,
       attempt_count                     as attemptcount,
       spawned_runs::TEXT                as spawnedruns,
       run_exceptions::TEXT              as runexceptions,
       active_deadline_seconds           as activedeadlineseconds,
       spark_extension::TEXT             as sparkextension,
       metrics_uri                       as metricsuri,
       description                       as description
from task t
`

//
// ListRunsSQL postgres specific query for listing runs
//
const ListRunsSQL = RunSelect + "\n%s %s limit $1 offset $2"

//
// GetRunSQL postgres specific query for getting a single run
//
const GetRunSQL = RunSelect + "\nwhere run_id = $1"

const GetRunSQLByEMRJobId = RunSelect + "\nwhere spark_extension->>'emr_job_id' = $1"

//
// GetRunSQLForUpdate postgres specific query for getting a single run
// for update
//
const GetRunSQLForUpdate = GetRunSQL + " for update"

//
// GroupsSelect postgres specific query for getting existing definition
// group_names
//
const GroupsSelect = `
select distinct group_name from task_def
`

//
// TagsSelect postgres specific query for getting existing definition tags
//
const TagsSelect = `
select distinct text from tags
`

//
// ListGroupsSQL postgres specific query for listing definition group_names
//
const ListGroupsSQL = GroupsSelect + "\n%s order by group_name asc limit $1 offset $2"

//
// ListTagsSQL postgres specific query for listing definition tags
//
const ListTagsSQL = TagsSelect + "\n%s order by text asc limit $1 offset $2"

//
// WorkerSelect postgres specific query for workers
//
const WorkerSelect = `
  select
    worker_type        as workertype,
    count_per_instance as countperinstance,
    engine
  from worker
`

//
// ListWorkersSQL postgres specific query for listing workers
//
const ListWorkersSQL = WorkerSelect

const GetWorkerEngine = WorkerSelect + "\nwhere engine = $1"

//
// GetWorkerSQL postgres specific query for retrieving data for a specific
// worker type.
//
const GetWorkerSQL = WorkerSelect + "\nwhere worker_type = $1 and engine = $2"

//
// GetWorkerSQLForUpdate postgres specific query for retrieving data for a specific
// worker type; locks the row.
//
const GetWorkerSQLForUpdate = GetWorkerSQL + " for update"

// TemplateSelect selects a template
const TemplateSelect = `
SELECT
  template_id as templateid,
  template_name as templatename,
  version,
  schema,
  command_template as commandtemplate,
  adaptive_resource_allocation as adaptiveresourceallocation,
  image,
  memory,
  env::TEXT as env,
  privileged,
  cpu,
  gpu,
  defaults,
  coalesce(avatar_uri, '') as avataruri
FROM template
`

// ListTemplatesSQL postgres specific query for listing templates
const ListTemplatesSQL = TemplateSelect + "\n%s limit $1 offset $2"

// GetTemplateByIDSQL postgres specific query for getting a single template
const GetTemplateByIDSQL = TemplateSelect + "\nwhere template_id = $1"

// ListTemplatesLatestOnlySQL lists the latest version of each distinct
// template name.
const ListTemplatesLatestOnlySQL = `
  SELECT DISTINCT ON (template_name)
    template_id as templateid,
    template_name as templatename,
    version,
    schema,
    command_template as commandtemplate,
    adaptive_resource_allocation as adaptiveresourceallocation,
    image,
    memory,
    env::TEXT as env,
    privileged,
    cpu,
    gpu,
    defaults,
    coalesce(avatar_uri, '') as avataruri
  FROM template
  ORDER BY template_name, version DESC, template_id
  LIMIT $1 OFFSET $2
`

// GetTemplateLatestOnlySQL get the latest version of a specific template name.
const GetTemplateLatestOnlySQL = TemplateSelect + "\nWHERE template_name = $1 ORDER BY version DESC LIMIT 1;"
const GetTemplateByVersionSQL = TemplateSelect + "\nWHERE template_name = $1 AND version = $2 ORDER BY version DESC LIMIT 1;"
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package doc

import (
	"encoding/xml"
	"sort"
	"strconv"

	"github.com/issue9/is"

	"github.com/caixw/apidoc/v5/internal/locale"
)

// Param 表示参数类型
//  <param name="user" deprecated="1.1.1" type="object" array="true">
//      <param name="name" type="string" />
//      <param name="sex" type="string">
//          <enum value="male" summary="male" />
//          <enum value="female" summary="female" />
//      </param>
//      <param name="age" type="number" />
//  </param>
type Param struct {
	Name        string   `xml:"name,attr"`
	Type        Type     `xml:"type,attr"`
	Deprecated  Version  `xml:"deprecated,attr,omitempty"`
	Default     string   `xml:"default,attr,omitempty"`
	Optional    bool     `xml:"optional,attr,omitempty"`
	Array       bool     `xml:"array,attr,omitempty"`
	Attr        bool     `xml:"attr,attr,omitempty"` // 作为父元素的一个属性，仅针对 xml 有效果
	Items       []*Param `xml:"param,omitempty"`
	Reference   string   `xml:"ref,attr,omitempty"`
	Summary     string   `xml:"summary,attr,omitempty"`
	Enums       []*Enum  `xml:"enum,omitempty"`
	Description Richtext `xml:"description,omitempty"`
}

// IsEnum 是否为一个枚举类型
func (p *Param) IsEnum() bool {
	return len(p.Enums) > 0
}

type shadowParam Param

// UnmarshalXML xml.Unmarshaler
func (p *Param) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	field := "/" + start.Name.Local
	shadow := (*shadowParam)(p)
	if err := d.DecodeElement(shadow, &start); err != nil {
		return fixedSyntaxError(err, "", field, 0)
	}

	if shadow.Name == "" {
		return newSyntaxError(field+"#name", locale.ErrRequired)
	}

	if shadow.Type == None {
		return newSyntaxError(field+"#type", locale.ErrRequired)
	}
	if shadow.Type == Object && len(shadow.Items) == 0 {
		return newSyntaxError(field+"/items", locale.ErrRequired)
	}

	// 判断 enums 的值是否相同
	if key := getDuplicateEnum(shadow.Enums); key != "" {
		return newSyntaxError(field+"/enum", locale.ErrDuplicateValue)
	}

	if err := chkEnumsType(shadow.Type, shadow.Enums, field); err != nil {
		return err
	}

	// 判断 items 的值是否相同
	if key := getDuplicateItems(shadow.Items); key != "" {
		return newSyntaxError(field+"/items", locale.ErrDuplicateValue)
	}

	if p.Summary == "" && p.Description.Text == "" {
		return newSyntaxError(field+"/summary", locale.ErrRequired)
	}

	return nil
}

// 检测 enums 中的类型是否符合 t 的标准，比如 Number 要求枚举值也都是数值
func chkEnumsType(t Type, enums []*Enum, field string) error {
	if len(enums) == 0 {
		return nil
	}

	switch t {
	case Number:
		for _, enum := range enums {
			if !is.Number(enum.Value) {
				return newSyntaxError(field+"/enum/#"+enum.Value, locale.ErrInvalidFormat)
			}
		}
	case Bool:
		for _, enum := range enums {
			if _, err := strconv.ParseBool(enum.Value); err != nil {
				return newSyntaxError(field+"/enum/#"+enum.Value, locale.ErrInvalidFormat)
			}
		}
	case Object, None:
		return newSyntaxError(field+"/enum", locale.ErrInvalidValue)
	}

	return nil
}

// 返回重复枚举的值
func getDuplicateEnum(enums []*Enum) string {
	if len(enums) == 0 {
		return ""
	}

	es := make([]string, 0, len(enums))
	for _, e := range enums {
		es = append(es, e.Value)
	}
	sort.SliceStable(es, func(i, j int) bool { return es[i] > es[j] })

	for i := 1; i < len(es); i++ {
		if es[i] == es[i-1] {
			return es[i]
		}
	}

	return ""
}

func getDuplicateItems(items []*Param) string {
	if len(items) == 0 {
		return ""
	}

	es := make([]string, 0, len(items))
	for _, e := range items {
		es = append(es, e.Name)
	}
	sort.SliceStable(es, func(i, j int) bool { return es[i] > es[j] })

	for i := 1; i < len(es); i++ {
		if es[i] == es[i-1] {
			return es[i]
		}
	}

	return ""
}
<commit_msg>修正部分错误提示依然采用 # 表示属性值的错误<commit_after>// SPDX-License-Identifier: MIT

package doc

import (
	"encoding/xml"
	"sort"
	"strconv"

	"github.com/issue9/is"

	"github.com/caixw/apidoc/v5/internal/locale"
)

// Param 表示参数类型
//  <param name="user" deprecated="1.1.1" type="object" array="true">
//      <param name="name" type="string" />
//      <param name="sex" type="string">
//          <enum value="male" summary="male" />
//          <enum value="female" summary="female" />
//      </param>
//      <param name="age" type="number" />
//  </param>
type Param struct {
	Name        string   `xml:"name,attr"`
	Type        Type     `xml:"type,attr"`
	Deprecated  Version  `xml:"deprecated,attr,omitempty"`
	Default     string   `xml:"default,attr,omitempty"`
	Optional    bool     `xml:"optional,attr,omitempty"`
	Array       bool     `xml:"array,attr,omitempty"`
	Attr        bool     `xml:"attr,attr,omitempty"` // 作为父元素的一个属性，仅针对 xml 有效果
	Items       []*Param `xml:"param,omitempty"`
	Reference   string   `xml:"ref,attr,omitempty"`
	Summary     string   `xml:"summary,attr,omitempty"`
	Enums       []*Enum  `xml:"enum,omitempty"`
	Description Richtext `xml:"description,omitempty"`
}

// IsEnum 是否为一个枚举类型
func (p *Param) IsEnum() bool {
	return len(p.Enums) > 0
}

type shadowParam Param

// UnmarshalXML xml.Unmarshaler
func (p *Param) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	field := "/" + start.Name.Local
	shadow := (*shadowParam)(p)
	if err := d.DecodeElement(shadow, &start); err != nil {
		return fixedSyntaxError(err, "", field, 0)
	}

	if shadow.Name == "" {
		return newSyntaxError(field+"/@name", locale.ErrRequired)
	}

	if shadow.Type == None {
		return newSyntaxError(field+"/@type", locale.ErrRequired)
	}
	if shadow.Type == Object && len(shadow.Items) == 0 {
		return newSyntaxError(field+"/items", locale.ErrRequired)
	}

	// 判断 enums 的值是否相同
	if key := getDuplicateEnum(shadow.Enums); key != "" {
		return newSyntaxError(field+"/enum", locale.ErrDuplicateValue)
	}

	if err := chkEnumsType(shadow.Type, shadow.Enums, field); err != nil {
		return err
	}

	// 判断 items 的值是否相同
	if key := getDuplicateItems(shadow.Items); key != "" {
		return newSyntaxError(field+"/items", locale.ErrDuplicateValue)
	}

	if p.Summary == "" && p.Description.Text == "" {
		return newSyntaxError(field+"/summary", locale.ErrRequired)
	}

	return nil
}

// 检测 enums 中的类型是否符合 t 的标准，比如 Number 要求枚举值也都是数值
func chkEnumsType(t Type, enums []*Enum, field string) error {
	if len(enums) == 0 {
		return nil
	}

	switch t {
	case Number:
		for _, enum := range enums {
			if !is.Number(enum.Value) {
				return newSyntaxError(field+"/enum/@"+enum.Value, locale.ErrInvalidFormat)
			}
		}
	case Bool:
		for _, enum := range enums {
			if _, err := strconv.ParseBool(enum.Value); err != nil {
				return newSyntaxError(field+"/enum/@"+enum.Value, locale.ErrInvalidFormat)
			}
		}
	case Object, None:
		return newSyntaxError(field+"/enum", locale.ErrInvalidValue)
	}

	return nil
}

// 返回重复枚举的值
func getDuplicateEnum(enums []*Enum) string {
	if len(enums) == 0 {
		return ""
	}

	es := make([]string, 0, len(enums))
	for _, e := range enums {
		es = append(es, e.Value)
	}
	sort.SliceStable(es, func(i, j int) bool { return es[i] > es[j] })

	for i := 1; i < len(es); i++ {
		if es[i] == es[i-1] {
			return es[i]
		}
	}

	return ""
}

func getDuplicateItems(items []*Param) string {
	if len(items) == 0 {
		return ""
	}

	es := make([]string, 0, len(items))
	for _, e := range items {
		es = append(es, e.Name)
	}
	sort.SliceStable(es, func(i, j int) bool { return es[i] > es[j] })

	for i := 1; i < len(es); i++ {
		if es[i] == es[i-1] {
			return es[i]
		}
	}

	return ""
}
<|endoftext|>


<commit_before>package logging

import (
	"expvar"
	"net/http"
	"strings"
	"time"

	"github.com/paulbellamy/ratecounter"
	"github.com/rs/zerolog"
)

var (
	proxyCounter *expvar.Map
	rps          *ratecounter.RateCounter
	bps          *ratecounter.RateCounter
	durAvg       *ratecounter.AvgRateCounter

	bytesSecond = expvar.NewInt(bytesRate)
	reqSecond   = expvar.NewInt(requestRate)
	avgDuration = expvar.NewFloat(avgDur)
)

func init() {
	proxyCounter = expvar.NewMap("proxyCounter")
	rps = ratecounter.NewRateCounter(time.Second)
	bps = ratecounter.NewRateCounter(time.Second)
	durAvg = ratecounter.NewAvgRateCounter(time.Minute)

}

const (
	// TimeFormat is the time format for logging.
	TimeFormat = time.RFC3339Nano

	avgDur      = "duration_1m_avg"
	bytesRate   = "bytes_second"
	requestRate = "requests_second"

	err400    = "400"
	err500    = "500"
	requests  = "requests"
	completed = "completed"
	bytes     = "bytes_transferred"
)

// byteCounter implements an io.Writer wrapping the http.ResponseWriter to
// count bytes written in the response.
type byteCounter struct {
	http.ResponseWriter

	status        int
	responseBytes int64
}

// Write writes bytes to the response writer.
func (ar *byteCounter) Write(b []byte) (int, error) {
	written, err := ar.ResponseWriter.Write(b)
	ar.responseBytes += int64(written)
	return written, err
}

// WriteHeader sets the response status.
func (ar *byteCounter) WriteHeader(status int) {
	ar.status = status
	ar.ResponseWriter.WriteHeader(status)
}

// NewAccessLogger returns a constructed AccessLogger pointer.
func NewAccessLogger(handler http.Handler, logger zerolog.Logger) http.Handler {
	logger = logger.With().Str("type", "access").Logger()
	return &AccessLogger{
		handler: handler,
		logger:  logger,
	}
}

// AccessLogger writes an NCSA combined-ish log record. Note this skips the
// rfc931 ident field and username as we aren't supporting either of those
// AFAIU. It also adds an elapsed time, and UTC time to Nanosecond precision
// in RFC3339 format.
type AccessLogger struct {
	handler http.Handler
	logger  zerolog.Logger
}

// ServeHTTP makes our type a http.HandlerFunc.
func (al *AccessLogger) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	proxyCounter.Add(requests, 1)
	clientIP := r.RemoteAddr
	if colon := strings.LastIndex(clientIP, ":"); colon != -1 {
		clientIP = clientIP[:colon]
	}
	bc := &byteCounter{
		ResponseWriter: w,
		status:         http.StatusOK,
	}
	start := time.Now()
	al.handler.ServeHTTP(bc, r)
	dur := time.Since(start)
	al.logger.Info().
		Str("client_ip", clientIP).
		Dur("duration", dur).
		Str("domain", r.Host).
		Str("method", r.Method).
		Str("uri", r.RequestURI).
		Str("protocol", r.Proto).
		Int("status", bc.status).
		Int64("reponse_bytes", bc.responseBytes).
		Str("referrer", r.Referer()).
		Str("user_agent", r.UserAgent()).Msg("")

	switch {
	case bc.status == 404:
		proxyCounter.Add("404", 1)
	case bc.status > 399 && bc.status < 500:
		proxyCounter.Add(err400, 1)
	case bc.status > 499:
		proxyCounter.Add(err500, 1)
	}

	proxyCounter.Add(bytes, bc.responseBytes)
	proxyCounter.Add(completed, 1)
	rps.Incr(1)
	bps.Incr(bc.responseBytes)
	durAvg.Incr(dur.Nanoseconds())

	reqSecond.Set(rps.Rate())
	bytesSecond.Set(bps.Rate())
	avgDuration.Set(durAvg.Rate())
}
<commit_msg>Use a variable rather than a magic string<commit_after>package logging

import (
	"expvar"
	"net/http"
	"strings"
	"time"

	"github.com/paulbellamy/ratecounter"
	"github.com/rs/zerolog"
)

var (
	proxyCounter *expvar.Map
	rps          *ratecounter.RateCounter
	bps          *ratecounter.RateCounter
	durAvg       *ratecounter.AvgRateCounter

	bytesSecond = expvar.NewInt(bytesRate)
	reqSecond   = expvar.NewInt(requestRate)
	avgDuration = expvar.NewFloat(avgDur)
)

func init() {
	proxyCounter = expvar.NewMap("proxyCounter")
	rps = ratecounter.NewRateCounter(time.Second)
	bps = ratecounter.NewRateCounter(time.Second)
	durAvg = ratecounter.NewAvgRateCounter(time.Minute)

}

const (
	// TimeFormat is the time format for logging.
	TimeFormat = time.RFC3339Nano

	avgDur      = "duration_1m_avg"
	bytesRate   = "bytes_second"
	requestRate = "requests_second"

	err400    = "400"
	err404    = "404"
	err500    = "500"
	requests  = "requests"
	completed = "completed"
	bytes     = "bytes_transferred"
)

// byteCounter implements an io.Writer wrapping the http.ResponseWriter to
// count bytes written in the response.
type byteCounter struct {
	http.ResponseWriter

	status        int
	responseBytes int64
}

// Write writes bytes to the response writer.
func (ar *byteCounter) Write(b []byte) (int, error) {
	written, err := ar.ResponseWriter.Write(b)
	ar.responseBytes += int64(written)
	return written, err
}

// WriteHeader sets the response status.
func (ar *byteCounter) WriteHeader(status int) {
	ar.status = status
	ar.ResponseWriter.WriteHeader(status)
}

// NewAccessLogger returns a constructed AccessLogger pointer.
func NewAccessLogger(handler http.Handler, logger zerolog.Logger) http.Handler {
	logger = logger.With().Str("type", "access").Logger()
	return &AccessLogger{
		handler: handler,
		logger:  logger,
	}
}

// AccessLogger writes an NCSA combined-ish log record. Note this skips the
// rfc931 ident field and username as we aren't supporting either of those
// AFAIU. It also adds an elapsed time, and UTC time to Nanosecond precision
// in RFC3339 format.
type AccessLogger struct {
	handler http.Handler
	logger  zerolog.Logger
}

// ServeHTTP makes our type a http.HandlerFunc.
func (al *AccessLogger) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	proxyCounter.Add(requests, 1)
	clientIP := r.RemoteAddr
	if colon := strings.LastIndex(clientIP, ":"); colon != -1 {
		clientIP = clientIP[:colon]
	}
	bc := &byteCounter{
		ResponseWriter: w,
		status:         http.StatusOK,
	}
	start := time.Now()
	al.handler.ServeHTTP(bc, r)
	dur := time.Since(start)
	al.logger.Info().
		Str("client_ip", clientIP).
		Dur("duration", dur).
		Str("domain", r.Host).
		Str("method", r.Method).
		Str("uri", r.RequestURI).
		Str("protocol", r.Proto).
		Int("status", bc.status).
		Int64("reponse_bytes", bc.responseBytes).
		Str("referrer", r.Referer()).
		Str("user_agent", r.UserAgent()).Msg("")

	switch {
	case bc.status == 404:
		proxyCounter.Add(err404, 1)
	case bc.status > 399 && bc.status < 500:
		proxyCounter.Add(err400, 1)
	case bc.status > 499:
		proxyCounter.Add(err500, 1)
	}

	proxyCounter.Add(bytes, bc.responseBytes)
	proxyCounter.Add(completed, 1)
	rps.Incr(1)
	bps.Incr(bc.responseBytes)
	durAvg.Incr(dur.Nanoseconds())

	reqSecond.Set(rps.Rate())
	bytesSecond.Set(bps.Rate())
	avgDuration.Set(durAvg.Rate())
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

package table

import (
	"testing"
	"bytes"
	"leveldb"
)

var cases = []string{
	"a",
	"aa",
	"aaa",
	"aab",
	"aac",
	"aacccc",
	"aacddddddd",
	"aacddddddde",
	"ba",
	"bb",
	"bc",
	"bcc",
	"bcd",
	"bcde",
	"ccde",
	"ccdef",
	"ccdefffffffuuuu",
	"f",
	"fz",
	"zzzz",
}

func TestBlock(t *testing.T) {
	o := &leveldb.Options{
		BlockRestartInterval: 3,
	}
	bb := newBlockBuilder(o, 0)
	for _, v := range cases {
		bb.Add([]byte(v), []byte(v))
	}
	ll := bb.Len()
	out := bb.Finish()
	if ll != len(out) {
		t.Error("guessed len doesn't equal with output len, ", ll, "!=", len(out))
	}
	br, err := newBlock(out)
	if err != nil {
		t.Error(err)
	}

	iter := br.NewIterator(o.GetComparator())

	for i := 1; i < len(cases)+1; i++ {
		for j := 0; j < i; j++ {
			if !iter.Next() {
				t.Error("early eof")
			}
			if !bytes.Equal(iter.Key(), iter.Value()) {
				t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
			}
			if !bytes.Equal(iter.Key(), []byte(cases[j])) {
				t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j]))
			}
		}

		for j := i; j > 0; j-- {
			if !iter.Prev() && j > 1 {
				t.Error("early eof")
			}
			if j > 1 {
				if !bytes.Equal(iter.Key(), iter.Value()) {
					t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
				}
				if !bytes.Equal(iter.Key(), []byte(cases[j-2])) {
					t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j-2]))
				}
			}
		}
	}

	if !iter.Last() {
		t.Error("early eof")
	}
	if iter.Next() {
		t.Error("should reaching eof")
	}
	for i := 1; i < len(cases)+1; i++ {
		nc := cases[len(cases)-i:]
		for j := i; j > 0; j-- {
			if !iter.Prev() {
				t.Error("early eof")
			}
			if !bytes.Equal(iter.Key(), iter.Value()) {
				t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
			}
			if !bytes.Equal(iter.Key(), []byte(nc[j-1])) {
				t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j-2]))
			}
		}

		for j := 0; j < i; j++ {
			if !iter.Next() && j < i-1 {
				t.Error("early eof")
			}
			if j < i-1 {
				if !bytes.Equal(iter.Key(), iter.Value()) {
					t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
				}
				if !bytes.Equal(iter.Key(), []byte(nc[j+1])) {
					t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j]))
				}
			}
		}
	}

	for i := 0; i < len(cases); i++ {
		if !iter.Seek([]byte(cases[i])) {
			t.Errorf("key '%s' is not found", cases[i])
		}
		if !bytes.Equal(iter.Key(), iter.Value()) {
			t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
		}
		if !bytes.Equal(iter.Key(), []byte(cases[i])) {
			t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[i]))
		}
	}
}
<commit_msg>table: block_test: better error reporting<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

package table

import (
	"testing"
	"bytes"
	"leveldb"
)

var cases = []string{
	"a",
	"aa",
	"aaa",
	"aab",
	"aac",
	"aacccc",
	"aacddddddd",
	"aacddddddde",
	"ba",
	"bb",
	"bc",
	"bcc",
	"bcd",
	"bcde",
	"ccde",
	"ccdef",
	"ccdefffffffuuuu",
	"f",
	"fz",
	"zzzz",
}

func TestBlock(t *testing.T) {
	o := &leveldb.Options{
		BlockRestartInterval: 3,
	}
	bb := newBlockBuilder(o, 0)
	for _, v := range cases {
		bb.Add([]byte(v), []byte(v))
	}
	ll := bb.Len()
	out := bb.Finish()
	if ll != len(out) {
		t.Error("guessed len doesn't equal with output len, ", ll, "!=", len(out))
	}
	br, err := newBlock(out)
	if err != nil {
		t.Error(err)
	}

	iter := br.NewIterator(o.GetComparator())

	for i := 1; i < len(cases)+1; i++ {
		for j := 0; j < i; j++ {
			if !iter.Next() {
				t.Error("early eof, err: '%v'", iter.Error())
			}
			if !bytes.Equal(iter.Key(), iter.Value()) {
				t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
			}
			if !bytes.Equal(iter.Key(), []byte(cases[j])) {
				t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j]))
			}
		}

		for j := i; j > 0; j-- {
			if !iter.Prev() && j > 1 {
				t.Error("early eof, err: '%v'", iter.Error())
			}
			if j > 1 {
				if !bytes.Equal(iter.Key(), iter.Value()) {
					t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
				}
				if !bytes.Equal(iter.Key(), []byte(cases[j-2])) {
					t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j-2]))
				}
			}
		}
	}

	if !iter.Last() {
		t.Error("early eof")
	}
	if iter.Next() {
		t.Error("should reaching eof")
	}
	for i := 1; i < len(cases)+1; i++ {
		nc := cases[len(cases)-i:]
		for j := i; j > 0; j-- {
			if !iter.Prev() {
				t.Errorf("early eof, err: '%v'", iter.Error())
			}
			if !bytes.Equal(iter.Key(), iter.Value()) {
				t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
			}
			if !bytes.Equal(iter.Key(), []byte(nc[j-1])) {
				t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j-2]))
			}
		}

		for j := 0; j < i; j++ {
			if !iter.Next() && j < i-1 {
				t.Errorf("early eof, err: '%v'", iter.Error())
			}
			if j < i-1 {
				if !bytes.Equal(iter.Key(), iter.Value()) {
					t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
				}
				if !bytes.Equal(iter.Key(), []byte(nc[j+1])) {
					t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[j]))
				}
			}
		}
	}

	for i := 0; i < len(cases); i++ {
		if !iter.Seek([]byte(cases[i])) {
			t.Errorf("key '%s' is not found, err: '%v'", cases[i], iter.Error())
		}
		if !bytes.Equal(iter.Key(), iter.Value()) {
			t.Error("key and value are not equal, ", iter.Key(), "!=", iter.Value())
		}
		if !bytes.Equal(iter.Key(), []byte(cases[i])) {
			t.Error("key and value are invalid, ", iter.Key(), "!=", []byte(cases[i]))
		}
	}
}
<|endoftext|>