<commit_before>package boomer

import (
	"time"
)

type requestSuccess struct {
	requestType    string
	name           string
	responseTime   int64
	responseLength int64
}

type requestFailure struct {
	requestType  string
	name         string
	responseTime int64
	error        string
}

type requestStats struct {
	entries   map[string]*statsEntry
	errors    map[string]*statsError
	total     *statsEntry
	startTime int64

	requestSuccessChan  chan *requestSuccess
	requestFailureChan  chan *requestFailure
	clearStatsChan      chan bool
	messageToRunnerChan chan map[string]interface{}
	shutdownChan        chan bool
}

func newRequestStats() (stats *requestStats) {
	entries := make(map[string]*statsEntry)
	errors := make(map[string]*statsError)

	stats = &requestStats{
		entries: entries,
		errors:  errors,
	}
	stats.requestSuccessChan = make(chan *requestSuccess, 100)
	stats.requestFailureChan = make(chan *requestFailure, 100)
	stats.clearStatsChan = make(chan bool)
	stats.messageToRunnerChan = make(chan map[string]interface{}, 10)
	stats.shutdownChan = make(chan bool)

	stats.total = &statsEntry{
		name:   "Total",
		method: "",
	}
	stats.total.reset()

	return stats
}

func (s *requestStats) logRequest(method, name string, responseTime int64, contentLength int64) {
	s.total.log(responseTime, contentLength)
	s.get(name, method).log(responseTime, contentLength)
}

func (s *requestStats) logError(method, name, err string) {
	s.total.logError(err)
	s.get(name, method).logError(err)

	// store error in errors map
	key := MD5(method, name, err)
	entry, ok := s.errors[key]
	if !ok {
		entry = &statsError{
			name:   name,
			method: method,
			error:  err,
		}
		s.errors[key] = entry
	}
	entry.occured()
}

func (s *requestStats) get(name string, method string) (entry *statsEntry) {
	entry, ok := s.entries[name+method]
	if !ok {
		newEntry := &statsEntry{
			name:          name,
			method:        method,
			numReqsPerSec: make(map[int64]int64),
			responseTimes: make(map[int64]int64),
		}
		newEntry.reset()
		s.entries[name+method] = newEntry
		return newEntry
	}
	return entry
}

func (s *requestStats) clearAll() {
	s.total = &statsEntry{
		name:   "Total",
		method: "",
	}
	s.total.reset()

	s.entries = make(map[string]*statsEntry)
	s.errors = make(map[string]*statsError)
	s.startTime = time.Now().Unix()
}

func (s *requestStats) serializeStats() []interface{} {
	entries := make([]interface{}, 0, len(s.entries))
	for _, v := range s.entries {
		if !(v.numRequests == 0 && v.numFailures == 0) {
			entries = append(entries, v.getStrippedReport())
		}
	}
	return entries
}

func (s *requestStats) serializeErrors() map[string]map[string]interface{} {
	errors := make(map[string]map[string]interface{})
	for k, v := range s.errors {
		errors[k] = v.toMap()
	}
	return errors
}

func (s *requestStats) collectReportData() map[string]interface{} {
	data := make(map[string]interface{})
	data["stats"] = s.serializeStats()
	data["stats_total"] = s.total.getStrippedReport()
	data["errors"] = s.serializeErrors()
	s.errors = make(map[string]*statsError)
	return data
}

func (s *requestStats) start() {
	go func() {
		var ticker = time.NewTicker(slaveReportInterval)
		for {
			select {
			case m := <-s.requestSuccessChan:
				s.logRequest(m.requestType, m.name, m.responseTime, m.responseLength)
			case n := <-s.requestFailureChan:
				s.logError(n.requestType, n.name, n.error)
			case <-s.clearStatsChan:
				s.clearAll()
			case <-ticker.C:
				data := s.collectReportData()
				// send data to channel, no network IO in this goroutine
				s.messageToRunnerChan <- data
			case <-s.shutdownChan:
				return
			}
		}
	}()
}

// close is used by unit tests to avoid leakage of goroutines
func (s *requestStats) close() {
	close(s.shutdownChan)
}

type statsEntry struct {
	name                 string
	method               string
	numRequests          int64
	numFailures          int64
	totalResponseTime    int64
	minResponseTime      int64
	maxResponseTime      int64
	numReqsPerSec        map[int64]int64
	responseTimes        map[int64]int64
	totalContentLength   int64
	startTime            int64
	lastRequestTimestamp int64
}

func (s *statsEntry) reset() {
	s.startTime = time.Now().Unix()
	s.numRequests = 0
	s.numFailures = 0
	s.totalResponseTime = 0
	s.responseTimes = make(map[int64]int64)
	s.minResponseTime = 0
	s.maxResponseTime = 0
	s.lastRequestTimestamp = time.Now().Unix()
	s.numReqsPerSec = make(map[int64]int64)
	s.totalContentLength = 0
}

func (s *statsEntry) log(responseTime int64, contentLength int64) {
	s.numRequests++

	s.logTimeOfRequest()
	s.logResponseTime(responseTime)

	s.totalContentLength += contentLength
}

func (s *statsEntry) logTimeOfRequest() {
	key := time.Now().Unix()
	_, ok := s.numReqsPerSec[key]
	if !ok {
		s.numReqsPerSec[key] = 1
	} else {
		s.numReqsPerSec[key]++
	}

	s.lastRequestTimestamp = key
}

func (s *statsEntry) logResponseTime(responseTime int64) {
	s.totalResponseTime += responseTime

	if s.minResponseTime == 0 {
		s.minResponseTime = responseTime
	}

	if responseTime < s.minResponseTime {
		s.minResponseTime = responseTime
	}

	if responseTime > s.maxResponseTime {
		s.maxResponseTime = responseTime
	}

	var roundedResponseTime int64

	// to avoid to much data that has to be transferred to the master node when
	// running in distributed mode, we save the response time rounded in a dict
	// so that 147 becomes 150, 3432 becomes 3400 and 58760 becomes 59000
	// see also locust's stats.py
	if responseTime < 100 {
		roundedResponseTime = responseTime
	} else if responseTime < 1000 {
		roundedResponseTime = int64(round(float64(responseTime), .5, -1))
	} else if responseTime < 10000 {
		roundedResponseTime = int64(round(float64(responseTime), .5, -2))
	} else {
		roundedResponseTime = int64(round(float64(responseTime), .5, -3))
	}

	_, ok := s.responseTimes[roundedResponseTime]
	if !ok {
		s.responseTimes[roundedResponseTime] = 1
	} else {
		s.responseTimes[roundedResponseTime]++
	}
}

func (s *statsEntry) logError(err string) {
	s.numFailures++
}

func (s *statsEntry) serialize() map[string]interface{} {
	result := make(map[string]interface{})
	result["name"] = s.name
	result["method"] = s.method
	result["last_request_timestamp"] = s.lastRequestTimestamp
	result["start_time"] = s.startTime
	result["num_requests"] = s.numRequests
	result["num_failures"] = s.numFailures
	result["total_response_time"] = s.totalResponseTime
	result["max_response_time"] = s.maxResponseTime
	result["min_response_time"] = s.minResponseTime
	result["total_content_length"] = s.totalContentLength
	result["response_times"] = s.responseTimes
	result["num_reqs_per_sec"] = s.numReqsPerSec
	return result
}

func (s *statsEntry) getStrippedReport() map[string]interface{} {
	report := s.serialize()
	s.reset()
	return report
}

type statsError struct {
	name        string
	method      string
	error       string
	occurrences int64
}

func (err *statsError) occured() {
	err.occurrences++
}

func (err *statsError) toMap() map[string]interface{} {
	m := make(map[string]interface{})

	m["method"] = err.method
	m["name"] = err.name
	m["error"] = err.error
	m["occurrences"] = err.occurrences

	// keep compatible with locust
	// https://github.com/locustio/locust/commit/f0a5f893734faeddb83860b2985010facc910d7d#diff-5d5f310549d6d596beaa43a1282ec49e
	m["occurences"] = err.occurrences
	return m
}
<commit_msg>FIX: add num_none_requests to keep compatible with locust<commit_after>package boomer

import (
	"time"
)

type requestSuccess struct {
	requestType    string
	name           string
	responseTime   int64
	responseLength int64
}

type requestFailure struct {
	requestType  string
	name         string
	responseTime int64
	error        string
}

type requestStats struct {
	entries   map[string]*statsEntry
	errors    map[string]*statsError
	total     *statsEntry
	startTime int64

	requestSuccessChan  chan *requestSuccess
	requestFailureChan  chan *requestFailure
	clearStatsChan      chan bool
	messageToRunnerChan chan map[string]interface{}
	shutdownChan        chan bool
}

func newRequestStats() (stats *requestStats) {
	entries := make(map[string]*statsEntry)
	errors := make(map[string]*statsError)

	stats = &requestStats{
		entries: entries,
		errors:  errors,
	}
	stats.requestSuccessChan = make(chan *requestSuccess, 100)
	stats.requestFailureChan = make(chan *requestFailure, 100)
	stats.clearStatsChan = make(chan bool)
	stats.messageToRunnerChan = make(chan map[string]interface{}, 10)
	stats.shutdownChan = make(chan bool)

	stats.total = &statsEntry{
		name:   "Total",
		method: "",
	}
	stats.total.reset()

	return stats
}

func (s *requestStats) logRequest(method, name string, responseTime int64, contentLength int64) {
	s.total.log(responseTime, contentLength)
	s.get(name, method).log(responseTime, contentLength)
}

func (s *requestStats) logError(method, name, err string) {
	s.total.logError(err)
	s.get(name, method).logError(err)

	// store error in errors map
	key := MD5(method, name, err)
	entry, ok := s.errors[key]
	if !ok {
		entry = &statsError{
			name:   name,
			method: method,
			error:  err,
		}
		s.errors[key] = entry
	}
	entry.occured()
}

func (s *requestStats) get(name string, method string) (entry *statsEntry) {
	entry, ok := s.entries[name+method]
	if !ok {
		newEntry := &statsEntry{
			name:          name,
			method:        method,
			numReqsPerSec: make(map[int64]int64),
			responseTimes: make(map[int64]int64),
		}
		newEntry.reset()
		s.entries[name+method] = newEntry
		return newEntry
	}
	return entry
}

func (s *requestStats) clearAll() {
	s.total = &statsEntry{
		name:   "Total",
		method: "",
	}
	s.total.reset()

	s.entries = make(map[string]*statsEntry)
	s.errors = make(map[string]*statsError)
	s.startTime = time.Now().Unix()
}

func (s *requestStats) serializeStats() []interface{} {
	entries := make([]interface{}, 0, len(s.entries))
	for _, v := range s.entries {
		if !(v.numRequests == 0 && v.numFailures == 0) {
			entries = append(entries, v.getStrippedReport())
		}
	}
	return entries
}

func (s *requestStats) serializeErrors() map[string]map[string]interface{} {
	errors := make(map[string]map[string]interface{})
	for k, v := range s.errors {
		errors[k] = v.toMap()
	}
	return errors
}

func (s *requestStats) collectReportData() map[string]interface{} {
	data := make(map[string]interface{})
	data["stats"] = s.serializeStats()
	data["stats_total"] = s.total.getStrippedReport()
	data["errors"] = s.serializeErrors()
	s.errors = make(map[string]*statsError)
	return data
}

func (s *requestStats) start() {
	go func() {
		var ticker = time.NewTicker(slaveReportInterval)
		for {
			select {
			case m := <-s.requestSuccessChan:
				s.logRequest(m.requestType, m.name, m.responseTime, m.responseLength)
			case n := <-s.requestFailureChan:
				s.logError(n.requestType, n.name, n.error)
			case <-s.clearStatsChan:
				s.clearAll()
			case <-ticker.C:
				data := s.collectReportData()
				// send data to channel, no network IO in this goroutine
				s.messageToRunnerChan <- data
			case <-s.shutdownChan:
				return
			}
		}
	}()
}

// close is used by unit tests to avoid leakage of goroutines
func (s *requestStats) close() {
	close(s.shutdownChan)
}

type statsEntry struct {
	name                 string
	method               string
	numRequests          int64
	numFailures          int64
	totalResponseTime    int64
	minResponseTime      int64
	maxResponseTime      int64
	numReqsPerSec        map[int64]int64
	responseTimes        map[int64]int64
	totalContentLength   int64
	startTime            int64
	lastRequestTimestamp int64
}

func (s *statsEntry) reset() {
	s.startTime = time.Now().Unix()
	s.numRequests = 0
	s.numFailures = 0
	s.totalResponseTime = 0
	s.responseTimes = make(map[int64]int64)
	s.minResponseTime = 0
	s.maxResponseTime = 0
	s.lastRequestTimestamp = time.Now().Unix()
	s.numReqsPerSec = make(map[int64]int64)
	s.totalContentLength = 0
}

func (s *statsEntry) log(responseTime int64, contentLength int64) {
	s.numRequests++

	s.logTimeOfRequest()
	s.logResponseTime(responseTime)

	s.totalContentLength += contentLength
}

