<commit_before>package server

import (
	"bytes"
	"context"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/fnproject/fn/api"
	"github.com/fnproject/fn/api/agent"
	"github.com/fnproject/fn/api/common"
	"github.com/fnproject/fn/api/models"
	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
)

type runnerResponse struct {
	RequestID string            `json:"request_id,omitempty"`
	Error     *models.ErrorBody `json:"error,omitempty"`
}

func (s *Server) handleRequest(c *gin.Context) {
	if strings.HasPrefix(c.Request.URL.Path, "/v1") {
		c.Status(http.StatusNotFound)
		return
	}

	ctx := c.Request.Context()

	r, routeExists := c.Get(api.Path)
	if !routeExists {
		r = "/"
	}

	reqRoute := &models.Route{
		AppName: c.MustGet(api.AppName).(string),
		Path:    path.Clean(r.(string)),
	}

	s.FireBeforeDispatch(ctx, reqRoute)

	s.serve(c, reqRoute.AppName, reqRoute.Path)

	s.FireAfterDispatch(ctx, reqRoute)
}

// TODO it would be nice if we could make this have nothing to do with the gin.Context but meh
// TODO make async store an *http.Request? would be sexy until we have different api format...
func (s *Server) serve(c *gin.Context, appName, path string) {
	// GetCall can mod headers, assign an id, look up the route/app (cached),
	// strip params, etc.
	call, err := s.Agent.GetCall(
		agent.WithWriter(c.Writer), // XXX (reed): order matters [for now]
		agent.FromRequest(appName, path, c.Request),
	)
	if err != nil {
		handleErrorResponse(c, err)
		return
	}

	// TODO we could add FireBeforeDispatch right here with Call in hand
	model := call.Model()
	{ // scope this, to disallow ctx use outside of this scope. add id for handleErrorResponse logger
		ctx, _ := common.LoggerWithFields(c.Request.Context(), logrus.Fields{"id": model.ID})
		c.Request = c.Request.WithContext(ctx)
	}

	if model.Type == "async" {
		// TODO we should push this into GetCall somehow (CallOpt maybe) or maybe agent.Queue(Call) ?
		buf := bytes.NewBuffer(make([]byte, c.Request.ContentLength)[:0]) // TODO sync.Pool me
		_, err := buf.ReadFrom(c.Request.Body)
		if err != nil {
			handleErrorResponse(c, models.ErrInvalidPayload)
			return
		}
		model.Payload = buf.String()

		// TODO we should probably add this to the datastore too. consider the plumber!
		_, err = s.MQ.Push(c.Request.Context(), model)
		if err != nil {
			handleErrorResponse(c, err)
			return
		}

		c.JSON(http.StatusAccepted, map[string]string{"call_id": model.ID})
		return
	}

	err = s.Agent.Submit(call)
	if err != nil {
		// NOTE if they cancel the request then it will stop the call (kind of cool),
		// we could filter that error out here too as right now it yells a little

		if err == context.DeadlineExceeded {
			// TODO maneuver
			// add this, since it means that start may not have been called [and it's relevant]
			c.Writer.Header().Add("XXX-FXLB-WAIT", time.Now().Sub(time.Time(model.CreatedAt)).String())

			err = models.ErrCallTimeout // 504 w/ friendly note
		}
		// NOTE: if the task wrote the headers already then this will fail to write
		// a 5xx (and log about it to us) -- that's fine (nice, even!)
		handleErrorResponse(c, err)
		return
	}

	// TODO plumb FXLB-WAIT somehow (api?)

	// TODO we need to watch the response writer and if no bytes written
	// then write a 200 at this point?
	// c.Data(http.StatusOK)
}
<commit_msg>fix neg content length<commit_after>package server

import (
	"bytes"
	"context"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/fnproject/fn/api"
	"github.com/fnproject/fn/api/agent"
	"github.com/fnproject/fn/api/common"
	"github.com/fnproject/fn/api/models"
	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
)

type runnerResponse struct {
	RequestID string            `json:"request_id,omitempty"`
	Error     *models.ErrorBody `json:"error,omitempty"`
}

func (s *Server) handleRequest(c *gin.Context) {
	if strings.HasPrefix(c.Request.URL.Path, "/v1") {
		c.Status(http.StatusNotFound)
		return
	}

	ctx := c.Request.Context()

	r, routeExists := c.Get(api.Path)
	if !routeExists {
		r = "/"
	}

	reqRoute := &models.Route{
		AppName: c.MustGet(api.AppName).(string),
		Path:    path.Clean(r.(string)),
	}

	s.FireBeforeDispatch(ctx, reqRoute)

	s.serve(c, reqRoute.AppName, reqRoute.Path)

	s.FireAfterDispatch(ctx, reqRoute)
}

// TODO it would be nice if we could make this have nothing to do with the gin.Context but meh
// TODO make async store an *http.Request? would be sexy until we have different api format...
func (s *Server) serve(c *gin.Context, appName, path string) {
	// GetCall can mod headers, assign an id, look up the route/app (cached),
	// strip params, etc.
	call, err := s.Agent.GetCall(
		agent.WithWriter(c.Writer), // XXX (reed): order matters [for now]
		agent.FromRequest(appName, path, c.Request),
	)
	if err != nil {
		handleErrorResponse(c, err)
		return
	}

	// TODO we could add FireBeforeDispatch right here with Call in hand
	model := call.Model()
	{ // scope this, to disallow ctx use outside of this scope. add id for handleErrorResponse logger
		ctx, _ := common.LoggerWithFields(c.Request.Context(), logrus.Fields{"id": model.ID})
		c.Request = c.Request.WithContext(ctx)
	}

	if model.Type == "async" {
		// TODO we should push this into GetCall somehow (CallOpt maybe) or maybe agent.Queue(Call) ?
		contentLength := c.Request.ContentLength
		if contentLength < 128 { // contentLength could be -1 or really small, sanitize
			contentLength = 128
		}
		buf := bytes.NewBuffer(make([]byte, int(contentLength))[:0]) // TODO sync.Pool me
		_, err := buf.ReadFrom(c.Request.Body)
		if err != nil {
			handleErrorResponse(c, models.ErrInvalidPayload)
			return
		}
		model.Payload = buf.String()

		// TODO we should probably add this to the datastore too. consider the plumber!
		_, err = s.MQ.Push(c.Request.Context(), model)
		if err != nil {
			handleErrorResponse(c, err)
			return
		}

		c.JSON(http.StatusAccepted, map[string]string{"call_id": model.ID})
		return
	}

	err = s.Agent.Submit(call)
	if err != nil {
		// NOTE if they cancel the request then it will stop the call (kind of cool),
		// we could filter that error out here too as right now it yells a little

		if err == context.DeadlineExceeded {
			// TODO maneuver
			// add this, since it means that start may not have been called [and it's relevant]
			c.Writer.Header().Add("XXX-FXLB-WAIT", time.Now().Sub(time.Time(model.CreatedAt)).String())

			err = models.ErrCallTimeout // 504 w/ friendly note
		}
		// NOTE: if the task wrote the headers already then this will fail to write
		// a 5xx (and log about it to us) -- that's fine (nice, even!)
		handleErrorResponse(c, err)
		return
	}

	// TODO plumb FXLB-WAIT somehow (api?)

	// TODO we need to watch the response writer and if no bytes written
	// then write a 200 at this point?
	// c.Data(http.StatusOK)
}
<|endoftext|>


<commit_before>/*
Package intern maps strings to integers for fast comparisons.

Description

Consider the string "This is string one" stored in variable x and the string
"This is string two" stored in variable y.  When a program tests if x == y, it
compares their first characters and finds that both are "T".  It then compares
their second characters and finds that both are "h".  This continues until the
16th character, where it finds that "o" is different from "t" and concludes
that the two strings are different.

The intern package provides two "symbol" data types to speed up such
comparisons.  The first, Eq, is an integer that is associated with particular
string contents.  For example, "This is string one" may be mapped to 52413
while "This is string two" may be mapped to 29488.  Comparing 52413 == 29488
fails in a single step, which can save a lot of time if such comparisons are
performed frequently.

The second data type, LGE, supports less than, greater than, and equal to
comparisons (and variations such as <=, >=, and !=).  For example, if x is "But
I have promises to keep", y is "And miles to go before I sleep", and z is also
"And miles to go before I sleep", then these might map to symbols 48712, 29519,
and 29519, respectively.  Let's call these sx, sy, and sz.  Then, x > y and sx
> sy; y == z and sy == sz; z < x and sz < sx; and any other such comparison of
x, y, and z, also holds for sx, sy, and sz.

A String method is defined for both Eq and LGE.  This maps a symbol back to its
original string.  Ergo, no information is lost.

Usage

NewEq maps a string to an Eq symbol, and NewLGE maps a string to an LGE symbol.
The former is faster and always succeeds (until the set of 64-bit integers is
exhausted).  NewLGE, in addition to being slower, can fail if earlier
assignments of integers to strings preclude a new string from being mapped to
an integer that respects comparisons with all existing symbols.  (In the
current implementation, a worst-case sequence of NewLGE calls will fail on the
65th call.)  As a workaround, the package provides a PreLGE function that
indicates an intention to invoke NewLGE on a particular string but without
actually assigning an integer.

Best practice is to pre-allocate as many LGE symbols as possible before calling
NewLGE.  When NewLGE is called, all strings previously passed to PreLGE are
assigned integers in an order that helps ensure that the desired relations are
preserved.  The process of pre-allocating LGE symbols with PreLGE and later
allocating them with NewLGE can be repeated as many times as necessary but with
increasing likelihood of failure with each repetition.

All functions in this package are thread-safe.
*/
package intern

import (
	"fmt"
	"sync"
)

// symbol represents either package symbol type (Eq or LGE).
type symbol uint64

// symbolList is a list of symbols.  It implements sort.Interface.
type symbolList []symbol

// Len returns the length of a symbolList.
func (sl symbolList) Len() int { return len(sl) }

// Less says is one symbolList element is less than another.
func (sl symbolList) Less(i, j int) bool { return sl[i] < sl[j] }

// Swap swaps two elements of a symbolList.
func (sl symbolList) Swap(i, j int) { sl[i], sl[j] = sl[j], sl[i] }

// state includes all the state needed to manipulate all interned-string types.
type state struct {
	symToStr     map[symbol]string // Mapping from symbols to strings
	strToSym     map[string]symbol // Mapping from strings to symbols
	tree         *tree             // Tree for maintaining symbols assignments
	pending      []string          // Strings not yet mapped to symbols
	sync.RWMutex                   // Mutex protecting all of the above
}

// forgetAll discards all extant string/symbol mappings and resets the
// assignment tables to their initial state.
func (st *state) forgetAll() {
	st.symToStr = make(map[symbol]string)
	st.strToSym = make(map[string]symbol)
	st.tree = nil
	st.pending = make([]string, 0, 100)
}

// toString converts a symbol back to a string.  It panics if given a symbol
// that was not created using New*.
func (st *state) toString(s symbol, ty string) string {
	st.RLock()
	defer st.RUnlock()
	if str, ok := st.symToStr[s]; ok {
		return str
	}
	panic(fmt.Sprintf("%d is not a valid intern.%s", s, ty))
}

// flushPending flushes all pending symbols, converting strings to symbols.
// The function returns an error status.
func (st *state) flushPending() error {
	var err error
	if len(st.pending) > 0 {
		var sMap map[string]symbol
		st.tree, sMap, err = st.tree.insertMany(st.pending)
		if err != nil {
			return err
		}
		st.pending = st.pending[:0]
		for k, v := range sMap {
			st.strToSym[k] = v
			st.symToStr[v] = k
		}
	}
	return nil
}

