<commit_before>package runner

import (
	"fmt"
	"html/template"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	"github.com/coreos/go-systemd/dbus"
	"github.com/dustin/go-humanize"
	"github.com/shirou/gopsutil/process"
	"github.com/sirupsen/logrus"
	"github.com/stampzilla/stampzilla-go/stampzilla/installer"
)

type Systemd struct {
	dbusConn *dbus.Conn
}

func (sd *Systemd) Close() {
	if sd.dbusConn != nil {
		sd.dbusConn.Close()
	}
}

func (sd *Systemd) dbus() *dbus.Conn {
	if sd.dbusConn != nil {
		return sd.dbusConn
	}

	var err error
	sd.dbusConn, err = dbus.NewSystemdConnection()
	if err != nil {
		logrus.Fatal("systemd init error: ", err)
	}

	return sd.dbusConn
}

func (sd *Systemd) Start(nodes ...string) error {

	conn := sd.dbus()

	if len(nodes) > 0 {
		for _, name := range nodes {
			name = getUnitName(name)
			logrus.Info("Starting ", name)
			ch := make(chan string)
			_, err := conn.StartUnit(name, "replace", ch)
			if err != nil {
				return err
			}
			<-ch
		}
		return nil
	}

	units, err := conn.ListUnitsByPatterns([]string{"inactive", "failed", "deactivating"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}
	for _, u := range units {
		logrus.Info("Starting ", u.Name)
		ch := make(chan string)
		_, err := conn.StartUnit(u.Name, "replace", ch)
		if err != nil {
			return err
		}
		<-ch
	}
	return nil
}

func (sd *Systemd) Stop(nodes ...string) error {
	conn := sd.dbus()
	if len(nodes) > 0 {
		for _, node := range nodes {
			name := getUnitName(node)
			logrus.Info("Stopping ", name)
			ch := make(chan string)
			_, err := conn.StopUnit(name, "replace", ch)
			if err != nil {
				return err
			}
			<-ch
		}
		return nil
	}

	units, err := conn.ListUnitsByPatterns([]string{"active", "activating", "failed"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}

	//stop all running stampzilla processes
	for _, p := range units {
		logrus.Info("Stopping ", p.Name)
		ch := make(chan string)
		_, err := conn.StopUnit(p.Name, "replace", ch)
		if err != nil {
			return err
		}
		<-ch
	}
	return nil
}

// Restart restarts currencly running nodes
func (sd *Systemd) Restart(nodes ...string) error {
	conn := sd.dbus()
	units, err := conn.ListUnitsByPatterns([]string{"active"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}
	ch := make(chan string)
	for _, p := range units {
		logrus.Info("Stopping ", p.Name)
		_, err := conn.StopUnit(p.Name, "replace", ch)
		if err != nil {
			return err
		}
		<-ch

		logrus.Info("Starting ", p.Name)
		_, err = conn.StartUnit(p.Name, "replace", ch)
		if err != nil {
			return err
		}
		<-ch
	}
	return nil
}

func (sd *Systemd) Status() error {
	conn := sd.dbus()
	units, err := conn.ListUnitsByPatterns([]string{"active", "activating", "failed"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 8, 0, '\t', 0)
	fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", "Name", "Active", "Pid", "CPU%", "Memory%", "Memory")
	for _, v := range units {
		prop, err := conn.GetServiceProperty(v.Name, "MainPID")
		if err != nil {
			logrus.Error(err)
			continue
		}
		pid := prop.Value.Value().(uint32)

		cpu, memPercent, mem := sd.getStats(pid)
		fmt.Fprintf(w, "%s\t%s\t%d\t%.2f%%\t%.2f%%\t%s\n", v.Name, v.ActiveState, pid, cpu, memPercent, humanize.Bytes(mem))
	}
	return w.Flush()
}

func (sd *Systemd) getStats(pid uint32) (cpu float64, memPercent float32, mem uint64) {
	if pid == 0 {
		return
	}

	p, err := process.NewProcess(int32(pid))
	if err != nil {
		logrus.Error(err)
	}
	cpu, err = p.CPUPercent()
	if err != nil {
		logrus.Error(err)
	}
	memPercent, err = p.MemoryPercent()
	if err != nil {
		logrus.Error(err)
	}
	m, err := p.MemoryInfo()
	if err != nil {
		logrus.Error(err)
	}
	mem = m.RSS
	return
}

func (sd *Systemd) GenerateUnit(name string) error {
	u := `[Unit]
Description={{.Name}}

[Service]
Type=simple

User=stampzilla
Group=stampzilla

Restart=always
RestartSec=5s

ExecStart=/home/stampzilla/go/bin/{{.Name}}
WorkingDirectory=/etc/stampzilla/nodes/{{.Name}}

[Install]
WantedBy=multi-user.target
`

	type td struct {
		Name string
	}
	tmplData := td{
		Name: GetProcessName(name),
	}
	t, err := template.New("t1").Parse(u)
	if err != nil {
		return err
	}

	unitName := getUnitName(name)
	installer.CreateDirAsUser(filepath.Join("/", "etc", "stampzilla", "nodes", strings.TrimSuffix(unitName, filepath.Ext(unitName))), "stampzilla")
	p := filepath.Join("/", "etc", "systemd", "system", unitName)
	f, err := os.Create(p)
	if err != nil {
		return fmt.Errorf("create file: %s", err)
	}
	defer f.Close()
	logrus.Infof("Generating %s", p)
	err = t.Execute(f, tmplData)
	if err != nil {
		return err
	}

	conn := sd.dbus()

	err = conn.Reload()
	if err != nil {
		return err
	}

	_, _, err = conn.EnableUnitFiles([]string{unitName}, false, false)
	return err
}

func (sd *Systemd) Disable(names ...string) error {
	conn := sd.dbus()
	ns := []string{}

	for _, name := range names {
		un := getUnitName(name)
		logrus.Infof("Disabling %s", un)
		ns = append(ns, un)
	}
	_, err := conn.DisableUnitFiles(ns, false)
	return err
}

func getUnitName(s string) string {
	return GetProcessName(s) + ".service"
}
<commit_msg>Restart specific nodes<commit_after>package runner

import (
	"fmt"
	"html/template"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	"github.com/coreos/go-systemd/dbus"
	"github.com/dustin/go-humanize"
	"github.com/shirou/gopsutil/process"
	"github.com/sirupsen/logrus"
	"github.com/stampzilla/stampzilla-go/stampzilla/installer"
)

type Systemd struct {
	dbusConn *dbus.Conn
}

func (sd *Systemd) Close() {
	if sd.dbusConn != nil {
		sd.dbusConn.Close()
	}
}

func (sd *Systemd) dbus() *dbus.Conn {
	if sd.dbusConn != nil {
		return sd.dbusConn
	}

	var err error
	sd.dbusConn, err = dbus.NewSystemdConnection()
	if err != nil {
		logrus.Fatal("systemd init error: ", err)
	}

	return sd.dbusConn
}

func (sd *Systemd) Start(nodes ...string) error {

	conn := sd.dbus()

	if len(nodes) > 0 {
		for _, name := range nodes {
			name = getUnitName(name)
			logrus.Info("Starting ", name)
			ch := make(chan string)
			_, err := conn.StartUnit(name, "replace", ch)
			if err != nil {
				return err
			}
			<-ch
		}
		return nil
	}

	units, err := conn.ListUnitsByPatterns([]string{"inactive", "failed", "deactivating"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}
	for _, u := range units {
		logrus.Info("Starting ", u.Name)
		ch := make(chan string)
		_, err := conn.StartUnit(u.Name, "replace", ch)
		if err != nil {
			return err
		}
		<-ch
	}
	return nil
}

func (sd *Systemd) Stop(nodes ...string) error {
	conn := sd.dbus()
	if len(nodes) > 0 {
		for _, node := range nodes {
			name := getUnitName(node)
			logrus.Info("Stopping ", name)
			ch := make(chan string)
			_, err := conn.StopUnit(name, "replace", ch)
			if err != nil {
				return err
			}
			<-ch
		}
		return nil
	}

	units, err := conn.ListUnitsByPatterns([]string{"active", "activating", "failed"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}

	//stop all running stampzilla processes
	for _, p := range units {
		logrus.Info("Stopping ", p.Name)
		ch := make(chan string)
		_, err := conn.StopUnit(p.Name, "replace", ch)
		if err != nil {
			return err
		}
		<-ch
	}
	return nil
}

// Restart restarts currencly running nodes
func (sd *Systemd) Restart(nodes ...string) error {
	conn := sd.dbus()
	units, err := conn.ListUnitsByPatterns([]string{"active", "activating", "failed"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}
	ch := make(chan string)

	toRestart := []string{}

	if len(nodes) > 0 {
		// restart requested nodes
		for _, unit := range units {
			for _, node := range nodes {
				if strings.HasPrefix(unit.Name, getUnitName(node)) {
					toRestart = append(toRestart, unit.Name)
				}
			}
		}
	} else {
		// restart all running
		for _, unit := range units {
			toRestart = append(toRestart, unit.Name)
		}
	}

	for _, unit := range toRestart {
		logrus.Info("Stopping ", unit)
		_, err := conn.StopUnit(unit, "replace", ch)
		if err != nil {
			return err
		}
		<-ch

		logrus.Info("Starting ", unit)
		_, err = conn.StartUnit(unit, "replace", ch)
		if err != nil {
			return err
		}
		<-ch
	}
	return nil
}

func (sd *Systemd) Status() error {
	conn := sd.dbus()
	units, err := conn.ListUnitsByPatterns([]string{"active", "activating", "failed"}, []string{"stampzilla-*"})
	if err != nil {
		return err
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 8, 0, '\t', 0)
	fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", "Name", "Active", "Pid", "CPU%", "Memory%", "Memory")
	for _, v := range units {
		prop, err := conn.GetServiceProperty(v.Name, "MainPID")
		if err != nil {
			logrus.Error(err)
			continue
		}
		pid := prop.Value.Value().(uint32)

		cpu, memPercent, mem := sd.getStats(pid)
		fmt.Fprintf(w, "%s\t%s\t%d\t%.2f%%\t%.2f%%\t%s\n", v.Name, v.ActiveState, pid, cpu, memPercent, humanize.Bytes(mem))
	}
	return w.Flush()
}

func (sd *Systemd) getStats(pid uint32) (cpu float64, memPercent float32, mem uint64) {
	if pid == 0 {
		return
	}

	p, err := process.NewProcess(int32(pid))
	if err != nil {
		logrus.Error(err)
	}
	cpu, err = p.CPUPercent()
	if err != nil {
		logrus.Error(err)
	}
	memPercent, err = p.MemoryPercent()
	if err != nil {
		logrus.Error(err)
	}
	m, err := p.MemoryInfo()
	if err != nil {
		logrus.Error(err)
	}
	mem = m.RSS
	return
}

func (sd *Systemd) GenerateUnit(name string) error {
	u := `[Unit]
Description={{.Name}}

[Service]
Type=simple

User=stampzilla
Group=stampzilla

Restart=always
RestartSec=5s

ExecStart=/home/stampzilla/go/bin/{{.Name}}
WorkingDirectory=/etc/stampzilla/nodes/{{.Name}}

[Install]
WantedBy=multi-user.target
`

	type td struct {
		Name string
	}
	tmplData := td{
		Name: GetProcessName(name),
	}
	t, err := template.New("t1").Parse(u)
	if err != nil {
		return err
	}

	unitName := getUnitName(name)
	installer.CreateDirAsUser(filepath.Join("/", "etc", "stampzilla", "nodes", strings.TrimSuffix(unitName, filepath.Ext(unitName))), "stampzilla")
	p := filepath.Join("/", "etc", "systemd", "system", unitName)
	f, err := os.Create(p)
	if err != nil {
		return fmt.Errorf("create file: %s", err)
	}
	defer f.Close()
	logrus.Infof("Generating %s", p)
	err = t.Execute(f, tmplData)
	if err != nil {
		return err
	}

	conn := sd.dbus()

	err = conn.Reload()
	if err != nil {
		return err
	}

	_, _, err = conn.EnableUnitFiles([]string{unitName}, false, false)
	return err
}

func (sd *Systemd) Disable(names ...string) error {
	conn := sd.dbus()
	ns := []string{}

	for _, name := range names {
		un := getUnitName(name)
		logrus.Infof("Disabling %s", un)
		ns = append(ns, un)
	}
	_, err := conn.DisableUnitFiles(ns, false)
	return err
}

func getUnitName(s string) string {
	return GetProcessName(s) + ".service"
}
<|endoftext|>


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

package apiserver

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/juju/loggo"
	"launchpad.net/tomb"

	"launchpad.net/juju-core/rpc"
	"launchpad.net/juju-core/rpc/jsoncodec"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/apiserver/common"
	"launchpad.net/juju-core/utils"
)

var logger = loggo.GetLogger("juju.state.apiserver")

// loginRateLimit defines how many concurrent Login requests we will
// accept
const loginRateLimit = 10

// Server holds the server side of the API.
type Server struct {
	tomb    tomb.Tomb
	wg      sync.WaitGroup
	state   *state.State
	addr    net.Addr
	dataDir string
	logDir  string
	limiter utils.Limiter
}

// Serve serves the given state by accepting requests on the given
// listener, using the given certificate and key (in PEM format) for
// authentication.
func NewServer(s *state.State, addr string, cert, key []byte, datadir, logDir string) (*Server, error) {
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	logger.Infof("listening on %q", lis.Addr())
	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return nil, err
	}
	srv := &Server{
		state:   s,
		addr:    lis.Addr(),
		dataDir: datadir,
		logDir:  logDir,
		limiter: utils.NewLimiter(loginRateLimit),
	}
	// TODO(rog) check that *srvRoot is a valid type for using
	// as an RPC server.
	lis = tls.NewListener(lis, &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
	})
	go srv.run(lis)
	return srv, nil
}

// Dead returns a channel that signals when the server has exited.
func (srv *Server) Dead() <-chan struct{} {
	return srv.tomb.Dead()
}

// Stop stops the server and returns when all running requests
// have completed.
func (srv *Server) Stop() error {
	srv.tomb.Kill(nil)
	return srv.tomb.Wait()
}

// Kill implements worker.Worker.Kill.
func (srv *Server) Kill() {
	srv.tomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (srv *Server) Wait() error {
	return srv.tomb.Wait()
}

type requestNotifier struct {
	id    int64
	start time.Time

	mu   sync.Mutex
	tag_ string
}

var globalCounter int64

func newRequestNotifier() *requestNotifier {
	return &requestNotifier{
		id:    atomic.AddInt64(&globalCounter, 1),
		tag_:  "<unknown>",
		start: time.Now(),
	}
}

func (n *requestNotifier) login(tag string) {
	n.mu.Lock()
	n.tag_ = tag
	n.mu.Unlock()
}

func (n *requestNotifier) tag() (tag string) {
	n.mu.Lock()
	tag = n.tag_
	n.mu.Unlock()
	return
}

func (n *requestNotifier) ServerRequest(hdr *rpc.Header, body interface{}) {
	if hdr.Request.Type == "Pinger" && hdr.Request.Action == "Ping" {
		return
	}
	// TODO(rog) 2013-10-11 remove secrets from some requests.
	logger.Debugf("<- [%X] %s %s", n.id, n.tag(), jsoncodec.DumpRequest(hdr, body))
}

func (n *requestNotifier) ServerReply(req rpc.Request, hdr *rpc.Header, body interface{}, timeSpent time.Duration) {
	if req.Type == "Pinger" && req.Action == "Ping" {
		return
	}
	logger.Debugf("-> [%X] %s %s %s %s[%q].%s", n.id, n.tag(), timeSpent, jsoncodec.DumpRequest(hdr, body), req.Type, req.Id, req.Action)
}

func (n *requestNotifier) join(req *http.Request) {
	logger.Infof("[%X] API connection from %s", n.id, req.RemoteAddr)
}

func (n *requestNotifier) leave() {
	logger.Infof("[%X] %s API connection terminated after %v", n.id, n.tag(), time.Since(n.start))
}

func (n requestNotifier) ClientRequest(hdr *rpc.Header, body interface{}) {
}

func (n requestNotifier) ClientReply(req rpc.Request, hdr *rpc.Header, body interface{}) {
}

func (srv *Server) run(lis net.Listener) {
	defer srv.tomb.Done()
	defer srv.wg.Wait() // wait for any outstanding requests to complete.
	srv.wg.Add(1)
	go func() {
		<-srv.tomb.Dying()
		lis.Close()
		srv.wg.Done()
	}()
	srv.wg.Add(1)
	go func() {
		logger.Infof("mongo ping interval %v", mongoPingInterval)
		err := srv.mongoPinger()
		srv.tomb.Kill(err)
		srv.wg.Done()
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("/", srv.apiHandler)
	charmHandler := &charmsHandler{httpHandler: httpHandler{state: srv.state}, dataDir: srv.dataDir}
	// charmHandler itself provides the errorSender implementation for the embedded httpHandler.
	charmHandler.httpHandler.errorSender = charmHandler
	mux.Handle("/charms", charmHandler)
	toolsHandler := &toolsHandler{httpHandler{state: srv.state}}
	// toolsHandler itself provides the errorSender implementation for the embedded httpHandler.
	toolsHandler.httpHandler.errorSender = toolsHandler
	mux.Handle("/tools", toolsHandler)
	// The error from http.Serve is not interesting.
	http.Serve(lis, mux)
}

func (srv *Server) apiHandler(w http.ResponseWriter, req *http.Request) {
	reqNotifier := newRequestNotifier()
	reqNotifier.join(req)
	defer reqNotifier.leave()
	wsServer := websocket.Server{
		Handler: func(conn *websocket.Conn) {
			srv.wg.Add(1)
			defer srv.wg.Done()
			// If we've got to this stage and the tomb is still
			// alive, we know that any tomb.Kill must occur after we
			// have called wg.Add, so we avoid the possibility of a
			// handler goroutine running after Stop has returned.
			if srv.tomb.Err() != tomb.ErrStillAlive {
				return
			}
			if err := srv.serveConn(conn, reqNotifier); err != nil {
				logger.Errorf("error serving RPCs: %v", err)
			}
		},
	}
	wsServer.ServeHTTP(w, req)
}

// Addr returns the address that the server is listening on.
func (srv *Server) Addr() string {
	return srv.addr.String()
}

func (srv *Server) serveConn(wsConn *websocket.Conn, reqNotifier *requestNotifier) error {
	codec := jsoncodec.NewWebsocket(wsConn)
	if loggo.GetLogger("juju.rpc.jsoncodec").EffectiveLogLevel() <= loggo.TRACE {
		codec.SetLogging(true)
	}
	var notifier rpc.RequestNotifier
	if logger.EffectiveLogLevel() <= loggo.DEBUG {
		// Incur request monitoring overhead only if we
		// know we'll need it.
		notifier = reqNotifier
	}
	conn := rpc.NewConn(codec, notifier)
	conn.Serve(newStateServer(srv, conn, reqNotifier, srv.limiter), serverError)
	conn.Start()
	select {
	case <-conn.Dead():
	case <-srv.tomb.Dying():
	}
	return conn.Close()
}

func (srv *Server) mongoPinger() error {
	timer := time.NewTimer(0)
	session := srv.state.MongoSession()
	for {
		select {
		case <-timer.C:
		case <-srv.tomb.Dying():
			return tomb.ErrDying
		}
		logger.Infof("pinging mongo")
		if err := session.Ping(); err != nil {
			logger.Infof("got error pinging mongo: %v", err)
			return fmt.Errorf("error pinging mongo: %v", err)
		}
		timer.Reset(mongoPingInterval)
	}
}

func serverError(err error) error {
	if err := common.ServerError(err); err != nil {
		return err
	}
	return nil
}

var logRequests = true
<commit_msg>remove debug remnants<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package apiserver

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/juju/loggo"
	"launchpad.net/tomb"

	"launchpad.net/juju-core/rpc"
	"launchpad.net/juju-core/rpc/jsoncodec"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/apiserver/common"
	"launchpad.net/juju-core/utils"
)

var logger = loggo.GetLogger("juju.state.apiserver")

// loginRateLimit defines how many concurrent Login requests we will
// accept
const loginRateLimit = 10

// Server holds the server side of the API.
type Server struct {
	tomb    tomb.Tomb
	wg      sync.WaitGroup
	state   *state.State
	addr    net.Addr
	dataDir string
	logDir  string
	limiter utils.Limiter
}

// Serve serves the given state by accepting requests on the given
// listener, using the given certificate and key (in PEM format) for
// authentication.
func NewServer(s *state.State, addr string, cert, key []byte, datadir, logDir string) (*Server, error) {
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	logger.Infof("listening on %q", lis.Addr())
	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return nil, err
	}
	srv := &Server{
		state:   s,
		addr:    lis.Addr(),
		dataDir: datadir,
		logDir:  logDir,
		limiter: utils.NewLimiter(loginRateLimit),
	}
	// TODO(rog) check that *srvRoot is a valid type for using
	// as an RPC server.
	lis = tls.NewListener(lis, &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
	})
	go srv.run(lis)
	return srv, nil
}

// Dead returns a channel that signals when the server has exited.
func (srv *Server) Dead() <-chan struct{} {
	return srv.tomb.Dead()
}

// Stop stops the server and returns when all running requests
// have completed.
func (srv *Server) Stop() error {
	srv.tomb.Kill(nil)
	return srv.tomb.Wait()
}

// Kill implements worker.Worker.Kill.
func (srv *Server) Kill() {
	srv.tomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (srv *Server) Wait() error {
	return srv.tomb.Wait()
}

type requestNotifier struct {
	id    int64
	start time.Time

	mu   sync.Mutex
	tag_ string
}

var globalCounter int64

func newRequestNotifier() *requestNotifier {
	return &requestNotifier{
		id:    atomic.AddInt64(&globalCounter, 1),
		tag_:  "<unknown>",
		start: time.Now(),
	}
}

func (n *requestNotifier) login(tag string) {
	n.mu.Lock()
	n.tag_ = tag
	n.mu.Unlock()
}

func (n *requestNotifier) tag() (tag string) {
	n.mu.Lock()
	tag = n.tag_
	n.mu.Unlock()
	return
}

func (n *requestNotifier) ServerRequest(hdr *rpc.Header, body interface{}) {
	if hdr.Request.Type == "Pinger" && hdr.Request.Action == "Ping" {
		return
	}
	// TODO(rog) 2013-10-11 remove secrets from some requests.
	logger.Debugf("<- [%X] %s %s", n.id, n.tag(), jsoncodec.DumpRequest(hdr, body))
}

func (n *requestNotifier) ServerReply(req rpc.Request, hdr *rpc.Header, body interface{}, timeSpent time.Duration) {
	if req.Type == "Pinger" && req.Action == "Ping" {
		return
	}
	logger.Debugf("-> [%X] %s %s %s %s[%q].%s", n.id, n.tag(), timeSpent, jsoncodec.DumpRequest(hdr, body), req.Type, req.Id, req.Action)
}

func (n *requestNotifier) join(req *http.Request) {
	logger.Infof("[%X] API connection from %s", n.id, req.RemoteAddr)
}

func (n *requestNotifier) leave() {
	logger.Infof("[%X] %s API connection terminated after %v", n.id, n.tag(), time.Since(n.start))
}

func (n requestNotifier) ClientRequest(hdr *rpc.Header, body interface{}) {
}

func (n requestNotifier) ClientReply(req rpc.Request, hdr *rpc.Header, body interface{}) {
}

func (srv *Server) run(lis net.Listener) {
	defer srv.tomb.Done()
	defer srv.wg.Wait() // wait for any outstanding requests to complete.
	srv.wg.Add(1)
	go func() {
		<-srv.tomb.Dying()
		lis.Close()
		srv.wg.Done()
	}()
	srv.wg.Add(1)
	go func() {
		err := srv.mongoPinger()
		srv.tomb.Kill(err)
		srv.wg.Done()
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("/", srv.apiHandler)
	charmHandler := &charmsHandler{httpHandler: httpHandler{state: srv.state}, dataDir: srv.dataDir}
	// charmHandler itself provides the errorSender implementation for the embedded httpHandler.
	charmHandler.httpHandler.errorSender = charmHandler
	mux.Handle("/charms", charmHandler)
	toolsHandler := &toolsHandler{httpHandler{state: srv.state}}
	// toolsHandler itself provides the errorSender implementation for the embedded httpHandler.
	toolsHandler.httpHandler.errorSender = toolsHandler
	mux.Handle("/tools", toolsHandler)
	// The error from http.Serve is not interesting.
	http.Serve(lis, mux)
}

func (srv *Server) apiHandler(w http.ResponseWriter, req *http.Request) {
	reqNotifier := newRequestNotifier()
	reqNotifier.join(req)
	defer reqNotifier.leave()
	wsServer := websocket.Server{
		Handler: func(conn *websocket.Conn) {
			srv.wg.Add(1)
			defer srv.wg.Done()
			// If we've got to this stage and the tomb is still
			// alive, we know that any tomb.Kill must occur after we
			// have called wg.Add, so we avoid the possibility of a
			// handler goroutine running after Stop has returned.
			if srv.tomb.Err() != tomb.ErrStillAlive {
				return
			}
			if err := srv.serveConn(conn, reqNotifier); err != nil {
				logger.Errorf("error serving RPCs: %v", err)
			}
		},
	}
	wsServer.ServeHTTP(w, req)
}

// Addr returns the address that the server is listening on.
func (srv *Server) Addr() string {
	return srv.addr.String()
}

func (srv *Server) serveConn(wsConn *websocket.Conn, reqNotifier *requestNotifier) error {
	codec := jsoncodec.NewWebsocket(wsConn)
	if loggo.GetLogger("juju.rpc.jsoncodec").EffectiveLogLevel() <= loggo.TRACE {
		codec.SetLogging(true)
	}
	var notifier rpc.RequestNotifier
	if logger.EffectiveLogLevel() <= loggo.DEBUG {
		// Incur request monitoring overhead only if we
		// know we'll need it.
		notifier = reqNotifier
	}
	conn := rpc.NewConn(codec, notifier)
	conn.Serve(newStateServer(srv, conn, reqNotifier, srv.limiter), serverError)
	conn.Start()
	select {
	case <-conn.Dead():
	case <-srv.tomb.Dying():
	}
	return conn.Close()
}

func (srv *Server) mongoPinger() error {
	timer := time.NewTimer(0)
	session := srv.state.MongoSession()
	for {
		select {
		case <-timer.C:
		case <-srv.tomb.Dying():
			return tomb.ErrDying
		}
		if err := session.Ping(); err != nil {
			logger.Infof("got error pinging mongo: %v", err)
			return fmt.Errorf("error pinging mongo: %v", err)
		}
		timer.Reset(mongoPingInterval)
	}
}

func serverError(err error) error {
	if err := common.ServerError(err); err != nil {
		return err
	}
	return nil
}

var logRequests = true
<|endoftext|>


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

package apiserver

import (
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/loggo/loggo"
	"launchpad.net/tomb"

	"launchpad.net/juju-core/rpc"
	"launchpad.net/juju-core/rpc/jsoncodec"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/apiserver/common"
)

var logger = loggo.GetLogger("juju.state.apiserver")

// Server holds the server side of the API.
type Server struct {
	tomb    tomb.Tomb
	wg      sync.WaitGroup
	state   *state.State
	addr    net.Addr
	dataDir string
}

// Serve serves the given state by accepting requests on the given
// listener, using the given certificate and key (in PEM format) for
// authentication.
func NewServer(s *state.State, addr string, cert, key []byte, datadir string) (*Server, error) {
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	logger.Infof("listening on %q", lis.Addr())
	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return nil, err
	}
	srv := &Server{
		state:   s,
		addr:    lis.Addr(),
		dataDir: datadir,
	}
	// TODO(rog) check that *srvRoot is a valid type for using
	// as an RPC server.
	lis = tls.NewListener(lis, &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
	})
	go srv.run(lis)
	return srv, nil
}

// Dead returns a channel that signals when the server has exited.
func (srv *Server) Dead() <-chan struct{} {
	return srv.tomb.Dead()
}

// Stop stops the server and returns when all running requests
// have completed.
func (srv *Server) Stop() error {
	srv.tomb.Kill(nil)
	return srv.tomb.Wait()
}

// Kill implements worker.Worker.Kill.
func (srv *Server) Kill() {
	srv.tomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (srv *Server) Wait() error {
	return srv.tomb.Wait()
}

type requestNotifier struct {
	id    int64
	start time.Time

	mu   sync.Mutex
	tag_ string
}

var globalCounter int64

func newRequestNotifier() *requestNotifier {
	return &requestNotifier{
		id:    atomic.AddInt64(&globalCounter, 1),
		tag_:  "<unknown>",
		start: time.Now(),
	}
}

func (n *requestNotifier) login(tag string) {
	n.mu.Lock()
	n.tag_ = tag
	n.mu.Unlock()
}

func (n *requestNotifier) tag() (tag string) {
	n.mu.Lock()
	tag = n.tag_
	n.mu.Unlock()
	return
}

func (n *requestNotifier) ServerRequest(hdr *rpc.Header, body interface{}) {
	if hdr.Request.Type == "Pinger" && hdr.Request.Action == "Ping" {
		return
	}
	// TODO(rog) 2013-10-11 remove secrets from some requests.
	logger.Debugf("<- [%X] %s %s", n.id, n.tag(), jsoncodec.DumpRequest(hdr, body))
}

func (n *requestNotifier) ServerReply(req rpc.Request, hdr *rpc.Header, body interface{}, timeSpent time.Duration) {
	if req.Type == "Pinger" && req.Action == "Ping" {
		return
	}
	logger.Debugf("-> [%X] %s %s %s %s[%q].%s", n.id, n.tag(), timeSpent, jsoncodec.DumpRequest(hdr, body), req.Type, req.Id, req.Action)
}

func (n *requestNotifier) join(req *http.Request) {
	logger.Infof("[%X] API connection from %s", n.id, req.RemoteAddr)
}

func (n *requestNotifier) leave() {
	logger.Infof("[%X] API connection terminated after %v", n.id, time.Since(n.start))
}

func (n requestNotifier) ClientRequest(hdr *rpc.Header, body interface{}) {
}

func (n requestNotifier) ClientReply(req rpc.Request, hdr *rpc.Header, body interface{}) {
}

func (srv *Server) run(lis net.Listener) {
	defer srv.tomb.Done()
	defer srv.wg.Wait() // wait for any outstanding requests to complete.
	srv.wg.Add(1)
	go func() {
		<-srv.tomb.Dying()
		lis.Close()
		srv.wg.Done()
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("/", srv.apiHandler)
	mux.Handle("/charms", &charmsHandler{state: srv.state})
	// The error from http.Serve is not interesting.
	http.Serve(lis, mux)
}

func (srv *Server) apiHandler(w http.ResponseWriter, req *http.Request) {
	reqNotifier := newRequestNotifier()
	reqNotifier.join(req)
	defer reqNotifier.leave()
	wsServer := websocket.Server{
		Handler: func(conn *websocket.Conn) {
			srv.wg.Add(1)
			defer srv.wg.Done()
			// If we've got to this stage and the tomb is still
			// alive, we know that any tomb.Kill must occur after we
			// have called wg.Add, so we avoid the possibility of a
			// handler goroutine running after Stop has returned.
			if srv.tomb.Err() != tomb.ErrStillAlive {
				return
			}
			if err := srv.serveConn(conn, reqNotifier); err != nil {
				logger.Errorf("error serving RPCs: %v", err)
			}
		},
	}
	wsServer.ServeHTTP(w, req)
}

// Addr returns the address that the server is listening on.
func (srv *Server) Addr() string {
	return srv.addr.String()
}

func (srv *Server) serveConn(wsConn *websocket.Conn, reqNotifier *requestNotifier) error {
	codec := jsoncodec.NewWebsocket(wsConn)
	if loggo.GetLogger("juju.rpc.jsoncodec").EffectiveLogLevel() <= loggo.TRACE {
		codec.SetLogging(true)
	}
	var notifier rpc.RequestNotifier
	if logger.EffectiveLogLevel() <= loggo.DEBUG {
		// Incur request monitoring overhead only if we
		// know we'll need it.
		notifier = reqNotifier
	}
	conn := rpc.NewConn(codec, notifier)
	conn.Serve(newStateServer(srv, conn, reqNotifier), serverError)
	conn.Start()
	select {
	case <-conn.Dead():
	case <-srv.tomb.Dying():
	}
	return conn.Close()
}

func serverError(err error) error {
	if err := common.ServerError(err); err != nil {
		return err
	}
	return nil
}

var logRequests = true
<commit_msg>if we know the tag of the connection, use it when logging a disconnect<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package apiserver

import (
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/loggo/loggo"
	"launchpad.net/tomb"

	"launchpad.net/juju-core/rpc"
	"launchpad.net/juju-core/rpc/jsoncodec"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/apiserver/common"
)

var logger = loggo.GetLogger("juju.state.apiserver")

// Server holds the server side of the API.
type Server struct {
	tomb    tomb.Tomb
	wg      sync.WaitGroup
	state   *state.State
	addr    net.Addr
	dataDir string
}

// Serve serves the given state by accepting requests on the given
// listener, using the given certificate and key (in PEM format) for
// authentication.
func NewServer(s *state.State, addr string, cert, key []byte, datadir string) (*Server, error) {
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	logger.Infof("listening on %q", lis.Addr())
	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return nil, err
	}
	srv := &Server{
		state:   s,
		addr:    lis.Addr(),
		dataDir: datadir,
	}
	// TODO(rog) check that *srvRoot is a valid type for using
	// as an RPC server.
	lis = tls.NewListener(lis, &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
	})
	go srv.run(lis)
	return srv, nil
}

// Dead returns a channel that signals when the server has exited.
func (srv *Server) Dead() <-chan struct{} {
	return srv.tomb.Dead()
}

// Stop stops the server and returns when all running requests
// have completed.
func (srv *Server) Stop() error {
	srv.tomb.Kill(nil)
	return srv.tomb.Wait()
}

// Kill implements worker.Worker.Kill.
func (srv *Server) Kill() {
	srv.tomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (srv *Server) Wait() error {
	return srv.tomb.Wait()
}

type requestNotifier struct {
	id    int64
	start time.Time

	mu   sync.Mutex
	tag_ string
}

var globalCounter int64

func newRequestNotifier() *requestNotifier {
	return &requestNotifier{
		id:    atomic.AddInt64(&globalCounter, 1),
		tag_:  "<unknown>",
		start: time.Now(),
	}
}

func (n *requestNotifier) login(tag string) {
	n.mu.Lock()
	n.tag_ = tag
	n.mu.Unlock()
}

func (n *requestNotifier) tag() (tag string) {
	n.mu.Lock()
	tag = n.tag_
	n.mu.Unlock()
	return
}

func (n *requestNotifier) ServerRequest(hdr *rpc.Header, body interface{}) {
	if hdr.Request.Type == "Pinger" && hdr.Request.Action == "Ping" {
		return
	}
	// TODO(rog) 2013-10-11 remove secrets from some requests.
	logger.Debugf("<- [%X] %s %s", n.id, n.tag(), jsoncodec.DumpRequest(hdr, body))
}

func (n *requestNotifier) ServerReply(req rpc.Request, hdr *rpc.Header, body interface{}, timeSpent time.Duration) {
	if req.Type == "Pinger" && req.Action == "Ping" {
		return
	}
	logger.Debugf("-> [%X] %s %s %s %s[%q].%s", n.id, n.tag(), timeSpent, jsoncodec.DumpRequest(hdr, body), req.Type, req.Id, req.Action)
}

func (n *requestNotifier) join(req *http.Request) {
	logger.Infof("[%X] API connection from %s", n.id, req.RemoteAddr)
}

func (n *requestNotifier) leave() {
	logger.Infof("[%X] %s API connection terminated after %v", n.id, n.tag(), time.Since(n.start))
}

func (n requestNotifier) ClientRequest(hdr *rpc.Header, body interface{}) {
}

func (n requestNotifier) ClientReply(req rpc.Request, hdr *rpc.Header, body interface{}) {
}

func (srv *Server) run(lis net.Listener) {
	defer srv.tomb.Done()
	defer srv.wg.Wait() // wait for any outstanding requests to complete.
	srv.wg.Add(1)
	go func() {
		<-srv.tomb.Dying()
		lis.Close()
		srv.wg.Done()
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("/", srv.apiHandler)
	mux.Handle("/charms", &charmsHandler{state: srv.state})
	// The error from http.Serve is not interesting.
	http.Serve(lis, mux)
}

func (srv *Server) apiHandler(w http.ResponseWriter, req *http.Request) {
	reqNotifier := newRequestNotifier()
	reqNotifier.join(req)
	defer reqNotifier.leave()
	wsServer := websocket.Server{
		Handler: func(conn *websocket.Conn) {
			srv.wg.Add(1)
			defer srv.wg.Done()
			// If we've got to this stage and the tomb is still
			// alive, we know that any tomb.Kill must occur after we
			// have called wg.Add, so we avoid the possibility of a
			// handler goroutine running after Stop has returned.
			if srv.tomb.Err() != tomb.ErrStillAlive {
				return
			}
			if err := srv.serveConn(conn, reqNotifier); err != nil {
				logger.Errorf("error serving RPCs: %v", err)
			}
		},
	}
	wsServer.ServeHTTP(w, req)
}

// Addr returns the address that the server is listening on.
func (srv *Server) Addr() string {
	return srv.addr.String()
}

func (srv *Server) serveConn(wsConn *websocket.Conn, reqNotifier *requestNotifier) error {
	codec := jsoncodec.NewWebsocket(wsConn)
	if loggo.GetLogger("juju.rpc.jsoncodec").EffectiveLogLevel() <= loggo.TRACE {
		codec.SetLogging(true)
	}
	var notifier rpc.RequestNotifier
	if logger.EffectiveLogLevel() <= loggo.DEBUG {
		// Incur request monitoring overhead only if we
		// know we'll need it.
		notifier = reqNotifier
	}
	conn := rpc.NewConn(codec, notifier)
	conn.Serve(newStateServer(srv, conn, reqNotifier), serverError)
	conn.Start()
	select {
	case <-conn.Dead():
	case <-srv.tomb.Dying():
	}
	return conn.Close()
}

func serverError(err error) error {
	if err := common.ServerError(err); err != nil {
		return err
	}
	return nil
}

var logRequests = true
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors.
// See https://code.google.com/p/go/source/browse/CONTRIBUTORS
// Licensed under the same terms as Go itself:
// https://code.google.com/p/go/source/browse/LICENSE

package http2

import (
	"bytes"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"testing"
)

var coverSpec = flag.Bool("coverspec", false, "Run spec coverage tests")

// The global map of sentence coverage for the http2 spec.
var defaultSpecCoverage specCoverage

func init() {
	f, err := os.Open("testdata/draft-ietf-httpbis-http2.xml")
	if err != nil {
		panic(err)
	}
	defaultSpecCoverage = readSpecCov(f)
}

// specCover marks all sentences for section sec in defaultSpecCoverage. Sentences not
// "covered" will be included in report outputed by TestSpecCoverage.
func specCover(sec, sentences string) {
	defaultSpecCoverage.cover(sec, sentences)
}

type specPart struct {
	section  string
	sentence string
}

func (ss specPart) Less(oo specPart) bool {
	atoi := func(s string) int {
		n, err := strconv.Atoi(s)
		if err != nil {
			panic(err)
		}
		return n
	}
	a := strings.Split(ss.section, ".")
	b := strings.Split(oo.section, ".")
	for len(a) > 0 {
		if len(b) == 0 {
			return false
		}
		x, y := atoi(a[0]), atoi(b[0])
		if x == y {
			a, b = a[1:], b[1:]
			continue
		}
		return x < y
	}
	if len(b) > 0 {
		return true
	}
	return false
}

type bySpecSection []specPart

func (a bySpecSection) Len() int           { return len(a) }
func (a bySpecSection) Less(i, j int) bool { return a[i].Less(a[j]) }
func (a bySpecSection) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

type specCoverage struct {
	coverage map[specPart]bool
	d        *xml.Decoder
}

func joinSection(sec []int) string {
	s := fmt.Sprintf("%d", sec[0])
	for _, n := range sec[1:] {
		s = fmt.Sprintf("%s.%d", s, n)
	}
	return s
}

func (sc specCoverage) readSection(sec []int) {
	var (
		buf = new(bytes.Buffer)
		sub = 0
	)
	for {
		tk, err := sc.d.Token()
		if err != nil {
			if err == io.EOF {
				return
			}
			panic(err)
		}
		switch v := tk.(type) {
		case xml.StartElement:
			if skipElement(v) {
				if err := sc.d.Skip(); err != nil {
					panic(err)
				}
				break
			}
			switch v.Name.Local {
			case "section":
				sub++
				sc.readSection(append(sec, sub))
			case "xref":
				buf.Write(sc.readXRef(v))
			}
		case xml.CharData:
			if len(sec) == 0 {
				break
			}
			buf.Write(v)
		case xml.EndElement:
			if v.Name.Local == "section" {
				sc.addSentences(joinSection(sec), buf.String())
				return
			}
		}
	}
}

func (sc specCoverage) readXRef(se xml.StartElement) []byte {
	var b []byte
	for {
		tk, err := sc.d.Token()
		if err != nil {
			panic(err)
		}
		switch v := tk.(type) {
		case xml.CharData:
			if b != nil {
				panic("unexpected CharData")
			}
			b = []byte(string(v))
		case xml.EndElement:
			if v.Name.Local != "xref" {
				panic("expected </xref>")
			}
			if b != nil {
				return b
			}
			sig := attrSig(se)
			switch sig {
			case "target":
				return []byte(fmt.Sprintf("[%s]", attrValue(se, "target")))
			case "fmt-of,rel,target", "fmt-,,rel,target":
				return []byte(fmt.Sprintf("[%s, %s]", attrValue(se, "target"), attrValue(se, "rel")))
			case "fmt-of,sec,target", "fmt-,,sec,target":
				return []byte(fmt.Sprintf("[section %s of %s]", attrValue(se, "sec"), attrValue(se, "target")))
			case "fmt-of,rel,sec,target":
				return []byte(fmt.Sprintf("[section %s of %s, %s]", attrValue(se, "sec"), attrValue(se, "target"), attrValue(se, "rel")))
			default:
				panic(fmt.Sprintf("unknown attribute signature %q in %#v", sig, fmt.Sprintf("%#v", se)))
			}
		default:
			panic(fmt.Sprintf("unexpected tag %q", v))
		}
	}
}

var skipAnchor = map[string]bool{
	"intro":    true,
	"Overview": true,
}

var skipTitle = map[string]bool{
	"Acknowledgements":            true,
	"Change Log":                  true,
	"Document Organization":       true,
	"Conventions and Terminology": true,
}

func skipElement(s xml.StartElement) bool {
	switch s.Name.Local {
	case "artwork":
		return true
	case "section":
		for _, attr := range s.Attr {
			switch attr.Name.Local {
			case "anchor":
				if skipAnchor[attr.Value] || strings.HasPrefix(attr.Value, "changes.since.") {
					return true
				}
			case "title":
				if skipTitle[attr.Value] {
					return true
				}
			}
		}
	}
	return false
}

func readSpecCov(r io.Reader) specCoverage {
	sc := specCoverage{
		coverage: map[specPart]bool{},
		d:        xml.NewDecoder(r)}
	sc.readSection(nil)
	return sc
}

func (sc specCoverage) addSentences(sec string, sentence string) {
	for _, s := range parseSentences(sentence) {
		sc.coverage[specPart{sec, s}] = false
	}
}

func (sc specCoverage) cover(sec string, sentence string) {
	for _, s := range parseSentences(sentence) {
		p := specPart{sec, s}
		if _, ok := sc.coverage[p]; !ok {
			panic(fmt.Sprintf("Not found in spec: %q, %q", sec, s))
		}
		sc.coverage[specPart{sec, s}] = true
	}

}

var whitespaceRx = regexp.MustCompile(`\s+`)

func parseSentences(sens string) []string {
	sens = strings.TrimSpace(sens)
	if sens == "" {
		return nil
	}
	ss := strings.Split(whitespaceRx.ReplaceAllString(sens, " "), ". ")
	for i, s := range ss {
		s = strings.TrimSpace(s)
		if !strings.HasSuffix(s, ".") {
			s += "."
		}
		ss[i] = s
	}
	return ss
}

func TestSpecParseSentences(t *testing.T) {
	tests := []struct {
		ss   string
		want []string
	}{
		{"Sentence 1. Sentence 2.",
			[]string{
				"Sentence 1.",
				"Sentence 2.",
			}},
		{"Sentence 1.  \nSentence 2.\tSentence 3.",
			[]string{
				"Sentence 1.",
				"Sentence 2.",
				"Sentence 3.",
			}},
	}

	for i, tt := range tests {
		got := parseSentences(tt.ss)
		if !reflect.DeepEqual(got, tt.want) {
			t.Errorf("%d: got = %q, want %q", i, got, tt.want)
		}
	}
}

func TestSpecCoverage(t *testing.T) {
	if !*coverSpec {
		t.Skip()
	}

	var (
		list     []specPart
		cv       = defaultSpecCoverage.coverage
		total    = len(cv)
		complete = 0
	)

	for sp, touched := range defaultSpecCoverage.coverage {
		if touched {
			complete++
		} else {
			list = append(list, sp)
		}
	}
	sort.Stable(bySpecSection(list))

	if testing.Short() {
		list = list[:5]
	}

	for _, p := range list {
		t.Errorf("\tSECTION %s: %s", p.section, p.sentence)
	}

	t.Logf("%d/%d (%d%%) sentances covered", complete, total, (complete/total)*100)
}

func attrSig(se xml.StartElement) string {
	var names []string
	for _, attr := range se.Attr {
		if attr.Name.Local == "fmt" {
			names = append(names, "fmt-"+attr.Value)
		} else {
			names = append(names, attr.Name.Local)
		}
	}
	sort.Strings(names)
	return strings.Join(names, ",")
}

func attrValue(se xml.StartElement, attr string) string {
	for _, a := range se.Attr {
		if a.Name.Local == attr {
			return a.Value
		}
	}
	panic("unknown attribute " + attr)
}

func TestSpecPartLess(t *testing.T) {
	tests := []struct {
		sec1, sec2 string
		want       bool
	}{
		{"6.2.1", "6.2", false},
		{"6.2", "6.2.1", true},
		{"6.10", "6.10.1", true},
		{"6.10", "6.1.1", false}, // 10, not 1
		{"6.1", "6.1", false},    // equal, so not less
	}
	for _, tt := range tests {
		got := (specPart{tt.sec1, "foo"}).Less(specPart{tt.sec2, "foo"})
		if got != tt.want {
			t.Errorf("Less(%q, %q) = %v; want %v", tt.sec1, tt.sec2, got, tt.want)
		}
	}
}
<commit_msg>spec coverage: specCover -> covers<commit_after>// Copyright 2014 The Go Authors.
// See https://code.google.com/p/go/source/browse/CONTRIBUTORS
// Licensed under the same terms as Go itself:
// https://code.google.com/p/go/source/browse/LICENSE

package http2

import (
	"bytes"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"testing"
)

var coverSpec = flag.Bool("coverspec", false, "Run spec coverage tests")

// The global map of sentence coverage for the http2 spec.
var defaultSpecCoverage specCoverage

func init() {
	f, err := os.Open("testdata/draft-ietf-httpbis-http2.xml")
	if err != nil {
		panic(err)
	}
	defaultSpecCoverage = readSpecCov(f)
}

// covers marks all sentences for section sec in defaultSpecCoverage. Sentences not
// "covered" will be included in report outputed by TestSpecCoverage.
func covers(sec, sentences string) {
	defaultSpecCoverage.cover(sec, sentences)
}

type specPart struct {
	section  string
	sentence string
}

func (ss specPart) Less(oo specPart) bool {
	atoi := func(s string) int {
		n, err := strconv.Atoi(s)
		if err != nil {
			panic(err)
		}
		return n
	}
	a := strings.Split(ss.section, ".")
	b := strings.Split(oo.section, ".")
	for len(a) > 0 {
		if len(b) == 0 {
			return false
		}
		x, y := atoi(a[0]), atoi(b[0])
		if x == y {
			a, b = a[1:], b[1:]
			continue
		}
		return x < y
	}
	if len(b) > 0 {
		return true
	}
	return false
}

type bySpecSection []specPart

func (a bySpecSection) Len() int           { return len(a) }
func (a bySpecSection) Less(i, j int) bool { return a[i].Less(a[j]) }
func (a bySpecSection) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

type specCoverage struct {
	coverage map[specPart]bool
	d        *xml.Decoder
}

func joinSection(sec []int) string {
	s := fmt.Sprintf("%d", sec[0])
	for _, n := range sec[1:] {
		s = fmt.Sprintf("%s.%d", s, n)
	}
	return s
}

func (sc specCoverage) readSection(sec []int) {
	var (
		buf = new(bytes.Buffer)
		sub = 0
	)
	for {
		tk, err := sc.d.Token()
		if err != nil {
			if err == io.EOF {
				return
			}
			panic(err)
		}
		switch v := tk.(type) {
		case xml.StartElement:
			if skipElement(v) {
				if err := sc.d.Skip(); err != nil {
					panic(err)
				}
				if v.Name.Local == "section" {
					sub++
				}
				break
			}
			switch v.Name.Local {
			case "section":
				sub++
				sc.readSection(append(sec, sub))
			case "xref":
				buf.Write(sc.readXRef(v))
			}
		case xml.CharData:
			if len(sec) == 0 {
				break
			}
			buf.Write(v)
		case xml.EndElement:
			if v.Name.Local == "section" {
				sc.addSentences(joinSection(sec), buf.String())
				return
			}
		}
	}
}

func (sc specCoverage) readXRef(se xml.StartElement) []byte {
	var b []byte
	for {
		tk, err := sc.d.Token()
		if err != nil {
			panic(err)
		}
		switch v := tk.(type) {
		case xml.CharData:
			if b != nil {
				panic("unexpected CharData")
			}
			b = []byte(string(v))
		case xml.EndElement:
			if v.Name.Local != "xref" {
				panic("expected </xref>")
			}
			if b != nil {
				return b
			}
			sig := attrSig(se)
			switch sig {
			case "target":
				return []byte(fmt.Sprintf("[%s]", attrValue(se, "target")))
			case "fmt-of,rel,target", "fmt-,,rel,target":
				return []byte(fmt.Sprintf("[%s, %s]", attrValue(se, "target"), attrValue(se, "rel")))
			case "fmt-of,sec,target", "fmt-,,sec,target":
				return []byte(fmt.Sprintf("[section %s of %s]", attrValue(se, "sec"), attrValue(se, "target")))
			case "fmt-of,rel,sec,target":
				return []byte(fmt.Sprintf("[section %s of %s, %s]", attrValue(se, "sec"), attrValue(se, "target"), attrValue(se, "rel")))
			default:
				panic(fmt.Sprintf("unknown attribute signature %q in %#v", sig, fmt.Sprintf("%#v", se)))
			}
		default:
			panic(fmt.Sprintf("unexpected tag %q", v))
		}
	}
}

var skipAnchor = map[string]bool{
	"intro":    true,
	"Overview": true,
}

var skipTitle = map[string]bool{
	"Acknowledgements":            true,
	"Change Log":                  true,
	"Document Organization":       true,
	"Conventions and Terminology": true,
}

func skipElement(s xml.StartElement) bool {
	switch s.Name.Local {
	case "artwork":
		return true
	case "section":
		for _, attr := range s.Attr {
			switch attr.Name.Local {
			case "anchor":
				if skipAnchor[attr.Value] || strings.HasPrefix(attr.Value, "changes.since.") {
					return true
				}
			case "title":
				if skipTitle[attr.Value] {
					return true
				}
			}
		}
	}
	return false
}

func readSpecCov(r io.Reader) specCoverage {
	sc := specCoverage{
		coverage: map[specPart]bool{},
		d:        xml.NewDecoder(r)}
	sc.readSection(nil)
	return sc
}

func (sc specCoverage) addSentences(sec string, sentence string) {
	for _, s := range parseSentences(sentence) {
		sc.coverage[specPart{sec, s}] = false
	}
}

func (sc specCoverage) cover(sec string, sentence string) {
	for _, s := range parseSentences(sentence) {
		p := specPart{sec, s}
		if _, ok := sc.coverage[p]; !ok {
			panic(fmt.Sprintf("Not found in spec: %q, %q", sec, s))
		}
		sc.coverage[specPart{sec, s}] = true
	}

}

var whitespaceRx = regexp.MustCompile(`\s+`)

func parseSentences(sens string) []string {
	sens = strings.TrimSpace(sens)
	if sens == "" {
		return nil
	}
	ss := strings.Split(whitespaceRx.ReplaceAllString(sens, " "), ". ")
	for i, s := range ss {
		s = strings.TrimSpace(s)
		if !strings.HasSuffix(s, ".") {
			s += "."
		}
		ss[i] = s
	}
	return ss
}

func TestSpecParseSentences(t *testing.T) {
	tests := []struct {
		ss   string
		want []string
	}{
		{"Sentence 1. Sentence 2.",
			[]string{
				"Sentence 1.",
				"Sentence 2.",
			}},
		{"Sentence 1.  \nSentence 2.\tSentence 3.",
			[]string{
				"Sentence 1.",
				"Sentence 2.",
				"Sentence 3.",
			}},
	}

	for i, tt := range tests {
		got := parseSentences(tt.ss)
		if !reflect.DeepEqual(got, tt.want) {
			t.Errorf("%d: got = %q, want %q", i, got, tt.want)
		}
	}
}

func TestSpecCoverage(t *testing.T) {
	if !*coverSpec {
		t.Skip()
	}

	var (
		list     []specPart
		cv       = defaultSpecCoverage.coverage
		total    = len(cv)
		complete = 0
	)

	for sp, touched := range defaultSpecCoverage.coverage {
		if touched {
			complete++
		} else {
			list = append(list, sp)
		}
	}
	sort.Stable(bySpecSection(list))

	if testing.Short() && len(list) > 5 {
		list = list[:5]
	}

	for _, p := range list {
		t.Errorf("\tSECTION %s: %s", p.section, p.sentence)
	}

	t.Logf("%d/%d (%d%%) sentances covered", complete, total, (complete/total)*100)
}

func attrSig(se xml.StartElement) string {
	var names []string
	for _, attr := range se.Attr {
		if attr.Name.Local == "fmt" {
			names = append(names, "fmt-"+attr.Value)
		} else {
			names = append(names, attr.Name.Local)
		}
	}
	sort.Strings(names)
	return strings.Join(names, ",")
}

func attrValue(se xml.StartElement, attr string) string {
	for _, a := range se.Attr {
		if a.Name.Local == attr {
			return a.Value
		}
	}
	panic("unknown attribute " + attr)
}

func TestSpecPartLess(t *testing.T) {
	tests := []struct {
		sec1, sec2 string
		want       bool
	}{
		{"6.2.1", "6.2", false},
		{"6.2", "6.2.1", true},
		{"6.10", "6.10.1", true},
		{"6.10", "6.1.1", false}, // 10, not 1
		{"6.1", "6.1", false},    // equal, so not less
	}
	for _, tt := range tests {
		got := (specPart{tt.sec1, "foo"}).Less(specPart{tt.sec2, "foo"})
		if got != tt.want {
			t.Errorf("Less(%q, %q) = %v; want %v", tt.sec1, tt.sec2, got, tt.want)
		}
	}
}
<|endoftext|>


<commit_before>package zenodb

import (
	"context"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/encoding"
	. "github.com/getlantern/zenodb/expr"
	"github.com/getlantern/zenodb/planner"

	"github.com/stretchr/testify/assert"
	"testing"
)

func TestRoundTime(t *testing.T) {
	ts := time.Date(2015, 5, 6, 7, 8, 9, 10, time.UTC)
	rounded := encoding.RoundTime(ts, time.Second)
	expected := time.Date(2015, 5, 6, 7, 8, 10, 0, time.UTC)
	assert.Equal(t, expected, rounded)
}

func TestSingleDB(t *testing.T) {
	doTest(t, false, func(tmpDir string, tmpFile string) (*DB, func(time.Time), func(string, func(*table))) {
		db, err := NewDB(&DBOpts{
			Dir:            filepath.Join(tmpDir, "leader"),
			SchemaFile:     tmpFile,
			VirtualTime:    true,
			MaxMemoryBytes: 1,
		})
		if !assert.NoError(t, err, "Unable to create leader DB") {
			t.Fatal()
		}
		return db, func(t time.Time) {
				db.clock.Advance(t)
			}, func(tableName string, cb func(tbl *table)) {
				cb(db.getTable(tableName))
			}
	})
}

func TestCluster(t *testing.T) {
	numPartitions := 1 + rand.Intn(10)
	numPartitions = 1

	doTest(t, true, func(tmpDir string, tmpFile string) (*DB, func(time.Time), func(string, func(*table))) {
		leader, err := NewDB(&DBOpts{
			Dir:            filepath.Join(tmpDir, "leader"),
			SchemaFile:     tmpFile,
			VirtualTime:    true,
			Passthrough:    true,
			PartitionBy:    []string{"r", "u"},
			NumPartitions:  numPartitions,
			MaxMemoryBytes: 1,
		})
		if !assert.NoError(t, err, "Unable to create leader DB") {
			t.Fatal()
		}

		followers := make([]*DB, 0, numPartitions)
		for i := 0; i < numPartitions; i++ {
			part := i
			follower, followerErr := NewDB(&DBOpts{
				Dir:            filepath.Join(tmpDir, fmt.Sprintf("follower%d", i)),
				SchemaFile:     tmpFile,
				VirtualTime:    true,
				Partition:      part,
				MaxMemoryBytes: 1,
				Follow: func(f *Follow, cb func(data []byte, newOffset wal.Offset) error) {
					leader.Follow(f, cb)
				},
				RegisterRemoteQueryHandler: func(partition int, query planner.QueryClusterFN) {
					var register func()
					register = func() {
						leader.RegisterQueryHandler(partition, func(ctx context.Context, sqlString string, isSubQuery bool, subQueryResults [][]interface{}, onRow core.OnFlatRow) error {
							// Re-register when finished
							defer register()
							return query(ctx, sqlString, isSubQuery, subQueryResults, onRow)
						})
					}

					register()
				},
			})
			if !assert.NoError(t, followerErr, "Unable to create follower DB") {
				t.Fatal()
			}

			followers = append(followers, follower)
		}

		return leader, func(t time.Time) {
				leader.clock.Advance(t)
				for _, follower := range followers {
					follower.clock.Advance(t)
				}
			}, func(tableName string, cb func(tbl *table)) {
				cb(leader.getTable(tableName))
				for _, follower := range followers {
					cb(follower.getTable(tableName))
				}
			}
	})
}

func doTest(t *testing.T, isClustered bool, buildDB func(tmpDir string, tmpFile string) (*DB, func(time.Time), func(string, func(*table)))) {
	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "zenodbtest")
	if !assert.NoError(t, err, "Unable to create temp directory") {
		return
	}
	defer os.RemoveAll(tmpDir)

	tmpFile, err := ioutil.TempFile("", "zenodbschema")
	if !assert.NoError(t, err, "Unable to create temp file") {
		return
	}
	tmpFile.Close()

	resolution := time.Millisecond

	schemaA := `
Test_a:
  maxflushlatency: 1ms
  retentionperiod: 200ms
  sql: >
    SELECT
      IF(md = 'glub', SUM(i)) AS i,
      ii,
      i * ii / COUNT(ii) AS iii,
      z
    FROM inbound
    WHERE r = 'A'
    GROUP BY *, period(1ms)
`
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaA), 0644)
	if !assert.NoError(t, err, "Unable to write schemaA") {
		return
	}

	db, advanceClock, modifyTable := buildDB(tmpDir, tmpFile.Name())

	schemaB := schemaA + `
view_a:
  view: true
  maxflushlatency: 1ms
  retentionperiod: 200ms
  sql: >
    SELECT *
    FROM teSt_a
    WHERE r = 'A'
    GROUP BY u, b`
	log.Debug("Writing schemaB")
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaB), 0644)
	if !assert.NoError(t, err, "Unable to write schemaB") {
		return
	}

	viewCreated := false
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if db.getTable("view_a") != nil {
			viewCreated = true
		}
	}
	if !assert.True(t, viewCreated, "View failed to create within 5 seconds") {
		return
	}

	now := epoch
	advance := func(d time.Duration) {
		time.Sleep(250 * time.Millisecond)
		now = now.Add(d)
		advanceClock(now)
		time.Sleep(250 * time.Millisecond)
		for _, table := range []string{"test_a", "view_a"} {
			log.Debug(db.PrintTableStats(table))
		}
	}

	// This shuffles around fields to make sure that we're reading them correctly
	// from the file stores.
	shuffleFields := func() {
		time.Sleep(100 * time.Millisecond)
		if true {
			// TODO: make sure that schema evolution works.
			return
		}
		modifyTable("test_a", func(tab *table) {
			tab.Fields[0], tab.Fields[1], tab.Fields[2] = tab.Fields[1], tab.Fields[2], tab.Fields[0]
		})
		time.Sleep(100 * time.Millisecond)
	}

	randAboveRes := func() time.Duration {
		return time.Duration(1 * rand.Intn(int(resolution)))
	}

	db.Insert("inbound",
		now.Add(randAboveRes()),
		map[string]interface{}{
			"r":  "A",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  1,
			"ii": 2,
		})
	shuffleFields()

	// This should get excluded by the filter
	db.Insert("inbound",
		now.Add(randAboveRes()),
		map[string]interface{}{
			"r":  "B",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  1,
			"ii": 2,
		})
	shuffleFields()

	db.Insert("inbound",
		now.Add(randAboveRes()),
		map[string]interface{}{
			"r":  "A",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  10,
			"ii": 20,
		})
	shuffleFields()

	// Change the schema a bit
	if false {
		// TODO: make sure that schema evolution works.
		modifyTable("test_a", func(tab *table) {
			newFields := make([]core.Field, 0, len(tab.Fields)+1)
			newFields = append(newFields, core.NewField("newfield", AVG("h")))
			for _, field := range tab.Fields {
				newFields = append(newFields, field)
			}
			tab.Fields = newFields
		})
	}

	advance(resolution)

	nextTS := now.Add(randAboveRes())
	advanceClock(nextTS)
	db.Insert("inbound",
		nextTS,
		map[string]interface{}{
			"r":  "A",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  111,
			"ii": 222,
		})
	shuffleFields()

	nextTS = now.Add(randAboveRes())
	advanceClock(nextTS)
	db.Insert("inbound",
		nextTS,
		map[string]interface{}{
			"r":  "A",
			"u":  2,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  31,
			"ii": 42,
			"z":  53,
		})
	shuffleFields()

	nextTS = now.Add(randAboveRes())
	advanceClock(nextTS)
	db.Insert("inbound",
		nextTS,
		map[string]interface{}{
			"r":  "A",
			"u":  2,
			"b":  true,
			"md": "glub",
		},
		map[string]float64{
			"i":  30000,
			"ii": 40000,
		})
	shuffleFields()

	// Give archiver time to catch up
	time.Sleep(2 * time.Second)

	testAggregateQuery(t, db, now, epoch, resolution, modifyTable)
}