func (s *statsEntry) logTimeOfRequest() {
	key := time.Now().Unix()
	_, ok := s.numReqsPerSec[key]
	if !ok {
		s.numReqsPerSec[key] = 1
	} else {
		s.numReqsPerSec[key]++
	}

	s.lastRequestTimestamp = key
}

func (s *statsEntry) logResponseTime(responseTime int64) {
	s.totalResponseTime += responseTime

	if s.minResponseTime == 0 {
		s.minResponseTime = responseTime
	}

	if responseTime < s.minResponseTime {
		s.minResponseTime = responseTime
	}

	if responseTime > s.maxResponseTime {
		s.maxResponseTime = responseTime
	}

	var roundedResponseTime int64

	// to avoid to much data that has to be transferred to the master node when
	// running in distributed mode, we save the response time rounded in a dict
	// so that 147 becomes 150, 3432 becomes 3400 and 58760 becomes 59000
	// see also locust's stats.py
	if responseTime < 100 {
		roundedResponseTime = responseTime
	} else if responseTime < 1000 {
		roundedResponseTime = int64(round(float64(responseTime), .5, -1))
	} else if responseTime < 10000 {
		roundedResponseTime = int64(round(float64(responseTime), .5, -2))
	} else {
		roundedResponseTime = int64(round(float64(responseTime), .5, -3))
	}

	_, ok := s.responseTimes[roundedResponseTime]
	if !ok {
		s.responseTimes[roundedResponseTime] = 1
	} else {
		s.responseTimes[roundedResponseTime]++
	}
}

func (s *statsEntry) logError(err string) {
	s.numFailures++
}

func (s *statsEntry) serialize() map[string]interface{} {
	result := make(map[string]interface{})
	result["name"] = s.name
	result["method"] = s.method
	result["last_request_timestamp"] = s.lastRequestTimestamp
	result["start_time"] = s.startTime
	result["num_requests"] = s.numRequests
	// Boomer doesn't allow None response time for requests like locust.
	// num_none_requests is added to keep compatible with locust.
	result["num_none_requests"] = 0
	result["num_failures"] = s.numFailures
	result["total_response_time"] = s.totalResponseTime
	result["max_response_time"] = s.maxResponseTime
	result["min_response_time"] = s.minResponseTime
	result["total_content_length"] = s.totalContentLength
	result["response_times"] = s.responseTimes
	result["num_reqs_per_sec"] = s.numReqsPerSec
	return result
}

func (s *statsEntry) getStrippedReport() map[string]interface{} {
	report := s.serialize()
	s.reset()
	return report
}

type statsError struct {
	name        string
	method      string
	error       string
	occurrences int64
}

func (err *statsError) occured() {
	err.occurrences++
}

func (err *statsError) toMap() map[string]interface{} {
	m := make(map[string]interface{})

	m["method"] = err.method
	m["name"] = err.name
	m["error"] = err.error
	m["occurrences"] = err.occurrences

	// keep compatible with locust
	// https://github.com/locustio/locust/commit/f0a5f893734faeddb83860b2985010facc910d7d#diff-5d5f310549d6d596beaa43a1282ec49e
	m["occurences"] = err.occurrences
	return m
}
<|endoftext|>


<commit_before>// An hour is splitted into mailing periods. If it is 16:02, then the current time period is 2.
// Depending on this when a user receives a message in 16:02, they will be notified in 10 minutes (16:12),
// where the mailing period will be 12
// This periods are circular, which leads when the time is 16:56 next mailing period is going to be 6

// Three different redis key is needed here
// 1- AccountNextPeriod hashset (AccountMailingPeriod): It stores notification mailing period for each account
// 2- PeriodAccountId set (Notifiee Queue): It stores notified accounts for each given time period
// 3- AccountId (AccountChannelNotifications): stores account ChannelId:CreatedAt information for each message
package feeder

import (
	"errors"
	"socialapi/models"
	"socialapi/request"
	"socialapi/workers/email/chatemail/common"
	"socialapi/workers/email/emailmodels"
	"strconv"

	"github.com/koding/logging"
	"github.com/koding/redis"
	"github.com/streadway/amqp"
)

var (
	ErrInvalidPeriod  = errors.New("invalid period")
	ErrPeriodNotFound = errors.New("period not found")
)

type Controller struct {
	log   logging.Logger
	redis *redis.RedisSession
}

func New(log logging.Logger, redis *redis.RedisSession) *Controller {
	return &Controller{
		log:   log,
		redis: redis,
	}
}

func (n *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	n.log.Error("an error occurred: %s", err)
	delivery.Ack(false)

	return false
}

var isEligibleToNotify = func(accountId int64) (bool, error) {
	uc, err := emailmodels.FetchUserContact(accountId)
	if err != nil {
		return false, err
	}

	if !uc.EmailSettings.Global {
		return false, nil
	}

	return uc.EmailSettings.PrivateMessage, nil
}

// AddMessageToQueue adds a new arrival message into participants notification queues
func (c *Controller) AddMessageToQueue(cm *models.ChannelMessage) error {
	if cm.TypeConstant != models.ChannelMessage_TYPE_PRIVATE_MESSAGE {
		return nil
	}

	// TODO later on fetch this from cache.
	participantIds, err := c.fetchParticipantIds(cm)
	if err != nil {
		return err
	}

	for _, participantId := range participantIds {
		if participantId == cm.AccountId {
			continue
		}

		if err := c.notifyAccount(participantId, cm); err != nil {
			c.log.Error("Could not add message %d to queue for account %d: %s", cm.Id, participantId, err)
		}

	}

	return nil
}

// GlanceChannel removes a channel from awaiting notification channel hash set and
// when none other channels are awaiting resets Account information from AccountPeriod hash
func (c *Controller) GlanceChannel(cp *models.ChannelParticipant) error {
	a := models.NewAccount()
	a.Id = cp.AccountId

	ch := models.NewChannel()
	ch.Id = cp.ChannelId
	nextPeriod, err := c.getMailingPeriod(a)
	// no awaiting notifications
	if err == ErrPeriodNotFound {
		return nil
	}

	if err != nil {
		return err
	}

	hasPendingNotification, err := c.deletePendingNotifications(a, ch, nextPeriod)
	if err != nil {
		return err
	}

	// there is not any pending notification mail for the channel
	if !hasPendingNotification {
		return nil
	}

	hasUnglancedChannels, err := c.hasUnglancedChannels(a, nextPeriod)
	if err != nil {
		return err
	}

	if hasUnglancedChannels {
		return nil
	}

	// when there is not any pending notifications (unglanced active channels),
	// just reset the account mailing period value
	return common.ResetMailingPeriodForAccount(c.redis, a)
}

// deletePendingNotification deletes notification information for given account, and channel. It returns
// true when there are pending notifications, false otherwise
func (c *Controller) deletePendingNotifications(a *models.Account, ch *models.Channel, nextPeriod string) (bool, error) {
	count, err := c.redis.DeleteHashSetField(common.AccountChannelHashSetKey(a.Id, nextPeriod), strconv.FormatInt(ch.Id, 10))
	if err != nil {
		return false, err
	}

	return count == 1, nil
}

func (c *Controller) hasUnglancedChannels(a *models.Account, nextPeriod string) (bool, error) {
	count, err := c.redis.GetHashLength(common.AccountChannelHashSetKey(a.Id, nextPeriod))
	if err != nil {
		return false, err
	}

	return count == 1, nil
}

func (c *Controller) fetchParticipantIds(cm *models.ChannelMessage) ([]int64, error) {
	ch := models.NewChannel()
	ch.Id = cm.InitialChannelId

	return ch.FetchParticipantIds(&request.Query{})
}

func (c *Controller) notifyAccount(accountId int64, cm *models.ChannelMessage) error {
	a := models.NewAccount()
	a.Id = accountId

	eligible, err := isEligibleToNotify(accountId)
	if err != nil {
		return err
	}

	if !eligible {
		return nil
	}

	nextPeriod, err := c.getOrCreateMailingPeriod(a)
	if err != nil {
		return err
	}

	if err := c.addAccountToNotifieeQueue(nextPeriod, accountId); err != nil {
		return err
	}

	return c.addMessageToAccountChannelNotifications(nextPeriod, accountId, cm)
}

// getOrCreateMailingPeriod updates the Account-Segment hash set and returns
// the next mailing period of the account
func (c *Controller) getOrCreateMailingPeriod(a *models.Account) (string, error) {
	field := strconv.FormatInt(a.Id, 10)

	nextPeriod, err := c.getMailingPeriod(a)
	// if not exist get new mailing period for the account
	if err == ErrPeriodNotFound {
		nextPeriod = common.GetNextMailPeriod()
		err := c.redis.HashMultipleSet(common.AccountNextPeriodHashSetKey(), map[string]interface{}{field: nextPeriod})
		if err != nil {
			return "", err
		}

		return nextPeriod, nil
	}

	if err != nil {
		return "", err
	}

	return nextPeriod, nil
}

func (c *Controller) getMailingPeriod(a *models.Account) (string, error) {
	values, err := c.redis.GetHashMultipleSet(common.AccountNextPeriodHashSetKey(), a.Id)
	if err != nil {
		return "", err
	}

	if len(values) == 0 || values[0] == nil {
		return "", ErrPeriodNotFound
	}

	nextPeriod, err := c.redis.String(values[0])
	if err != nil {
		return "", err
	}

	return nextPeriod, nil
}

func (c *Controller) addAccountToNotifieeQueue(period string, accountId int64) error {
	_, err := c.redis.AddSetMembers(common.PeriodAccountSetKey(period), strconv.FormatInt(accountId, 10))
	if err != nil {
		return err
	}

	return nil
}

