<commit_before>package main

import (
	"net/http"
	"time"

	"github.com/TykTechnologies/leakybucket"
	"github.com/TykTechnologies/leakybucket/memorycache"
	"github.com/TykTechnologies/tyk/config"
	"github.com/TykTechnologies/tyk/storage"
	"github.com/TykTechnologies/tyk/user"
)

type PublicSession struct {
	Quota struct {
		QuotaMax       int64 `json:"quota_max"`
		QuotaRemaining int64 `json:"quota_remaining"`
		QuotaRenews    int64 `json:"quota_renews"`
	} `json:"quota"`
	RateLimit struct {
		Rate float64 `json:"requests"`
		Per  float64 `json:"per_unit"`
	} `json:"rate_limit"`
}

const (
	QuotaKeyPrefix     = "quota-"
	RateLimitKeyPrefix = "rate-limit-"
)

// SessionLimiter is the rate limiter for the API, use ForwardMessage() to
// check if a message should pass through or not
type SessionLimiter struct {
	bucketStore leakybucket.Storage
}

func (l *SessionLimiter) doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey string,
	currentSession *user.SessionState,
	store storage.Handler,
	globalConf *config.Config,
	apiLimit *user.APILimit, dryRun bool) bool {

	var per, rate float64

	if apiLimit != nil { // respect limit on API level
		per = apiLimit.Per
		rate = apiLimit.Rate
	} else {
		per = currentSession.Per
		rate = currentSession.Rate
	}

	log.Debug("[RATELIMIT] Inbound raw key is: ", key)
	log.Debug("[RATELIMIT] Rate limiter key is: ", rateLimiterKey)
	pipeline := globalConf.EnableNonTransactionalRateLimiter

	var ratePerPeriodNow int
	if dryRun {
		ratePerPeriodNow, _ = store.GetRollingWindow(rateLimiterKey, int64(per), pipeline)
	} else {
		ratePerPeriodNow, _ = store.SetRollingWindow(rateLimiterKey, int64(per), "-1", pipeline)
	}

	//log.Info("Num Requests: ", ratePerPeriodNow)

	// Subtract by 1 because of the delayed add in the window
	subtractor := 1
	if globalConf.EnableSentinelRateLimiter {
		// and another subtraction because of the preemptive limit
		subtractor = 2
	}
	// The test TestRateLimitForAPIAndRateLimitAndQuotaCheck
	// will only work with ththese two lines here
	//log.Info("break: ", (int(currentSession.Rate) - subtractor))
	if ratePerPeriodNow > int(rate)-subtractor {
		// Set a sentinel value with expire
		if globalConf.EnableSentinelRateLimiter {
			if !dryRun {
				store.SetRawKey(rateLimiterSentinelKey, "1", int64(per))
			}
		}
		return true
	}

	return false
}

type sessionFailReason uint

const (
	sessionFailNone sessionFailReason = iota
	sessionFailRateLimit
	sessionFailQuota
)

// ForwardMessage will enforce rate limiting, returning a non-zero
// sessionFailReason if session limits have been exceeded.
// Key values to manage rate are Rate and Per, e.g. Rate of 10 messages
// Per 10 seconds
func (l *SessionLimiter) ForwardMessage(r *http.Request, currentSession *user.SessionState, key string, store storage.Handler, enableRL, enableQ bool, globalConf *config.Config, apiID string, dryRun bool) sessionFailReason {
	if enableRL {
		// check for limit on API level (set to session by ApplyPolicies)
		var apiLimit *user.APILimit
		if len(currentSession.AccessRights) > 0 {
			if rights, ok := currentSession.AccessRights[apiID]; !ok {
				log.WithField("apiID", apiID).Debug("[RATE] unexpected apiID")
				return sessionFailRateLimit
			} else {
				apiLimit = rights.Limit
			}
		}

		if globalConf.EnableSentinelRateLimiter {
			rateLimiterKey := RateLimitKeyPrefix + currentSession.KeyHash()
			rateLimiterSentinelKey := RateLimitKeyPrefix + currentSession.KeyHash() + ".BLOCKED"
			if apiLimit != nil {
				rateLimiterKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash()
				rateLimiterSentinelKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash() + ".BLOCKED"
			}

			go l.doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey, currentSession, store, globalConf, apiLimit, dryRun)

			// Check sentinel
			_, sentinelActive := store.GetRawKey(rateLimiterSentinelKey)
			if sentinelActive == nil {
				// Sentinel is set, fail
				return sessionFailRateLimit
			}
		} else if globalConf.EnableRedisRollingLimiter {
			rateLimiterKey := RateLimitKeyPrefix + currentSession.KeyHash()
			rateLimiterSentinelKey := RateLimitKeyPrefix + currentSession.KeyHash() + ".BLOCKED"
			if apiLimit != nil {
				rateLimiterKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash()
				rateLimiterSentinelKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash() + ".BLOCKED"
			}

			if l.doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey, currentSession, store, globalConf, apiLimit, dryRun) {
				return sessionFailRateLimit
			}
		} else {
			// In-memory limiter
			if l.bucketStore == nil {
				l.bucketStore = memorycache.New()
			}

			// If a token has been updated, we must ensure we don't use
			// an old bucket an let the cache deal with it
			bucketKey := ""
			var currRate float64
			var per float64
			if apiLimit == nil {
				bucketKey = key + ":" + currentSession.LastUpdated
				currRate = currentSession.Rate
				per = currentSession.Per
			} else { // respect limit on API level
				bucketKey = apiID + ":" + key + ":" + currentSession.LastUpdated
				currRate = apiLimit.Rate
				per = apiLimit.Per
			}

			// DRL will always overflow with more servers on low rates
			rate := uint(currRate * float64(DRLManager.RequestTokenValue))
			if rate < uint(DRLManager.CurrentTokenValue) {
				rate = uint(DRLManager.CurrentTokenValue)
			}

			userBucket, err := l.bucketStore.Create(bucketKey, rate, time.Duration(per)*time.Second)
			if err != nil {
				log.Error("Failed to create bucket!")
				return sessionFailRateLimit
			}

			if dryRun {
				// if userBucket is empty and not expired.
				if userBucket.Remaining() == 0 && time.Now().Before(userBucket.Reset()) {
					return sessionFailRateLimit
				}
			} else {
				_, errF := userBucket.Add(uint(DRLManager.CurrentTokenValue))
				if errF != nil {
					return sessionFailRateLimit
				}
			}
		}
	}

	if enableQ {
		if globalConf.LegacyEnableAllowanceCountdown {
			currentSession.Allowance--
		}

		if l.RedisQuotaExceeded(r, currentSession, key, store, apiID) {
			return sessionFailQuota
		}
	}

	return sessionFailNone

}

func (l *SessionLimiter) RedisQuotaExceeded(r *http.Request, currentSession *user.SessionState, key string, store storage.Handler, apiID string) bool {
	log.Debug("[QUOTA] Inbound raw key is: ", key)

	// check for limit on API level (set to session by ApplyPolicies)
	var apiLimit *user.APILimit
	if len(currentSession.AccessRights) > 0 {
		if rights, ok := currentSession.AccessRights[apiID]; !ok {
			log.WithField("apiID", apiID).Debug("[QUOTA] unexpected apiID")
			return false
		} else {
			apiLimit = rights.Limit
		}
	}

	// Are they unlimited?
	if apiLimit == nil {
		if currentSession.QuotaMax == -1 {
			// No quota set
			return false
		}
	} else if apiLimit.QuotaMax == -1 {
		// No quota set
		return false
	}

	rawKey := ""
	var quotaRenewalRate int64
	var quotaRenews int64
	var quotaMax int64
	if apiLimit == nil {
		rawKey = QuotaKeyPrefix + currentSession.KeyHash()
		quotaRenewalRate = currentSession.QuotaRenewalRate
		quotaRenews = currentSession.QuotaRenews
		quotaMax = currentSession.QuotaMax
	} else {
		rawKey = QuotaKeyPrefix + apiID + "-" + currentSession.KeyHash()
		quotaRenewalRate = apiLimit.QuotaRenewalRate
		quotaRenews = apiLimit.QuotaRenews
		quotaMax = apiLimit.QuotaMax
	}

	log.Debug("[QUOTA] Quota limiter key is: ", rawKey)
	log.Debug("Renewing with TTL: ", quotaRenewalRate)
	// INCR the key (If it equals 1 - set EXPIRE)
	qInt := store.IncrememntWithExpire(rawKey, quotaRenewalRate)

	// if the returned val is >= quota: block
	if qInt-1 >= quotaMax {
		renewalDate := time.Unix(quotaRenews, 0)
		log.Debug("Renewal Date is: ", renewalDate)
		log.Debug("As epoch: ", quotaRenews)
		log.Debug("Session: ", currentSession)
		log.Debug("Now:", time.Now())
		if time.Now().After(renewalDate) {
			// The renewal date is in the past, we should update the quota!
			// Also, this fixes legacy issues where there is no TTL on quota buckets
			log.Warning("Incorrect key expiry setting detected, correcting")
			go store.DeleteRawKey(rawKey)
			qInt = 1
		} else {
			// Renewal date is in the future and the quota is exceeded
			return true
		}

	}

	// If this is a new Quota period, ensure we let the end user know
	if qInt == 1 {
		current := time.Now().Unix()
		if apiLimit == nil {
			currentSession.QuotaRenews = current + quotaRenewalRate
		} else {
			apiLimit.QuotaRenews = current + quotaRenewalRate
		}
		ctxScheduleSessionUpdate(r)
	}

	// If not, pass and set the values of the session to quotamax - counter
	remaining := quotaMax - qInt
	if remaining < 0 {
		remaining = 0
	}

	if apiLimit == nil {
		currentSession.QuotaRemaining = remaining
	} else {
		apiLimit.QuotaRemaining = remaining
	}

	return false
}
<commit_msg>Debug instead of Warning (#2220)<commit_after>package main

import (
	"net/http"
	"time"

	"github.com/TykTechnologies/leakybucket"
	"github.com/TykTechnologies/leakybucket/memorycache"
	"github.com/TykTechnologies/tyk/config"
	"github.com/TykTechnologies/tyk/storage"
	"github.com/TykTechnologies/tyk/user"
)

type PublicSession struct {
	Quota struct {
		QuotaMax       int64 `json:"quota_max"`
		QuotaRemaining int64 `json:"quota_remaining"`
		QuotaRenews    int64 `json:"quota_renews"`
	} `json:"quota"`
	RateLimit struct {
		Rate float64 `json:"requests"`
		Per  float64 `json:"per_unit"`
	} `json:"rate_limit"`
}

const (
	QuotaKeyPrefix     = "quota-"
	RateLimitKeyPrefix = "rate-limit-"
)

// SessionLimiter is the rate limiter for the API, use ForwardMessage() to
// check if a message should pass through or not
type SessionLimiter struct {
	bucketStore leakybucket.Storage
}

func (l *SessionLimiter) doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey string,
	currentSession *user.SessionState,
	store storage.Handler,
	globalConf *config.Config,
	apiLimit *user.APILimit, dryRun bool) bool {

	var per, rate float64

	if apiLimit != nil { // respect limit on API level
		per = apiLimit.Per
		rate = apiLimit.Rate
	} else {
		per = currentSession.Per
		rate = currentSession.Rate
	}

	log.Debug("[RATELIMIT] Inbound raw key is: ", key)
	log.Debug("[RATELIMIT] Rate limiter key is: ", rateLimiterKey)
	pipeline := globalConf.EnableNonTransactionalRateLimiter

	var ratePerPeriodNow int
	if dryRun {
		ratePerPeriodNow, _ = store.GetRollingWindow(rateLimiterKey, int64(per), pipeline)
	} else {
		ratePerPeriodNow, _ = store.SetRollingWindow(rateLimiterKey, int64(per), "-1", pipeline)
	}

	//log.Info("Num Requests: ", ratePerPeriodNow)

	// Subtract by 1 because of the delayed add in the window
	subtractor := 1
	if globalConf.EnableSentinelRateLimiter {
		// and another subtraction because of the preemptive limit
		subtractor = 2
	}
	// The test TestRateLimitForAPIAndRateLimitAndQuotaCheck
	// will only work with ththese two lines here
	//log.Info("break: ", (int(currentSession.Rate) - subtractor))
	if ratePerPeriodNow > int(rate)-subtractor {
		// Set a sentinel value with expire
		if globalConf.EnableSentinelRateLimiter {
			if !dryRun {
				store.SetRawKey(rateLimiterSentinelKey, "1", int64(per))
			}
		}
		return true
	}

	return false
}

type sessionFailReason uint

const (
	sessionFailNone sessionFailReason = iota
	sessionFailRateLimit
	sessionFailQuota
)