// assignSymbol assigns a new symbol to a string.  It takes as arguments the
// string to assign and a Boolean that indicates whether to use a tree to
// preserve order.  This method returns the assigned symbol and an error value.
func (st *state) assignSymbol(s string, useTree bool) (symbol, error) {
	// Check if the string was already assigned a symbol.
	sym, ok := st.strToSym[s]
	if ok {
		// The string was already assigned a symbol.
		return sym, nil
	}

	// We haven't seen this string before.  Find a symbol for it.
	if useTree {
		// Assign the symbol using a tree to maintain order.
		var err error
		st.tree, sym, err = st.tree.insert(s)
		if err != nil {
			return 0, err
		}
	} else {
		// Assign the next available number, starting at 1 to ensure
		// that an uninitialized symbol is treated as invalid.
		sym = symbol(len(st.symToStr)) + 1
	}
	st.symToStr[sym] = s
	st.strToSym[s] = sym
	return sym, nil
}
<commit_msg>Mention RemapAllLGEs in the package-level documentation<commit_after>/*
Package intern maps strings to integers for fast comparisons.

Description

Consider the string "This is string one" stored in variable x and the string
"This is string two" stored in variable y.  When a program tests if x == y, it
compares their first characters and finds that both are "T".  It then compares
their second characters and finds that both are "h".  This continues until the
16th character, where it finds that "o" is different from "t" and concludes
that the two strings are different.

The intern package provides two "symbol" data types to speed up such
comparisons.  The first, Eq, is an integer that is associated with particular
string contents.  For example, "This is string one" may be mapped to 52413
while "This is string two" may be mapped to 29488.  Comparing 52413 == 29488
fails in a single step, which can save a lot of time if such comparisons are
performed frequently.

The second data type, LGE, supports less than, greater than, and equal to
comparisons (and variations such as <=, >=, and !=).  For example, if x is "But
I have promises to keep", y is "And miles to go before I sleep", and z is also
"And miles to go before I sleep", then these might map to symbols 48712, 29519,
and 29519, respectively.  Let's call these sx, sy, and sz.  Then, x > y and sx
> sy; y == z and sy == sz; z < x and sz < sx; and any other such comparison of
x, y, and z, also holds for sx, sy, and sz.

A String method is defined for both Eq and LGE.  This maps a symbol back to its
original string.  Ergo, no information is lost when mapping strings to symbols.

Usage

NewEq maps a string to an Eq symbol, and NewLGE maps a string to an LGE symbol.
The former is faster and always succeeds (until the set of 64-bit integers is
exhausted).  NewLGE, in addition to being slower, can fail if earlier
assignments of integers to strings preclude a new string from being mapped to
an integer that respects comparisons with all existing symbols.  (In the
current implementation, a worst-case sequence of NewLGE calls will fail on the
65th call.)  As a workaround, the package provides a PreLGE function that
indicates an intention to invoke NewLGE on a particular string but without
actually assigning an integer.

Best practice is to pre-allocate as many LGE symbols as possible before calling
NewLGE.  When NewLGE is called, all strings previously passed to PreLGE are
assigned integers in an order that helps ensure that the desired relations are
preserved.  The process of pre-allocating LGE symbols with PreLGE and later
allocating them with NewLGE can be repeated as many times as necessary but with
increasing likelihood of failure with each repetition.  If NewLGE does fail,
the RemapAllLGEs function can be called to completely redo the mapping from
strings to LGE symbols.  The program will need to update any live LGE symbols
it has stored in data structures.

All functions in this package are thread-safe.
*/
package intern

import (
	"fmt"
	"sync"
)

// symbol represents either package symbol type (Eq or LGE).
type symbol uint64

// symbolList is a list of symbols.  It implements sort.Interface.
type symbolList []symbol

// Len returns the length of a symbolList.
func (sl symbolList) Len() int { return len(sl) }

// Less says is one symbolList element is less than another.
func (sl symbolList) Less(i, j int) bool { return sl[i] < sl[j] }

// Swap swaps two elements of a symbolList.
func (sl symbolList) Swap(i, j int) { sl[i], sl[j] = sl[j], sl[i] }

// state includes all the state needed to manipulate all interned-string types.
type state struct {
	symToStr     map[symbol]string // Mapping from symbols to strings
	strToSym     map[string]symbol // Mapping from strings to symbols
	tree         *tree             // Tree for maintaining symbols assignments
	pending      []string          // Strings not yet mapped to symbols
	sync.RWMutex                   // Mutex protecting all of the above
}

// forgetAll discards all extant string/symbol mappings and resets the
// assignment tables to their initial state.
func (st *state) forgetAll() {
	st.symToStr = make(map[symbol]string)
	st.strToSym = make(map[string]symbol)
	st.tree = nil
	st.pending = make([]string, 0, 100)
}

// toString converts a symbol back to a string.  It panics if given a symbol
// that was not created using New*.
func (st *state) toString(s symbol, ty string) string {
	st.RLock()
	defer st.RUnlock()
	if str, ok := st.symToStr[s]; ok {
		return str
	}
	panic(fmt.Sprintf("%d is not a valid intern.%s", s, ty))
}

// flushPending flushes all pending symbols, converting strings to symbols.
// The function returns an error status.
func (st *state) flushPending() error {
	var err error
	if len(st.pending) > 0 {
		var sMap map[string]symbol
		st.tree, sMap, err = st.tree.insertMany(st.pending)
		if err != nil {
			return err
		}
		st.pending = st.pending[:0]
		for k, v := range sMap {
			st.strToSym[k] = v
			st.symToStr[v] = k
		}
	}
	return nil
}

// assignSymbol assigns a new symbol to a string.  It takes as arguments the
// string to assign and a Boolean that indicates whether to use a tree to
// preserve order.  This method returns the assigned symbol and an error value.
func (st *state) assignSymbol(s string, useTree bool) (symbol, error) {
	// Check if the string was already assigned a symbol.
	sym, ok := st.strToSym[s]
	if ok {
		// The string was already assigned a symbol.
		return sym, nil
	}

	// We haven't seen this string before.  Find a symbol for it.
	if useTree {
		// Assign the symbol using a tree to maintain order.
		var err error
		st.tree, sym, err = st.tree.insert(s)
		if err != nil {
			return 0, err
		}
	} else {
		// Assign the next available number, starting at 1 to ensure
		// that an uninitialized symbol is treated as invalid.
		sym = symbol(len(st.symToStr)) + 1
	}
	st.symToStr[sym] = s
	st.strToSym[s] = sym
	return sym, nil
}
<|endoftext|>


<commit_before>package resource

import (
	"compress/gzip"
	"encoding/json"
	"github.com/golang/glog"
	"net/http"
	"strings"
)

const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT"
	DELETE = "DELETE"
	HEAD   = "HEAD"
	PATCH  = "PATCH"
)

type GetSupported interface {
	Get(*http.Request) (interface{}, int)
}

type PostSupported interface {
	Post(*http.Request) (interface{}, int)
}

type PutSupported interface {
	Put(*http.Request) (interface{}, int)
}

type DeleteSupported interface {
	Delete(*http.Request) (interface{}, int)
}

type HeadSupported interface {
	Head(*http.Request) (interface{}, int)
}

type PatchSupported interface {
	Patch(*http.Request) (interface{}, int)
}

func SupportedMethods(r interface{}) []string {
	var list []string

	if _, ok := r.(GetSupported); ok {
		list = append(list, GET)
	}
	if _, ok := r.(PostSupported); ok {
		list = append(list, POST)
	}
	if _, ok := r.(PutSupported); ok {
		list = append(list, PUT)
	}
	if _, ok := r.(DeleteSupported); ok {
		list = append(list, DELETE)
	}
	if _, ok := r.(HeadSupported); ok {
		list = append(list, HEAD)
	}
	if _, ok := r.(PatchSupported); ok {
		list = append(list, PATCH)
	}

	return list
}

func HandlerFunc(r interface{}) http.HandlerFunc {
	return func(rw http.ResponseWriter, req *http.Request) {

		var handler func(*http.Request) (interface{}, int)

		switch req.Method {
		case GET:
			if r, ok := r.(GetSupported); ok {
				handler = r.Get
			}
		case POST:
			if r, ok := r.(PostSupported); ok {
				handler = r.Post
			}
		case PUT:
			if r, ok := r.(PutSupported); ok {
				handler = r.Put
			}
		case DELETE:
			if r, ok := r.(DeleteSupported); ok {
				handler = r.Delete
			}
		case HEAD:
			if r, ok := r.(HeadSupported); ok {
				handler = r.Head
			}
		case PATCH:
			if r, ok := r.(PatchSupported); ok {
				handler = r.Patch
			}
		}

		if handler == nil {
			if supported := SupportedMethods(r); len(supported) == 0 {
				// The resource does not exist at all.
				rw.WriteHeader(http.StatusNotFound)
			} else {
				// We have other methods available, tell client.
				rw.WriteHeader(http.StatusMethodNotAllowed)
				rw.Header().Set("Allow", strings.Join(supported, ", "))
			}
			return
		}
		data, code := handler(req)

		// Return gzipped data.
		rw.Header().Set("Content-Type", "application/json; charset=utf-8")
		rw.Header().Set("Content-Encoding", "gzip")
		encoder := json.NewEncoder(gzip.NewWriter(rw))

		if err := encoder.Encode(data); err != nil {
			glog.Errorf("Failed to marshal data into JSON: %s", err)
			rw.WriteHeader(http.StatusInternalServerError)
		} else {
			rw.WriteHeader(code)
		}
	}
}
<commit_msg>Fix gzip.<commit_after>package resource

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"github.com/golang/glog"
	"net/http"
	"strings"
)

const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT"
	DELETE = "DELETE"
	HEAD   = "HEAD"
	PATCH  = "PATCH"
)

type GetSupported interface {
	Get(*http.Request) (interface{}, int)
}

type PostSupported interface {
	Post(*http.Request) (interface{}, int)
}

type PutSupported interface {
	Put(*http.Request) (interface{}, int)
}

type DeleteSupported interface {
	Delete(*http.Request) (interface{}, int)
}

type HeadSupported interface {
	Head(*http.Request) (interface{}, int)
}

type PatchSupported interface {
	Patch(*http.Request) (interface{}, int)
}

func SupportedMethods(r interface{}) []string {
	var list []string

	if _, ok := r.(GetSupported); ok {
		list = append(list, GET)
	}
	if _, ok := r.(PostSupported); ok {
		list = append(list, POST)
	}
	if _, ok := r.(PutSupported); ok {
		list = append(list, PUT)
	}
	if _, ok := r.(DeleteSupported); ok {
		list = append(list, DELETE)
	}
	if _, ok := r.(HeadSupported); ok {
		list = append(list, HEAD)
	}
	if _, ok := r.(PatchSupported); ok {
		list = append(list, PATCH)
	}

	return list
}

func HandlerFunc(r interface{}) http.HandlerFunc {
	return func(rw http.ResponseWriter, req *http.Request) {

		var handler func(*http.Request) (interface{}, int)

		switch req.Method {
		case GET:
			if r, ok := r.(GetSupported); ok {
				handler = r.Get
			}
		case POST:
			if r, ok := r.(PostSupported); ok {
				handler = r.Post
			}
		case PUT:
			if r, ok := r.(PutSupported); ok {
				handler = r.Put
			}
		case DELETE:
			if r, ok := r.(DeleteSupported); ok {
				handler = r.Delete
			}
		case HEAD:
			if r, ok := r.(HeadSupported); ok {
				handler = r.Head
			}
		case PATCH:
			if r, ok := r.(PatchSupported); ok {
				handler = r.Patch
			}
		}

		if handler == nil {
			if supported := SupportedMethods(r); len(supported) == 0 {
				// The resource does not exist at all.
				rw.WriteHeader(http.StatusNotFound)
			} else {
				// We have other methods available, tell client.
				rw.Header().Set("Allow", strings.Join(supported, ", "))
				rw.WriteHeader(http.StatusMethodNotAllowed)
			}
			return
		}
		data, code := handler(req)

		// Return gzipped data.
		rw.Header().Set("Content-Type", "application/json; charset=utf-8")
		rw.Header().Set("Content-Encoding", "gzip")

		buffer := new(bytes.Buffer)
		gz := gzip.NewWriter(buffer)
		err := json.NewEncoder(gz).Encode(data)
		gz.Close()

		if err != nil {
			glog.Errorf("Failed to marshal data into JSON: %s", err)
			rw.WriteHeader(http.StatusInternalServerError)
			return
		}

		// Encoded, send it.
		rw.WriteHeader(code)
		rw.Write(buffer.Bytes())
	}
}
<|endoftext|>


<commit_before><commit_msg>Use finalizer to close the DB. Do not open/close the DB for multiple times.<commit_after><|endoftext|>


<commit_before><commit_msg>golint<commit_after><|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/coreos/go-systemd/unit"
	"github.com/urfave/cli"
)

type FourGState struct {
	Enabled bool   `json:"enabled"`
	Configg *FourG `json:"config"`
}

type FourG struct {
	Network
}

//ToSystemdUnit implement UnitFile interface
func (f FourG) ToSystemdUnit() ([]*unit.UnitOption, error) {
	if f.Interface == "" {
		f.Interface = "eth1"
	}
	return f.Network.ToSystemdUnit()
}

func fourGState(i string) (*FourGState, error) {
	dir := os.Getenv("FCONF_CONFIGDIR")
	if dir == "" {
		dir = fconfConfigDir
	}
	b, err := ioutil.ReadFile(filepath.Join(dir,
		fmt.Sprintf(defaultFougGConfig, i)))
	if err != nil {
		return nil, err
	}
	f := &FourGState{}
	err = json.Unmarshal(b, f)
	if err != nil {
		return nil, err
	}
	if f.Configg == nil {
		return nil, ErrWrongStateFile
	}
	return f, nil
}

func FourgCMD(ctx *cli.Context) error {
	if ctx.IsSet(enableFlag) {
		return EnableFourg(ctx)
	}
	if ctx.IsSet(disableFlag) {
		return DisableFourg(ctx)
	}
	if ctx.IsSet(removeFlag) {
		return RemoveFourg(ctx)
	}
	if ctx.IsSet(configFlag) {
		return configFourgCMD(ctx)
	}
	return nil
}