func (c *Controller) addMessageToAccountChannelNotifications(period string, accountId int64, cm *models.ChannelMessage) error {
	key := common.AccountChannelHashSetKey(accountId, period)
	channelId := strconv.FormatInt(cm.InitialChannelId, 10)
	awaySince := strconv.FormatInt(cm.CreatedAt.UnixNano(), 10)
	// add the first received message for channel
	_, err := c.redis.HashSetIfNotExists(key, channelId, awaySince)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>email: use HGET for getMailingPeriod<commit_after>// An hour is splitted into mailing periods. If it is 16:02, then the current time period is 2.
// Depending on this when a user receives a message in 16:02, they will be notified in 10 minutes (16:12),
// where the mailing period will be 12
// This periods are circular, which leads when the time is 16:56 next mailing period is going to be 6

// Three different redis key is needed here
// 1- AccountNextPeriod hashset (AccountMailingPeriod): It stores notification mailing period for each account
// 2- PeriodAccountId set (Notifiee Queue): It stores notified accounts for each given time period
// 3- AccountId (AccountChannelNotifications): stores account ChannelId:CreatedAt information for each message
package feeder

import (
	"errors"
	"socialapi/models"
	"socialapi/request"
	"socialapi/workers/email/chatemail/common"
	"socialapi/workers/email/emailmodels"
	"strconv"

	"github.com/koding/logging"
	"github.com/koding/redis"
	"github.com/streadway/amqp"
)

var (
	ErrInvalidPeriod  = errors.New("invalid period")
	ErrPeriodNotFound = errors.New("period not found")
)

type Controller struct {
	log   logging.Logger
	redis *redis.RedisSession
}

func New(log logging.Logger, redis *redis.RedisSession) *Controller {
	return &Controller{
		log:   log,
		redis: redis,
	}
}

func (n *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	n.log.Error("an error occurred: %s", err)
	delivery.Ack(false)

	return false
}

var isEligibleToNotify = func(accountId int64) (bool, error) {
	uc, err := emailmodels.FetchUserContact(accountId)
	if err != nil {
		return false, err
	}

	if !uc.EmailSettings.Global {
		return false, nil
	}

	return uc.EmailSettings.PrivateMessage, nil
}

// AddMessageToQueue adds a new arrival message into participants notification queues
func (c *Controller) AddMessageToQueue(cm *models.ChannelMessage) error {
	if cm.TypeConstant != models.ChannelMessage_TYPE_PRIVATE_MESSAGE {
		return nil
	}

	// TODO later on fetch this from cache.
	participantIds, err := c.fetchParticipantIds(cm)
	if err != nil {
		return err
	}

	for _, participantId := range participantIds {
		if participantId == cm.AccountId {
			continue
		}

		if err := c.notifyAccount(participantId, cm); err != nil {
			c.log.Error("Could not add message %d to queue for account %d: %s", cm.Id, participantId, err)
		}

	}

	return nil
}

// GlanceChannel removes a channel from awaiting notification channel hash set and
// when none other channels are awaiting resets Account information from AccountPeriod hash
func (c *Controller) GlanceChannel(cp *models.ChannelParticipant) error {
	a := models.NewAccount()
	a.Id = cp.AccountId

	ch := models.NewChannel()
	ch.Id = cp.ChannelId
	nextPeriod, err := c.getMailingPeriod(a)
	// no awaiting notifications
	if err == ErrPeriodNotFound {
		return nil
	}

	if err != nil {
		return err
	}

	hasPendingNotification, err := c.deletePendingNotifications(a, ch, nextPeriod)
	if err != nil {
		return err
	}

	// there is not any pending notification mail for the channel
	if !hasPendingNotification {
		return nil
	}

	hasUnglancedChannels, err := c.hasUnglancedChannels(a, nextPeriod)
	if err != nil {
		return err
	}

	if hasUnglancedChannels {
		return nil
	}

	// when there is not any pending notifications (unglanced active channels),
	// just reset the account mailing period value
	return common.ResetMailingPeriodForAccount(c.redis, a)
}

// deletePendingNotification deletes notification information for given account, and channel. It returns
// true when there are pending notifications, false otherwise
func (c *Controller) deletePendingNotifications(a *models.Account, ch *models.Channel, nextPeriod string) (bool, error) {
	count, err := c.redis.DeleteHashSetField(common.AccountChannelHashSetKey(a.Id, nextPeriod), strconv.FormatInt(ch.Id, 10))
	if err != nil {
		return false, err
	}

	return count == 1, nil
}

func (c *Controller) hasUnglancedChannels(a *models.Account, nextPeriod string) (bool, error) {
	count, err := c.redis.GetHashLength(common.AccountChannelHashSetKey(a.Id, nextPeriod))
	if err != nil {
		return false, err
	}

	return count == 1, nil
}

func (c *Controller) fetchParticipantIds(cm *models.ChannelMessage) ([]int64, error) {
	ch := models.NewChannel()
	ch.Id = cm.InitialChannelId

	return ch.FetchParticipantIds(&request.Query{})
}

func (c *Controller) notifyAccount(accountId int64, cm *models.ChannelMessage) error {
	a := models.NewAccount()
	a.Id = accountId

	eligible, err := isEligibleToNotify(accountId)
	if err != nil {
		return err
	}

	if !eligible {
		return nil
	}

	nextPeriod, err := c.getOrCreateMailingPeriod(a)
	if err != nil {
		return err
	}

	if err := c.addAccountToNotifieeQueue(nextPeriod, accountId); err != nil {
		return err
	}

	return c.addMessageToAccountChannelNotifications(nextPeriod, accountId, cm)
}

// getOrCreateMailingPeriod updates the Account-Segment hash set and returns
// the next mailing period of the account
func (c *Controller) getOrCreateMailingPeriod(a *models.Account) (string, error) {
	field := strconv.FormatInt(a.Id, 10)

	nextPeriod, err := c.getMailingPeriod(a)
	// if not exist get new mailing period for the account
	if err == ErrPeriodNotFound {
		nextPeriod = common.GetNextMailPeriod()
		err := c.redis.HashMultipleSet(common.AccountNextPeriodHashSetKey(), map[string]interface{}{field: nextPeriod})
		if err != nil {
			return "", err
		}

		return nextPeriod, nil
	}

	if err != nil {
		return "", err
	}

	return nextPeriod, nil
}

func (c *Controller) getMailingPeriod(a *models.Account) (string, error) {
	period, err := c.redis.GetHashSetField(common.AccountNextPeriodHashSetKey(), strconv.FormatInt(a.Id, 10))
	if err == redis.ErrNil {
		return "", ErrPeriodNotFound
	}

	if err != nil {
		return "", err
	}

	if period == "" {
		return "", ErrPeriodNotFound
	}

	return period, nil
}

func (c *Controller) addAccountToNotifieeQueue(period string, accountId int64) error {
	_, err := c.redis.AddSetMembers(common.PeriodAccountSetKey(period), strconv.FormatInt(accountId, 10))
	if err != nil {
		return err
	}

	return nil
}

func (c *Controller) addMessageToAccountChannelNotifications(period string, accountId int64, cm *models.ChannelMessage) error {
	key := common.AccountChannelHashSetKey(accountId, period)
	channelId := strconv.FormatInt(cm.InitialChannelId, 10)
	awaySince := strconv.FormatInt(cm.CreatedAt.UnixNano(), 10)
	// add the first received message for channel
	_, err := c.redis.HashSetIfNotExists(key, channelId, awaySince)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package ratelimiter

import (
	"fmt"
	"github.com/Clever/leakybucket"
	"github.com/Clever/sphinx/common"
	"github.com/Clever/sphinx/config"
	"github.com/Clever/sphinx/limit"
	"testing"
)

func returnLastAddStatus(rateLimiter RateLimiter, request common.Request, numAdds int) ([]Status, error) {
	statuses := []Status{}
	var err error
	for i := 0; i < numAdds; i++ {
		if statuses, err = rateLimiter.Add(request); err != nil {
			return statuses, err
		}
	}
	return statuses, nil
}

func checkLastStatusForRequests(ratelimiter RateLimiter,
	request common.Request, numAdds int, expectedStatuses []Status) error {

	if statuses, err := returnLastAddStatus(ratelimiter, request, numAdds); err != nil {
		return err
	} else if len(statuses) != len(expectedStatuses) {
		return fmt.Errorf("expected to match %d buckets. Got: %d", len(expectedStatuses),
			len(statuses))
	} else {
		for i, status := range expectedStatuses {
			if status.Remaining != statuses[i].Remaining && status.Name != statuses[i].Name {
				return fmt.Errorf("expected %d remaining for the %s limit. Found: %d Remaining, %s Limit",
					statuses[i].Remaining, statuses[i].Name, status.Remaining, status.Name)
			}
		}
	}

	return nil
}

// ratelimiter is initialized properly based on config
func TestNew(t *testing.T) {

	config, err := config.New("../example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}

	rater, err := New(config)
	ratelimiter := rater.(*rateLimiter)
	if err != nil {
		t.Errorf("Error while instantiating ratelimiter: %s", err.Error())
	}
	if len(ratelimiter.limits) != len(config.Limits) {
		t.Error("expected number of limits in configuration to match instantiated limits")
	}
}

// adds different kinds of requests and checks limit Status
// focusses on single bucket adds
func TestSimpleAdd(t *testing.T) {
	config, err := config.New("../example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}
	ratelimiter, err := New(config)

	request := common.Request{
		"path": "/special/resources/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization":   []string{"Bearer 12345"},
			"X-Forwarded-For": []string{"IP1", "IP2"},
		}).Header,
		"remoteaddr": "127.0.0.1",
	}
	if err = checkLastStatusForRequests(
		ratelimiter, request, 5, []Status{
			Status{Remaining: 195, Name: "bearer-special"}}); err != nil {
		t.Error(err)
	}

	request = common.Request{
		"path": "/resources/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization": []string{"Basic 12345"},
		}).Header,
	}

	if err = checkLastStatusForRequests(
		ratelimiter, request, 1, []Status{
			Status{Remaining: 195, Name: "basic-simple"}}); err != nil {
		t.Error(err)
	}

	if status, err := returnLastAddStatus(ratelimiter, request, 200); err == nil {
		t.Fatal("expected error")
	} else if len(status) != 1 {
		t.Fatalf("expected one status, found %d", len(status))
	} else if status[0].Remaining != 0 {
		t.Fatalf("expected 0 remaining, found %d", status[0].Remaining)
	} else if status[0].Name != "basic-simple" {
		t.Fatalf("expected 'basic-simple' limit, found '%s'", status[0].Name)
	}
}

type NeverMatch struct{}

func (m NeverMatch) Name() string {
	return "name"
}
func (m NeverMatch) Match(common.Request) bool {
	return false
}
func (m NeverMatch) Add(common.Request) (leakybucket.BucketState, error) {
	return leakybucket.BucketState{}, nil
}

func createRateLimiter(numLimits int) RateLimiter {
	rateLimiter := &rateLimiter{}
	limits := []limit.Limit{}
	limit := &NeverMatch{}
	for i := 0; i < numLimits; i++ {
		limits = append(limits, limit)
	}
	rateLimiter.limits = limits
	return rateLimiter
}

var benchAdd = func(b *testing.B, numLimits int) {
	rateLimiter := createRateLimiter(numLimits)
	request := common.Request{}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rateLimiter.Add(request)
	}
}

func BenchmarkAdd1(b *testing.B) {
	benchAdd(b, 1)
}

func BenchmarkAdd100(b *testing.B) {
	benchAdd(b, 100)
}
<commit_msg>ratelimiter: test rate limit error is ErrorFull<commit_after>package ratelimiter

import (
	"fmt"
	"github.com/Clever/leakybucket"
	"github.com/Clever/sphinx/common"
	"github.com/Clever/sphinx/config"
	"github.com/Clever/sphinx/limit"
	"testing"
)

func returnLastAddStatus(rateLimiter RateLimiter, request common.Request, numAdds int) ([]Status, error) {
	statuses := []Status{}
	var err error
	for i := 0; i < numAdds; i++ {
		if statuses, err = rateLimiter.Add(request); err != nil {
			return statuses, err
		}
	}
	return statuses, nil
}

func checkLastStatusForRequests(ratelimiter RateLimiter,
	request common.Request, numAdds int, expectedStatuses []Status) error {

	if statuses, err := returnLastAddStatus(ratelimiter, request, numAdds); err != nil {
		return err
	} else if len(statuses) != len(expectedStatuses) {
		return fmt.Errorf("expected to match %d buckets. Got: %d", len(expectedStatuses),
			len(statuses))
	} else {
		for i, status := range expectedStatuses {
			if status.Remaining != statuses[i].Remaining && status.Name != statuses[i].Name {
				return fmt.Errorf("expected %d remaining for the %s limit. Found: %d Remaining, %s Limit",
					statuses[i].Remaining, statuses[i].Name, status.Remaining, status.Name)
			}
		}
	}

	return nil
}

// ratelimiter is initialized properly based on config
func TestNew(t *testing.T) {

	config, err := config.New("../example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}

	rater, err := New(config)
	ratelimiter := rater.(*rateLimiter)
	if err != nil {
		t.Errorf("Error while instantiating ratelimiter: %s", err.Error())
	}
	if len(ratelimiter.limits) != len(config.Limits) {
		t.Error("expected number of limits in configuration to match instantiated limits")
	}
}

// adds different kinds of requests and checks limit Status
// focusses on single bucket adds
func TestSimpleAdd(t *testing.T) {
	config, err := config.New("../example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}
	ratelimiter, err := New(config)

	request := common.Request{
		"path": "/special/resources/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization":   []string{"Bearer 12345"},
			"X-Forwarded-For": []string{"IP1", "IP2"},
		}).Header,
		"remoteaddr": "127.0.0.1",
	}
	if err = checkLastStatusForRequests(
		ratelimiter, request, 5, []Status{
			Status{Remaining: 195, Name: "bearer-special"}}); err != nil {
		t.Error(err)
	}

	request = common.Request{
		"path": "/resources/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization": []string{"Basic 12345"},
		}).Header,
	}

	if err = checkLastStatusForRequests(
		ratelimiter, request, 1, []Status{
			Status{Remaining: 195, Name: "basic-simple"}}); err != nil {
		t.Error(err)
	}

	if status, err := returnLastAddStatus(ratelimiter, request, 200); err == nil {
		t.Fatal("expected error")
	} else if err != leakybucket.ErrorFull {
		t.Fatalf("expected ErrorFull, received %#v", leakybucket.ErrorFull)
	} else if len(status) != 1 {
		t.Fatalf("expected one status, found %d", len(status))
	} else if status[0].Remaining != 0 {
		t.Fatalf("expected 0 remaining, found %d", status[0].Remaining)
	} else if status[0].Name != "basic-simple" {
		t.Fatalf("expected 'basic-simple' limit, found '%s'", status[0].Name)
	}
}

type NeverMatch struct{}

func (m NeverMatch) Name() string {
	return "name"
}
func (m NeverMatch) Match(common.Request) bool {
	return false
}
func (m NeverMatch) Add(common.Request) (leakybucket.BucketState, error) {
	return leakybucket.BucketState{}, nil
}

func createRateLimiter(numLimits int) RateLimiter {
	rateLimiter := &rateLimiter{}
	limits := []limit.Limit{}
	limit := &NeverMatch{}
	for i := 0; i < numLimits; i++ {
		limits = append(limits, limit)
	}
	rateLimiter.limits = limits
	return rateLimiter
}

var benchAdd = func(b *testing.B, numLimits int) {
	rateLimiter := createRateLimiter(numLimits)
	request := common.Request{}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rateLimiter.Add(request)
	}
}

func BenchmarkAdd1(b *testing.B) {
	benchAdd(b, 1)
}

func BenchmarkAdd100(b *testing.B) {
	benchAdd(b, 100)
}
<|endoftext|>


<commit_before>package round

import (
	"fmt"
	"time"
)

// Style is a spinner style. Any number of frames is allowed, each frame can be any length.
// The following Styles are supplied: Pipe, Block, Moon, Hearts.
type Style struct {
	Frames []string
	Rate   time.Duration
}

// Pipe is guaranteed to work. | 0020—007F Basic Latin.
var Pipe = Style{[]string{"|", "/", "-", "\\"}, 60 * time.Millisecond}

// Moon is a series of moon phases. 🌕 1F300—1F5FF Misc Symbols and Pictographs.
var Moon = Style{[]string{"🌑", "🌒", "🌓", "🌔", "🌕", "🌖", "🌗", "🌘"}, 90 * time.Millisecond}