// ForwardMessage will enforce rate limiting, returning a non-zero
// sessionFailReason if session limits have been exceeded.
// Key values to manage rate are Rate and Per, e.g. Rate of 10 messages
// Per 10 seconds
func (l *SessionLimiter) ForwardMessage(r *http.Request, currentSession *user.SessionState, key string, store storage.Handler, enableRL, enableQ bool, globalConf *config.Config, apiID string, dryRun bool) sessionFailReason {
	if enableRL {
		// check for limit on API level (set to session by ApplyPolicies)
		var apiLimit *user.APILimit
		if len(currentSession.AccessRights) > 0 {
			if rights, ok := currentSession.AccessRights[apiID]; !ok {
				log.WithField("apiID", apiID).Debug("[RATE] unexpected apiID")
				return sessionFailRateLimit
			} else {
				apiLimit = rights.Limit
			}
		}

		if globalConf.EnableSentinelRateLimiter {
			rateLimiterKey := RateLimitKeyPrefix + currentSession.KeyHash()
			rateLimiterSentinelKey := RateLimitKeyPrefix + currentSession.KeyHash() + ".BLOCKED"
			if apiLimit != nil {
				rateLimiterKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash()
				rateLimiterSentinelKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash() + ".BLOCKED"
			}

			go l.doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey, currentSession, store, globalConf, apiLimit, dryRun)

			// Check sentinel
			_, sentinelActive := store.GetRawKey(rateLimiterSentinelKey)
			if sentinelActive == nil {
				// Sentinel is set, fail
				return sessionFailRateLimit
			}
		} else if globalConf.EnableRedisRollingLimiter {
			rateLimiterKey := RateLimitKeyPrefix + currentSession.KeyHash()
			rateLimiterSentinelKey := RateLimitKeyPrefix + currentSession.KeyHash() + ".BLOCKED"
			if apiLimit != nil {
				rateLimiterKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash()
				rateLimiterSentinelKey = RateLimitKeyPrefix + apiID + "-" + currentSession.KeyHash() + ".BLOCKED"
			}

			if l.doRollingWindowWrite(key, rateLimiterKey, rateLimiterSentinelKey, currentSession, store, globalConf, apiLimit, dryRun) {
				return sessionFailRateLimit
			}
		} else {
			// In-memory limiter
			if l.bucketStore == nil {
				l.bucketStore = memorycache.New()
			}

			// If a token has been updated, we must ensure we don't use
			// an old bucket an let the cache deal with it
			bucketKey := ""
			var currRate float64
			var per float64
			if apiLimit == nil {
				bucketKey = key + ":" + currentSession.LastUpdated
				currRate = currentSession.Rate
				per = currentSession.Per
			} else { // respect limit on API level
				bucketKey = apiID + ":" + key + ":" + currentSession.LastUpdated
				currRate = apiLimit.Rate
				per = apiLimit.Per
			}

			// DRL will always overflow with more servers on low rates
			rate := uint(currRate * float64(DRLManager.RequestTokenValue))
			if rate < uint(DRLManager.CurrentTokenValue) {
				rate = uint(DRLManager.CurrentTokenValue)
			}

			userBucket, err := l.bucketStore.Create(bucketKey, rate, time.Duration(per)*time.Second)
			if err != nil {
				log.Error("Failed to create bucket!")
				return sessionFailRateLimit
			}

			if dryRun {
				// if userBucket is empty and not expired.
				if userBucket.Remaining() == 0 && time.Now().Before(userBucket.Reset()) {
					return sessionFailRateLimit
				}
			} else {
				_, errF := userBucket.Add(uint(DRLManager.CurrentTokenValue))
				if errF != nil {
					return sessionFailRateLimit
				}
			}
		}
	}

	if enableQ {
		if globalConf.LegacyEnableAllowanceCountdown {
			currentSession.Allowance--
		}

		if l.RedisQuotaExceeded(r, currentSession, key, store, apiID) {
			return sessionFailQuota
		}
	}

	return sessionFailNone

}

func (l *SessionLimiter) RedisQuotaExceeded(r *http.Request, currentSession *user.SessionState, key string, store storage.Handler, apiID string) bool {
	log.Debug("[QUOTA] Inbound raw key is: ", key)

	// check for limit on API level (set to session by ApplyPolicies)
	var apiLimit *user.APILimit
	if len(currentSession.AccessRights) > 0 {
		if rights, ok := currentSession.AccessRights[apiID]; !ok {
			log.WithField("apiID", apiID).Debug("[QUOTA] unexpected apiID")
			return false
		} else {
			apiLimit = rights.Limit
		}
	}

	// Are they unlimited?
	if apiLimit == nil {
		if currentSession.QuotaMax == -1 {
			// No quota set
			return false
		}
	} else if apiLimit.QuotaMax == -1 {
		// No quota set
		return false
	}

	rawKey := ""
	var quotaRenewalRate int64
	var quotaRenews int64
	var quotaMax int64
	if apiLimit == nil {
		rawKey = QuotaKeyPrefix + currentSession.KeyHash()
		quotaRenewalRate = currentSession.QuotaRenewalRate
		quotaRenews = currentSession.QuotaRenews
		quotaMax = currentSession.QuotaMax
	} else {
		rawKey = QuotaKeyPrefix + apiID + "-" + currentSession.KeyHash()
		quotaRenewalRate = apiLimit.QuotaRenewalRate
		quotaRenews = apiLimit.QuotaRenews
		quotaMax = apiLimit.QuotaMax
	}

	log.Debug("[QUOTA] Quota limiter key is: ", rawKey)
	log.Debug("Renewing with TTL: ", quotaRenewalRate)
	// INCR the key (If it equals 1 - set EXPIRE)
	qInt := store.IncrememntWithExpire(rawKey, quotaRenewalRate)

	// if the returned val is >= quota: block
	if qInt-1 >= quotaMax {
		renewalDate := time.Unix(quotaRenews, 0)
		log.Debug("Renewal Date is: ", renewalDate)
		log.Debug("As epoch: ", quotaRenews)
		log.Debug("Session: ", currentSession)
		log.Debug("Now:", time.Now())
		if time.Now().After(renewalDate) {
			// The renewal date is in the past, we should update the quota!
			// Also, this fixes legacy issues where there is no TTL on quota buckets
			log.Debug("Incorrect key expiry setting detected, correcting")
			go store.DeleteRawKey(rawKey)
			qInt = 1
		} else {
			// Renewal date is in the future and the quota is exceeded
			return true
		}

	}

	// If this is a new Quota period, ensure we let the end user know
	if qInt == 1 {
		current := time.Now().Unix()
		if apiLimit == nil {
			currentSession.QuotaRenews = current + quotaRenewalRate
		} else {
			apiLimit.QuotaRenews = current + quotaRenewalRate
		}
		ctxScheduleSessionUpdate(r)
	}

	// If not, pass and set the values of the session to quotamax - counter
	remaining := quotaMax - qInt
	if remaining < 0 {
		remaining = 0
	}

	if apiLimit == nil {
		currentSession.QuotaRemaining = remaining
	} else {
		apiLimit.QuotaRemaining = remaining
	}

	return false
}
<|endoftext|>


<commit_before>package registry

import (
	"errors"
	"sort"
	"sync"

	"github.com/coreos/fleet/Godeps/_workspace/src/github.com/coreos/go-semver/semver"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/unit"
)

func NewFakeRegistry() *FakeRegistry {
	return &FakeRegistry{
		machines:  []machine.MachineState{},
		jobStates: map[string]*unit.UnitState{},
		jobs:      map[string]job.Job{},
		units:     []unit.Unit{},
		version:   nil,
		bids:      map[string]pkg.Set{},
	}
}

type FakeRegistry struct {
	// Not all methods of required by the Registry interface are implemented
	// by the TestRegistry. Any calls to these unimplemented methods will
	// result in a panic.
	Registry
	sync.RWMutex

	machines  []machine.MachineState
	jobStates map[string]*unit.UnitState
	jobs      map[string]job.Job
	units     []unit.Unit
	version   *semver.Version
	bids      map[string]pkg.Set
}

func (f *FakeRegistry) SetMachines(machines []machine.MachineState) {
	f.Lock()
	defer f.Unlock()

	f.machines = machines
}

func (f *FakeRegistry) SetJobs(jobs []job.Job) {
	f.Lock()
	defer f.Unlock()

	f.jobs = make(map[string]job.Job, len(jobs))
	for _, j := range jobs {
		f.jobs[j.Name] = j
	}
}

func (f *FakeRegistry) SetUnitStates(jobStates map[string]*unit.UnitState) {
	f.Lock()
	defer f.Unlock()

	f.jobStates = jobStates
}

func (f *FakeRegistry) SetUnits(units []unit.Unit) {
	f.Lock()
	defer f.Unlock()

	f.units = units
}

func (f *FakeRegistry) SetLatestVersion(v semver.Version) {
	f.Lock()
	defer f.Unlock()

	f.version = &v
}

func (f *FakeRegistry) Machines() ([]machine.MachineState, error) {
	f.RLock()
	defer f.RUnlock()

	return f.machines, nil
}

func (f *FakeRegistry) Jobs() ([]job.Job, error) {
	f.RLock()
	defer f.RUnlock()

	var sorted sort.StringSlice
	for _, j := range f.jobs {
		sorted = append(sorted, j.Name)
	}
	sorted.Sort()

	jobs := make([]job.Job, 0, len(f.jobs))
	for _, jName := range sorted {
		jobs = append(jobs, f.jobs[jName])
	}

	return jobs, nil
}

func (f *FakeRegistry) Job(name string) (*job.Job, error) {
	f.RLock()
	defer f.RUnlock()

	j, ok := f.jobs[name]
	if !ok {
		return nil, nil
	}

	j.UnitState = f.jobStates[name]
	return &j, nil
}

func (f *FakeRegistry) CreateJob(j *job.Job) error {
	f.Lock()
	defer f.Unlock()

	_, ok := f.jobs[j.Name]
	if ok {
		return errors.New("Job already exists")
	}

	f.jobs[j.Name] = *j
	return nil
}

func (f *FakeRegistry) DestroyJob(name string) error {
	f.Lock()
	defer f.Unlock()

	delete(f.jobs, name)
	return nil
}

func (f *FakeRegistry) Bids(jo *job.JobOffer) (pkg.Set, error) {
	f.RLock()
	defer f.RUnlock()

	return f.bids[jo.Job.Name], nil
}

func (f *FakeRegistry) SubmitJobBid(jName, machID string) {
	f.Lock()
	defer f.Unlock()

	_, ok := f.bids[jName]
	if !ok {
		f.bids[jName] = pkg.NewUnsafeSet()
	}
	f.bids[jName].Add(machID)
}

func (f *FakeRegistry) SetJobTargetState(name string, target job.JobState) error {
	f.Lock()
	defer f.Unlock()

	j, ok := f.jobs[name]

	if !ok {
		return errors.New("job does not exist")
	}

	j.TargetState = target
	f.jobs[name] = j

	return nil
}

func (f *FakeRegistry) ScheduleJob(name string, machID string) error {
	f.Lock()
	defer f.Unlock()

	j, ok := f.jobs[name]

	if !ok {
		return errors.New("job does not exist")
	}

	j.TargetMachineID = machID
	f.jobs[name] = j

	return nil
}

func (f *FakeRegistry) SaveUnitState(jobName string, unitState *unit.UnitState) {
	f.Lock()
	defer f.Unlock()

	f.jobStates[jobName] = unitState
}

func (f *FakeRegistry) LatestVersion() (*semver.Version, error) {
	f.RLock()
	defer f.RUnlock()

	return f.version, nil
}
<commit_msg>registry: fill in more FakeRegistry methods<commit_after>package registry

import (
	"errors"
	"sort"
	"sync"
	"time"

	"github.com/coreos/fleet/Godeps/_workspace/src/github.com/coreos/go-semver/semver"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/unit"
)

func NewFakeRegistry() *FakeRegistry {
	return &FakeRegistry{
		machines:  []machine.MachineState{},
		jobStates: map[string]*unit.UnitState{},
		jobs:      map[string]job.Job{},
		units:     []unit.Unit{},
		version:   nil,
		bids:      map[string]pkg.Set{},
	}
}

type FakeRegistry struct {
	// Not all methods of required by the Registry interface are implemented
	// by the TestRegistry. Any calls to these unimplemented methods will
	// result in a panic.
	Registry
	sync.RWMutex

	machines  []machine.MachineState
	jobStates map[string]*unit.UnitState
	jobs      map[string]job.Job
	units     []unit.Unit
	version   *semver.Version
	bids      map[string]pkg.Set
}

func (f *FakeRegistry) SetMachines(machines []machine.MachineState) {
	f.Lock()
	defer f.Unlock()

	f.machines = machines
}