func testAggregateQuery(t *testing.T, db *DB, now time.Time, epoch time.Time, resolution time.Duration, modifyTable func(string, func(*table))) {
	scalingFactor := 5

	sqlString := fmt.Sprintf(`
SELECT
	iii / 2 AS ciii,
	IF(b != true, ii) AS ii,
	*,
	IF(b = true, i) AS i_filtered,
	_points
FROM test_a
ASOF '%s' UNTIL '%s'
WHERE b != true AND r IN (SELECT r FROM test_a)
GROUP BY r, u, period(%v)
HAVING ii * 2 = 488 OR ii = 42 OR unknown = 12
ORDER BY u DESC
`, time.Duration(1-scalingFactor)*resolution, resolution, resolution*time.Duration(scalingFactor))

	var rows []*core.FlatRow
	source, err := db.Query(sqlString, false, nil, randomlyIncludeMemStore())
	if !assert.NoError(t, err, "Unable to plan SQL query") {
		return
	}
	err = source.Iterate(context.Background(), func(row *core.FlatRow) (bool, error) {
		rows = append(rows, row)
		return true, nil
	})
	if !assert.NoError(t, err, "Unable to plan SQL query") {
		return
	}

	md := MetaDataFor(source)
	log.Debugf("%v -> %v", md.AsOf, md.Until)
	log.Debug(spew.Sdump(rows))
	if !assert.Len(t, rows, 2, "Wrong number of rows, perhaps HAVING isn't working") {
		return
	}
	if !assert.EqualValues(t, 2, rows[0].Key.Get("u"), "Wrong dim, result may be sorted incorrectly") {
		return
	}
	if !assert.EqualValues(t, 1, rows[1].Key.Get("u"), "Wrong dim, result may be sorted incorrectly") {
		return
	}
	// TODO: _having shouldn't bleed through like that
	assert.Equal(t, []string{"ciii", "ii", "_points", "i", "iii", "z", "i_filtered"}, md.FieldNames)

	fieldIdx := func(name string) int {
		for i, candidate := range md.FieldNames {
			if candidate == name {
				return i
			}
		}
		return -1
	}

	pointsIdx := fieldIdx("_points")
	iIdx := fieldIdx("i")
	iiIdx := fieldIdx("ii")
	ciiiIdx := fieldIdx("ciii")
	zIdx := fieldIdx("z")
	iFilteredIdx := fieldIdx("i_filtered")

	assert.EqualValues(t, 3, rows[1].Values[pointsIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 0, rows[1].Values[iFilteredIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 122, rows[1].Values[iIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 244, rows[1].Values[iiIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(122*244)/float64(3)/float64(2), rows[1].Values[ciiiIdx], "Wrong derived value, bucketing may not be working correctly")

	assert.EqualValues(t, 1, rows[0].Values[pointsIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 0, rows[0].Values[iFilteredIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 31, rows[0].Values[iIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 42, rows[0].Values[iiIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 53, rows[0].Values[zIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(31*42)/float64(1)/float64(2), rows[0].Values[ciiiIdx], "Wrong derived value, bucketing may not be working correctly")
}

func randomlyIncludeMemStore() bool {
	return rand.Float64() > 0.5
}
<commit_msg>Fixed unit test<commit_after>package zenodb

import (
	"context"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/encoding"
	. "github.com/getlantern/zenodb/expr"
	"github.com/getlantern/zenodb/planner"

	"github.com/stretchr/testify/assert"
	"testing"
)

func TestRoundTime(t *testing.T) {
	ts := time.Date(2015, 5, 6, 7, 8, 9, 10, time.UTC)
	rounded := encoding.RoundTime(ts, time.Second)
	expected := time.Date(2015, 5, 6, 7, 8, 10, 0, time.UTC)
	assert.Equal(t, expected, rounded)
}

func TestSingleDB(t *testing.T) {
	doTest(t, false, func(tmpDir string, tmpFile string) (*DB, func(time.Time), func(string, func(*table))) {
		db, err := NewDB(&DBOpts{
			Dir:            filepath.Join(tmpDir, "leader"),
			SchemaFile:     tmpFile,
			VirtualTime:    true,
			MaxMemoryRatio: 0.00001,
		})
		if !assert.NoError(t, err, "Unable to create leader DB") {
			t.Fatal()
		}
		return db, func(t time.Time) {
				db.clock.Advance(t)
			}, func(tableName string, cb func(tbl *table)) {
				cb(db.getTable(tableName))
			}
	})
}

func TestCluster(t *testing.T) {
	numPartitions := 1 + rand.Intn(10)
	numPartitions = 1

	doTest(t, true, func(tmpDir string, tmpFile string) (*DB, func(time.Time), func(string, func(*table))) {
		leader, err := NewDB(&DBOpts{
			Dir:            filepath.Join(tmpDir, "leader"),
			SchemaFile:     tmpFile,
			VirtualTime:    true,
			Passthrough:    true,
			PartitionBy:    []string{"r", "u"},
			NumPartitions:  numPartitions,
			MaxMemoryRatio: 0.00001,
		})
		if !assert.NoError(t, err, "Unable to create leader DB") {
			t.Fatal()
		}

		followers := make([]*DB, 0, numPartitions)
		for i := 0; i < numPartitions; i++ {
			part := i
			follower, followerErr := NewDB(&DBOpts{
				Dir:            filepath.Join(tmpDir, fmt.Sprintf("follower%d", i)),
				SchemaFile:     tmpFile,
				VirtualTime:    true,
				Partition:      part,
				MaxMemoryRatio: 0.00001,
				Follow: func(f *Follow, cb func(data []byte, newOffset wal.Offset) error) {
					leader.Follow(f, cb)
				},
				RegisterRemoteQueryHandler: func(partition int, query planner.QueryClusterFN) {
					var register func()
					register = func() {
						leader.RegisterQueryHandler(partition, func(ctx context.Context, sqlString string, isSubQuery bool, subQueryResults [][]interface{}, onRow core.OnFlatRow) error {
							// Re-register when finished
							defer register()
							return query(ctx, sqlString, isSubQuery, subQueryResults, onRow)
						})
					}

					register()
				},
			})
			if !assert.NoError(t, followerErr, "Unable to create follower DB") {
				t.Fatal()
			}

			followers = append(followers, follower)
		}

		return leader, func(t time.Time) {
				leader.clock.Advance(t)
				for _, follower := range followers {
					follower.clock.Advance(t)
				}
			}, func(tableName string, cb func(tbl *table)) {
				cb(leader.getTable(tableName))
				for _, follower := range followers {
					cb(follower.getTable(tableName))
				}
			}
	})
}

func doTest(t *testing.T, isClustered bool, buildDB func(tmpDir string, tmpFile string) (*DB, func(time.Time), func(string, func(*table)))) {
	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "zenodbtest")
	if !assert.NoError(t, err, "Unable to create temp directory") {
		return
	}
	defer os.RemoveAll(tmpDir)

	tmpFile, err := ioutil.TempFile("", "zenodbschema")
	if !assert.NoError(t, err, "Unable to create temp file") {
		return
	}
	tmpFile.Close()

	resolution := time.Millisecond

	schemaA := `
Test_a:
  maxflushlatency: 1ms
  retentionperiod: 200ms
  sql: >
    SELECT
      IF(md = 'glub', SUM(i)) AS i,
      ii,
      i * ii / COUNT(ii) AS iii,
      z
    FROM inbound
    WHERE r = 'A'
    GROUP BY *, period(1ms)
`
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaA), 0644)
	if !assert.NoError(t, err, "Unable to write schemaA") {
		return
	}

	db, advanceClock, modifyTable := buildDB(tmpDir, tmpFile.Name())

	schemaB := schemaA + `
view_a:
  view: true
  maxflushlatency: 1ms
  retentionperiod: 200ms
  sql: >
    SELECT *
    FROM teSt_a
    WHERE r = 'A'
    GROUP BY u, b`
	log.Debug("Writing schemaB")
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaB), 0644)
	if !assert.NoError(t, err, "Unable to write schemaB") {
		return
	}

	viewCreated := false
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if db.getTable("view_a") != nil {
			viewCreated = true
		}
	}
	if !assert.True(t, viewCreated, "View failed to create within 5 seconds") {
		return
	}

	now := epoch
	advance := func(d time.Duration) {
		time.Sleep(250 * time.Millisecond)
		now = now.Add(d)
		advanceClock(now)
		time.Sleep(250 * time.Millisecond)
		for _, table := range []string{"test_a", "view_a"} {
			log.Debug(db.PrintTableStats(table))
		}
	}

	// This shuffles around fields to make sure that we're reading them correctly
	// from the file stores.
	shuffleFields := func() {
		time.Sleep(100 * time.Millisecond)
		if true {
			// TODO: make sure that schema evolution works.
			return
		}
		modifyTable("test_a", func(tab *table) {
			tab.Fields[0], tab.Fields[1], tab.Fields[2] = tab.Fields[1], tab.Fields[2], tab.Fields[0]
		})
		time.Sleep(100 * time.Millisecond)
	}

	randAboveRes := func() time.Duration {
		return time.Duration(1 * rand.Intn(int(resolution)))
	}

	db.Insert("inbound",
		now.Add(randAboveRes()),
		map[string]interface{}{
			"r":  "A",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  1,
			"ii": 2,
		})
	shuffleFields()

	// This should get excluded by the filter
	db.Insert("inbound",
		now.Add(randAboveRes()),
		map[string]interface{}{
			"r":  "B",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  1,
			"ii": 2,
		})
	shuffleFields()

	db.Insert("inbound",
		now.Add(randAboveRes()),
		map[string]interface{}{
			"r":  "A",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  10,
			"ii": 20,
		})
	shuffleFields()

	// Change the schema a bit
	if false {
		// TODO: make sure that schema evolution works.
		modifyTable("test_a", func(tab *table) {
			newFields := make([]core.Field, 0, len(tab.Fields)+1)
			newFields = append(newFields, core.NewField("newfield", AVG("h")))
			for _, field := range tab.Fields {
				newFields = append(newFields, field)
			}
			tab.Fields = newFields
		})
	}

	advance(resolution)

	nextTS := now.Add(randAboveRes())
	advanceClock(nextTS)
	db.Insert("inbound",
		nextTS,
		map[string]interface{}{
			"r":  "A",
			"u":  1,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  111,
			"ii": 222,
		})
	shuffleFields()

	nextTS = now.Add(randAboveRes())
	advanceClock(nextTS)
	db.Insert("inbound",
		nextTS,
		map[string]interface{}{
			"r":  "A",
			"u":  2,
			"b":  false,
			"md": "glub",
		},
		map[string]float64{
			"i":  31,
			"ii": 42,
			"z":  53,
		})
	shuffleFields()

	nextTS = now.Add(randAboveRes())
	advanceClock(nextTS)
	db.Insert("inbound",
		nextTS,
		map[string]interface{}{
			"r":  "A",
			"u":  2,
			"b":  true,
			"md": "glub",
		},
		map[string]float64{
			"i":  30000,
			"ii": 40000,
		})
	shuffleFields()

	// Give archiver time to catch up
	time.Sleep(2 * time.Second)

	testAggregateQuery(t, db, now, epoch, resolution, modifyTable)
}

func testAggregateQuery(t *testing.T, db *DB, now time.Time, epoch time.Time, resolution time.Duration, modifyTable func(string, func(*table))) {
	scalingFactor := 5

	sqlString := fmt.Sprintf(`
SELECT
	iii / 2 AS ciii,
	IF(b != true, ii) AS ii,
	*,
	IF(b = true, i) AS i_filtered,
	_points
FROM test_a
ASOF '%s' UNTIL '%s'
WHERE b != true AND r IN (SELECT r FROM test_a)
GROUP BY r, u, period(%v)
HAVING ii * 2 = 488 OR ii = 42 OR unknown = 12
ORDER BY u DESC
`, time.Duration(1-scalingFactor)*resolution, resolution, resolution*time.Duration(scalingFactor))

	var rows []*core.FlatRow
	source, err := db.Query(sqlString, false, nil, randomlyIncludeMemStore())
	if !assert.NoError(t, err, "Unable to plan SQL query") {
		return
	}
	err = source.Iterate(context.Background(), func(row *core.FlatRow) (bool, error) {
		rows = append(rows, row)
		return true, nil
	})
	if !assert.NoError(t, err, "Unable to plan SQL query") {
		return
	}

	md := MetaDataFor(source)
	log.Debugf("%v -> %v", md.AsOf, md.Until)
	log.Debug(spew.Sdump(rows))
	if !assert.Len(t, rows, 2, "Wrong number of rows, perhaps HAVING isn't working") {
		return
	}
	if !assert.EqualValues(t, 2, rows[0].Key.Get("u"), "Wrong dim, result may be sorted incorrectly") {
		return
	}
	if !assert.EqualValues(t, 1, rows[1].Key.Get("u"), "Wrong dim, result may be sorted incorrectly") {
		return
	}
	// TODO: _having shouldn't bleed through like that
	assert.Equal(t, []string{"ciii", "ii", "_points", "i", "iii", "z", "i_filtered"}, md.FieldNames)

	fieldIdx := func(name string) int {
		for i, candidate := range md.FieldNames {
			if candidate == name {
				return i
			}
		}
		return -1
	}

	pointsIdx := fieldIdx("_points")
	iIdx := fieldIdx("i")
	iiIdx := fieldIdx("ii")
	ciiiIdx := fieldIdx("ciii")
	zIdx := fieldIdx("z")
	iFilteredIdx := fieldIdx("i_filtered")

	assert.EqualValues(t, 3, rows[1].Values[pointsIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 0, rows[1].Values[iFilteredIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 122, rows[1].Values[iIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 244, rows[1].Values[iiIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(122*244)/float64(3)/float64(2), rows[1].Values[ciiiIdx], "Wrong derived value, bucketing may not be working correctly")

	assert.EqualValues(t, 1, rows[0].Values[pointsIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 0, rows[0].Values[iFilteredIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 31, rows[0].Values[iIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 42, rows[0].Values[iiIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 53, rows[0].Values[zIdx], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(31*42)/float64(1)/float64(2), rows[0].Values[ciiiIdx], "Wrong derived value, bucketing may not be working correctly")
}

func randomlyIncludeMemStore() bool {
	return rand.Float64() > 0.5
}
<|endoftext|>


<commit_before>package zenodb

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

	"github.com/Knetic/govaluate"
	"github.com/getlantern/bytemap"
	. "github.com/getlantern/zenodb/expr"
	"github.com/getlantern/zenodb/sql"

	"github.com/stretchr/testify/assert"
	"testing"
)

func TestRoundTime(t *testing.T) {
	ts := time.Date(2015, 5, 6, 7, 8, 9, 10, time.UTC)
	rounded := roundTime(ts, time.Second)
	expected := time.Date(2015, 5, 6, 7, 8, 9, 0, time.UTC)
	assert.Equal(t, expected, rounded)
}

func TestIntegration(t *testing.T) {
	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "zenodbtest")
	if !assert.NoError(t, err, "Unable to create temp directory") {
		return
	}
	defer os.RemoveAll(tmpDir)

	tmpFile, err := ioutil.TempFile("", "zenodbschema")
	if !assert.NoError(t, err, "Unable to create temp file") {
		return
	}
	tmpFile.Close()

	resolution := time.Millisecond

	schemaA := `
Test_a:
  maxmemstorebytes: 1
  retentionperiod: 200ms
  sql: >
    SELECT
      SUM(i) AS i,
      ii,
      i * ii / COUNT(ii) AS iii
    FROM inbound
    WHERE r = 'A'
    GROUP BY *, period(1ms)
`
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaA), 0644)
	if !assert.NoError(t, err, "Unable to write schemaA") {
		return
	}

	db, err := NewDB(&DBOpts{
		Dir:         tmpDir,
		SchemaFile:  tmpFile.Name(),
		VirtualTime: true,
	})
	if !assert.NoError(t, err, "Unable to create DB") {
		return
	}

	schemaB := schemaA + `
view_a:
  view: true
  maxmemstorebytes: 1
  retentionperiod: 200ms
  sql: >
    SELECT
      i,
      ii,
      iii
    FROM teSt_a
    WHERE r = 'A'
    GROUP BY u, b, period(1ms)`
	log.Debug("Writing schemaB")
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaB), 0644)
	if !assert.NoError(t, err, "Unable to write schemaB") {
		return
	}

	viewCreated := false
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if db.getTable("view_a") != nil {
			viewCreated = true
		}
	}
	if !assert.True(t, viewCreated, "View failed to create within 5 seconds") {
		return
	}

	now := epoch
	advance := func(d time.Duration) {
		time.Sleep(250 * time.Millisecond)
		now = now.Add(d)
		db.clock.Advance(now)
		time.Sleep(250 * time.Millisecond)
		for _, table := range []string{"test_a", "view_a"} {
			log.Debug(db.PrintTableStats(table))
		}
	}

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	// This should get excluded by the filter
	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "B",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  10,
			"ii": 20,
		},
	})

	advance(resolution)

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  111,
			"ii": 222,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  31,
			"ii": 42,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": true,
		},
		Vals: map[string]float64{
			"i":  30000,
			"ii": 40000,
		},
	})

	query := func(table string, from time.Time, to time.Time, dim string, field string) (map[int][]float64, error) {
		filter, queryErr := govaluate.NewEvaluableExpression("b != true")
		if queryErr != nil {
			return nil, queryErr
		}
		fromOffset := from.Sub(now)
		toOffset := to.Sub(now)
		result := make(map[int][]float64, 0)
		q := &query{
			table:       table,
			fields:      []string{field},
			asOfOffset:  fromOffset,
			untilOffset: toOffset,
			filter:      filter,
			onValues: func(keybytes bytemap.ByteMap, resultField string, e Expr, seq sequence, startOffset int) {
				key := keybytes.AsMap()
				log.Debugf("%v : %v : %v : %d : %v", key, field, resultField, startOffset, seq.String(e))
				if field == resultField {
					numPeriods := seq.numPeriods(e.EncodedWidth())
					vals := make([]float64, 0, numPeriods-startOffset)
					for i := 0; i < numPeriods-startOffset; i++ {
						val, wasSet := seq.ValueAt(i+startOffset, e)
						if wasSet {
							vals = append(vals, val)
						}
					}
					if len(vals) > 0 {
						result[key[dim].(int)] = vals
					}
				}
			},
		}
		stats, err := q.run(db)
		log.Debugf("Query stats - scanned: %d    filterpass: %d    datavalid: %d    intimerange: %d", stats.Scanned, stats.FilterPass, stats.DataValid, stats.InTimeRange)
		log.Debugf("Result: %v", result)
		return result, err
	}

	// Give archiver time to catch up
	time.Sleep(2 * time.Second)

	test := func(from time.Time, to time.Time, dim string, field string, onResult func(map[int][]float64)) {
		for _, table := range []string{"Test_A", "view_A"} {
			result, err := query(table, from, to, dim, field)
			if assert.NoError(t, err, "Unable to run query") {
				t.Logf("Table: %v", table)
				onResult(result)
			}
		}
	}

	log.Debug("A")
	test(epoch.Add(-1*resolution), epoch, "u", "i", func(result map[int][]float64) {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{11}, result[1])
		}
	})

	log.Debug("B")
	test(epoch.Add(-1*resolution), epoch, "u", "ii", func(result map[int][]float64) {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{22}, result[1])
		}
	})

	log.Debug("C")
	test(epoch.Add(-1*resolution), epoch, "u", "iii", func(result map[int][]float64) {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{121}, result[1])
		}
	})

	log.Debug("D")
	test(epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii", func(result map[int][]float64) {
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22}, result[1])
			assert.Equal(t, []float64{42}, result[2])
		}
	})

	log.Debug("E")
	test(epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii", func(result map[int][]float64) {
		log.Debug(result)
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22}, result[1])
			assert.Equal(t, []float64{42}, result[2])
		}
	})

	testAggregateQuery(t, db, now, epoch, resolution)
}

