<commit_before>package goboots

import (
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"errors"
	"github.com/gabstv/i18ngo"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"
	"time"
)

const (
	ErrNil = 0
	// TLS
	ErrTLSNil = 100
)

var (
	staticControllers []IController
	sessionDbs        map[string]ISessionDBEngine
	curSessionDb      ISessionDBEngine
	httpErrorStrings  = map[int]string{
		400: "Bad Request",
		401: "Unauthorized",
		403: "Forbidden",
		404: "Not Found",
		405: "Method Not Allowed",
		406: "Not Acceptable",
		501: "Internal Server Error",
	}
)

type Params map[string]string

func hostonly(hostport string) string {
	h, _, _ := net.SplitHostPort(hostport)
	return h
}

func trimhost(hostport string) string {
	return strings.Split(hostport, ":")[0]
}

func getTLSRedirectURL(hostaddrtls string, uri *url.URL) (string, error) {
	if uri == nil {
		return "", errors.New("url is null")
	}
	_, p, err := net.SplitHostPort(hostaddrtls)
	if err != nil {
		return "", err
	}
	if uri.Scheme == "ws" {
		uri.Scheme = "wss"
	} else if uri.Scheme == "http" {
		uri.Scheme = "https"
	}
	if p != "443" {
		uri.Host = uri.Host + ":" + p
	}
	return uri.String(), nil
}

func RegisterSessionStorageDriver(name string, engine ISessionDBEngine) {
	if sessionDbs == nil {
		sessionDbs = make(map[string]ISessionDBEngine, 0)
	}
	sessionDbs[name] = engine
}

func RegisterControllerGlobal(controller IController) {
	if staticControllers == nil {
		staticControllers = make([]IController, 0)
	}
	staticControllers = append(staticControllers, controller)
}

func (app *App) InitSessionStorage(driver string) error {
	if sessionDbs == nil {
		return errors.New("No session storage registered.")
	}
	if _, ok := sessionDbs[driver]; !ok {
		return errors.New("The session storage driver " + driver + " does not exist.")
	}
	curSessionDb = sessionDbs[driver]
	curSessionDb.SetApp(app)
	return nil
}

func CloseSessionStorage() {
	if curSessionDb != nil {
		curSessionDb.Close()
	}
}

func initAnySessionStorage() {
	if curSessionDb != nil {
		return
	}
	if sessionDbs == nil {
		curSessionDb = &SessionDevNull{}
		return
	}
	for _, v := range sessionDbs {
		curSessionDb = v
		break
	}
}

type AppError struct {
	Id      int
	Message string
}

func (e *AppError) Error() string {
	return e.Message
}

type env struct {
	path       string
	controller string
	action     string
	params     []string
	request    *http.Request
	body       string
	form       map[string][]string
}

type ISession interface {
	GetData() (string, map[string]interface{}, time.Time)
}

type Session struct {
	SID     string
	Data    map[string]interface{}
	Flash   SessFlash `json:"-" bson:"-"` // never save flash
	Time    time.Time
	Updated time.Time
	r       *http.Request
	w       http.ResponseWriter
}

type SessFlash struct {
	vals map[string]interface{}
}

func (s *SessFlash) init() {
	if s.vals == nil {
		s.vals = make(map[string]interface{})
	}
}

func (s *SessFlash) Set(key string, val interface{}) {
	s.init()
	s.vals[key] = val
}

func (s *SessFlash) Get(key string) interface{} {
	s.init()
	return s.vals[key]
}

func (s *SessFlash) Get2(key string) (interface{}, bool) {
	s.init()
	v, ok := s.vals[key]
	return v, ok
}

func (s *SessFlash) Del(key string) {
	s.init()
	delete(s.vals, key)
}

func (s *SessFlash) Clear() {
	s.vals = nil
}

func (s *SessFlash) All() map[string]interface{} {
	s.init()
	return s.vals
}

func (s *Session) GetData() (string, map[string]interface{}, time.Time) {
	return s.SID, s.Data, s.Time
}

func (s *Session) Flush() {
	FlushSession(s)
}

func (s *Session) GetBool(key string) (bool, bool) {
	iface, ok := s.Data[key]
	if !ok {
		return false, false
	}
	return bool(reflect.ValueOf(iface).Bool()), true
}

func (s *Session) GetBoolD(key string, defaultValue bool) bool {
	iface, ok := s.Data[key]
	if !ok {
		return defaultValue
	}
	return bool(reflect.ValueOf(iface).Bool())
}

func (s *Session) GetInt32(key string) (int, bool) {
	if s.Data == nil {
		return 0, false
	}
	iface, ok := s.Data[key]
	if !ok {
		return 0, false
	}
	return int(reflect.ValueOf(iface).Int()), true
}

func (s *Session) GetInt32D(key string, defaultValue int) int {
	if s.Data == nil {
		return defaultValue
	}
	iface, ok := s.Data[key]
	if !ok {
		return defaultValue
	}
	return int(reflect.ValueOf(iface).Int())
}

func (s *Session) GetString(key string) (string, bool) {
	if s.Data == nil {
		return "", false
	}
	iface, ok := s.Data[key]
	if !ok {
		return "", false
	}
	return reflect.ValueOf(iface).String(), true
}

func (s *Session) GetStringD(key string, defaultValue string) string {
	if s.Data == nil {
		return defaultValue
	}
	iface, ok := s.Data[key]
	if !ok {
		return defaultValue
	}
	return reflect.ValueOf(iface).String()
}

func (s *Session) Expire(t time.Time) {
	SetCookieAdv(s.w, "goboots_sessid", s.SID, "/", "", t, 0, false, true)
}

func (s *Session) GetExpires() time.Time {
	cookie, err := s.r.Cookie("goboots_sessid")
	if err != nil {
		log.Println("Error @ (s *Session) GetExpires()", err)
		return time.Now()
	}
	return cookie.Expires
}

func (app *App) GetSession(w http.ResponseWriter, r *http.Request) *Session {
	if w == nil || r == nil {
		return nil
	}
	initAnySessionStorage()
	var cookie *http.Cookie
	var sid string
	var err error
	//LOAD
	cookie, err = r.Cookie("goboots_sessid")
	if err == nil {
		sid = cookie.Value
		if _validateSidString(sid) {
			var msession *Session
			// get from a saved location
			msession, err = curSessionDb.GetSession(sid)
			if err == nil {
				msession.r = r
				msession.w = w
				return msession
			}
			app.Logger.Printf("SESSION ERROR :( [%s] %s\n", sid, err.Error())
			// not found, generate a new sid
		} else {
			app.Logger.Printf("COULD NOT VALIDATE :( [%s]\n", sid)
		}
	}
	// gen session
	sid = uuid.New()

	session := &Session{
		SID:     sid,
		Data:    make(map[string]interface{}),
		Time:    time.Now(),
		Updated: time.Now(),
	}
	err = curSessionDb.NewSession(session)
	if err != nil {
		log.Println("[FATAL] [curSessionDb.NewSession] Could not get session!", err)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("The server encountered an error while processing your request [ERR_CONN_NEW_SESSION]."))
		return nil
	}
	SetCookieAdv(w, "goboots_sessid", sid, "/", "", time.Now().AddDate(0, 1, 0), 0, false, true)
	session.w = w
	session.r = r
	return session
}

func GetUserLang(w http.ResponseWriter, r *http.Request) string {
	var l string
	// get lang from cookies
	cookie, err := r.Cookie("lang")
	if err != nil && cookie != nil {
		l = cookie.Value
		return l
	}
	// get lang from http request
	alh := r.Header.Get("Accept-Language")
	// 2013-07-29 : not all clients actually send this header (googlebot/wget etc)
	if len(alh) > 1 {
		validlangs := i18ngo.GetLanguageCodes()
		// break alh
		alh0 := strings.Split(alh, ",")
		for _, v := range alh0 {
			lcd := strings.Split(v, ";")[0]
			var lc string
			if len(lcd) >= 2 {
				lc = lcd[0:2]
			}
			lc = lcd
			if StringIndexOf(validlangs, lc) != -1 {
				l = lc
				break
			}
		}
		if len(l) > 0 {
			// just update the cookies
			if w != nil {
				SetCookieSimple(w, "lang", l)
			}
		} else {
			l = i18ngo.GetDefaultLanguageCode()
		}
		return l
	}
	return i18ngo.GetDefaultLanguageCode()
}

func SetUserLang(w http.ResponseWriter, r *http.Request, langcode string) {
	SetCookieSimple(w, "lang", langcode)
}

func FlushSession(s *Session) error {
	return curSessionDb.PutSession(s)
}

func DestroySession(w http.ResponseWriter, r *http.Request, s *Session) {
	curSessionDb.RemoveSession(s)
	SetCookieAdv(w, "goboots_sessid", "", "/", "", time.Now(), 1, false, true)
}