func (f *FakeRegistry) SetJobs(jobs []job.Job) {
	f.Lock()
	defer f.Unlock()

	f.jobs = make(map[string]job.Job, len(jobs))
	for _, j := range jobs {
		f.jobs[j.Name] = j
	}
}

func (f *FakeRegistry) SetUnitStates(jobStates map[string]*unit.UnitState) {
	f.Lock()
	defer f.Unlock()

	f.jobStates = jobStates
}

func (f *FakeRegistry) SetUnits(units []unit.Unit) {
	f.Lock()
	defer f.Unlock()

	f.units = units
}

func (f *FakeRegistry) SetLatestVersion(v semver.Version) {
	f.Lock()
	defer f.Unlock()

	f.version = &v
}

func (f *FakeRegistry) Machines() ([]machine.MachineState, error) {
	f.RLock()
	defer f.RUnlock()

	return f.machines, nil
}

func (f *FakeRegistry) Jobs() ([]job.Job, error) {
	f.RLock()
	defer f.RUnlock()

	var sorted sort.StringSlice
	for _, j := range f.jobs {
		sorted = append(sorted, j.Name)
	}
	sorted.Sort()

	jobs := make([]job.Job, 0, len(f.jobs))
	for _, jName := range sorted {
		jobs = append(jobs, f.jobs[jName])
	}

	return jobs, nil
}

func (f *FakeRegistry) Job(name string) (*job.Job, error) {
	f.RLock()
	defer f.RUnlock()

	j, ok := f.jobs[name]
	if !ok {
		return nil, nil
	}

	j.UnitState = f.jobStates[name]
	return &j, nil
}

func (f *FakeRegistry) CreateJob(j *job.Job) error {
	f.Lock()
	defer f.Unlock()

	_, ok := f.jobs[j.Name]
	if ok {
		return errors.New("Job already exists")
	}

	f.jobs[j.Name] = *j
	return nil
}

func (f *FakeRegistry) DestroyJob(name string) error {
	f.Lock()
	defer f.Unlock()

	delete(f.jobs, name)
	return nil
}

func (f *FakeRegistry) Bids(jo *job.JobOffer) (pkg.Set, error) {
	f.RLock()
	defer f.RUnlock()

	return f.bids[jo.Job.Name], nil
}

func (f *FakeRegistry) SubmitJobBid(jName, machID string) {
	f.Lock()
	defer f.Unlock()

	_, ok := f.bids[jName]
	if !ok {
		f.bids[jName] = pkg.NewUnsafeSet()
	}
	f.bids[jName].Add(machID)
}

func (f *FakeRegistry) SetJobTargetState(name string, target job.JobState) error {
	f.Lock()
	defer f.Unlock()

	j, ok := f.jobs[name]

	if !ok {
		return errors.New("job does not exist")
	}

	j.TargetState = target
	f.jobs[name] = j

	return nil
}

func (f *FakeRegistry) ScheduleJob(name string, machID string) error {
	f.Lock()
	defer f.Unlock()

	j, ok := f.jobs[name]

	if !ok {
		return errors.New("job does not exist")
	}

	j.TargetMachineID = machID
	f.jobs[name] = j

	return nil
}

func (f *FakeRegistry) SaveUnitState(jobName string, unitState *unit.UnitState) {
	f.Lock()
	defer f.Unlock()

	f.jobStates[jobName] = unitState
}

func (f *FakeRegistry) RemoveUnitState(jobName string) error {
	delete(f.jobStates, jobName)
	return nil
}

func (f *FakeRegistry) LatestVersion() (*semver.Version, error) {
	f.RLock()
	defer f.RUnlock()

	return f.version, nil
}

func (f *FakeRegistry) JobHeartbeat(jobName, agentMachID string, ttl time.Duration) error {
	return nil
}

func (f *FakeRegistry) ClearJobHeartbeat(string) {}
<|endoftext|>


<commit_before>package registry

import (
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
)

type Files []*File

func (fs *Files) Add(f *File) {
	for _, c := range *fs {
		if c.Name == f.Name {
			return
		}
	}
	(*fs) = append((*fs), f)
}

func (fs *Files) Get(key string) (*File, bool) {
	for _, c := range *fs {
		if c.Name == key {
			return c, true
		}
	}
	return nil, false
}

type File struct {
	File     *descriptor.FileDescriptorProto
	Name     string
	Messages []*Message
	Package  string
	Registry *Registry
	Options  map[string]string
}

func NewFile(f *descriptor.FileDescriptorProto, r *Registry) *File {
	file := &File{
		File:     f,
		Name:     *f.Name,
		Package:  f.GetPackage(),
		Registry: r,
	}

	// isExtension := false
	var messages []*Message

	for j, m := range f.MessageType {
		// for _, e := range f.Extension {
		// if *e.TypeName == fmt.Sprintf(".%s.%s", file.Package, *m.Name) {
		// 	isExtension = true

		// 	break
		// }
		// }

		// if !isExtension {
		messages = append(messages, NewMessage(m, file, j))
		// }

		// isExtension = false
	}

	file.Messages = messages

	return file
}
<commit_msg>fixed panic<commit_after>package registry

import (
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
)

type Files []*File

func (fs *Files) Add(f *File) {
	for _, c := range *fs {
		if c.Name == f.Name {
			return
		}
	}
	(*fs) = append((*fs), f)
}

func (fs *Files) Get(key string) (*File, bool) {
	for _, c := range *fs {
		if c.Name == key {
			return c, true
		}
	}
	return nil, false
}

type File struct {
	File     *descriptor.FileDescriptorProto
	Name     string
	Messages []*Message
	Package  string
	Registry *Registry
	Options  map[string]string
}

func NewFile(f *descriptor.FileDescriptorProto, r *Registry) *File {
	file := &File{
		File:     f,
		Name:     *f.Name,
		Package:  f.GetPackage(),
		Registry: r,
	}

	var messages []*Message

	for j, m := range f.MessageType {
		messages = append(messages, NewMessage(m, file, j))
	}

	file.Messages = messages

	return file
}
<|endoftext|>


<commit_before>// Copyright 2014 Ardan Studios. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE handle.

// Package work manages a pool of routines to perform work.
package work

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const (
	addRoutine = 1
	rmvRoutine = 3
)

// ErrorInvalidMinRoutines is the error for the invalid minRoutine parameter.
var ErrorInvalidMinRoutines = errors.New("Invalid minimum number of routines")

// ErrorInvalidIdleTime is the error for the invalid idle time parameter.
var ErrorInvalidIdleTime = errors.New("Invalid duration for idle time")

// ErrorInvalidStatTime is the error for the invalid stat time parameter.
var ErrorInvalidStatTime = errors.New("Invalid duration for stat time")

// Worker must be implemented by types that want to use
// this worker processes.
type Worker interface {
	Work(id int)
}

// Work provides a pool of routines that can execute any Worker
// tasks that are submitted.
type Work struct {
	minRoutines int            // Minumum number of routines always in the pool.
	idleTime    time.Duration  // Time for routines to die on idle.
	statTime    time.Duration  // Time to display stats.
	counter     int            // Maintains a running total number of routines ever created.
	tasks       chan Worker    // Unbuffered channel that work is sent into.
	control     chan int       // Unbuffered channel that work for the manager is send into.
	kill        chan struct{}  // Unbuffered channel to signal for a goroutine to die.
	shutdown    chan struct{}  // Closed when the Work pool is being shutdown.
	wg          sync.WaitGroup // Manages the number of routines for shutdown.
	routines    int64          // Number of routines
	active      int64          // Active number of routines in the work pool.
	pending     int64          // Pending number of routines waiting to submit work.
}

// manager controls changes to the work pool including stats
// and shutting down.
func (w *Work) manager() {
	w.wg.Add(1)

	go func() {
		log.Println("Work : manager : Started")

		// Create a timer to run stats.
		var stats <-chan time.Time
		timer := time.NewTimer(w.statTime)
		stats = timer.C

		for {
			select {
			case <-w.shutdown:
				// Capture the current number of routines.
				routines := int(atomic.LoadInt64(&w.routines))

				// Send a kill to all the existing routines.
				for i := 0; i < routines; i++ {
					w.kill <- struct{}{}
				}

				// Decrement the waitgroup and kill the manager.
				w.wg.Done()
				return

			case c := <-w.control:
				switch c {
				case addRoutine:
					log.Println("Work : manager : Info : Add Routine")

					// Capture a unique id.
					w.counter++

					// Add to the counts.
					w.wg.Add(1)
					atomic.AddInt64(&w.routines, 1)

					// Create the routine.
					go w.work(w.counter)

				case rmvRoutine:
					log.Println("Work : manager : Info : Remove Routine")

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&w.routines))

					// Are there routines to remove.
					if routines <= w.minRoutines {
						break
					}

					// Send a kill signal to remove a routine.
					w.kill <- struct{}{}
				}

			case <-stats:
				// Capture the stats.
				routines := atomic.LoadInt64(&w.routines)
				pending := atomic.LoadInt64(&w.pending)
				active := atomic.LoadInt64(&w.active)

				// Display the stats.
				fmt.Printf("Work : manager : Stats : G[%d] P[%d] A[%d]\n", routines, pending, active)

				// Reset the clock.
				timer.Reset(w.statTime)
			}
		}
	}()
}

// New creates a new Worker.
func New(minRoutines int, idleTime time.Duration, statTime time.Duration) (*Work, error) {
	if minRoutines < 0 {
		return nil, ErrorInvalidMinRoutines
	}

	if idleTime < time.Minute {
		return nil, ErrorInvalidIdleTime
	}

	if statTime < time.Millisecond {
		return nil, ErrorInvalidStatTime
	}

	w := Work{
		minRoutines: minRoutines,
		idleTime:    idleTime,
		statTime:    statTime,
		tasks:       make(chan Worker),
		control:     make(chan int),
		kill:        make(chan struct{}),
		shutdown:    make(chan struct{}),
	}

	// Start the manager.
	w.manager()

	// Add the routines.
	w.Add(minRoutines)

	return &w, nil
}

// Add creates routines to process work or sets a count for
// routines to terminate.
func (w *Work) Add(routines int) {
	if routines == 0 {
		return
	}

	cmd := addRoutine
	if routines < 0 {
		routines = routines * -1
		cmd = rmvRoutine
	}

	for i := 0; i < routines; i++ {
		w.control <- cmd
	}
}

// work performs the users work and keeps stats.
func (w *Work) work(id int) {
	// Create a timer to track idle time.
	var idle <-chan time.Time
	var timer *time.Timer

	// Set the timer for routines about the min mark.
	if id > w.minRoutines {
		timer = time.NewTimer(w.idleTime)
		idle = timer.C
	}

done:
	for {
		select {
		case t := <-w.tasks:
			atomic.AddInt64(&w.active, 1)
			{
				// Perform the work.
				t.Work(id)
			}
			atomic.AddInt64(&w.active, -1)

		case <-w.kill:
			break done

		case <-idle:
			break done
		}

		// If this goroutine can die on idle time
		// then reset the timer.
		if timer != nil {
			timer.Reset(w.idleTime)
		}
	}

	// Decrement the counts.
	atomic.AddInt64(&w.routines, -1)
	w.wg.Done()

	log.Println("Work : gr : Info : Shutdown")
}

// Run wait for the goroutine pool to take the work
// to be executed.
func (w *Work) Run(work Worker) {
	atomic.AddInt64(&w.pending, 1)
	{
		w.tasks <- work
	}
	atomic.AddInt64(&w.pending, -1)
}

// Shutdown waits for all the workers to finish.
func (w *Work) Shutdown() {
	close(w.shutdown)
	w.wg.Wait()
}
<commit_msg>Updated logging.<commit_after>// Copyright 2014 Ardan Studios. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE handle.

// Package work manages a pool of routines to perform work.
package work

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const (
	addRoutine = 1
	rmvRoutine = 3
)

// ErrorInvalidMinRoutines is the error for the invalid minRoutine parameter.
var ErrorInvalidMinRoutines = errors.New("Invalid minimum number of routines")

// ErrorInvalidIdleTime is the error for the invalid idle time parameter.
var ErrorInvalidIdleTime = errors.New("Invalid duration for idle time")

// ErrorInvalidStatTime is the error for the invalid stat time parameter.
var ErrorInvalidStatTime = errors.New("Invalid duration for stat time")

// Worker must be implemented by types that want to use
// this worker processes.
type Worker interface {
	Work(id int)
}