func testAggregateQuery(t *testing.T, db *DB, now time.Time, epoch time.Time, resolution time.Duration) {
	scalingFactor := 5

	aq, err := db.SQLQuery(fmt.Sprintf(`
SELECT
	iii / 2 AS ciii,
	IF(b != true, ii) AS ii,
	i,
	IF(b = true, i) AS i_filtered
FROM test_a
ASOF '%v' UNTIL '%v'
WHERE b != true
GROUP BY r, u, period('%v')
HAVING ii * 2 = 488 OR ii = 42
ORDER BY u DESC
`, epoch.Add(-1*resolution).Sub(now), epoch.Add(3*resolution).Sub(now), resolution*time.Duration(scalingFactor)))
	if !assert.NoError(t, err, "Unable to create SQL query") {
		return
	}

	result, err := aq.Run()
	if !assert.NoError(t, err, "Unable to run query") {
		return
	}

	log.Debugf("%v -> %v", result.AsOf, result.Until)
	if !assert.Equal(t, 1, result.NumPeriods, "Wrong number of periods, bucketing may not be working correctly") {
		return
	}
	rows := result.Rows
	if !assert.Len(t, rows, 2, "Wrong number of rows, perhaps HAVING isn't working") {
		return
	}
	if !assert.EqualValues(t, 2, result.Rows[0].Dims[1], "Wrong dim, result may be sorted incorrectly") {
		return
	}
	if !assert.EqualValues(t, 1, result.Rows[1].Dims[1], "Wrong dim, result may be sorted incorrectly") {
		return
	}
	fields := make([]string, 0, len(result.FieldNames))
	for _, field := range result.FieldNames {
		fields = append(fields, field)
	}
	assert.Equal(t, []string{"ciii", "ii", "i", "i_filtered"}, fields)

	assert.EqualValues(t, 0, rows[1].Values[3], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 122, rows[1].Values[2], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 244, rows[1].Values[1], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(122*244)/float64(3)/float64(2), rows[1].Values[0], "Wrong derived value, bucketing may not be working correctly")

	assert.EqualValues(t, 0, rows[0].Values[3], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 31, rows[0].Values[2], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 42, rows[0].Values[1], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(31*42)/float64(1)/float64(2), rows[0].Values[0], "Wrong derived value, bucketing may not be working correctly")

	// Test having on non-existent field
	aq, err = db.SQLQuery(fmt.Sprintf(`
SELECT i
FROM test_a
GROUP BY period('%v')
HAVING unknown = 5
`, resolution*time.Duration(scalingFactor)))
	if !assert.NoError(t, err, "Unable to creat query") {
		return
	}
	result, err = aq.Run()
	if !assert.NoError(t, err, "Unable to run query") {
		return
	}
	assert.Len(t, result.Rows, 0)

	// Test defaults
	aq = db.Query(&sql.Query{
		From:       "test_a",
		Fields:     []sql.Field{sql.Field{Expr: SUM("ii"), Name: "ii"}},
		GroupByAll: true,
		AsOfOffset: epoch.Add(-1 * resolution).Sub(now),
	})

	result, err = aq.Run()
	if assert.NoError(t, err, "Unable to run query with defaults") {
		assert.Equal(t, []string{"b", "r", "u"}, result.GroupBy)
		assert.NotNil(t, result.Until)
	}
}
<commit_msg>Integration test exercises schema evolution in files<commit_after>package zenodb

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

	"github.com/Knetic/govaluate"
	"github.com/getlantern/bytemap"
	. "github.com/getlantern/zenodb/expr"
	"github.com/getlantern/zenodb/sql"

	"github.com/stretchr/testify/assert"
	"testing"
)