func getInterface(ctx *cli.Context) string {
	var i string
	if ctx.GlobalIsSet("interface") {
		i = ctx.GlobalString("interface")
	} else {
		i = ctx.Args().First()
	}
	return i
}
func RemoveFourg(ctx *cli.Context) error {
	err := DisableFourg(ctx)
	if err != nil {
		return err
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	f, err := fourGState(i)
	if err != nil {
		return err
	}

	// removestate file
	stateFile := filepath.Join(stateDir(),
		fmt.Sprintf(defaultFougGConfig, i))
	err = removeFile(stateFile)
	if err != nil {
		return err
	}
	// remove systemd file
	unit := filepath.Join(networkBase,
		fmt.Sprintf(fourgService, i))
	err = removeFile(unit)
	if err != nil {
		return err
	}

	err = FlushInterface(f.Configg.Interface)
	if err != nil {
		return err
	}

	// reload systemd-networkd
	return restartService("systemd-networkd")
}
func EnableFourg(ctx *cli.Context) error {
	if ctx.IsSet(configFlag) {
		err := configFourgCMD(ctx)
		if err != nil {
			return err
		}
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	e, err := fourGState(i)
	if err != nil {
		return err
	}
	unit := filepath.Join(networkBase,
		fmt.Sprintf(fourgService, e.Configg.Interface))
	_, err = os.Stat(unit)
	if os.IsNotExist(err) {
		err = CreateSystemdFile(e.Configg, unit, 0644)
		if err != nil {
			return err
		}
	}
	_, err = exec.Command("ip", "link", "set", "up", e.Configg.Interface).Output()
	if err != nil {
		return fmt.Errorf("ERROR: runnin ip link set up %s %v",
			e.Configg.Interface, err,
		)
	}
	err = restartService("systemd-networkd")
	if err != nil {
		return fmt.Errorf("ERROR: restarting systemd %v ", err)
	}
	e.Enabled = true
	data, err := json.Marshal(e)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultFougGConfig, i), data)
}

func DisableFourg(ctx *cli.Context) error {
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	e, err := fourGState(i)
	if err != nil {
		return err
	}
	_, err = exec.Command("ip", "link", "set", "down", "dev", e.Configg.Interface).Output()
	if err != nil {
		return err
	}
	unit := filepath.Join(networkBase,
		fmt.Sprintf(fourgService, e.Configg.Interface))
	err = removeFile(unit)
	if err != nil {
		return err
	}
	fmt.Println("successfully disabled 4G")
	e.Enabled = false
	data, err := json.Marshal(e)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultFougGConfig, i), data)
}

func configFourgCMD(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.String("config")
	if src == "" {
		return errors.New("fconf: missing configuration source file")
	}
	var b []byte
	var err error
	if src == "stdin" {
		b, err = ReadFromStdin()
		if err != nil {
			return err
		}
	} else {
		b, err = ioutil.ReadFile(src)
		if err != nil {
			return err
		}
	}
	e := FourG{}
	err = json.Unmarshal(b, &e)
	if err != nil {
		return err
	}
	if e.Interface == "" {
		e.Interface = "eth1"
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	if strings.Contains(name, "%s") {
		name = fmt.Sprintf(name, e.Interface)
	}
	filename := filepath.Join(base, name)
	err = CreateSystemdFile(e, filename, 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written 4G configuration to %s \n", filename)
	state := &FourGState{Configg: &e}
	ctx.GlobalSet("interface", e.Interface)
	b, _ = json.Marshal(state)
	return keepState(
		fmt.Sprintf(defaultFougGConfig, e.Interface), b)
}
<commit_msg>4g: ignore ip command<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/coreos/go-systemd/unit"
	"github.com/urfave/cli"
)

type FourGState struct {
	Enabled bool   `json:"enabled"`
	Configg *FourG `json:"config"`
}

type FourG struct {
	Network
}

//ToSystemdUnit implement UnitFile interface
func (f FourG) ToSystemdUnit() ([]*unit.UnitOption, error) {
	if f.Interface == "" {
		f.Interface = "eth1"
	}
	return f.Network.ToSystemdUnit()
}

func fourGState(i string) (*FourGState, error) {
	dir := os.Getenv("FCONF_CONFIGDIR")
	if dir == "" {
		dir = fconfConfigDir
	}
	b, err := ioutil.ReadFile(filepath.Join(dir,
		fmt.Sprintf(defaultFougGConfig, i)))
	if err != nil {
		return nil, err
	}
	f := &FourGState{}
	err = json.Unmarshal(b, f)
	if err != nil {
		return nil, err
	}
	if f.Configg == nil {
		return nil, ErrWrongStateFile
	}
	return f, nil
}

func FourgCMD(ctx *cli.Context) error {
	if ctx.IsSet(enableFlag) {
		return EnableFourg(ctx)
	}
	if ctx.IsSet(disableFlag) {
		return DisableFourg(ctx)
	}
	if ctx.IsSet(removeFlag) {
		return RemoveFourg(ctx)
	}
	if ctx.IsSet(configFlag) {
		return configFourgCMD(ctx)
	}
	return nil
}

func getInterface(ctx *cli.Context) string {
	var i string
	if ctx.GlobalIsSet("interface") {
		i = ctx.GlobalString("interface")
	} else {
		i = ctx.Args().First()
	}
	return i
}
func RemoveFourg(ctx *cli.Context) error {
	err := DisableFourg(ctx)
	if err != nil {
		return err
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	f, err := fourGState(i)
	if err != nil {
		return err
	}

	// removestate file
	stateFile := filepath.Join(stateDir(),
		fmt.Sprintf(defaultFougGConfig, i))
	err = removeFile(stateFile)
	if err != nil {
		return err
	}
	// remove systemd file
	unit := filepath.Join(networkBase,
		fmt.Sprintf(fourgService, i))
	err = removeFile(unit)
	if err != nil {
		return err
	}

	err = FlushInterface(f.Configg.Interface)
	if err != nil {
		return err
	}

	// reload systemd-networkd
	return restartService("systemd-networkd")
}
func EnableFourg(ctx *cli.Context) error {
	if ctx.IsSet(configFlag) {
		err := configFourgCMD(ctx)
		if err != nil {
			return err
		}
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	e, err := fourGState(i)
	if err != nil {
		return err
	}
	unit := filepath.Join(networkBase,
		fmt.Sprintf(fourgService, e.Configg.Interface))
	_, err = os.Stat(unit)
	if os.IsNotExist(err) {
		err = CreateSystemdFile(e.Configg, unit, 0644)
		if err != nil {
			return err
		}
	}
	_, err = exec.Command("ip", "link", "set", "up", e.Configg.Interface).Output()
	if err != nil {
		return fmt.Errorf("ERROR: runnin ip link set up %s %v",
			e.Configg.Interface, err,
		)
	}
	err = restartService("systemd-networkd")
	if err != nil {
		return fmt.Errorf("ERROR: restarting systemd %v ", err)
	}
	e.Enabled = true
	data, err := json.Marshal(e)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultFougGConfig, i), data)
}

func DisableFourg(ctx *cli.Context) error {
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	e, err := fourGState(i)
	if err != nil {
		return err
	}
	//_, err = exec.Command("ip", "link", "set", "down", "dev", e.Configg.Interface).Output()
	//if err != nil {
	//return err
	//}
	unit := filepath.Join(networkBase,
		fmt.Sprintf(fourgService, e.Configg.Interface))
	err = removeFile(unit)
	if err != nil {
		return err
	}
	fmt.Println("successfully disabled 4G")
	e.Enabled = false
	data, err := json.Marshal(e)
	if err != nil {
		return err
	}
	err = restartService("systemd-networkd")
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultFougGConfig, i), data)
}

func configFourgCMD(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.String("config")
	if src == "" {
		return errors.New("fconf: missing configuration source file")
	}
	var b []byte
	var err error
	if src == "stdin" {
		b, err = ReadFromStdin()
		if err != nil {
			return err
		}
	} else {
		b, err = ioutil.ReadFile(src)
		if err != nil {
			return err
		}
	}
	e := FourG{}
	err = json.Unmarshal(b, &e)
	if err != nil {
		return err
	}
	if e.Interface == "" {
		e.Interface = "eth1"
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	if strings.Contains(name, "%s") {
		name = fmt.Sprintf(name, e.Interface)
	}
	filename := filepath.Join(base, name)
	err = CreateSystemdFile(e, filename, 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written 4G configuration to %s \n", filename)
	state := &FourGState{Configg: &e}
	ctx.GlobalSet("interface", e.Interface)
	b, _ = json.Marshal(state)
	return keepState(
		fmt.Sprintf(defaultFougGConfig, e.Interface), b)
}
<|endoftext|>


<commit_before>package query

import (
	"fmt"
	pb "github.com/mediachain/concat/proto"
	"strings"
)

// query evaluation: very primitive eval with only simple statements
// will have to do until we have an index and we can compile to sql.
func EvalQuery(query *Query, stmts []*pb.Statement) ([]interface{}, error) {
	nsfilter := makeNamespaceFilter(query)

	cfilter, err := makeCriteriaFilter(query)
	if err != nil {
		return nil, err
	}

	rs, err := makeResultSet(query)
	if err != nil {
		return nil, err
	}

	rs.begin(len(stmts))
	for _, stmt := range stmts {
		if nsfilter(stmt) && cfilter(stmt) {
			rs.add(stmt)
		}
	}
	rs.end()

	return rs.result(), nil
}

type QueryResultSet interface {
	begin(hint int)
	add(*pb.Statement)
	end()
	result() []interface{}
}

type QueryEvalError string

func (e QueryEvalError) Error() string {
	return string(e)
}

type StatementFilter func(*pb.Statement) bool
type ValueCriteriaFilter func(*pb.Statement, string) bool
type TimeCriteriaFilter func(*pb.Statement, int64) bool
type CompoundCriteriaFilter func(stmt *pb.Statement, left, right StatementFilter) bool

func idCriteriaFilter(stmt *pb.Statement, id string) bool {
	return stmt.Id == id
}

func publisherCriteriaFilter(stmt *pb.Statement, pub string) bool {
	return stmt.Publisher == pub
}

func sourceCriteriaFilter(stmt *pb.Statement, src string) bool {
	// only support simple statements for now, so src = publisher
	return stmt.Publisher == src
}

var valueCriteriaFilters = map[string]ValueCriteriaFilter{
	"id":        idCriteriaFilter,
	"publisher": publisherCriteriaFilter,
	"source":    sourceCriteriaFilter}

func timestampFilterLTEQ(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp <= ts
}

func timestampFilterLT(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp < ts
}

func timestampFilterEQ(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp == ts
}

func timestampFilterGTEQ(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp >= ts
}

func timestampFilterGT(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp > ts
}

var timeCriteriaFilters = map[string]TimeCriteriaFilter{
	"<=": timestampFilterLTEQ,
	"<":  timestampFilterLT,
	"=":  timestampFilterEQ,
	">=": timestampFilterGTEQ,
	">":  timestampFilterGT}

func compoundCriteriaAND(stmt *pb.Statement, left, right StatementFilter) bool {
	return left(stmt) && right(stmt)
}

func compoundCriteriaOR(stmt *pb.Statement, left, right StatementFilter) bool {
	return left(stmt) || right(stmt)
}

var compoundCriteriaFilters = map[string]CompoundCriteriaFilter{
	"AND": compoundCriteriaAND,
	"OR":  compoundCriteriaOR}

func emptyFilter(*pb.Statement) bool {
	return true
}

func makeNamespaceFilter(query *Query) StatementFilter {
	ns := query.namespace
	switch {
	case ns == "*":
		return emptyFilter

	case ns[len(ns)-1] == '*':
		prefix := ns[:len(ns)-2]
		return func(stmt *pb.Statement) bool {
			return strings.HasPrefix(stmt.Namespace, prefix)
		}

	default:
		return func(stmt *pb.Statement) bool {
			return stmt.Namespace == ns
		}
	}
}

func makeCriteriaFilter(query *Query) (StatementFilter, error) {
	c := query.criteria
	if c == nil {
		return emptyFilter, nil
	}

	return makeCriteriaFilterF(c)
}

func makeCriteriaFilterF(c QueryCriteria) (StatementFilter, error) {
	switch c := c.(type) {
	case *ValueCriteria:
		filter, ok := valueCriteriaFilters[c.sel]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria selector: %s", c.sel))
		}

		return func(stmt *pb.Statement) bool {
			return filter(stmt, c.val)
		}, nil

	case *TimeCriteria:
		filter, ok := timeCriteriaFilters[c.op]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria time op: %s", c.op))
		}

		return func(stmt *pb.Statement) bool {
			return filter(stmt, c.ts)
		}, nil

	case *CompoundCriteria:
		filter, ok := compoundCriteriaFilters[c.op]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria combinator: %s", c.op))
		}

		left, err := makeCriteriaFilterF(c.left)
		if err != nil {
			return nil, err
		}

		right, err := makeCriteriaFilterF(c.right)
		if err != nil {
			return nil, err
		}

		return func(stmt *pb.Statement) bool {
			return filter(stmt, left, right)
		}, nil

	default:
		return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria type: %T", c))
	}
}