// Work provides a pool of routines that can execute any Worker
// tasks that are submitted.
type Work struct {
	minRoutines int            // Minumum number of routines always in the pool.
	idleTime    time.Duration  // Time for routines to die on idle.
	statTime    time.Duration  // Time to display stats.
	counter     int            // Maintains a running total number of routines ever created.
	tasks       chan Worker    // Unbuffered channel that work is sent into.
	control     chan int       // Unbuffered channel that work for the manager is send into.
	kill        chan struct{}  // Unbuffered channel to signal for a goroutine to die.
	shutdown    chan struct{}  // Closed when the Work pool is being shutdown.
	wg          sync.WaitGroup // Manages the number of routines for shutdown.
	routines    int64          // Number of routines
	active      int64          // Active number of routines in the work pool.
	pending     int64          // Pending number of routines waiting to submit work.
}

// manager controls changes to the work pool including stats
// and shutting down.
func (w *Work) manager() {
	w.wg.Add(1)

	go func() {
		log.Println("Work : manager : Started")

		// Create a timer to run stats.
		var stats <-chan time.Time
		timer := time.NewTimer(w.statTime)
		stats = timer.C

		for {
			select {
			case <-w.shutdown:
				// Capture the current number of routines.
				routines := int(atomic.LoadInt64(&w.routines))

				// Send a kill to all the existing routines.
				for i := 0; i < routines; i++ {
					w.kill <- struct{}{}
				}

				// Decrement the waitgroup and kill the manager.
				w.wg.Done()
				return

			case c := <-w.control:
				switch c {
				case addRoutine:
					log.Println("Work : manager : Info : Add Routine")

					// Capture a unique id.
					w.counter++

					// Add to the counts.
					w.wg.Add(1)
					atomic.AddInt64(&w.routines, 1)

					// Create the routine.
					go w.work(w.counter)

				case rmvRoutine:
					log.Println("Work : manager : Info : Remove Routine")

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&w.routines))

					// Are there routines to remove.
					if routines <= w.minRoutines {
						log.Println("Work : manager : Info : Remove Routine Cancelled")
						break
					}

					// Send a kill signal to remove a routine.
					w.kill <- struct{}{}
				}

			case <-stats:
				// Capture the stats.
				routines := atomic.LoadInt64(&w.routines)
				pending := atomic.LoadInt64(&w.pending)
				active := atomic.LoadInt64(&w.active)

				// Display the stats.
				fmt.Printf("Work : manager : Stats : G[%d] P[%d] A[%d]\n", routines, pending, active)

				// Reset the clock.
				timer.Reset(w.statTime)
			}
		}
	}()
}

// New creates a new Worker.
func New(minRoutines int, idleTime time.Duration, statTime time.Duration) (*Work, error) {
	if minRoutines < 0 {
		return nil, ErrorInvalidMinRoutines
	}

	if idleTime < time.Minute {
		return nil, ErrorInvalidIdleTime
	}

	if statTime < time.Millisecond {
		return nil, ErrorInvalidStatTime
	}

	w := Work{
		minRoutines: minRoutines,
		idleTime:    idleTime,
		statTime:    statTime,
		tasks:       make(chan Worker),
		control:     make(chan int),
		kill:        make(chan struct{}),
		shutdown:    make(chan struct{}),
	}

	// Start the manager.
	w.manager()

	// Add the routines.
	w.Add(minRoutines)

	return &w, nil
}

// Add creates routines to process work or sets a count for
// routines to terminate.
func (w *Work) Add(routines int) {
	if routines == 0 {
		return
	}

	cmd := addRoutine
	if routines < 0 {
		routines = routines * -1
		cmd = rmvRoutine
	}

	for i := 0; i < routines; i++ {
		w.control <- cmd
	}
}

// work performs the users work and keeps stats.
func (w *Work) work(id int) {
	// Create a timer to track idle time.
	var idle <-chan time.Time
	var timer *time.Timer

	// Set the timer for routines about the min mark.
	if id > w.minRoutines {
		timer = time.NewTimer(w.idleTime)
		idle = timer.C
	}

done:
	for {
		select {
		case t := <-w.tasks:
			atomic.AddInt64(&w.active, 1)
			{
				// Perform the work.
				t.Work(id)
			}
			atomic.AddInt64(&w.active, -1)

		case <-w.kill:
			break done

		case <-idle:
			break done
		}

		// If this goroutine can die on idle time
		// then reset the timer.
		if timer != nil {
			timer.Reset(w.idleTime)
		}
	}

	// Decrement the counts.
	atomic.AddInt64(&w.routines, -1)
	w.wg.Done()

	log.Println("Work : gr : Info : Shutdown")
}

// Run wait for the goroutine pool to take the work
// to be executed.
func (w *Work) Run(work Worker) {
	atomic.AddInt64(&w.pending, 1)
	{
		w.tasks <- work
	}
	atomic.AddInt64(&w.pending, -1)
}

// Shutdown waits for all the workers to finish.
func (w *Work) Shutdown() {
	close(w.shutdown)
	w.wg.Wait()
}
<|endoftext|>


<commit_before>package gokalman

import (
	"fmt"
	"math"
	"os"
	"strings"
	"time"
)

// Exporter defines an export interface.
type Exporter interface {
	Write(Estimate) error
	Close() error
}

// CSVExporter returns a new CSV exporter.
type CSVExporter struct {
	delimiter string
	hdlr      *os.File
}

// Close closes the file.
func (e CSVExporter) Close() (err error) {
	err = e.WriteRawLn(fmt.Sprintf("# Closing date (UTC): %s\n", time.Now().UTC()))
	if err != nil {
		return
	}
	return e.hdlr.Close()
}

// Write writes the estimate to the CSV file.
func (e CSVExporter) Write(est Estimate) error {
	r, _ := est.State().Dims()
	vals := make([]string, r*3)
	for i := 0; i < r*3; i += 3 {
		vals[i] = fmt.Sprintf("%f", est.State().At(i/3, 0))
		covar := 2 * math.Sqrt(est.Covariance().At(i/3, i/3))
		vals[i+1] = fmt.Sprintf("%f", covar)
		vals[i+2] = fmt.Sprintf("%f", -1*covar)
	}
	_, err := e.hdlr.WriteString(strings.Join(vals, e.delimiter) + "\n")
	return err
}

// WriteRawLn writes a raw line to the CSV file.
func (e CSVExporter) WriteRawLn(s string) error {
	_, err := e.hdlr.WriteString(s + "\n")
	return err
}

// NewCSVExporter initializes a new CSV export.
func NewCSVExporter(headers []string, filepath, filename string) (e *CSVExporter, err error) {
	f, err := os.Create(fmt.Sprintf("%s/%s", filepath, filename))
	if err != nil {
		return
	}
	delimiter := ","
	// Header
	hdr := make([]string, len(headers)*3)
	for i := 0; i < len(headers)*3; i += 3 {
		hdr[i] = headers[i/3]
		hdr[i+1] = hdr[i] + "+2s"
		hdr[i+2] = hdr[i] + "-2s"
	}
	f.WriteString(fmt.Sprintf("# Creation date (UTC): %s\n%s\n", time.Now(), strings.Join(hdr, delimiter)))
	e = &CSVExporter{delimiter, f}
	return
}
<commit_msg>Minor change to exporter<commit_after>package gokalman

import (
	"fmt"
	"math"
	"os"
	"strings"
	"time"
)

// Exporter defines an export interface.
type Exporter interface {
	Write(Estimate) error
	Close() error
}

// CSVExporter returns a new CSV exporter.
type CSVExporter struct {
	covarBound float64
	delimiter  string
	hdlr       *os.File
}

// Close closes the file.
func (e CSVExporter) Close() (err error) {
	err = e.WriteRawLn(fmt.Sprintf("# Closing date (UTC): %s\n", time.Now().UTC()))
	if err != nil {
		return
	}
	return e.hdlr.Close()
}

// Write writes the estimate to the CSV file.
func (e CSVExporter) Write(est Estimate) error {
	r, _ := est.State().Dims()
	vals := make([]string, r*3)
	for i := 0; i < r*3; i += 3 {
		vals[i] = fmt.Sprintf("%f", est.State().At(i/3, 0))
		covar := 2 * math.Sqrt(est.Covariance().At(i/3, i/3))
		vals[i+1] = fmt.Sprintf("%f", covar)
		vals[i+2] = fmt.Sprintf("%f", -1*covar)
	}
	_, err := e.hdlr.WriteString(strings.Join(vals, e.delimiter) + "\n")
	return err
}

// WriteRawLn writes a raw line to the CSV file.
func (e CSVExporter) WriteRawLn(s string) error {
	_, err := e.hdlr.WriteString(s + "\n")
	return err
}

// NewCustomCSVExporter initializes a new CSV export.
func NewCustomCSVExporter(headers []string, filepath, filename string, covarBound float64) (e *CSVExporter, err error) {
	f, err := os.Create(fmt.Sprintf("%s/%s", filepath, filename))
	if err != nil {
		return
	}
	delimiter := ","
	// Header
	hdr := make([]string, len(headers)*3)
	bhdr := fmt.Sprintf("%.0fs", covarBound)
	for i := 0; i < len(headers)*3; i += 3 {
		hdr[i] = headers[i/3]
		hdr[i+1] = hdr[i] + "+" + bhdr
		hdr[i+2] = hdr[i] + "-" + bhdr
	}
	f.WriteString(fmt.Sprintf("# Creation date (UTC): %s\n%s\n", time.Now(), strings.Join(hdr, delimiter)))
	e = &CSVExporter{covarBound, delimiter, f}
	return
}

// NewCSVExporter initializes a new CSV export.
func NewCSVExporter(headers []string, filepath, filename string) (e *CSVExporter, err error) {
	return NewCustomCSVExporter(headers, filepath, filename, 2)
}
<|endoftext|>


<commit_before>// Copyright 2018, OpenCensus 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 aws

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/xray"
	"github.com/aws/aws-sdk-go/service/xray/xrayiface"
	"go.opencensus.io/trace"
)

// OnExport structure passed when a root segment is published
type OnExport struct {
	// TraceID holds the raw aws traceID e.g. 1-581cf771-a006649127e371903a2de979
	TraceID string
}

const (
	// defaultInterval - segments will be published at this frequency
	defaultInterval = time.Second

	// maxBufferSize - maximum number of segments that can be published in a single call
	maxBufferSize = 50
)

type config struct {
	region     string                // aws region; for sending to another region's xray
	output     io.Writer             // output error messages
	api        xrayiface.XRayAPI     // use specific api instance
	onExport   func(export OnExport) // callback on publish
	origin     origin                // origin of span
	service    *service              // contains embedded version info
	interval   time.Duration         // interval spans are published to aws
	bufferSize int                   // bufSize represents max number of spans before forcing publish
	blacklist  []*regexp.Regexp
}

type Option interface {
	apply(c *config)
}

type optionFunc func(c *config)

func (fn optionFunc) apply(c *config) {
	fn(c)
}

// WithRegion - optional aws region to send xray messages to
func WithRegion(region string) Option {
	return optionFunc(func(c *config) {
		c.region = region
	})
}

// WithOutput - optional writer for error messages
func WithOutput(w io.Writer) Option {
	return optionFunc(func(c *config) {
		c.output = w
	})
}

// WithAPI - optional manually constructed api instance
func WithAPI(api xrayiface.XRayAPI) Option {
	return optionFunc(func(c *config) {
		c.api = api
	})
}

// WithOnExport - function to be called when spans are published to AWS; useful
// if you would like the traceID used for AWS
func WithOnExport(fn func(OnExport)) Option {
	return optionFunc(func(c *config) {
		c.onExport = fn
	})
}

// WithOrigin - specifies the aws origin of the span; currently supported are
// OriginEC2, OriginECS, and OriginEB
func WithOrigin(origin origin) Option {
	return optionFunc(func(c *config) {
		c.origin = origin
	})
}

// WithVersion - specifies the version of the application running
func WithVersion(version string) Option {
	return optionFunc(func(c *config) {
		c.service = &service{
			Version: version,
		}
	})
}

// WithInterval - specifies longest time before buffered spans are published; defaults to 1s
func WithInterval(interval time.Duration) Option {
	return optionFunc(func(c *config) {
		c.interval = interval
	})
}

// WithBufferSize - specifies the maximum number of spans to buffer before publishing them; defaults to 100
func WithBufferSize(bufferSize int) Option {
	return optionFunc(func(c *config) {
		c.bufferSize = bufferSize
	})
}

// WithBlacklist filters out the spans those names match one of the given regexps.
func WithBlacklist(blacklist []*regexp.Regexp) Option {
	return optionFunc(func(c *config) {
		c.blacklist = blacklist
	})
}

// Exporter is an implementation of trace.Exporter that uploads spans to AWS XRay
type Exporter struct {
	api      xrayiface.XRayAPI
	onExport func(export OnExport)
	logger   *log.Logger
	service  *service
	origin   string
	wg       sync.WaitGroup // wg holds number of publishers in flight

	ctx    context.Context    // ctx cancels the child goroutine
	cancel context.CancelFunc // cancels the child goroutine; idempotent
	done   chan struct{}      // done returns immediately once the child goroutine has finished

	mutex     sync.Mutex        // mutex protects offset, buffer, and closed
	offset    int               // offset holds position of next SpanData within queue
	buffer    []*trace.SpanData // buffer of spans not yet published
	closed    bool              // indicates Exporter is closed.  any additional spans will be dropped
	blacklist []*regexp.Regexp
}

