<commit_before>package http2

import (
	"math/rand"
	"net/http"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dropbox/godropbox/errors"
)

const (
	// Number of attempts to try to connect to a target host.
	connectionAttempts = 3

	// Default instance mark down duration.
	markDownDuration = 10 * time.Second
)

type LBStrategy int

const (
	// In 'RoundRobin' load balancing strategy requests are sent to
	// different hosts in round robin fashion.
	LBRoundRobin LBStrategy = 0
	// In 'Fixed' load balancing strategy requests are routed to same host,
	// others are used only in case of failover.
	LBFixed LBStrategy = 1
)

type LoadBalancedPool struct {
	lock sync.RWMutex

	// Maps "host:port" -> instancePool.
	instances    map[string]*instancePool
	instanceList instancePoolSlice
	// Atomic counter that is used for round robining instances
	// from instanceList.
	instanceIdx uint64

	// UNIX epoch time in seconds that represents time till address is considered
	// as down and unusable.
	markDownUntil []int64

	params   ConnectionParams // Parameters for creating SimplePool-s.
	strategy LBStrategy       // Load balancing strategy.
}

type instancePool struct {
	SimplePool
	instanceId int
}

type instancePoolSlice []*instancePool

func (s instancePoolSlice) Len() int      { return len(s) }
func (s instancePoolSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// instancePoolSlice sorts by instanceId in descending order.
func (s instancePoolSlice) Less(i, j int) bool { return s[i].instanceId > s[j].instanceId }

type LBPoolInstanceInfo struct {
	InstanceId int
	Addr       string
}

func NewLoadBalancedPool(params ConnectionParams) *LoadBalancedPool {
	return &LoadBalancedPool{
		instances:     make(map[string]*instancePool),
		instanceList:  make(instancePoolSlice, 0),
		markDownUntil: make([]int64, 0),
		params:        params,
		strategy:      LBRoundRobin,
	}
}

// Sets Load Balancing strategy. Must be called before pool is actually put to use.
func (pool *LoadBalancedPool) SetStrategy(strategy LBStrategy) {
	pool.strategy = strategy
}

func (pool *LoadBalancedPool) newInstancePool(info LBPoolInstanceInfo) *instancePool {
	simplePool := NewSimplePool(info.Addr, pool.params)
	return &instancePool{SimplePool: *simplePool, instanceId: info.InstanceId}
}

func (pool *LoadBalancedPool) Update(instanceInfos []LBPoolInstanceInfo) {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	newInstances := make(map[string]*instancePool)
	var newInstanceList instancePoolSlice
	for _, instanceInfo := range instanceInfos {
		if _, ok := newInstances[instanceInfo.Addr]; !ok {
                        var instance *instancePool
                        if instance, ok = pool.instances[instanceInfo.Addr]; !ok {
                                instance = pool.newInstancePool(instanceInfo)
                        }
			newInstances[instanceInfo.Addr] = instance
			newInstanceList = append(newInstanceList, instance)
		}
	}
	switch pool.strategy {
	case LBRoundRobin:
		// In RoundRobin strategy, InstanceList is a randomly shuffled list of instances.
		for i, _ := range newInstanceList {
			randIdx := rand.Intn(i + 1)
			newInstanceList.Swap(i, randIdx)
		}
	case LBFixed:
		// In Fixed strategy, InstanceList is a sorted list, sorted by instanceId.
		sort.Sort(newInstanceList)
	}

	for addr, instancePool := range pool.instances {
		// Close out all InstancePools that are not needed anymore.
		if _, ok := newInstances[addr]; !ok {
			instancePool.Close()
		}
	}
	pool.instances = newInstances
	pool.instanceList = newInstanceList
	pool.markDownUntil = make([]int64, len(newInstanceList))
}

//
// Pool interface methods
//

// Issues an HTTP request, distributing more load to relatively unloaded instances.
func (pool *LoadBalancedPool) Do(req *http.Request) (*http.Response, error) {
	var requestErr error = nil
	for i := 0; ; i++ {
		idx, instance, isDown, err := pool.getInstance()
		if err != nil {
			return nil, errors.Wrap(err, "can't get HTTP connection")
		}
		if isDown && requestErr != nil {
			// If current pool instance is marked down, that means all instances in the pool
			// are most likely marked down, thus avoid performing any connect retries, to fail
			// faster.
			return nil, requestErr
		}

		resp, err := instance.Do(req)
		if err != nil || resp.StatusCode == 500 {
			// 500s are also treated as service being down momentarily,
			// note that even if all servers get marked down LBPool continues
			// to send requests in round robin manner, thus this provides extra
			// protection when service may still be up but have higher rate of
			// 500s for whatever reason.
			pool.markInstanceDown(idx, instance, time.Now().Add(markDownDuration).Unix())
		} else if isDown {
			// If an instance was marked as down, but succeeded, reset the mark down timer, so
			// instance is treated as healthy right away.
			pool.markInstanceUp(idx, instance)
		}
		if err != nil {
			if _, ok := err.(DialError); !ok {
				return resp, err
			}

			if (i + 1) < connectionAttempts {
				requestErr = err
				continue
			}
		}
		return resp, err
	}
}

// Checks out an HTTP connection from an instance pool, favoring less loaded instances.
func (pool *LoadBalancedPool) Get() (*http.Client, error) {
	_, instance, _, err := pool.getInstance()
	if err != nil {
		return nil, errors.Wrap(err, "can't get HTTP connection")
	}
	conn, err := instance.Get()
	if err != nil {
		return nil, errors.Wrap(err, "couldn't Get from LoadBalancedPool")
	}
	return conn, err
}

// Returns instance that isn't marked down, if all instances are
// marked as down it will just choose a next one.
func (pool *LoadBalancedPool) getInstance() (
	idx int,
	instance *instancePool,
	isDown bool,
	err error) {
	pool.lock.RLock()
	defer pool.lock.RUnlock()
	if len(pool.instanceList) == 0 {
		return 0, nil, false, errors.Newf("no available instances")
	}
	now := time.Now().Unix()
	for i := 0; i < len(pool.instanceList); i++ {
		switch pool.strategy {
		case LBRoundRobin:
			// In RoundRobin strategy instanceIdx keeps changing, to
			// achieve round robin load balancing.
			instanceIdx := atomic.AddUint64(&pool.instanceIdx, 1)
			idx = int(instanceIdx % uint64(len(pool.instanceList)))
		case LBFixed:
			// In Fixed strategy instances are always traversed in same
			// exact order.
			idx = i
		}

		if pool.markDownUntil[idx] < now {
			break
		}
	}
	return idx, pool.instanceList[idx], (pool.markDownUntil[idx] >= now), nil
}

// Returns a SimplePool for given instanceId, or an error if it does not exist.
// TODO(zviad): right now this scans all instances, thus if there are a lot of
// instances per partition it can become very slow. If it becomes a problem, fix it!
func (pool *LoadBalancedPool) GetInstancePool(instanceId int) (*SimplePool, error) {
	pool.lock.RLock()
	defer pool.lock.RUnlock()
	for _, instancePool := range pool.instanceList {
		if instancePool.instanceId == instanceId {
			return &instancePool.SimplePool, nil
		}
	}
	return nil, errors.Newf("InstanceId: %v not found in the pool", instanceId)
}

// Marks instance down till downUntil epoch in seconds.
func (pool *LoadBalancedPool) markInstanceDown(
	idx int, instance *instancePool, downUntil int64) {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	if idx < len(pool.instanceList) && pool.instanceList[idx] == instance {
		pool.markDownUntil[idx] = downUntil
	}
}

// Marks instance as ready to be used.
func (pool *LoadBalancedPool) markInstanceUp(
	idx int, instance *instancePool) {
	pool.markInstanceDown(idx, instance, 0)
}

func (pool *LoadBalancedPool) Close() {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	for _, instance := range pool.instances {
		instance.Close()
	}
}
<commit_msg>run gofmt<commit_after>package http2

import (
	"math/rand"
	"net/http"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dropbox/godropbox/errors"
)

const (
	// Number of attempts to try to connect to a target host.
	connectionAttempts = 3

	// Default instance mark down duration.
	markDownDuration = 10 * time.Second
)

type LBStrategy int

const (
	// In 'RoundRobin' load balancing strategy requests are sent to
	// different hosts in round robin fashion.
	LBRoundRobin LBStrategy = 0
	// In 'Fixed' load balancing strategy requests are routed to same host,
	// others are used only in case of failover.
	LBFixed LBStrategy = 1
)

type LoadBalancedPool struct {
	lock sync.RWMutex

	// Maps "host:port" -> instancePool.
	instances    map[string]*instancePool
	instanceList instancePoolSlice
	// Atomic counter that is used for round robining instances
	// from instanceList.
	instanceIdx uint64

	// UNIX epoch time in seconds that represents time till address is considered
	// as down and unusable.
	markDownUntil []int64

	params   ConnectionParams // Parameters for creating SimplePool-s.
	strategy LBStrategy       // Load balancing strategy.
}

type instancePool struct {
	SimplePool
	instanceId int
}

type instancePoolSlice []*instancePool

func (s instancePoolSlice) Len() int      { return len(s) }
func (s instancePoolSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// instancePoolSlice sorts by instanceId in descending order.
func (s instancePoolSlice) Less(i, j int) bool { return s[i].instanceId > s[j].instanceId }

type LBPoolInstanceInfo struct {
	InstanceId int
	Addr       string
}

func NewLoadBalancedPool(params ConnectionParams) *LoadBalancedPool {
	return &LoadBalancedPool{
		instances:     make(map[string]*instancePool),
		instanceList:  make(instancePoolSlice, 0),
		markDownUntil: make([]int64, 0),
		params:        params,
		strategy:      LBRoundRobin,
	}
}

// Sets Load Balancing strategy. Must be called before pool is actually put to use.
func (pool *LoadBalancedPool) SetStrategy(strategy LBStrategy) {
	pool.strategy = strategy
}

func (pool *LoadBalancedPool) newInstancePool(info LBPoolInstanceInfo) *instancePool {
	simplePool := NewSimplePool(info.Addr, pool.params)
	return &instancePool{SimplePool: *simplePool, instanceId: info.InstanceId}
}

func (pool *LoadBalancedPool) Update(instanceInfos []LBPoolInstanceInfo) {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	newInstances := make(map[string]*instancePool)
	var newInstanceList instancePoolSlice
	for _, instanceInfo := range instanceInfos {
		if _, ok := newInstances[instanceInfo.Addr]; !ok {
			var instance *instancePool
			if instance, ok = pool.instances[instanceInfo.Addr]; !ok {
				instance = pool.newInstancePool(instanceInfo)
			}
			newInstances[instanceInfo.Addr] = instance
			newInstanceList = append(newInstanceList, instance)
		}
	}
	switch pool.strategy {
	case LBRoundRobin:
		// In RoundRobin strategy, InstanceList is a randomly shuffled list of instances.
		for i, _ := range newInstanceList {
			randIdx := rand.Intn(i + 1)
			newInstanceList.Swap(i, randIdx)
		}
	case LBFixed:
		// In Fixed strategy, InstanceList is a sorted list, sorted by instanceId.
		sort.Sort(newInstanceList)
	}

	for addr, instancePool := range pool.instances {
		// Close out all InstancePools that are not needed anymore.
		if _, ok := newInstances[addr]; !ok {
			instancePool.Close()
		}
	}
	pool.instances = newInstances
	pool.instanceList = newInstanceList
	pool.markDownUntil = make([]int64, len(newInstanceList))
}

//
// Pool interface methods
//

// Issues an HTTP request, distributing more load to relatively unloaded instances.
func (pool *LoadBalancedPool) Do(req *http.Request) (*http.Response, error) {
	var requestErr error = nil
	for i := 0; ; i++ {
		idx, instance, isDown, err := pool.getInstance()
		if err != nil {
			return nil, errors.Wrap(err, "can't get HTTP connection")
		}
		if isDown && requestErr != nil {
			// If current pool instance is marked down, that means all instances in the pool
			// are most likely marked down, thus avoid performing any connect retries, to fail
			// faster.
			return nil, requestErr
		}

		resp, err := instance.Do(req)
		if err != nil || resp.StatusCode == 500 {
			// 500s are also treated as service being down momentarily,
			// note that even if all servers get marked down LBPool continues
			// to send requests in round robin manner, thus this provides extra
			// protection when service may still be up but have higher rate of
			// 500s for whatever reason.
			pool.markInstanceDown(idx, instance, time.Now().Add(markDownDuration).Unix())
		} else if isDown {
			// If an instance was marked as down, but succeeded, reset the mark down timer, so
			// instance is treated as healthy right away.
			pool.markInstanceUp(idx, instance)
		}
		if err != nil {
			if _, ok := err.(DialError); !ok {
				return resp, err
			}

			if (i + 1) < connectionAttempts {
				requestErr = err
				continue
			}
		}
		return resp, err
	}
}

// Checks out an HTTP connection from an instance pool, favoring less loaded instances.
func (pool *LoadBalancedPool) Get() (*http.Client, error) {
	_, instance, _, err := pool.getInstance()
	if err != nil {
		return nil, errors.Wrap(err, "can't get HTTP connection")
	}
	conn, err := instance.Get()
	if err != nil {
		return nil, errors.Wrap(err, "couldn't Get from LoadBalancedPool")
	}
	return conn, err
}

// Returns instance that isn't marked down, if all instances are
// marked as down it will just choose a next one.
func (pool *LoadBalancedPool) getInstance() (
	idx int,
	instance *instancePool,
	isDown bool,
	err error) {
	pool.lock.RLock()
	defer pool.lock.RUnlock()
	if len(pool.instanceList) == 0 {
		return 0, nil, false, errors.Newf("no available instances")
	}
	now := time.Now().Unix()
	for i := 0; i < len(pool.instanceList); i++ {
		switch pool.strategy {
		case LBRoundRobin:
			// In RoundRobin strategy instanceIdx keeps changing, to
			// achieve round robin load balancing.
			instanceIdx := atomic.AddUint64(&pool.instanceIdx, 1)
			idx = int(instanceIdx % uint64(len(pool.instanceList)))
		case LBFixed:
			// In Fixed strategy instances are always traversed in same
			// exact order.
			idx = i
		}

		if pool.markDownUntil[idx] < now {
			break
		}
	}
	return idx, pool.instanceList[idx], (pool.markDownUntil[idx] >= now), nil
}

// Returns a SimplePool for given instanceId, or an error if it does not exist.
// TODO(zviad): right now this scans all instances, thus if there are a lot of
// instances per partition it can become very slow. If it becomes a problem, fix it!
func (pool *LoadBalancedPool) GetInstancePool(instanceId int) (*SimplePool, error) {
	pool.lock.RLock()
	defer pool.lock.RUnlock()
	for _, instancePool := range pool.instanceList {
		if instancePool.instanceId == instanceId {
			return &instancePool.SimplePool, nil
		}
	}
	return nil, errors.Newf("InstanceId: %v not found in the pool", instanceId)
}

// Marks instance down till downUntil epoch in seconds.
func (pool *LoadBalancedPool) markInstanceDown(
	idx int, instance *instancePool, downUntil int64) {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	if idx < len(pool.instanceList) && pool.instanceList[idx] == instance {
		pool.markDownUntil[idx] = downUntil
	}
}

// Marks instance as ready to be used.
func (pool *LoadBalancedPool) markInstanceUp(
	idx int, instance *instancePool) {
	pool.markInstanceDown(idx, instance, 0)
}

func (pool *LoadBalancedPool) Close() {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	for _, instance := range pool.instances {
		instance.Close()
	}
}
<|endoftext|>


<commit_before>package extra

import (
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/golang/protobuf/proto"
	"github.com/name5566/leaf/log"
	"github.com/name5566/leaf/util"
	"math"
	"reflect"
)

// -------------------------
// | id | protobuf message |
// -------------------------
type ProtobufRouter struct {
	littleEndian bool
	msgInfo      []*ProtobufMsgInfo
	msgID        map[reflect.Type]uint16
}

type ProtobufMsgInfo struct {
	msgType    reflect.Type
	msgRouter  *util.CallRouter
	msgHandler ProtobufMsgHandler
}

type ProtobufMsgHandler func([]interface{})

func NewProtobufRouter() *ProtobufRouter {
	r := new(ProtobufRouter)
	r.littleEndian = false
	r.msgID = make(map[reflect.Type]uint16)
	return r
}

// It's dangerous to call the method on routing or marshaling (unmarshaling)
func (r *ProtobufRouter) SetByteOrder(littleEndian bool) {
	r.littleEndian = littleEndian
}

// It's dangerous to call the method on routing or marshaling (unmarshaling)
func (r *ProtobufRouter) RegisterRouter(msg proto.Message, msgRouter *util.CallRouter) {
	protobufMsgInfo(msg).msgRouter = msgRouter
}

// It's dangerous to call the method on routing or marshaling (unmarshaling)
func (r *ProtobufRouter) RegisterHandler(msg proto.Message, msgHandler ProtobufMsgHandler) {
	protobufMsgInfo(msg).msgHandler = msgHandler
}

func (r *ProtobufRouter) protobufMsgInfo(msg proto.Message) *ProtobufMsgInfo {
	msgType := reflect.TypeOf(msg)
	if msgType == nil || msgType.Kind() != reflect.Ptr {
		log.Fatal("protobuf message pointer required")
	}

	if id, ok := r.msgID[msgType]; ok {
		return r.msgInfo[id]
	}

	if len(r.msgInfo) >= math.MaxUint16 {
		log.Fatal("too many protobuf messages (max = %v)", math.MaxUint16)
	}

	i := new(ProtobufMsgInfo)
	i.msgType = msgType
	r.msgInfo = append(r.msgInfo, i)
	r.msgID[msgType] = uint16(len(r.msgInfo) - 1)
	return i
}

// goroutine safe
func (r *ProtobufRouter) Route(msg proto.Message, userData interface{}) error {
	msgType := reflect.TypeOf(msg)

	id, ok := r.msgID[msgType]
	if !ok {
		return errors.New(fmt.Sprintf("message %s not registered", msgType))
	}

	i := r.msgInfo[id]
	if i.msgHandler != nil {
		i.msgHandler(msgType, msg, userData)
	}
	if i.msgRouter != nil {
		i.msgRouter.AsynCall0(msgType, msg, userData)
	}
	return nil
}

// goroutine safe
func (r *ProtobufRouter) Unmarshal(data []byte) (proto.Message, error) {
	if len(data) < 2 {
		return nil, errors.New("protobuf data too short")
	}

	// id
	var id uint16
	if r.littleEndian {
		id = binary.LittleEndian.Uint16(data)
	} else {
		id = binary.BigEndian.Uint16(data)
	}

	// msg
	if id >= uint16(len(r.msgInfo)) {
		return nil, errors.New(fmt.Sprintf("message id %v not registered", id))
	}
	msg := reflect.New(r.msgInfo[id].msgType.Elem()).Interface().(proto.Message)
	return msg, proto.UnmarshalMerge(data[2:], msg)
}

// goroutine safe
func (r *ProtobufRouter) Marshal(msg proto.Message) (id []byte, data []byte, err error) {
	msgType := reflect.TypeOf(msg)

	// id
	_id, ok := r.msgID[msgType]
	if !ok {
		err = errors.New(fmt.Sprintf("message %s not registered", msgType))
		return
	}

	id = make([]byte, 2)
	if r.littleEndian {
		binary.LittleEndian.PutUint16(id, _id)
	} else {
		binary.BigEndian.PutUint16(id, _id)
	}

	// data
	data, err = proto.Marshal(msg)
	return
}
<commit_msg>add msg handler support.<commit_after>package extra

import (
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/golang/protobuf/proto"
	"github.com/name5566/leaf/log"
	"github.com/name5566/leaf/util"
	"math"
	"reflect"
)

// -------------------------
// | id | protobuf message |
// -------------------------
type ProtobufRouter struct {
	littleEndian bool
	msgInfo      []*ProtobufMsgInfo
	msgID        map[reflect.Type]uint16
}

type ProtobufMsgInfo struct {
	msgType    reflect.Type
	msgRouter  *util.CallRouter
	msgHandler ProtobufMsgHandler
}

type ProtobufMsgHandler func([]interface{})

func NewProtobufRouter() *ProtobufRouter {
	r := new(ProtobufRouter)
	r.littleEndian = false
	r.msgID = make(map[reflect.Type]uint16)
	return r
}

// It's dangerous to call the method on routing or marshaling (unmarshaling)
func (r *ProtobufRouter) SetByteOrder(littleEndian bool) {
	r.littleEndian = littleEndian
}

// It's dangerous to call the method on routing or marshaling (unmarshaling)
func (r *ProtobufRouter) RegisterRouter(msg proto.Message, msgRouter *util.CallRouter) {
	r.protobufMsgInfo(msg).msgRouter = msgRouter
}

// It's dangerous to call the method on routing or marshaling (unmarshaling)
func (r *ProtobufRouter) RegisterHandler(msg proto.Message, msgHandler ProtobufMsgHandler) {
	r.protobufMsgInfo(msg).msgHandler = msgHandler
}

func (r *ProtobufRouter) protobufMsgInfo(msg proto.Message) *ProtobufMsgInfo {
	msgType := reflect.TypeOf(msg)
	if msgType == nil || msgType.Kind() != reflect.Ptr {
		log.Fatal("protobuf message pointer required")
	}

	if id, ok := r.msgID[msgType]; ok {
		return r.msgInfo[id]
	}

	if len(r.msgInfo) >= math.MaxUint16 {
		log.Fatal("too many protobuf messages (max = %v)", math.MaxUint16)
	}

	i := new(ProtobufMsgInfo)
	i.msgType = msgType
	r.msgInfo = append(r.msgInfo, i)
	r.msgID[msgType] = uint16(len(r.msgInfo) - 1)
	return i
}

// goroutine safe
func (r *ProtobufRouter) Route(msg proto.Message, userData interface{}) error {
	msgType := reflect.TypeOf(msg)

	id, ok := r.msgID[msgType]
	if !ok {
		return errors.New(fmt.Sprintf("message %s not registered", msgType))
	}

	i := r.msgInfo[id]
	if i.msgHandler != nil {
		i.msgHandler([]interface{}{msgType, msg, userData})
	}
	if i.msgRouter != nil {
		i.msgRouter.AsynCall0(msgType, msg, userData)
	}
	return nil
}

// goroutine safe
func (r *ProtobufRouter) Unmarshal(data []byte) (proto.Message, error) {
	if len(data) < 2 {
		return nil, errors.New("protobuf data too short")
	}

	// id
	var id uint16
	if r.littleEndian {
		id = binary.LittleEndian.Uint16(data)
	} else {
		id = binary.BigEndian.Uint16(data)
	}

	// msg
	if id >= uint16(len(r.msgInfo)) {
		return nil, errors.New(fmt.Sprintf("message id %v not registered", id))
	}
	msg := reflect.New(r.msgInfo[id].msgType.Elem()).Interface().(proto.Message)
	return msg, proto.UnmarshalMerge(data[2:], msg)
}

// goroutine safe
func (r *ProtobufRouter) Marshal(msg proto.Message) (id []byte, data []byte, err error) {
	msgType := reflect.TypeOf(msg)

	// id
	_id, ok := r.msgID[msgType]
	if !ok {
		err = errors.New(fmt.Sprintf("message %s not registered", msgType))
		return
	}

	id = make([]byte, 2)
	if r.littleEndian {
		binary.LittleEndian.PutUint16(id, _id)
	} else {
		binary.BigEndian.PutUint16(id, _id)
	}

	// data
	data, err = proto.Marshal(msg)
	return
}
<|endoftext|>


<commit_before>package ash

import (
	"reflect"

	"github.com/256dpi/xo"
	"go.mongodb.org/mongo-driver/bson"

	"github.com/256dpi/fire"
	"github.com/256dpi/fire/coal"
	"github.com/256dpi/fire/stick"
)

// PolicyDataKey is the key used to store policies.
const PolicyDataKey = "ash:policy"

// Policy defines an authorization policy.
type Policy struct {
	// Access defines the general access.
	Access Access

	// Actions defines the allowed actions.
	Actions map[string]bool

	// The default fields used to determine the field access level. If the
	// getter is set, these will only be used to establish valid filters and
	// sorters during the fire.List operation authorizer stage, as well as the
	// writable fields during the fire.Create operation, otherwise the model
	// specific fields are used instead.
	Fields AccessTable

	// GetFilter is called to obtain the general resource access filter. This
	// filter is used to narrow down accessible resources in all operations
	// except fire.Create and fire.CollectionAction operations.
	GetFilter func(ctx *fire.Context) bson.M

	// VerifyID is called to for every direct model lookup to verify resource
	// level access. This function is called for all operations except fire.List,
	// fire.Create and fire.CollectionAction.
	VerifyID func(ctx *fire.Context, id coal.ID) Access

	// VerifyModel is called for every model load from the database to determine
	// resource level access. This function is called for all operations except
	// fire.Create and fire.CollectionAction.
	//
	// Note: The verification is deferred to the fire.Verifier stage.
	VerifyModel func(ctx *fire.Context, model coal.Model) Access

	// VerifyCreate and VerifyUpdate determine resource level access after all
	// modification have been applied. This function is called for the
	// fire.Create and fire.Update operation.
	//
	// Note: The verification is deferred to the fire.Validator stage.
	VerifyCreate func(ctx *fire.Context, model coal.Model) bool
	VerifyUpdate func(ctx *fire.Context, model coal.Model) bool

	// GetFields is called for every model to determine the field level access.
	// The policy should refrain from creating a new map for every request and
	// instead pre-allocate possible combinations and return those. The function
	// is called for all operations except fire.Delete, fire.CollectionAction and
	// fire.ResourceAction.
	GetFields func(ctx *fire.Context, model coal.Model) AccessTable

	// GetProperties is called for every model to determine the property level
	// access. The policy should refrain from creating a map for every request
	// and instead pre-allocate possible combinations and return those. The
	// function is called for all operations except fire.Delete,
	// fire.CollectionAction and fire.ResourceAction.
	GetProperties func(ctx *fire.Context, model coal.Model) AccessTable
}

// Selector is the function run to select a policy.
type Selector func(Identity) *Policy

// Select will run the provided function to select a policy for the supplied
// identity.
func Select(selector Selector) *fire.Callback {
	return fire.C("ash/Selector", fire.Authorizer, fire.All(), func(ctx *fire.Context) error {
		// get identity
		identity, _ := ctx.Data[IdentityDataKey]
		if identity == nil {
			return fire.ErrAccessDenied.Wrap()
		}

		// run selector
		policy := selector(identity)

		// check policy
		if policy == nil {
			return nil
		}

		// check stored
		if ctx.Data[PolicyDataKey] != nil {
			return xo.F("existing policy")
		}

		// store policy
		ctx.Data[PolicyDataKey] = policy

		return nil
	})
}

// SelectMatch will match the provided identity and on success use the provided
// factory to create a policy.
func SelectMatch(identity Identity, policy func(Identity) *Policy) *fire.Callback {
	// get type
	typ := reflect.TypeOf(identity)

	return Select(func(identity Identity) *Policy {
		// check type
		if typ != reflect.TypeOf(identity) {
			return nil
		}

		return policy(identity)
	})
}

// SelectPublic will math the public identity and use the provided factory to
// create a policy.
func SelectPublic(fn func() *Policy) *fire.Callback {
	return SelectMatch(&PublicIdentity{}, func(Identity) *Policy {
		return fn()
	})
}

// Execute will execute the selected policy or deny access.
func Execute() *fire.Callback {
	// prepare matchers
	getFilterMatcher := fire.Except(fire.Create | fire.CollectionAction)
	verifyIDMatcher := fire.Except(fire.List | fire.Create | fire.CollectionAction)
	verifyModelMatcher := fire.Except(fire.Create | fire.CollectionAction)
	verifyCreateMatcher := fire.Only(fire.Create)
	verifyUpdateMatcher := fire.Only(fire.Update)
	getFieldsMatcher := fire.Except(fire.Delete | fire.CollectionAction | fire.ResourceAction)

	// prepare access tables
	genericAccess := map[fire.Operation]Access{
		fire.List:           List,
		fire.Find:           Find,
		fire.Create:         Create,
		fire.Update:         Update,
		fire.Delete:         Delete,
		fire.ResourceAction: Find,
	}
	readAccess := map[fire.Operation]Access{
		fire.List:   List,
		fire.Find:   Find,
		fire.Create: Find,
		fire.Update: Find,
	}
	writeAccess := map[fire.Operation]Access{
		fire.Create: Create,
		fire.Update: Update,
	}

	return fire.C("ash/Execute", fire.Authorizer, fire.All(), func(ctx *fire.Context) error {
		// get policy
		policy, _ := ctx.Data[PolicyDataKey].(*Policy)
		if policy == nil {
			return fire.ErrAccessDenied.Wrap()
		}

		// check access
		if policy.Access&genericAccess[ctx.Operation] == 0 {
			return fire.ErrAccessDenied.Wrap()
		}

		// apply filter if available
		if getFilterMatcher(ctx) && policy.GetFilter != nil {
			ctx.Filters = append(ctx.Filters, policy.GetFilter(ctx))
		}

		// verify action access
		if ctx.Operation.Action() {
			// get action
			action := ctx.JSONAPIRequest.CollectionAction
			if ctx.Operation == fire.ResourceAction {
				action = ctx.JSONAPIRequest.ResourceAction
			}

			// check action
			if !policy.Actions[action] {
				return fire.ErrAccessDenied.Wrap()
			}
		}

		// verify id if available
		if verifyIDMatcher(ctx) && policy.VerifyID != nil {
			// get access
			access := policy.VerifyID(ctx, ctx.Selector["_id"].(coal.ID))

			// check access
			if access&genericAccess[ctx.Operation] == 0 {
				return fire.ErrAccessDenied.Wrap()
			}
		}

		// verify model if available
		if verifyModelMatcher(ctx) && policy.VerifyModel != nil {
			ctx.Defer(fire.C("ash/Execute-VerifyModel", fire.Verifier, verifyModelMatcher, func(ctx *fire.Context) error {
				// get required access
				reqAccess := genericAccess[ctx.Operation]

				// check access
				if ctx.Operation == fire.List {
					for _, model := range ctx.Models {
						if policy.VerifyModel(ctx, model)&reqAccess == 0 {
							return fire.ErrAccessDenied.Wrap()
						}
					}
				} else {
					if policy.VerifyModel(ctx, ctx.Model)&reqAccess == 0 {
						return fire.ErrAccessDenied.Wrap()
					}
				}

				return nil
			}))
		}

		// verify create if available
		if verifyCreateMatcher(ctx) && policy.VerifyCreate != nil {
			ctx.Defer(fire.C("ash/Execute-VerifyCreate", fire.Validator, verifyCreateMatcher, func(ctx *fire.Context) error {
				// check access
				if !policy.VerifyCreate(ctx, ctx.Model) {
					return fire.ErrAccessDenied.Wrap()
				}

				return nil
			}))
		}

		// verify update if available
		if verifyUpdateMatcher(ctx) && policy.VerifyUpdate != nil {
			ctx.Defer(fire.C("ash/Execute-VerifyUpdate", fire.Validator, verifyUpdateMatcher, func(ctx *fire.Context) error {
				// check access
				if !policy.VerifyUpdate(ctx, ctx.Model) {
					return fire.ErrAccessDenied.Wrap()
				}

				return nil
			}))
		}

		// collect fields
		readableFields := policy.Fields.Collect(readAccess[ctx.Operation])
		writableFields := policy.Fields.Collect(writeAccess[ctx.Operation])

		// set intersections of fields
		ctx.ReadableFields = stick.Intersect(ctx.ReadableFields, readableFields)
		ctx.WritableFields = stick.Intersect(ctx.WritableFields, writableFields)

		// set fields getters if available
		if getFieldsMatcher(ctx) && policy.GetFields != nil {
			ctx.GetReadableFields = func(model coal.Model) []string {
				if model == nil {
					return readableFields
				}
				return policy.GetFields(ctx, model).Collect(readAccess[ctx.Operation])
			}
			ctx.GetWritableFields = func(model coal.Model) []string {
				if ctx.Operation == fire.Create {
					return writableFields
				}
				return policy.GetFields(ctx, model).Collect(writeAccess[ctx.Operation])
			}
		}

		// set properties getter if available
		if getFieldsMatcher(ctx) && policy.GetProperties != nil {
			ctx.GetReadableProperties = func(model coal.Model) []string {
				return policy.GetProperties(ctx, model).Collect(readAccess[ctx.Operation])
			}
		}

		return nil
	})
}
<commit_msg>added todo<commit_after>package ash

import (
	"reflect"

	"github.com/256dpi/xo"
	"go.mongodb.org/mongo-driver/bson"

	"github.com/256dpi/fire"
	"github.com/256dpi/fire/coal"
	"github.com/256dpi/fire/stick"
)

// TODO: Field access authorization should be clearer, especially the static
//  fields used for List and Create operations.

// PolicyDataKey is the key used to store policies.
const PolicyDataKey = "ash:policy"

// Policy defines an authorization policy.
type Policy struct {
	// Access defines the general access.
	Access Access

	// Actions defines the allowed actions.
	Actions map[string]bool

	// The default fields used to determine the field access level. If the
	// getter is set, these will only be used to establish valid filters and
	// sorters during the fire.List operation authorizer stage, as well as the
	// writable fields during the fire.Create operation, otherwise the model
	// specific fields are used instead.
	Fields AccessTable

	// GetFilter is called to obtain the general resource access filter. This
	// filter is used to narrow down accessible resources in all operations
	// except fire.Create and fire.CollectionAction operations.
	GetFilter func(ctx *fire.Context) bson.M

	// VerifyID is called to for every direct model lookup to verify resource
	// level access. This function is called for all operations except fire.List,
	// fire.Create and fire.CollectionAction.
	VerifyID func(ctx *fire.Context, id coal.ID) Access

	// VerifyModel is called for every model load from the database to determine
	// resource level access. This function is called for all operations except
	// fire.Create and fire.CollectionAction.
	//
	// Note: The verification is deferred to the fire.Verifier stage.
	VerifyModel func(ctx *fire.Context, model coal.Model) Access

	// VerifyCreate and VerifyUpdate determine resource level access after all
	// modification have been applied. This function is called for the
	// fire.Create and fire.Update operation.
	//
	// Note: The verification is deferred to the fire.Validator stage.
	VerifyCreate func(ctx *fire.Context, model coal.Model) bool
	VerifyUpdate func(ctx *fire.Context, model coal.Model) bool

	// GetFields is called for every model to determine the field level access.
	// The policy should refrain from creating a new map for every request and
	// instead pre-allocate possible combinations and return those. The function
	// is called for all operations except fire.Delete, fire.CollectionAction and
	// fire.ResourceAction.
	GetFields func(ctx *fire.Context, model coal.Model) AccessTable

	// GetProperties is called for every model to determine the property level
	// access. The policy should refrain from creating a map for every request
	// and instead pre-allocate possible combinations and return those. The
	// function is called for all operations except fire.Delete,
	// fire.CollectionAction and fire.ResourceAction.
	GetProperties func(ctx *fire.Context, model coal.Model) AccessTable
}

// Selector is the function run to select a policy.
type Selector func(Identity) *Policy

// Select will run the provided function to select a policy for the supplied
// identity.
func Select(selector Selector) *fire.Callback {
	return fire.C("ash/Selector", fire.Authorizer, fire.All(), func(ctx *fire.Context) error {
		// get identity
		identity, _ := ctx.Data[IdentityDataKey]
		if identity == nil {
			return fire.ErrAccessDenied.Wrap()
		}

		// run selector
		policy := selector(identity)

		// check policy
		if policy == nil {
			return nil
		}

		// check stored
		if ctx.Data[PolicyDataKey] != nil {
			return xo.F("existing policy")
		}

		// store policy
		ctx.Data[PolicyDataKey] = policy

		return nil
	})
}

// SelectMatch will match the provided identity and on success use the provided
// factory to create a policy.
func SelectMatch(identity Identity, policy func(Identity) *Policy) *fire.Callback {
	// get type
	typ := reflect.TypeOf(identity)

	return Select(func(identity Identity) *Policy {
		// check type
		if typ != reflect.TypeOf(identity) {
			return nil
		}

		return policy(identity)
	})
}

// SelectPublic will math the public identity and use the provided factory to
// create a policy.
func SelectPublic(fn func() *Policy) *fire.Callback {
	return SelectMatch(&PublicIdentity{}, func(Identity) *Policy {
		return fn()
	})
}

// Execute will execute the selected policy or deny access.
func Execute() *fire.Callback {
	// prepare matchers
	getFilterMatcher := fire.Except(fire.Create | fire.CollectionAction)
	verifyIDMatcher := fire.Except(fire.List | fire.Create | fire.CollectionAction)
	verifyModelMatcher := fire.Except(fire.Create | fire.CollectionAction)
	verifyCreateMatcher := fire.Only(fire.Create)
	verifyUpdateMatcher := fire.Only(fire.Update)
	getFieldsMatcher := fire.Except(fire.Delete | fire.CollectionAction | fire.ResourceAction)

	// prepare access tables
	genericAccess := map[fire.Operation]Access{
		fire.List:           List,
		fire.Find:           Find,
		fire.Create:         Create,
		fire.Update:         Update,
		fire.Delete:         Delete,
		fire.ResourceAction: Find,
	}
	readAccess := map[fire.Operation]Access{
		fire.List:   List,
		fire.Find:   Find,
		fire.Create: Find,
		fire.Update: Find,
	}
	writeAccess := map[fire.Operation]Access{
		fire.Create: Create,
		fire.Update: Update,
	}

	return fire.C("ash/Execute", fire.Authorizer, fire.All(), func(ctx *fire.Context) error {
		// get policy
		policy, _ := ctx.Data[PolicyDataKey].(*Policy)
		if policy == nil {
			return fire.ErrAccessDenied.Wrap()
		}

		// check access
		if policy.Access&genericAccess[ctx.Operation] == 0 {
			return fire.ErrAccessDenied.Wrap()
		}

		// apply filter if available
		if getFilterMatcher(ctx) && policy.GetFilter != nil {
			ctx.Filters = append(ctx.Filters, policy.GetFilter(ctx))
		}

		// verify action access
		if ctx.Operation.Action() {
			// get action
			action := ctx.JSONAPIRequest.CollectionAction
			if ctx.Operation == fire.ResourceAction {
				action = ctx.JSONAPIRequest.ResourceAction
			}

			// check action
			if !policy.Actions[action] {
				return fire.ErrAccessDenied.Wrap()
			}
		}

		// verify id if available
		if verifyIDMatcher(ctx) && policy.VerifyID != nil {
			// get access
			access := policy.VerifyID(ctx, ctx.Selector["_id"].(coal.ID))

			// check access
			if access&genericAccess[ctx.Operation] == 0 {
				return fire.ErrAccessDenied.Wrap()
			}
		}

		// verify model if available
		if verifyModelMatcher(ctx) && policy.VerifyModel != nil {
			ctx.Defer(fire.C("ash/Execute-VerifyModel", fire.Verifier, verifyModelMatcher, func(ctx *fire.Context) error {
				// get required access
				reqAccess := genericAccess[ctx.Operation]

				// check access
				if ctx.Operation == fire.List {
					for _, model := range ctx.Models {
						if policy.VerifyModel(ctx, model)&reqAccess == 0 {
							return fire.ErrAccessDenied.Wrap()
						}
					}
				} else {
					if policy.VerifyModel(ctx, ctx.Model)&reqAccess == 0 {
						return fire.ErrAccessDenied.Wrap()
					}
				}

				return nil
			}))
		}

		// verify create if available
		if verifyCreateMatcher(ctx) && policy.VerifyCreate != nil {
			ctx.Defer(fire.C("ash/Execute-VerifyCreate", fire.Validator, verifyCreateMatcher, func(ctx *fire.Context) error {
				// check access
				if !policy.VerifyCreate(ctx, ctx.Model) {
					return fire.ErrAccessDenied.Wrap()
				}

				return nil
			}))
		}

		// verify update if available
		if verifyUpdateMatcher(ctx) && policy.VerifyUpdate != nil {
			ctx.Defer(fire.C("ash/Execute-VerifyUpdate", fire.Validator, verifyUpdateMatcher, func(ctx *fire.Context) error {
				// check access
				if !policy.VerifyUpdate(ctx, ctx.Model) {
					return fire.ErrAccessDenied.Wrap()
				}

				return nil
			}))
		}

		// collect fields
		readableFields := policy.Fields.Collect(readAccess[ctx.Operation])
		writableFields := policy.Fields.Collect(writeAccess[ctx.Operation])

		// set intersections of fields
		ctx.ReadableFields = stick.Intersect(ctx.ReadableFields, readableFields)
		ctx.WritableFields = stick.Intersect(ctx.WritableFields, writableFields)

		// set fields getters if available
		if getFieldsMatcher(ctx) && policy.GetFields != nil {
			ctx.GetReadableFields = func(model coal.Model) []string {
				if model == nil {
					return readableFields
				}
				return policy.GetFields(ctx, model).Collect(readAccess[ctx.Operation])
			}
			ctx.GetWritableFields = func(model coal.Model) []string {
				if ctx.Operation == fire.Create {
					return writableFields
				}
				return policy.GetFields(ctx, model).Collect(writeAccess[ctx.Operation])
			}
		}

		// set properties getter if available
		if getFieldsMatcher(ctx) && policy.GetProperties != nil {
			ctx.GetReadableProperties = func(model coal.Model) []string {
				return policy.GetProperties(ctx, model).Collect(readAccess[ctx.Operation])
			}
		}

		return nil
	})
}
<|endoftext|>


<commit_before>package httptesting

import (
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/buger/jsonparser"
	"github.com/golib/assert"
)

// AssertOK tests that the response status code is 200.
func (test *Client) AssertOK() {
	test.AssertStatus(http.StatusOK)
}

// AssertNotFound tests that the response status code is 404.
func (test *Client) AssertNotFound() {
	test.AssertStatus(http.StatusNotFound)
}

// AssertStatus tests that the response status code is equal with the given.
func (test *Client) AssertStatus(status int) {
	assert.EqualValues(test.t, status, test.Response.StatusCode, "Expected response status code "+strconv.Itoa(status)+", but got "+test.Response.Status+".")
}

// AssertContentType tests that the response includes Content-Type header with the given value.
func (test *Client) AssertContentType(contentType string) {
	test.AssertHeader("Content-Type", contentType)
}

// AssertHeader tests that the response includes named header with the given value.
func (test *Client) AssertHeader(name, value string) {
	name = http.CanonicalHeaderKey(name)
	actual := test.Response.Header.Get(name)

	assert.EqualValues(test.t, value, actual, "Expected response header "+name+" with "+value+", but got "+actual+".")
}

// AssertExistHeader tests that the response includes named header.
func (test *Client) AssertExistHeader(name string) {
	name = http.CanonicalHeaderKey(name)

	_, ok := test.Response.Header[name]
	if !ok {
		assert.Fail(test.t, "Response header: "+name+" (*required)", "Expected response header includes "+name+".")
	}
}

// AssertNotExistHeader tests that the response does not include named header.
func (test *Client) AssertNotExistHeader(name string) {
	name = http.CanonicalHeaderKey(name)

	_, ok := test.Response.Header[name]
	if ok {
		assert.Fail(test.t, "Response header: "+name+" (*not required)", "Expected response header does not include "+name+".")
	}
}

// AssertEmpty tests that the response is empty.
func (test *Client) AssertEmpty() {
	assert.Empty(test.t, string(test.ResponseBody))
}

// AssertNotEmpty tests that the response is not empty.
func (test *Client) AssertNotEmpty() {
	assert.NotEmpty(test.t, string(test.ResponseBody))
}

// AssertContains tests that the response contains the given string.
func (test *Client) AssertContains(s string) {
	assert.Contains(test.t, string(test.ResponseBody), s, "Expected response body contains "+s+".")
}

// AssertNotContains tests that the response does not contain the given string.
func (test *Client) AssertNotContains(s string) {
	assert.NotContains(test.t, string(test.ResponseBody), s, "Expected response body does not contain "+s+".")
}

// AssertMatch tests that the response matches the given regular expression.
func (test *Client) AssertMatch(re string) {
	r := regexp.MustCompile(re)

	if !r.Match(test.ResponseBody) {
		test.t.Errorf("Expected response body to match regexp %s", re)
	}
}

// AssertNotMatch tests that the response does not match the given regular expression.
func (test *Client) AssertNotMatch(re string) {
	r := regexp.MustCompile(re)

	if r.Match(test.ResponseBody) {
		test.t.Errorf("Expected response body does not match regexp %s", re)
	}
}

func (test *Client) AssertContainsJSON(key string, value interface{}) {
	var (
		buf = test.ResponseBody
		err error
	)

	keys := strings.Split(key, ".")
	for _, yek := range keys {
		// is the yek a array subscript?
		n, e := strconv.ParseInt(yek, 10, 32)
		if e != nil {
			buf, _, _, err = jsonparser.Get(buf, yek)
		} else {
			var i int64 = 0
			_, err = jsonparser.ArrayEach(buf, func(arrBuf []byte, arrType jsonparser.ValueType, arrOffset int, arrErr error) {
				if i == n {
					buf = arrBuf
					err = arrErr
				}

				i += 1
			})
		}

		if err != nil {
			test.t.Errorf("Expected response body contains json key %s with %s, but got Errr(%v)", key, value, err)
			break
		}
	}

	actual := string(buf)
	switch value.(type) {
	case []byte:
		expected := string(value.([]byte))
		assert.EqualValues(test.t, expected, actual, "Expected response body contains json key "+key+" with "+expected+", but got "+actual+".")

	case string:
		expected := value.(string)
		assert.EqualValues(test.t, expected, actual, "Expected response body contains json key "+key+" with "+expected+", but got "+actual+".")

	case int8:
		expected := int(value.(int8))
		actualInt, _ := strconv.Atoi(actual)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.Itoa(expected)+", but got "+actual+".")

	case int:
		expected := value.(int)
		actualInt, _ := strconv.Atoi(actual)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.Itoa(expected)+", but got "+actual+".")

	case int16:
		expected := int64(value.(int16))
		actualInt, _ := strconv.ParseInt(actual, 10, 16)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatInt(expected, 10)+", but got "+actual+".")

	case int32:
		expected := int64(value.(int32))
		actualInt, _ := strconv.ParseInt(actual, 10, 32)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatInt(expected, 10)+", but got "+actual+".")

	case int64:
		expected := value.(int64)
		actualInt, _ := strconv.ParseInt(actual, 10, 64)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatInt(expected, 10)+", but got "+actual+".")

	case float32:
		expected := float64(value.(float32))
		actualInt, _ := strconv.ParseFloat(actual, 32)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatFloat(expected, 'f', 5, 32)+", but got "+actual+".")

	case float64:
		expected := value.(float64)
		actualInt, _ := strconv.ParseFloat(actual, 64)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatFloat(expected, 'f', 5, 64)+", but got "+actual+".")

	case bool:
		expected := value.(bool)
		switch actual {
		case "true", "True", "1", "on":
			assert.True(test.t, expected, "Expected response body contains json key "+key+" with [true|True|1|on], but got "+actual+".")

		default:
			assert.False(test.t, expected, "Expected response body contains json key "+key+" with [false|False|0|off], but got "+actual+".")
		}
	}
}