type StatementSelector func(*pb.Statement) interface{}

func simpleSelectorAll(stmt *pb.Statement) interface{} {
	return stmt
}

func simpleSelectorBody(stmt *pb.Statement) interface{} {
	return stmt.Body
}

func simpleSelectorId(stmt *pb.Statement) interface{} {
	return stmt.Id
}

func simpleSelectorPublisher(stmt *pb.Statement) interface{} {
	return stmt.Publisher
}

func simpleSelectorNamespace(stmt *pb.Statement) interface{} {
	return stmt.Namespace
}

func simpleSelectorSource(stmt *pb.Statement) interface{} {
	// simple statements, source = publisher
	return stmt.Publisher
}

func simpleSelectorTimestamp(stmt *pb.Statement) interface{} {
	return stmt.Timestamp
}

var simpleSelectors = map[string]StatementSelector{
	"*":         simpleSelectorAll,
	"body":      simpleSelectorBody,
	"id":        simpleSelectorId,
	"publisher": simpleSelectorPublisher,
	"namespace": simpleSelectorNamespace,
	"source":    simpleSelectorSource,
	"timestamp": simpleSelectorTimestamp}

type FunctionStatementSelector func([]interface{}) []interface{}

func countFunctionSelector(res []interface{}) []interface{} {
	return []interface{}{len(res)}
}

var functionSelectors = map[string]FunctionStatementSelector{
	"COUNT": countFunctionSelector}

func makeResultSet(query *Query) (QueryResultSet, error) {
	sel := query.selector
	switch sel := sel.(type) {
	case SimpleSelector:
		getf, ok := simpleSelectors[string(sel)]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", sel))
		}

		return makeSimpleResultSet(getf, query.limit), nil

	case CompoundSelector:
		getfs := make([]StatementSelector, len(sel))
		for x, key := range sel {
			getf, ok := simpleSelectors[string(key)]
			if !ok {
				return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", key))
			}
			getfs[x] = getf
		}

		return makeCompoundResultSet(sel, getfs, query.limit), nil

	case *FunctionSelector:
		fun, ok := functionSelectors[sel.op]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", sel.op))
		}

		getf, ok := simpleSelectors[string(sel.sel)]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", sel.sel))
		}

		return makeFunctionResultSet(fun, getf, query.limit), nil

	default:
		return nil, QueryEvalError(fmt.Sprintf("Unexpected selector type: %T", sel))
	}
}

func makeSimpleResultSet(getf StatementSelector, limit int) QueryResultSet {
	return &SimpleResultSet{rset: make(map[interface{}]bool), getf: getf, limit: limit}
}

type SimpleResultSet struct {
	rset  map[interface{}]bool
	res   []interface{}
	getf  StatementSelector
	limit int
}

func (rs *SimpleResultSet) begin(hint int) {}

func (rs *SimpleResultSet) add(stmt *pb.Statement) {
	if rs.limit == 0 || len(rs.rset) < rs.limit {
		rs.rset[rs.getf(stmt)] = true
	}
}

func (rs *SimpleResultSet) end() {
	rs.res = make([]interface{}, len(rs.rset))
	x := 0
	for obj, _ := range rs.rset {
		rs.res[x] = obj
		x++
	}
	rs.rset = nil
}

func (rs *SimpleResultSet) result() []interface{} {
	return rs.res
}

func makeCompoundResultSet(sels []SimpleSelector, getfs []StatementSelector, limit int) QueryResultSet {
	compf := makeCompoundStatementSelector(sels, getfs)
	return &CompoundResultSet{getf: compf, limit: limit}
}

func makeCompoundStatementSelector(sels []SimpleSelector, getfs []StatementSelector) StatementSelector {
	return func(stmt *pb.Statement) interface{} {
		val := make(map[string]interface{})
		for x, sel := range sels {
			val[string(sel)] = getfs[x](stmt)
		}
		return val
	}
}

type CompoundResultSet struct {
	rset  []interface{}
	getf  StatementSelector
	limit int
}

func (rs *CompoundResultSet) begin(hint int) {
	rs.rset = make([]interface{}, 0, hint)
}

func (rs *CompoundResultSet) add(stmt *pb.Statement) {
	rs.rset = append(rs.rset, rs.getf(stmt))
}

func (rs *CompoundResultSet) end() {}

func (rs *CompoundResultSet) result() []interface{} {
	return rs.rset
}

func makeFunctionResultSet(fun FunctionStatementSelector, getf StatementSelector, limit int) QueryResultSet {
	return &FunctionResultSet{rset: makeSimpleResultSet(getf, limit), fun: fun}
}

type FunctionResultSet struct {
	rset QueryResultSet
	fun  FunctionStatementSelector
	res  []interface{}
}

func (rs *FunctionResultSet) begin(hint int) {
	rs.rset.begin(hint)
}

func (rs *FunctionResultSet) add(stmt *pb.Statement) {
	rs.rset.add(stmt)
}

func (rs *FunctionResultSet) end() {
	rs.rset.end()
	rs.res = rs.fun(rs.rset.result())
}

func (rs *FunctionResultSet) result() []interface{} {
	return rs.res
}
<commit_msg>query/eval: Add comments explaining the differences of the possible result sets<commit_after>package query

import (
	"fmt"
	pb "github.com/mediachain/concat/proto"
	"strings"
)

// query evaluation: very primitive eval with only simple statements
// will have to do until we have an index and we can compile to sql.
func EvalQuery(query *Query, stmts []*pb.Statement) ([]interface{}, error) {
	nsfilter := makeNamespaceFilter(query)

	cfilter, err := makeCriteriaFilter(query)
	if err != nil {
		return nil, err
	}

	rs, err := makeResultSet(query)
	if err != nil {
		return nil, err
	}

	rs.begin(len(stmts))
	for _, stmt := range stmts {
		if nsfilter(stmt) && cfilter(stmt) {
			rs.add(stmt)
		}
	}
	rs.end()

	return rs.result(), nil
}

type QueryResultSet interface {
	begin(hint int)
	add(*pb.Statement)
	end()
	result() []interface{}
}

type QueryEvalError string

func (e QueryEvalError) Error() string {
	return string(e)
}

type StatementFilter func(*pb.Statement) bool
type ValueCriteriaFilter func(*pb.Statement, string) bool
type TimeCriteriaFilter func(*pb.Statement, int64) bool
type CompoundCriteriaFilter func(stmt *pb.Statement, left, right StatementFilter) bool

func idCriteriaFilter(stmt *pb.Statement, id string) bool {
	return stmt.Id == id
}

func publisherCriteriaFilter(stmt *pb.Statement, pub string) bool {
	return stmt.Publisher == pub
}

func sourceCriteriaFilter(stmt *pb.Statement, src string) bool {
	// only support simple statements for now, so src = publisher
	return stmt.Publisher == src
}

var valueCriteriaFilters = map[string]ValueCriteriaFilter{
	"id":        idCriteriaFilter,
	"publisher": publisherCriteriaFilter,
	"source":    sourceCriteriaFilter}

func timestampFilterLTEQ(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp <= ts
}

func timestampFilterLT(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp < ts
}

func timestampFilterEQ(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp == ts
}

func timestampFilterGTEQ(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp >= ts
}

func timestampFilterGT(stmt *pb.Statement, ts int64) bool {
	return stmt.Timestamp > ts
}

var timeCriteriaFilters = map[string]TimeCriteriaFilter{
	"<=": timestampFilterLTEQ,
	"<":  timestampFilterLT,
	"=":  timestampFilterEQ,
	">=": timestampFilterGTEQ,
	">":  timestampFilterGT}

func compoundCriteriaAND(stmt *pb.Statement, left, right StatementFilter) bool {
	return left(stmt) && right(stmt)
}

func compoundCriteriaOR(stmt *pb.Statement, left, right StatementFilter) bool {
	return left(stmt) || right(stmt)
}

var compoundCriteriaFilters = map[string]CompoundCriteriaFilter{
	"AND": compoundCriteriaAND,
	"OR":  compoundCriteriaOR}

func emptyFilter(*pb.Statement) bool {
	return true
}

func makeNamespaceFilter(query *Query) StatementFilter {
	ns := query.namespace
	switch {
	case ns == "*":
		return emptyFilter

	case ns[len(ns)-1] == '*':
		prefix := ns[:len(ns)-2]
		return func(stmt *pb.Statement) bool {
			return strings.HasPrefix(stmt.Namespace, prefix)
		}

	default:
		return func(stmt *pb.Statement) bool {
			return stmt.Namespace == ns
		}
	}
}

func makeCriteriaFilter(query *Query) (StatementFilter, error) {
	c := query.criteria
	if c == nil {
		return emptyFilter, nil
	}

	return makeCriteriaFilterF(c)
}

func makeCriteriaFilterF(c QueryCriteria) (StatementFilter, error) {
	switch c := c.(type) {
	case *ValueCriteria:
		filter, ok := valueCriteriaFilters[c.sel]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria selector: %s", c.sel))
		}

		return func(stmt *pb.Statement) bool {
			return filter(stmt, c.val)
		}, nil

	case *TimeCriteria:
		filter, ok := timeCriteriaFilters[c.op]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria time op: %s", c.op))
		}

		return func(stmt *pb.Statement) bool {
			return filter(stmt, c.ts)
		}, nil

	case *CompoundCriteria:
		filter, ok := compoundCriteriaFilters[c.op]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria combinator: %s", c.op))
		}

		left, err := makeCriteriaFilterF(c.left)
		if err != nil {
			return nil, err
		}

		right, err := makeCriteriaFilterF(c.right)
		if err != nil {
			return nil, err
		}

		return func(stmt *pb.Statement) bool {
			return filter(stmt, left, right)
		}, nil

	default:
		return nil, QueryEvalError(fmt.Sprintf("Unexpected criteria type: %T", c))
	}
}

type StatementSelector func(*pb.Statement) interface{}

func simpleSelectorAll(stmt *pb.Statement) interface{} {
	return stmt
}

func simpleSelectorBody(stmt *pb.Statement) interface{} {
	return stmt.Body
}

func simpleSelectorId(stmt *pb.Statement) interface{} {
	return stmt.Id
}

func simpleSelectorPublisher(stmt *pb.Statement) interface{} {
	return stmt.Publisher
}

func simpleSelectorNamespace(stmt *pb.Statement) interface{} {
	return stmt.Namespace
}

func simpleSelectorSource(stmt *pb.Statement) interface{} {
	// simple statements, source = publisher
	return stmt.Publisher
}

func simpleSelectorTimestamp(stmt *pb.Statement) interface{} {
	return stmt.Timestamp
}

var simpleSelectors = map[string]StatementSelector{
	"*":         simpleSelectorAll,
	"body":      simpleSelectorBody,
	"id":        simpleSelectorId,
	"publisher": simpleSelectorPublisher,
	"namespace": simpleSelectorNamespace,
	"source":    simpleSelectorSource,
	"timestamp": simpleSelectorTimestamp}

type FunctionStatementSelector func([]interface{}) []interface{}

func countFunctionSelector(res []interface{}) []interface{} {
	return []interface{}{len(res)}
}

var functionSelectors = map[string]FunctionStatementSelector{
	"COUNT": countFunctionSelector}

// The difference between the types of result set:
//  Simple selectors (SimpleResultSet) have unique (set) semantics.
//  Compound selectors (CompoundResultSet) create objects with fields named by
//   their selector, and have distinct semantics.
//  Function selectors (FunctionResultSet) perform a selection and apply a
//   function on the simple result set.
// The difference is illustrated with these two expressions
//  SELECT namespace FROM *
//  SELECT (namespace) FROM *
//  SELECT COUNT(namespace) FROM *
// The first form, will return a set of unique namespaces from all statements
// as flat strings
// The second form will return a list of objects, with a field named namespace
//  containing the namespace of some statement. There are as many objects in
//  the result set as statements.
// The third form will return a list with one element, which will be the count
//  of distinct namespaces.
func makeResultSet(query *Query) (QueryResultSet, error) {
	sel := query.selector
	switch sel := sel.(type) {
	case SimpleSelector:
		getf, ok := simpleSelectors[string(sel)]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", sel))
		}

		return makeSimpleResultSet(getf, query.limit), nil

	case CompoundSelector:
		getfs := make([]StatementSelector, len(sel))
		for x, key := range sel {
			getf, ok := simpleSelectors[string(key)]
			if !ok {
				return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", key))
			}
			getfs[x] = getf
		}

		return makeCompoundResultSet(sel, getfs, query.limit), nil

	case *FunctionSelector:
		fun, ok := functionSelectors[sel.op]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", sel.op))
		}

		getf, ok := simpleSelectors[string(sel.sel)]
		if !ok {
			return nil, QueryEvalError(fmt.Sprintf("Unexpected selector: %s", sel.sel))
		}

		return makeFunctionResultSet(fun, getf, query.limit), nil

	default:
		return nil, QueryEvalError(fmt.Sprintf("Unexpected selector type: %T", sel))
	}
}