var (
	// extract region from aws availability zone
	reRegion = regexp.MustCompile(`([^\-]+-[^\-]+-\d+)`)

	// client exposed to simplify testing
	clientDo = http.DefaultClient.Do
)

// lookupRegionFromMetaData attempts to determine region from aws metadata endpoint
func lookupRegionFromMetaData() (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	req, _ := http.NewRequest(http.MethodGet, "http://169.254.169.254/latest/meta-data/placement/availability-zone", nil)
	req = req.WithContext(ctx)

	resp, err := clientDo(req)
	if err != nil {
		return "", fmt.Errorf("unable lookup region via http://169.254.169.254/latest/meta-data/placement/availability-zone - %v", err)
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("unable to read contents from http://169.254.169.254/latest/meta-data/placement/availability-zone - %v", err)
	}

	segments := reRegion.FindSubmatch(data)
	if len(segments) == 0 {
		return "", fmt.Errorf("unable to determine aws region from availability zone, %v", string(data))
	}

	return string(segments[1]), nil
}

func lookupRegion() (string, error) {
	region := os.Getenv("AWS_DEFAULT_REGION")
	if region == "" {
		region = os.Getenv("AWS_REGION")
	}
	if region == "" {
		v, err := lookupRegionFromMetaData()
		if err != nil {
			return "", err
		}
		region = v
	}

	return region, nil
}

// makeApi constructs an instance of the aws xray api
func makeApi(region string) (*xray.XRay, error) {
	if region == "" {
		v, err := lookupRegion()
		if err != nil {
			return nil, err
		}
		region = v
	}

	// Region value will instruct the SDK where to make service API requests to. If is
	// not provided in the environment the region must be provided before a service
	// client request is made.
	//
	// https://docs.aws.amazon.com/sdk-for-go/api/aws/session/
	config := aws.Config{Region: aws.String(region)}
	s, err := session.NewSession(&config)
	if err != nil {
		return nil, err
	}
	return xray.New(s), nil
}

// buildConfig from options
func buildConfig(opts ...Option) (config, error) {
	var c config
	for _, opt := range opts {
		opt.apply(&c)
	}

	if c.output == nil {
		c.output = os.Stderr
	}
	if c.onExport == nil {
		c.onExport = func(export OnExport) {}
	}
	if c.interval <= 0 {
		c.interval = defaultInterval
	}
	if c.bufferSize <= 0 || c.bufferSize > maxBufferSize {
		c.bufferSize = maxBufferSize
	}
	if c.api == nil {
		api, err := makeApi(c.region)
		if err != nil {
			return config{}, err
		}
		c.api = api
	}

	return c, nil
}

// NewExporter returns an implementation of trace.Exporter that uploads spans
// to AWS X-Ray
func NewExporter(opts ...Option) (*Exporter, error) {
	c, err := buildConfig(opts...)
	if err != nil {
		return nil, err
	}

	var (
		logger = log.New(c.output, "XRAY ", log.LstdFlags)
		buffer = make([]*trace.SpanData, c.bufferSize)
		done   = make(chan struct{})
	)

	ctx, cancel := context.WithCancel(context.Background())

	exporter := &Exporter{
		api:      c.api,
		onExport: c.onExport,
		service:  c.service,
		logger:   logger,

		ctx:    ctx,
		cancel: cancel,
		done:   done,

		buffer:    buffer,
		origin:    string(c.origin),
		blacklist: c.blacklist,
	}
	go exporter.publishAtInterval(c.interval)

	return exporter, nil
}

func (e *Exporter) isBlacklisted(span *trace.SpanData) bool {
	for _, blacklistEntry := range e.blacklist {
		if blacklistEntry.MatchString(span.Name) {
			return true
		}
	}
	return false
}

func (e *Exporter) makeInput(spans []*trace.SpanData) (xray.PutTraceSegmentsInput, []string) {
	var (
		traceIDs []string
		docs     []*string
		w        = borrow()
	)
	defer release(w)

	for _, span := range spans {
		if e.isBlacklisted(span) {
			continue
		}
		var segment = e.makeSegment(span)
		if segment.ParentID == "" {
			traceIDs = append(traceIDs, segment.TraceID)
		}

		w.Reset()
		if err := w.Encode(segment); err != nil {
			e.logger.Printf("unable to encode segment, %v", err)
			continue
		}

		docs = append(docs, aws.String(w.String()))
	}

	input := xray.PutTraceSegmentsInput{
		TraceSegmentDocuments: docs,
	}
	return input, traceIDs
}

func (e *Exporter) publish(spans []*trace.SpanData) {
	defer e.wg.Done()

	var (
		input, traceIDs = e.makeInput(spans)
	)

	for attempt := 0; attempt < 3; attempt++ {
		_, err := e.api.PutTraceSegments(&input)
		if err == nil {
			for _, traceID := range traceIDs {
				go e.onExport(OnExport{TraceID: traceID})
			}
			return
		}

		e.logger.Printf("attempt %v failed to publish spans, %v\n", attempt, err)
		time.Sleep(500 * time.Millisecond)
	}

	e.logger.Println("all attempts to publish span failed.  giving up.")
}

func (e *Exporter) flush() {
	if e.offset == 0 {
		return
	}

	var spans []*trace.SpanData
	spans = append(spans, e.buffer[0:e.offset]...)

	// reset buffer
	e.offset = 0

	e.wg.Add(1)
	go e.publish(spans)
}

func (e *Exporter) publishAtInterval(interval time.Duration) {
	defer close(e.done)

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-e.ctx.Done():
			return

		case <-ticker.C:
			e.mutex.Lock()
			e.flush()
			e.mutex.Unlock()
		}
	}
}

// Close this exporter and publish any spans that may have been buffered
func (e *Exporter) Close() error {
	e.cancel()
	<-e.done // wait for goroutine to shut down

	e.mutex.Lock()
	if !e.closed {
		e.closed = true
		e.flush()
		e.wg.Wait()
	}
	e.mutex.Unlock()

	return nil
}

func (e *Exporter) Flush() {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	e.flush()
}

// ExportSpan exports a span to AWS X-Ray
func (e *Exporter) ExportSpan(s *trace.SpanData) {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	if e.closed {
		e.logger.Println("ExportSpan called on a closed Exporter.  SpanData will be dropped.")
		return
	}

	e.buffer[e.offset] = s
	e.offset++

	if e.offset == cap(e.buffer) {
		e.flush()
	}
}

func (e *Exporter) makeSegment(span *trace.SpanData) segment {
	var (
		s = rawSegment(span)
	)

	if isRootSpan := span.ParentSpanID == zeroSpanID; isRootSpan {
		s.Origin = e.origin
		s.Service = e.service

	} else {
		s.Type = "subsegment"
	}

	return s
}
<commit_msg>Don't send empty traces. (#27)<commit_after>// Copyright 2018, OpenCensus 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 aws

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/xray"
	"github.com/aws/aws-sdk-go/service/xray/xrayiface"
	"go.opencensus.io/trace"
)

// OnExport structure passed when a root segment is published
type OnExport struct {
	// TraceID holds the raw aws traceID e.g. 1-581cf771-a006649127e371903a2de979
	TraceID string
}

const (
	// defaultInterval - segments will be published at this frequency
	defaultInterval = time.Second

	// maxBufferSize - maximum number of segments that can be published in a single call
	maxBufferSize = 50
)

type config struct {
	region     string                // aws region; for sending to another region's xray
	output     io.Writer             // output error messages
	api        xrayiface.XRayAPI     // use specific api instance
	onExport   func(export OnExport) // callback on publish
	origin     origin                // origin of span
	service    *service              // contains embedded version info
	interval   time.Duration         // interval spans are published to aws
	bufferSize int                   // bufSize represents max number of spans before forcing publish
	blacklist  []*regexp.Regexp
}

type Option interface {
	apply(c *config)
}

type optionFunc func(c *config)

func (fn optionFunc) apply(c *config) {
	fn(c)
}

// WithRegion - optional aws region to send xray messages to
func WithRegion(region string) Option {
	return optionFunc(func(c *config) {
		c.region = region
	})
}

// WithOutput - optional writer for error messages
func WithOutput(w io.Writer) Option {
	return optionFunc(func(c *config) {
		c.output = w
	})
}

// WithAPI - optional manually constructed api instance
func WithAPI(api xrayiface.XRayAPI) Option {
	return optionFunc(func(c *config) {
		c.api = api
	})
}

// WithOnExport - function to be called when spans are published to AWS; useful
// if you would like the traceID used for AWS
func WithOnExport(fn func(OnExport)) Option {
	return optionFunc(func(c *config) {
		c.onExport = fn
	})
}

// WithOrigin - specifies the aws origin of the span; currently supported are
// OriginEC2, OriginECS, and OriginEB
func WithOrigin(origin origin) Option {
	return optionFunc(func(c *config) {
		c.origin = origin
	})
}

// WithVersion - specifies the version of the application running
func WithVersion(version string) Option {
	return optionFunc(func(c *config) {
		c.service = &service{
			Version: version,
		}
	})
}

// WithInterval - specifies longest time before buffered spans are published; defaults to 1s
func WithInterval(interval time.Duration) Option {
	return optionFunc(func(c *config) {
		c.interval = interval
	})
}

// WithBufferSize - specifies the maximum number of spans to buffer before publishing them; defaults to 100
func WithBufferSize(bufferSize int) Option {
	return optionFunc(func(c *config) {
		c.bufferSize = bufferSize
	})
}

// WithBlacklist filters out the spans those names match one of the given regexps.
func WithBlacklist(blacklist []*regexp.Regexp) Option {
	return optionFunc(func(c *config) {
		c.blacklist = blacklist
	})
}

// Exporter is an implementation of trace.Exporter that uploads spans to AWS XRay
type Exporter struct {
	api      xrayiface.XRayAPI
	onExport func(export OnExport)
	logger   *log.Logger
	service  *service
	origin   string
	wg       sync.WaitGroup // wg holds number of publishers in flight

	ctx    context.Context    // ctx cancels the child goroutine
	cancel context.CancelFunc // cancels the child goroutine; idempotent
	done   chan struct{}      // done returns immediately once the child goroutine has finished

	mutex     sync.Mutex        // mutex protects offset, buffer, and closed
	offset    int               // offset holds position of next SpanData within queue
	buffer    []*trace.SpanData // buffer of spans not yet published
	closed    bool              // indicates Exporter is closed.  any additional spans will be dropped
	blacklist []*regexp.Regexp
}

var (
	// extract region from aws availability zone
	reRegion = regexp.MustCompile(`([^\-]+-[^\-]+-\d+)`)

	// client exposed to simplify testing
	clientDo = http.DefaultClient.Do
)

// lookupRegionFromMetaData attempts to determine region from aws metadata endpoint
func lookupRegionFromMetaData() (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	req, _ := http.NewRequest(http.MethodGet, "http://169.254.169.254/latest/meta-data/placement/availability-zone", nil)
	req = req.WithContext(ctx)

	resp, err := clientDo(req)
	if err != nil {
		return "", fmt.Errorf("unable lookup region via http://169.254.169.254/latest/meta-data/placement/availability-zone - %v", err)
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("unable to read contents from http://169.254.169.254/latest/meta-data/placement/availability-zone - %v", err)
	}

	segments := reRegion.FindSubmatch(data)
	if len(segments) == 0 {
		return "", fmt.Errorf("unable to determine aws region from availability zone, %v", string(data))
	}

	return string(segments[1]), nil
}

func lookupRegion() (string, error) {
	region := os.Getenv("AWS_DEFAULT_REGION")
	if region == "" {
		region = os.Getenv("AWS_REGION")
	}
	if region == "" {
		v, err := lookupRegionFromMetaData()
		if err != nil {
			return "", err
		}
		region = v
	}

	return region, nil
}

// makeApi constructs an instance of the aws xray api
func makeApi(region string) (*xray.XRay, error) {
	if region == "" {
		v, err := lookupRegion()
		if err != nil {
			return nil, err
		}
		region = v
	}

	// Region value will instruct the SDK where to make service API requests to. If is
	// not provided in the environment the region must be provided before a service
	// client request is made.
	//
	// https://docs.aws.amazon.com/sdk-for-go/api/aws/session/
	config := aws.Config{Region: aws.String(region)}
	s, err := session.NewSession(&config)
	if err != nil {
		return nil, err
	}
	return xray.New(s), nil
}