func (test *Client) AssertNotContainsJSON(key string) {
	var (
		buf = test.ResponseBody
		err error
	)

	keys := strings.Split(key, ".")
	for _, yek := range keys {
		// is the yek a array subscript?
		n, e := strconv.ParseInt(yek, 10, 32)
		if e != nil {
			buf, _, _, err = jsonparser.Get(buf, yek)
		} else {
			var i int64 = 0
			_, err = jsonparser.ArrayEach(buf, func(arrBuf []byte, arrType jsonparser.ValueType, arrOffset int, arrErr error) {
				if i == n {
					buf = arrBuf
					err = arrErr
				}

				i += 1
			})
		}
	}

	if err == nil {
		test.t.Errorf("Expected response body does not contain json key %s, but got %s", key, string(buf))
	}
}
<commit_msg>What: refactor json parser logic<commit_after>package httptesting

import (
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/buger/jsonparser"
	"github.com/golib/assert"
)

// AssertOK tests that the response status code is 200.
func (test *Client) AssertOK() {
	test.AssertStatus(http.StatusOK)
}

// AssertNotFound tests that the response status code is 404.
func (test *Client) AssertNotFound() {
	test.AssertStatus(http.StatusNotFound)
}

// AssertStatus tests that the response status code is equal with the given.
func (test *Client) AssertStatus(status int) {
	assert.EqualValues(test.t, status, test.Response.StatusCode, "Expected response status code "+strconv.Itoa(status)+", but got "+test.Response.Status+".")
}