func TestRoundTime(t *testing.T) {
	ts := time.Date(2015, 5, 6, 7, 8, 9, 10, time.UTC)
	rounded := roundTime(ts, time.Second)
	expected := time.Date(2015, 5, 6, 7, 8, 9, 0, time.UTC)
	assert.Equal(t, expected, rounded)
}

func TestIntegration(t *testing.T) {
	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "zenodbtest")
	if !assert.NoError(t, err, "Unable to create temp directory") {
		return
	}
	defer os.RemoveAll(tmpDir)

	tmpFile, err := ioutil.TempFile("", "zenodbschema")
	if !assert.NoError(t, err, "Unable to create temp file") {
		return
	}
	tmpFile.Close()

	resolution := time.Millisecond

	schemaA := `
Test_a:
  maxmemstorebytes: 1
  retentionperiod: 200ms
  sql: >
    SELECT
      SUM(i) AS i,
      ii,
      i * ii / COUNT(ii) AS iii
    FROM inbound
    WHERE r = 'A'
    GROUP BY *, period(1ms)
`
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaA), 0644)
	if !assert.NoError(t, err, "Unable to write schemaA") {
		return
	}

	db, err := NewDB(&DBOpts{
		Dir:         tmpDir,
		SchemaFile:  tmpFile.Name(),
		VirtualTime: true,
	})
	if !assert.NoError(t, err, "Unable to create DB") {
		return
	}

	schemaB := schemaA + `
view_a:
  view: true
  maxmemstorebytes: 1
  retentionperiod: 200ms
  sql: >
    SELECT
      i,
      ii,
      iii
    FROM teSt_a
    WHERE r = 'A'
    GROUP BY u, b, period(1ms)`
	log.Debug("Writing schemaB")
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaB), 0644)
	if !assert.NoError(t, err, "Unable to write schemaB") {
		return
	}

	viewCreated := false
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if db.getTable("view_a") != nil {
			viewCreated = true
		}
	}
	if !assert.True(t, viewCreated, "View failed to create within 5 seconds") {
		return
	}

	now := epoch
	advance := func(d time.Duration) {
		time.Sleep(250 * time.Millisecond)
		now = now.Add(d)
		db.clock.Advance(now)
		time.Sleep(250 * time.Millisecond)
		for _, table := range []string{"test_a", "view_a"} {
			log.Debug(db.PrintTableStats(table))
		}
	}

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	// This should get excluded by the filter
	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "B",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  10,
			"ii": 20,
		},
	})

	// Change the schema a bit
	tab := db.getTable("test_a")
	newFields := make([]sql.Field, 0, len(tab.Fields)+1)
	newFields = append(newFields, sql.Field{AVG("h"), "newfield"})
	for _, field := range tab.Fields {
		newFields = append(newFields, field)
	}
	tab.Fields = newFields

	advance(resolution)

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  111,
			"ii": 222,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  31,
			"ii": 42,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": true,
		},
		Vals: map[string]float64{
			"i":  30000,
			"ii": 40000,
		},
	})

	query := func(table string, from time.Time, to time.Time, dim string, field string) (map[int][]float64, error) {
		filter, queryErr := govaluate.NewEvaluableExpression("b != true")
		if queryErr != nil {
			return nil, queryErr
		}
		fromOffset := from.Sub(now)
		toOffset := to.Sub(now)
		result := make(map[int][]float64, 0)
		q := &query{
			table:       table,
			fields:      []string{field},
			asOfOffset:  fromOffset,
			untilOffset: toOffset,
			filter:      filter,
			onValues: func(keybytes bytemap.ByteMap, resultField string, e Expr, seq sequence, startOffset int) {
				key := keybytes.AsMap()
				log.Debugf("%v : %v : %v : %d : %v", key, field, resultField, startOffset, seq.String(e))
				if field == resultField {
					numPeriods := seq.numPeriods(e.EncodedWidth())
					vals := make([]float64, 0, numPeriods-startOffset)
					for i := 0; i < numPeriods-startOffset; i++ {
						val, wasSet := seq.ValueAt(i+startOffset, e)
						if wasSet {
							vals = append(vals, val)
						}
					}
					if len(vals) > 0 {
						result[key[dim].(int)] = vals
					}
				}
			},
		}
		stats, err := q.run(db)
		log.Debugf("Query stats - scanned: %d    filterpass: %d    datavalid: %d    intimerange: %d", stats.Scanned, stats.FilterPass, stats.DataValid, stats.InTimeRange)
		log.Debugf("Result: %v", result)
		return result, err
	}

	// Give archiver time to catch up
	time.Sleep(2 * time.Second)

	test := func(from time.Time, to time.Time, dim string, field string, onResult func(map[int][]float64)) {
		for _, table := range []string{"Test_A", "view_A"} {
			result, err := query(table, from, to, dim, field)
			if assert.NoError(t, err, "Unable to run query") {
				t.Logf("Table: %v", table)
				onResult(result)
			}
		}
	}

	log.Debug("A")
	test(epoch.Add(-1*resolution), epoch, "u", "i", func(result map[int][]float64) {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{11}, result[1])
		}
	})

	log.Debug("B")
	test(epoch.Add(-1*resolution), epoch, "u", "ii", func(result map[int][]float64) {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{22}, result[1])
		}
	})

	log.Debug("C")
	test(epoch.Add(-1*resolution), epoch, "u", "iii", func(result map[int][]float64) {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{121}, result[1])
		}
	})

	log.Debug("D")
	test(epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii", func(result map[int][]float64) {
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22}, result[1])
			assert.Equal(t, []float64{42}, result[2])
		}
	})

	log.Debug("E")
	test(epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii", func(result map[int][]float64) {
		log.Debug(result)
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22}, result[1])
			assert.Equal(t, []float64{42}, result[2])
		}
	})

	testAggregateQuery(t, db, now, epoch, resolution)
}