// Block is a good old trusty block thing. █  2580—259F Block Elements.
var Block = Style{
	[]string{"▏", "▎", "▍", "▌", "▋", "▊", "▉", "█", "▇", "▆", "▅", "▄", "▃", "▂", "▁", ""},
	60 * time.Millisecond,
}

// Hearts is clearly the best style! 💜 1F300—1F5FF Misc Symbols and Pictographs.
var Hearts = Style{
	[]string{"💖💛💚💙💜", "💜💖💛💚💙", "💙💜💖💛💚", "💚💙💜💖💛", "💛💚💙💜💖"},
	90 * time.Millisecond,
}

// NewScroller creates a Style for a text scroller with the specified width and format.
func NewScroller(width int, format, text string) Style {
	text = fmt.Sprintf(fmt.Sprintf("%%%vv%%v%%%[1]v[1]v", width), "", text)
	s := Style{[]string{}, 90 * time.Millisecond}
	for i := 0; i < len(text)-width; i++ {
		s.Frames = append(s.Frames, fmt.Sprintf(format, text[i:i+width]))
	}
	return s
}
<commit_msg>add new style cylon plus style helpers<commit_after>package round

import (
	"fmt"
	"time"
)

// Style is a spinner style. Any number of frames is allowed, each frame can be any length.
// The following Styles are supplied: Pipe, Block, Moon, Hearts.
type Style struct {
	Frames []string
	Rate   time.Duration
}

// Block is a good old trusty block thing. █  2580—259F Block Elements.
var Block = Style{
	[]string{"▏", "▎", "▍", "▌", "▋", "▊", "▉", "█", "▇", "▆", "▅", "▄", "▃", "▂", "▁", ""},
	60 * time.Millisecond,
}

// Cylon is ominous. @ 0020—007F Basic Latin.
var Cylon = NewBounce(7, "(\x1b[31m%v\x1b[0m)", "(@)")

// Hearts is clearly the best style! 💜 1F300—1F5FF Misc Symbols and Pictographs.
var Hearts = Style{
	[]string{"💖💛💚💙💜", "💜💖💛💚💙", "💙💜💖💛💚", "💚💙💜💖💛", "💛💚💙💜💖"},
	90 * time.Millisecond,
}

// Moon is a series of moon phases. 🌕 1F300—1F5FF Misc Symbols and Pictographs.
var Moon = Style{[]string{"🌑", "🌒", "🌓", "🌔", "🌕", "🌖", "🌗", "🌘"}, 90 * time.Millisecond}

// Pipe is guaranteed to work. | 0020—007F Basic Latin.
var Pipe = Style{[]string{"|", "/", "-", "\\"}, 60 * time.Millisecond}

// NewScroller creates a Style for a text scroller with the specified width
// and format. It scrolls from right to left.
func NewScroller(width int, format, text string) Style {
	text = fmt.Sprintf(fmt.Sprintf("%%%vv%%v%%%[1]v[1]v", width), "", text)
	s := Style{[]string{}, 90 * time.Millisecond}
	for i := 0; i < len(text)-width; i++ {
		s.Frames = append(s.Frames, fmt.Sprintf(format, text[i:i+width]))
	}
	return s
}

// NewInvertedScroller creates a Style for a text scroller with the specified
// width and format. It scrolls from left to right.
func NewInvertedScroller(width int, format, text string) Style {
	text = fmt.Sprintf(fmt.Sprintf("%%%vv%%v%%%[1]v[1]v", width), "", text)
	s := Style{[]string{}, 90 * time.Millisecond}
	for i := len(text) - width; i >= 0; i-- {
		s.Frames = append(s.Frames, fmt.Sprintf(format, text[i:i+width]))
	}
	return s
}

// NewBounce creates a Style with some text that bounces back and forth.
func NewBounce(width int, format, text string) Style {
	a := NewScroller(width, format, text)
	b := NewInvertedScroller(width, format, text)
	a.Frames = append(a.Frames[1:len(a.Frames)-1], b.Frames[1:len(b.Frames)-1]...)
	return a
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/json"
	"net/http"
	"sync/atomic"
	"time"

	"../rpc"
	"../util"
)

func (s *ProxyServer) StatsIndex(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)

	hashrate, hashrate24h, totalOnline, miners := s.collectMinersStats()
	stats := map[string]interface{}{
		"miners":      miners,
		"hashrate":    hashrate,
		"hashrate24h": hashrate24h,
		"totalMiners": len(miners),
		"totalOnline": totalOnline,
		"timedOut":    len(miners) - totalOnline,
	}

	var upstreams []interface{}
	current := atomic.LoadInt32(&s.upstream)

	for i, u := range s.upstreams {
		upstream := convertUpstream(u)
		upstream["current"] = current == int32(i)
		upstreams = append(upstreams, upstream)
	}
	stats["upstreams"] = upstreams
	stats["current"] = convertUpstream(s.rpc())
	stats["url"] = "http://" + s.config.Proxy.Listen + "/miner/<diff>/<id>"

	t := s.currentBlockTemplate()
	stats["height"] = t.Height
	stats["diff"] = t.Difficulty
	stats["now"] = util.MakeTimestamp()
	stats["luck"] = s.getLuckStats()
	json.NewEncoder(w).Encode(stats)
}

func convertUpstream(u *rpc.RPCClient) map[string]interface{} {
	upstream := map[string]interface{}{
		"name":             u.Name,
		"url":              u.Url.String(),
		"pool":             u.Pool,
		"sick":             u.Sick(),
		"accepts":          atomic.LoadUint64(&u.Accepts),
		"rejects":          atomic.LoadUint64(&u.Rejects),
		"lastSubmissionAt": atomic.LoadInt64(&u.LastSubmissionAt),
	}
	return upstream
}

func (s *ProxyServer) collectMinersStats() (int64, int64, int, []interface{}) {
	now := util.MakeTimestamp()
	var result []interface{}
	totalHashrate := int64(0)
	totalHashrate24h := int64(0)
	totalOnline := 0
	window24h := 24 * time.Hour

	for m := range s.miners.Iter() {
		stats := make(map[string]interface{})
		lastBeat := m.Val.getLastBeat()
		hashrate := m.Val.hashrate(s.hashrateWindow)
		hashrate24h := m.Val.hashrate(window24h)
		totalHashrate += hashrate
		totalHashrate24h += hashrate24h
		stats["name"] = m.Key
		stats["hashrate"] = hashrate
		stats["hashrate24h"] = hashrate24h
		stats["lastBeat"] = lastBeat
		stats["validShares"] = atomic.LoadUint64(&m.Val.validShares)
		stats["invalidShares"] = atomic.LoadUint64(&m.Val.invalidShares)
		stats["accepts"] = atomic.LoadUint64(&m.Val.accepts)
		stats["rejects"] = atomic.LoadUint64(&m.Val.rejects)
		stats["ip"] = m.Val.IP

		if now-lastBeat > (int64(s.timeout/2) / 1000000) {
			stats["warning"] = true
		}
		if now-lastBeat > (int64(s.timeout) / 1000000) {
			stats["timeout"] = true
		} else {
			totalOnline++
		}
		result = append(result, stats)
	}
	return totalHashrate, totalHashrate24h, totalOnline, result
}

func (s *ProxyServer) getLuckStats() map[string]interface{} {
	now := util.MakeTimestamp()
	var variance float64
	var totalVariance float64
	var blocksCount int
	var totalBlocksCount int

	s.blocksMu.Lock()
	defer s.blocksMu.Unlock()

	for k, v := range s.blockStats {
		if k >= now-int64(s.luckWindow) {
			blocksCount++
			variance += v
		}
		if k >= now-int64(s.luckLargeWindow) {
			totalBlocksCount++
			totalVariance += v
		} else {
			delete(s.blockStats, k)
		}
	}
	if blocksCount != 0 {
		variance = variance / float64(blocksCount)
	}
	if totalBlocksCount != 0 {
		totalVariance = totalVariance / float64(totalBlocksCount)
	}
	result := make(map[string]interface{})
	result["variance"] = variance
	result["blocksCount"] = blocksCount
	result["window"] = s.config.Proxy.LuckWindow
	result["totalVariance"] = totalVariance
	result["totalBlocksCount"] = totalBlocksCount
	result["largeWindow"] = s.config.Proxy.LargeLuckWindow
	return result
}
<commit_msg>Add timestamp after all stats<commit_after>package proxy

import (
	"encoding/json"
	"net/http"
	"sync/atomic"
	"time"

	"../rpc"
	"../util"
)

func (s *ProxyServer) StatsIndex(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)

	hashrate, hashrate24h, totalOnline, miners := s.collectMinersStats()
	stats := map[string]interface{}{
		"miners":      miners,
		"hashrate":    hashrate,
		"hashrate24h": hashrate24h,
		"totalMiners": len(miners),
		"totalOnline": totalOnline,
		"timedOut":    len(miners) - totalOnline,
	}

	var upstreams []interface{}
	current := atomic.LoadInt32(&s.upstream)

	for i, u := range s.upstreams {
		upstream := convertUpstream(u)
		upstream["current"] = current == int32(i)
		upstreams = append(upstreams, upstream)
	}
	stats["upstreams"] = upstreams
	stats["current"] = convertUpstream(s.rpc())
	stats["url"] = "http://" + s.config.Proxy.Listen + "/miner/<diff>/<id>"

	t := s.currentBlockTemplate()
	stats["height"] = t.Height
	stats["diff"] = t.Difficulty
	stats["luck"] = s.getLuckStats()
	stats["now"] = util.MakeTimestamp()
	json.NewEncoder(w).Encode(stats)
}

func convertUpstream(u *rpc.RPCClient) map[string]interface{} {
	upstream := map[string]interface{}{
		"name":             u.Name,
		"url":              u.Url.String(),
		"pool":             u.Pool,
		"sick":             u.Sick(),
		"accepts":          atomic.LoadUint64(&u.Accepts),
		"rejects":          atomic.LoadUint64(&u.Rejects),
		"lastSubmissionAt": atomic.LoadInt64(&u.LastSubmissionAt),
	}
	return upstream
}

func (s *ProxyServer) collectMinersStats() (int64, int64, int, []interface{}) {
	now := util.MakeTimestamp()
	var result []interface{}
	totalHashrate := int64(0)
	totalHashrate24h := int64(0)
	totalOnline := 0
	window24h := 24 * time.Hour

	for m := range s.miners.Iter() {
		stats := make(map[string]interface{})
		lastBeat := m.Val.getLastBeat()
		hashrate := m.Val.hashrate(s.hashrateWindow)
		hashrate24h := m.Val.hashrate(window24h)
		totalHashrate += hashrate
		totalHashrate24h += hashrate24h
		stats["name"] = m.Key
		stats["hashrate"] = hashrate
		stats["hashrate24h"] = hashrate24h
		stats["lastBeat"] = lastBeat
		stats["validShares"] = atomic.LoadUint64(&m.Val.validShares)
		stats["invalidShares"] = atomic.LoadUint64(&m.Val.invalidShares)
		stats["accepts"] = atomic.LoadUint64(&m.Val.accepts)
		stats["rejects"] = atomic.LoadUint64(&m.Val.rejects)
		stats["ip"] = m.Val.IP

		if now-lastBeat > (int64(s.timeout/2) / 1000000) {
			stats["warning"] = true
		}
		if now-lastBeat > (int64(s.timeout) / 1000000) {
			stats["timeout"] = true
		} else {
			totalOnline++
		}
		result = append(result, stats)
	}
	return totalHashrate, totalHashrate24h, totalOnline, result
}

func (s *ProxyServer) getLuckStats() map[string]interface{} {
	now := util.MakeTimestamp()
	var variance float64
	var totalVariance float64
	var blocksCount int
	var totalBlocksCount int

	s.blocksMu.Lock()
	defer s.blocksMu.Unlock()

	for k, v := range s.blockStats {
		if k >= now-int64(s.luckWindow) {
			blocksCount++
			variance += v
		}
		if k >= now-int64(s.luckLargeWindow) {
			totalBlocksCount++
			totalVariance += v
		} else {
			delete(s.blockStats, k)
		}
	}
	if blocksCount != 0 {
		variance = variance / float64(blocksCount)
	}
	if totalBlocksCount != 0 {
		totalVariance = totalVariance / float64(totalBlocksCount)
	}
	result := make(map[string]interface{})
	result["variance"] = variance
	result["blocksCount"] = blocksCount
	result["window"] = s.config.Proxy.LuckWindow
	result["totalVariance"] = totalVariance
	result["totalBlocksCount"] = totalBlocksCount
	result["largeWindow"] = s.config.Proxy.LargeLuckWindow
	return result
}
<|endoftext|>


<commit_before>package migrations

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
)

const (
	am01EthereumRegistryAddressMainnet = "0x5c69ccf91eab4ef80d9929b3c1b4d5bc03eb0981"
	am01EthereumRegistryAddressRinkeby = "0x5cEF053c7b383f430FC4F4e1ea2F7D31d8e2D16C"
	am01EthereumRegistryAddressRopsten = "0x403d907982474cdd51687b09a8968346159378f3"
	am01UpVersion                      = 31
	am01DownVersion                    = 30
)

// am01 - required migration struct
type am01 struct{}

type Migration031 struct{ am01 }