// AssertContentType tests that the response includes Content-Type header with the given value.
func (test *Client) AssertContentType(contentType string) {
	test.AssertHeader("Content-Type", contentType)
}

// AssertHeader tests that the response includes named header with the given value.
func (test *Client) AssertHeader(name, value string) {
	name = http.CanonicalHeaderKey(name)
	actual := test.Response.Header.Get(name)

	assert.EqualValues(test.t, value, actual, "Expected response header "+name+" with "+value+", but got "+actual+".")
}

// AssertExistHeader tests that the response includes named header.
func (test *Client) AssertExistHeader(name string) {
	name = http.CanonicalHeaderKey(name)

	_, ok := test.Response.Header[name]
	if !ok {
		assert.Fail(test.t, "Response header: "+name+" (*required)", "Expected response header includes "+name+".")
	}
}

// AssertNotExistHeader tests that the response does not include named header.
func (test *Client) AssertNotExistHeader(name string) {
	name = http.CanonicalHeaderKey(name)

	_, ok := test.Response.Header[name]
	if ok {
		assert.Fail(test.t, "Response header: "+name+" (*not required)", "Expected response header does not include "+name+".")
	}
}

// AssertEmpty tests that the response is empty.
func (test *Client) AssertEmpty() {
	assert.Empty(test.t, string(test.ResponseBody))
}