// buildConfig from options
func buildConfig(opts ...Option) (config, error) {
	var c config
	for _, opt := range opts {
		opt.apply(&c)
	}

	if c.output == nil {
		c.output = os.Stderr
	}
	if c.onExport == nil {
		c.onExport = func(export OnExport) {}
	}
	if c.interval <= 0 {
		c.interval = defaultInterval
	}
	if c.bufferSize <= 0 || c.bufferSize > maxBufferSize {
		c.bufferSize = maxBufferSize
	}
	if c.api == nil {
		api, err := makeApi(c.region)
		if err != nil {
			return config{}, err
		}
		c.api = api
	}

	return c, nil
}

// NewExporter returns an implementation of trace.Exporter that uploads spans
// to AWS X-Ray
func NewExporter(opts ...Option) (*Exporter, error) {
	c, err := buildConfig(opts...)
	if err != nil {
		return nil, err
	}

	var (
		logger = log.New(c.output, "XRAY ", log.LstdFlags)
		buffer = make([]*trace.SpanData, c.bufferSize)
		done   = make(chan struct{})
	)

	ctx, cancel := context.WithCancel(context.Background())

	exporter := &Exporter{
		api:      c.api,
		onExport: c.onExport,
		service:  c.service,
		logger:   logger,

		ctx:    ctx,
		cancel: cancel,
		done:   done,

		buffer:    buffer,
		origin:    string(c.origin),
		blacklist: c.blacklist,
	}
	go exporter.publishAtInterval(c.interval)

	return exporter, nil
}

func (e *Exporter) isBlacklisted(span *trace.SpanData) bool {
	for _, blacklistEntry := range e.blacklist {
		if blacklistEntry.MatchString(span.Name) {
			return true
		}
	}
	return false
}

func (e *Exporter) makeInput(spans []*trace.SpanData) (xray.PutTraceSegmentsInput, []string) {
	var (
		traceIDs []string
		docs     []*string
		w        = borrow()
	)
	defer release(w)

	for _, span := range spans {
		if e.isBlacklisted(span) {
			continue
		}
		var segment = e.makeSegment(span)
		if segment.ParentID == "" {
			traceIDs = append(traceIDs, segment.TraceID)
		}

		w.Reset()
		if err := w.Encode(segment); err != nil {
			e.logger.Printf("unable to encode segment, %v", err)
			continue
		}

		docs = append(docs, aws.String(w.String()))
	}

	input := xray.PutTraceSegmentsInput{
		TraceSegmentDocuments: docs,
	}
	return input, traceIDs
}

func (e *Exporter) publish(spans []*trace.SpanData) {
	defer e.wg.Done()

	var (
		input, traceIDs = e.makeInput(spans)
	)
	if len(input.TraceSegmentDocuments) == 0 {
		return
	}

	for attempt := 0; attempt < 3; attempt++ {
		_, err := e.api.PutTraceSegments(&input)
		if err == nil {
			for _, traceID := range traceIDs {
				go e.onExport(OnExport{TraceID: traceID})
			}
			return
		}

		e.logger.Printf("attempt %v failed to publish spans, %v\n", attempt, err)
		time.Sleep(500 * time.Millisecond)
	}

	e.logger.Println("all attempts to publish span failed.  giving up.")
}

func (e *Exporter) flush() {
	if e.offset == 0 {
		return
	}

	var spans []*trace.SpanData
	spans = append(spans, e.buffer[0:e.offset]...)

	// reset buffer
	e.offset = 0

	e.wg.Add(1)
	go e.publish(spans)
}

func (e *Exporter) publishAtInterval(interval time.Duration) {
	defer close(e.done)

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-e.ctx.Done():
			return

		case <-ticker.C:
			e.mutex.Lock()
			e.flush()
			e.mutex.Unlock()
		}
	}
}

// Close this exporter and publish any spans that may have been buffered
func (e *Exporter) Close() error {
	e.cancel()
	<-e.done // wait for goroutine to shut down

	e.mutex.Lock()
	if !e.closed {
		e.closed = true
		e.flush()
		e.wg.Wait()
	}
	e.mutex.Unlock()

	return nil
}

func (e *Exporter) Flush() {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	e.flush()
}

// ExportSpan exports a span to AWS X-Ray
func (e *Exporter) ExportSpan(s *trace.SpanData) {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	if e.closed {
		e.logger.Println("ExportSpan called on a closed Exporter.  SpanData will be dropped.")
		return
	}

	e.buffer[e.offset] = s
	e.offset++

	if e.offset == cap(e.buffer) {
		e.flush()
	}
}

func (e *Exporter) makeSegment(span *trace.SpanData) segment {
	var (
		s = rawSegment(span)
	)

	if isRootSpan := span.ParentSpanID == zeroSpanID; isRootSpan {
		s.Origin = e.origin
		s.Service = e.service

	} else {
		s.Type = "subsegment"
	}

	return s
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/jroimartin/gocui"
	"gopkg.in/alecthomas/kingpin.v1"
)

func main() {
	name := "xtop"
	help := `A top-like tool to monitor responses from a target URL. This tool periodically collects and prints out response statuses and a custom response header received from the url.`

	app := kingpin.New(name, help)
	app.Version("0.0.1")

	var (
		url        = app.Arg("url", "target URL").Required().URL()
		concurrent = app.Flag("concurrent", "number of concurrent requests sent in one batch").Short('c').Default("3").Int()
		header     = app.Flag("header", "custom header name to collect").Short('x').Default("X-Server").String()
	)

	kingpin.MustParse(app.Parse(os.Args[1:]))

	x := NewXTop((*url).String(), *concurrent, *header)
	x.Start()
}

// Pair is a data structure to hold a key/value pair.
type Pair struct {
	Key   string
	Value int
}

// PairList is a slice of Pairs that implements sort.Interface to sort by Value.
type PairList []Pair

func (p PairList) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p PairList) Len() int           { return len(p) }
func (p PairList) Less(i, j int) bool { return p[i].Value > p[j].Value }

// sortMapByValue turns a map into a PairList, then sort and return it.
func sortMapByValue(m map[string]int) PairList {
	p := make(PairList, len(m))
	i := 0
	for k, v := range m {
		p[i] = Pair{k, v}
		i++
	}
	sort.Sort(p)
	return p
}

// appendIfMissing appends a string into a slice if it does not exist yet
func appendIfMissing(a []string, item string) []string {
	for _, v := range a {
		if item == v {
			return a
		}
	}
	return append(a, item)
}

// XTop is the main struct
type XTop struct {
	URL        string
	Concurrent int
	Header     string

	TotalRequestsSent int
	StatusMap         map[string]int
	ServerMap         map[string]int
	G                 *gocui.Gui
}

// NewXTop returns a new XTop instance
func NewXTop(url string, concurrent int, header string) XTop {
	url = strings.ToLower(url)
	if strings.Index(url, "http://") == -1 && strings.Index(url, "https://") == -1 {
		url = "http://" + url
	}

	x := XTop{}
	x.URL = url
	x.Concurrent = concurrent
	x.Header = header
	x.StatusMap = make(map[string]int)
	x.ServerMap = make(map[string]int)

	var err error = nil
	g := gocui.NewGui()
	if err := g.Init(); err != nil {
		panic(err)
	}
	g.FgColor = gocui.ColorGreen
	g.SetLayout(x.layout)
	if err = g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, func(g *gocui.Gui, v *gocui.View) error {
		return gocui.ErrorQuit
	}); err != nil {
		log.Panicln(err)
	}

	x.G = g
	return x
}

// Run the app
func (r *XTop) Start() {
	defer r.G.Close()

	// One thread (that spawns several goroutines) to send requests and collect data
	go r.run()

	// One thread to periodically print out results to the screen
	go r.updateView()

	// Main loop of gocui
	if err := r.G.MainLoop(); err != nil && err != gocui.ErrorQuit {
		log.Panicln(err)
	}
}

// run has an infinite loop, spawning several goroutines to send http requests and collect data
func (r *XTop) run() {

	var (
		status, header string
		wg             sync.WaitGroup
	)

	for { // loop forever, send N requests each loop
		for i := 0; i < r.Concurrent; i++ {
			wg.Add(1)
			go func() {
				// decrement the counter when the goroutine completes.
				defer wg.Done()

				r.TotalRequestsSent++

				resp, err := http.Get(r.URL)
				if err != nil {
					// todo better error handling
					return
				}

				// must close body after finish
				defer resp.Body.Close()

				// collect response status code
				status = resp.Status
				_, exist := r.StatusMap[status]
				if exist {
					r.StatusMap[status]++
				} else {
					r.StatusMap[status] = 1
				}

				// collect custom response header
				header = resp.Header.Get(r.Header)
				_, exist = r.ServerMap[header]
				if exist {
					r.ServerMap[header]++
				} else {
					r.ServerMap[header] = 1
				}
			}()
		}
		wg.Wait()
	}
}

func (r *XTop) layout(*gocui.Gui) error {
	maxX, maxY := r.G.Size()
	if v, err := r.G.SetView("center", 3, 0, maxX, maxY); err != nil {
		if err != gocui.ErrorUnkView {
			return err
		}
		v.Frame = false
		r.display(v)
	}
	return nil
}

// updateView runs in a goroutine to periodically print out results to the screen
// it calls display() to actually print to a view
func (r *XTop) updateView() {
	for {
		time.Sleep(time.Second)
		v, err := r.G.View("center")
		if err != nil {
			panic(err)
		}
		v.Clear()
		r.display(v)
		r.G.Flush()
	}
}

// display prints output to the view
func (r *XTop) display(v io.Writer) error {

	output := fmt.Sprintf("Target: %s\n", r.URL)
	output += fmt.Sprintf("Header to check: %s\n", r.Header)
	output += fmt.Sprintf("Concurrent requests: %d\n\n", r.Concurrent)

	// Statuses
	sorted := sortMapByValue(r.StatusMap)
	output += fmt.Sprintf("=== Response status ===\n")
	for _, v := range sorted {
		output += fmt.Sprintf("%6s %8s %s\n",
			fmt.Sprintf("[%d%s]", v.Value*100/r.TotalRequestsSent, "%%"),
			fmt.Sprintf("[%d/%d]", v.Value, r.TotalRequestsSent),
			v.Key)
	}
	output += fmt.Sprintf("\n")

	// Custom header
	a := []string{}

	for k := range r.ServerMap {
		a = append(a, k)
	}
	sort.Strings(a)

	output += fmt.Sprintf("=== Response header %s ===\n", r.Header)
	for i, v := range a {
		output += fmt.Sprintf("%6s %8s %3d %s\n",
			fmt.Sprintf("[%d%s]", r.ServerMap[v]*100/r.TotalRequestsSent, "%%"),
			fmt.Sprintf("[%d/%d]", r.ServerMap[v], r.TotalRequestsSent),
			i+1,
			v)
	}
	output += fmt.Sprintf("\n")
	fmt.Fprintf(v, output)
	return nil
}
<commit_msg>Improve concurrency model, use buffered channel with transfer struct for communication<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/jroimartin/gocui"
	"gopkg.in/alecthomas/kingpin.v1"
)

func main() {
	name := "xtop"
	help := `A top-like tool to monitor responses from a target URL. This tool periodically collects and prints out response statuses and a custom response header received from the url.`

	app := kingpin.New(name, help)
	app.Version("0.0.2")

	var (
		url         = app.Arg("url", "target URL").Required().URL()
		concurrency = app.Flag("concurrency", "max number of concurrent requests").Short('c').Default("10").Int()
		header      = app.Flag("header", "custom header name to collect").Short('x').Default("X-Server").String()
	)

	kingpin.MustParse(app.Parse(os.Args[1:]))

	x := NewXTop((*url).String(), *concurrency, *header)
	x.Start()
}

// Pair is a data structure to hold a key/value pair.
type Pair struct {
	Key   string
	Value int
}

// PairList is a slice of Pairs that implements sort.Interface to sort by Value.
type PairList []Pair

func (p PairList) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p PairList) Len() int           { return len(p) }
func (p PairList) Less(i, j int) bool { return p[i].Value > p[j].Value }

// sortMapByValue turns a map into a PairList, then sort and return it.
func sortMapByValue(m map[string]int) PairList {
	p := make(PairList, len(m))
	i := 0
	for k, v := range m {
		p[i] = Pair{k, v}
		i++
	}
	sort.Sort(p)
	return p
}

// appendIfMissing appends a string into a slice if it does not exist yet
func appendIfMissing(a []string, item string) []string {
	for _, v := range a {
		if item == v {
			return a
		}
	}
	return append(a, item)
}

// XTop is the main struct
type XTop struct {
	URL         string
	Concurrency int
	Header      string

	TotalRequestsSent int
	StatusMap         map[string]int
	HeaderMap         map[string]int
	G                 *gocui.Gui
}

// Data is the transfer object for communication from our background worker threads
type Data struct {
	RespStatus string
	RespHeader string
	Error      error
}