func makeSimpleResultSet(getf StatementSelector, limit int) QueryResultSet {
	return &SimpleResultSet{rset: make(map[interface{}]bool), getf: getf, limit: limit}
}

type SimpleResultSet struct {
	rset  map[interface{}]bool
	res   []interface{}
	getf  StatementSelector
	limit int
}

func (rs *SimpleResultSet) begin(hint int) {}

func (rs *SimpleResultSet) add(stmt *pb.Statement) {
	if rs.limit == 0 || len(rs.rset) < rs.limit {
		rs.rset[rs.getf(stmt)] = true
	}
}

func (rs *SimpleResultSet) end() {
	rs.res = make([]interface{}, len(rs.rset))
	x := 0
	for obj, _ := range rs.rset {
		rs.res[x] = obj
		x++
	}
	rs.rset = nil
}

func (rs *SimpleResultSet) result() []interface{} {
	return rs.res
}

func makeCompoundResultSet(sels []SimpleSelector, getfs []StatementSelector, limit int) QueryResultSet {
	compf := makeCompoundStatementSelector(sels, getfs)
	return &CompoundResultSet{getf: compf, limit: limit}
}

func makeCompoundStatementSelector(sels []SimpleSelector, getfs []StatementSelector) StatementSelector {
	return func(stmt *pb.Statement) interface{} {
		val := make(map[string]interface{})
		for x, sel := range sels {
			val[string(sel)] = getfs[x](stmt)
		}
		return val
	}
}

type CompoundResultSet struct {
	rset  []interface{}
	getf  StatementSelector
	limit int
}

func (rs *CompoundResultSet) begin(hint int) {
	rs.rset = make([]interface{}, 0, hint)
}

func (rs *CompoundResultSet) add(stmt *pb.Statement) {
	rs.rset = append(rs.rset, rs.getf(stmt))
}

func (rs *CompoundResultSet) end() {}

func (rs *CompoundResultSet) result() []interface{} {
	return rs.rset
}

func makeFunctionResultSet(fun FunctionStatementSelector, getf StatementSelector, limit int) QueryResultSet {
	return &FunctionResultSet{rset: makeSimpleResultSet(getf, limit), fun: fun}
}

type FunctionResultSet struct {
	rset QueryResultSet
	fun  FunctionStatementSelector
	res  []interface{}
}

func (rs *FunctionResultSet) begin(hint int) {
	rs.rset.begin(hint)
}

func (rs *FunctionResultSet) add(stmt *pb.Statement) {
	rs.rset.add(stmt)
}

func (rs *FunctionResultSet) end() {
	rs.rset.end()
	rs.res = rs.fun(rs.rset.result())
}

func (rs *FunctionResultSet) result() []interface{} {
	return rs.res
}
<|endoftext|>


<commit_before>package khepri

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path"
	"path/filepath"
)

const (
	dirMode  = 0700
	blobPath = "blobs"
	refPath  = "refs"
	tempPath = "tmp"
)

var (
	ErrIDDoesNotExist = errors.New("ID does not exist")
)

// Name stands for the alias given to an ID.
type Name string

func (n Name) Encode() string {
	return url.QueryEscape(string(n))
}

type Repository struct {
	path string
	hash func() hash.Hash
}

type Type int

const (
	TYPE_BLOB = iota
	TYPE_REF
)

func NewTypeFromString(s string) Type {
	switch s {
	case "blob":
		return TYPE_BLOB
	case "ref":
		return TYPE_REF
	}

	panic(fmt.Sprintf("unknown type %q", s))
}

func (t Type) String() string {
	switch t {
	case TYPE_BLOB:
		return "blob"
	case TYPE_REF:
		return "ref"
	}

	panic(fmt.Sprintf("unknown type %d", t))
}

// NewDirRepository creates a new dir-baked repository at the given path.
func NewRepository(path string) (*Repository, error) {
	d := &Repository{
		path: path,
		hash: sha256.New,
	}

	err := d.create()

	if err != nil {
		return nil, err
	}

	return d, nil
}

func (r *Repository) create() error {
	dirs := []string{
		r.path,
		path.Join(r.path, blobPath),
		path.Join(r.path, refPath),
		path.Join(r.path, tempPath),
	}

	for _, dir := range dirs {
		err := os.MkdirAll(dir, dirMode)
		if err != nil {
			return err
		}
	}

	return nil
}

// SetHash changes the hash function used for deriving IDs. Default is SHA256.
func (r *Repository) SetHash(h func() hash.Hash) {
	r.hash = h
}

// Path returns the directory used for this repository.
func (r *Repository) Path() string {
	return r.path
}

// Return temp directory in correct directory for this repository.
func (r *Repository) tempFile() (*os.File, error) {
	return ioutil.TempFile(path.Join(r.path, tempPath), "temp-")
}

// Rename temp file to final name according to type and ID.
func (r *Repository) renameFile(file *os.File, t Type, id ID) error {
	filename := path.Join(r.dir(t), id.String())
	return os.Rename(file.Name(), filename)
}

// Construct directory for given Type.
func (r *Repository) dir(t Type) string {
	switch t {
	case TYPE_BLOB:
		return path.Join(r.path, blobPath)
	case TYPE_REF:
		return path.Join(r.path, refPath)
	}

	panic(fmt.Sprintf("unknown type %d", t))
}

// Construct path for given Type and ID.
func (r *Repository) filename(t Type, id ID) string {
	return path.Join(r.dir(t), id.String())
}

// Test returns true if the given ID exists in the repository.
func (r *Repository) Test(t Type, id ID) (bool, error) {
	// try to open file
	file, err := os.Open(r.filename(t, id))
	defer func() {
		file.Close()
	}()

	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}

	return true, nil
}

// Get returns a reader for the content stored under the given ID.
func (r *Repository) Get(t Type, id ID) (io.Reader, error) {
	// try to open file
	file, err := os.Open(r.filename(t, id))
	if err != nil {
		return nil, err
	}

	return file, nil
}

// Remove removes the content stored at ID.
func (r *Repository) Remove(t Type, id ID) error {
	return os.Remove(r.filename(t, id))
}

type IDs []ID

// Lists all objects of a given type.
func (r *Repository) ListIDs(t Type) (IDs, error) {
	// TODO: use os.Open() and d.Readdirnames() instead of Glob()
	pattern := path.Join(r.dir(t), "*")

	matches, err := filepath.Glob(pattern)
	if err != nil {
		return nil, err
	}

	ids := make(IDs, 0, len(matches))

	for _, m := range matches {
		base := filepath.Base(m)

		if base == "" {
			continue
		}
		id, err := ParseID(base)

		if err != nil {
			continue
		}

		ids = append(ids, id)
	}

	return ids, nil
}

func (ids IDs) Len() int {
	return len(ids)
}

func (ids IDs) Less(i, j int) bool {
	if len(ids[i]) < len(ids[j]) {
		return true
	}

	for k, b := range ids[i] {
		if b == ids[j][k] {
			continue
		}

		if b < ids[j][k] {
			return true
		} else {
			return false
		}
	}

	return false
}

func (ids IDs) Swap(i, j int) {
	ids[i], ids[j] = ids[j], ids[i]
}
<commit_msg>Remove SetHash()<commit_after>package khepri

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path"
	"path/filepath"
)

const (
	dirMode  = 0700
	blobPath = "blobs"
	refPath  = "refs"
	tempPath = "tmp"
)

var (
	ErrIDDoesNotExist = errors.New("ID does not exist")
)

// Name stands for the alias given to an ID.
type Name string

func (n Name) Encode() string {
	return url.QueryEscape(string(n))
}

type Repository struct {
	path string
	hash func() hash.Hash
}

type Type int

const (
	TYPE_BLOB = iota
	TYPE_REF
)

func NewTypeFromString(s string) Type {
	switch s {
	case "blob":
		return TYPE_BLOB
	case "ref":
		return TYPE_REF
	}

	panic(fmt.Sprintf("unknown type %q", s))
}

func (t Type) String() string {
	switch t {
	case TYPE_BLOB:
		return "blob"
	case TYPE_REF:
		return "ref"
	}

	panic(fmt.Sprintf("unknown type %d", t))
}

// NewDirRepository creates a new dir-baked repository at the given path.
func NewRepository(path string) (*Repository, error) {
	d := &Repository{
		path: path,
		hash: sha256.New,
	}

	err := d.create()

	if err != nil {
		return nil, err
	}

	return d, nil
}

func (r *Repository) create() error {
	dirs := []string{
		r.path,
		path.Join(r.path, blobPath),
		path.Join(r.path, refPath),
		path.Join(r.path, tempPath),
	}

	for _, dir := range dirs {
		err := os.MkdirAll(dir, dirMode)
		if err != nil {
			return err
		}
	}

	return nil
}

// Path returns the directory used for this repository.
func (r *Repository) Path() string {
	return r.path
}

// Return temp directory in correct directory for this repository.
func (r *Repository) tempFile() (*os.File, error) {
	return ioutil.TempFile(path.Join(r.path, tempPath), "temp-")
}

// Rename temp file to final name according to type and ID.
func (r *Repository) renameFile(file *os.File, t Type, id ID) error {
	filename := path.Join(r.dir(t), id.String())
	return os.Rename(file.Name(), filename)
}

// Construct directory for given Type.
func (r *Repository) dir(t Type) string {
	switch t {
	case TYPE_BLOB:
		return path.Join(r.path, blobPath)
	case TYPE_REF:
		return path.Join(r.path, refPath)
	}

	panic(fmt.Sprintf("unknown type %d", t))
}

// Construct path for given Type and ID.
func (r *Repository) filename(t Type, id ID) string {
	return path.Join(r.dir(t), id.String())
}

// Test returns true if the given ID exists in the repository.
func (r *Repository) Test(t Type, id ID) (bool, error) {
	// try to open file
	file, err := os.Open(r.filename(t, id))
	defer func() {
		file.Close()
	}()

	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}

	return true, nil
}

// Get returns a reader for the content stored under the given ID.
func (r *Repository) Get(t Type, id ID) (io.Reader, error) {
	// try to open file
	file, err := os.Open(r.filename(t, id))
	if err != nil {
		return nil, err
	}

	return file, nil
}

// Remove removes the content stored at ID.
func (r *Repository) Remove(t Type, id ID) error {
	return os.Remove(r.filename(t, id))
}

type IDs []ID

// Lists all objects of a given type.
func (r *Repository) ListIDs(t Type) (IDs, error) {
	// TODO: use os.Open() and d.Readdirnames() instead of Glob()
	pattern := path.Join(r.dir(t), "*")

	matches, err := filepath.Glob(pattern)
	if err != nil {
		return nil, err
	}

	ids := make(IDs, 0, len(matches))

	for _, m := range matches {
		base := filepath.Base(m)

		if base == "" {
			continue
		}
		id, err := ParseID(base)

		if err != nil {
			continue
		}

		ids = append(ids, id)
	}

	return ids, nil
}

func (ids IDs) Len() int {
	return len(ids)
}

func (ids IDs) Less(i, j int) bool {
	if len(ids[i]) < len(ids[j]) {
		return true
	}

	for k, b := range ids[i] {
		if b == ids[j][k] {
			continue
		}

		if b < ids[j][k] {
			return true
		} else {
			return false
		}
	}

	return false
}

func (ids IDs) Swap(i, j int) {
	ids[i], ids[j] = ids[j], ids[i]
}
<|endoftext|>


<commit_before>package scene

import (
	"github.com/pankona/gomo-simra/simra"

	"golang.org/x/mobile/asset"
)

type sound struct {
	game *game
}

func (s *sound) play(assetName string) {
	a := simra.NewAudio()
	resource, err := asset.Open(assetName)
	if err != nil {
		panic(err.Error())
	}
	err = a.Play(resource, false, func(err error) {
		if err != nil {
			simra.LogError(err.Error())
		}
	})
	if err != nil {
		panic(err.Error())
	}
}

func (s *sound) OnEvent(i interface{}) {
	c, ok := i.(*command)
	if !ok {
		panic("unexpected command received. fatal.")
	}
	switch c.commandtype {
	case commandSpawn:
		d := c.data.(uniter)
		if d.IsAlly() {
			s.play("summoning.mp3")
		} else {
			s.play("enemy_spawn.mp3")
		}

	case commandDead:
		d := c.data.(uniter)
		if d.IsAlly() {
			s.play("player_dead.mp3")
		} else {
			s.play("enemy_dead.mp3")
		}

	case commandRecall:
		s.play("summoning.mp3")

	case commandRecalled:
		s.play("recall.mp3")

	case commandUpdateSelection:
		s.play("select.mp3")

	case commandAttack:
		d := c.data.(uniter)
		switch d.GetUnitType() {
		case "player1":
			s.play("atk_sword.mp3")
		case "player2":
			s.play("atk_hammer.mp3")
		case "player3":
			s.play("atk_fire.mp3")
		case "enemy1":
			s.play("enemy_atk.mp3")
		case "enemy2":
			s.play("enemy_atk.mp3")
		}
	}
}
<commit_msg>sync attack timing and attack sound<commit_after>package scene