// AssertNotEmpty tests that the response is not empty.
func (test *Client) AssertNotEmpty() {
	assert.NotEmpty(test.t, string(test.ResponseBody))
}

// AssertContains tests that the response contains the given string.
func (test *Client) AssertContains(s string) {
	assert.Contains(test.t, string(test.ResponseBody), s, "Expected response body contains "+s+".")
}

// AssertNotContains tests that the response does not contain the given string.
func (test *Client) AssertNotContains(s string) {
	assert.NotContains(test.t, string(test.ResponseBody), s, "Expected response body does not contain "+s+".")
}

// AssertMatch tests that the response matches the given regular expression.
func (test *Client) AssertMatch(re string) {
	r := regexp.MustCompile(re)

	if !r.Match(test.ResponseBody) {
		test.t.Errorf("Expected response body to match regexp %s", re)
	}
}

// AssertNotMatch tests that the response does not match the given regular expression.
func (test *Client) AssertNotMatch(re string) {
	r := regexp.MustCompile(re)

	if r.Match(test.ResponseBody) {
		test.t.Errorf("Expected response body does not match regexp %s", re)
	}
}

func (test *Client) AssertContainsJSON(key string, value interface{}) {
	var (
		buf  = test.ResponseBody
		data []byte
		err  error
	)

	for _, yek := range strings.Split(key, ".") {
		data, _, _, err = jsonparser.Get(buf, yek)
		if err == nil {
			buf = data

			continue
		}

		// is the yek an array subscript?
		n, e := strconv.ParseInt(yek, 10, 32)
		if e != nil {
			break
		}

		var i int64 = 0
		jsonparser.ArrayEach(buf, func(arrBuf []byte, arrType jsonparser.ValueType, arrOffset int, arrErr error) {
			if i == n {
				buf = arrBuf
				err = arrErr
			}

			i++
		})
		if err != nil {
			break
		}
	}
	if err != nil {
		test.t.Errorf("Expected response body contains json key %s with %s, but got Errr(%v)", key, value, err)
	}

	actual := string(buf)
	switch value.(type) {
	case []byte:
		expected := string(value.([]byte))
		assert.EqualValues(test.t, expected, actual, "Expected response body contains json key "+key+" with "+expected+", but got "+actual+".")

	case string:
		expected := value.(string)
		assert.EqualValues(test.t, expected, actual, "Expected response body contains json key "+key+" with "+expected+", but got "+actual+".")

	case int8:
		expected := int(value.(int8))
		actualInt, _ := strconv.Atoi(actual)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.Itoa(expected)+", but got "+actual+".")

	case int:
		expected := value.(int)
		actualInt, _ := strconv.Atoi(actual)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.Itoa(expected)+", but got "+actual+".")

	case int16:
		expected := int64(value.(int16))
		actualInt, _ := strconv.ParseInt(actual, 10, 16)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatInt(expected, 10)+", but got "+actual+".")

	case int32:
		expected := int64(value.(int32))
		actualInt, _ := strconv.ParseInt(actual, 10, 32)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatInt(expected, 10)+", but got "+actual+".")

	case int64:
		expected := value.(int64)
		actualInt, _ := strconv.ParseInt(actual, 10, 64)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatInt(expected, 10)+", but got "+actual+".")

	case float32:
		expected := float64(value.(float32))
		actualInt, _ := strconv.ParseFloat(actual, 32)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatFloat(expected, 'f', 5, 32)+", but got "+actual+".")

	case float64:
		expected := value.(float64)
		actualInt, _ := strconv.ParseFloat(actual, 64)
		assert.EqualValues(test.t, expected, actualInt, "Expected response body contains json key "+key+" with "+strconv.FormatFloat(expected, 'f', 5, 64)+", but got "+actual+".")

	case bool:
		expected := value.(bool)
		switch actual {
		case "true", "True", "1", "on":
			assert.True(test.t, expected, "Expected response body contains json key "+key+" with [true|True|1|on], but got "+actual+".")

		default:
			assert.False(test.t, expected, "Expected response body contains json key "+key+" with [false|False|0|off], but got "+actual+".")
		}
	}
}