// NewXTop returns a new XTop instance
func NewXTop(url string, concurrency int, header string) XTop {
	url = strings.ToLower(url)
	if strings.Index(url, "http://") == -1 && strings.Index(url, "https://") == -1 {
		url = "http://" + url
	}

	x := XTop{}
	x.URL = url
	x.Concurrency = concurrency
	x.Header = header
	x.StatusMap = make(map[string]int)
	x.HeaderMap = make(map[string]int)

	var err error
	g := gocui.NewGui()
	if err := g.Init(); err != nil {
		panic(err)
	}
	g.FgColor = gocui.ColorGreen
	g.SetLayout(x.layout)
	if err = g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, func(g *gocui.Gui, v *gocui.View) error {
		return gocui.ErrorQuit
	}); err != nil {
		log.Panicln(err)
	}

	x.G = g
	return x
}

// Start starts everything
func (r *XTop) Start() {
	defer r.G.Close()

	// Collect data
	go r.run()

	// Periodically print out results to the screen
	go r.updateView()

	// Main loop of gocui
	if err := r.G.MainLoop(); err != nil && err != gocui.ErrorQuit {
		log.Panicln(err)
	}
}

// run spawns multiple worker goroutines making requests to target URL in background
// then it keeps listening for results sent back from those workers
func (r *XTop) run() {
	N := r.Concurrency       // max concurrency
	ch := make(chan Data, N) // channel for workers to send back results

	// Spawn worker goroutines
	for i := 0; i < N; i++ {
		go func(ch chan Data) {
			for {
				resp, err := http.Get(r.URL)
				if err != nil {
					ch <- Data{Error: err}
					continue
				}
				ch <- Data{resp.Status, resp.Header.Get(r.Header), nil}
			}
		}(ch)
	}

	// Collecting results from workers
	for {
		data := <-ch
		r.TotalRequestsSent++

		if data.Error != nil {
			// TODO collect errors
			continue
		}

		s := data.RespStatus
		if _, exist := r.StatusMap[s]; exist {
			r.StatusMap[s]++
		} else {
			r.StatusMap[s] = 1
		}

		h := data.RespHeader
		if _, exist := r.HeaderMap[h]; exist {
			r.HeaderMap[h]++
		} else {
			r.HeaderMap[h] = 1
		}
	}
}

func (r *XTop) layout(*gocui.Gui) error {
	maxX, maxY := r.G.Size()
	if v, err := r.G.SetView("center", 3, 0, maxX, maxY); err != nil {
		if err != gocui.ErrorUnkView {
			return err
		}
		v.Frame = false
		r.display(v)
	}
	return nil
}

// updateView runs in a goroutine to periodically print out results to the screen
// it calls display() to actually print to a view
func (r *XTop) updateView() {
	for {
		time.Sleep(time.Second)
		v, err := r.G.View("center")
		if err != nil {
			panic(err)
		}
		v.Clear()
		r.display(v)
		r.G.Flush()
	}
}

// display prints output to the view
func (r *XTop) display(v io.Writer) error {

	output := fmt.Sprintf("Target: %s\n", r.URL)
	output += fmt.Sprintf("Header: %s\n", r.Header)
	output += fmt.Sprintf("Max concurrency: %d\n\n", r.Concurrency)

	// Response status
	sorted := sortMapByValue(r.StatusMap)
	output += fmt.Sprintf("=== Response status ===\n")
	for _, v := range sorted {
		output += fmt.Sprintf("%6s %8s %s\n",
			fmt.Sprintf("[%d%s]", v.Value*100/r.TotalRequestsSent, "%%"),
			fmt.Sprintf("[%d/%d]", v.Value, r.TotalRequestsSent),
			v.Key)
	}
	output += fmt.Sprintf("\n")

	// Response header
	a := []string{}

	for k := range r.HeaderMap {
		a = append(a, k)
	}
	sort.Strings(a)

	output += fmt.Sprintf("=== Response header %s ===\n", r.Header)
	for i, v := range a {
		output += fmt.Sprintf("%6s %8s %3d %s\n",
			fmt.Sprintf("[%d%s]", r.HeaderMap[v]*100/r.TotalRequestsSent, "%%"),
			fmt.Sprintf("[%d/%d]", r.HeaderMap[v], r.TotalRequestsSent),
			i+1,
			v)
	}
	output += fmt.Sprintf("\n")
	fmt.Fprintf(v, output)
	return nil
}
<|endoftext|>


<commit_before>// Package web is a lightweight web framework for Go. It's ideal for
// writing simple, performant backend web services.
package xweb

import (
	"crypto/tls"
	"io/ioutil"
	"log"
	"net/http"
	//"os"
	"path"
	//"reflect"
	//"strings"
	"fmt"
)

const (
	version = "0.1.2"
)

// small optimization: cache the context type instead of repeteadly calling reflect.Typeof
//var contextType reflect.Type

func init() {
	//contextType = reflect.TypeOf(Context{})
	//find the location of the exe file
	/*wd, _ := os.Getwd()
	arg0 := path.Clean(os.Args[0])
	var exeFile string
	if strings.HasPrefix(arg0, "/") {
		exeFile = arg0
	} else {
		//TODO for robustness, search each directory in $PATH
		exeFile = path.Join(wd, arg0)
	}
	_, _ := path.Split(exeFile)*/
	return
}

func Redirect(w http.ResponseWriter, url string) error {
	w.Header().Set("Location", url)
	w.WriteHeader(302)
	_, err := w.Write([]byte("Redirecting to: " + url))
	return err
}

func Download(w http.ResponseWriter, fpath string) error {
	data, err := ioutil.ReadFile(fpath)
	if err != nil {
		return err
	}
	fName := fpath[len(path.Dir(fpath))+1:]
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%v\"", fName))
	_, err = w.Write(data)
	return err
}

const (
	errorTmpl = `<title>%d - %s</title>

%d - %s

	%s

	xweb v%s <input type="button" style="" value="返回" onclick="history.go(-1)"/>

	`
)

func Error(w http.ResponseWriter, status int, content string) error {
	w.WriteHeader(status)
	res := fmt.Sprintf(errorTmpl, status, statusText[status],
		status, statusText[status], content, version)
	_, err := w.Write([]byte(res))
	return err
}

// Process invokes the main server's routing system.
func Process(c http.ResponseWriter, req *http.Request) {
	mainServer.Process(c, req)
}

// Run starts the web application and serves HTTP requests for the main server.
func Run(addr string) {
	mainServer.Run(addr)
}

// RunTLS starts the web application and serves HTTPS requests for the main server.
func RunTLS(addr string, config *tls.Config) {
	mainServer.RunTLS(addr, config)
}

// RunScgi starts the web application and serves SCGI requests for the main server.
func RunScgi(addr string) {
	mainServer.RunScgi(addr)
}

// RunFcgi starts the web application and serves FastCGI requests for the main server.
func RunFcgi(addr string) {
	mainServer.RunFcgi(addr)
}

// Close stops the main server.
func Close() {
	mainServer.Close()
}

func AutoAction(c ...interface{}) {
	mainServer.AutoAction(c...)
}

func AddAction(c ...interface{}) {
	mainServer.AddAction(c...)
}

func AddRouter(url string, c interface{}) {
	mainServer.AddRouter(url, c)
}

func AddFilter(filter Filter) {
	mainServer.AddFilter(filter)
}

func AddApp(a *App) {
	mainServer.AddApp(a)
}

func AddConfig(name string, value interface{}) {
	mainServer.AddConfig(name, value)
}

func SetTemplateDir(dir string) {
	mainServer.SetTemplateDir(dir)
}

func SetStaticDir(dir string) {
	mainServer.SetStaticDir(dir)
}

// SetLogger sets the logger for the main server.
func SetLogger(logger *log.Logger) {
	mainServer.Logger = logger
}

func MainServer() *Server {
	return mainServer
}

func RootApp() *App {
	return MainServer().RootApp
}

// Config is the configuration of the main server.
var (
	Config *ServerConfig = &ServerConfig{
		RecoverPanic: true,
		EnableGzip:   true,
		//Profiler: true,
		StaticExtensionsToGzip: []string{".css", ".js"},
	}
	Servers    map[string]*Server = make(map[string]*Server) //[SWH|+]
	mainServer *Server            = NewServer("main")
)
<commit_msg>美化错误页面模板和增加对自定义错误页面模板的支持<commit_after>// Package web is a lightweight web framework for Go. It's ideal for
// writing simple, performant backend web services.
package xweb

import (
	"crypto/tls"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	//"reflect"
	"fmt"
	//"strings"
)

const (
	version = "0.1.2"
)

// small optimization: cache the context type instead of repeteadly calling reflect.Typeof
//var contextType reflect.Type

func init() {
	//contextType = reflect.TypeOf(Context{})
	//find the location of the exe file
	/*wd, _ := os.Getwd()
	arg0 := path.Clean(os.Args[0])
	var exeFile string
	if strings.HasPrefix(arg0, "/") {
		exeFile = arg0
	} else {
		//TODO for robustness, search each directory in $PATH
		exeFile = path.Join(wd, arg0)
	}
	_, _ := path.Split(exeFile)*/
	return
}

func Redirect(w http.ResponseWriter, url string) error {
	w.Header().Set("Location", url)
	w.WriteHeader(302)
	_, err := w.Write([]byte("Redirecting to: " + url))
	return err
}

func Download(w http.ResponseWriter, fpath string) error {
	data, err := ioutil.ReadFile(fpath)
	if err != nil {
		return err
	}
	fName := fpath[len(path.Dir(fpath))+1:]
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%v\"", fName))
	_, err = w.Write(data)
	return err
}

const (
	defaultErrorTmpl        = `<!DOCTYPE html>

	<meta charset="UTF-8" />
	<title>%d - %s</title>

%d - %s

%s

			<input type="button" title="Back" value="返回" onclick="history.go(-1)"/>
			(xweb v%s)

`
)
var errorTmpl string = ""

func Error(w http.ResponseWriter, status int, content string) error {
	w.WriteHeader(status)
	if errorTmpl == "" {
		errTmplFile := RootApp().AppConfig.TemplateDir + "/_error.html"
		if file, err := os.Stat(errTmplFile); err == nil && !file.IsDir() {
			if b, e := ioutil.ReadFile(errTmplFile); e == nil {
				errorTmpl = string(b)
			}
		}
		if errorTmpl == "" {
			errorTmpl = defaultErrorTmpl
		}
	}
	res := fmt.Sprintf(errorTmpl, status, statusText[status],
		status, statusText[status], content, version)
	_, err := w.Write([]byte(res))
	return err
}

// Process invokes the main server's routing system.
func Process(c http.ResponseWriter, req *http.Request) {
	mainServer.Process(c, req)
}

// Run starts the web application and serves HTTP requests for the main server.
func Run(addr string) {
	mainServer.Run(addr)
}

// RunTLS starts the web application and serves HTTPS requests for the main server.
func RunTLS(addr string, config *tls.Config) {
	mainServer.RunTLS(addr, config)
}

// RunScgi starts the web application and serves SCGI requests for the main server.
func RunScgi(addr string) {
	mainServer.RunScgi(addr)
}

// RunFcgi starts the web application and serves FastCGI requests for the main server.
func RunFcgi(addr string) {
	mainServer.RunFcgi(addr)
}

// Close stops the main server.
func Close() {
	mainServer.Close()
}

func AutoAction(c ...interface{}) {
	mainServer.AutoAction(c...)
}

func AddAction(c ...interface{}) {
	mainServer.AddAction(c...)
}

func AddRouter(url string, c interface{}) {
	mainServer.AddRouter(url, c)
}

func AddFilter(filter Filter) {
	mainServer.AddFilter(filter)
}

func AddApp(a *App) {
	mainServer.AddApp(a)
}

func AddConfig(name string, value interface{}) {
	mainServer.AddConfig(name, value)
}

func SetTemplateDir(dir string) {
	mainServer.SetTemplateDir(dir)
}

func SetStaticDir(dir string) {
	mainServer.SetStaticDir(dir)
}

// SetLogger sets the logger for the main server.
func SetLogger(logger *log.Logger) {
	mainServer.Logger = logger
}

func MainServer() *Server {
	return mainServer
}

func RootApp() *App {
	return MainServer().RootApp
}

// Config is the configuration of the main server.
var (
	Config *ServerConfig = &ServerConfig{
		RecoverPanic: true,
		EnableGzip:   true,
		//Profiler: true,
		StaticExtensionsToGzip: []string{".css", ".js"},
	}
	Servers    map[string]*Server = make(map[string]*Server) //[SWH|+]
	mainServer *Server            = NewServer("main")
)
<|endoftext|>


<commit_before>// Copyright © 2017 NAME HERE <EMAIL ADDRESS>
//
// 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 cmd