func testAggregateQuery(t *testing.T, db *DB, now time.Time, epoch time.Time, resolution time.Duration) {
	scalingFactor := 5

	aq, err := db.SQLQuery(fmt.Sprintf(`
SELECT
	iii / 2 AS ciii,
	IF(b != true, ii) AS ii,
	i,
	IF(b = true, i) AS i_filtered
FROM test_a
ASOF '%v' UNTIL '%v'
WHERE b != true
GROUP BY r, u, period('%v')
HAVING ii * 2 = 488 OR ii = 42
ORDER BY u DESC
`, epoch.Add(-1*resolution).Sub(now), epoch.Add(3*resolution).Sub(now), resolution*time.Duration(scalingFactor)))
	if !assert.NoError(t, err, "Unable to create SQL query") {
		return
	}

	result, err := aq.Run()
	if !assert.NoError(t, err, "Unable to run query") {
		return
	}

	log.Debugf("%v -> %v", result.AsOf, result.Until)
	if !assert.Equal(t, 1, result.NumPeriods, "Wrong number of periods, bucketing may not be working correctly") {
		return
	}
	rows := result.Rows
	if !assert.Len(t, rows, 2, "Wrong number of rows, perhaps HAVING isn't working") {
		return
	}
	if !assert.EqualValues(t, 2, result.Rows[0].Dims[1], "Wrong dim, result may be sorted incorrectly") {
		return
	}
	if !assert.EqualValues(t, 1, result.Rows[1].Dims[1], "Wrong dim, result may be sorted incorrectly") {
		return
	}
	fields := make([]string, 0, len(result.FieldNames))
	for _, field := range result.FieldNames {
		fields = append(fields, field)
	}
	assert.Equal(t, []string{"ciii", "ii", "i", "i_filtered"}, fields)

	assert.EqualValues(t, 0, rows[1].Values[3], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 122, rows[1].Values[2], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 244, rows[1].Values[1], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(122*244)/float64(3)/float64(2), rows[1].Values[0], "Wrong derived value, bucketing may not be working correctly")

	assert.EqualValues(t, 0, rows[0].Values[3], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 31, rows[0].Values[2], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, 42, rows[0].Values[1], "Wrong derived value, bucketing may not be working correctly")
	assert.EqualValues(t, float64(31*42)/float64(1)/float64(2), rows[0].Values[0], "Wrong derived value, bucketing may not be working correctly")

	// Test having on non-existent field
	aq, err = db.SQLQuery(fmt.Sprintf(`
SELECT i
FROM test_a
GROUP BY period('%v')
HAVING unknown = 5
`, resolution*time.Duration(scalingFactor)))
	if !assert.NoError(t, err, "Unable to creat query") {
		return
	}
	result, err = aq.Run()
	if !assert.NoError(t, err, "Unable to run query") {
		return
	}
	assert.Len(t, result.Rows, 0)

	// Test defaults
	aq = db.Query(&sql.Query{
		From:       "test_a",
		Fields:     []sql.Field{sql.Field{Expr: SUM("ii"), Name: "ii"}},
		GroupByAll: true,
		AsOfOffset: epoch.Add(-1 * resolution).Sub(now),
	})

	result, err = aq.Run()
	if assert.NoError(t, err, "Unable to run query with defaults") {
		assert.Equal(t, []string{"b", "r", "u"}, result.GroupBy)
		assert.NotNil(t, result.Until)
	}

	testMissingField(t, db, epoch, resolution, now)
}

func testMissingField(t *testing.T, db *DB, epoch time.Time, resolution time.Duration, now time.Time) {
	defer func() {
		assert.NotNil(t, recover(), "Query after removing fields should have panicked")
	}()

	tab := db.getTable("test_a")
	for i, field := range tab.Fields {
		tab.Fields[i] = sql.Field{field.Expr, "_" + field.Name}
	}

	aq := db.Query(&sql.Query{
		From:       "test_a",
		Fields:     []sql.Field{sql.Field{Expr: SUM("ii"), Name: "ii"}},
		GroupByAll: true,
		AsOfOffset: epoch.Add(-1 * resolution).Sub(now),
	})

	aq.Run()
}
<|endoftext|>