// Up - upgrade the state
func (am01) Up(repoPath, dbPassword string, testnet bool) error {
	var (
		configMap        = map[string]interface{}{}
		configBytes, err = ioutil.ReadFile(path.Join(repoPath, "config"))
	)
	if err != nil {
		return fmt.Errorf("reading config: %s", err.Error())
	}

	if err = json.Unmarshal(configBytes, &configMap); err != nil {
		return fmt.Errorf("unmarshal config: %s", err.Error())
	}

	c, ok := configMap["Wallets"]
	if !ok {
		return errors.New("invalid config: missing key Wallets")
	}

	walletCfg, ok := c.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid key Wallets")
	}

	btc, ok := walletCfg["BTC"]
	if !ok {
		return errors.New("invalid config: missing BTC Wallet")
	}

	btcWalletCfg, ok := btc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BTC Wallet")
	}

	btcWalletCfg["APIPool"] = []string{"https://btc.api.openbazaar.org/api"}
	btcWalletCfg["APITestnetPool"] = []string{"https://tbtc.api.openbazaar.org/api"}

	bch, ok := walletCfg["BCH"]
	if !ok {
		return errors.New("invalid config: missing BCH Wallet")
	}

	bchWalletCfg, ok := bch.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BCH Wallet")
	}

	bchWalletCfg["APIPool"] = []string{"https://bch.api.openbazaar.org/api"}
	bchWalletCfg["APITestnetPool"] = []string{"https://tbch.api.openbazaar.org/api"}

	ltc, ok := walletCfg["LTC"]
	if !ok {
		return errors.New("invalid config: missing LTC Wallet")
	}

	ltcWalletCfg, ok := ltc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid LTC Wallet")
	}

	ltcWalletCfg["APIPool"] = []string{"https://ltc.api.openbazaar.org/api"}
	ltcWalletCfg["APITestnetPool"] = []string{"https://tltc.api.openbazaar.org/api"}

	zec, ok := walletCfg["ZEC"]
	if !ok {
		return errors.New("invalid config: missing ZEC Wallet")
	}

	zecWalletCfg, ok := zec.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ZEC Wallet")
	}

	zecWalletCfg["APIPool"] = []string{"https://zec.api.openbazaar.org/api"}
	zecWalletCfg["APITestnetPool"] = []string{"https://tzec.api.openbazaar.org/api"}

	eth, ok := walletCfg["ETH"]
	if !ok {
		return errors.New("invalid config: missing ETH Wallet")
	}

	ethWalletCfg, ok := eth.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ETH Wallet")
	}

	ethWalletCfg["APIPool"] = []string{"https://mainnet.infura.io"}
	ethWalletCfg["APITestnetPool"] = []string{"https://rinkeby.infura.io"}
	ethWalletCfg["WalletOptions"] = map[string]interface{}{
		"RegistryAddress":        am01EthereumRegistryAddressMainnet,
		"RinkebyRegistryAddress": am01EthereumRegistryAddressRinkeby,
		"RopstenRegistryAddress": am01EthereumRegistryAddressRopsten,
	}

	newConfigBytes, err := json.MarshalIndent(configMap, "", "    ")
	if err != nil {
		return fmt.Errorf("marshal migrated config: %s", err.Error())
	}

	if err := ioutil.WriteFile(path.Join(repoPath, "config"), newConfigBytes, os.ModePerm); err != nil {
		return fmt.Errorf("writing migrated config: %s", err.Error())
	}

	if err := writeRepoVer(repoPath, am01UpVersion); err != nil {
		return fmt.Errorf("bumping repover to %d: %s", am01UpVersion, err.Error())
	}
	return nil
}

// Down - downgrade/restore the state
func (am01) Down(repoPath, dbPassword string, testnet bool) error {
	var (
		configMap        = map[string]interface{}{}
		configBytes, err = ioutil.ReadFile(path.Join(repoPath, "config"))
	)
	if err != nil {
		return fmt.Errorf("reading config: %s", err.Error())
	}

	if err = json.Unmarshal(configBytes, &configMap); err != nil {
		return fmt.Errorf("unmarshal config: %s", err.Error())
	}

	c, ok := configMap["Wallets"]
	if !ok {
		return errors.New("invalid config: missing key Wallets")
	}

	walletCfg, ok := c.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid key Wallets")
	}

	btc, ok := walletCfg["BTC"]
	if !ok {
		return errors.New("invalid config: missing BTC Wallet")
	}

	btcWalletCfg, ok := btc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BTC Wallet")
	}

	btcWalletCfg["APIPool"] = []string{"https://btc.blockbook.api.openbazaar.org/api"}
	btcWalletCfg["APITestnetPool"] = []string{"https://tbtc.blockbook.api.openbazaar.org/api"}

	bch, ok := walletCfg["BCH"]
	if !ok {
		return errors.New("invalid config: missing BCH Wallet")
	}

	bchWalletCfg, ok := bch.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BCH Wallet")
	}

	bchWalletCfg["APIPool"] = []string{"https://bch.blockbook.api.openbazaar.org/api"}
	bchWalletCfg["APITestnetPool"] = []string{"https://tbch.blockbook.api.openbazaar.org/api"}

	ltc, ok := walletCfg["LTC"]
	if !ok {
		return errors.New("invalid config: missing LTC Wallet")
	}

	ltcWalletCfg, ok := ltc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid LTC Wallet")
	}

	ltcWalletCfg["APIPool"] = []string{"https://ltc.blockbook.api.openbazaar.org/api"}
	ltcWalletCfg["APITestnetPool"] = []string{"https://tltc.blockbook.api.openbazaar.org/api"}

	zec, ok := walletCfg["ZEC"]
	if !ok {
		return errors.New("invalid config: missing ZEC Wallet")
	}

	zecWalletCfg, ok := zec.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ZEC Wallet")
	}

	zecWalletCfg["APIPool"] = []string{"https://zec.blockbook.api.openbazaar.org/api"}
	zecWalletCfg["APITestnetPool"] = []string{"https://tzec.blockbook.api.openbazaar.org/api"}

	eth, ok := walletCfg["ETH"]
	if !ok {
		return errors.New("invalid config: missing ETH Wallet")
	}

	ethWalletCfg, ok := eth.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ETH Wallet")
	}

	ethWalletCfg["APIPool"] = []string{"https://mainnet.infura.io"}
	ethWalletCfg["APITestnetPool"] = []string{"https://rinkeby.infura.io"}
	ethWalletCfg["WalletOptions"] = map[string]interface{}{
		"RegistryAddress":        am01EthereumRegistryAddressMainnet,
		"RinkebyRegistryAddress": am01EthereumRegistryAddressRinkeby,
		"RopstenRegistryAddress": am01EthereumRegistryAddressRopsten,
	}

	newConfigBytes, err := json.MarshalIndent(configMap, "", "    ")
	if err != nil {
		return fmt.Errorf("marshal migrated config: %s", err.Error())
	}

	if err := ioutil.WriteFile(path.Join(repoPath, "config"), newConfigBytes, os.ModePerm); err != nil {
		return fmt.Errorf("writing migrated config: %s", err.Error())
	}

	if err := writeRepoVer(repoPath, am01DownVersion); err != nil {
		return fmt.Errorf("dropping repover to %d: %s", am01DownVersion, err.Error())
	}
	return nil
}
<commit_msg>Update the API URL for infura<commit_after>package migrations

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
)

const (
	am01EthereumRegistryAddressMainnet = "0x5c69ccf91eab4ef80d9929b3c1b4d5bc03eb0981"
	am01EthereumRegistryAddressRinkeby = "0x5cEF053c7b383f430FC4F4e1ea2F7D31d8e2D16C"
	am01EthereumRegistryAddressRopsten = "0x403d907982474cdd51687b09a8968346159378f3"
	am01UpVersion                      = 31
	am01DownVersion                    = 30
)

// am01 - required migration struct
type am01 struct{}

type Migration031 struct{ am01 }

// Up - upgrade the state
func (am01) Up(repoPath, dbPassword string, testnet bool) error {
	var (
		configMap        = map[string]interface{}{}
		configBytes, err = ioutil.ReadFile(path.Join(repoPath, "config"))
	)
	if err != nil {
		return fmt.Errorf("reading config: %s", err.Error())
	}

	if err = json.Unmarshal(configBytes, &configMap); err != nil {
		return fmt.Errorf("unmarshal config: %s", err.Error())
	}

	c, ok := configMap["Wallets"]
	if !ok {
		return errors.New("invalid config: missing key Wallets")
	}

	walletCfg, ok := c.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid key Wallets")
	}

	btc, ok := walletCfg["BTC"]
	if !ok {
		return errors.New("invalid config: missing BTC Wallet")
	}

	btcWalletCfg, ok := btc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BTC Wallet")
	}

	btcWalletCfg["APIPool"] = []string{"https://btc.api.openbazaar.org/api"}
	btcWalletCfg["APITestnetPool"] = []string{"https://tbtc.api.openbazaar.org/api"}

	bch, ok := walletCfg["BCH"]
	if !ok {
		return errors.New("invalid config: missing BCH Wallet")
	}

	bchWalletCfg, ok := bch.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BCH Wallet")
	}

	bchWalletCfg["APIPool"] = []string{"https://bch.api.openbazaar.org/api"}
	bchWalletCfg["APITestnetPool"] = []string{"https://tbch.api.openbazaar.org/api"}

	ltc, ok := walletCfg["LTC"]
	if !ok {
		return errors.New("invalid config: missing LTC Wallet")
	}

	ltcWalletCfg, ok := ltc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid LTC Wallet")
	}

	ltcWalletCfg["APIPool"] = []string{"https://ltc.api.openbazaar.org/api"}
	ltcWalletCfg["APITestnetPool"] = []string{"https://tltc.api.openbazaar.org/api"}

	zec, ok := walletCfg["ZEC"]
	if !ok {
		return errors.New("invalid config: missing ZEC Wallet")
	}

	zecWalletCfg, ok := zec.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ZEC Wallet")
	}

	zecWalletCfg["APIPool"] = []string{"https://zec.api.openbazaar.org/api"}
	zecWalletCfg["APITestnetPool"] = []string{"https://tzec.api.openbazaar.org/api"}

	eth, ok := walletCfg["ETH"]
	if !ok {
		return errors.New("invalid config: missing ETH Wallet")
	}

	ethWalletCfg, ok := eth.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ETH Wallet")
	}

	ethWalletCfg["API"] = []string{"https://mainnet.infura.io"}
	ethWalletCfg["APIPool"] = []string{"https://mainnet.infura.io"}
	ethWalletCfg["APITestnetPool"] = []string{"https://rinkeby.infura.io"}
	ethWalletCfg["WalletOptions"] = map[string]interface{}{
		"RegistryAddress":        am01EthereumRegistryAddressMainnet,
		"RinkebyRegistryAddress": am01EthereumRegistryAddressRinkeby,
		"RopstenRegistryAddress": am01EthereumRegistryAddressRopsten,
	}

	newConfigBytes, err := json.MarshalIndent(configMap, "", "    ")
	if err != nil {
		return fmt.Errorf("marshal migrated config: %s", err.Error())
	}

	if err := ioutil.WriteFile(path.Join(repoPath, "config"), newConfigBytes, os.ModePerm); err != nil {
		return fmt.Errorf("writing migrated config: %s", err.Error())
	}

	if err := writeRepoVer(repoPath, am01UpVersion); err != nil {
		return fmt.Errorf("bumping repover to %d: %s", am01UpVersion, err.Error())
	}
	return nil
}