import (
	"github.com/pankona/gomo-simra/simra"

	"golang.org/x/mobile/asset"
)

type sound struct {
	game *game
}

func (s *sound) play(assetName string) {
	a := simra.NewAudio()
	resource, err := asset.Open(assetName)
	if err != nil {
		panic(err.Error())
	}
	err = a.Play(resource, false, func(err error) {
		if err != nil {
			simra.LogError(err.Error())
		}
	})
	if err != nil {
		panic(err.Error())
	}
}

func (s *sound) OnEvent(i interface{}) {
	c, ok := i.(*command)
	if !ok {
		panic("unexpected command received. fatal.")
	}
	switch c.commandtype {
	case commandSpawn:
		d := c.data.(uniter)
		if d.IsAlly() {
			s.play("summoning.mp3")
		} else {
			s.play("enemy_spawn.mp3")
		}

	case commandDead:
		d := c.data.(uniter)
		if d.IsAlly() {
			s.play("player_dead.mp3")
		} else {
			s.play("enemy_dead.mp3")
		}

	case commandRecall:
		s.play("summoning.mp3")

	case commandRecalled:
		s.play("recall.mp3")

	case commandUpdateSelection:
		s.play("select.mp3")

	case commandAttacking:
		d := c.data.(uniter)
		switch d.GetUnitType() {
		case "player1":
			s.play("atk_sword.mp3")
		case "player2":
			s.play("atk_hammer.mp3")
		case "player3":
			s.play("atk_fire.mp3")
		case "enemy1":
			s.play("enemy_atk.mp3")
		case "enemy2":
			s.play("enemy_atk.mp3")
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) Liam Stanley <me@liamstanley.io>. All rights reserved. Use
// of this source code is governed by the MIT license that can be found in
// the LICENSE file.

package girc

import "strings"

// CMode represents a single step of a given mode change.
type CMode struct {
	add     bool   // if it's a +, or -.
	name    byte   // character representation of the given mode.
	setting bool   // if it's a setting (should be stored) or temporary (op/voice/etc).
	args    string // arguments to the mode, if arguments are supported.
}

// Short returns a short representation of a mode without arguments. E.g. "+a",
// or "-b".
func (c *CMode) Short() string {
	var status string
	if c.add {
		status = "+"
	} else {
		status = "-"
	}

	return status + string(c.name)
}

// String returns a string representation of a mode, including optional
// arguments. E.g. "+b user*!ident@host.*.com"
func (c *CMode) String() string {
	if len(c.args) == 0 {
		return c.Short()
	}

	return c.Short() + " " + c.args
}

// CModes is a representation of a set of modes. This may be the given state
// of a channel/user, or the given state changes to a given channel/user.
type CModes struct {
	raw           string // raw supported modes.
	modesListArgs string // modes that add/remove users from lists and support args.
	modesArgs     string // modes that support args.
	modesSetArgs  string // modes that support args ONLY when set.
	modesNoArgs   string // modes that do not support args.

	prefixes string  // user permission prefixes. these aren't a CMode.setting.
	modes    []CMode // the list of modes for this given state.
}

// String returns a complete set of modes for this given state (change?). For
// example, "+a-b+cde some-arg".
func (c *CModes) String() string {
	var out string
	var args string

	if len(c.modes) > 0 {
		out += "+"
	}

	for i := 0; i < len(c.modes); i++ {
		out += string(c.modes[i].name)

		if len(c.modes[i].args) > 0 {
			args += " " + c.modes[i].args
		}
	}

	return out + args
}

// HasMode checks if the CModes state has a given mode. E.g. "m", or "I".
func (c *CModes) HasMode(mode string) bool {
	for i := 0; i < len(c.modes); i++ {
		if string(c.modes[i].name) == mode {
			return true
		}
	}

	return false
}

// hasArg checks to see if the mode supports arguments. What ones support this?:
//   A = Mode that adds or removes a nick or address to a list. Always has a parameter.
//   B = Mode that changes a setting and always has a parameter.
//   C = Mode that changes a setting and only has a parameter when set.
//   D = Mode that changes a setting and never has a parameter.
//   Note: Modes of type A return the list when there is no parameter present.
//   Note: Some clients assumes that any mode not listed is of type D.
//   Note: Modes in PREFIX are not listed but could be considered type B.
func (c *CModes) hasArg(set bool, mode byte) (hasArgs, isSetting bool) {
	if len(c.raw) < 1 {
		return false, true
	}

	if strings.IndexByte(c.modesListArgs, mode) > -1 {
		return true, false
	}

	if strings.IndexByte(c.modesArgs, mode) > -1 {
		return true, true
	}

	if strings.IndexByte(c.modesSetArgs, mode) > -1 {
		if set {
			return true, true
		}

		return false, true
	}

	if strings.IndexByte(c.prefixes, mode) > -1 {
		return true, false
	}

	return false, true
}

// apply merges two state changes, or one state change into a state of modes.
// For example, the latter would mean applying an incoming MODE with the modes
// stored for a channel.
func (c *CModes) apply(modes []CMode) {
	var new []CMode

	for j := 0; j < len(c.modes); j++ {
		isin := false
		for i := 0; i < len(modes); i++ {
			if !modes[i].setting {
				continue
			}
			if c.modes[j].name == modes[i].name && modes[i].add {
				new = append(new, modes[i])
				isin = true
				break
			}
		}

		if !isin {
			new = append(new, c.modes[j])
		}
	}

	for i := 0; i < len(modes); i++ {
		if !modes[i].setting || !modes[i].add {
			continue
		}

		isin := false
		for j := 0; j < len(new); j++ {
			if modes[i].name == new[j].name {
				isin = true
				break
			}
		}

		if !isin {
			new = append(new, modes[i])
		}
	}

	c.modes = new
}

// parse parses a set of flags and args, returning the necessary list of
// mappings for the mode flags.
func (c *CModes) parse(flags string, args []string) (out []CMode) {
	// add is the mode state we're currently in. Adding, or removing modes.
	add := true
	var argCount int

	for i := 0; i < len(flags); i++ {
		if flags[i] == 0x2B {
			add = true
			continue
		}
		if flags[i] == 0x2D {
			add = false
			continue
		}

		mode := CMode{
			name: flags[i],
			add:  add,
		}

		hasArgs, isSetting := c.hasArg(add, flags[i])
		if hasArgs && len(args) >= argCount+1 {
			mode.args = args[argCount]
			argCount++
		}
		mode.setting = isSetting

		out = append(out, mode)
	}

	return out
}

// newCModes returns a new CModes reference. channelModes and userPrefixes
// would be something you see from the server's "CHANMODES" and "PREFIX"
// ISUPPORT capability messages (alternatively, fall back to the standard)
// DefaultPrefixes and ModeDefaults.
func newCModes(channelModes, userPrefixes string) CModes {
	split := strings.SplitN(channelModes, ",", 4)
	if len(split) != 4 {
		for i := len(split); i < 4; i++ {
			split = append(split, "")
		}
	}

	return CModes{
		raw:           channelModes,
		modesListArgs: split[0],
		modesArgs:     split[1],
		modesSetArgs:  split[2],
		modesNoArgs:   split[3],

		prefixes: userPrefixes,
		modes:    []CMode{},
	}
}

// isValidChannelMode validates a channel mode (CHANMODES).
func isValidChannelMode(raw string) bool {
	if len(raw) < 1 {
		return false
	}

	for i := 0; i < len(raw); i++ {
		// Allowed are: ",", A-Z and a-z.
		if raw[i] != 0x2C && (raw[i] < 0x41 || raw[i] > 0x5A) && (raw[i] < 0x61 || raw[i] > 0x7A) {
			return false
		}
	}

	return true
}

// isValidUserPrefix validates a list of ISUPPORT-style user prefixes (PREFIX).
func isValidUserPrefix(raw string) bool {
	if len(raw) < 1 {
		return false
	}

	if raw[0] != 0x28 { // (.
		return false
	}

	var keys, rep int
	var passedKeys bool

	// Skip the first one as we know it's (.
	for i := 1; i < len(raw); i++ {
		if raw[i] == 0x29 { // ).
			passedKeys = true
			continue
		}

		if passedKeys {
			rep++
		} else {
			keys++
		}
	}

	return keys == rep
}

// parsePrefixes parses the mode character mappings from the symbols of a
// ISUPPORT-style user prefixes list (PREFIX).
func parsePrefixes(raw string) (modes, prefixes string) {
	if !isValidUserPrefix(raw) {
		return modes, prefixes
	}

	i := strings.Index(raw, ")")
	if i < 1 {
		return modes, prefixes
	}

	return raw[1:i], raw[i+1:]
}

// handleMODE handles incoming MODE messages, and updates the tracking
// information for each channel, as well as if any of the modes affect user
// permissions.
func handleMODE(c *Client, e Event) {
	// Check if it's a RPL_CHANNELMODEIS.
	if e.Command == RPL_CHANNELMODEIS && len(e.Params) > 2 {
		// RPL_CHANNELMODEIS sends the user as the first param, skip it.
		e.Params = e.Params[1:]
	}
	// Should be at least MODE <target> <flags>, to be useful. As well, only
	// tracking channel modes at the moment.
	if len(e.Params) < 2 || !IsValidChannel(e.Params[0]) {
		return
	}

	c.state.mu.Lock()
	channel := c.state.lookupChannel(e.Params[0])
	if channel == nil {
		c.state.mu.Unlock()
		return
	}

	flags := e.Params[1]
	var args []string
	if len(e.Params) > 2 {
		args = append(args, e.Params[2:]...)
	}

	modes := channel.Modes.parse(flags, args)
	channel.Modes.apply(modes)

	// Loop through and update users modes as necessary.
	for i := 0; i < len(modes); i++ {
		if modes[i].setting || len(modes[i].args) == 0 {
			continue
		}

		users := c.state.lookupUsers("nick", modes[i].args)
		for j := 0; j < len(users); j++ {
			users[j].Perms.setFromMode(modes[i])
		}
	}

	c.state.mu.Unlock()
}

// chanModes returns the ISUPPORT list of server-supported channel modes,
// alternatively falling back to ModeDefaults.
func (s *state) chanModes() string {
	if modes, ok := s.serverOptions["CHANMODES"]; ok && isValidChannelMode(modes) {
		return modes
	}

	return ModeDefaults
}

// userPrefixes returns the ISUPPORT list of server-supported user prefixes.
// This includes mode characters, as well as user prefix symbols. Falls back
// to DefaultPrefixes if not server-supported.
func (s *state) userPrefixes() string {
	if prefix, ok := s.serverOptions["PREFIX"]; ok && isValidUserPrefix(prefix) {
		return prefix
	}

	return DefaultPrefixes
}

// UserPerms contains all channel-based user permissions. The minimum op, and
// voice should be supported on all networks. This also supports non-rfc
// Owner, Admin, and HalfOp, if the network has support for it.
type UserPerms struct {
	// Owner (non-rfc) indicates that the user has full permissions to the
	// channel. More than one user can have owner permission.
	Owner bool
	// Admin (non-rfc) is commonly given to users that are trusted enough
	// to manage channel permissions, as well as higher level service settings.
	Admin bool
	// Op is commonly given to trusted users who can manage a given channel
	// by kicking, and banning users.
	Op bool
	// HalfOp (non-rfc) is commonly used to give users permissions like the
	// ability to kick, without giving them greater abilities to ban all users.
	HalfOp bool
	// Voice indicates the user has voice permissions, commonly given to known
	// users, wih very light trust, or to indicate a user is active.
	Voice bool
}

// IsAdmin indicates that the user has banning abilities, and are likely a
// very trustable user (e.g. op+).
func (m UserPerms) IsAdmin() bool {
	if m.Owner || m.Admin || m.Op {
		return true
	}

	return false
}

// IsAdmin indicates that the user at least has modes set upon them, higher
// than a regular joining user.
func (m UserPerms) IsTrusted() bool {
	if m.IsAdmin() || m.HalfOp || m.Voice {
		return true
	}

	return false
}

// reset resets the modes of a user.
func (m *UserPerms) reset() {
	m.Owner = false
	m.Admin = false
	m.Op = false
	m.HalfOp = false
	m.Voice = false
}

// set translates raw prefix characters into proper permissions. Only
// use this function when you have a session lock.
func (m *UserPerms) set(prefix string, append bool) {
	if !append {
		m.reset()
	}

	for i := 0; i < len(prefix); i++ {
		switch string(prefix[i]) {
		case OwnerPrefix:
			m.Owner = true
		case AdminPrefix:
			m.Admin = true
		case OperatorPrefix:
			m.Op = true
		case HalfOperatorPrefix:
			m.HalfOp = true
		case VoicePrefix:
			m.Voice = true
		}
	}
}

// setFromMode sets user-permissions based on channel user mode chars. E.g.
// "o" being oper, "v" being voice, etc.
func (m *UserPerms) setFromMode(mode CMode) {
	switch string(mode.name) {
	case ModeOwner:
		m.Owner = mode.add
	case ModeAdmin:
		m.Admin = mode.add
	case ModeOperator:
		m.Op = mode.add
	case ModeHalfOperator:
		m.HalfOp = mode.add
	case ModeVoice:
		m.Voice = mode.add
	}
}

// parseUserPrefix parses a raw mode line, like "@user" or "@+user".
func parseUserPrefix(raw string) (modes, nick string, success bool) {
	for i := 0; i < len(raw); i++ {
		char := string(raw[i])

		if char == OwnerPrefix || char == AdminPrefix || char == HalfOperatorPrefix ||
			char == OperatorPrefix || char == VoicePrefix {
			modes += char
			continue
		}

		// Assume we've gotten to the nickname part.
		return modes, raw[i:], true
	}

	return
}
<commit_msg>implement CModes.Get()<commit_after>// Copyright (c) Liam Stanley <me@liamstanley.io>. All rights reserved. Use
// of this source code is governed by the MIT license that can be found in
// the LICENSE file.

package girc

import "strings"

// CMode represents a single step of a given mode change.
type CMode struct {
	add     bool   // if it's a +, or -.
	name    byte   // character representation of the given mode.
	setting bool   // if it's a setting (should be stored) or temporary (op/voice/etc).
	args    string // arguments to the mode, if arguments are supported.
}

// Short returns a short representation of a mode without arguments. E.g. "+a",
// or "-b".
func (c *CMode) Short() string {
	var status string
	if c.add {
		status = "+"
	} else {
		status = "-"
	}

	return status + string(c.name)
}

// String returns a string representation of a mode, including optional
// arguments. E.g. "+b user*!ident@host.*.com"
func (c *CMode) String() string {
	if len(c.args) == 0 {
		return c.Short()
	}

	return c.Short() + " " + c.args
}

// CModes is a representation of a set of modes. This may be the given state
// of a channel/user, or the given state changes to a given channel/user.
type CModes struct {
	raw           string // raw supported modes.
	modesListArgs string // modes that add/remove users from lists and support args.
	modesArgs     string // modes that support args.
	modesSetArgs  string // modes that support args ONLY when set.
	modesNoArgs   string // modes that do not support args.

	prefixes string  // user permission prefixes. these aren't a CMode.setting.
	modes    []CMode // the list of modes for this given state.
}

// String returns a complete set of modes for this given state (change?). For
// example, "+a-b+cde some-arg".
func (c *CModes) String() string {
	var out string
	var args string

	if len(c.modes) > 0 {
		out += "+"
	}

	for i := 0; i < len(c.modes); i++ {
		out += string(c.modes[i].name)

		if len(c.modes[i].args) > 0 {
			args += " " + c.modes[i].args
		}
	}

	return out + args
}

// HasMode checks if the CModes state has a given mode. E.g. "m", or "I".
func (c *CModes) HasMode(mode string) bool {
	for i := 0; i < len(c.modes); i++ {
		if string(c.modes[i].name) == mode {
			return true
		}
	}

	return false
}

// Get returns the arguments for a given mode within this session, if it
// supports args.
func (c *CModes) Get(mode string) (args string, ok bool) {
	for i := 0; i < len(c.modes); i++ {
		if string(c.modes[i].name) == mode {
			if len(c.modes[i].args) == 0 {
				return "", false
			}

			return c.modes[i].args, true
		}
	}

	return "", false
}

// hasArg checks to see if the mode supports arguments. What ones support this?:
//   A = Mode that adds or removes a nick or address to a list. Always has a parameter.
//   B = Mode that changes a setting and always has a parameter.
//   C = Mode that changes a setting and only has a parameter when set.
//   D = Mode that changes a setting and never has a parameter.
//   Note: Modes of type A return the list when there is no parameter present.
//   Note: Some clients assumes that any mode not listed is of type D.
//   Note: Modes in PREFIX are not listed but could be considered type B.
func (c *CModes) hasArg(set bool, mode byte) (hasArgs, isSetting bool) {
	if len(c.raw) < 1 {
		return false, true
	}

	if strings.IndexByte(c.modesListArgs, mode) > -1 {
		return true, false
	}

	if strings.IndexByte(c.modesArgs, mode) > -1 {
		return true, true
	}

	if strings.IndexByte(c.modesSetArgs, mode) > -1 {
		if set {
			return true, true
		}

		return false, true
	}

	if strings.IndexByte(c.prefixes, mode) > -1 {
		return true, false
	}

	return false, true
}

// apply merges two state changes, or one state change into a state of modes.
// For example, the latter would mean applying an incoming MODE with the modes
// stored for a channel.
func (c *CModes) apply(modes []CMode) {
	var new []CMode

	for j := 0; j < len(c.modes); j++ {
		isin := false
		for i := 0; i < len(modes); i++ {
			if !modes[i].setting {
				continue
			}
			if c.modes[j].name == modes[i].name && modes[i].add {
				new = append(new, modes[i])
				isin = true
				break
			}
		}

		if !isin {
			new = append(new, c.modes[j])
		}
	}

	for i := 0; i < len(modes); i++ {
		if !modes[i].setting || !modes[i].add {
			continue
		}

		isin := false
		for j := 0; j < len(new); j++ {
			if modes[i].name == new[j].name {
				isin = true
				break
			}
		}

		if !isin {
			new = append(new, modes[i])
		}
	}

	c.modes = new
}

// parse parses a set of flags and args, returning the necessary list of
// mappings for the mode flags.
func (c *CModes) parse(flags string, args []string) (out []CMode) {
	// add is the mode state we're currently in. Adding, or removing modes.
	add := true
	var argCount int

	for i := 0; i < len(flags); i++ {
		if flags[i] == 0x2B {
			add = true
			continue
		}
		if flags[i] == 0x2D {
			add = false
			continue
		}

		mode := CMode{
			name: flags[i],
			add:  add,
		}

		hasArgs, isSetting := c.hasArg(add, flags[i])
		if hasArgs && len(args) >= argCount+1 {
			mode.args = args[argCount]
			argCount++
		}
		mode.setting = isSetting

		out = append(out, mode)
	}

	return out
}

// newCModes returns a new CModes reference. channelModes and userPrefixes
// would be something you see from the server's "CHANMODES" and "PREFIX"
// ISUPPORT capability messages (alternatively, fall back to the standard)
// DefaultPrefixes and ModeDefaults.
func newCModes(channelModes, userPrefixes string) CModes {
	split := strings.SplitN(channelModes, ",", 4)
	if len(split) != 4 {
		for i := len(split); i < 4; i++ {
			split = append(split, "")
		}
	}

	return CModes{
		raw:           channelModes,
		modesListArgs: split[0],
		modesArgs:     split[1],
		modesSetArgs:  split[2],
		modesNoArgs:   split[3],

		prefixes: userPrefixes,
		modes:    []CMode{},
	}
}

// isValidChannelMode validates a channel mode (CHANMODES).
func isValidChannelMode(raw string) bool {
	if len(raw) < 1 {
		return false
	}

	for i := 0; i < len(raw); i++ {
		// Allowed are: ",", A-Z and a-z.
		if raw[i] != 0x2C && (raw[i] < 0x41 || raw[i] > 0x5A) && (raw[i] < 0x61 || raw[i] > 0x7A) {
			return false
		}
	}

	return true
}

// isValidUserPrefix validates a list of ISUPPORT-style user prefixes (PREFIX).
func isValidUserPrefix(raw string) bool {
	if len(raw) < 1 {
		return false
	}

	if raw[0] != 0x28 { // (.
		return false
	}

	var keys, rep int
	var passedKeys bool

	// Skip the first one as we know it's (.
	for i := 1; i < len(raw); i++ {
		if raw[i] == 0x29 { // ).
			passedKeys = true
			continue
		}

		if passedKeys {
			rep++
		} else {
			keys++
		}
	}

	return keys == rep
}

// parsePrefixes parses the mode character mappings from the symbols of a
// ISUPPORT-style user prefixes list (PREFIX).
func parsePrefixes(raw string) (modes, prefixes string) {
	if !isValidUserPrefix(raw) {
		return modes, prefixes
	}

	i := strings.Index(raw, ")")
	if i < 1 {
		return modes, prefixes
	}

	return raw[1:i], raw[i+1:]
}

// handleMODE handles incoming MODE messages, and updates the tracking
// information for each channel, as well as if any of the modes affect user
// permissions.
func handleMODE(c *Client, e Event) {
	// Check if it's a RPL_CHANNELMODEIS.
	if e.Command == RPL_CHANNELMODEIS && len(e.Params) > 2 {
		// RPL_CHANNELMODEIS sends the user as the first param, skip it.
		e.Params = e.Params[1:]
	}
	// Should be at least MODE <target> <flags>, to be useful. As well, only
	// tracking channel modes at the moment.
	if len(e.Params) < 2 || !IsValidChannel(e.Params[0]) {
		return
	}

	c.state.mu.Lock()
	channel := c.state.lookupChannel(e.Params[0])
	if channel == nil {
		c.state.mu.Unlock()
		return
	}

	flags := e.Params[1]
	var args []string
	if len(e.Params) > 2 {
		args = append(args, e.Params[2:]...)
	}

	modes := channel.Modes.parse(flags, args)
	channel.Modes.apply(modes)

	// Loop through and update users modes as necessary.
	for i := 0; i < len(modes); i++ {
		if modes[i].setting || len(modes[i].args) == 0 {
			continue
		}

		users := c.state.lookupUsers("nick", modes[i].args)
		for j := 0; j < len(users); j++ {
			users[j].Perms.setFromMode(modes[i])
		}
	}

	c.state.mu.Unlock()
}

// chanModes returns the ISUPPORT list of server-supported channel modes,
// alternatively falling back to ModeDefaults.
func (s *state) chanModes() string {
	if modes, ok := s.serverOptions["CHANMODES"]; ok && isValidChannelMode(modes) {
		return modes
	}

	return ModeDefaults
}

// userPrefixes returns the ISUPPORT list of server-supported user prefixes.
// This includes mode characters, as well as user prefix symbols. Falls back
// to DefaultPrefixes if not server-supported.
func (s *state) userPrefixes() string {
	if prefix, ok := s.serverOptions["PREFIX"]; ok && isValidUserPrefix(prefix) {
		return prefix
	}

	return DefaultPrefixes
}

// UserPerms contains all channel-based user permissions. The minimum op, and
// voice should be supported on all networks. This also supports non-rfc
// Owner, Admin, and HalfOp, if the network has support for it.
type UserPerms struct {
	// Owner (non-rfc) indicates that the user has full permissions to the
	// channel. More than one user can have owner permission.
	Owner bool
	// Admin (non-rfc) is commonly given to users that are trusted enough
	// to manage channel permissions, as well as higher level service settings.
	Admin bool
	// Op is commonly given to trusted users who can manage a given channel
	// by kicking, and banning users.
	Op bool
	// HalfOp (non-rfc) is commonly used to give users permissions like the
	// ability to kick, without giving them greater abilities to ban all users.
	HalfOp bool
	// Voice indicates the user has voice permissions, commonly given to known
	// users, wih very light trust, or to indicate a user is active.
	Voice bool
}

// IsAdmin indicates that the user has banning abilities, and are likely a
// very trustable user (e.g. op+).
func (m UserPerms) IsAdmin() bool {
	if m.Owner || m.Admin || m.Op {
		return true
	}

	return false
}

// IsAdmin indicates that the user at least has modes set upon them, higher
// than a regular joining user.
func (m UserPerms) IsTrusted() bool {
	if m.IsAdmin() || m.HalfOp || m.Voice {
		return true
	}

	return false
}

// reset resets the modes of a user.
func (m *UserPerms) reset() {
	m.Owner = false
	m.Admin = false
	m.Op = false
	m.HalfOp = false
	m.Voice = false
}

// set translates raw prefix characters into proper permissions. Only
// use this function when you have a session lock.
func (m *UserPerms) set(prefix string, append bool) {
	if !append {
		m.reset()
	}

	for i := 0; i < len(prefix); i++ {
		switch string(prefix[i]) {
		case OwnerPrefix:
			m.Owner = true
		case AdminPrefix:
			m.Admin = true
		case OperatorPrefix:
			m.Op = true
		case HalfOperatorPrefix:
			m.HalfOp = true
		case VoicePrefix:
			m.Voice = true
		}
	}
}

// setFromMode sets user-permissions based on channel user mode chars. E.g.
// "o" being oper, "v" being voice, etc.
func (m *UserPerms) setFromMode(mode CMode) {
	switch string(mode.name) {
	case ModeOwner:
		m.Owner = mode.add
	case ModeAdmin:
		m.Admin = mode.add
	case ModeOperator:
		m.Op = mode.add
	case ModeHalfOperator:
		m.HalfOp = mode.add
	case ModeVoice:
		m.Voice = mode.add
	}
}

// parseUserPrefix parses a raw mode line, like "@user" or "@+user".
func parseUserPrefix(raw string) (modes, nick string, success bool) {
	for i := 0; i < len(raw); i++ {
		char := string(raw[i])

		if char == OwnerPrefix || char == AdminPrefix || char == HalfOperatorPrefix ||
			char == OperatorPrefix || char == VoicePrefix {
			modes += char
			continue
		}

		// Assume we've gotten to the nickname part.
		return modes, raw[i:], true
	}

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package endtoend

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/stretchr/testify/assert"

	"vitess.io/vitess/go/sqltypes"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/endtoend/framework"
)

func TestStreamUnion(t *testing.T) {
	qr, err := framework.NewClient().StreamExecute("select 1 from dual union select 1 from dual", nil)
	if err != nil {
		t.Error(err)
		return
	}
	assert.Equal(t, 1, len(qr.Rows))
}

func TestStreamConsolidation(t *testing.T) {
	const Workers = 50
	const RowCount = 1100
	const RowContent = "abcdefghijklmnopqrstuvwxyz"

	populate := func(client *framework.QueryClient) error {
		err := client.Begin(false)
		if err != nil {
			return err
		}
		defer client.Rollback()

		for i := 0; i < RowCount; i++ {
			query := fmt.Sprintf("insert into vitess_stress values (%d, '%s')", i, strings.Repeat(RowContent, 2048/len(RowContent)))
			_, err := client.Execute(query, nil)
			if err != nil {
				return err
			}
		}
		return client.Commit()
	}

	client := framework.NewClient()
	err := populate(client)
	require.NoError(t, err)
	defer client.Execute("delete from vitess_stress", nil)

	defaultPoolSize := framework.Server.StreamPoolSize()

	framework.Server.SetStreamPoolSize(4)
	framework.Server.SetStreamConsolidationBlocking(true)

	defer func() {
		framework.Server.SetStreamPoolSize(defaultPoolSize)
		framework.Server.SetStreamConsolidationBlocking(false)
	}()

	var start = make(chan struct{})
	var finish sync.WaitGroup

	// Spawn N workers at the same time to stress test the stream consolidator
	for i := 0; i < Workers; i++ {
		finish.Add(1)
		go func() {
			defer finish.Done()

			// block all the workers so they all perform their queries at the same time
			<-start

			var rowCount int
			err := client.Stream("select * from vitess_stress", nil, func(result *sqltypes.Result) error {
				for _, r := range result.Rows {
					rowCount += len(r)
				}
				return nil
			})
			require.NoError(t, err)
			require.Equal(t, 2200, rowCount)
		}()
	}

	// wait until all the goroutines have spawned and are blocked before we unblock them at once
	time.Sleep(500 * time.Millisecond)
	close(start)
	finish.Wait()
}

func TestStreamBigData(t *testing.T) {
	client := framework.NewClient()
	err := populateBigData(client)
	if err != nil {
		t.Error(err)
		return
	}
	defer client.Execute("delete from vitess_big", nil)

	qr, err := client.StreamExecute("select * from vitess_big b1, vitess_big b2 order by b1.id, b2.id", nil)
	if err != nil {
		t.Error(err)
		return
	}
	row10 := framework.RowsToStrings(qr)[10]
	want := []string{
		"0",
		"AAAAAAAAAAAAAAAAAA 0",
		"BBBBBBBBBBBBBBBBBB 0",
		"C",
		"DDDDDDDDDDDDDDDDDD 0",
		"EEEEEEEEEEEEEEEEEE 0",
		"FF 0",
		"GGGGGGGGGGGGGGGGGG 0",
		"0",
		"0",
		"0",
		"0",
		"10",
		"AAAAAAAAAAAAAAAAAA 10",
		"BBBBBBBBBBBBBBBBBB 10",
		"C",
		"DDDDDDDDDDDDDDDDDD 10",
		"EEEEEEEEEEEEEEEEEE 10",
		"FF 10",
		"GGGGGGGGGGGGGGGGGG 10",
		"10",
		"10",
		"10",
		"10"}
	if !reflect.DeepEqual(row10, want) {
		t.Errorf("Row10: \n%#v, want \n%#v", row10, want)
	}
}

func TestStreamTerminate(t *testing.T) {
	client := framework.NewClient()
	err := populateBigData(client)
	if err != nil {
		t.Error(err)
		return
	}
	defer client.Execute("delete from vitess_big", nil)

	called := false
	err = client.Stream(
		"select * from vitess_big b1, vitess_big b2 order by b1.id, b2.id",
		nil,
		func(*sqltypes.Result) error {
			if !called {
				queries := framework.LiveQueryz()
				if l := len(queries); l != 1 {
					t.Errorf("len(queries): %d, want 1", l)
					return errors.New("no queries from LiveQueryz")
				}
				err := framework.StreamTerminate(queries[0].ConnID)
				if err != nil {
					return err
				}
				called = true
			}
			time.Sleep(10 * time.Millisecond)
			return nil
		},
	)
	if code := vterrors.Code(err); code != vtrpcpb.Code_CANCELED {
		t.Errorf("Errorcode: %v, want %v", code, vtrpcpb.Code_CANCELED)
	}
}

func populateBigData(client *framework.QueryClient) error {
	err := client.Begin(false)
	if err != nil {
		return err
	}
	defer client.Rollback()

	for i := 0; i < 100; i++ {
		stri := strconv.Itoa(i)
		query := "insert into vitess_big values " +
			"(" + stri + ", " +
			"'AAAAAAAAAAAAAAAAAA " + stri + "', " +
			"'BBBBBBBBBBBBBBBBBB " + stri + "', " +
			"'C', " +
			"'DDDDDDDDDDDDDDDDDD " + stri + "', " +
			"'EEEEEEEEEEEEEEEEEE " + stri + "', " +
			"'FF " + stri + "', " +
			"'GGGGGGGGGGGGGGGGGG " + stri + "', " +
			stri + ", " +
			stri + ", " +
			stri + ", " +
			stri + ")"
		_, err := client.Execute(query, nil)
		if err != nil {
			return err
		}
	}
	return client.Commit()
}

func TestStreamError(t *testing.T) {
	_, err := framework.NewClient().StreamExecute("select count(abcd) from vitess_big", nil)
	want := "Unknown column"
	if err == nil || !strings.HasPrefix(err.Error(), want) {
		t.Errorf("Error: %v, must start with %s", err, want)
	}
}
<commit_msg>endtoend: add streaming benchmark<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package endtoend

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/stretchr/testify/assert"

	"vitess.io/vitess/go/sqltypes"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/endtoend/framework"
)

func TestStreamUnion(t *testing.T) {
	qr, err := framework.NewClient().StreamExecute("select 1 from dual union select 1 from dual", nil)
	if err != nil {
		t.Error(err)
		return
	}
	assert.Equal(t, 1, len(qr.Rows))
}

func populateStressQuery(client *framework.QueryClient, rowCount int, rowContent string) error {
	err := client.Begin(false)
	if err != nil {
		return err
	}
	defer client.Rollback()

	for i := 0; i < rowCount; i++ {
		query := fmt.Sprintf("insert into vitess_stress values (%d, '%s')", i, strings.Repeat(rowContent, 2048/len(rowContent)))
		_, err := client.Execute(query, nil)
		if err != nil {
			return err
		}
	}
	return client.Commit()
}

func BenchmarkStreamQuery(b *testing.B) {
	const RowCount = 1100
	const RowContent = "abcdefghijklmnopqrstuvwxyz"

	client := framework.NewClient()
	err := populateStressQuery(client, RowCount, RowContent)
	if err != nil {
		b.Fatal(err)
	}
	defer client.Execute("delete from vitess_stress", nil)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		err := client.Stream("select * from vitess_stress", nil, func(result *sqltypes.Result) error {
			return nil
		})
		if err != nil {
			b.Fatal(err)
		}
	}
}