import (
	"fmt"
	"log"
    "agenda/entity"
	"github.com/spf13/cobra"
)

// RemoveParticipatorCmd represents the RemoveParticipator command
var RemoveParticipatorCmd = &cobra.Command{
	Use:   "remove -p [Participator] -t [Title]",
	Short: "To remove Participator from the meeting",
	Long: `remove [Participator] from the meeting with the title of [Title]:

attention:If there is no Participators in the meeting,the meeting will be deleted`,
	Run: func(cmd *cobra.Command, args []string) {
		debugLog := log.New(logFile,"[Result]", log.Ldate|log.Ltime|log.Lshortfile)
		if entity.StartAgenda() == false {
			debugLog.Println("Fail,please log in")
			fmt.Println("Fail,please log in")
		}

		arg_p, _ := cmd.Flags().GetStringSlice("Participator")
		arg_t, _ := cmd.Flags().GetString("Title")
		if entity.Removeparticipator(arg_t, arg_p) {
			debugLog.Println("Remove participators successfully")
			fmt.Println("Remove participators successfully")
		} else {
			debugLog.Println("Fail to remove participators")
			fmt.Println("Fail to remove participators")
		}
		entity.QuitAgenda()
	},
}
func init() {
	RootCmd.AddCommand(RemoveParticipatorCmd)
	RemoveParticipatorCmd.Flags().StringSliceP("Participator", "p", []string{}, "meeting's participator")
	RemoveParticipatorCmd.Flags().StringP("Title", "t", "", "meeting title")

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// RemoveParticipatorCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// RemoveParticipatorCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}
<commit_msg>Update RemoveParticipator.go<commit_after>// Copyright © 2017 NAME HERE <EMAIL ADDRESS>
//
// 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 cmd

import (
	"fmt"
	"log"
	entity "github.com/LeungChiHo/agenda/tree/master/entity"
	"github.com/spf13/cobra"
)

// RemoveParticipatorCmd represents the RemoveParticipator command
var RemoveParticipatorCmd = &cobra.Command{
	Use:   "remove -p [Participator] -t [Title]",
	Short: "To remove Participator from the meeting",
	Long: `remove [Participator] from the meeting with the title of [Title]:

attention:If there is no Participators in the meeting,the meeting will be deleted`,
	Run: func(cmd *cobra.Command, args []string) {
		debugLog := log.New(logFile,"[Result]", log.Ldate|log.Ltime|log.Lshortfile)
		if entity.StartAgenda() == false {
			debugLog.Println("Fail,please log in")
			fmt.Println("Fail,please log in")
		}

		arg_p, _ := cmd.Flags().GetStringSlice("Participator")
		arg_t, _ := cmd.Flags().GetString("Title")
		if entity.Removeparticipator(arg_t, arg_p) {
			debugLog.Println("Remove participators successfully")
			fmt.Println("Remove participators successfully")
		} else {
			debugLog.Println("Fail to remove participators")
			fmt.Println("Fail to remove participators")
		}
		entity.QuitAgenda()
	},
}
func init() {
	RootCmd.AddCommand(RemoveParticipatorCmd)
	RemoveParticipatorCmd.Flags().StringSliceP("Participator", "p", []string{}, "meeting's participator")
	RemoveParticipatorCmd.Flags().StringP("Title", "t", "", "meeting title")

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// RemoveParticipatorCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// RemoveParticipatorCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}
<|endoftext|>


<commit_before>//  Copyright (c) 2017 Couchbase, Inc.
//
// 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 zap

import (
	"encoding/binary"
	"fmt"
	"math"

	"github.com/blevesearch/bleve/index/scorch/segment/zap"
	"github.com/couchbase/vellum"
	"github.com/spf13/cobra"
)

// dictCmd represents the dict command
var dictCmd = &cobra.Command{
	Use:   "dict [path] [field]",
	Short: "dict prints the term dictionary for the specified field",
	Long:  `The dict command lets you print the term dictionary for the specified field.`,
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) < 2 {
			return fmt.Errorf("must specify field")
		}

		data := segment.Data()

		addr, err := segment.DictAddr(args[1])
		if err != nil {
			return fmt.Errorf("error determining address: %v", err)
		}
		fmt.Printf("dictionary for field starts at %d (%x)\n", addr, addr)

		vellumLen, read := binary.Uvarint(data[addr : addr+binary.MaxVarintLen64])
		fmt.Printf("vellum length: %d\n", vellumLen)
		fstBytes := data[addr+uint64(read) : addr+uint64(read)+vellumLen]
		fmt.Printf("raw vellum data:\n % x\n", fstBytes)
		fmt.Printf("dictionary:\n")
		if fstBytes != nil {
			fst, err := vellum.Load(fstBytes)
			if err != nil {
				return fmt.Errorf("dictionary field %s vellum err: %v", args[1], err)
			}

			itr, err := fst.Iterator(nil, nil)
			for err == nil {
				currTerm, currVal := itr.Current()
				extra := ""
				if currVal&zap.FSTValEncodingMask == zap.FSTValEncoding1Hit {
					docNum, normBits := zap.FSTValDecode1Hit(currVal)
					norm := math.Float32frombits(uint32(normBits))
					extra = fmt.Sprintf("-- docNum: %d, norm: %f", docNum, norm)
				}

				fmt.Printf(" %s - %d (%x) %s\n", currTerm, currVal, currVal, extra)
				err = itr.Next()
			}
			if err != nil && err != vellum.ErrIteratorDone {
				return fmt.Errorf("error iterating dictionary: %v", err)
			}

		}

		return nil
	},
}

func init() {
	RootCmd.AddCommand(dictCmd)
}
<commit_msg>dict cmd additions<commit_after>//  Copyright (c) 2017 Couchbase, Inc.
//
// 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 zap

import (
	"encoding/binary"
	"fmt"
	"math"

	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/index/scorch/segment/zap"
	"github.com/couchbase/vellum"
	"github.com/spf13/cobra"
)

// dictCmd represents the dict command
var dictCmd = &cobra.Command{
	Use:   "dict [path] [field]",
	Short: "dict prints the term dictionary for the specified field",
	Long:  `The dict command lets you print the term dictionary for the specified field.`,
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) < 2 {
			return fmt.Errorf("must specify field")
		}

		data := segment.Data()

		addr, err := segment.DictAddr(args[1])
		if err != nil {
			return fmt.Errorf("error determining address: %v", err)
		}
		fmt.Printf("dictionary for field starts at %d (%x)\n", addr, addr)

		vellumLen, read := binary.Uvarint(data[addr : addr+binary.MaxVarintLen64])
		fmt.Printf("vellum length: %d\n", vellumLen)
		fstBytes := data[addr+uint64(read) : addr+uint64(read)+vellumLen]
		fmt.Printf("raw vellum data:\n % x\n", fstBytes)
		fmt.Printf("dictionary:\n")
		var termsCount, hit1Count int64
		if fstBytes != nil {
			fst, err := vellum.Load(fstBytes)
			if err != nil {
				return fmt.Errorf("dictionary field %s vellum err: %v", args[1], err)
			}

			itr, err := fst.Iterator(nil, nil)
			for err == nil {
				currTerm, currVal := itr.Current()
				extra := ""
				if currVal&zap.FSTValEncodingMask == zap.FSTValEncoding1Hit {
					docNum, normBits := zap.FSTValDecode1Hit(currVal)
					norm := math.Float32frombits(uint32(normBits))
					extra = fmt.Sprintf("-- docNum: %d, norm: %f", docNum, norm)
					fmt.Printf(" %s - %d (%x) %s\n", currTerm, currVal, currVal, extra)
					hit1Count++
				} else {
					// fetch the postings size, cardinality in case of non 1 hits
					l, c := readPostingCardinality(currVal, data)
					fmt.Printf(" %s - %d (%x) posting byteSize: %d cardinality: %d\n",
						currTerm, currVal, currVal, l, c)
				}
				termsCount++
				err = itr.Next()
			}
			if err != nil && err != vellum.ErrIteratorDone {
				return fmt.Errorf("error iterating dictionary: %v", err)
			}
			fmt.Printf("Total terms in dictionary : %d 1hit count: %d\n", termsCount, hit1Count)
		}

		return nil
	},
}

func init() {
	RootCmd.AddCommand(dictCmd)
}

func readPostingCardinality(postingsOffset uint64, data []byte) (int, uint64) {
	// read the location of the freq/norm details
	var n uint64
	var read int

	_, read = binary.Uvarint(data[postingsOffset+n : postingsOffset+binary.MaxVarintLen64])
	n += uint64(read)

	_, read = binary.Uvarint(data[postingsOffset+n : postingsOffset+n+binary.MaxVarintLen64])
	n += uint64(read)

	var postingsLen uint64
	postingsLen, read = binary.Uvarint(data[postingsOffset+n : postingsOffset+n+binary.MaxVarintLen64])
	n += uint64(read)

	roaringBytes := data[postingsOffset+n : postingsOffset+n+postingsLen]

	r := roaring.NewBitmap()

	_, err := r.FromBuffer(roaringBytes)
	if err != nil {
		fmt.Printf("error loading roaring bitmap: %v", err)
	}

	return len(roaringBytes), r.GetCardinality()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/labstack/echo"
	"github.com/monochromegane/gannoy"
)

var (
	dataDir string
)

func init() {
	flag.StringVar(&dataDir, "d", ".", "Data directory.")
	flag.Parse()
}

type Feature struct {
	W []float64 `json:"features"`
}

func main() {
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	databases := map[string]gannoy.GannoyIndex{}
	for _, file := range files {
		if file.IsDir() || filepath.Ext(file.Name()) != ".meta" {
			continue
		}
		key := strings.TrimSuffix(file.Name(), ".meta")
		gannoy, err := gannoy.NewGannoyIndex("hoge.meta", gannoy.Angular{}, gannoy.RandRandom{})
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		databases[key] = gannoy
	}

	e := echo.New()
	e.GET("/search", func(c echo.Context) error {
		database := c.QueryParam("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusNotFound)
		}
		id, err := strconv.Atoi(c.QueryParam("id"))
		if err != nil {
			id = -1
		}
		limit, err := strconv.Atoi(c.QueryParam("limit"))
		if err != nil {
			limit = 10
		}

		gannoy := databases[database]
		r, err := gannoy.GetNnsByKey(id, limit, -1)
		if err != nil || len(r) == 0 {
			return c.NoContent(http.StatusNotFound)
		}

		return c.JSON(http.StatusOK, r)
	})

	e.PUT("/databases/:database/features/:id", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		id, err := strconv.Atoi(c.Param("id"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		feature := new(Feature)
		if err := c.Bind(feature); err != nil {
			return err
		}

		gannoy := databases[database]
		err = gannoy.AddItem(id, feature.W)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		return c.NoContent(http.StatusOK)
	})

	e.Start(":1323")
}
<commit_msg>Added delete API in gannoy-server.<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/labstack/echo"
	"github.com/monochromegane/gannoy"
)

var (
	dataDir string
)

func init() {
	flag.StringVar(&dataDir, "d", ".", "Data directory.")
	flag.Parse()
}

type Feature struct {
	W []float64 `json:"features"`
}

func main() {
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	databases := map[string]gannoy.GannoyIndex{}
	for _, file := range files {
		if file.IsDir() || filepath.Ext(file.Name()) != ".meta" {
			continue
		}
		key := strings.TrimSuffix(file.Name(), ".meta")
		gannoy, err := gannoy.NewGannoyIndex("hoge.meta", gannoy.Angular{}, gannoy.RandRandom{})
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		databases[key] = gannoy
	}

	e := echo.New()
	e.GET("/search", func(c echo.Context) error {
		database := c.QueryParam("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusNotFound)
		}
		id, err := strconv.Atoi(c.QueryParam("id"))
		if err != nil {
			id = -1
		}
		limit, err := strconv.Atoi(c.QueryParam("limit"))
		if err != nil {
			limit = 10
		}

		gannoy := databases[database]
		r, err := gannoy.GetNnsByKey(id, limit, -1)
		if err != nil || len(r) == 0 {
			return c.NoContent(http.StatusNotFound)
		}

		return c.JSON(http.StatusOK, r)
	})

	e.PUT("/databases/:database/features/:id", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		id, err := strconv.Atoi(c.Param("id"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		feature := new(Feature)
		if err := c.Bind(feature); err != nil {
			return err
		}

		gannoy := databases[database]
		err = gannoy.AddItem(id, feature.W)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		return c.NoContent(http.StatusOK)
	})

	e.DELETE("/databases/:database/features/:id", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		id, err := strconv.Atoi(c.Param("id"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		gannoy := databases[database]
		err = gannoy.RemoveItem(id)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}

		return c.NoContent(http.StatusOK)
	})

	e.Start(":1323")
}
<|endoftext|>