// Down - downgrade/restore the state
func (am01) Down(repoPath, dbPassword string, testnet bool) error {
	var (
		configMap        = map[string]interface{}{}
		configBytes, err = ioutil.ReadFile(path.Join(repoPath, "config"))
	)
	if err != nil {
		return fmt.Errorf("reading config: %s", err.Error())
	}

	if err = json.Unmarshal(configBytes, &configMap); err != nil {
		return fmt.Errorf("unmarshal config: %s", err.Error())
	}

	c, ok := configMap["Wallets"]
	if !ok {
		return errors.New("invalid config: missing key Wallets")
	}

	walletCfg, ok := c.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid key Wallets")
	}

	btc, ok := walletCfg["BTC"]
	if !ok {
		return errors.New("invalid config: missing BTC Wallet")
	}

	btcWalletCfg, ok := btc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BTC Wallet")
	}

	btcWalletCfg["APIPool"] = []string{"https://btc.blockbook.api.openbazaar.org/api"}
	btcWalletCfg["APITestnetPool"] = []string{"https://tbtc.blockbook.api.openbazaar.org/api"}

	bch, ok := walletCfg["BCH"]
	if !ok {
		return errors.New("invalid config: missing BCH Wallet")
	}

	bchWalletCfg, ok := bch.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid BCH Wallet")
	}

	bchWalletCfg["APIPool"] = []string{"https://bch.blockbook.api.openbazaar.org/api"}
	bchWalletCfg["APITestnetPool"] = []string{"https://tbch.blockbook.api.openbazaar.org/api"}

	ltc, ok := walletCfg["LTC"]
	if !ok {
		return errors.New("invalid config: missing LTC Wallet")
	}

	ltcWalletCfg, ok := ltc.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid LTC Wallet")
	}

	ltcWalletCfg["APIPool"] = []string{"https://ltc.blockbook.api.openbazaar.org/api"}
	ltcWalletCfg["APITestnetPool"] = []string{"https://tltc.blockbook.api.openbazaar.org/api"}

	zec, ok := walletCfg["ZEC"]
	if !ok {
		return errors.New("invalid config: missing ZEC Wallet")
	}

	zecWalletCfg, ok := zec.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ZEC Wallet")
	}

	zecWalletCfg["APIPool"] = []string{"https://zec.blockbook.api.openbazaar.org/api"}
	zecWalletCfg["APITestnetPool"] = []string{"https://tzec.blockbook.api.openbazaar.org/api"}

	eth, ok := walletCfg["ETH"]
	if !ok {
		return errors.New("invalid config: missing ETH Wallet")
	}

	ethWalletCfg, ok := eth.(map[string]interface{})
	if !ok {
		return errors.New("invalid config: invalid ETH Wallet")
	}

	ethWalletCfg["APIPool"] = []string{"https://mainnet.infura.io"}
	ethWalletCfg["APITestnetPool"] = []string{"https://rinkeby.infura.io"}
	ethWalletCfg["WalletOptions"] = map[string]interface{}{
		"RegistryAddress":        am01EthereumRegistryAddressMainnet,
		"RinkebyRegistryAddress": am01EthereumRegistryAddressRinkeby,
		"RopstenRegistryAddress": am01EthereumRegistryAddressRopsten,
	}

	newConfigBytes, err := json.MarshalIndent(configMap, "", "    ")
	if err != nil {
		return fmt.Errorf("marshal migrated config: %s", err.Error())
	}

	if err := ioutil.WriteFile(path.Join(repoPath, "config"), newConfigBytes, os.ModePerm); err != nil {
		return fmt.Errorf("writing migrated config: %s", err.Error())
	}

	if err := writeRepoVer(repoPath, am01DownVersion); err != nil {
		return fmt.Errorf("dropping repover to %d: %s", am01DownVersion, err.Error())
	}
	return nil
}
<|endoftext|>


<commit_before><commit_msg>Remove random sleeps<commit_after><|endoftext|>


<commit_before>package parser

import (
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/m-lab/etl/web100"

	"cloud.google.com/go/bigquery"

	"github.com/m-lab/etl/annotation"
	"github.com/m-lab/etl/metrics"
	"github.com/m-lab/etl/schema"
	"github.com/prometheus/client_golang/prometheus"
)

// AddGeoDataSSConnSpec takes a pointer to a
// Web100ConnectionSpecification struct and a timestamp. With these,
// it will fetch the appropriate geo data and add it to the hop struct
// referenced by the pointer.
func AddGeoDataSSConnSpec(spec *schema.Web100ConnectionSpecification, timestamp time.Time) {
	if spec == nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "SS ConnSpec was nil!!!"}).Inc()
		return
	}
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "SS"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)

	ipSlice := []string{spec.Local_ip, spec.Remote_ip}
	geoSlice := []*annotation.GeolocationIP{&spec.Local_geolocation, &spec.Remote_geolocation}
	annotation.FetchGeoAnnotations(ipSlice, timestamp, geoSlice)
}

// AddGeoDataPTConnSpec takes a pointer to a
// MLabConnectionSpecification struct and a timestamp. With these, it
// will fetch the appropriate geo data and add it to the hop struct
// referenced by the pointer.
func AddGeoDataPTConnSpec(spec *schema.MLabConnectionSpecification, timestamp time.Time) {
	if spec == nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT ConnSpec was nil!!!"}).Inc()
		return
	}
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "PT"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	ipSlice := []string{spec.Server_ip, spec.Client_ip}
	geoSlice := []*annotation.GeolocationIP{&spec.Server_geolocation, &spec.Client_geolocation}
	annotation.FetchGeoAnnotations(ipSlice, timestamp, geoSlice)
}

// AddGeoDataPTHopBatch takes a slice of pointers to
// schema.ParisTracerouteHops and will annotate all of them or fail
// silently. It sends them all in a single remote request.
func AddGeoDataPTHopBatch(hops []*schema.ParisTracerouteHop, timestamp time.Time) {
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "PT-HOP Batch"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	requestSlice := CreateRequestDataFromPTHops(hops, timestamp)
	annotationData := annotation.GetBatchGeoData(annotation.BatchURL, requestSlice)
	AnnotatePTHops(hops, annotationData, timestamp)
}

// AnnotatePTHops takes a slice of hop pointers, the annotation data
// mapping ip addresses to geo data and a timestamp. It will then use
// these to attach the appropriate geo data to the PT hops.
func AnnotatePTHops(hops []*schema.ParisTracerouteHop, annotationData map[string]annotation.GeoData, timestamp time.Time) {
	if annotationData == nil {
		return
	}
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for _, hop := range hops {
		if hop == nil {
			continue
		}

		if data, ok := annotationData[hop.Src_ip+timeString]; ok && data.Geo != nil {
			hop.Src_geolocation = *data.Geo
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Couldn't get geo data for PT Hop!"}).Inc()
		}

		if data, ok := annotationData[hop.Dest_ip+timeString]; ok && data.Geo != nil {
			hop.Dest_geolocation = *data.Geo
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Couldn't get geo data for PT Hop!"}).Inc()
		}
	}
}

// CreateRequestDataFromPTHops will take a slice of PT hop pointers
// and the associate timestamp. From those, it will create a slice of
// requests to send to the annotation service, removing duplicates
// along the way.
func CreateRequestDataFromPTHops(hops []*schema.ParisTracerouteHop, timestamp time.Time) []annotation.RequestData {
	hopMap := map[string]annotation.RequestData{}
	for _, hop := range hops {
		if hop == nil {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "PT Hop was nil!!!"}).Inc()
			continue
		}
		if hop.Src_ip != "" {
			hopMap[hop.Src_ip] = annotation.RequestData{hop.Src_ip, 0, timestamp}
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "PT Hop was missing an IP!!!"}).Inc()
		}

		if hop.Dest_ip != "" {
			hopMap[hop.Dest_ip] = annotation.RequestData{hop.Dest_ip, 0, timestamp}
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "PT Hop was missing an IP!!!"}).Inc()
		}
	}

	requestSlice := make([]annotation.RequestData, 0, len(hopMap))
	for _, req := range hopMap {
		requestSlice = append(requestSlice, req)
	}
	return requestSlice
}

// AddGeoDataPTHop takes a pointer to a ParisTracerouteHop and a
// timestamp. With these, it will fetch the appropriate geo data and
// add it to the hop struct referenced by the pointer.
func AddGeoDataPTHop(hop *schema.ParisTracerouteHop, timestamp time.Time) {
	if hop == nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT Hop was nil!!!"}).Inc()
		return
	}
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "PT-HOP"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	if hop.Src_ip != "" {
		annotation.GetAndInsertGeolocationIPStruct(&hop.Src_geolocation, hop.Src_ip, timestamp)
	} else {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT Hop had no src_ip!"}).Inc()
	}
	if hop.Dest_ip != "" {
		annotation.GetAndInsertGeolocationIPStruct(&hop.Dest_geolocation, hop.Dest_ip, timestamp)
	} else {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT Hop had no dest_ip!"}).Inc()
	}
}

// AddGeoDataNDTConnSpec takes a connection spec and a timestamp and
// annotates the connection spec with geo data associated with each IP
// Address. It will either sucessfully add the geo data or fail
// silently and make no changes.
func AddGeoDataNDTConnSpec(spec schema.Web100ValueMap, timestamp time.Time) {
	// Only annotate if flag enabled...
	// TODO(gfr) - should propogate this to other pipelines, or push to a common
	// intercept point.
	if !annotation.IPAnnotationEnabled {
		metrics.AnnotationErrorCount.With(prometheus.Labels{
			"source": "IP Annotation Disabled."}).Inc()
		return
	}

	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "NDT"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)

	GetAndInsertTwoSidedGeoIntoNDTConnSpec(spec, timestamp)
}

// CopyStructToMap takes a POINTER to an arbitrary struct and copies
// it's fields into a value map. It will also make fields entirely
// lower case, for convienece when working with exported structs. Also,
// NEVER pass in something that is not a pointer to a struct, as this
// will cause a panic.
func CopyStructToMap(sourceStruct interface{}, destinationMap map[string]bigquery.Value) {
	structToCopy := reflect.ValueOf(sourceStruct).Elem()
	typeOfStruct := structToCopy.Type()
	for i := 0; i < typeOfStruct.NumField(); i++ {
		v := structToCopy.Field(i).Interface()
		switch t := v.(type) {
		case string:
			if t == "" {
				continue
			}
		case int64:
			if t == 0 {
				continue
			}
		}
		destinationMap[strings.ToLower(typeOfStruct.Field(i).Name)] = v
	}
}

// GetAndInsertTwoSidedGeoIntoNDTConnSpec takes a timestamp and an
// NDT connection spec. It will either insert the data into the
// connection spec or silently fail.
func GetAndInsertTwoSidedGeoIntoNDTConnSpec(spec schema.Web100ValueMap, timestamp time.Time) {
	// TODO: Make metrics for sok and cok failures. And double check metrics for cleanliness.
	cip, cok := spec.GetString([]string{"client_ip"})
	sip, sok := spec.GetString([]string{"server_ip"})
	reqData := []annotation.RequestData{}
	if cok {
		cip, _ := web100.NormalizeIPv6(cip)
		reqData = append(reqData, annotation.RequestData{IP: cip, Timestamp: timestamp})
	} else {
		metrics.AnnotationWarningCount.With(prometheus.
			Labels{"source": "Missing client side IP."}).Inc()
	}
	if sok {
		sip, _ := web100.NormalizeIPv6(sip)
		reqData = append(reqData, annotation.RequestData{IP: sip, Timestamp: timestamp})
	} else {
		metrics.AnnotationWarningCount.With(prometheus.
			Labels{"source": "Missing server side IP."}).Inc()
	}
	if cok || sok {
		annotationDataMap := annotation.GetBatchGeoData(annotation.BatchURL, reqData)
		// TODO: Revisit decision to use base36 for
		// encoding, rather than base64. (It had to do with
		// library support.)
		timeString := strconv.FormatInt(timestamp.Unix(), 36)
		if cok {
			if data, ok := annotationDataMap[cip+timeString]; ok && data.Geo != nil {
				CopyStructToMap(data.Geo, spec.Get("client_geolocation"))
			} else {
				metrics.AnnotationErrorCount.With(prometheus.
					Labels{"source": "Couldn't get geo data for the client side."}).Inc()
			}
		}
		if sok {
			if data, ok := annotationDataMap[sip+timeString]; ok && data.Geo != nil {
				CopyStructToMap(data.Geo, spec.Get("server_geolocation"))
			} else {
				metrics.AnnotationErrorCount.With(prometheus.
					Labels{"source": "Couldn't get geo data for the server side."}).Inc()
			}

		}
	}

}
<commit_msg>Use NormalizeIPv6 for PT<commit_after>package parser

import (
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/m-lab/etl/web100"

	"cloud.google.com/go/bigquery"

	"github.com/m-lab/etl/annotation"
	"github.com/m-lab/etl/metrics"
	"github.com/m-lab/etl/schema"
	"github.com/prometheus/client_golang/prometheus"
)

// AddGeoDataSSConnSpec takes a pointer to a
// Web100ConnectionSpecification struct and a timestamp. With these,
// it will fetch the appropriate geo data and add it to the hop struct
// referenced by the pointer.
func AddGeoDataSSConnSpec(spec *schema.Web100ConnectionSpecification, timestamp time.Time) {
	if spec == nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "SS ConnSpec was nil!!!"}).Inc()
		return
	}
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "SS"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)

	ipSlice := []string{spec.Local_ip, spec.Remote_ip}
	geoSlice := []*annotation.GeolocationIP{&spec.Local_geolocation, &spec.Remote_geolocation}
	annotation.FetchGeoAnnotations(ipSlice, timestamp, geoSlice)
}

// AddGeoDataPTConnSpec takes a pointer to a
// MLabConnectionSpecification struct and a timestamp. With these, it
// will fetch the appropriate geo data and add it to the hop struct
// referenced by the pointer.
func AddGeoDataPTConnSpec(spec *schema.MLabConnectionSpecification, timestamp time.Time) {
	if spec == nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT ConnSpec was nil!!!"}).Inc()
		return
	}
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "PT"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	ipSlice := []string{spec.Server_ip, spec.Client_ip}
	geoSlice := []*annotation.GeolocationIP{&spec.Server_geolocation, &spec.Client_geolocation}
	annotation.FetchGeoAnnotations(ipSlice, timestamp, geoSlice)
}

// AddGeoDataPTHopBatch takes a slice of pointers to
// schema.ParisTracerouteHops and will annotate all of them or fail
// silently. It sends them all in a single remote request.
func AddGeoDataPTHopBatch(hops []*schema.ParisTracerouteHop, timestamp time.Time) {
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "PT-HOP Batch"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	requestSlice := CreateRequestDataFromPTHops(hops, timestamp)
	annotationData := annotation.GetBatchGeoData(annotation.BatchURL, requestSlice)
	AnnotatePTHops(hops, annotationData, timestamp)
}