func (test *Client) AssertNotContainsJSON(key string) {
	var (
		buf  = test.ResponseBody
		data []byte
		err  error
	)

	for _, yek := range strings.Split(key, ".") {
		data, _, _, err = jsonparser.Get(buf, yek)
		if err == nil {
			buf = data

			continue
		}

		// is the yek an array subscript?
		n, e := strconv.ParseInt(yek, 10, 32)
		if e != nil {
			break
		}

		var i int64 = 0
		jsonparser.ArrayEach(buf, func(arrBuf []byte, arrType jsonparser.ValueType, arrOffset int, arrErr error) {
			if i == n {
				buf = arrBuf
				err = arrErr
			}

			i++
		})
		if err != nil {
			break
		}
	}

	if err == nil {
		test.t.Errorf("Expected response body does not contain json key %s, but got %s", key, string(buf))
	}
}
<|endoftext|>


<commit_before>package indexer

import (
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"sync"

	"strings"

	"github.com/deepfabric/pilosa"
	"github.com/pkg/errors"
)

// Frame represents a string field of an index. Refers to pilosa.Frame and pilosa.View.
type Frame struct {
	mu       sync.Mutex
	path     string
	index    string
	name     string
	maxSlice uint64

	fragments map[uint64]*pilosa.Fragment //map slice to Fragment
	td        *TermDict
}

// NewFrame returns a new instance of frame, and initializes it.
func NewFrame(path, index, name string) (f *Frame, err error) {
	var td *TermDict
	if td, err = NewTermDict(path); err != nil {
		return
	}
	f = &Frame{
		path:      path,
		index:     index,
		name:      name,
		td:        td,
		fragments: make(map[uint64]*pilosa.Fragment),
	}
	err = f.Open()
	return
}