func TestStreamConsolidation(t *testing.T) {
	const Workers = 50
	const RowCount = 1100
	const RowContent = "abcdefghijklmnopqrstuvwxyz"

	client := framework.NewClient()
	err := populateStressQuery(client, RowCount, RowContent)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Execute("delete from vitess_stress", nil)

	defaultPoolSize := framework.Server.StreamPoolSize()

	framework.Server.SetStreamPoolSize(4)
	framework.Server.SetStreamConsolidationBlocking(true)

	defer func() {
		framework.Server.SetStreamPoolSize(defaultPoolSize)
		framework.Server.SetStreamConsolidationBlocking(false)
	}()

	var start = make(chan struct{})
	var finish sync.WaitGroup

	// Spawn N workers at the same time to stress test the stream consolidator
	for i := 0; i < Workers; i++ {
		finish.Add(1)
		go func() {
			defer finish.Done()

			// block all the workers so they all perform their queries at the same time
			<-start

			var rowCount int
			err := client.Stream("select * from vitess_stress", nil, func(result *sqltypes.Result) error {
				for _, r := range result.Rows {
					rowCount += len(r)
				}
				return nil
			})
			require.NoError(t, err)
			require.Equal(t, 2200, rowCount)
		}()
	}

	// wait until all the goroutines have spawned and are blocked before we unblock them at once
	time.Sleep(500 * time.Millisecond)
	close(start)
	finish.Wait()
}