// AnnotatePTHops takes a slice of hop pointers, the annotation data
// mapping ip addresses to geo data and a timestamp. It will then use
// these to attach the appropriate geo data to the PT hops.
func AnnotatePTHops(hops []*schema.ParisTracerouteHop, annotationData map[string]annotation.GeoData, timestamp time.Time) {
	if annotationData == nil {
		return
	}
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for _, hop := range hops {
		if hop == nil {
			continue
		}

		if data, ok := annotationData[hop.Src_ip+timeString]; ok && data.Geo != nil {
			hop.Src_geolocation = *data.Geo
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Couldn't get geo data for PT Hop!"}).Inc()
		}

		if data, ok := annotationData[hop.Dest_ip+timeString]; ok && data.Geo != nil {
			hop.Dest_geolocation = *data.Geo
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Couldn't get geo data for PT Hop!"}).Inc()
		}
	}
}

// CreateRequestDataFromPTHops will take a slice of PT hop pointers
// and the associate timestamp. From those, it will create a slice of
// requests to send to the annotation service, removing duplicates
// along the way.
func CreateRequestDataFromPTHops(hops []*schema.ParisTracerouteHop, timestamp time.Time) []annotation.RequestData {
	hopMap := map[string]annotation.RequestData{}
	for _, hop := range hops {
		if hop == nil {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "PT Hop was nil!!!"}).Inc()
			continue
		}
		if hop.Src_ip != "" {
			hop.Src_ip, _ = web100.NormalizeIPv6(hop.Src_ip)
			hopMap[hop.Src_ip] = annotation.RequestData{hop.Src_ip, 0, timestamp}
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "PT Hop was missing an IP!!!"}).Inc()
		}

		if hop.Dest_ip != "" {
			hop.Dest_ip, _ = web100.NormalizeIPv6(hop.Dest_ip)
			hopMap[hop.Dest_ip] = annotation.RequestData{hop.Dest_ip, 0, timestamp}
		} else {
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "PT Hop was missing an IP!!!"}).Inc()
		}
	}

	requestSlice := make([]annotation.RequestData, 0, len(hopMap))
	for _, req := range hopMap {
		requestSlice = append(requestSlice, req)
	}
	return requestSlice
}

// AddGeoDataPTHop takes a pointer to a ParisTracerouteHop and a
// timestamp. With these, it will fetch the appropriate geo data and
// add it to the hop struct referenced by the pointer.
func AddGeoDataPTHop(hop *schema.ParisTracerouteHop, timestamp time.Time) {
	if hop == nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT Hop was nil!!!"}).Inc()
		return
	}
	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "PT-HOP"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	if hop.Src_ip != "" {
		annotation.GetAndInsertGeolocationIPStruct(&hop.Src_geolocation, hop.Src_ip, timestamp)
	} else {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT Hop had no src_ip!"}).Inc()
	}
	if hop.Dest_ip != "" {
		annotation.GetAndInsertGeolocationIPStruct(&hop.Dest_geolocation, hop.Dest_ip, timestamp)
	} else {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "PT Hop had no dest_ip!"}).Inc()
	}
}

// AddGeoDataNDTConnSpec takes a connection spec and a timestamp and
// annotates the connection spec with geo data associated with each IP
// Address. It will either sucessfully add the geo data or fail
// silently and make no changes.
func AddGeoDataNDTConnSpec(spec schema.Web100ValueMap, timestamp time.Time) {
	// Only annotate if flag enabled...
	// TODO(gfr) - should propogate this to other pipelines, or push to a common
	// intercept point.
	if !annotation.IPAnnotationEnabled {
		metrics.AnnotationErrorCount.With(prometheus.Labels{
			"source": "IP Annotation Disabled."}).Inc()
		return
	}

	// Time the response
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.AnnotationTimeSummary.
			With(prometheus.Labels{"test_type": "NDT"}).
			Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)

	GetAndInsertTwoSidedGeoIntoNDTConnSpec(spec, timestamp)
}

// CopyStructToMap takes a POINTER to an arbitrary struct and copies
// it's fields into a value map. It will also make fields entirely
// lower case, for convienece when working with exported structs. Also,
// NEVER pass in something that is not a pointer to a struct, as this
// will cause a panic.
func CopyStructToMap(sourceStruct interface{}, destinationMap map[string]bigquery.Value) {
	structToCopy := reflect.ValueOf(sourceStruct).Elem()
	typeOfStruct := structToCopy.Type()
	for i := 0; i < typeOfStruct.NumField(); i++ {
		v := structToCopy.Field(i).Interface()
		switch t := v.(type) {
		case string:
			if t == "" {
				continue
			}
		case int64:
			if t == 0 {
				continue
			}
		}
		destinationMap[strings.ToLower(typeOfStruct.Field(i).Name)] = v
	}
}

// GetAndInsertTwoSidedGeoIntoNDTConnSpec takes a timestamp and an
// NDT connection spec. It will either insert the data into the
// connection spec or silently fail.
func GetAndInsertTwoSidedGeoIntoNDTConnSpec(spec schema.Web100ValueMap, timestamp time.Time) {
	// TODO: Make metrics for sok and cok failures. And double check metrics for cleanliness.
	cip, cok := spec.GetString([]string{"client_ip"})
	sip, sok := spec.GetString([]string{"server_ip"})
	reqData := []annotation.RequestData{}
	if cok {
		cip, _ := web100.NormalizeIPv6(cip)
		reqData = append(reqData, annotation.RequestData{IP: cip, Timestamp: timestamp})
	} else {
		metrics.AnnotationWarningCount.With(prometheus.
			Labels{"source": "Missing client side IP."}).Inc()
	}
	if sok {
		sip, _ := web100.NormalizeIPv6(sip)
		reqData = append(reqData, annotation.RequestData{IP: sip, Timestamp: timestamp})
	} else {
		metrics.AnnotationWarningCount.With(prometheus.
			Labels{"source": "Missing server side IP."}).Inc()
	}
	if cok || sok {
		annotationDataMap := annotation.GetBatchGeoData(annotation.BatchURL, reqData)
		// TODO: Revisit decision to use base36 for
		// encoding, rather than base64. (It had to do with
		// library support.)
		timeString := strconv.FormatInt(timestamp.Unix(), 36)
		if cok {
			if data, ok := annotationDataMap[cip+timeString]; ok && data.Geo != nil {
				CopyStructToMap(data.Geo, spec.Get("client_geolocation"))
			} else {
				metrics.AnnotationErrorCount.With(prometheus.
					Labels{"source": "Couldn't get geo data for the client side."}).Inc()
			}
		}
		if sok {
			if data, ok := annotationDataMap[sip+timeString]; ok && data.Geo != nil {
				CopyStructToMap(data.Geo, spec.Get("server_geolocation"))
			} else {
				metrics.AnnotationErrorCount.With(prometheus.
					Labels{"source": "Couldn't get geo data for the server side."}).Inc()
			}

		}
	}

}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 flow