//Open opens an existing frame
func (f *Frame) Open() (err error) {
	var sliceList []uint64
	if sliceList, err = getSliceList(f.path); err != nil {
		return
	}
	for _, slice := range sliceList {
		fp := f.FragmentPath(slice)
		fragment := pilosa.NewFragment(fp, f.index, f.name, pilosa.ViewStandard, slice)
		if err = fragment.Open(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
		f.fragments[slice] = fragment
		if f.maxSlice < slice {
			f.maxSlice = slice
		}
	}
	return
}

func getSliceList(dir string) (numList []uint64, err error) {
	var d *os.File
	var fns []string
	var num uint64
	fragDir := filepath.Join(dir, "fragments")
	if err = os.MkdirAll(fragDir, 0700); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	d, err = os.Open(fragDir)
	if err != nil {
		err = errors.Wrap(err, "")
		return
	}
	fns, err = d.Readdirnames(0)
	if err != nil {
		err = errors.Wrap(err, "")
		return
	}
	re := regexp.MustCompile("(?P<num>[0-9]+)")
	for _, fn := range fns {
		subs := re.FindStringSubmatch(fn)
		if subs == nil {
			continue
		}
		num, err = strconv.ParseUint(subs[1], 10, 64)
		if err != nil {
			err = errors.Wrap(err, "")
			return
		}
		numList = append(numList, num)
	}
	return
}

// Close closes all fragments without removing files on disk.
func (f *Frame) Close() (err error) {
	for _, fragment := range f.fragments {
		if err = fragment.Close(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
	}
	err = f.td.Close()
	return
}

// Destroy closes all fragments, removes all files on disk.
func (f *Frame) Destroy() (err error) {
	for _, fragment := range f.fragments {
		if err = fragment.Close(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
	}
	if err = os.RemoveAll(filepath.Join(f.path, "fragments")); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	err = f.td.Destroy()
	return
}

// FragmentPath returns the path to a fragment
func (f *Frame) FragmentPath(slice uint64) string {
	return filepath.Join(f.path, "fragments", strconv.FormatUint(slice, 10))
}

// Fragment returns a fragment in the view by slice.
func (f *Frame) Fragment(slice uint64) *pilosa.Fragment {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.fragments[slice]
}

// Name returns the name the frame was initialized with.
func (f *Frame) Name() string { return f.name }

// Index returns the index name the frame was initialized with.
func (f *Frame) Index() string { return f.index }

// Path returns the path the frame was initialized with.
func (f *Frame) Path() string { return f.path }

// MaxSlice returns the max slice in the frame.
func (f *Frame) MaxSlice() uint64 { return f.maxSlice }

// SetBit sets a bit within the frame, and expands fragments if necessary.
func (f *Frame) SetBit(rowID, colID uint64) (changed bool, err error) {
	slice := colID / pilosa.SliceWidth
	fragment, ok := f.fragments[slice]
	if !ok {
		fp := f.FragmentPath(slice)
		fragment = pilosa.NewFragment(fp, f.index, f.name, pilosa.ViewStandard, slice)
		if err = fragment.Open(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
		f.fragments[slice] = fragment
	}
	changed, err = fragment.SetBit(rowID, colID)
	if f.maxSlice < slice {
		f.maxSlice = slice
	}
	return
}

// ClearBit clears a bit within the frame.
func (f *Frame) ClearBit(rowID, colID uint64) (changed bool, err error) {
	slice := colID / pilosa.SliceWidth
	fragment, ok := f.fragments[slice]
	if !ok {
		err = errors.New("column out of bounds")
		return
	}
	changed, err = fragment.ClearBit(rowID, colID)
	return
}

// ParseAndIndex parses and index a field.
func (f *Frame) ParseAndIndex(docID uint64, text string) (err error) {
	terms := strings.SplitN(text, " ", -1)
	ids, err := f.td.CreateTermsIfNotExist(terms)
	if err != nil {
		return
	}
	for _, termID := range ids {
		if _, err = f.SetBit(termID, docID); err != nil {
			return
		}
	}
	return
}

func (f *Frame) Query(term string) (bm *pilosa.Bitmap, err error) {
	bm = pilosa.NewBitmap()
	termID, found := f.td.GetTermID(term)
	if !found {
		return
	}
	for _, fragment := range f.fragments {
		bm2 := fragment.Row(termID)
		bm.Merge(bm2)
	}
	return
}
<commit_msg>fixed getSliceList<commit_after>package indexer

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"sync"

	"strings"

	"github.com/deepfabric/pilosa"
	"github.com/pkg/errors"
)

// Frame represents a string field of an index. Refers to pilosa.Frame and pilosa.View.
type Frame struct {
	mu       sync.Mutex
	path     string
	index    string
	name     string
	maxSlice uint64

	fragments map[uint64]*pilosa.Fragment //map slice to Fragment
	td        *TermDict
}

// NewFrame returns a new instance of frame, and initializes it.
func NewFrame(path, index, name string) (f *Frame, err error) {
	var td *TermDict
	if td, err = NewTermDict(path); err != nil {
		return
	}
	f = &Frame{
		path:      path,
		index:     index,
		name:      name,
		td:        td,
		fragments: make(map[uint64]*pilosa.Fragment),
	}
	err = f.Open()
	return
}

//Open opens an existing frame
func (f *Frame) Open() (err error) {
	var sliceList []uint64
	if sliceList, err = getSliceList(f.path); err != nil {
		return
	}
	fmt.Printf("sliceList: %v\n", sliceList)
	for _, slice := range sliceList {
		fp := f.FragmentPath(slice)
		fragment := pilosa.NewFragment(fp, f.index, f.name, pilosa.ViewStandard, slice)
		if err = fragment.Open(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
		f.fragments[slice] = fragment
		if f.maxSlice < slice {
			f.maxSlice = slice
		}
	}
	return
}

func getSliceList(dir string) (numList []uint64, err error) {
	var d *os.File
	var fns []string
	var num uint64
	fragDir := filepath.Join(dir, "fragments")
	if err = os.MkdirAll(fragDir, 0700); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	d, err = os.Open(fragDir)
	if err != nil {
		err = errors.Wrap(err, "")
		return
	}
	fns, err = d.Readdirnames(0)
	if err != nil {
		err = errors.Wrap(err, "")
		return
	}
	re := regexp.MustCompile("^(?P<num>[0-9]+)$")
	for _, fn := range fns {
		subs := re.FindStringSubmatch(fn)
		if subs == nil {
			continue
		}
		num, err = strconv.ParseUint(subs[1], 10, 64)
		if err != nil {
			err = errors.Wrap(err, "")
			return
		}
		numList = append(numList, num)
	}
	return
}

// Close closes all fragments without removing files on disk.
func (f *Frame) Close() (err error) {
	for _, fragment := range f.fragments {
		if err = fragment.Close(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
	}
	err = f.td.Close()
	return
}

// Destroy closes all fragments, removes all files on disk.
func (f *Frame) Destroy() (err error) {
	for _, fragment := range f.fragments {
		if err = fragment.Close(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
	}
	if err = os.RemoveAll(filepath.Join(f.path, "fragments")); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	err = f.td.Destroy()
	return
}

// FragmentPath returns the path to a fragment
func (f *Frame) FragmentPath(slice uint64) string {
	return filepath.Join(f.path, "fragments", strconv.FormatUint(slice, 10))
}

// Fragment returns a fragment in the view by slice.
func (f *Frame) Fragment(slice uint64) *pilosa.Fragment {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.fragments[slice]
}

// Name returns the name the frame was initialized with.
func (f *Frame) Name() string { return f.name }

// Index returns the index name the frame was initialized with.
func (f *Frame) Index() string { return f.index }

// Path returns the path the frame was initialized with.
func (f *Frame) Path() string { return f.path }

// MaxSlice returns the max slice in the frame.
func (f *Frame) MaxSlice() uint64 { return f.maxSlice }

// SetBit sets a bit within the frame, and expands fragments if necessary.
func (f *Frame) SetBit(rowID, colID uint64) (changed bool, err error) {
	slice := colID / pilosa.SliceWidth
	fragment, ok := f.fragments[slice]
	if !ok {
		fp := f.FragmentPath(slice)
		fragment = pilosa.NewFragment(fp, f.index, f.name, pilosa.ViewStandard, slice)
		if err = fragment.Open(); err != nil {
			err = errors.Wrap(err, "")
			return
		}
		f.fragments[slice] = fragment
	}
	changed, err = fragment.SetBit(rowID, colID)
	if f.maxSlice < slice {
		f.maxSlice = slice
	}
	return
}

// ClearBit clears a bit within the frame.
func (f *Frame) ClearBit(rowID, colID uint64) (changed bool, err error) {
	slice := colID / pilosa.SliceWidth
	fragment, ok := f.fragments[slice]
	if !ok {
		err = errors.New("column out of bounds")
		return
	}
	changed, err = fragment.ClearBit(rowID, colID)
	return
}

// ParseAndIndex parses and index a field.
func (f *Frame) ParseAndIndex(docID uint64, text string) (err error) {
	terms := strings.SplitN(text, " ", -1)
	ids, err := f.td.CreateTermsIfNotExist(terms)
	if err != nil {
		return
	}
	for _, termID := range ids {
		if _, err = f.SetBit(termID, docID); err != nil {
			return
		}
	}
	return
}

func (f *Frame) Query(term string) (bm *pilosa.Bitmap, err error) {
	bm = pilosa.NewBitmap()
	termID, found := f.td.GetTermID(term)
	if !found {
		return
	}
	for _, fragment := range f.fragments {
		bm2 := fragment.Row(termID)
		bm.Merge(bm2)
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 fs

import (
	"fmt"
	"reflect"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"github.com/jacobsa/gcsfuse/fs/inode"
	"github.com/jacobsa/gcsfuse/timeutil"
	"golang.org/x/net/context"
)

type fileSystem struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Dependencies
	/////////////////////////

	clock  timeutil.Clock
	bucket gcs.Bucket

	/////////////////////////
	// Mutable state
	/////////////////////////

	// When acquiring this lock, the caller must hold no inode or dir handle
	// locks.
	mu syncutil.InvariantMutex

	// The collection of live inodes, keyed by inode ID. No ID less than
	// fuse.RootInodeID is ever used.
	//
	// INVARIANT: All values are of type *inode.DirInode or *inode.FileInode
	// INVARIANT: For all keys k, k >= fuse.RootInodeID
	// INVARIANT: inodes[fuse.RootInodeID] is of type *inode.DirInode
	//
	// GUARDED_BY(mu)
	inodes map[fuse.InodeID]interface{}

	// The next inode ID to hand out. We assume that this will never overflow,
	// since even if we were handing out inode IDs at 4 GHz, it would still take
	// over a century to do so.
	//
	// INVARIANT: For all keys k in inodes, k < nextInodeID
	nextInodeID fuse.InodeID

	// The collection of live handles, keyed by handle ID.
	//
	// INVARIANT: All values are of type *dirHandle
	//
	// GUARDED_BY(mu)
	handles map[fuse.HandleID]interface{}

	// The next handle ID to hand out. We assume that this will never overflow.
	//
	// INVARIANT: For all keys k in handles, k < nextHandleID
	nextHandleID fuse.HandleID
}

// Create a fuse file system whose root directory is the root of the supplied
// bucket. The supplied clock will be used for cache invalidation, modification
// times, etc.
func NewFileSystem(
	clock timeutil.Clock,
	bucket gcs.Bucket) (ffs fuse.FileSystem, err error) {
	// Set up the basic struct.
	fs := &fileSystem{
		clock:       clock,
		bucket:      bucket,
		inodes:      make(map[fuse.InodeID]interface{}),
		nextInodeID: fuse.RootInodeID + 1,
	}

	// Set up the root inode.
	fs.inodes[fuse.RootInodeID] = inode.NewDirInode(bucket, "")

	// Set up invariant checking.
	fs.mu = syncutil.NewInvariantMutex(fs.checkInvariants)

	ffs = fs
	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func (fs *fileSystem) checkInvariants() {
	// Check inode keys.
	for id, _ := range fs.inodes {
		if id < fuse.RootInodeID || id >= fs.nextInodeID {
			panic(fmt.Sprintf("Illegal inode ID: %v", id))
		}
	}

	// Check the root inode.
	_ = fs.inodes[fuse.RootInodeID].(*inode.DirInode)

	// Check the type of each inode.
	for _, in := range fs.inodes {
		switch in.(type) {
		case *inode.DirInode:
		case *inode.FileInode:

		default:
			panic(fmt.Sprintf("Unexpected inode type: %v", reflect.TypeOf(in)))
		}
	}

	// Check handles.
	for id, h := range fs.handles {
		if id >= fs.nextHandleID {
			panic(fmt.Sprintf("Illegal handle ID: %v", id))
		}

		_ = h.(*dirHandle)
	}
}

// Find the given inode and return it with its lock held for reading. Panic if
// it doesn't exist or is the wrong type.
//
// SHARED_LOCKS_REQUIRED(fs.mu)
// SHARED_LOCK_FUNCTION(inode.mu)
func (fs *fileSystem) getDirForReadingOrDie(
	id fuse.InodeID) (in *inode.DirInode) {
	in = fs.inodes[id].(*inode.DirInode)
	in.Mu.RLock()
	return
}

////////////////////////////////////////////////////////////////////////
// fuse.FileSystem methods
////////////////////////////////////////////////////////////////////////

func (fs *fileSystem) Init(
	ctx context.Context,
	req *fuse.InitRequest) (resp *fuse.InitResponse, err error) {
	// Nothing interesting to do.
	resp = &fuse.InitResponse{}
	return
}

func (fs *fileSystem) OpenDir(
	ctx context.Context,
	req *fuse.OpenDirRequest) (resp *fuse.OpenDirResponse, err error) {
	resp = &fuse.OpenDirResponse{}

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Make sure the inode still exists and is a directory. If not, something has
	// screwed up because the VFS layer shouldn't have let us forget the inode
	// before opening it.
	in := fs.getDirForReadingOrDie(req.Inode)
	defer in.Mu.RUnlock()

	// Allocate a handle.
	handleID := fs.nextHandleID
	fs.nextHandleID++

	fs.handles[handleID] = newDirHandle(in)
	resp.Handle = handleID

	return
}

func (fs *fileSystem) ReleaseDirHandle(
	ctx context.Context,
	req *fuse.ReleaseDirHandleRequest) (
	resp *fuse.ReleaseDirHandleResponse, err error) {
	resp = &fuse.ReleaseDirHandleResponse{}

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Sanity check that this handle exists and is of the correct type.
	_ = fs.handles[req.Handle].(*dirHandle)

	// Clear the entry from the map.
	delete(fs.handles, req.Handle)

	return
}
<commit_msg>Fixed a panic.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 fs

import (
	"fmt"
	"reflect"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"github.com/jacobsa/gcsfuse/fs/inode"
	"github.com/jacobsa/gcsfuse/timeutil"
	"golang.org/x/net/context"
)

type fileSystem struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Dependencies
	/////////////////////////

	clock  timeutil.Clock
	bucket gcs.Bucket

	/////////////////////////
	// Mutable state
	/////////////////////////

	// When acquiring this lock, the caller must hold no inode or dir handle
	// locks.
	mu syncutil.InvariantMutex

	// The collection of live inodes, keyed by inode ID. No ID less than
	// fuse.RootInodeID is ever used.
	//
	// INVARIANT: All values are of type *inode.DirInode or *inode.FileInode
	// INVARIANT: For all keys k, k >= fuse.RootInodeID
	// INVARIANT: inodes[fuse.RootInodeID] is of type *inode.DirInode
	//
	// GUARDED_BY(mu)
	inodes map[fuse.InodeID]interface{}

	// The next inode ID to hand out. We assume that this will never overflow,
	// since even if we were handing out inode IDs at 4 GHz, it would still take
	// over a century to do so.
	//
	// INVARIANT: For all keys k in inodes, k < nextInodeID
	nextInodeID fuse.InodeID

	// The collection of live handles, keyed by handle ID.
	//
	// INVARIANT: All values are of type *dirHandle
	//
	// GUARDED_BY(mu)
	handles map[fuse.HandleID]interface{}

	// The next handle ID to hand out. We assume that this will never overflow.
	//
	// INVARIANT: For all keys k in handles, k < nextHandleID
	nextHandleID fuse.HandleID
}

// Create a fuse file system whose root directory is the root of the supplied
// bucket. The supplied clock will be used for cache invalidation, modification
// times, etc.
func NewFileSystem(
	clock timeutil.Clock,
	bucket gcs.Bucket) (ffs fuse.FileSystem, err error) {
	// Set up the basic struct.
	fs := &fileSystem{
		clock:       clock,
		bucket:      bucket,
		inodes:      make(map[fuse.InodeID]interface{}),
		nextInodeID: fuse.RootInodeID + 1,
		handles:     make(map[fuse.HandleID]interface{}),
	}

	// Set up the root inode.
	fs.inodes[fuse.RootInodeID] = inode.NewDirInode(bucket, "")

	// Set up invariant checking.
	fs.mu = syncutil.NewInvariantMutex(fs.checkInvariants)

	ffs = fs
	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func (fs *fileSystem) checkInvariants() {
	// Check inode keys.
	for id, _ := range fs.inodes {
		if id < fuse.RootInodeID || id >= fs.nextInodeID {
			panic(fmt.Sprintf("Illegal inode ID: %v", id))
		}
	}

	// Check the root inode.
	_ = fs.inodes[fuse.RootInodeID].(*inode.DirInode)

	// Check the type of each inode.
	for _, in := range fs.inodes {
		switch in.(type) {
		case *inode.DirInode:
		case *inode.FileInode:

		default:
			panic(fmt.Sprintf("Unexpected inode type: %v", reflect.TypeOf(in)))
		}
	}

	// Check handles.
	for id, h := range fs.handles {
		if id >= fs.nextHandleID {
			panic(fmt.Sprintf("Illegal handle ID: %v", id))
		}

		_ = h.(*dirHandle)
	}
}

// Find the given inode and return it with its lock held for reading. Panic if
// it doesn't exist or is the wrong type.
//
// SHARED_LOCKS_REQUIRED(fs.mu)
// SHARED_LOCK_FUNCTION(inode.mu)
func (fs *fileSystem) getDirForReadingOrDie(
	id fuse.InodeID) (in *inode.DirInode) {
	in = fs.inodes[id].(*inode.DirInode)
	in.Mu.RLock()
	return
}

////////////////////////////////////////////////////////////////////////
// fuse.FileSystem methods
////////////////////////////////////////////////////////////////////////

func (fs *fileSystem) Init(
	ctx context.Context,
	req *fuse.InitRequest) (resp *fuse.InitResponse, err error) {
	// Nothing interesting to do.
	resp = &fuse.InitResponse{}
	return
}

func (fs *fileSystem) OpenDir(
	ctx context.Context,
	req *fuse.OpenDirRequest) (resp *fuse.OpenDirResponse, err error) {
	resp = &fuse.OpenDirResponse{}

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Make sure the inode still exists and is a directory. If not, something has
	// screwed up because the VFS layer shouldn't have let us forget the inode
	// before opening it.
	in := fs.getDirForReadingOrDie(req.Inode)
	defer in.Mu.RUnlock()

	// Allocate a handle.
	handleID := fs.nextHandleID
	fs.nextHandleID++

	fs.handles[handleID] = newDirHandle(in)
	resp.Handle = handleID

	return
}

func (fs *fileSystem) ReleaseDirHandle(
	ctx context.Context,
	req *fuse.ReleaseDirHandleRequest) (
	resp *fuse.ReleaseDirHandleResponse, err error) {
	resp = &fuse.ReleaseDirHandleResponse{}

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Sanity check that this handle exists and is of the correct type.
	_ = fs.handles[req.Handle].(*dirHandle)

	// Clear the entry from the map.
	delete(fs.handles, req.Handle)

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2015 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 (
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"sync"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/intstr"
	"k8s.io/kubernetes/test/e2e/framework"

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

const (
	smallRCSize       = 5
	mediumRCSize      = 30
	bigRCSize         = 250
	smallRCGroupName  = "load-small-rc"
	mediumRCGroupName = "load-medium-rc"
	bigRCGroupName    = "load-big-rc"
	smallRCBatchSize  = 30
	mediumRCBatchSize = 5
	bigRCBatchSize    = 1
	// We start RCs/Services/pods/... in different namespace in this test.
	// nodeCountPerNamespace determines how many namespaces we will be using
	// depending on the number of nodes in the underlying cluster.
	nodeCountPerNamespace = 250
)

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = framework.KubeDescribe("Load capacity", func() {
	var c *client.Client
	var nodeCount int
	var ns string
	var configs []*framework.RCConfig
	var namespaces []*api.Namespace

	// Gathers metrics before teardown
	// TODO add flag that allows to skip cleanup on failure
	AfterEach(func() {
		// Verify latency metrics
		highLatencyRequests, err := framework.HighLatencyRequests(c)
		framework.ExpectNoError(err, "Too many instances metrics above the threshold")
		Expect(highLatencyRequests).NotTo(BeNumerically(">", 0))
	})

	// Explicitly put here, to delete namespace at the end of the test
	// (after measuring latency metrics, etc.).
	options := framework.FrameworkOptions{
		ClientQPS:   50,
		ClientBurst: 100,
	}
	f := framework.NewFramework("load", options, nil)
	f.NamespaceDeletionTimeout = time.Hour

	BeforeEach(func() {
		c = f.Client

		// In large clusters we may get to this point but still have a bunch
		// of nodes without Routes created. Since this would make a node
		// unschedulable, we need to wait until all of them are schedulable.
		framework.ExpectNoError(framework.WaitForAllNodesSchedulable(c))

		ns = f.Namespace.Name
		nodes := framework.GetReadySchedulableNodesOrDie(c)
		nodeCount = len(nodes.Items)
		Expect(nodeCount).NotTo(BeZero())

		// Terminating a namespace (deleting the remaining objects from it - which
		// generally means events) can affect the current run. Thus we wait for all
		// terminating namespace to be finally deleted before starting this test.
		err := framework.CheckTestingNSDeletedExcept(c, ns)
		framework.ExpectNoError(err)

		framework.ExpectNoError(framework.ResetMetrics(c))
	})

	type Load struct {
		podsPerNode int
		image       string
		command     []string
	}

	loadTests := []Load{
		// The container will consume 1 cpu and 512mb of memory.
		{podsPerNode: 3, image: "jess/stress", command: []string{"stress", "-c", "1", "-m", "2"}},
		{podsPerNode: 30, image: "gcr.io/google_containers/serve_hostname:v1.4"},
	}

	for _, testArg := range loadTests {
		name := fmt.Sprintf("should be able to handle %v pods per node", testArg.podsPerNode)
		if testArg.podsPerNode == 30 {
			name = "[Feature:Performance] " + name
		} else {
			name = "[Feature:ManualPerformance] " + name
		}
		itArg := testArg

		It(name, func() {
			// Create a number of namespaces.
			namespaces = createNamespaces(f, nodeCount, itArg.podsPerNode)

			totalPods := itArg.podsPerNode * nodeCount
			configs = generateRCConfigs(totalPods, itArg.image, itArg.command, c, namespaces)
			var services []*api.Service
			// Read the environment variable to see if we want to create services
			createServices := os.Getenv("CREATE_SERVICES")
			if createServices == "true" {
				framework.Logf("Creating services")
				services := generateServicesForConfigs(configs)
				for _, service := range services {
					_, err := c.Services(service.Namespace).Create(service)
					framework.ExpectNoError(err)
				}
			} else {
				framework.Logf("Skipping service creation")
			}

			// Simulate lifetime of RC:
			//  * create with initial size
			//  * scale RC to a random size and list all pods
			//  * scale RC to a random size and list all pods
			//  * delete it
			//
			// This will generate ~5 creations/deletions per second assuming:
			//  - X small RCs each 5 pods   [ 5 * X = totalPods / 2 ]
			//  - Y medium RCs each 30 pods [ 30 * Y = totalPods / 4 ]
			//  - Z big RCs each 250 pods   [ 250 * Z = totalPods / 4]

			// We would like to spread creating replication controllers over time
			// to make it possible to create/schedule them in the meantime.
			// Currently we assume 10 pods/second average throughput.
			// We may want to revisit it in the future.
			creatingTime := time.Duration(totalPods/10) * time.Second
			createAllRC(configs, creatingTime)
			By("============================================================================")

			// We would like to spread scaling replication controllers over time
			// to make it possible to create/schedule & delete them in the meantime.
			// Currently we assume that 10 pods/second average throughput.
			// The expected number of created/deleted pods is less than totalPods/3.
			scalingTime := time.Duration(totalPods/30) * time.Second
			scaleAllRC(configs, scalingTime)
			By("============================================================================")

			scaleAllRC(configs, scalingTime)
			By("============================================================================")

			// Cleanup all created replication controllers.
			// Currently we assume 10 pods/second average deletion throughput.
			// We may want to revisit it in the future.
			deletingTime := time.Duration(totalPods/10) * time.Second
			deleteAllRC(configs, deletingTime)
			if createServices == "true" {
				for _, service := range services {
					err := c.Services(ns).Delete(service.Name)
					framework.ExpectNoError(err)
				}
				framework.Logf("%v Services created.", len(services))
			}
		})
	}
})

func createNamespaces(f *framework.Framework, nodeCount, podsPerNode int) []*api.Namespace {
	namespaceCount := (nodeCount + nodeCountPerNamespace - 1) / nodeCountPerNamespace
	namespaces := []*api.Namespace{}
	for i := 1; i <= namespaceCount; i++ {
		namespace, err := f.CreateNamespace(fmt.Sprintf("load-%d-nodepods-%d", podsPerNode, i), nil)
		framework.ExpectNoError(err)
		namespaces = append(namespaces, namespace)
	}
	return namespaces
}

func computeRCCounts(total int) (int, int, int) {
	// Small RCs owns ~0.5 of total number of pods, medium and big RCs ~0.25 each.
	// For example for 3000 pods (100 nodes, 30 pods per node) there are:
	//  - 300 small RCs each 5 pods
	//  - 25 medium RCs each 30 pods
	//  - 3 big RCs each 250 pods
	bigRCCount := total / 4 / bigRCSize
	total -= bigRCCount * bigRCSize
	mediumRCCount := total / 3 / mediumRCSize
	total -= mediumRCCount * mediumRCSize
	smallRCCount := total / smallRCSize
	return smallRCCount, mediumRCCount, bigRCCount
}

func generateRCConfigs(totalPods int, image string, command []string, c *client.Client, nss []*api.Namespace) []*framework.RCConfig {
	configs := make([]*framework.RCConfig, 0)

	smallRCCount, mediumRCCount, bigRCCount := computeRCCounts(totalPods)
	configs = append(configs, generateRCConfigsForGroup(c, nss, smallRCGroupName, smallRCSize, smallRCCount, image, command)...)
	configs = append(configs, generateRCConfigsForGroup(c, nss, mediumRCGroupName, mediumRCSize, mediumRCCount, image, command)...)
	configs = append(configs, generateRCConfigsForGroup(c, nss, bigRCGroupName, bigRCSize, bigRCCount, image, command)...)

	return configs
}

func generateRCConfigsForGroup(c *client.Client, nss []*api.Namespace, groupName string, size, count int, image string, command []string) []*framework.RCConfig {
	configs := make([]*framework.RCConfig, 0, count)
	for i := 1; i <= count; i++ {
		config := &framework.RCConfig{
			Client:     c,
			Name:       groupName + "-" + strconv.Itoa(i),
			Namespace:  nss[i%len(nss)].Name,
			Timeout:    10 * time.Minute,
			Image:      image,
			Command:    command,
			Replicas:   size,
			CpuRequest: 10,       // 0.01 core
			MemRequest: 26214400, // 25MB
		}
		configs = append(configs, config)
	}
	return configs
}

func generateServicesForConfigs(configs []*framework.RCConfig) []*api.Service {
	services := make([]*api.Service, 0, len(configs))
	for _, config := range configs {
		serviceName := config.Name + "-svc"
		labels := map[string]string{"name": config.Name}
		service := &api.Service{
			ObjectMeta: api.ObjectMeta{
				Name:      serviceName,
				Namespace: config.Namespace,
			},
			Spec: api.ServiceSpec{
				Selector: labels,
				Ports: []api.ServicePort{{
					Port:       80,
					TargetPort: intstr.FromInt(80),
				}},
			},
		}
		services = append(services, service)
	}
	return services
}

func sleepUpTo(d time.Duration) {
	time.Sleep(time.Duration(rand.Int63n(d.Nanoseconds())))
}

func createAllRC(configs []*framework.RCConfig, creatingTime time.Duration) {
	var wg sync.WaitGroup
	wg.Add(len(configs))
	for _, config := range configs {
		go createRC(&wg, config, creatingTime)
	}
	wg.Wait()
}

func createRC(wg *sync.WaitGroup, config *framework.RCConfig, creatingTime time.Duration) {
	defer GinkgoRecover()
	defer wg.Done()

	sleepUpTo(creatingTime)
	framework.ExpectNoError(framework.RunRC(*config), fmt.Sprintf("creating rc %s", config.Name))
}

func scaleAllRC(configs []*framework.RCConfig, scalingTime time.Duration) {
	var wg sync.WaitGroup
	wg.Add(len(configs))
	for _, config := range configs {
		go scaleRC(&wg, config, scalingTime)
	}
	wg.Wait()
}

// Scales RC to a random size within [0.5*size, 1.5*size] and lists all the pods afterwards.
// Scaling happens always based on original size, not the current size.
func scaleRC(wg *sync.WaitGroup, config *framework.RCConfig, scalingTime time.Duration) {
	defer GinkgoRecover()
	defer wg.Done()

	sleepUpTo(scalingTime)
	newSize := uint(rand.Intn(config.Replicas) + config.Replicas/2)
	framework.ExpectNoError(framework.ScaleRC(config.Client, config.Namespace, config.Name, newSize, true),
		fmt.Sprintf("scaling rc %s for the first time", config.Name))
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": config.Name}))
	options := api.ListOptions{
		LabelSelector:   selector,
		ResourceVersion: "0",
	}
	_, err := config.Client.Pods(config.Namespace).List(options)
	framework.ExpectNoError(err, fmt.Sprintf("listing pods from rc %v", config.Name))
}

func deleteAllRC(configs []*framework.RCConfig, deletingTime time.Duration) {
	var wg sync.WaitGroup
	wg.Add(len(configs))
	for _, config := range configs {
		go deleteRC(&wg, config, deletingTime)
	}
	wg.Wait()
}

func deleteRC(wg *sync.WaitGroup, config *framework.RCConfig, deletingTime time.Duration) {
	defer GinkgoRecover()
	defer wg.Done()

	sleepUpTo(deletingTime)
	framework.ExpectNoError(framework.DeleteRC(config.Client, config.Namespace, config.Name), fmt.Sprintf("deleting rc %s", config.Name))
}
<commit_msg>Allow for overriding throughput in load test<commit_after>/*
Copyright 2015 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 (
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"sync"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/intstr"
	"k8s.io/kubernetes/test/e2e/framework"

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

const (
	smallRCSize       = 5
	mediumRCSize      = 30
	bigRCSize         = 250
	smallRCGroupName  = "load-small-rc"
	mediumRCGroupName = "load-medium-rc"
	bigRCGroupName    = "load-big-rc"
	smallRCBatchSize  = 30
	mediumRCBatchSize = 5
	bigRCBatchSize    = 1
	// We start RCs/Services/pods/... in different namespace in this test.
	// nodeCountPerNamespace determines how many namespaces we will be using
	// depending on the number of nodes in the underlying cluster.
	nodeCountPerNamespace = 250
)

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = framework.KubeDescribe("Load capacity", func() {
	var c *client.Client
	var nodeCount int
	var ns string
	var configs []*framework.RCConfig
	var namespaces []*api.Namespace

	// Gathers metrics before teardown
	// TODO add flag that allows to skip cleanup on failure
	AfterEach(func() {
		// Verify latency metrics
		highLatencyRequests, err := framework.HighLatencyRequests(c)
		framework.ExpectNoError(err, "Too many instances metrics above the threshold")
		Expect(highLatencyRequests).NotTo(BeNumerically(">", 0))
	})

	// Explicitly put here, to delete namespace at the end of the test
	// (after measuring latency metrics, etc.).
	options := framework.FrameworkOptions{
		ClientQPS:   50,
		ClientBurst: 100,
	}
	f := framework.NewFramework("load", options, nil)
	f.NamespaceDeletionTimeout = time.Hour

	BeforeEach(func() {
		c = f.Client

		// In large clusters we may get to this point but still have a bunch
		// of nodes without Routes created. Since this would make a node
		// unschedulable, we need to wait until all of them are schedulable.
		framework.ExpectNoError(framework.WaitForAllNodesSchedulable(c))

		ns = f.Namespace.Name
		nodes := framework.GetReadySchedulableNodesOrDie(c)
		nodeCount = len(nodes.Items)
		Expect(nodeCount).NotTo(BeZero())

		// Terminating a namespace (deleting the remaining objects from it - which
		// generally means events) can affect the current run. Thus we wait for all
		// terminating namespace to be finally deleted before starting this test.
		err := framework.CheckTestingNSDeletedExcept(c, ns)
		framework.ExpectNoError(err)

		framework.ExpectNoError(framework.ResetMetrics(c))
	})

	type Load struct {
		podsPerNode int
		image       string
		command     []string
	}

	loadTests := []Load{
		// The container will consume 1 cpu and 512mb of memory.
		{podsPerNode: 3, image: "jess/stress", command: []string{"stress", "-c", "1", "-m", "2"}},
		{podsPerNode: 30, image: "gcr.io/google_containers/serve_hostname:v1.4"},
	}

	for _, testArg := range loadTests {
		name := fmt.Sprintf("should be able to handle %v pods per node", testArg.podsPerNode)
		if testArg.podsPerNode == 30 {
			name = "[Feature:Performance] " + name
		} else {
			name = "[Feature:ManualPerformance] " + name
		}
		itArg := testArg

		It(name, func() {
			// Create a number of namespaces.
			namespaces = createNamespaces(f, nodeCount, itArg.podsPerNode)

			totalPods := itArg.podsPerNode * nodeCount
			configs = generateRCConfigs(totalPods, itArg.image, itArg.command, c, namespaces)
			var services []*api.Service
			// Read the environment variable to see if we want to create services
			createServices := os.Getenv("CREATE_SERVICES")
			if createServices == "true" {
				framework.Logf("Creating services")
				services := generateServicesForConfigs(configs)
				for _, service := range services {
					_, err := c.Services(service.Namespace).Create(service)
					framework.ExpectNoError(err)
				}
			} else {
				framework.Logf("Skipping service creation")
			}

			// We assume a default throughput of 10 pods/second throughput.
			// We may want to revisit it in the future.
			// However, this can be overriden by LOAD_TEST_THROUGHPUT env var.
			throughput := 10
			if throughputEnv := os.Getenv("LOAD_TEST_THROUGHPUT"); throughputEnv != "" {
				if newThroughput, err := strconv.Atoi(throughputEnv); err == nil {
					throughput = newThroughput
				}
			}

			// Simulate lifetime of RC:
			//  * create with initial size
			//  * scale RC to a random size and list all pods
			//  * scale RC to a random size and list all pods
			//  * delete it
			//
			// This will generate ~5 creations/deletions per second assuming:
			//  - X small RCs each 5 pods   [ 5 * X = totalPods / 2 ]
			//  - Y medium RCs each 30 pods [ 30 * Y = totalPods / 4 ]
			//  - Z big RCs each 250 pods   [ 250 * Z = totalPods / 4]

			// We would like to spread creating replication controllers over time
			// to make it possible to create/schedule them in the meantime.
			// Currently we assume <throughput> pods/second average throughput.
			// We may want to revisit it in the future.
			creatingTime := time.Duration(totalPods/throughput) * time.Second
			createAllRC(configs, creatingTime)
			By("============================================================================")

			// We would like to spread scaling replication controllers over time
			// to make it possible to create/schedule & delete them in the meantime.
			// Currently we assume that <throughput> pods/second average throughput.
			// The expected number of created/deleted pods is less than totalPods/3.
			scalingTime := time.Duration(totalPods/(3*throughput)) * time.Second
			scaleAllRC(configs, scalingTime)
			By("============================================================================")

			scaleAllRC(configs, scalingTime)
			By("============================================================================")

			// Cleanup all created replication controllers.
			// Currently we assume <throughput> pods/second average deletion throughput.
			// We may want to revisit it in the future.
			deletingTime := time.Duration(totalPods/throughput) * time.Second
			deleteAllRC(configs, deletingTime)
			if createServices == "true" {
				for _, service := range services {
					err := c.Services(ns).Delete(service.Name)
					framework.ExpectNoError(err)
				}
				framework.Logf("%v Services created.", len(services))
			}
		})
	}
})

func createNamespaces(f *framework.Framework, nodeCount, podsPerNode int) []*api.Namespace {
	namespaceCount := (nodeCount + nodeCountPerNamespace - 1) / nodeCountPerNamespace
	namespaces := []*api.Namespace{}
	for i := 1; i <= namespaceCount; i++ {
		namespace, err := f.CreateNamespace(fmt.Sprintf("load-%d-nodepods-%d", podsPerNode, i), nil)
		framework.ExpectNoError(err)
		namespaces = append(namespaces, namespace)
	}
	return namespaces
}

func computeRCCounts(total int) (int, int, int) {
	// Small RCs owns ~0.5 of total number of pods, medium and big RCs ~0.25 each.
	// For example for 3000 pods (100 nodes, 30 pods per node) there are:
	//  - 300 small RCs each 5 pods
	//  - 25 medium RCs each 30 pods
	//  - 3 big RCs each 250 pods
	bigRCCount := total / 4 / bigRCSize
	total -= bigRCCount * bigRCSize
	mediumRCCount := total / 3 / mediumRCSize
	total -= mediumRCCount * mediumRCSize
	smallRCCount := total / smallRCSize
	return smallRCCount, mediumRCCount, bigRCCount
}

func generateRCConfigs(totalPods int, image string, command []string, c *client.Client, nss []*api.Namespace) []*framework.RCConfig {
	configs := make([]*framework.RCConfig, 0)

	smallRCCount, mediumRCCount, bigRCCount := computeRCCounts(totalPods)
	configs = append(configs, generateRCConfigsForGroup(c, nss, smallRCGroupName, smallRCSize, smallRCCount, image, command)...)
	configs = append(configs, generateRCConfigsForGroup(c, nss, mediumRCGroupName, mediumRCSize, mediumRCCount, image, command)...)
	configs = append(configs, generateRCConfigsForGroup(c, nss, bigRCGroupName, bigRCSize, bigRCCount, image, command)...)

	return configs
}

func generateRCConfigsForGroup(c *client.Client, nss []*api.Namespace, groupName string, size, count int, image string, command []string) []*framework.RCConfig {
	configs := make([]*framework.RCConfig, 0, count)
	for i := 1; i <= count; i++ {
		config := &framework.RCConfig{
			Client:     c,
			Name:       groupName + "-" + strconv.Itoa(i),
			Namespace:  nss[i%len(nss)].Name,
			Timeout:    10 * time.Minute,
			Image:      image,
			Command:    command,
			Replicas:   size,
			CpuRequest: 10,       // 0.01 core
			MemRequest: 26214400, // 25MB
		}
		configs = append(configs, config)
	}
	return configs
}

func generateServicesForConfigs(configs []*framework.RCConfig) []*api.Service {
	services := make([]*api.Service, 0, len(configs))
	for _, config := range configs {
		serviceName := config.Name + "-svc"
		labels := map[string]string{"name": config.Name}
		service := &api.Service{
			ObjectMeta: api.ObjectMeta{
				Name:      serviceName,
				Namespace: config.Namespace,
			},
			Spec: api.ServiceSpec{
				Selector: labels,
				Ports: []api.ServicePort{{
					Port:       80,
					TargetPort: intstr.FromInt(80),
				}},
			},
		}
		services = append(services, service)
	}
	return services
}

func sleepUpTo(d time.Duration) {
	time.Sleep(time.Duration(rand.Int63n(d.Nanoseconds())))
}

func createAllRC(configs []*framework.RCConfig, creatingTime time.Duration) {
	var wg sync.WaitGroup
	wg.Add(len(configs))
	for _, config := range configs {
		go createRC(&wg, config, creatingTime)
	}
	wg.Wait()
}

func createRC(wg *sync.WaitGroup, config *framework.RCConfig, creatingTime time.Duration) {
	defer GinkgoRecover()
	defer wg.Done()

	sleepUpTo(creatingTime)
	framework.ExpectNoError(framework.RunRC(*config), fmt.Sprintf("creating rc %s", config.Name))
}

func scaleAllRC(configs []*framework.RCConfig, scalingTime time.Duration) {
	var wg sync.WaitGroup
	wg.Add(len(configs))
	for _, config := range configs {
		go scaleRC(&wg, config, scalingTime)
	}
	wg.Wait()
}

// Scales RC to a random size within [0.5*size, 1.5*size] and lists all the pods afterwards.
// Scaling happens always based on original size, not the current size.
func scaleRC(wg *sync.WaitGroup, config *framework.RCConfig, scalingTime time.Duration) {
	defer GinkgoRecover()
	defer wg.Done()

	sleepUpTo(scalingTime)
	newSize := uint(rand.Intn(config.Replicas) + config.Replicas/2)
	framework.ExpectNoError(framework.ScaleRC(config.Client, config.Namespace, config.Name, newSize, true),
		fmt.Sprintf("scaling rc %s for the first time", config.Name))
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": config.Name}))
	options := api.ListOptions{
		LabelSelector:   selector,
		ResourceVersion: "0",
	}
	_, err := config.Client.Pods(config.Namespace).List(options)
	framework.ExpectNoError(err, fmt.Sprintf("listing pods from rc %v", config.Name))
}

func deleteAllRC(configs []*framework.RCConfig, deletingTime time.Duration) {
	var wg sync.WaitGroup
	wg.Add(len(configs))
	for _, config := range configs {
		go deleteRC(&wg, config, deletingTime)
	}
	wg.Wait()
}

func deleteRC(wg *sync.WaitGroup, config *framework.RCConfig, deletingTime time.Duration) {
	defer GinkgoRecover()
	defer wg.Done()

	sleepUpTo(deletingTime)
	framework.ExpectNoError(framework.DeleteRC(config.Client, config.Namespace, config.Name), fmt.Sprintf("deleting rc %s", config.Name))
}
<|endoftext|>