// DEPRECATED
func _validateSidString(sid string) bool {
	if len(sid) < 32 {
		return false
	}
	for _, c := range sid {
		// check if sid contains invalid chars
		if !((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || (c == '-')) {
			return false
		}
	}
	return true
}

type templateInfo struct {
	path       string
	data       *template.Template
	lastUpdate time.Time
}

func SetCookieAdv(w http.ResponseWriter, name string, value string, path string, domain string, expires time.Time, maxage int, secure bool, httpOnly bool) {
	cookie := &http.Cookie{
		Name:     name,
		Value:    value,
		Path:     path,
		Domain:   domain,
		Expires:  expires,
		MaxAge:   maxage,
		Secure:   secure,
		HttpOnly: httpOnly,
	}
	http.SetCookie(w, cookie)
}

func SetCookieSimple(w http.ResponseWriter, key string, value string) {
	SetCookieAdv(w, key, value, "/", "", time.Now().AddDate(100, 0, 0), 0, false, false)
}

func StrConcat(strings ...string) string {
	var buffer bytes.Buffer
	for _, str := range strings {
		buffer.WriteString(str)
	}
	return buffer.String()
}

func DeleteItem(slice interface{}, i int) {
	v := reflect.ValueOf(slice).Elem()
	v.Set(reflect.AppendSlice(v.Slice(0, i), v.Slice(i+1, v.Len())))
}

func InsertItem(slice interface{}, i int, val interface{}) {
	v := reflect.ValueOf(slice).Elem()
	v.Set(reflect.AppendSlice(v.Slice(0, i+1), v.Slice(i, v.Len())))
	v.Index(i).Set(reflect.ValueOf(val))
}

func StringIndexOf(haystack []string, needle string) int {
	for i, v := range haystack {
		if v == needle {
			return i
		}
	}
	return -1
}

func IntMax(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func FormatPath(rawpath string) string {
	cwd, _ := os.Getwd()
	if strings.HasPrefix(rawpath, "./") {
		return filepath.Clean(cwd + rawpath[1:])
	} else if strings.HasPrefix(rawpath, "../") {
		return filepath.Clean(cwd + "/" + rawpath)
	}
	return filepath.Clean(rawpath)
}
<commit_msg>fixed redirect when on localhost<commit_after>package goboots

import (
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"errors"
	"github.com/gabstv/i18ngo"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"
	"time"
)

const (
	ErrNil = 0
	// TLS
	ErrTLSNil = 100
)

var (
	staticControllers []IController
	sessionDbs        map[string]ISessionDBEngine
	curSessionDb      ISessionDBEngine
	httpErrorStrings  = map[int]string{
		400: "Bad Request",
		401: "Unauthorized",
		403: "Forbidden",
		404: "Not Found",
		405: "Method Not Allowed",
		406: "Not Acceptable",
		501: "Internal Server Error",
	}
)

type Params map[string]string

func hostonly(hostport string) string {
	h, _, _ := net.SplitHostPort(hostport)
	return h
}

func trimhost(hostport string) string {
	return strings.Split(hostport, ":")[0]
}

func getTLSRedirectURL(hostaddrtls string, uri *url.URL) (string, error) {
	if uri == nil {
		return "", errors.New("url is null")
	}
	_, p, err := net.SplitHostPort(hostaddrtls)
	if err != nil {
		return "", err
	}
	if uri.Scheme == "ws" {
		uri.Scheme = "wss"
	} else if uri.Scheme == "http" {
		uri.Scheme = "https"
	} else {
		// localhost may leave this empty
		uri.Scheme = "https"
	}
	if uri.Host == "" {
		// it's localhost
		uri.Host = "localhost"
	}
	if p != "443" {
		uri.Host = uri.Host + ":" + p
	}
	return uri.String(), nil
}

func RegisterSessionStorageDriver(name string, engine ISessionDBEngine) {
	if sessionDbs == nil {
		sessionDbs = make(map[string]ISessionDBEngine, 0)
	}
	sessionDbs[name] = engine
}

func RegisterControllerGlobal(controller IController) {
	if staticControllers == nil {
		staticControllers = make([]IController, 0)
	}
	staticControllers = append(staticControllers, controller)
}

func (app *App) InitSessionStorage(driver string) error {
	if sessionDbs == nil {
		return errors.New("No session storage registered.")
	}
	if _, ok := sessionDbs[driver]; !ok {
		return errors.New("The session storage driver " + driver + " does not exist.")
	}
	curSessionDb = sessionDbs[driver]
	curSessionDb.SetApp(app)
	return nil
}

func CloseSessionStorage() {
	if curSessionDb != nil {
		curSessionDb.Close()
	}
}

func initAnySessionStorage() {
	if curSessionDb != nil {
		return
	}
	if sessionDbs == nil {
		curSessionDb = &SessionDevNull{}
		return
	}
	for _, v := range sessionDbs {
		curSessionDb = v
		break
	}
}

type AppError struct {
	Id      int
	Message string
}

func (e *AppError) Error() string {
	return e.Message
}

type env struct {
	path       string
	controller string
	action     string
	params     []string
	request    *http.Request
	body       string
	form       map[string][]string
}

type ISession interface {
	GetData() (string, map[string]interface{}, time.Time)
}

type Session struct {
	SID     string
	Data    map[string]interface{}
	Flash   SessFlash `json:"-" bson:"-"` // never save flash
	Time    time.Time
	Updated time.Time
	r       *http.Request
	w       http.ResponseWriter
}

type SessFlash struct {
	vals map[string]interface{}
}

func (s *SessFlash) init() {
	if s.vals == nil {
		s.vals = make(map[string]interface{})
	}
}

func (s *SessFlash) Set(key string, val interface{}) {
	s.init()
	s.vals[key] = val
}

func (s *SessFlash) Get(key string) interface{} {
	s.init()
	return s.vals[key]
}

func (s *SessFlash) Get2(key string) (interface{}, bool) {
	s.init()
	v, ok := s.vals[key]
	return v, ok
}

func (s *SessFlash) Del(key string) {
	s.init()
	delete(s.vals, key)
}

func (s *SessFlash) Clear() {
	s.vals = nil
}

func (s *SessFlash) All() map[string]interface{} {
	s.init()
	return s.vals
}

func (s *Session) GetData() (string, map[string]interface{}, time.Time) {
	return s.SID, s.Data, s.Time
}

func (s *Session) Flush() {
	FlushSession(s)
}

func (s *Session) GetBool(key string) (bool, bool) {
	iface, ok := s.Data[key]
	if !ok {
		return false, false
	}
	return bool(reflect.ValueOf(iface).Bool()), true
}

func (s *Session) GetBoolD(key string, defaultValue bool) bool {
	iface, ok := s.Data[key]
	if !ok {
		return defaultValue
	}
	return bool(reflect.ValueOf(iface).Bool())
}

func (s *Session) GetInt32(key string) (int, bool) {
	if s.Data == nil {
		return 0, false
	}
	iface, ok := s.Data[key]
	if !ok {
		return 0, false
	}
	return int(reflect.ValueOf(iface).Int()), true
}

func (s *Session) GetInt32D(key string, defaultValue int) int {
	if s.Data == nil {
		return defaultValue
	}
	iface, ok := s.Data[key]
	if !ok {
		return defaultValue
	}
	return int(reflect.ValueOf(iface).Int())
}

func (s *Session) GetString(key string) (string, bool) {
	if s.Data == nil {
		return "", false
	}
	iface, ok := s.Data[key]
	if !ok {
		return "", false
	}
	return reflect.ValueOf(iface).String(), true
}

func (s *Session) GetStringD(key string, defaultValue string) string {
	if s.Data == nil {
		return defaultValue
	}
	iface, ok := s.Data[key]
	if !ok {
		return defaultValue
	}
	return reflect.ValueOf(iface).String()
}

func (s *Session) Expire(t time.Time) {
	SetCookieAdv(s.w, "goboots_sessid", s.SID, "/", "", t, 0, false, true)
}

func (s *Session) GetExpires() time.Time {
	cookie, err := s.r.Cookie("goboots_sessid")
	if err != nil {
		log.Println("Error @ (s *Session) GetExpires()", err)
		return time.Now()
	}
	return cookie.Expires
}

func (app *App) GetSession(w http.ResponseWriter, r *http.Request) *Session {
	if w == nil || r == nil {
		return nil
	}
	initAnySessionStorage()
	var cookie *http.Cookie
	var sid string
	var err error
	//LOAD
	cookie, err = r.Cookie("goboots_sessid")
	if err == nil {
		sid = cookie.Value
		if _validateSidString(sid) {
			var msession *Session
			// get from a saved location
			msession, err = curSessionDb.GetSession(sid)
			if err == nil {
				msession.r = r
				msession.w = w
				return msession
			}
			app.Logger.Printf("SESSION ERROR :( [%s] %s\n", sid, err.Error())
			// not found, generate a new sid
		} else {
			app.Logger.Printf("COULD NOT VALIDATE :( [%s]\n", sid)
		}
	}
	// gen session
	sid = uuid.New()

	session := &Session{
		SID:     sid,
		Data:    make(map[string]interface{}),
		Time:    time.Now(),
		Updated: time.Now(),
	}
	err = curSessionDb.NewSession(session)
	if err != nil {
		log.Println("[FATAL] [curSessionDb.NewSession] Could not get session!", err)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("The server encountered an error while processing your request [ERR_CONN_NEW_SESSION]."))
		return nil
	}
	SetCookieAdv(w, "goboots_sessid", sid, "/", "", time.Now().AddDate(0, 1, 0), 0, false, true)
	session.w = w
	session.r = r
	return session
}

func GetUserLang(w http.ResponseWriter, r *http.Request) string {
	var l string
	// get lang from cookies
	cookie, err := r.Cookie("lang")
	if err != nil && cookie != nil {
		l = cookie.Value
		return l
	}
	// get lang from http request
	alh := r.Header.Get("Accept-Language")
	// 2013-07-29 : not all clients actually send this header (googlebot/wget etc)
	if len(alh) > 1 {
		validlangs := i18ngo.GetLanguageCodes()
		// break alh
		alh0 := strings.Split(alh, ",")
		for _, v := range alh0 {
			lcd := strings.Split(v, ";")[0]
			var lc string
			if len(lcd) >= 2 {
				lc = lcd[0:2]
			}
			lc = lcd
			if StringIndexOf(validlangs, lc) != -1 {
				l = lc
				break
			}
		}
		if len(l) > 0 {
			// just update the cookies
			if w != nil {
				SetCookieSimple(w, "lang", l)
			}
		} else {
			l = i18ngo.GetDefaultLanguageCode()
		}
		return l
	}
	return i18ngo.GetDefaultLanguageCode()
}

func SetUserLang(w http.ResponseWriter, r *http.Request, langcode string) {
	SetCookieSimple(w, "lang", langcode)
}

func FlushSession(s *Session) error {
	return curSessionDb.PutSession(s)
}

func DestroySession(w http.ResponseWriter, r *http.Request, s *Session) {
	curSessionDb.RemoveSession(s)
	SetCookieAdv(w, "goboots_sessid", "", "/", "", time.Now(), 1, false, true)
}

// DEPRECATED
func _validateSidString(sid string) bool {
	if len(sid) < 32 {
		return false
	}
	for _, c := range sid {
		// check if sid contains invalid chars
		if !((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || (c == '-')) {
			return false
		}
	}
	return true
}

type templateInfo struct {
	path       string
	data       *template.Template
	lastUpdate time.Time
}

func SetCookieAdv(w http.ResponseWriter, name string, value string, path string, domain string, expires time.Time, maxage int, secure bool, httpOnly bool) {
	cookie := &http.Cookie{
		Name:     name,
		Value:    value,
		Path:     path,
		Domain:   domain,
		Expires:  expires,
		MaxAge:   maxage,
		Secure:   secure,
		HttpOnly: httpOnly,
	}
	http.SetCookie(w, cookie)
}

func SetCookieSimple(w http.ResponseWriter, key string, value string) {
	SetCookieAdv(w, key, value, "/", "", time.Now().AddDate(100, 0, 0), 0, false, false)
}

func StrConcat(strings ...string) string {
	var buffer bytes.Buffer
	for _, str := range strings {
		buffer.WriteString(str)
	}
	return buffer.String()
}

func DeleteItem(slice interface{}, i int) {
	v := reflect.ValueOf(slice).Elem()
	v.Set(reflect.AppendSlice(v.Slice(0, i), v.Slice(i+1, v.Len())))
}

func InsertItem(slice interface{}, i int, val interface{}) {
	v := reflect.ValueOf(slice).Elem()
	v.Set(reflect.AppendSlice(v.Slice(0, i+1), v.Slice(i, v.Len())))
	v.Index(i).Set(reflect.ValueOf(val))
}

func StringIndexOf(haystack []string, needle string) int {
	for i, v := range haystack {
		if v == needle {
			return i
		}
	}
	return -1
}

func IntMax(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func FormatPath(rawpath string) string {
	cwd, _ := os.Getwd()
	if strings.HasPrefix(rawpath, "./") {
		return filepath.Clean(cwd + rawpath[1:])
	} else if strings.HasPrefix(rawpath, "../") {
		return filepath.Clean(cwd + "/" + rawpath)
	}
	return filepath.Clean(rawpath)
}
<|endoftext|>


<commit_before>// Package quantile computes approximate quantiles over an unbounded data
// stream within low memory and CPU bounds.
//
// A small amount of accuracy is traded to achieve the above properties.
//
// Multiple streams can be merged before calling Query to generate a single set
// of results. This is meaningful when the streams represent the same type of
// data. See Merge and Samples.
//
// For more detailed information about the algorithm used, see:
// Effective Computation of Biased Quantiles over Data Streams - http://www.cs.rutgers.edu/~muthu/bquant.pdf
package quantile

import (
	"container/list"
	"math"
	"sort"
)

// Sample holds an observed value and meta information for compression. JSON
// tags have been added for convenience.
type Sample struct {
	Value float64 `json:",string"`
	Width float64 `json:",string"`
	Delta float64 `json:",string"`
}

// Samples represents a slice of samples. It implements sort.Interface.
type Samples []Sample

func (a Samples) Len() int {
	return len(a)
}

func (a Samples) Less(i, j int) bool {
	return a[i].Value < a[j].Value
}

func (a Samples) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

type invariant func(s *stream, r float64) float64

// NewBiased returns an initialized Stream for high-biased quantiles (e.g.
// 50th, 90th, 99th) not known a priori with ﬁner error guarantees for the
// higher ranks of the data distribution.
// See http://www.cs.rutgers.edu/~muthu/bquant.pdf for time, space, and error properties.
func NewBiased() *Stream {
	ƒ := func(s *stream, r float64) float64 {
		return 2 * s.epsilon * r
	}
	return newStream(ƒ)
}

// NewTargeted returns an initialized Stream concerned with a particular set of
// quantile values that are supplied a priori. Knowing these a priori reduces
// space and computation time.
// See http://www.cs.rutgers.edu/~muthu/bquant.pdf for time, space, and error properties.
func NewTargeted(quantiles ...float64) *Stream {
	ƒ := func(s *stream, r float64) float64 {
		var m float64 = math.MaxFloat64
		var f float64
		for _, q := range quantiles {
			if q*s.n <= r {
				f = (2 * s.epsilon * r) / q
			} else {
				f = (2 * s.epsilon * (s.n - r)) / (1 - q)
			}
			m = math.Min(m, f)
		}
		return m
	}
	return newStream(ƒ)
}

// Stream computes quantiles for a stream of float64s. It is not thread-safe.
type Stream struct {
	*stream
	b Samples
}

func newStream(ƒ invariant) *Stream {
	const defaultEpsilon = 0.01
	x := &stream{epsilon: defaultEpsilon, ƒ: ƒ, l: list.New()}
	return &Stream{x, make(Samples, 0, 500)}
}

// Insert inserts v into the stream.
func (s *Stream) Insert(v float64) {
	s.insert(Sample{Value: v, Width: 1})
}

func (s *Stream) insert(sample Sample) {
	s.b = append(s.b, sample)
	if len(s.b) == cap(s.b) {
		s.flush()
		s.compress()
	}
}

// Query returns the computed qth percentiles value. If s was created with
// NewTargeted, and q is not in the set of quantiles provided a priori, Query
// will return an unspecified result.
func (s *Stream) Query(q float64) float64 {
	if s.flushed() {
		// Fast path when there hasn't been enough data for a flush;
		// this also yeilds better accuracy for small sets of data.
		i := float64(len(s.b)) * q
		return s.b[int(i)-1].Value
	}
	s.flush()
	return s.stream.query(q)
}

// Merge merges samples into the underlying streams samples. This is handy when
// merging multiple streams from separate threads, database shards, etc.
func (s *Stream) Merge(samples Samples) {
	s.stream.merge(samples)
}

// Reset reinitializes and clears the list reusing the samples buffer memory.
func (s *Stream) Reset() {
	s.stream.reset()
	s.b = s.b[:0]
}

// Samples returns stream samples held by s.
func (s *Stream) Samples() Samples {
	if !s.flushed() {
		return s.b
	}
	return s.stream.samples()
}

// Count returns the total number of samples observed in the stream
// since initialization.
func (s *Stream) Count() int {
	return len(s.b) + s.stream.count()
}

func (s *Stream) flush() {
	sort.Sort(s.b)
	s.stream.merge(s.b)
	s.b = s.b[:0]
}

func (s *Stream) flushed() bool {
	return s.stream.l.Len() == 0
}

type stream struct {
	epsilon float64
	n       float64
	l       *list.List
	ƒ       invariant
}

// SetEpsilon sets the error epsilon for the Stream. The default epsilon is
// 0.01 and is usually satisfactory. If needed, this must be called before all
// Inserts.
// To learn more, see: http://www.cs.rutgers.edu/~muthu/bquant.pdf
func (s *stream) SetEpsilon(epsilon float64) {
	s.epsilon = epsilon
}

func (s *stream) reset() {
	s.l.Init()
	s.n = 0
}

func (s *stream) insert(v float64) {
	fn := s.mergeFunc()
	fn(v, 1)
}

func (s *stream) merge(samples Samples) {
	fn := s.mergeFunc()
	for _, s := range samples {
		fn(s.Value, s.Width)
	}
}

func (s *stream) mergeFunc() func(v, w float64) {
	// NOTE: I used a goto over defer because it bought me a few extra
	// nanoseconds. I know. I know.
	var r float64
	e := s.l.Front()
	return func(v, w float64) {
		for ; e != nil; e = e.Next() {
			c := e.Value.(*Sample)
			if c.Value > v {
				sm := &Sample{v, w, math.Floor(s.ƒ(s, r)) - 1}
				s.l.InsertBefore(sm, e)
				goto inserted
			}
			r += c.Width
		}
		s.l.PushBack(&Sample{v, w, 0})
	inserted:
		s.n += w
	}
}

func (s *stream) count() int {
	return int(s.n)
}

func (s *stream) query(q float64) float64 {
	e := s.l.Front()
	t := math.Ceil(q * s.n)
	t += math.Ceil(s.ƒ(s, t) / 2)
	p := e.Value.(*Sample)
	e = e.Next()
	r := float64(0)
	for e != nil {
		c := e.Value.(*Sample)
		if r+c.Width+c.Delta > t {
			return p.Value
		}
		r += p.Width
		p = c
		e = e.Next()
	}
	return p.Value
}

func (s *stream) compress() {
	if s.l.Len() < 2 {
		return
	}
	e := s.l.Back()
	x := e.Value.(*Sample)
	r := s.n - 1 - x.Width
	e = e.Prev()
	for e != nil {
		c := e.Value.(*Sample)
		if c.Width+x.Width+x.Delta <= s.ƒ(s, r) {
			x.Width += c.Width
			o := e
			e = e.Prev()
			s.l.Remove(o)
		} else {
			x = c
			e = e.Prev()
		}
		r -= c.Width
	}
}

func (s *stream) samples() Samples {
	samples := make(Samples, 0, s.l.Len())
	for e := s.l.Front(); e != nil; e = e.Next() {
		samples = append(samples, *e.Value.(*Sample))
	}
	return samples
}
<commit_msg>fix formatting<commit_after>// Package quantile computes approximate quantiles over an unbounded data
// stream within low memory and CPU bounds.
//
// A small amount of accuracy is traded to achieve the above properties.
//
// Multiple streams can be merged before calling Query to generate a single set
// of results. This is meaningful when the streams represent the same type of
// data. See Merge and Samples.
//
// For more detailed information about the algorithm used, see:
//
// Effective Computation of Biased Quantiles over Data Streams
//
// http://www.cs.rutgers.edu/~muthu/bquant.pdf
package quantile

import (
	"container/list"
	"math"
	"sort"
)

// Sample holds an observed value and meta information for compression. JSON
// tags have been added for convenience.
type Sample struct {
	Value float64 `json:",string"`
	Width float64 `json:",string"`
	Delta float64 `json:",string"`
}

// Samples represents a slice of samples. It implements sort.Interface.
type Samples []Sample

func (a Samples) Len() int {
	return len(a)
}

func (a Samples) Less(i, j int) bool {
	return a[i].Value < a[j].Value
}

func (a Samples) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

type invariant func(s *stream, r float64) float64

// NewBiased returns an initialized Stream for high-biased quantiles (e.g.
// 50th, 90th, 99th) not known a priori with ﬁner error guarantees for the
// higher ranks of the data distribution.
// See http://www.cs.rutgers.edu/~muthu/bquant.pdf for time, space, and error properties.
func NewBiased() *Stream {
	ƒ := func(s *stream, r float64) float64 {
		return 2 * s.epsilon * r
	}
	return newStream(ƒ)
}

// NewTargeted returns an initialized Stream concerned with a particular set of
// quantile values that are supplied a priori. Knowing these a priori reduces
// space and computation time.
// See http://www.cs.rutgers.edu/~muthu/bquant.pdf for time, space, and error properties.
func NewTargeted(quantiles ...float64) *Stream {
	ƒ := func(s *stream, r float64) float64 {
		var m float64 = math.MaxFloat64
		var f float64
		for _, q := range quantiles {
			if q*s.n <= r {
				f = (2 * s.epsilon * r) / q
			} else {
				f = (2 * s.epsilon * (s.n - r)) / (1 - q)
			}
			m = math.Min(m, f)
		}
		return m
	}
	return newStream(ƒ)
}

// Stream computes quantiles for a stream of float64s. It is not thread-safe.
type Stream struct {
	*stream
	b Samples
}

func newStream(ƒ invariant) *Stream {
	const defaultEpsilon = 0.01
	x := &stream{epsilon: defaultEpsilon, ƒ: ƒ, l: list.New()}
	return &Stream{x, make(Samples, 0, 500)}
}

// Insert inserts v into the stream.
func (s *Stream) Insert(v float64) {
	s.insert(Sample{Value: v, Width: 1})
}

func (s *Stream) insert(sample Sample) {
	s.b = append(s.b, sample)
	if len(s.b) == cap(s.b) {
		s.flush()
		s.compress()
	}
}

// Query returns the computed qth percentiles value. If s was created with
// NewTargeted, and q is not in the set of quantiles provided a priori, Query
// will return an unspecified result.
func (s *Stream) Query(q float64) float64 {
	if s.flushed() {
		// Fast path when there hasn't been enough data for a flush;
		// this also yeilds better accuracy for small sets of data.
		i := float64(len(s.b)) * q
		return s.b[int(i)-1].Value
	}
	s.flush()
	return s.stream.query(q)
}

// Merge merges samples into the underlying streams samples. This is handy when
// merging multiple streams from separate threads, database shards, etc.
func (s *Stream) Merge(samples Samples) {
	s.stream.merge(samples)
}

// Reset reinitializes and clears the list reusing the samples buffer memory.
func (s *Stream) Reset() {
	s.stream.reset()
	s.b = s.b[:0]
}

// Samples returns stream samples held by s.
func (s *Stream) Samples() Samples {
	if !s.flushed() {
		return s.b
	}
	return s.stream.samples()
}

// Count returns the total number of samples observed in the stream
// since initialization.
func (s *Stream) Count() int {
	return len(s.b) + s.stream.count()
}

func (s *Stream) flush() {
	sort.Sort(s.b)
	s.stream.merge(s.b)
	s.b = s.b[:0]
}

func (s *Stream) flushed() bool {
	return s.stream.l.Len() == 0
}

type stream struct {
	epsilon float64
	n       float64
	l       *list.List
	ƒ       invariant
}

// SetEpsilon sets the error epsilon for the Stream. The default epsilon is
// 0.01 and is usually satisfactory. If needed, this must be called before all
// Inserts.
// To learn more, see: http://www.cs.rutgers.edu/~muthu/bquant.pdf
func (s *stream) SetEpsilon(epsilon float64) {
	s.epsilon = epsilon
}

func (s *stream) reset() {
	s.l.Init()
	s.n = 0
}

func (s *stream) insert(v float64) {
	fn := s.mergeFunc()
	fn(v, 1)
}

func (s *stream) merge(samples Samples) {
	fn := s.mergeFunc()
	for _, s := range samples {
		fn(s.Value, s.Width)
	}
}

func (s *stream) mergeFunc() func(v, w float64) {
	// NOTE: I used a goto over defer because it bought me a few extra
	// nanoseconds. I know. I know.
	var r float64
	e := s.l.Front()
	return func(v, w float64) {
		for ; e != nil; e = e.Next() {
			c := e.Value.(*Sample)
			if c.Value > v {
				sm := &Sample{v, w, math.Floor(s.ƒ(s, r)) - 1}
				s.l.InsertBefore(sm, e)
				goto inserted
			}
			r += c.Width
		}
		s.l.PushBack(&Sample{v, w, 0})
	inserted:
		s.n += w
	}
}

func (s *stream) count() int {
	return int(s.n)
}

func (s *stream) query(q float64) float64 {
	e := s.l.Front()
	t := math.Ceil(q * s.n)
	t += math.Ceil(s.ƒ(s, t) / 2)
	p := e.Value.(*Sample)
	e = e.Next()
	r := float64(0)
	for e != nil {
		c := e.Value.(*Sample)
		if r+c.Width+c.Delta > t {
			return p.Value
		}
		r += p.Width
		p = c
		e = e.Next()
	}
	return p.Value
}

func (s *stream) compress() {
	if s.l.Len() < 2 {
		return
	}
	e := s.l.Back()
	x := e.Value.(*Sample)
	r := s.n - 1 - x.Width
	e = e.Prev()
	for e != nil {
		c := e.Value.(*Sample)
		if c.Width+x.Width+x.Delta <= s.ƒ(s, r) {
			x.Width += c.Width
			o := e
			e = e.Prev()
			s.l.Remove(o)
		} else {
			x = c
			e = e.Prev()
		}
		r -= c.Width
	}
}

func (s *stream) samples() Samples {
	samples := make(Samples, 0, s.l.Len())
	for e := s.l.Front(); e != nil; e = e.Next() {
		samples = append(samples, *e.Value.(*Sample))
	}
	return samples
}
<|endoftext|>


<commit_before>package lib

import (
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/cavaliercoder/grab"
	"github.com/flavioribeiro/gonfig"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

// HTTPDownload function downloads sources using
// http protocol.
func HTTPDownload(jobID string) error {
	currentDir, _ := os.Getwd()
	cfg, _ := gonfig.FromJsonFile(currentDir + "/config.json")
	swapDir, _ := cfg.GetString("SWAP_DIRECTORY", "")
	dbInstance, _ := db.GetDatabase()
	job, _ := dbInstance.RetrieveJob(jobID)

	basePath := swapDir + string(job.ID)

	sourceDir := basePath + "/src/"
	job.LocalSource = sourceDir + path.Base(job.Source)

	outputDir := basePath + "/dst/"
	os.MkdirAll(sourceDir, 0777)
	os.MkdirAll(outputDir, 0777)
	outputFilename := strings.Split(path.Base(job.Source), ".")[0] + "_" + job.Preset.Name + "." + job.Preset.Container
	job.LocalDestination = outputDir + outputFilename
	job.Destination = job.Destination + outputFilename
	job.Status = types.JobDownloading
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	respch, _ := grab.GetAsync(basePath+"/src/", job.Source)

	resp := <-respch
	for !resp.IsComplete() {
		job, _ = dbInstance.RetrieveJob(jobID)
		percentage := strconv.FormatInt(int64(resp.BytesTransferred()*100/resp.Size), 10)
		if job.Details != percentage {
			job.Details = percentage
			dbInstance.UpdateJob(job.ID, job)
		}
	}

	if resp.Error != nil {
		return resp.Error
	}

	return nil
}
<commit_msg>http_downloader: add % on job progress<commit_after>package lib

import (
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/cavaliercoder/grab"
	"github.com/flavioribeiro/gonfig"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

// HTTPDownload function downloads sources using
// http protocol.
func HTTPDownload(jobID string) error {
	currentDir, _ := os.Getwd()
	cfg, _ := gonfig.FromJsonFile(currentDir + "/config.json")
	swapDir, _ := cfg.GetString("SWAP_DIRECTORY", "")
	dbInstance, _ := db.GetDatabase()
	job, _ := dbInstance.RetrieveJob(jobID)

	basePath := swapDir + string(job.ID)

	sourceDir := basePath + "/src/"
	job.LocalSource = sourceDir + path.Base(job.Source)

	outputDir := basePath + "/dst/"
	os.MkdirAll(sourceDir, 0777)
	os.MkdirAll(outputDir, 0777)
	outputFilename := strings.Split(path.Base(job.Source), ".")[0] + "_" + job.Preset.Name + "." + job.Preset.Container
	job.LocalDestination = outputDir + outputFilename
	job.Destination = job.Destination + outputFilename
	job.Status = types.JobDownloading
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	respch, _ := grab.GetAsync(basePath+"/src/", job.Source)

	resp := <-respch
	for !resp.IsComplete() {
		job, _ = dbInstance.RetrieveJob(jobID)
		percentage := strconv.FormatInt(int64(resp.BytesTransferred()*100/resp.Size), 10)
		if job.Details != percentage {
			job.Details = percentage + "%"
			dbInstance.UpdateJob(job.ID, job)
		}
	}

	if resp.Error != nil {
		return resp.Error
	}

	return nil
}
<|endoftext|>


<commit_before>// +build linux

package journald

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/go-systemd/sdjournal"
	"github.com/segmentio/ecs-logs/lib"
)

func NewReader() (r ecslogs.Reader, err error) {
	var j *sdjournal.Journal

	if j, err = sdjournal.NewJournal(); err != nil {
		return
	}

	if err = j.SeekTail(); err != nil {
		j.Close()
		return
	}

	r = reader{j}
	return
}

type reader struct {
	*sdjournal.Journal
}

func (r reader) ReadMessage() (msg ecslogs.Message, err error) {
	for {
		var cur int
		var ok bool

		if cur, err = r.Next(); err != nil {
			return
		}

		if cur == 0 {
			r.Wait(sdjournal.IndefiniteWait)
			continue
		}

		if msg, ok, err = r.getMessage(); ok || err != nil {
			return
		}
	}
}

func (r reader) getMessage() (msg ecslogs.Message, ok bool, err error) {
	if msg.Group, err = r.GetDataValue("CONTAINER_TAG"); len(msg.Group) == 0 {
		// No CONTAINER_TAG, this must be a journal message from a process that
		// isn't running in a docker container.
		err = nil
		return
	}

	if msg.Stream, err = r.GetDataValue("CONTAINER_NAME"); err != nil {
		// There's a CONTAINER_TAG but no CONTAINER_NAME, something is seriously
		// wrong here, the log docker log driver is misbehaving.
		err = fmt.Errorf("missing CONTAINER_NAME in message with CONTAINER_TAG=%s", msg.Group)
		return
	}

	if s := r.getString("MESSAGE"); len(s) != 0 {
		d := json.NewDecoder(strings.NewReader(s))
		d.UseNumber()

		if d.Decode(&msg.Event) != nil {
			msg.Event.Message = s
		}
	}

	if msg.Event.Level == ecslogs.NONE {
		msg.Event.Level = r.getPriority()
	}

	if len(msg.Event.Info.Host) == 0 {
		msg.Event.Info.Host = r.getString("_HOSTNAME")
	}

	if len(msg.Event.Info.Source) == 0 {
		msg.Event.Info.Source = (ecslogs.FuncInfo{
			File: r.getString("CODE_FILE"),
			Func: r.getString("CODE_FUNC"),
			Line: r.getInt("CODE_LINE"),
		}).String()
	}

	if len(msg.Event.Info.ID) == 0 {
		msg.Event.Info.ID = r.getString("MESSAGE_ID")
	}

	if msg.Event.Info.PID == 0 {
		msg.Event.Info.PID = r.getInt("_PID")
	}

	if msg.Event.Info.GID == 0 {
		msg.Event.Info.GID = r.getInt("_GID")
	}

	if msg.Event.Info.UID == 0 {
		msg.Event.Info.UID = r.getInt("_UID")
	}

	if msg.Event.Time == (time.Time{}) {
		msg.Event.Time = r.getTime()
	}

	ok = true
	return
}

func (r reader) getInt(k string) (v int) {
	v, _ = strconv.Atoi(r.getString(k))
	return
}

func (r reader) getTime() (t time.Time) {
	if u, e := r.GetRealtimeUsec(); e == nil {
		t = time.Unix(int64(u/1000000), int64((u%1000000)*1000))
	}
	return
}

func (r reader) getPriority() (p ecslogs.Level) {
	if v, e := strconv.Atoi(r.getString("PRIORITY")); e != nil {
		p = ecslogs.INFO
	} else {
		p = ecslogs.Level(v + 1)
	}
	return
}

func (r reader) getString(k string) (s string) {
	s, _ = r.GetDataValue(k)
	return
}
<commit_msg>always set a valid time in messages returned by the journal reader<commit_after>// +build linux

package journald

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/go-systemd/sdjournal"
	"github.com/segmentio/ecs-logs/lib"
)

func NewReader() (r ecslogs.Reader, err error) {
	var j *sdjournal.Journal

	if j, err = sdjournal.NewJournal(); err != nil {
		return
	}

	if err = j.SeekTail(); err != nil {
		j.Close()
		return
	}

	r = reader{j}
	return
}

type reader struct {
	*sdjournal.Journal
}

func (r reader) ReadMessage() (msg ecslogs.Message, err error) {
	for {
		var cur int
		var ok bool

		if cur, err = r.Next(); err != nil {
			return
		}

		if cur == 0 {
			r.Wait(sdjournal.IndefiniteWait)
			continue
		}

		if msg, ok, err = r.getMessage(); ok || err != nil {
			return
		}
	}
}

func (r reader) getMessage() (msg ecslogs.Message, ok bool, err error) {
	if msg.Group, err = r.GetDataValue("CONTAINER_TAG"); len(msg.Group) == 0 {
		// No CONTAINER_TAG, this must be a journal message from a process that
		// isn't running in a docker container.
		err = nil
		return
	}

	if msg.Stream, err = r.GetDataValue("CONTAINER_NAME"); err != nil {
		// There's a CONTAINER_TAG but no CONTAINER_NAME, something is seriously
		// wrong here, the log docker log driver is misbehaving.
		err = fmt.Errorf("missing CONTAINER_NAME in message with CONTAINER_TAG=%s", msg.Group)
		return
	}

	if s := r.getString("MESSAGE"); len(s) != 0 {
		d := json.NewDecoder(strings.NewReader(s))
		d.UseNumber()

		if d.Decode(&msg.Event) != nil {
			msg.Event.Message = s
		}
	}

	if msg.Event.Level == ecslogs.NONE {
		msg.Event.Level = r.getPriority()
	}

	if len(msg.Event.Info.Host) == 0 {
		msg.Event.Info.Host = r.getString("_HOSTNAME")
	}

	if len(msg.Event.Info.Source) == 0 {
		msg.Event.Info.Source = (ecslogs.FuncInfo{
			File: r.getString("CODE_FILE"),
			Func: r.getString("CODE_FUNC"),
			Line: r.getInt("CODE_LINE"),
		}).String()
	}

	if len(msg.Event.Info.ID) == 0 {
		msg.Event.Info.ID = r.getString("MESSAGE_ID")
	}

	if msg.Event.Info.PID == 0 {
		msg.Event.Info.PID = r.getInt("_PID")
	}

	if msg.Event.Info.GID == 0 {
		msg.Event.Info.GID = r.getInt("_GID")
	}

	if msg.Event.Info.UID == 0 {
		msg.Event.Info.UID = r.getInt("_UID")
	}

	if msg.Event.Time == (time.Time{}) {
		msg.Event.Time = r.getTime()
	}

	ok = true
	return
}

func (r reader) getInt(k string) (v int) {
	v, _ = strconv.Atoi(r.getString(k))
	return
}

func (r reader) getTime() (t time.Time) {
	if u, e := r.GetRealtimeUsec(); e == nil {
		t = time.Unix(int64(u/1000000), int64((u%1000000)*1000))
	} else {
		t = time.Now()
	}
	return
}

func (r reader) getPriority() (p ecslogs.Level) {
	if v, e := strconv.Atoi(r.getString("PRIORITY")); e != nil {
		p = ecslogs.INFO
	} else {
		p = ecslogs.Level(v + 1)
	}
	return
}

func (r reader) getString(k string) (s string) {
	s, _ = r.GetDataValue(k)
	return
}
<|endoftext|>


<commit_before>package core

import (
	"context"
	"sync"
	"time"

	"v2ray.com/core/common"
	"v2ray.com/core/common/platform"
)

// TimeoutPolicy contains limits for connection timeout.
type TimeoutPolicy struct {
	// Timeout for handshake phase in a connection.
	Handshake time.Duration
	// Timeout for connection being idle, i.e., there is no egress or ingress traffic in this connection.
	ConnectionIdle time.Duration
	// Timeout for an uplink only connection, i.e., the downlink of the connection has been closed.
	UplinkOnly time.Duration
	// Timeout for an downlink only connection, i.e., the uplink of the connection has been closed.
	DownlinkOnly time.Duration
}

// StatsPolicy contains settings for stats counters.
type StatsPolicy struct {
	// Whether or not to enable stat counter for user uplink traffic.
	UserUplink bool
	// Whether or not to enable stat counter for user downlink traffic.
	UserDownlink bool
}

// BufferPolicy contains settings for internal buffer.
type BufferPolicy struct {
	// Size of buffer per connection, in bytes. -1 for unlimited buffer.
	PerConnection int32
}

// SystemStatsPolicy contains stat policy settings on system level.
type SystemStatsPolicy struct {
	// Whether or not to enable stat counter for uplink traffic in inbound handlers.
	InboundUplink bool
	// Whether or not to enable stat counter for downlink traffic in inbound handlers.
	InboundDownlink bool
}

// SystemPolicy contains policy settings at system level.
type SystemPolicy struct {
	Stats  SystemStatsPolicy
	Buffer BufferPolicy
}

// Policy is session based settings for controlling V2Ray requests. It contains various settings (or limits) that may differ for different users in the context.
type Policy struct {
	Timeouts TimeoutPolicy // Timeout settings
	Stats    StatsPolicy
	Buffer   BufferPolicy
}

// PolicyManager is a feature that provides Policy for the given user by its id or level.
type PolicyManager interface {
	Feature

	// ForLevel returns the Policy for the given user level.
	ForLevel(level uint32) Policy

	// ForSystem returns the Policy for V2Ray system.
	ForSystem() SystemPolicy
}

var defaultBufferSize int32 = 10 * 1024 * 1024

func init() {
	const key = "v2ray.ray.buffer.size"
	size := platform.EnvFlag{
		Name:    key,
		AltName: platform.NormalizeEnvName(key),
	}.GetValueAsInt(10)
	if size == 0 {
		defaultBufferSize = -1
	} else {
		defaultBufferSize = int32(size) * 1024 * 1024
	}
}

func defaultBufferPolicy() BufferPolicy {
	return BufferPolicy{
		PerConnection: defaultBufferSize,
	}
}

// DefaultPolicy returns the Policy when user is not specified.
func DefaultPolicy() Policy {
	return Policy{
		Timeouts: TimeoutPolicy{
			Handshake:      time.Second * 4,
			ConnectionIdle: time.Second * 300,
			UplinkOnly:     time.Second * 2,
			DownlinkOnly:   time.Second * 5,
		},
		Stats: StatsPolicy{
			UserUplink:   false,
			UserDownlink: false,
		},
		Buffer: defaultBufferPolicy(),
	}
}

type policyKey int

const (
	bufferPolicyKey policyKey = 0
)

func ContextWithBufferPolicy(ctx context.Context, p BufferPolicy) context.Context {
	return context.WithValue(ctx, bufferPolicyKey, p)
}

func BufferPolicyFromContext(ctx context.Context) BufferPolicy {
	pPolicy := ctx.Value(bufferPolicyKey)
	if pPolicy == nil {
		return defaultBufferPolicy()
	}
	return pPolicy.(BufferPolicy)
}

type syncPolicyManager struct {
	sync.RWMutex
	PolicyManager
}

func (m *syncPolicyManager) ForLevel(level uint32) Policy {
	m.RLock()
	defer m.RUnlock()

	if m.PolicyManager == nil {
		p := DefaultPolicy()
		if level == 1 {
			p.Timeouts.ConnectionIdle = time.Second * 600
		}
		return p
	}

	return m.PolicyManager.ForLevel(level)
}

func (m *syncPolicyManager) ForSystem() SystemPolicy {
	m.RLock()
	defer m.RUnlock()

	if m.PolicyManager == nil {
		return SystemPolicy{}
	}

	return m.PolicyManager.ForSystem()
}

func (m *syncPolicyManager) Start() error {
	m.RLock()
	defer m.RUnlock()

	if m.PolicyManager == nil {
		return nil
	}

	return m.PolicyManager.Start()
}

func (m *syncPolicyManager) Close() error {
	m.RLock()
	defer m.RUnlock()

	return common.Close(m.PolicyManager)
}

func (m *syncPolicyManager) Set(manager PolicyManager) {
	if manager == nil {
		return
	}

	m.Lock()
	defer m.Unlock()

	common.Close(m.PolicyManager) // nolint: errcheck
	m.PolicyManager = manager
}
<commit_msg>reduce default buffer size to 2MB<commit_after>package core

import (
	"context"
	"sync"
	"time"

	"v2ray.com/core/common"
	"v2ray.com/core/common/platform"
)

// TimeoutPolicy contains limits for connection timeout.
type TimeoutPolicy struct {
	// Timeout for handshake phase in a connection.
	Handshake time.Duration
	// Timeout for connection being idle, i.e., there is no egress or ingress traffic in this connection.
	ConnectionIdle time.Duration
	// Timeout for an uplink only connection, i.e., the downlink of the connection has been closed.
	UplinkOnly time.Duration
	// Timeout for an downlink only connection, i.e., the uplink of the connection has been closed.
	DownlinkOnly time.Duration
}

// StatsPolicy contains settings for stats counters.
type StatsPolicy struct {
	// Whether or not to enable stat counter for user uplink traffic.
	UserUplink bool
	// Whether or not to enable stat counter for user downlink traffic.
	UserDownlink bool
}

// BufferPolicy contains settings for internal buffer.
type BufferPolicy struct {
	// Size of buffer per connection, in bytes. -1 for unlimited buffer.
	PerConnection int32
}

// SystemStatsPolicy contains stat policy settings on system level.
type SystemStatsPolicy struct {
	// Whether or not to enable stat counter for uplink traffic in inbound handlers.
	InboundUplink bool
	// Whether or not to enable stat counter for downlink traffic in inbound handlers.
	InboundDownlink bool
}

// SystemPolicy contains policy settings at system level.
type SystemPolicy struct {
	Stats  SystemStatsPolicy
	Buffer BufferPolicy
}

// Policy is session based settings for controlling V2Ray requests. It contains various settings (or limits) that may differ for different users in the context.
type Policy struct {
	Timeouts TimeoutPolicy // Timeout settings
	Stats    StatsPolicy
	Buffer   BufferPolicy
}

// PolicyManager is a feature that provides Policy for the given user by its id or level.
type PolicyManager interface {
	Feature

	// ForLevel returns the Policy for the given user level.
	ForLevel(level uint32) Policy

	// ForSystem returns the Policy for V2Ray system.
	ForSystem() SystemPolicy
}

var defaultBufferSize int32 = 2 * 1024 * 1024

func init() {
	const key = "v2ray.ray.buffer.size"
	size := platform.EnvFlag{
		Name:    key,
		AltName: platform.NormalizeEnvName(key),
	}.GetValueAsInt(10)
	if size == 0 {
		defaultBufferSize = -1
	} else {
		defaultBufferSize = int32(size) * 1024 * 1024
	}
}

func defaultBufferPolicy() BufferPolicy {
	return BufferPolicy{
		PerConnection: defaultBufferSize,
	}
}

// DefaultPolicy returns the Policy when user is not specified.
func DefaultPolicy() Policy {
	return Policy{
		Timeouts: TimeoutPolicy{
			Handshake:      time.Second * 4,
			ConnectionIdle: time.Second * 300,
			UplinkOnly:     time.Second * 2,
			DownlinkOnly:   time.Second * 5,
		},
		Stats: StatsPolicy{
			UserUplink:   false,
			UserDownlink: false,
		},
		Buffer: defaultBufferPolicy(),
	}
}

type policyKey int

const (
	bufferPolicyKey policyKey = 0
)

func ContextWithBufferPolicy(ctx context.Context, p BufferPolicy) context.Context {
	return context.WithValue(ctx, bufferPolicyKey, p)
}

func BufferPolicyFromContext(ctx context.Context) BufferPolicy {
	pPolicy := ctx.Value(bufferPolicyKey)
	if pPolicy == nil {
		return defaultBufferPolicy()
	}
	return pPolicy.(BufferPolicy)
}

type syncPolicyManager struct {
	sync.RWMutex
	PolicyManager
}

func (m *syncPolicyManager) ForLevel(level uint32) Policy {
	m.RLock()
	defer m.RUnlock()

	if m.PolicyManager == nil {
		p := DefaultPolicy()
		if level == 1 {
			p.Timeouts.ConnectionIdle = time.Second * 600
		}
		return p
	}

	return m.PolicyManager.ForLevel(level)
}

func (m *syncPolicyManager) ForSystem() SystemPolicy {
	m.RLock()
	defer m.RUnlock()

	if m.PolicyManager == nil {
		return SystemPolicy{}
	}

	return m.PolicyManager.ForSystem()
}

func (m *syncPolicyManager) Start() error {
	m.RLock()
	defer m.RUnlock()

	if m.PolicyManager == nil {
		return nil
	}

	return m.PolicyManager.Start()
}

func (m *syncPolicyManager) Close() error {
	m.RLock()
	defer m.RUnlock()

	return common.Close(m.PolicyManager)
}

func (m *syncPolicyManager) Set(manager PolicyManager) {
	if manager == nil {
		return
	}

	m.Lock()
	defer m.Unlock()

	common.Close(m.PolicyManager) // nolint: errcheck
	m.PolicyManager = manager
}
<|endoftext|>


<commit_before>/*
 * This file contains a library that implements Bit Slice Indexing or BSI.
 * A BSI at its simplest is an array of bitmaps that represent an encoded
 * binary value.  The advantage of a BSI is that comparisons can be made
 * across ranges of values whereas a bitmap can only represent the existance
 * of a single value for a given column ID.  Another usage scenario involves
 * storage of high cardinality values.
 *
 * It depends upon the roaring bitmap libraries.  It is not thread safe, so
 * upstream concurrency guards must be provided.
 *
 * Author: Guy A. Molinari, The Walt Disney Company
 *
 */
package roaring

import (
    _ "fmt"
    "math"
    "math/bits"
    "runtime"
    "sync"
    "github.com/RoaringBitmap/roaring"
)

type BSI struct {
    bA          []*roaring.Bitmap
    eBM         *roaring.Bitmap        // Existence BitMap
    MaxValue    int64
    MinValue    int64
}

/*
 * Construct a new BSI.  Min/Max values are optional.  If set to 0
 * then the underlying BSI will be automatically sized.
 */
func NewBSI(maxValue int64, minValue int64) *BSI {

    ba := make([]*roaring.Bitmap, bits.Len64(uint64(maxValue)))
    for i := 0; i < len(ba); i++ {
        ba[i] = roaring.NewBitmap()
    }
    return &BSI{bA: ba, eBM: roaring.NewBitmap(), MaxValue: maxValue, MinValue: minValue}
}

// Construct an auto-sized BSI
func NewDefaultBSI() *BSI {
    return NewBSI(int64(0), int64(0))
}

// Value existance test.
func (b *BSI) ValueExists(columnID uint64) bool {

    return b.eBM.Contains(uint32(columnID))
}

// Return a count of unique column IDs for which a value has been set.
func (b *BSI) GetCardinality() uint64 {

    return b.eBM.GetCardinality()
}

// Number of bits needed to represent values.
func (b *BSI) BitCount() int {

    return len(b.bA)
}

// Set a value for a given columnID.
func (b *BSI) SetValue(columnID uint64, value int64) {

    // If max/min values are set to zero then automatically determine bit array size
    if b.MaxValue == 0 && b.MinValue == 0 {
        ba := make([]*roaring.Bitmap, bits.Len64(uint64(value)))
        for i := len(ba) - b.BitCount(); i > 0; i-- {
            b.bA = append(b.bA, roaring.NewBitmap())
        }
    }

    var wg sync.WaitGroup

    for i := 0; i < b.BitCount(); i++ {
        wg.Add(1)
        go func(j int) {
            defer wg.Done()
            if uint64(value) & (1 << uint64(j)) > 0 {
                b.bA[j].Add(uint32(columnID))
            } else {
                b.bA[j].Remove(uint32(columnID))
            }
        }(i)
    }
    wg.Wait()
    b.eBM.Add(uint32(columnID))
}

// Get the value at the column ID.  Second param will be false for non-existant values.
func (b *BSI) GetValue(columnID uint64) (int64, bool) {

    value := int64(0)
    exists := b.eBM.Contains(uint32(columnID))
    if !exists {
        return value, exists
    }
    for i := 0; i < b.BitCount(); i++ {
        if b.bA[i].Contains(uint32(columnID)) {
            value |= (1 << uint64(i))
        }
    }
    return int64(value), exists
}

type action func(t *task, batch []uint32, resutsChan chan *roaring.Bitmap, wg *sync.WaitGroup)

func parallelExecutor (parallelism int, t *task, e action,
        foundSet *roaring.Bitmap) *roaring.Bitmap {

    var n int = parallelism
    if n == 0 {
        n = runtime.NumCPU()
    }

    resultsChan := make(chan *roaring.Bitmap, n)

    card := foundSet.GetCardinality()
    x := uint64(math.Floor(float64(card) / float64(n)))

    remainder := card - (x * uint64(n))
    var batch []uint32
    var wg sync.WaitGroup
    iter := foundSet.ManyIterator()
    for i := 0; i < n; i++ {
        if i == n - 1 {
            batch = make([]uint32, x + remainder)
        } else {
            batch = make([]uint32, x)
        }
        iter.NextMany(batch)
        wg.Add(1)
        go e(t, batch, resultsChan, &wg)
    }

    wg.Wait()

    close(resultsChan)

    ba := make([]*roaring.Bitmap, 0)
    for bm := range resultsChan {
        ba = append(ba, bm)
    }

    return roaring.ParOr(0, ba...)

}

type Operation int

const (
 LT Operation = 1 + iota
 LE
 EQ
 GE
 GT
 RANGE
)

type task struct {
    bsi           *BSI
    op            Operation
    valueOrStart  int64
    end           int64
    values        map[int64]struct{}
    bits          *roaring.Bitmap
}

/*
 * For all operations with the exception of RANGE, the value to be compared is specified by valueOrStart.
 * For the RANGE parameter the comparison criteria is >= valueOrStart and <= end.
 * The parallelism parameter indicates the number of CPU threads to be applied for processing.  A value
 * of zero indicates that all available CPU resources will be potentially utilized.
 */
func (b *BSI) CompareValue(parallelism int, op Operation, valueOrStart, end int64,
        foundSet *roaring.Bitmap) *roaring.Bitmap {

    comp := &task{bsi: b, op: op, valueOrStart: valueOrStart, end: end}
    if foundSet == nil {
        return parallelExecutor(parallelism, comp, compareValue, b.eBM)
    } else {
        return parallelExecutor(parallelism, comp, compareValue, foundSet)
    }
}

func compareValue(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg *sync.WaitGroup) {

    defer wg.Done()

    results := roaring.NewBitmap()

    for i := 0; i < len(batch); i++ {
        cID := batch[i]
        eq1, eq2 := true, true
        lt1, lt2, gt1 := false, false, false
        for j := e.bsi.BitCount() - 1; j >= 0; j-- {
            sliceContainsBit := e.bsi.bA[j].Contains(cID)

            if uint64(e.valueOrStart) & (1 << uint64(j)) > 0 {
                // BIT in value is SET
                if ! sliceContainsBit {
                    if eq1 {
                        if e.op == LT || e.op == LE {
                            lt1 = true
                        }
                        eq1 = false
                        break
                    }
                }
            } else {
                // BIT in value is CLEAR
                if sliceContainsBit {
                    if eq1 {
                        if e.op == GT || e.op == GE || e.op == RANGE {
                            gt1 = true
                        }
                        eq1 = false
                        if e.op != RANGE {
                            break
                        }
                    }
                }
            }

            if e.op == RANGE && uint64(e.end) & (1 << uint64(j)) > 0 {
                // BIT in value is SET
                if ! sliceContainsBit {
                    if eq2 {
                        lt2 = true
                        eq2 = false
                    }
                }
            } else {
                // BIT in value is CLEAR
                if sliceContainsBit {
                    if eq2 {
                        eq2 = false
                    }
                }
            }

        }

        switch e.op {
        case LT:
            if lt1 {
                results.Add(cID)
            }
        case LE:
            if eq1 || lt1 {
                results.Add(cID)
            }
        case EQ:
            if eq1 {
                results.Add(cID)
            }
        case GE:
            if eq1 || gt1 {
                results.Add(cID)
            }
        case GT:
            if gt1 {
                results.Add(cID)
            }
        case RANGE:
            if (eq1 || gt1) && (eq2 || lt2) {
                results.Add(cID)
            }
        default:
            if eq1 {
                results.Add(cID)
            }
        }
    }

    resultsChan <- results
}

/*
 * Sum all values contained within the foundSet.   As a convenience, the cardinality of the foundSet
 * is also returned (for calculating the average).
 */
func (b *BSI) Sum(foundSet *roaring.Bitmap) (sum int64, count uint64) {

    sum = int64(0)
    count = foundSet.GetCardinality()

    iter := foundSet.Iterator()
    for iter.HasNext() {
        cID := iter.Next()
        if value, ok := b.GetValue(uint64(cID)); ok {
            sum += value
        }
    }
    return
}

func (b *BSI) Transpose() *roaring.Bitmap {

    return b.IntersectAndTranspose(0, b.eBM)
}

/*
 * Matrix Transpose function.  Return a bitmap such that the values are represented as column IDs
 * in the returned bitmap. This is accomplished by iterating over the foundSet and only including
 * the column IDs in the source (foundSet) as compared with this BSI.  This can be useful for
 * vectoring one set of integers to another.
 */
func (b *BSI) IntersectAndTranspose(parallelism int, foundSet *roaring.Bitmap) *roaring.Bitmap {

    trans := &task{bsi: b}
    return parallelExecutor(parallelism, trans, transpose, foundSet)
}

func transpose(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg *sync.WaitGroup) {

    defer wg.Done()

    results := roaring.NewBitmap()
    for _, cID := range batch {
        if value, ok := e.bsi.GetValue(uint64(cID)); ok {
            results.Add(uint32(value))
        }
    }
    resultsChan <- results
}

/*
 * ParOr is intended primarily to be a concatenation function to be used during bulk load operations.
 * Care should be taken to make sure that columnIDs do not overlap (unless overlapping values are
 * identical).
 */
func (b *BSI) ParOr(parallelism int, bsis ...*BSI) {

    // Consolidate sets
    bits := len(b.bA)
    for i := 0; i < len(bsis); i++ {
        if len(bsis[i].bA) > bits {
            bits = bsis[i].BitCount()
        }
    }

    // Make sure we have enough bit slices
    for bits > b.BitCount() {
        b.bA = append(b.bA, roaring.NewBitmap())
    }

    a := make([][]*roaring.Bitmap, bits)
    for i, _ := range a {
        a[i] = make([]*roaring.Bitmap, 0)
        for _, x := range bsis {
            if len(x.bA) > i {
                a[i] = append(a[i], x.bA[i])
            } else {
                a[i] = []*roaring.Bitmap{roaring.NewBitmap()}
            }
        }
    }

    // Consolidate existence bit maps
    ebms := make([]*roaring.Bitmap, len(bsis))
    for i, _ := range ebms {
        ebms[i] = bsis[i].eBM
    }

    // First merge all the bit slices from all bsi maps that exist in target
    var wg sync.WaitGroup
    for i := 0; i < bits; i++ {
        wg.Add(1)
        go func(j int) {
            defer wg.Done()
            x := []*roaring.Bitmap{b.bA[j]}
            x = append(x, a[j]...)
            b.bA[j] = roaring.ParOr(parallelism, x...)
        }(i)
    }
    wg.Wait()

    // merge all the EBM maps
    x := []*roaring.Bitmap{b.eBM}
    x = append(x, ebms...)
    b.eBM = roaring.ParOr(parallelism, x...)
}

/*
 * De-serialize a BSI.  The value at bitData[0] is the EBM.  Other indices are in least to most
 * significance order starting at bitData[1] (bit position 0).
 */
func (b *BSI) UnmarshalBinary(bitData [][]byte) error {

    for i := 1; i < len(bitData); i++ {
        if bitData == nil || len(bitData[i]) == 0 {
            continue
        }
        if b.BitCount() < i {
            b.bA = append(b.bA, roaring.NewBitmap())
        }
        if err := b.bA[i - 1].UnmarshalBinary(bitData[i]); err != nil {
            return err
        }
    }
    // First element of bitData is the EBM
    if bitData[0] == nil {
        b.eBM = roaring.NewBitmap()
        return nil
    }
    if err := b.eBM.UnmarshalBinary(bitData[0]); err != nil {
        return err
    }
    return nil
}

/*
 * Serialize a BSI.
 */
func (b *BSI) MarshalBinary() ([][]byte, error) {

    var err error
    data := make([][]byte, b.BitCount() + 1)
    // Add extra element for EBM (BitCount() + 1)
    for i := 1; i < b.BitCount() + 1; i++ {
        data[i], err = b.bA[i - 1].MarshalBinary()
        if err != nil {
            return nil, err
        }
    }
    // Marshal EBM
    data[0], err = b.eBM.MarshalBinary()
    if err != nil {
        return nil, err
    }
    return data, nil
}

/*
 * Return a bitmap containing the column IDs where the values are contained within the list of values
 * provided.
 */
func (b *BSI) BatchEqual(parallelism int, values []int64) *roaring.Bitmap {

    valMap := make(map[int64]struct{}, len(values))
    for i := 0; i < len(values); i++ {
        valMap[values[i]] = struct{}{}
    }
    comp := &task{bsi: b, values: valMap}
    return parallelExecutor(parallelism, comp, batchEqual, b.eBM)
}

func batchEqual(e *task, batch []uint32, resultsChan chan *roaring.Bitmap,
        wg *sync.WaitGroup) {

    defer wg.Done()

    results := roaring.NewBitmap()

    for i := 0; i < len(batch); i++ {
        cID := batch[i]
        if value, ok := e.bsi.GetValue(uint64(cID)); ok {
            if _, yes := e.values[int64(value)]; yes {
                results.Add(cID)
            }
        }
    }
    resultsChan <- results
}

/*
 * Bits that exist in the target are cleared if they are also in the found set.
 */
func ClearBits(foundSet, target *roaring.Bitmap) {
    iter := foundSet.Iterator()
    for iter.HasNext() {
       cID := iter.Next()
       target.Remove(cID)
    }
}

func (b *BSI) ClearValues(foundSet *roaring.Bitmap) {

    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        defer wg.Done()
        ClearBits(foundSet, b.eBM)
    }()
    for i := 0; i < b.BitCount(); i++ {
        wg.Add(1)
        go func(j int) {
            defer wg.Done()
            ClearBits(foundSet, b.bA[j])
        }(i)
    }
    wg.Wait()
}
<commit_msg>Failed my spelling test.   Sorry prof Lemire. :-)<commit_after>/*
 * This file contains a library that implements Bit Slice Indexing or BSI.
 * A BSI at its simplest is an array of bitmaps that represent an encoded
 * binary value.  The advantage of a BSI is that comparisons can be made
 * across ranges of values whereas a bitmap can only represent the existence
 * of a single value for a given column ID.  Another usage scenario involves
 * storage of high cardinality values.
 *
 * It depends upon the roaring bitmap libraries.  It is not thread safe, so
 * upstream concurrency guards must be provided.
 *
 * Author: Guy A. Molinari, The Walt Disney Company
 *
 */
package roaring

import (
    _ "fmt"
    "math"
    "math/bits"
    "runtime"
    "sync"
    "github.com/RoaringBitmap/roaring"
)

type BSI struct {
    bA          []*roaring.Bitmap
    eBM         *roaring.Bitmap        // Existence BitMap
    MaxValue    int64
    MinValue    int64
}

/*
 * Construct a new BSI.  Min/Max values are optional.  If set to 0
 * then the underlying BSI will be automatically sized.
 */
func NewBSI(maxValue int64, minValue int64) *BSI {

    ba := make([]*roaring.Bitmap, bits.Len64(uint64(maxValue)))
    for i := 0; i < len(ba); i++ {
        ba[i] = roaring.NewBitmap()
    }
    return &BSI{bA: ba, eBM: roaring.NewBitmap(), MaxValue: maxValue, MinValue: minValue}
}

// Construct an auto-sized BSI
func NewDefaultBSI() *BSI {
    return NewBSI(int64(0), int64(0))
}

// Value existence test.
func (b *BSI) ValueExists(columnID uint64) bool {

    return b.eBM.Contains(uint32(columnID))
}

// Return a count of unique column IDs for which a value has been set.
func (b *BSI) GetCardinality() uint64 {

    return b.eBM.GetCardinality()
}

// Number of bits needed to represent values.
func (b *BSI) BitCount() int {

    return len(b.bA)
}

// Set a value for a given columnID.
func (b *BSI) SetValue(columnID uint64, value int64) {

    // If max/min values are set to zero then automatically determine bit array size
    if b.MaxValue == 0 && b.MinValue == 0 {
        ba := make([]*roaring.Bitmap, bits.Len64(uint64(value)))
        for i := len(ba) - b.BitCount(); i > 0; i-- {
            b.bA = append(b.bA, roaring.NewBitmap())
        }
    }

    var wg sync.WaitGroup

    for i := 0; i < b.BitCount(); i++ {
        wg.Add(1)
        go func(j int) {
            defer wg.Done()
            if uint64(value) & (1 << uint64(j)) > 0 {
                b.bA[j].Add(uint32(columnID))
            } else {
                b.bA[j].Remove(uint32(columnID))
            }
        }(i)
    }
    wg.Wait()
    b.eBM.Add(uint32(columnID))
}

// Get the value at the column ID.  Second param will be false for non-existant values.
func (b *BSI) GetValue(columnID uint64) (int64, bool) {

    value := int64(0)
    exists := b.eBM.Contains(uint32(columnID))
    if !exists {
        return value, exists
    }
    for i := 0; i < b.BitCount(); i++ {
        if b.bA[i].Contains(uint32(columnID)) {
            value |= (1 << uint64(i))
        }
    }
    return int64(value), exists
}

type action func(t *task, batch []uint32, resutsChan chan *roaring.Bitmap, wg *sync.WaitGroup)

func parallelExecutor (parallelism int, t *task, e action,
        foundSet *roaring.Bitmap) *roaring.Bitmap {

    var n int = parallelism
    if n == 0 {
        n = runtime.NumCPU()
    }

    resultsChan := make(chan *roaring.Bitmap, n)

    card := foundSet.GetCardinality()
    x := uint64(math.Floor(float64(card) / float64(n)))

    remainder := card - (x * uint64(n))
    var batch []uint32
    var wg sync.WaitGroup
    iter := foundSet.ManyIterator()
    for i := 0; i < n; i++ {
        if i == n - 1 {
            batch = make([]uint32, x + remainder)
        } else {
            batch = make([]uint32, x)
        }
        iter.NextMany(batch)
        wg.Add(1)
        go e(t, batch, resultsChan, &wg)
    }

    wg.Wait()

    close(resultsChan)

    ba := make([]*roaring.Bitmap, 0)
    for bm := range resultsChan {
        ba = append(ba, bm)
    }

    return roaring.ParOr(0, ba...)

}

type Operation int

const (
 LT Operation = 1 + iota
 LE
 EQ
 GE
 GT
 RANGE
)

type task struct {
    bsi           *BSI
    op            Operation
    valueOrStart  int64
    end           int64
    values        map[int64]struct{}
    bits          *roaring.Bitmap
}

/*
 * For all operations with the exception of RANGE, the value to be compared is specified by valueOrStart.
 * For the RANGE parameter the comparison criteria is >= valueOrStart and <= end.
 * The parallelism parameter indicates the number of CPU threads to be applied for processing.  A value
 * of zero indicates that all available CPU resources will be potentially utilized.
 */
func (b *BSI) CompareValue(parallelism int, op Operation, valueOrStart, end int64,
        foundSet *roaring.Bitmap) *roaring.Bitmap {

    comp := &task{bsi: b, op: op, valueOrStart: valueOrStart, end: end}
    if foundSet == nil {
        return parallelExecutor(parallelism, comp, compareValue, b.eBM)
    } else {
        return parallelExecutor(parallelism, comp, compareValue, foundSet)
    }
}

func compareValue(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg *sync.WaitGroup) {

    defer wg.Done()

    results := roaring.NewBitmap()

    for i := 0; i < len(batch); i++ {
        cID := batch[i]
        eq1, eq2 := true, true
        lt1, lt2, gt1 := false, false, false
        for j := e.bsi.BitCount() - 1; j >= 0; j-- {
            sliceContainsBit := e.bsi.bA[j].Contains(cID)

            if uint64(e.valueOrStart) & (1 << uint64(j)) > 0 {
                // BIT in value is SET
                if ! sliceContainsBit {
                    if eq1 {
                        if e.op == LT || e.op == LE {
                            lt1 = true
                        }
                        eq1 = false
                        break
                    }
                }
            } else {
                // BIT in value is CLEAR
                if sliceContainsBit {
                    if eq1 {
                        if e.op == GT || e.op == GE || e.op == RANGE {
                            gt1 = true
                        }
                        eq1 = false
                        if e.op != RANGE {
                            break
                        }
                    }
                }
            }

            if e.op == RANGE && uint64(e.end) & (1 << uint64(j)) > 0 {
                // BIT in value is SET
                if ! sliceContainsBit {
                    if eq2 {
                        lt2 = true
                        eq2 = false
                    }
                }
            } else {
                // BIT in value is CLEAR
                if sliceContainsBit {
                    if eq2 {
                        eq2 = false
                    }
                }
            }

        }

        switch e.op {
        case LT:
            if lt1 {
                results.Add(cID)
            }
        case LE:
            if eq1 || lt1 {
                results.Add(cID)
            }
        case EQ:
            if eq1 {
                results.Add(cID)
            }
        case GE:
            if eq1 || gt1 {
                results.Add(cID)
            }
        case GT:
            if gt1 {
                results.Add(cID)
            }
        case RANGE:
            if (eq1 || gt1) && (eq2 || lt2) {
                results.Add(cID)
            }
        default:
            if eq1 {
                results.Add(cID)
            }
        }
    }

    resultsChan <- results
}

/*
 * Sum all values contained within the foundSet.   As a convenience, the cardinality of the foundSet
 * is also returned (for calculating the average).
 */
func (b *BSI) Sum(foundSet *roaring.Bitmap) (sum int64, count uint64) {

    sum = int64(0)
    count = foundSet.GetCardinality()

    iter := foundSet.Iterator()
    for iter.HasNext() {
        cID := iter.Next()
        if value, ok := b.GetValue(uint64(cID)); ok {
            sum += value
        }
    }
    return
}

func (b *BSI) Transpose() *roaring.Bitmap {

    return b.IntersectAndTranspose(0, b.eBM)
}

/*
 * Matrix Transpose function.  Return a bitmap such that the values are represented as column IDs
 * in the returned bitmap. This is accomplished by iterating over the foundSet and only including
 * the column IDs in the source (foundSet) as compared with this BSI.  This can be useful for
 * vectoring one set of integers to another.
 */
func (b *BSI) IntersectAndTranspose(parallelism int, foundSet *roaring.Bitmap) *roaring.Bitmap {

    trans := &task{bsi: b}
    return parallelExecutor(parallelism, trans, transpose, foundSet)
}

func transpose(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg *sync.WaitGroup) {

    defer wg.Done()

    results := roaring.NewBitmap()
    for _, cID := range batch {
        if value, ok := e.bsi.GetValue(uint64(cID)); ok {
            results.Add(uint32(value))
        }
    }
    resultsChan <- results
}

/*
 * ParOr is intended primarily to be a concatenation function to be used during bulk load operations.
 * Care should be taken to make sure that columnIDs do not overlap (unless overlapping values are
 * identical).
 */
func (b *BSI) ParOr(parallelism int, bsis ...*BSI) {

    // Consolidate sets
    bits := len(b.bA)
    for i := 0; i < len(bsis); i++ {
        if len(bsis[i].bA) > bits {
            bits = bsis[i].BitCount()
        }
    }

    // Make sure we have enough bit slices
    for bits > b.BitCount() {
        b.bA = append(b.bA, roaring.NewBitmap())
    }

    a := make([][]*roaring.Bitmap, bits)
    for i, _ := range a {
        a[i] = make([]*roaring.Bitmap, 0)
        for _, x := range bsis {
            if len(x.bA) > i {
                a[i] = append(a[i], x.bA[i])
            } else {
                a[i] = []*roaring.Bitmap{roaring.NewBitmap()}
            }
        }
    }

    // Consolidate existence bit maps
    ebms := make([]*roaring.Bitmap, len(bsis))
    for i, _ := range ebms {
        ebms[i] = bsis[i].eBM
    }

    // First merge all the bit slices from all bsi maps that exist in target
    var wg sync.WaitGroup
    for i := 0; i < bits; i++ {
        wg.Add(1)
        go func(j int) {
            defer wg.Done()
            x := []*roaring.Bitmap{b.bA[j]}
            x = append(x, a[j]...)
            b.bA[j] = roaring.ParOr(parallelism, x...)
        }(i)
    }
    wg.Wait()

    // merge all the EBM maps
    x := []*roaring.Bitmap{b.eBM}
    x = append(x, ebms...)
    b.eBM = roaring.ParOr(parallelism, x...)
}

/*
 * De-serialize a BSI.  The value at bitData[0] is the EBM.  Other indices are in least to most
 * significance order starting at bitData[1] (bit position 0).
 */
func (b *BSI) UnmarshalBinary(bitData [][]byte) error {

    for i := 1; i < len(bitData); i++ {
        if bitData == nil || len(bitData[i]) == 0 {
            continue
        }
        if b.BitCount() < i {
            b.bA = append(b.bA, roaring.NewBitmap())
        }
        if err := b.bA[i - 1].UnmarshalBinary(bitData[i]); err != nil {
            return err
        }
    }
    // First element of bitData is the EBM
    if bitData[0] == nil {
        b.eBM = roaring.NewBitmap()
        return nil
    }
    if err := b.eBM.UnmarshalBinary(bitData[0]); err != nil {
        return err
    }
    return nil
}

/*
 * Serialize a BSI.
 */
func (b *BSI) MarshalBinary() ([][]byte, error) {

    var err error
    data := make([][]byte, b.BitCount() + 1)
    // Add extra element for EBM (BitCount() + 1)
    for i := 1; i < b.BitCount() + 1; i++ {
        data[i], err = b.bA[i - 1].MarshalBinary()
        if err != nil {
            return nil, err
        }
    }
    // Marshal EBM
    data[0], err = b.eBM.MarshalBinary()
    if err != nil {
        return nil, err
    }
    return data, nil
}

/*
 * Return a bitmap containing the column IDs where the values are contained within the list of values
 * provided.
 */
func (b *BSI) BatchEqual(parallelism int, values []int64) *roaring.Bitmap {

    valMap := make(map[int64]struct{}, len(values))
    for i := 0; i < len(values); i++ {
        valMap[values[i]] = struct{}{}
    }
    comp := &task{bsi: b, values: valMap}
    return parallelExecutor(parallelism, comp, batchEqual, b.eBM)
}

func batchEqual(e *task, batch []uint32, resultsChan chan *roaring.Bitmap,
        wg *sync.WaitGroup) {

    defer wg.Done()

    results := roaring.NewBitmap()

    for i := 0; i < len(batch); i++ {
        cID := batch[i]
        if value, ok := e.bsi.GetValue(uint64(cID)); ok {
            if _, yes := e.values[int64(value)]; yes {
                results.Add(cID)
            }
        }
    }
    resultsChan <- results
}

/*
 * Bits that exist in the target are cleared if they are also in the found set.
 */
func ClearBits(foundSet, target *roaring.Bitmap) {
    iter := foundSet.Iterator()
    for iter.HasNext() {
       cID := iter.Next()
       target.Remove(cID)
    }
}

func (b *BSI) ClearValues(foundSet *roaring.Bitmap) {

    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        defer wg.Done()
        ClearBits(foundSet, b.eBM)
    }()
    for i := 0; i < b.BitCount(); i++ {
        wg.Add(1)
        go func(j int) {
            defer wg.Done()
            ClearBits(foundSet, b.bA[j])
        }(i)
    }
    wg.Wait()
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 e2e

import (
	"time"

	"k8s.io/kubernetes/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api/v1"
	metav1 "k8s.io/kubernetes/pkg/apis/meta/v1"
	storage "k8s.io/kubernetes/pkg/apis/storage/v1beta1"
	storageutil "k8s.io/kubernetes/pkg/apis/storage/v1beta1/util"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	// Requested size of the volume
	requestedSize = "1500Mi"
	// Expected size of the volume is 2GiB, because all three supported cloud
	// providers allocate volumes in 1GiB chunks.
	expectedSize = "2Gi"
)

func testDynamicProvisioning(client clientset.Interface, claim *v1.PersistentVolumeClaim) {
	err := framework.WaitForPersistentVolumeClaimPhase(v1.ClaimBound, client, claim.Namespace, claim.Name, framework.Poll, framework.ClaimProvisionTimeout)
	Expect(err).NotTo(HaveOccurred())

	By("checking the claim")
	// Get new copy of the claim
	claim, err = client.Core().PersistentVolumeClaims(claim.Namespace).Get(claim.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	// Get the bound PV
	pv, err := client.Core().PersistentVolumes().Get(claim.Spec.VolumeName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	// Check sizes
	expectedCapacity := resource.MustParse(expectedSize)
	pvCapacity := pv.Spec.Capacity[v1.ResourceName(v1.ResourceStorage)]
	Expect(pvCapacity.Value()).To(Equal(expectedCapacity.Value()))

	requestedCapacity := resource.MustParse(requestedSize)
	claimCapacity := claim.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
	Expect(claimCapacity.Value()).To(Equal(requestedCapacity.Value()))

	// Check PV properties
	Expect(pv.Spec.PersistentVolumeReclaimPolicy).To(Equal(v1.PersistentVolumeReclaimDelete))
	expectedAccessModes := []v1.PersistentVolumeAccessMode{v1.ReadWriteOnce}
	Expect(pv.Spec.AccessModes).To(Equal(expectedAccessModes))
	Expect(pv.Spec.ClaimRef.Name).To(Equal(claim.ObjectMeta.Name))
	Expect(pv.Spec.ClaimRef.Namespace).To(Equal(claim.ObjectMeta.Namespace))

	// We start two pods:
	// - The first writes 'hello word' to the /mnt/test (= the volume).
	// - The second one runs grep 'hello world' on /mnt/test.
	// If both succeed, Kubernetes actually allocated something that is
	// persistent across pods.
	By("checking the created volume is writable")
	runInPodWithVolume(client, claim.Namespace, claim.Name, "echo 'hello world' > /mnt/test/data")

	By("checking the created volume is readable and retains data")
	runInPodWithVolume(client, claim.Namespace, claim.Name, "grep 'hello world' /mnt/test/data")

	// Ugly hack: if we delete the AWS/GCE/OpenStack volume here, it will
	// probably collide with destruction of the pods above - the pods
	// still have the volume attached (kubelet is slow...) and deletion
	// of attached volume is not allowed by AWS/GCE/OpenStack.
	// Kubernetes *will* retry deletion several times in
	// pvclaimbinder-sync-period.
	// So, technically, this sleep is not needed. On the other hand,
	// the sync perion is 10 minutes and we really don't want to wait
	// 10 minutes here. There is no way how to see if kubelet is
	// finished with cleaning volumes. A small sleep here actually
	// speeds up the test!
	// Three minutes should be enough to clean up the pods properly.
	// We've seen GCE PD detach to take more than 1 minute.
	By("Sleeping to let kubelet destroy all pods")
	time.Sleep(3 * time.Minute)

	By("deleting the claim")
	framework.ExpectNoError(client.Core().PersistentVolumeClaims(claim.Namespace).Delete(claim.Name, nil))

	// Wait for the PV to get deleted too.
	framework.ExpectNoError(framework.WaitForPersistentVolumeDeleted(client, pv.Name, 5*time.Second, 20*time.Minute))
}

var _ = framework.KubeDescribe("Dynamic provisioning", func() {
	f := framework.NewDefaultFramework("volume-provisioning")

	// filled in BeforeEach
	var c clientset.Interface
	var ns string

	BeforeEach(func() {
		c = f.ClientSet
		ns = f.Namespace.Name
	})

	framework.KubeDescribe("DynamicProvisioner", func() {
		It("should create and delete persistent volumes [Slow]", func() {
			framework.SkipUnlessProviderIs("openstack", "gce", "aws", "gke")

			By("creating a StorageClass")
			class := newStorageClass()
			_, err := c.Storage().StorageClasses().Create(class)
			defer c.Storage().StorageClasses().Delete(class.Name, nil)
			Expect(err).NotTo(HaveOccurred())

			By("creating a claim with a dynamic provisioning annotation")
			claim := newClaim(ns, false)
			defer func() {
				c.Core().PersistentVolumeClaims(ns).Delete(claim.Name, nil)
			}()
			claim, err = c.Core().PersistentVolumeClaims(ns).Create(claim)
			Expect(err).NotTo(HaveOccurred())

			testDynamicProvisioning(c, claim)
		})
	})

	framework.KubeDescribe("DynamicProvisioner Alpha", func() {
		It("should create and delete alpha persistent volumes [Slow]", func() {
			framework.SkipUnlessProviderIs("openstack", "gce", "aws", "gke")

			By("creating a claim with an alpha dynamic provisioning annotation")
			claim := newClaim(ns, true)
			defer func() {
				c.Core().PersistentVolumeClaims(ns).Delete(claim.Name, nil)
			}()
			claim, err := c.Core().PersistentVolumeClaims(ns).Create(claim)
			Expect(err).NotTo(HaveOccurred())

			testDynamicProvisioning(c, claim)
		})
	})
})

func newClaim(ns string, alpha bool) *v1.PersistentVolumeClaim {
	claim := v1.PersistentVolumeClaim{
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "pvc-",
			Namespace:    ns,
		},
		Spec: v1.PersistentVolumeClaimSpec{
			AccessModes: []v1.PersistentVolumeAccessMode{
				v1.ReadWriteOnce,
			},
			Resources: v1.ResourceRequirements{
				Requests: v1.ResourceList{
					v1.ResourceName(v1.ResourceStorage): resource.MustParse(requestedSize),
				},
			},
		},
	}

	if alpha {
		claim.Annotations = map[string]string{
			storageutil.AlphaStorageClassAnnotation: "",
		}
	} else {
		claim.Annotations = map[string]string{
			storageutil.StorageClassAnnotation: "fast",
		}

	}

	return &claim
}

// runInPodWithVolume runs a command in a pod with given claim mounted to /mnt directory.
func runInPodWithVolume(c clientset.Interface, ns, claimName, command string) {
	pod := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "pvc-volume-tester-",
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:    "volume-tester",
					Image:   "gcr.io/google_containers/busybox:1.24",
					Command: []string{"/bin/sh"},
					Args:    []string{"-c", command},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      "my-volume",
							MountPath: "/mnt/test",
						},
					},
				},
			},
			RestartPolicy: v1.RestartPolicyNever,
			Volumes: []v1.Volume{
				{
					Name: "my-volume",
					VolumeSource: v1.VolumeSource{
						PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
							ClaimName: claimName,
							ReadOnly:  false,
						},
					},
				},
			},
		},
	}
	pod, err := c.Core().Pods(ns).Create(pod)
	defer func() {
		framework.ExpectNoError(c.Core().Pods(ns).Delete(pod.Name, nil))
	}()
	framework.ExpectNoError(err, "Failed to create pod: %v", err)
	framework.ExpectNoError(framework.WaitForPodSuccessInNamespaceSlow(c, pod.Name, pod.Namespace))
}