import (
	"math/rand"
	"net"
	"testing"
	"time"

	"encoding/json"
	j "github.com/gima/jsonv/src"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

func forgeEthIPTCP(t *testing.T, seed int64) *gopacket.Packet {
	var options gopacket.SerializeOptions
	rnd := rand.New(rand.NewSource(seed))

	rawBytes := []byte{10, 20, 30}
	ethernetLayer := &layers.Ethernet{
		SrcMAC: net.HardwareAddr{0x00, 0x0F, 0xAA, 0xFA, 0xAA, byte(rnd.Intn(0x100))},
		DstMAC: net.HardwareAddr{0x00, 0x0D, 0xBD, 0xBD, byte(rnd.Intn(0x100)), 0xBD},
	}
	ipLayer := &layers.IPv4{
		SrcIP: net.IP{127, 0, 0, byte(rnd.Intn(0x100))},
		DstIP: net.IP{byte(rnd.Intn(0x100)), 8, 8, 8},
	}
	tcpLayer := &layers.TCP{
		SrcPort: layers.TCPPort(byte(rnd.Intn(0x10000))),
		DstPort: layers.TCPPort(byte(rnd.Intn(0x10000))),
	}
	// And create the packet with the layers
	buffer := gopacket.NewSerializeBuffer()
	err := gopacket.SerializeLayers(buffer, options,
		ethernetLayer,
		ipLayer,
		tcpLayer,
		gopacket.Payload(rawBytes),
	)
	if err != nil {
		t.Fail()
	}

	gpacket := gopacket.NewPacket(buffer.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
	return &gpacket

}

func generateFlows(t *testing.T, ft *FlowTable) []*Flow {
	flows := []*Flow{}
	for i := 0; i < 10; i++ {
		packet := forgeEthIPTCP(t, int64(i))
		flow, new := ft.GetFlow(string(i))
		if !new {
			t.Fail()
		}
		err := flow.fillFromGoPacket(packet)
		if err != nil {
			t.Error("fillFromGoPacket : " + err.Error())
		}
		flows = append(flows, flow)
	}
	return flows
}

func checkFlowTable(t *testing.T, ft *FlowTable) {
	for uuid, f := range ft.table {
		if uuid != f.UUID {
			t.Error("FlowTable Collision ", uuid, f.UUID)
		}
	}
}

func NewFlowTableSimple(t *testing.T) *FlowTable {
	ft := NewFlowTable()
	var flows []*Flow
	flow := &Flow{}
	flow.UUID = "1234"
	flows = append(flows, flow)
	ft.Update(flows)
	checkFlowTable(t, ft)
	ft.Update(flows)
	checkFlowTable(t, ft)
	if "1 flows" != ft.String() {
		t.Error("We should got only 1 flow")
	}
	flow = &Flow{}
	flow.UUID = "4567"
	flows = append(flows, flow)
	ft.Update(flows)
	checkFlowTable(t, ft)
	ft.Update(flows)
	checkFlowTable(t, ft)
	if "2 flows" != ft.String() {
		t.Error("We should got only 2 flows")
	}
	return ft
}
func NewFlowTableComplex(t *testing.T) *FlowTable {
	ft := NewFlowTable()
	flows := generateFlows(t, ft)
	ft = ft.NewFlowTableFromFlows(flows)
	checkFlowTable(t, ft)
	return ft
}

func TestNewFlowTable(t *testing.T) {
	ft := NewFlowTable()
	if ft == nil {
		t.Error("new FlowTable return null")
	}
}
func TestFlowTable_String(t *testing.T) {
	ft := NewFlowTable()
	if "0 flows" != ft.String() {
		t.Error("FlowTable too big")
	}
	ft = NewFlowTableSimple(t)
	if "2 flows" != ft.String() {
		t.Error("FlowTable too big")
	}
}
func TestFlowTable_Update(t *testing.T) {
	ft := NewFlowTableSimple(t)
	/* simulate a collision */
	f := &Flow{}
	ft.table["789"] = f
	ft.table["789"].UUID = "78910"
	f = &Flow{}
	f.UUID = "789"
	ft.Update([]*Flow{f})
	NewFlowTableComplex(t)
}

type MyTestFlowCounter struct {
	NbFlow int
}

func (fo *MyTestFlowCounter) expireCallback(f *Flow) {
	fo.NbFlow++
}

func TestFlowTable_expire(t *testing.T) {
	const MaxInt64 = int64(^uint64(0) >> 1)
	ft := NewFlowTableComplex(t)

	fc := MyTestFlowCounter{}
	beforeNbFlow := fc.NbFlow
	ft.expire(fc.expireCallback, 0)
	afterNbFlow := fc.NbFlow
	if beforeNbFlow != 0 && afterNbFlow != 0 {
		t.Error("we should not expire a flow")
	}

	fc = MyTestFlowCounter{}
	beforeNbFlow = fc.NbFlow
	ft.expire(fc.expireCallback, MaxInt64)
	afterNbFlow = fc.NbFlow
	if beforeNbFlow != 0 && afterNbFlow != 10 {
		t.Error("we should expire all flows")
	}
}

func TestFlowTable_AsyncExpire(t *testing.T) {
	t.Skip()
}

func TestFlowTable_IsExist(t *testing.T) {
	ft := NewFlowTableSimple(t)
	flow := &Flow{}
	flow.UUID = "1234"
	if ft.IsExist(flow) == false {
		t.Fail()
	}
	flow.UUID = "12345"
	if ft.IsExist(flow) == true {
		t.Fail()
	}
}

func TestFlowTable_GetFlow(t *testing.T) {
	ft := NewFlowTableComplex(t)
	flows := generateFlows(t, ft)
	if len(flows) != 10 {
		t.Error("missing some flows ", len(flows))
	}
	forgeEthIPTCP(t, int64(1234))
	_, new := ft.GetFlow("abcd")
	if !new {
		t.Error("Collision in the FlowTable, should be new")
	}
	forgeEthIPTCP(t, int64(1234))
	_, new = ft.GetFlow("abcd")
	if new {
		t.Error("Collision in the FlowTable, should be an update")
	}
	forgeEthIPTCP(t, int64(1234))
	_, new = ft.GetFlow("abcde")
	if !new {
		t.Error("Collision in the FlowTable, should be a new flow")
	}
}

func isJSON(s string) bool {
	var js map[string]interface{}
	return json.Unmarshal([]byte(s), &js) == nil
}

func TestFlowTable_JSONFlowConversationEthernetPath(t *testing.T) {
	ft := NewFlowTableComplex(t)
	statStr := ft.JSONFlowConversationEthernetPath()
	if statStr == `{"nodes":[],"links":[]}` {
		t.Error("stat should not be empty")
	}
	if !isJSON(statStr) {
		t.Error("stat should be JSON format")
	}

	decoded := new(interface{})
	if err := json.Unmarshal([]byte(statStr), decoded); err != nil {
		t.Error("JSON parsing failed:", err)
	}

	schema := &j.Object{Properties: []j.ObjectItem{
		{"nodes", &j.Array{Each: &j.Object{Properties: []j.ObjectItem{
			{"name", &j.String{MinLen: 1}},
			{"group", &j.Number{Min: 0, Max: 20}}}}},
		},
		{"links", &j.Array{Each: &j.Object{Properties: []j.ObjectItem{
			{"source", &j.Number{Min: 0, Max: 20}},
			{"target", &j.Number{Min: 0, Max: 20}},
			{"value", &j.Number{Min: 0, Max: 9999}}}}},
		},
	}}
	if path, err := schema.Validate(decoded); err != nil {
		t.Errorf("Failed (%s). Path: %s", err, path)
	}
}

func TestFlowTable_NewFlowTableFromFlows(t *testing.T) {
	ft := NewFlowTableComplex(t)
	var flows []*Flow
	for _, f := range ft.table {
		flow := *f
		flows = append(flows, &flow)
	}
	ft2 := ft.NewFlowTableFromFlows(flows)
	if len(ft.table) != len(ft2.table) {
		t.Error("NewFlowTable(copy) are not the same size")
	}
	flows = flows[:0]
	for _, f := range ft.table {
		flows = append(flows, f)
	}
	ft3 := ft.NewFlowTableFromFlows(flows)
	if len(ft.table) != len(ft3.table) {
		t.Error("NewFlowTable(ref) are not the same size")
	}
}

func TestFlowTable_FilterLast(t *testing.T) {
	ft := NewFlowTableComplex(t)
	/* hack to put the FlowTable 1 second older */
	for _, f := range ft.table {
		fs := f.GetStatistics()
		fs.Start -= int64(1)
		fs.Last -= int64(1)
	}
	flows := ft.FilterLast(10 * time.Minute)
	if len(flows) != 10 {
		t.Error("FilterLast should return more/less flows", len(flows), 10)
	}
	flows = ft.FilterLast(0 * time.Minute)
	if len(flows) != 0 {
		t.Error("FilterLast should return less flows", len(flows), 0)
	}
}
<commit_msg>[tests] FlowTable remove unused isJSON()<commit_after>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 flow

import (
	"math/rand"
	"net"
	"testing"
	"time"

	"encoding/json"
	j "github.com/gima/jsonv/src"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

func forgeEthIPTCP(t *testing.T, seed int64) *gopacket.Packet {
	var options gopacket.SerializeOptions
	rnd := rand.New(rand.NewSource(seed))

	rawBytes := []byte{10, 20, 30}
	ethernetLayer := &layers.Ethernet{
		SrcMAC: net.HardwareAddr{0x00, 0x0F, 0xAA, 0xFA, 0xAA, byte(rnd.Intn(0x100))},
		DstMAC: net.HardwareAddr{0x00, 0x0D, 0xBD, 0xBD, byte(rnd.Intn(0x100)), 0xBD},
	}
	ipLayer := &layers.IPv4{
		SrcIP: net.IP{127, 0, 0, byte(rnd.Intn(0x100))},
		DstIP: net.IP{byte(rnd.Intn(0x100)), 8, 8, 8},
	}
	tcpLayer := &layers.TCP{
		SrcPort: layers.TCPPort(byte(rnd.Intn(0x10000))),
		DstPort: layers.TCPPort(byte(rnd.Intn(0x10000))),
	}
	// And create the packet with the layers
	buffer := gopacket.NewSerializeBuffer()
	err := gopacket.SerializeLayers(buffer, options,
		ethernetLayer,
		ipLayer,
		tcpLayer,
		gopacket.Payload(rawBytes),
	)
	if err != nil {
		t.Fail()
	}

	gpacket := gopacket.NewPacket(buffer.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
	return &gpacket

}

func generateFlows(t *testing.T, ft *FlowTable) []*Flow {
	flows := []*Flow{}
	for i := 0; i < 10; i++ {
		packet := forgeEthIPTCP(t, int64(i))
		flow, new := ft.GetFlow(string(i))
		if !new {
			t.Fail()
		}
		err := flow.fillFromGoPacket(packet)
		if err != nil {
			t.Error("fillFromGoPacket : " + err.Error())
		}
		flows = append(flows, flow)
	}
	return flows
}

func checkFlowTable(t *testing.T, ft *FlowTable) {
	for uuid, f := range ft.table {
		if uuid != f.UUID {
			t.Error("FlowTable Collision ", uuid, f.UUID)
		}
	}
}

func NewFlowTableSimple(t *testing.T) *FlowTable {
	ft := NewFlowTable()
	var flows []*Flow
	flow := &Flow{}
	flow.UUID = "1234"
	flows = append(flows, flow)
	ft.Update(flows)
	checkFlowTable(t, ft)
	ft.Update(flows)
	checkFlowTable(t, ft)
	if "1 flows" != ft.String() {
		t.Error("We should got only 1 flow")
	}
	flow = &Flow{}
	flow.UUID = "4567"
	flows = append(flows, flow)
	ft.Update(flows)
	checkFlowTable(t, ft)
	ft.Update(flows)
	checkFlowTable(t, ft)
	if "2 flows" != ft.String() {
		t.Error("We should got only 2 flows")
	}
	return ft
}
func NewFlowTableComplex(t *testing.T) *FlowTable {
	ft := NewFlowTable()
	flows := generateFlows(t, ft)
	ft = ft.NewFlowTableFromFlows(flows)
	checkFlowTable(t, ft)
	return ft
}

func TestNewFlowTable(t *testing.T) {
	ft := NewFlowTable()
	if ft == nil {
		t.Error("new FlowTable return null")
	}
}
func TestFlowTable_String(t *testing.T) {
	ft := NewFlowTable()
	if "0 flows" != ft.String() {
		t.Error("FlowTable too big")
	}
	ft = NewFlowTableSimple(t)
	if "2 flows" != ft.String() {
		t.Error("FlowTable too big")
	}
}
func TestFlowTable_Update(t *testing.T) {
	ft := NewFlowTableSimple(t)
	/* simulate a collision */
	f := &Flow{}
	ft.table["789"] = f
	ft.table["789"].UUID = "78910"
	f = &Flow{}
	f.UUID = "789"
	ft.Update([]*Flow{f})
	NewFlowTableComplex(t)
}

type MyTestFlowCounter struct {
	NbFlow int
}

func (fo *MyTestFlowCounter) expireCallback(f *Flow) {
	fo.NbFlow++
}

func TestFlowTable_expire(t *testing.T) {
	const MaxInt64 = int64(^uint64(0) >> 1)
	ft := NewFlowTableComplex(t)

	fc := MyTestFlowCounter{}
	beforeNbFlow := fc.NbFlow
	ft.expire(fc.expireCallback, 0)
	afterNbFlow := fc.NbFlow
	if beforeNbFlow != 0 && afterNbFlow != 0 {
		t.Error("we should not expire a flow")
	}

	fc = MyTestFlowCounter{}
	beforeNbFlow = fc.NbFlow
	ft.expire(fc.expireCallback, MaxInt64)
	afterNbFlow = fc.NbFlow
	if beforeNbFlow != 0 && afterNbFlow != 10 {
		t.Error("we should expire all flows")
	}
}

func TestFlowTable_AsyncExpire(t *testing.T) {
	t.Skip()
}

func TestFlowTable_IsExist(t *testing.T) {
	ft := NewFlowTableSimple(t)
	flow := &Flow{}
	flow.UUID = "1234"
	if ft.IsExist(flow) == false {
		t.Fail()
	}
	flow.UUID = "12345"
	if ft.IsExist(flow) == true {
		t.Fail()
	}
}

func TestFlowTable_GetFlow(t *testing.T) {
	ft := NewFlowTableComplex(t)
	flows := generateFlows(t, ft)
	if len(flows) != 10 {
		t.Error("missing some flows ", len(flows))
	}
	forgeEthIPTCP(t, int64(1234))
	_, new := ft.GetFlow("abcd")
	if !new {
		t.Error("Collision in the FlowTable, should be new")
	}
	forgeEthIPTCP(t, int64(1234))
	_, new = ft.GetFlow("abcd")
	if new {
		t.Error("Collision in the FlowTable, should be an update")
	}
	forgeEthIPTCP(t, int64(1234))
	_, new = ft.GetFlow("abcde")
	if !new {
		t.Error("Collision in the FlowTable, should be a new flow")
	}
}

func TestFlowTable_JSONFlowConversationEthernetPath(t *testing.T) {
	ft := NewFlowTableComplex(t)
	statStr := ft.JSONFlowConversationEthernetPath()
	if statStr == `{"nodes":[],"links":[]}` {
		t.Error("stat should not be empty")
	}

	decoded := new(interface{})
	if err := json.Unmarshal([]byte(statStr), decoded); err != nil {
		t.Error("JSON parsing failed:", err)
	}

	schema := &j.Object{Properties: []j.ObjectItem{
		{"nodes", &j.Array{Each: &j.Object{Properties: []j.ObjectItem{
			{"name", &j.String{MinLen: 1}},
			{"group", &j.Number{Min: 0, Max: 20}}}}},
		},
		{"links", &j.Array{Each: &j.Object{Properties: []j.ObjectItem{
			{"source", &j.Number{Min: 0, Max: 20}},
			{"target", &j.Number{Min: 0, Max: 20}},
			{"value", &j.Number{Min: 0, Max: 9999}}}}},
		},
	}}
	if path, err := schema.Validate(decoded); err != nil {
		t.Errorf("Failed (%s). Path: %s", err, path)
	}
}

func TestFlowTable_NewFlowTableFromFlows(t *testing.T) {
	ft := NewFlowTableComplex(t)
	var flows []*Flow
	for _, f := range ft.table {
		flow := *f
		flows = append(flows, &flow)
	}
	ft2 := ft.NewFlowTableFromFlows(flows)
	if len(ft.table) != len(ft2.table) {
		t.Error("NewFlowTable(copy) are not the same size")
	}
	flows = flows[:0]
	for _, f := range ft.table {
		flows = append(flows, f)
	}
	ft3 := ft.NewFlowTableFromFlows(flows)
	if len(ft.table) != len(ft3.table) {
		t.Error("NewFlowTable(ref) are not the same size")
	}
}

func TestFlowTable_FilterLast(t *testing.T) {
	ft := NewFlowTableComplex(t)
	/* hack to put the FlowTable 1 second older */
	for _, f := range ft.table {
		fs := f.GetStatistics()
		fs.Start -= int64(1)
		fs.Last -= int64(1)
	}
	flows := ft.FilterLast(10 * time.Minute)
	if len(flows) != 10 {
		t.Error("FilterLast should return more/less flows", len(flows), 10)
	}
	flows = ft.FilterLast(0 * time.Minute)
	if len(flows) != 0 {
		t.Error("FilterLast should return less flows", len(flows), 0)
	}
}
<|endoftext|>


<commit_before>package buildserver

import (
	"net/http"
	"strconv"

	"github.com/pivotal-golang/lager"
)

func (s *Server) AbortBuild(w http.ResponseWriter, r *http.Request) {
	buildID, err := strconv.Atoi(r.FormValue(":build_id"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	aLog := s.logger.Session("abort", lager.Data{
		"build": buildID,
	})

	build, found, err := s.db.GetBuild(buildID)
	if err != nil {
		aLog.Error("failed-to-get-build", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if !found {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	engineBuild, err := s.engine.LookupBuild(aLog, build)
	if err != nil {
		aLog.Error("failed-to-lookup-build", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	err = engineBuild.Abort(aLog)
	if err != nil {
		aLog.Error("failed-to-unmarshal-metadata", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
<commit_msg>fix misleading log line<commit_after>package buildserver

import (
	"net/http"
	"strconv"

	"github.com/pivotal-golang/lager"
)

func (s *Server) AbortBuild(w http.ResponseWriter, r *http.Request) {
	buildID, err := strconv.Atoi(r.FormValue(":build_id"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	aLog := s.logger.Session("abort", lager.Data{
		"build": buildID,
	})

	build, found, err := s.db.GetBuild(buildID)
	if err != nil {
		aLog.Error("failed-to-get-build", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if !found {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	engineBuild, err := s.engine.LookupBuild(aLog, build)
	if err != nil {
		aLog.Error("failed-to-lookup-build", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	err = engineBuild.Abort(aLog)
	if err != nil {
		aLog.Error("failed-to-abort-build", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
<|endoftext|>