func TestStreamBigData(t *testing.T) {
	client := framework.NewClient()
	err := populateBigData(client)
	if err != nil {
		t.Error(err)
		return
	}
	defer client.Execute("delete from vitess_big", nil)

	qr, err := client.StreamExecute("select * from vitess_big b1, vitess_big b2 order by b1.id, b2.id", nil)
	if err != nil {
		t.Error(err)
		return
	}
	row10 := framework.RowsToStrings(qr)[10]
	want := []string{
		"0",
		"AAAAAAAAAAAAAAAAAA 0",
		"BBBBBBBBBBBBBBBBBB 0",
		"C",
		"DDDDDDDDDDDDDDDDDD 0",
		"EEEEEEEEEEEEEEEEEE 0",
		"FF 0",
		"GGGGGGGGGGGGGGGGGG 0",
		"0",
		"0",
		"0",
		"0",
		"10",
		"AAAAAAAAAAAAAAAAAA 10",
		"BBBBBBBBBBBBBBBBBB 10",
		"C",
		"DDDDDDDDDDDDDDDDDD 10",
		"EEEEEEEEEEEEEEEEEE 10",
		"FF 10",
		"GGGGGGGGGGGGGGGGGG 10",
		"10",
		"10",
		"10",
		"10"}
	if !reflect.DeepEqual(row10, want) {
		t.Errorf("Row10: \n%#v, want \n%#v", row10, want)
	}
}

func TestStreamTerminate(t *testing.T) {
	client := framework.NewClient()
	err := populateBigData(client)
	if err != nil {
		t.Error(err)
		return
	}
	defer client.Execute("delete from vitess_big", nil)

	called := false
	err = client.Stream(
		"select * from vitess_big b1, vitess_big b2 order by b1.id, b2.id",
		nil,
		func(*sqltypes.Result) error {
			if !called {
				queries := framework.LiveQueryz()
				if l := len(queries); l != 1 {
					t.Errorf("len(queries): %d, want 1", l)
					return errors.New("no queries from LiveQueryz")
				}
				err := framework.StreamTerminate(queries[0].ConnID)
				if err != nil {
					return err
				}
				called = true
			}
			time.Sleep(10 * time.Millisecond)
			return nil
		},
	)
	if code := vterrors.Code(err); code != vtrpcpb.Code_CANCELED {
		t.Errorf("Errorcode: %v, want %v", code, vtrpcpb.Code_CANCELED)
	}
}

func populateBigData(client *framework.QueryClient) error {
	err := client.Begin(false)
	if err != nil {
		return err
	}
	defer client.Rollback()

	for i := 0; i < 100; i++ {
		stri := strconv.Itoa(i)
		query := "insert into vitess_big values " +
			"(" + stri + ", " +
			"'AAAAAAAAAAAAAAAAAA " + stri + "', " +
			"'BBBBBBBBBBBBBBBBBB " + stri + "', " +
			"'C', " +
			"'DDDDDDDDDDDDDDDDDD " + stri + "', " +
			"'EEEEEEEEEEEEEEEEEE " + stri + "', " +
			"'FF " + stri + "', " +
			"'GGGGGGGGGGGGGGGGGG " + stri + "', " +
			stri + ", " +
			stri + ", " +
			stri + ", " +
			stri + ")"
		_, err := client.Execute(query, nil)
		if err != nil {
			return err
		}
	}
	return client.Commit()
}

func TestStreamError(t *testing.T) {
	_, err := framework.NewClient().StreamExecute("select count(abcd) from vitess_big", nil)
	want := "Unknown column"
	if err == nil || !strings.HasPrefix(err.Error(), want) {
		t.Errorf("Error: %v, must start with %s", err, want)
	}
}
<|endoftext|>