func newStorageClass() *storage.StorageClass {
	var pluginName string

	switch {
	case framework.ProviderIs("gke"), framework.ProviderIs("gce"):
		pluginName = "kubernetes.io/gce-pd"
	case framework.ProviderIs("aws"):
		pluginName = "kubernetes.io/aws-ebs"
	case framework.ProviderIs("openstack"):
		pluginName = "kubernetes.io/cinder"
	}

	return &storage.StorageClass{
		TypeMeta: metav1.TypeMeta{
			Kind: "StorageClass",
		},
		ObjectMeta: v1.ObjectMeta{
			Name: "fast",
		},
		Provisioner: pluginName,
	}
}
<commit_msg>Add e2e test for external pv provisioning<commit_after>/*
Copyright 2016 The Kubernetes 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 e2e

import (
	"time"

	"k8s.io/kubernetes/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api/v1"
	metav1 "k8s.io/kubernetes/pkg/apis/meta/v1"
	storage "k8s.io/kubernetes/pkg/apis/storage/v1beta1"
	storageutil "k8s.io/kubernetes/pkg/apis/storage/v1beta1/util"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	// Requested size of the volume
	requestedSize = "1500Mi"
	// Expected size of the volume is 2GiB, because all three supported cloud
	// providers allocate volumes in 1GiB chunks.
	expectedSize = "2Gi"
	// Expected size of the externally provisioned volume depends on the external
	// provisioner: for nfs-provisioner used here, it's equal to requested
	externalExpectedSize = "1500Mi"
	// Plugin name of the external provisioner
	externalPluginName = "example.com/nfs"
)

func testDynamicProvisioning(client clientset.Interface, claim *v1.PersistentVolumeClaim, external bool) {
	err := framework.WaitForPersistentVolumeClaimPhase(v1.ClaimBound, client, claim.Namespace, claim.Name, framework.Poll, framework.ClaimProvisionTimeout)
	Expect(err).NotTo(HaveOccurred())

	By("checking the claim")
	// Get new copy of the claim
	claim, err = client.Core().PersistentVolumeClaims(claim.Namespace).Get(claim.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	// Get the bound PV
	pv, err := client.Core().PersistentVolumes().Get(claim.Spec.VolumeName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	// Check sizes
	expectedCapacity := resource.MustParse(expectedSize)
	if external {
		expectedCapacity = resource.MustParse(externalExpectedSize)
	}
	pvCapacity := pv.Spec.Capacity[v1.ResourceName(v1.ResourceStorage)]
	Expect(pvCapacity.Value()).To(Equal(expectedCapacity.Value()))

	requestedCapacity := resource.MustParse(requestedSize)
	claimCapacity := claim.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
	Expect(claimCapacity.Value()).To(Equal(requestedCapacity.Value()))

	// Check PV properties
	Expect(pv.Spec.PersistentVolumeReclaimPolicy).To(Equal(v1.PersistentVolumeReclaimDelete))
	expectedAccessModes := []v1.PersistentVolumeAccessMode{v1.ReadWriteOnce}
	Expect(pv.Spec.AccessModes).To(Equal(expectedAccessModes))
	Expect(pv.Spec.ClaimRef.Name).To(Equal(claim.ObjectMeta.Name))
	Expect(pv.Spec.ClaimRef.Namespace).To(Equal(claim.ObjectMeta.Namespace))

	// We start two pods:
	// - The first writes 'hello word' to the /mnt/test (= the volume).
	// - The second one runs grep 'hello world' on /mnt/test.
	// If both succeed, Kubernetes actually allocated something that is
	// persistent across pods.
	By("checking the created volume is writable")
	runInPodWithVolume(client, claim.Namespace, claim.Name, "echo 'hello world' > /mnt/test/data")

	By("checking the created volume is readable and retains data")
	runInPodWithVolume(client, claim.Namespace, claim.Name, "grep 'hello world' /mnt/test/data")

	if !external {
		// Ugly hack: if we delete the AWS/GCE/OpenStack volume here, it will
		// probably collide with destruction of the pods above - the pods
		// still have the volume attached (kubelet is slow...) and deletion
		// of attached volume is not allowed by AWS/GCE/OpenStack.
		// Kubernetes *will* retry deletion several times in
		// pvclaimbinder-sync-period.
		// So, technically, this sleep is not needed. On the other hand,
		// the sync perion is 10 minutes and we really don't want to wait
		// 10 minutes here. There is no way how to see if kubelet is
		// finished with cleaning volumes. A small sleep here actually
		// speeds up the test!
		// Three minutes should be enough to clean up the pods properly.
		// We've seen GCE PD detach to take more than 1 minute.
		By("Sleeping to let kubelet destroy all pods")
		time.Sleep(3 * time.Minute)
	}

	By("deleting the claim")
	framework.ExpectNoError(client.Core().PersistentVolumeClaims(claim.Namespace).Delete(claim.Name, nil))

	// Wait for the PV to get deleted too.
	framework.ExpectNoError(framework.WaitForPersistentVolumeDeleted(client, pv.Name, 5*time.Second, 20*time.Minute))
}

var _ = framework.KubeDescribe("Dynamic provisioning", func() {
	f := framework.NewDefaultFramework("volume-provisioning")

	// filled in BeforeEach
	var c clientset.Interface
	var ns string

	BeforeEach(func() {
		c = f.ClientSet
		ns = f.Namespace.Name
	})

	framework.KubeDescribe("DynamicProvisioner", func() {
		It("should create and delete persistent volumes [Slow]", func() {
			framework.SkipUnlessProviderIs("openstack", "gce", "aws", "gke")

			By("creating a StorageClass")
			class := newStorageClass("")
			_, err := c.Storage().StorageClasses().Create(class)
			defer c.Storage().StorageClasses().Delete(class.Name, nil)
			Expect(err).NotTo(HaveOccurred())

			By("creating a claim with a dynamic provisioning annotation")
			claim := newClaim(ns, false)
			defer func() {
				c.Core().PersistentVolumeClaims(ns).Delete(claim.Name, nil)
			}()
			claim, err = c.Core().PersistentVolumeClaims(ns).Create(claim)
			Expect(err).NotTo(HaveOccurred())

			testDynamicProvisioning(c, claim, false)
		})
	})

	framework.KubeDescribe("DynamicProvisioner Alpha", func() {
		It("should create and delete alpha persistent volumes [Slow]", func() {
			framework.SkipUnlessProviderIs("openstack", "gce", "aws", "gke")

			By("creating a claim with an alpha dynamic provisioning annotation")
			claim := newClaim(ns, true)
			defer func() {
				c.Core().PersistentVolumeClaims(ns).Delete(claim.Name, nil)
			}()
			claim, err := c.Core().PersistentVolumeClaims(ns).Create(claim)
			Expect(err).NotTo(HaveOccurred())

			testDynamicProvisioning(c, claim, false)
		})
	})

	framework.KubeDescribe("DynamicProvisioner External", func() {
		It("should let an external dynamic provisioner create and delete persistent volumes [Slow]", func() {
			By("creating an external dynamic provisioner pod")
			pod := startExternalProvisioner(c, ns)
			defer c.Core().Pods(ns).Delete(pod.Name, nil)

			By("creating a StorageClass")
			class := newStorageClass(externalPluginName)
			_, err := c.Storage().StorageClasses().Create(class)
			defer c.Storage().StorageClasses().Delete(class.Name, nil)
			Expect(err).NotTo(HaveOccurred())

			By("creating a claim with a dynamic provisioning annotation")
			claim := newClaim(ns, false)
			defer func() {
				c.Core().PersistentVolumeClaims(ns).Delete(claim.Name, nil)
			}()
			claim, err = c.Core().PersistentVolumeClaims(ns).Create(claim)
			Expect(err).NotTo(HaveOccurred())

			testDynamicProvisioning(c, claim, true)
		})
	})
})

func newClaim(ns string, alpha bool) *v1.PersistentVolumeClaim {
	claim := v1.PersistentVolumeClaim{
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "pvc-",
			Namespace:    ns,
		},
		Spec: v1.PersistentVolumeClaimSpec{
			AccessModes: []v1.PersistentVolumeAccessMode{
				v1.ReadWriteOnce,
			},
			Resources: v1.ResourceRequirements{
				Requests: v1.ResourceList{
					v1.ResourceName(v1.ResourceStorage): resource.MustParse(requestedSize),
				},
			},
		},
	}

	if alpha {
		claim.Annotations = map[string]string{
			storageutil.AlphaStorageClassAnnotation: "",
		}
	} else {
		claim.Annotations = map[string]string{
			storageutil.StorageClassAnnotation: "fast",
		}

	}

	return &claim
}

// runInPodWithVolume runs a command in a pod with given claim mounted to /mnt directory.
func runInPodWithVolume(c clientset.Interface, ns, claimName, command string) {
	pod := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "pvc-volume-tester-",
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:    "volume-tester",
					Image:   "gcr.io/google_containers/busybox:1.24",
					Command: []string{"/bin/sh"},
					Args:    []string{"-c", command},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      "my-volume",
							MountPath: "/mnt/test",
						},
					},
				},
			},
			RestartPolicy: v1.RestartPolicyNever,
			Volumes: []v1.Volume{
				{
					Name: "my-volume",
					VolumeSource: v1.VolumeSource{
						PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
							ClaimName: claimName,
							ReadOnly:  false,
						},
					},
				},
			},
		},
	}
	pod, err := c.Core().Pods(ns).Create(pod)
	defer func() {
		framework.ExpectNoError(c.Core().Pods(ns).Delete(pod.Name, nil))
	}()
	framework.ExpectNoError(err, "Failed to create pod: %v", err)
	framework.ExpectNoError(framework.WaitForPodSuccessInNamespaceSlow(c, pod.Name, pod.Namespace))
}

func newStorageClass(pluginName string) *storage.StorageClass {
	if pluginName == "" {
		switch {
		case framework.ProviderIs("gke"), framework.ProviderIs("gce"):
			pluginName = "kubernetes.io/gce-pd"
		case framework.ProviderIs("aws"):
			pluginName = "kubernetes.io/aws-ebs"
		case framework.ProviderIs("openstack"):
			pluginName = "kubernetes.io/cinder"
		}
	}

	return &storage.StorageClass{
		TypeMeta: metav1.TypeMeta{
			Kind: "StorageClass",
		},
		ObjectMeta: v1.ObjectMeta{
			Name: "fast",
		},
		Provisioner: pluginName,
	}
}

func startExternalProvisioner(c clientset.Interface, ns string) *v1.Pod {
	podClient := c.Core().Pods(ns)

	provisionerPod := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "external-provisioner-",
		},

		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "nfs-provisioner",
					Image: "quay.io/kubernetes_incubator/nfs-provisioner:v1.0.1",
					SecurityContext: &v1.SecurityContext{
						Capabilities: &v1.Capabilities{
							Add: []v1.Capability{"DAC_READ_SEARCH"},
						},
					},
					Args: []string{
						"-provisioner=" + externalPluginName,
						"-grace-period=0",
					},
					Ports: []v1.ContainerPort{
						{Name: "nfs", ContainerPort: 2049},
						{Name: "mountd", ContainerPort: 20048},
						{Name: "rpcbind", ContainerPort: 111},
						{Name: "rpcbind-udp", ContainerPort: 111, Protocol: v1.ProtocolUDP},
					},
					Env: []v1.EnvVar{
						{
							Name: "POD_IP",
							ValueFrom: &v1.EnvVarSource{
								FieldRef: &v1.ObjectFieldSelector{
									FieldPath: "status.podIP",
								},
							},
						},
					},
					ImagePullPolicy: v1.PullIfNotPresent,
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      "export-volume",
							MountPath: "/export",
						},
					},
				},
			},
			Volumes: []v1.Volume{
				{
					Name: "export-volume",
					VolumeSource: v1.VolumeSource{
						EmptyDir: &v1.EmptyDirVolumeSource{},
					},
				},
			},
		},
	}
	provisionerPod, err := podClient.Create(provisionerPod)
	framework.ExpectNoError(err, "Failed to create %s pod: %v", provisionerPod.Name, err)

	framework.ExpectNoError(framework.WaitForPodRunningInNamespace(c, provisionerPod))

	By("locating the provisioner pod")
	pod, err := podClient.Get(provisionerPod.Name, metav1.GetOptions{})
	framework.ExpectNoError(err, "Cannot locate the provisioner pod %v: %v", provisionerPod.Name, err)

	By("sleeping a bit to give the provisioner time to start")
	time.Sleep(5 * time.Second)
	return pod
}
<|endoftext|>