<commit_before>package session

/*
Pacage Session contains an API Binding, and
interfaces that will assist in accessing the
default Session or implementing your own Custom Session.
*/
import (
	"encoding/json"
	"fmt"

	. "github.com/miketheprogrammer/go-thrust/commands"
	. "github.com/miketheprogrammer/go-thrust/common"
	"github.com/miketheprogrammer/go-thrust/connection"
	"github.com/miketheprogrammer/go-thrust/dispatcher"
)

/*
Session is the core API Binding object used to communicate with Thrust.
*/
type Session struct {
	TargetID                 uint
	CookieStore              bool
	OffTheRecord             bool
	Path                     bool
	Ready                    bool
	CommandHistory           []*Command
	ResponseHistory          []*CommandResponse
	WaitingResponses         []*Command
	CommandQueue             []*Command
	SendChannel              *connection.In
	SessionOverrideInterface SessionInvokable
}

/*
NewSession is a constructor that takes 3 arguments,
incognito which is a boolean, meaning dont persist session state
after close.
overrideDefaultSession which is a boolean that till tell thrust core
to try to invoke session methods from us.
path string is the path to store session data.
*/
func NewSession(incognito, overrideDefaultSession bool, path string) *Session {
	session := Session{
		CookieStore:  overrideDefaultSession,
		OffTheRecord: incognito,
		Path:         path,
	}
	if overrideDefaultSession == true {
		session.SetInvokable(*NewDummySession())
	}
	command := Command{
		Action:     "create",
		ObjectType: "session",
		Args: CommandArguments{
			CookieStore:  session.CookieStore,
			OffTheRecord: session.OffTheRecord,
			Path:         session.Path,
		},
	}
	session.SendChannel = connection.GetInputChannels()
	session.WaitingResponses = append(session.WaitingResponses, &command)
	session.Send(&command)
	dispatcher.RegisterHandler(session.DispatchResponse)
	return &session
}

func (session *Session) HandleInvoke(reply CommandResponse) {
	if reply.TargetID == session.TargetID {
		response := &CommandResponse{
			Action: "reply",
			ID:     reply.ID,
			Result: ReplyResult{},
		}
		switch reply.Method {
		case "cookies_load":
			cookies := session.SessionOverrideInterface.InvokeCookiesLoad(&reply.Args, session)
			marshaledCookies, err := json.Marshal(cookies)
			if err != nil {
				Log.Error(err)
			} else {
				response.Result.Cookies = marshaledCookies
			}
		case "cookies_load_for_key":
			cookies := session.SessionOverrideInterface.InvokeCookiesLoadForKey(&reply.Args, session)
			marshaledCookies, err := json.Marshal(cookies)
			if err != nil {
				Log.Error(err)
			} else {
				response.Result.Cookies = marshaledCookies
			}
		case "cookies_flush":
			session.SessionOverrideInterface.InvokeCookiesFlush(&reply.Args, session)
		case "cookies_add":
			session.SessionOverrideInterface.InvokeCookiesAdd(&reply.Args, session)
		case "cookies_delete":
			session.SessionOverrideInterface.InvokeCookiesDelete(&reply.Args, session)
		case "cookies_update_access_time":
			session.SessionOverrideInterface.InvokeCookiesUpdateAccessTime(&reply.Args, session)
		case "cookies_force_keep_session_state":
			session.SessionOverrideInterface.InvokeCookieForceKeepSessionState(&reply.Args, session)
		}
		Log.Info("Sending Response to Invoke")
		session.SendChannel.CommandResponses <- response
	}
}

func (session *Session) HandleReply(reply CommandResponse) {
	fmt.Println(reply)
	for k, command := range session.WaitingResponses {
		if command.ID != reply.ID {
			continue
		}
		if command.ID == reply.ID {
			Log.Debug("Window(", session.TargetID, ")::Handling Reply::", reply)
			if len(session.WaitingResponses) > 1 {
				// Remove the element at index k
				session.WaitingResponses = session.WaitingResponses[:k+copy(session.WaitingResponses[k:], session.WaitingResponses[k+1:])]
			} else {
				// Just initialize to empty splice literal
				session.WaitingResponses = []*Command{}
			}
			fmt.Println("session", session.TargetID, command.Action, reply.Result.TargetID)
			if session.TargetID == 0 && command.Action == "create" {
				if reply.Result.TargetID != 0 {
					session.TargetID = reply.Result.TargetID
					Log.Debug("Session:: Received TargetID(", session.TargetID, ") :: Setting Ready State")
					session.Ready = true
				}
			}
		}
	}
}

func (session *Session) DispatchResponse(response CommandResponse) {
	switch response.Action {
	case "invoke":
		session.HandleInvoke(response)
	case "reply":
		session.HandleReply(response)

	}
}

func (session *Session) Send(command *Command) {
	session.SendChannel.Commands <- command
}

func (session *Session) SetInvokable(si SessionInvokable) {
	session.SessionOverrideInterface = si
}
<commit_msg>Fix session problem<commit_after>package session

/*
Pacage Session contains an API Binding, and
interfaces that will assist in accessing the
default Session or implementing your own Custom Session.
*/
import (
	"encoding/json"
	"fmt"

	. "github.com/miketheprogrammer/go-thrust/commands"
	. "github.com/miketheprogrammer/go-thrust/common"
	"github.com/miketheprogrammer/go-thrust/connection"
	"github.com/miketheprogrammer/go-thrust/dispatcher"
)

/*
Session is the core API Binding object used to communicate with Thrust.
*/
type Session struct {
	TargetID                 uint
	CookieStore              bool
	OffTheRecord             bool
	Path                     string
	Ready                    bool
	CommandHistory           []*Command
	ResponseHistory          []*CommandResponse
	WaitingResponses         []*Command
	CommandQueue             []*Command
	SendChannel              *connection.In
	SessionOverrideInterface SessionInvokable
}

/*
NewSession is a constructor that takes 3 arguments,
incognito which is a boolean, meaning dont persist session state
after close.
overrideDefaultSession which is a boolean that till tell thrust core
to try to invoke session methods from us.
path string is the path to store session data.
*/
func NewSession(incognito, overrideDefaultSession bool, path string) *Session {
	session := Session{
		CookieStore:  overrideDefaultSession,
		OffTheRecord: incognito,
		Path:         path,
	}
	if overrideDefaultSession == true {
		session.SetInvokable(*NewDummySession())
	}
	command := Command{
		Action:     "create",
		ObjectType: "session",
		Args: CommandArguments{
			CookieStore:  session.CookieStore,
			OffTheRecord: session.OffTheRecord,
			Path:         session.Path,
		},
	}
	session.SendChannel = connection.GetInputChannels()
	session.WaitingResponses = append(session.WaitingResponses, &command)
	session.Send(&command)
	dispatcher.RegisterHandler(session.DispatchResponse)
	return &session
}

func (session *Session) HandleInvoke(reply CommandResponse) {
	if reply.TargetID == session.TargetID {
		response := &CommandResponse{
			Action: "reply",
			ID:     reply.ID,
			Result: ReplyResult{},
		}
		switch reply.Method {
		case "cookies_load":
			cookies := session.SessionOverrideInterface.InvokeCookiesLoad(&reply.Args, session)
			marshaledCookies, err := json.Marshal(cookies)
			if err != nil {
				Log.Error(err)
			} else {
				response.Result.Cookies = marshaledCookies
			}
		case "cookies_load_for_key":
			cookies := session.SessionOverrideInterface.InvokeCookiesLoadForKey(&reply.Args, session)
			marshaledCookies, err := json.Marshal(cookies)
			if err != nil {
				Log.Error(err)
			} else {
				response.Result.Cookies = marshaledCookies
			}
		case "cookies_flush":
			session.SessionOverrideInterface.InvokeCookiesFlush(&reply.Args, session)
		case "cookies_add":
			session.SessionOverrideInterface.InvokeCookiesAdd(&reply.Args, session)
		case "cookies_delete":
			session.SessionOverrideInterface.InvokeCookiesDelete(&reply.Args, session)
		case "cookies_update_access_time":
			session.SessionOverrideInterface.InvokeCookiesUpdateAccessTime(&reply.Args, session)
		case "cookies_force_keep_session_state":
			session.SessionOverrideInterface.InvokeCookieForceKeepSessionState(&reply.Args, session)
		}
		Log.Info("Sending Response to Invoke")
		session.SendChannel.CommandResponses <- response
	}
}

func (session *Session) HandleReply(reply CommandResponse) {
	fmt.Println(reply)
	for k, command := range session.WaitingResponses {
		if command.ID != reply.ID {
			continue
		}
		if command.ID == reply.ID {
			Log.Debug("Window(", session.TargetID, ")::Handling Reply::", reply)
			if len(session.WaitingResponses) > 1 {
				// Remove the element at index k
				session.WaitingResponses = session.WaitingResponses[:k+copy(session.WaitingResponses[k:], session.WaitingResponses[k+1:])]
			} else {
				// Just initialize to empty splice literal
				session.WaitingResponses = []*Command{}
			}
			fmt.Println("session", session.TargetID, command.Action, reply.Result.TargetID)
			if session.TargetID == 0 && command.Action == "create" {
				if reply.Result.TargetID != 0 {
					session.TargetID = reply.Result.TargetID
					Log.Debug("Session:: Received TargetID(", session.TargetID, ") :: Setting Ready State")
					session.Ready = true
				}
			}
		}
	}
}

func (session *Session) DispatchResponse(response CommandResponse) {
	switch response.Action {
	case "invoke":
		session.HandleInvoke(response)
	case "reply":
		session.HandleReply(response)

	}
}

func (session *Session) Send(command *Command) {
	session.SendChannel.Commands <- command
}

func (session *Session) SetInvokable(si SessionInvokable) {
	session.SessionOverrideInterface = si
}
<|endoftext|>


<commit_before>package sudoku

import (
	"math/rand"
)

type RankedObject interface {
	Rank() int
}

type FiniteQueue struct {
	min     int
	max     int
	objects []*finiteQueueBucket
}

type finiteQueueBucket struct {
	objects []RankedObject
	numNils int
	rank    int
}

const _REALLOCATE_PROPORTION = 0.20

type SyncedFiniteQueue struct {
	queue FiniteQueue
	In    chan RankedObject
	Out   chan RankedObject
	Exit  chan bool
}

type FiniteQueueGetter struct {
	queue         *FiniteQueue
	ignoreObjects map[RankedObject]bool
	currentList   *finiteQueueBucket
}

//Returns a new queue that will work for items with a rank as low as min or as high as max (inclusive)
func NewFiniteQueue(min int, max int) *FiniteQueue {
	result := FiniteQueue{min, max, make([]*finiteQueueBucket, max-min+1)}
	for i := 0; i < max-min+1; i++ {
		result.objects[i] = &finiteQueueBucket{make([]RankedObject, 1), 0, i + result.min}
	}
	return &result
}

func NewSyncedFiniteQueue(min int, max int) *SyncedFiniteQueue {
	result := &SyncedFiniteQueue{*NewFiniteQueue(min, max), make(chan RankedObject), make(chan RankedObject), make(chan bool, 1)}
	go result.workLoop()
	return result
}

func (self *finiteQueueBucket) removeNils() {
	result := make([]RankedObject, len(self.objects)-self.numNils)
	targetIndex := 0
	for i := 0; i < len(self.objects); i++ {
		if self.objects[i] == nil {
			continue
		}
		result[targetIndex] = self.objects[i]
		targetIndex++
	}
	self.objects = result
	self.numNils = 0
}

func (self *finiteQueueBucket) trimNils() {
	if float64(self.numNils)/float64(len(self.objects)) > _REALLOCATE_PROPORTION {
		self.removeNils()
	}
	for len(self.objects) > 0 && self.objects[0] == nil {
		self.objects = self.objects[1:]
		self.numNils--
	}

}

func (self *finiteQueueBucket) setNil(index int) {
	if self.objects[index] == nil {
		//We don't want to double-count nils.
		return
	}
	self.objects[index] = nil
	self.numNils++
}

func (self *finiteQueueBucket) getItem() RankedObject {
	self.trimNils()
	for len(self.objects) > 0 {
		index := rand.Intn(len(self.objects))
		obj := self.objects[index]
		//Mark this object as gathered.
		self.setNil(index)
		if obj != nil && obj.Rank() == self.rank {
			return obj
		}
		self.trimNils()
	}
	return nil
}

func (self *finiteQueueBucket) addItem(item RankedObject) {
	//Scrub the list for this item.
	for _, obj := range self.objects {
		//Structs will compare equal if all of their fields are the same.
		if item == obj {
			//It's already there, just return.
			return
		}
	}
	self.objects = append(self.objects, item)
}

func (self *finiteQueueBucket) compact() {

	//First, find any i's that don't have the right rank anymore and mark as nils.
	for i, obj := range self.objects {
		if obj == nil {
			continue
		}
		if obj.Rank() != self.rank {
			self.setNil(i)
		}
	}
	self.removeNils()
}

func (self *SyncedFiniteQueue) workLoop() {
	for {
		firstItem := self.queue.Get()
		itemSent := false
		if firstItem == nil {
			//We can take in new things or accept an exit.
			select {
			case <-self.Exit:
				return
			case incoming := <-self.In:
				self.queue.Insert(incoming)
			}
		} else {
			//We can take in new things, send out smallest one, or exit.
			select {
			case <-self.Exit:
				return
			case incoming := <-self.In:
				self.queue.Insert(incoming)
			case self.Out <- firstItem:
				itemSent = true
			}
			//If we didn't send the item out, we need to put it back in.
			if !itemSent {
				self.queue.Insert(firstItem)
			}
		}
	}
}

func (self *FiniteQueue) NewGetter() *FiniteQueueGetter {
	list, _ := self.getList(self.min)
	return &FiniteQueueGetter{self, make(map[RankedObject]bool), list}
}

func (self *FiniteQueue) Min() int {
	return self.min
}

func (self *FiniteQueue) Max() int {
	return self.max
}

func (self *FiniteQueue) Insert(obj RankedObject) {
	rank := obj.Rank()
	list, ok := self.getList(rank)
	if !ok {
		//Apparently rank wasn't legal.
		return
	}
	list.addItem(obj)
}

func (self *FiniteQueue) Get() RankedObject {
	return self.GetSmallerThan(self.max + 1)
}

func (self *FiniteQueue) GetSmallerThan(max int) RankedObject {
	return self.getSmallerThan(max, make(map[RankedObject]int))
}

func (self *FiniteQueue) getSmallerThan(max int, ignoredObjects map[RankedObject]int) RankedObject {
	//TOOD: actually respect ignoredObjects
	for i, list := range self.objects {
		if i+self.min >= max {
			return nil
		}
		result := list.getItem()
		if result != nil {
			return result
		}
	}
	return nil
}

func (self *FiniteQueue) legalRank(rank int) bool {
	return rank >= self.min && rank <= self.max
}

func (self *FiniteQueue) getList(rank int) (*finiteQueueBucket, bool) {
	if !self.legalRank(rank) {
		return nil, false
	}
	return self.objects[rank-self.min], true
}
<commit_msg>getList --> getBucket<commit_after>package sudoku

import (
	"math/rand"
)

type RankedObject interface {
	Rank() int
}

type FiniteQueue struct {
	min     int
	max     int
	objects []*finiteQueueBucket
}

type finiteQueueBucket struct {
	objects []RankedObject
	numNils int
	rank    int
}

const _REALLOCATE_PROPORTION = 0.20

type SyncedFiniteQueue struct {
	queue FiniteQueue
	In    chan RankedObject
	Out   chan RankedObject
	Exit  chan bool
}

type FiniteQueueGetter struct {
	queue         *FiniteQueue
	ignoreObjects map[RankedObject]bool
	currentList   *finiteQueueBucket
}

//Returns a new queue that will work for items with a rank as low as min or as high as max (inclusive)
func NewFiniteQueue(min int, max int) *FiniteQueue {
	result := FiniteQueue{min, max, make([]*finiteQueueBucket, max-min+1)}
	for i := 0; i < max-min+1; i++ {
		result.objects[i] = &finiteQueueBucket{make([]RankedObject, 1), 0, i + result.min}
	}
	return &result
}

func NewSyncedFiniteQueue(min int, max int) *SyncedFiniteQueue {
	result := &SyncedFiniteQueue{*NewFiniteQueue(min, max), make(chan RankedObject), make(chan RankedObject), make(chan bool, 1)}
	go result.workLoop()
	return result
}

func (self *finiteQueueBucket) removeNils() {
	result := make([]RankedObject, len(self.objects)-self.numNils)
	targetIndex := 0
	for i := 0; i < len(self.objects); i++ {
		if self.objects[i] == nil {
			continue
		}
		result[targetIndex] = self.objects[i]
		targetIndex++
	}
	self.objects = result
	self.numNils = 0
}

func (self *finiteQueueBucket) trimNils() {
	if float64(self.numNils)/float64(len(self.objects)) > _REALLOCATE_PROPORTION {
		self.removeNils()
	}
	for len(self.objects) > 0 && self.objects[0] == nil {
		self.objects = self.objects[1:]
		self.numNils--
	}

}

func (self *finiteQueueBucket) setNil(index int) {
	if self.objects[index] == nil {
		//We don't want to double-count nils.
		return
	}
	self.objects[index] = nil
	self.numNils++
}

func (self *finiteQueueBucket) getItem() RankedObject {
	self.trimNils()
	for len(self.objects) > 0 {
		index := rand.Intn(len(self.objects))
		obj := self.objects[index]
		//Mark this object as gathered.
		self.setNil(index)
		if obj != nil && obj.Rank() == self.rank {
			return obj
		}
		self.trimNils()
	}
	return nil
}

func (self *finiteQueueBucket) addItem(item RankedObject) {
	//Scrub the list for this item.
	for _, obj := range self.objects {
		//Structs will compare equal if all of their fields are the same.
		if item == obj {
			//It's already there, just return.
			return
		}
	}
	self.objects = append(self.objects, item)
}

func (self *finiteQueueBucket) compact() {

	//First, find any i's that don't have the right rank anymore and mark as nils.
	for i, obj := range self.objects {
		if obj == nil {
			continue
		}
		if obj.Rank() != self.rank {
			self.setNil(i)
		}
	}
	self.removeNils()
}

func (self *SyncedFiniteQueue) workLoop() {
	for {
		firstItem := self.queue.Get()
		itemSent := false
		if firstItem == nil {
			//We can take in new things or accept an exit.
			select {
			case <-self.Exit:
				return
			case incoming := <-self.In:
				self.queue.Insert(incoming)
			}
		} else {
			//We can take in new things, send out smallest one, or exit.
			select {
			case <-self.Exit:
				return
			case incoming := <-self.In:
				self.queue.Insert(incoming)
			case self.Out <- firstItem:
				itemSent = true
			}
			//If we didn't send the item out, we need to put it back in.
			if !itemSent {
				self.queue.Insert(firstItem)
			}
		}
	}
}

func (self *FiniteQueue) NewGetter() *FiniteQueueGetter {
	list, _ := self.getBucket(self.min)
	return &FiniteQueueGetter{self, make(map[RankedObject]bool), list}
}

func (self *FiniteQueue) Min() int {
	return self.min
}

func (self *FiniteQueue) Max() int {
	return self.max
}

func (self *FiniteQueue) Insert(obj RankedObject) {
	rank := obj.Rank()
	list, ok := self.getBucket(rank)
	if !ok {
		//Apparently rank wasn't legal.
		return
	}
	list.addItem(obj)
}

func (self *FiniteQueue) Get() RankedObject {
	return self.GetSmallerThan(self.max + 1)
}

func (self *FiniteQueue) GetSmallerThan(max int) RankedObject {
	return self.getSmallerThan(max, make(map[RankedObject]int))
}

func (self *FiniteQueue) getSmallerThan(max int, ignoredObjects map[RankedObject]int) RankedObject {
	//TOOD: actually respect ignoredObjects
	for i, list := range self.objects {
		if i+self.min >= max {
			return nil
		}
		result := list.getItem()
		if result != nil {
			return result
		}
	}
	return nil
}

func (self *FiniteQueue) legalRank(rank int) bool {
	return rank >= self.min && rank <= self.max
}

func (self *FiniteQueue) getBucket(rank int) (*finiteQueueBucket, bool) {
	if !self.legalRank(rank) {
		return nil, false
	}
	return self.objects[rank-self.min], true
}
<|endoftext|>


<commit_before>package app

import (
	"fmt"
	"github.com/Tapjoy/dynamiq/app/stats"
	"github.com/hashicorp/memberlist"
	"github.com/tpjg/goriakpbc"
	"log"
	"math"
	"math/rand"
	"strconv"
	"sync"
	"time"
)

// Define statistics keys suffixes

const QUEUE_SENT_STATS_SUFFIX = "sent.count"
const QUEUE_RECEIVED_STATS_SUFFIX = "received.count"
const QUEUE_DELETED_STATS_SUFFIX = "deleted.count"
const QUEUE_DEPTH_STATS_SUFFIX = "depth.count"
const QUEUE_INFLIGHT_STATS_SUFFIX = "inflight.count"
const QUEUE_DEPTHAPR_STATS_SUFFIX = "approximate_depth.count"

type Queues struct {
	// a container for all queues
	QueueMap map[string]*Queue
	// Settings for Queues in general, ie queue list
	Config *riak.RDtMap
	// Mutex for protecting rw access to the Config object
	sync.RWMutex
}

type Queue struct {
	// the definition of a queue
	// name of the queue
	Name string
	// the partitions of the queue
	Parts *Partitions
	// Individual settings for the queue
	Config *riak.RDtMap
	// Mutex for protecting rw access to the Config object
	sync.RWMutex
}

func incrementMessageCount(c stats.StatsClient, queueName string, numberOfMessages int64) error {
	// Increment # Sent
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_SENT_STATS_SUFFIX)
	err := c.Incr(key, numberOfMessages)
	// Increment Depth count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_DEPTH_STATS_SUFFIX)
	err = c.IncrGauge(key, numberOfMessages)
	return err
}

func decrementMessageCount(c stats.StatsClient, queueName string, numberOfMessages int64) error {
	// Increment # Deleted
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_DELETED_STATS_SUFFIX)
	err := c.Incr(key, numberOfMessages)
	// Decrement Inflight count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_INFLIGHT_STATS_SUFFIX)
	err = c.DecrGauge(key, numberOfMessages)
	// Decrement Depth count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_DEPTH_STATS_SUFFIX)
	err = c.DecrGauge(key, numberOfMessages)
	return err
}

func incrementReceiveCount(c stats.StatsClient, queueName string, numberOfMessages int64) error {
	// Increment # Received
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_RECEIVED_STATS_SUFFIX)
	err := c.Incr(key, numberOfMessages)
	// Increment Inflight count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_INFLIGHT_STATS_SUFFIX)
	err = c.IncrGauge(key, numberOfMessages)
	return err
}
func (queue *Queue) setQueueDepthApr(c stats.StatsClient, list *memberlist.Memberlist, queueName string, ids []string) error {
	// set  depth
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_DEPTHAPR_STATS_SUFFIX)
	// find the difference between the first messages id and the last messages id

	first, _ := strconv.ParseInt(ids[0], 10, 64)
	last, _ := strconv.ParseInt(ids[len(ids)-1], 10, 64)
	difference := last - first
	//find the density of messages
	density := float64(len(ids)) / float64(difference)
	// find the total count of messages
	count := density * math.MaxInt64

	// for small queues where we only return 1 message guesstimate ( or should we return 0? )
	multiplier := queue.Parts.PartitionCount() * len(list.Members())
	if len(ids) == 1 {
		return c.SetGauge(key, int64(len(ids)*multiplier))
	} else {
		return c.SetGauge(key, int64(count))
	}
}

func (queues *Queues) Exists(cfg *Config, queueName string) bool {
	// For now, lets go right to Riak for this
	// Because of the config delay, we don't wanna check the memory values
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	m, _ := bucket.FetchMap(QUEUE_CONFIG_NAME)
	set := m.AddSet(QUEUE_SET_NAME)

	for _, value := range set.GetValue() {
		log.Printf("Looking for %s, found %s", queueName, string(value[:]))
		if string(value[:]) == queueName {
			return true
		}
	}
	return false
}

// get a message from the queue
func (queue *Queue) Get(cfg *Config, list *memberlist.Memberlist, batchsize uint32) ([]riak.RObject, error) {
	// get the top and bottom partitions
	partBottom, partTop, err := queue.Parts.GetPartition(cfg, queue.Name, list)

	if err != nil {
		return nil, err
	}
	// grab a riak client
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	//set the bucket
	bucket, err := client.NewBucket(queue.Name)
	if err != nil {
		log.Printf("Error%v", err)
	}
	//get a list of batchsize message ids
	messageIds, _, err := bucket.IndexQueryRangePage("id_int", strconv.Itoa(partBottom), strconv.Itoa(partTop), batchsize, "")
	defer queue.setQueueDepthApr(cfg.Stats.Client, list, queue.Name, messageIds)

	if err != nil {
		log.Printf("Error%v", err)
	}
	messageCount := len(messageIds)
	defer incrementReceiveCount(cfg.Stats.Client, queue.Name, int64(messageCount))
	log.Println("Message retrieved ", messageCount)
	return queue.RetrieveMessages(messageIds, cfg), err
}

// Put a Message onto the queue
func (queue *Queue) Put(cfg *Config, message string) string {
	//Grab our bucket
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		//Retrieve a UUID
		rand.Seed(time.Now().UnixNano())
		randInt := rand.Int63n(math.MaxInt64)
		uuid := strconv.FormatInt(randInt, 10)

		messageObj := bucket.NewObject(uuid)
		messageObj.Indexes["id_int"] = []string{uuid}
		messageObj.Data = []byte(message)
		messageObj.Store()

		defer incrementMessageCount(cfg.Stats.Client, queue.Name, 1)
		return uuid
	} else {
		//Actually want to handle this in some other way
		return ""
	}
}

// Delete a Message from the queue
func (queue *Queue) Delete(cfg *Config, id string) bool {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		log.Println("Deleting: ", id)
		err = bucket.Delete(id)
		if err == nil {
			defer decrementMessageCount(cfg.Stats.Client, queue.Name, 1)
			return true
		}
	}
	if err != nil {
		log.Println(err)
	}
	// if we got here we're borked
	// TODO stats cleanup? Possibility that this gets us out of sync
	return false
}

// helpers
func (queue *Queue) RetrieveMessages(ids []string, cfg *Config) []riak.RObject {
	var rObjectArrayChan = make(chan []riak.RObject, len(ids))
	var rKeys = make(chan string, len(ids))

	start := time.Now()
	//fmt.Println("In get multi")
	for i := 0; i < len(ids); i++ {
		go func() {
			var riakKey string
			client := cfg.RiakConnection()
			defer cfg.ReleaseRiakConnection(client)
			//fmt.Println("Getting bucket")
			bucket, _ := client.NewBucket(queue.Name)
			riakKey = <-rKeys
			//fmt.Println("Getting value")
			rObject, _ := bucket.Get(riakKey)

			rObjectArrayChan <- []riak.RObject{*rObject}

			//fmt.Println("Returning value")
		}()
		rKeys <- ids[i]
	}
	returnVals := make([]riak.RObject, 0)
	for i := 0; i < len(ids); i++ {
		var rObjectArray = <-rObjectArrayChan
		//If the key isn't blank, we've got a meaningful object to deal with
		if len(rObjectArray) == 1 {
			returnVals = append(returnVals, rObjectArray[0])
		}
	}
	elapsed := time.Since(start)
	log.Printf("Get Multi Took %s\n", elapsed)
	return returnVals
}

func (queues *Queues) syncConfig(cfg *Config) {
	for {
		log.Println("syncing Queue config with Riak")
		client := cfg.RiakConnection()
		bucket, err := client.NewBucketType("maps", CONFIGURATION_BUCKET)
		if err != nil {
			log.Println(err)
		}

		queuesConfig, err := bucket.FetchMap(QUEUE_CONFIG_NAME)
		if err != nil {
			log.Println(err)
		}
		queues.updateQueuesConfig(queuesConfig)

		//iterate the map and add or remove topics that need to be destroyed
		queueSet := queues.getQueuesConfig().AddSet(QUEUE_SET_NAME)

		if queueSet == nil {
			//bail if there aren't any queues
			//but not before sleeping
			cfg.ReleaseRiakConnection(client)
			time.Sleep(cfg.Core.SyncConfigInterval * time.Second)
			continue
		}
		queueSlice := queueSet.GetValue()
		if queueSlice == nil {
			//bail if there aren't any queues
			//but not before sleeping
			cfg.ReleaseRiakConnection(client)
			time.Sleep(cfg.Core.SyncConfigInterval * time.Second)
			continue
		}

		//Is there a better way to do this?
		//iterate over the queues in riak and add the missing ones
		queuesToKeep := make(map[string]bool)
		for _, queue := range queueSlice {
			queueName := string(queue)
			var present bool
			_, present = queues.QueueMap[queueName]
			if present != true {
				initQueueFromRiak(cfg, queueName)
			}
			queuesToKeep[queueName] = true
		}

		//iterate over the topics in topics.TopicMap and delete the ones no longer used
		for queue, _ := range queues.QueueMap {
			var present bool
			_, present = queuesToKeep[queue]
			if present != true {
				delete(queues.QueueMap, queue)
			}
		}

		//sync all topics with riak

		for _, queue := range queues.QueueMap {
			queue.syncConfig(cfg)
		}
		//sleep for the configured interval
		cfg.ReleaseRiakConnection(client)
		time.Sleep(cfg.Core.SyncConfigInterval * time.Millisecond)
	}
}

func initQueueFromRiak(cfg *Config, queueName string) {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	config, _ := bucket.FetchMap(queueConfigRecordName(queueName))

	queue := Queue{
		Name:   queueName,
		Parts:  InitPartitions(cfg, queueName),
		Config: config,
	}

	// This is adding a new member to the collection, it shouldn't need a lock?
	// TODO Keep an eye on this for emergent issues
	cfg.Queues.QueueMap[queueName] = &queue
}

func (queue *Queue) syncConfig(cfg *Config) {
	//refresh the queue RDtMap
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)

	rCfg, _ := bucket.FetchMap(queueConfigRecordName(queue.Name))
	queue.updateQueueConfig(rCfg)
	queue.Parts.syncPartitions(cfg, queue.Name)
}

func (queue *Queue) updateQueueConfig(rCfg *riak.RDtMap) {
	queue.Lock()
	defer queue.Unlock()
	queue.Config = rCfg
}

func (queue *Queue) getQueueConfig() *riak.RDtMap {
	queue.RLock()
	defer queue.RUnlock()
	return queue.Config
}

func (queues *Queues) updateQueuesConfig(rCfg *riak.RDtMap) {
	queues.Lock()
	defer queues.Unlock()
	queues.Config = rCfg
}

func (queues *Queues) getQueuesConfig() *riak.RDtMap {
	queues.RLock()
	defer queues.RUnlock()
	return queues.Config
}
<commit_msg>avoid dividing by zero<commit_after>package app

import (
	"fmt"
	"github.com/Tapjoy/dynamiq/app/stats"
	"github.com/hashicorp/memberlist"
	"github.com/tpjg/goriakpbc"
	"log"
	"math"
	"math/rand"
	"strconv"
	"sync"
	"time"
)

// Define statistics keys suffixes

const QUEUE_SENT_STATS_SUFFIX = "sent.count"
const QUEUE_RECEIVED_STATS_SUFFIX = "received.count"
const QUEUE_DELETED_STATS_SUFFIX = "deleted.count"
const QUEUE_DEPTH_STATS_SUFFIX = "depth.count"
const QUEUE_INFLIGHT_STATS_SUFFIX = "inflight.count"
const QUEUE_DEPTHAPR_STATS_SUFFIX = "approximate_depth.count"

type Queues struct {
	// a container for all queues
	QueueMap map[string]*Queue
	// Settings for Queues in general, ie queue list
	Config *riak.RDtMap
	// Mutex for protecting rw access to the Config object
	sync.RWMutex
}

type Queue struct {
	// the definition of a queue
	// name of the queue
	Name string
	// the partitions of the queue
	Parts *Partitions
	// Individual settings for the queue
	Config *riak.RDtMap
	// Mutex for protecting rw access to the Config object
	sync.RWMutex
}

func incrementMessageCount(c stats.StatsClient, queueName string, numberOfMessages int64) error {
	// Increment # Sent
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_SENT_STATS_SUFFIX)
	err := c.Incr(key, numberOfMessages)
	// Increment Depth count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_DEPTH_STATS_SUFFIX)
	err = c.IncrGauge(key, numberOfMessages)
	return err
}

func decrementMessageCount(c stats.StatsClient, queueName string, numberOfMessages int64) error {
	// Increment # Deleted
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_DELETED_STATS_SUFFIX)
	err := c.Incr(key, numberOfMessages)
	// Decrement Inflight count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_INFLIGHT_STATS_SUFFIX)
	err = c.DecrGauge(key, numberOfMessages)
	// Decrement Depth count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_DEPTH_STATS_SUFFIX)
	err = c.DecrGauge(key, numberOfMessages)
	return err
}

func incrementReceiveCount(c stats.StatsClient, queueName string, numberOfMessages int64) error {
	// Increment # Received
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_RECEIVED_STATS_SUFFIX)
	err := c.Incr(key, numberOfMessages)
	// Increment Inflight count
	key = fmt.Sprintf("%s.%s", queueName, QUEUE_INFLIGHT_STATS_SUFFIX)
	err = c.IncrGauge(key, numberOfMessages)
	return err
}
func (queue *Queue) setQueueDepthApr(c stats.StatsClient, list *memberlist.Memberlist, queueName string, ids []string) error {
	// set  depth
	key := fmt.Sprintf("%s.%s", queueName, QUEUE_DEPTHAPR_STATS_SUFFIX)
	// find the difference between the first messages id and the last messages id

	if len(ids) > 1 {
		first, _ := strconv.ParseInt(ids[0], 10, 64)
		last, _ := strconv.ParseInt(ids[len(ids)-1], 10, 64)
		difference := last - first
		//find the density of messages
		density := float64(len(ids)) / float64(difference)
		// find the total count of messages by multiplying the density by the key range
		count := density * math.MaxInt64
		return c.SetGauge(key, int64(count))

	} else {
		// for small queues where we only return 1  message or no messages guesstimate ( or should we return 0? )
		multiplier := queue.Parts.PartitionCount() * len(list.Members())
		return c.SetGauge(key, int64(len(ids)*multiplier))
	}
}

func (queues *Queues) Exists(cfg *Config, queueName string) bool {
	// For now, lets go right to Riak for this
	// Because of the config delay, we don't wanna check the memory values
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	m, _ := bucket.FetchMap(QUEUE_CONFIG_NAME)
	set := m.AddSet(QUEUE_SET_NAME)

	for _, value := range set.GetValue() {
		log.Printf("Looking for %s, found %s", queueName, string(value[:]))
		if string(value[:]) == queueName {
			return true
		}
	}
	return false
}

// get a message from the queue
func (queue *Queue) Get(cfg *Config, list *memberlist.Memberlist, batchsize uint32) ([]riak.RObject, error) {
	// get the top and bottom partitions
	partBottom, partTop, err := queue.Parts.GetPartition(cfg, queue.Name, list)

	if err != nil {
		return nil, err
	}
	// grab a riak client
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	//set the bucket
	bucket, err := client.NewBucket(queue.Name)
	if err != nil {
		log.Printf("Error%v", err)
	}
	//get a list of batchsize message ids
	messageIds, _, err := bucket.IndexQueryRangePage("id_int", strconv.Itoa(partBottom), strconv.Itoa(partTop), batchsize, "")
	defer queue.setQueueDepthApr(cfg.Stats.Client, list, queue.Name, messageIds)

	if err != nil {
		log.Printf("Error%v", err)
	}
	messageCount := len(messageIds)
	defer incrementReceiveCount(cfg.Stats.Client, queue.Name, int64(messageCount))
	log.Println("Message retrieved ", messageCount)
	return queue.RetrieveMessages(messageIds, cfg), err
}

// Put a Message onto the queue
func (queue *Queue) Put(cfg *Config, message string) string {
	//Grab our bucket
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		//Retrieve a UUID
		rand.Seed(time.Now().UnixNano())
		randInt := rand.Int63n(math.MaxInt64)
		uuid := strconv.FormatInt(randInt, 10)

		messageObj := bucket.NewObject(uuid)
		messageObj.Indexes["id_int"] = []string{uuid}
		messageObj.Data = []byte(message)
		messageObj.Store()

		defer incrementMessageCount(cfg.Stats.Client, queue.Name, 1)
		return uuid
	} else {
		//Actually want to handle this in some other way
		return ""
	}
}

// Delete a Message from the queue
func (queue *Queue) Delete(cfg *Config, id string) bool {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		log.Println("Deleting: ", id)
		err = bucket.Delete(id)
		if err == nil {
			defer decrementMessageCount(cfg.Stats.Client, queue.Name, 1)
			return true
		}
	}
	if err != nil {
		log.Println(err)
	}
	// if we got here we're borked
	// TODO stats cleanup? Possibility that this gets us out of sync
	return false
}

// helpers
func (queue *Queue) RetrieveMessages(ids []string, cfg *Config) []riak.RObject {
	var rObjectArrayChan = make(chan []riak.RObject, len(ids))
	var rKeys = make(chan string, len(ids))

	start := time.Now()
	//fmt.Println("In get multi")
	for i := 0; i < len(ids); i++ {
		go func() {
			var riakKey string
			client := cfg.RiakConnection()
			defer cfg.ReleaseRiakConnection(client)
			//fmt.Println("Getting bucket")
			bucket, _ := client.NewBucket(queue.Name)
			riakKey = <-rKeys
			//fmt.Println("Getting value")
			rObject, _ := bucket.Get(riakKey)

			rObjectArrayChan <- []riak.RObject{*rObject}

			//fmt.Println("Returning value")
		}()
		rKeys <- ids[i]
	}
	returnVals := make([]riak.RObject, 0)
	for i := 0; i < len(ids); i++ {
		var rObjectArray = <-rObjectArrayChan
		//If the key isn't blank, we've got a meaningful object to deal with
		if len(rObjectArray) == 1 {
			returnVals = append(returnVals, rObjectArray[0])
		}
	}
	elapsed := time.Since(start)
	log.Printf("Get Multi Took %s\n", elapsed)
	return returnVals
}

func (queues *Queues) syncConfig(cfg *Config) {
	for {
		log.Println("syncing Queue config with Riak")
		client := cfg.RiakConnection()
		bucket, err := client.NewBucketType("maps", CONFIGURATION_BUCKET)
		if err != nil {
			log.Println(err)
		}

		queuesConfig, err := bucket.FetchMap(QUEUE_CONFIG_NAME)
		if err != nil {
			log.Println(err)
		}
		queues.updateQueuesConfig(queuesConfig)

		//iterate the map and add or remove topics that need to be destroyed
		queueSet := queues.getQueuesConfig().AddSet(QUEUE_SET_NAME)

		if queueSet == nil {
			//bail if there aren't any queues
			//but not before sleeping
			cfg.ReleaseRiakConnection(client)
			time.Sleep(cfg.Core.SyncConfigInterval * time.Second)
			continue
		}
		queueSlice := queueSet.GetValue()
		if queueSlice == nil {
			//bail if there aren't any queues
			//but not before sleeping
			cfg.ReleaseRiakConnection(client)
			time.Sleep(cfg.Core.SyncConfigInterval * time.Second)
			continue
		}

		//Is there a better way to do this?
		//iterate over the queues in riak and add the missing ones
		queuesToKeep := make(map[string]bool)
		for _, queue := range queueSlice {
			queueName := string(queue)
			var present bool
			_, present = queues.QueueMap[queueName]
			if present != true {
				initQueueFromRiak(cfg, queueName)
			}
			queuesToKeep[queueName] = true
		}

		//iterate over the topics in topics.TopicMap and delete the ones no longer used
		for queue, _ := range queues.QueueMap {
			var present bool
			_, present = queuesToKeep[queue]
			if present != true {
				delete(queues.QueueMap, queue)
			}
		}

		//sync all topics with riak

		for _, queue := range queues.QueueMap {
			queue.syncConfig(cfg)
		}
		//sleep for the configured interval
		cfg.ReleaseRiakConnection(client)
		time.Sleep(cfg.Core.SyncConfigInterval * time.Millisecond)
	}
}

func initQueueFromRiak(cfg *Config, queueName string) {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	config, _ := bucket.FetchMap(queueConfigRecordName(queueName))

	queue := Queue{
		Name:   queueName,
		Parts:  InitPartitions(cfg, queueName),
		Config: config,
	}

	// This is adding a new member to the collection, it shouldn't need a lock?
	// TODO Keep an eye on this for emergent issues
	cfg.Queues.QueueMap[queueName] = &queue
}

func (queue *Queue) syncConfig(cfg *Config) {
	//refresh the queue RDtMap
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)

	rCfg, _ := bucket.FetchMap(queueConfigRecordName(queue.Name))
	queue.updateQueueConfig(rCfg)
	queue.Parts.syncPartitions(cfg, queue.Name)
}

func (queue *Queue) updateQueueConfig(rCfg *riak.RDtMap) {
	queue.Lock()
	defer queue.Unlock()
	queue.Config = rCfg
}

func (queue *Queue) getQueueConfig() *riak.RDtMap {
	queue.RLock()
	defer queue.RUnlock()
	return queue.Config
}

func (queues *Queues) updateQueuesConfig(rCfg *riak.RDtMap) {
	queues.Lock()
	defer queues.Unlock()
	queues.Config = rCfg
}

func (queues *Queues) getQueuesConfig() *riak.RDtMap {
	queues.RLock()
	defer queues.RUnlock()
	return queues.Config
}
<|endoftext|>


<commit_before>package config

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
)

const (
	VERSION = "0.0.10"
)

type config struct {
	Name    string
	Version string

	LDAP struct {
		Host     string
		Port     uint
		Base     string
		BindDN   string
		Password string `ini:"pass"`
		Filter   string
	} `ini:"ldap"`

	Session struct {
		Name   string
		Domain string
		Secret string
		MaxAge int // cookie maxAge
	} `ini:"sess"`

	HttpListen   string
	ResUrl       string
	Backend      struct{ DSN string }
	Root         string
	Debug        bool
	UserLifetime int // secends, user online time
}

var (
	_configFile  string
	_loaded      bool
	fs           *flag.FlagSet
	Settings     *config = &config{}
	printVersion bool
)

func init() {
	Settings.Name = "staffio"
	Settings.Version = VERSION
	// Settings.HttpListen = "localhost:3000"
	// Settings.ResUrl = "/static/"
	// Settings.LDAP.Host = "localhost"
	// Settings.LDAP.Port = 389

	fs = flag.NewFlagSet("staffio", flag.ExitOnError)

	fs.StringVar(&Settings.LDAP.Host, "ldap-host", "localhost", "ldap hostname")
	fs.UintVar(&Settings.LDAP.Port, "ldap-port", 389, "ldap hostname")
	fs.StringVar(&Settings.LDAP.Base, "ldap-base", "", "ldap base")
	fs.StringVar(&Settings.LDAP.BindDN, "ldap-bind-dn", "", "ldap bind dn")
	fs.StringVar(&Settings.LDAP.Password, "ldap-pass", "", "ldap bind password")
	fs.StringVar(&Settings.LDAP.Filter, "ldap-user-filter", "(objectclass=inetOrgPerson)", "ldap search filter")
	fs.StringVar(&Settings.HttpListen, "http-listen", "localhost:5000", "bind address and port")
	fs.StringVar(&Settings.Session.Name, "sess-name", "staff_sess", "session name")
	fs.StringVar(&Settings.Session.Domain, "sess-domain", "", "session domain")
	fs.StringVar(&Settings.Session.Secret, "sess-secret", "very-secret", "session secret")
	fs.IntVar(&Settings.Session.MaxAge, "sess-maxage", 86400*30, "session cookie life time (in seconds)")
	fs.IntVar(&Settings.UserLifetime, "user-life", 2500, "user online life time (in seconds)")
	fs.StringVar(&Settings.ResUrl, "res-url", "/static/", "static resource url")
	fs.StringVar(&Settings.Backend.DSN, "backend-dsn", "postgres://staffio@localhost/staffio?sslmode=disable", "database dsn string for backend")
	fs.StringVar(&Settings.Root, "root", "./", "app root directory")
	fs.BoolVar(&Settings.Debug, "debug", false, "app in debug mode")
	fs.BoolVar(&printVersion, "version", false, "Print the version and exit")

}

func (c *config) Parse() {
	perr := fs.Parse(os.Args[1:])
	switch perr {
	case nil:
	case flag.ErrHelp:
		os.Exit(0)
	default:
		os.Exit(2)
	}
	if len(fs.Args()) != 0 {
		log.Fatalf("'%s' is not a valid flag", fs.Arg(0))
	}

	if printVersion {
		fmt.Println("staffio version", Settings.Version)
		os.Exit(0)
	}

	err := SetFlagsFromEnv(fs, Settings.Name+"_")
	if err != nil {
		log.Fatalf("staffio: %v", err)
	}

}

// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the prefix "PREFIX_", and any dashes are replaced by
// underscores - for example: some-flag => PREFIX_SOME_FLAG
func SetFlagsFromEnv(fs *flag.FlagSet, prefix string) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})

	if prefix == "" {
		prefix = "_"
	} else {
		prefix = strings.ToUpper(strings.Replace(prefix, "-", "_", -1))
	}

	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := prefix + strings.ToUpper(strings.Replace(f.Name, "-", "_", -1))
			val := os.Getenv(key)
			if val != "" {
				if serr := fs.Set(f.Name, val); serr != nil {
					err = fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
				}
			}
		}
	})
	return err
}
<commit_msg>upgrade version<commit_after>package config

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
)

const (
	VERSION = "0.0.11"
)

type config struct {
	Name    string
	Version string

	LDAP struct {
		Host     string
		Port     uint
		Base     string
		BindDN   string
		Password string `ini:"pass"`
		Filter   string
	} `ini:"ldap"`

	Session struct {
		Name   string
		Domain string
		Secret string
		MaxAge int // cookie maxAge
	} `ini:"sess"`

	HttpListen   string
	ResUrl       string
	Backend      struct{ DSN string }
	Root         string
	Debug        bool
	UserLifetime int // secends, user online time
}

var (
	_configFile  string
	_loaded      bool
	fs           *flag.FlagSet
	Settings     *config = &config{}
	printVersion bool
)

func init() {
	Settings.Name = "staffio"
	Settings.Version = VERSION
	// Settings.HttpListen = "localhost:3000"
	// Settings.ResUrl = "/static/"
	// Settings.LDAP.Host = "localhost"
	// Settings.LDAP.Port = 389

	fs = flag.NewFlagSet("staffio", flag.ExitOnError)

	fs.StringVar(&Settings.LDAP.Host, "ldap-host", "localhost", "ldap hostname")
	fs.UintVar(&Settings.LDAP.Port, "ldap-port", 389, "ldap hostname")
	fs.StringVar(&Settings.LDAP.Base, "ldap-base", "", "ldap base")
	fs.StringVar(&Settings.LDAP.BindDN, "ldap-bind-dn", "", "ldap bind dn")
	fs.StringVar(&Settings.LDAP.Password, "ldap-pass", "", "ldap bind password")
	fs.StringVar(&Settings.LDAP.Filter, "ldap-user-filter", "(objectclass=inetOrgPerson)", "ldap search filter")
	fs.StringVar(&Settings.HttpListen, "http-listen", "localhost:5000", "bind address and port")
	fs.StringVar(&Settings.Session.Name, "sess-name", "staff_sess", "session name")
	fs.StringVar(&Settings.Session.Domain, "sess-domain", "", "session domain")
	fs.StringVar(&Settings.Session.Secret, "sess-secret", "very-secret", "session secret")
	fs.IntVar(&Settings.Session.MaxAge, "sess-maxage", 86400*30, "session cookie life time (in seconds)")
	fs.IntVar(&Settings.UserLifetime, "user-life", 2500, "user online life time (in seconds)")
	fs.StringVar(&Settings.ResUrl, "res-url", "/static/", "static resource url")
	fs.StringVar(&Settings.Backend.DSN, "backend-dsn", "postgres://staffio@localhost/staffio?sslmode=disable", "database dsn string for backend")
	fs.StringVar(&Settings.Root, "root", "./", "app root directory")
	fs.BoolVar(&Settings.Debug, "debug", false, "app in debug mode")
	fs.BoolVar(&printVersion, "version", false, "Print the version and exit")

}

func (c *config) Parse() {
	perr := fs.Parse(os.Args[1:])
	switch perr {
	case nil:
	case flag.ErrHelp:
		os.Exit(0)
	default:
		os.Exit(2)
	}
	if len(fs.Args()) != 0 {
		log.Fatalf("'%s' is not a valid flag", fs.Arg(0))
	}

	if printVersion {
		fmt.Println("staffio version", Settings.Version)
		os.Exit(0)
	}

	err := SetFlagsFromEnv(fs, Settings.Name+"_")
	if err != nil {
		log.Fatalf("staffio: %v", err)
	}

}

// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the prefix "PREFIX_", and any dashes are replaced by
// underscores - for example: some-flag => PREFIX_SOME_FLAG
func SetFlagsFromEnv(fs *flag.FlagSet, prefix string) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})

	if prefix == "" {
		prefix = "_"
	} else {
		prefix = strings.ToUpper(strings.Replace(prefix, "-", "_", -1))
	}

	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := prefix + strings.ToUpper(strings.Replace(f.Name, "-", "_", -1))
			val := os.Getenv(key)
			if val != "" {
				if serr := fs.Set(f.Name, val); serr != nil {
					err = fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
				}
			}
		}
	})
	return err
}
<|endoftext|>


<commit_before>package status

import (
	"reflect"
	"testing"
)

func TestEqualRepoURLs(t *testing.T) {
	tests := []struct {
		rawurl1 string
		rawurl2 string
		want    bool
	}{
		{
			rawurl1: "https://github.com/user/repo",
			rawurl2: "https://github.com/user/repo",
			want:    true,
		},
		{
			rawurl1: "https://github.com/user/repo",
			rawurl2: "https://github.com/user/wrongrepo",
			want:    false,
		},
		{
			rawurl1: "https://github.com/user/repo",
			rawurl2: "git@github.com:user/repo",
			want:    true,
		},
	}
	for _, test := range tests {
		if got, want := EqualRepoURLs(test.rawurl1, test.rawurl2), test.want; got != want {
			t.Errorf("EqualRepoURLs(%q, %q): got %v, want %v", test.rawurl1, test.rawurl2, got, want)
		}
	}
}

func TestParseURL(t *testing.T) {
	tests := []struct {
		in      string
		want    string
		wantErr error
	}{
		{
			in:   "git@github.com:user/repo",
			want: "ssh://git@github.com/user/repo",
		},
		{
			in:   "https://github.com/user/repo",
			want: "https://github.com/user/repo",
		},
	}
	for _, test := range tests {
		u, err := parseURL(test.in)
		if got, want := err, test.wantErr; !reflect.DeepEqual(got, want) {
			t.Errorf("got %q, want %q", got, want)
		}
		if test.wantErr != nil {
			continue
		}

		if got, want := u.String(), test.want; got != want {
			t.Errorf("got %q, want %q", got, want)
		}
	}
}
<commit_msg>status: Make it clear that .git suffix is not ignored.<commit_after>package status

import (
	"reflect"
	"testing"
)

func TestEqualRepoURLs(t *testing.T) {
	tests := []struct {
		rawurl1 string
		rawurl2 string
		want    bool
	}{
		{
			rawurl1: "https://github.com/user/repo",
			rawurl2: "https://github.com/user/repo",
			want:    true,
		},
		{
			// .git suffix does not have special treatment, and these are not considered as equal URLs.
			// See https://github.com/shurcooL/gostatus/issues/37#issuecomment-225336823 for rationale.
			rawurl1: "https://github.com/user/repo",
			rawurl2: "https://github.com/user/repo.git",
			want:    false,
		},
		{
			rawurl1: "https://github.com/user/repo",
			rawurl2: "https://github.com/user/wrongrepo",
			want:    false,
		},
		{
			rawurl1: "https://github.com/user/repo",
			rawurl2: "git@github.com:user/repo",
			want:    true,
		},
	}
	for _, test := range tests {
		if got, want := EqualRepoURLs(test.rawurl1, test.rawurl2), test.want; got != want {
			t.Errorf("EqualRepoURLs(%q, %q): got %v, want %v", test.rawurl1, test.rawurl2, got, want)
		}
	}
}

func TestParseURL(t *testing.T) {
	tests := []struct {
		in      string
		want    string
		wantErr error
	}{
		{
			in:   "git@github.com:user/repo",
			want: "ssh://git@github.com/user/repo",
		},
		{
			in:   "https://github.com/user/repo",
			want: "https://github.com/user/repo",
		},
	}
	for _, test := range tests {
		u, err := parseURL(test.in)
		if got, want := err, test.wantErr; !reflect.DeepEqual(got, want) {
			t.Errorf("got %q, want %q", got, want)
		}
		if test.wantErr != nil {
			continue
		}

		if got, want := u.String(), test.want; got != want {
			t.Errorf("got %q, want %q", got, want)
		}
	}
}
<|endoftext|>


<commit_before>package config

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/Shopify/ejson"
	"github.com/spf13/viper"
	"go.mozilla.org/sops/v3/decrypt"

	"github.com/grafeas/voucher/v2/clair"
	"github.com/grafeas/voucher/v2/repository"
	"github.com/grafeas/voucher/v2/signer/pgp"
)

// Secrets represents the format that the ejson configuration is structured
// in.
type Secrets struct {
	Keys                     map[string]string  `json:"openpgpkeys"`
	ClairConfig              clair.Config       `json:"clair"`
	RepositoryAuthentication repository.KeyRing `json:"repositories"`
}

// ReadSecrets reads from the ejson file and populates the passed interface.
func ReadSecrets() (*Secrets, error) {
	decrypted, err := decryptSecrets()
	if err != nil {
		return nil, err
	}

	var data Secrets
	if err := json.Unmarshal(decrypted, &data); err != nil {
		return nil, err
	}
	return &data, nil
}

func decryptSecrets() ([]byte, error) {
	ejDir := viper.GetString("ejson.dir")
	ejSecrets := viper.GetString("ejson.secrets")
	if ejDir != "" && ejSecrets != "" {
		return ejson.DecryptFile(ejSecrets, ejDir, "")
	}

	sops := viper.GetString("sops.file")
	if sops != "" {
		return decrypt.File(sops, "json")
	}

	return nil, fmt.Errorf("secrets not provided via ejson or sops")
}

// getPGPKeyRing uses the Command's configured ejson file to populate a
// voucher.KeyRing.
func (s *Secrets) getPGPKeyRing() (*pgp.KeyRing, error) {
	newKeyRing := pgp.NewKeyRing()

	for name, key := range s.Keys {
		err := pgp.AddKeyToKeyRingFromReader(newKeyRing, name, bytes.NewReader([]byte(key)))
		if err != nil {
			return nil, err
		}
	}

	return newKeyRing, nil
}
<commit_msg>Prefer yoda condition<commit_after>package config

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/Shopify/ejson"
	"github.com/spf13/viper"
	"go.mozilla.org/sops/v3/decrypt"

	"github.com/grafeas/voucher/v2/clair"
	"github.com/grafeas/voucher/v2/repository"
	"github.com/grafeas/voucher/v2/signer/pgp"
)

// Secrets represents the format that the ejson configuration is structured
// in.
type Secrets struct {
	Keys                     map[string]string  `json:"openpgpkeys"`
	ClairConfig              clair.Config       `json:"clair"`
	RepositoryAuthentication repository.KeyRing `json:"repositories"`
}

// ReadSecrets reads from the ejson file and populates the passed interface.
func ReadSecrets() (*Secrets, error) {
	decrypted, err := decryptSecrets()
	if err != nil {
		return nil, err
	}

	var data Secrets
	if err := json.Unmarshal(decrypted, &data); err != nil {
		return nil, err
	}
	return &data, nil
}

func decryptSecrets() ([]byte, error) {
	ejDir := viper.GetString("ejson.dir")
	ejSecrets := viper.GetString("ejson.secrets")
	if ejDir != "" && ejSecrets != "" {
		return ejson.DecryptFile(ejSecrets, ejDir, "")
	}

	sops := viper.GetString("sops.file")
	if sops != "" {
		return decrypt.File(sops, "json")
	}

	return nil, fmt.Errorf("secrets not provided via ejson or sops")
}

// getPGPKeyRing uses the Command's configured ejson file to populate a
// voucher.KeyRing.
func (s *Secrets) getPGPKeyRing() (*pgp.KeyRing, error) {
	newKeyRing := pgp.NewKeyRing()

	for name, key := range s.Keys {
		err := pgp.AddKeyToKeyRingFromReader(newKeyRing, name, bytes.NewReader([]byte(key)))
		if nil != err {
			return nil, err
		}
	}

	return newKeyRing, nil
}
<|endoftext|>


<commit_before>package lang

import (
	"bytes"
	"fmt"
	"sync"

	"github.com/hashicorp/terraform/config/lang/ast"
)

// Engine is the execution engine for this language. It should be configured
// prior to running Execute.
type Engine struct {
	// GlobalScope is the global scope of execution for this engine.
	GlobalScope *Scope

	// SemanticChecks is a list of additional semantic checks that will be run
	// on the tree prior to executing it. The type checker, identifier checker,
	// etc. will be run before these.
	SemanticChecks []SemanticChecker
}

// SemanticChecker is the type that must be implemented to do a
// semantic check on an AST tree. This will be called with the root node.
type SemanticChecker func(ast.Node) error

// Execute executes the given ast.Node and returns its final value, its
// type, and an error if one exists.
func (e *Engine) Execute(root ast.Node) (interface{}, ast.Type, error) {
	// Run the type checker
	tv := &TypeVisitor{Scope: e.GlobalScope}
	if err := tv.Visit(root); err != nil {
		return nil, ast.TypeInvalid, err
	}

	// Execute
	v := &executeVisitor{Scope: e.GlobalScope}
	return v.Visit(root)
}

// executeVisitor is the visitor used to do the actual execution of
// a program. Note at this point it is assumed that the types check out
// and the identifiers exist.
type executeVisitor struct {
	Scope *Scope

	stack []*ast.LiteralNode
	err   error
	lock  sync.Mutex
}

func (v *executeVisitor) Visit(root ast.Node) (interface{}, ast.Type, error) {
	v.lock.Lock()
	defer v.lock.Unlock()

	// Run the actual visitor pattern
	root.Accept(v.visit)

	// Get our result and clear out everything else
	var result *ast.LiteralNode
	if len(v.stack) > 0 {
		result = v.stack[len(v.stack)-1]
	} else {
		result = new(ast.LiteralNode)
	}
	resultErr := v.err

	// Clear everything else so we aren't just dangling
	v.stack = nil
	v.err = nil

	return result.Value, result.Type, resultErr
}

func (v *executeVisitor) visit(raw ast.Node) {
	if v.err != nil {
		return
	}

	switch n := raw.(type) {
	case *ast.Call:
		v.visitCall(n)
	case *ast.Concat:
		v.visitConcat(n)
	case *ast.LiteralNode:
		v.visitLiteral(n)
	case *ast.VariableAccess:
		v.visitVariableAccess(n)
	default:
		v.err = fmt.Errorf("unknown node: %#v", raw)
	}
}

func (v *executeVisitor) visitCall(n *ast.Call) {
	// Look up the function in the map
	function, ok := v.Scope.FuncMap[n.Func]
	if !ok {
		v.err = fmt.Errorf("unknown function called: %s", n.Func)
		return
	}

	// The arguments are on the stack in reverse order, so pop them off.
	args := make([]interface{}, len(n.Args))
	for i, _ := range n.Args {
		node := v.stackPop()
		args[len(n.Args)-1-i] = node.Value
	}

	// Call the function
	result, err := function.Callback(args)
	if err != nil {
		v.err = fmt.Errorf("%s: %s", n.Func, err)
		return
	}

	// Push the result
	v.stackPush(&ast.LiteralNode{
		Value: result,
		Type:  function.ReturnType,
	})
}

func (v *executeVisitor) visitConcat(n *ast.Concat) {
	// The expressions should all be on the stack in reverse
	// order. So pop them off, reverse their order, and concatenate.
	nodes := make([]*ast.LiteralNode, 0, len(n.Exprs))
	for range n.Exprs {
		nodes = append(nodes, v.stackPop())
	}

	var buf bytes.Buffer
	for i := len(nodes) - 1; i >= 0; i-- {
		buf.WriteString(nodes[i].Value.(string))
	}

	v.stackPush(&ast.LiteralNode{
		Value: buf.String(),
		Type:  ast.TypeString,
	})
}

func (v *executeVisitor) visitLiteral(n *ast.LiteralNode) {
	v.stack = append(v.stack, n)
}

func (v *executeVisitor) visitVariableAccess(n *ast.VariableAccess) {
	// Look up the variable in the map
	variable, ok := v.Scope.VarMap[n.Name]
	if !ok {
		v.err = fmt.Errorf("unknown variable accessed: %s", n.Name)
		return
	}

	v.stack = append(v.stack, &ast.LiteralNode{
		Value: variable.Value,
		Type:  variable.Type,
	})
}

func (v *executeVisitor) stackPush(n *ast.LiteralNode) {
	v.stack = append(v.stack, n)
}

func (v *executeVisitor) stackPop() *ast.LiteralNode {
	var x *ast.LiteralNode
	x, v.stack = v.stack[len(v.stack)-1], v.stack[:len(v.stack)-1]
	return x
}

// Scope represents a lookup scope for execution.
type Scope struct {
	// VarMap and FuncMap are the mappings of identifiers to functions
	// and variable values.
	VarMap  map[string]Variable
	FuncMap map[string]Function
}

// Variable is a variable value for execution given as input to the engine.
// It records the value of a variables along with their type.
type Variable struct {
	Value interface{}
	Type  ast.Type
}

// Function defines a function that can be executed by the engine.
// The type checker will validate that the proper types will be called
// to the callback.
type Function struct {
	ArgTypes   []ast.Type
	ReturnType ast.Type
	Callback   func([]interface{}) (interface{}, error)
}

// LookupFunc will look up a variable by name.
// TODO test
func (s *Scope) LookupFunc(n string) (Function, bool) {
	if s == nil {
		return Function{}, false
	}

	v, ok := s.FuncMap[n]
	return v, ok
}

// LookupVar will look up a variable by name.
// TODO test
func (s *Scope) LookupVar(n string) (Variable, bool) {
	if s == nil {
		return Variable{}, false
	}

	v, ok := s.VarMap[n]
	return v, ok
}
<commit_msg>config/lang: make formal Stack object<commit_after>package lang

import (
	"bytes"
	"fmt"
	"sync"

	"github.com/hashicorp/terraform/config/lang/ast"
)

// Engine is the execution engine for this language. It should be configured
// prior to running Execute.
type Engine struct {
	// GlobalScope is the global scope of execution for this engine.
	GlobalScope *Scope

	// SemanticChecks is a list of additional semantic checks that will be run
	// on the tree prior to executing it. The type checker, identifier checker,
	// etc. will be run before these.
	SemanticChecks []SemanticChecker
}

// SemanticChecker is the type that must be implemented to do a
// semantic check on an AST tree. This will be called with the root node.
type SemanticChecker func(ast.Node) error

// Execute executes the given ast.Node and returns its final value, its
// type, and an error if one exists.
func (e *Engine) Execute(root ast.Node) (interface{}, ast.Type, error) {
	// Run the type checker
	tv := &TypeVisitor{Scope: e.GlobalScope}
	if err := tv.Visit(root); err != nil {
		return nil, ast.TypeInvalid, err
	}

	// Execute
	v := &executeVisitor{Scope: e.GlobalScope}
	return v.Visit(root)
}

// executeVisitor is the visitor used to do the actual execution of
// a program. Note at this point it is assumed that the types check out
// and the identifiers exist.
type executeVisitor struct {
	Scope *Scope

	stack EngineStack
	err   error
	lock  sync.Mutex
}

func (v *executeVisitor) Visit(root ast.Node) (interface{}, ast.Type, error) {
	v.lock.Lock()
	defer v.lock.Unlock()

	// Run the actual visitor pattern
	root.Accept(v.visit)

	// Get our result and clear out everything else
	var result *ast.LiteralNode
	if v.stack.Len() > 0 {
		result = v.stack.Pop()
	} else {
		result = new(ast.LiteralNode)
	}
	resultErr := v.err

	// Clear everything else so we aren't just dangling
	v.stack.Reset()
	v.err = nil

	return result.Value, result.Type, resultErr
}

func (v *executeVisitor) visit(raw ast.Node) {
	if v.err != nil {
		return
	}

	switch n := raw.(type) {
	case *ast.Call:
		v.visitCall(n)
	case *ast.Concat:
		v.visitConcat(n)
	case *ast.LiteralNode:
		v.visitLiteral(n)
	case *ast.VariableAccess:
		v.visitVariableAccess(n)
	default:
		v.err = fmt.Errorf("unknown node: %#v", raw)
	}
}

func (v *executeVisitor) visitCall(n *ast.Call) {
	// Look up the function in the map
	function, ok := v.Scope.FuncMap[n.Func]
	if !ok {
		v.err = fmt.Errorf("unknown function called: %s", n.Func)
		return
	}

	// The arguments are on the stack in reverse order, so pop them off.
	args := make([]interface{}, len(n.Args))
	for i, _ := range n.Args {
		node := v.stack.Pop()
		args[len(n.Args)-1-i] = node.Value
	}

	// Call the function
	result, err := function.Callback(args)
	if err != nil {
		v.err = fmt.Errorf("%s: %s", n.Func, err)
		return
	}

	// Push the result
	v.stack.Push(&ast.LiteralNode{
		Value: result,
		Type:  function.ReturnType,
	})
}

func (v *executeVisitor) visitConcat(n *ast.Concat) {
	// The expressions should all be on the stack in reverse
	// order. So pop them off, reverse their order, and concatenate.
	nodes := make([]*ast.LiteralNode, 0, len(n.Exprs))
	for range n.Exprs {
		nodes = append(nodes, v.stack.Pop())
	}

	var buf bytes.Buffer
	for i := len(nodes) - 1; i >= 0; i-- {
		buf.WriteString(nodes[i].Value.(string))
	}

	v.stack.Push(&ast.LiteralNode{
		Value: buf.String(),
		Type:  ast.TypeString,
	})
}

func (v *executeVisitor) visitLiteral(n *ast.LiteralNode) {
	v.stack.Push(n)
}

func (v *executeVisitor) visitVariableAccess(n *ast.VariableAccess) {
	// Look up the variable in the map
	variable, ok := v.Scope.VarMap[n.Name]
	if !ok {
		v.err = fmt.Errorf("unknown variable accessed: %s", n.Name)
		return
	}

	v.stack.Push(&ast.LiteralNode{
		Value: variable.Value,
		Type:  variable.Type,
	})
}

// EngineStack is a stack of ast.LiteralNodes that the Engine keeps track
// of during execution. This is currently backed by a dumb slice, but can be
// replaced with a better data structure at some point in the future if this
// turns out to require optimization.
type EngineStack struct {
	stack []*ast.LiteralNode
}

func (s *EngineStack) Len() int {
	return len(s.stack)
}

func (s *EngineStack) Push(n *ast.LiteralNode) {
	s.stack = append(s.stack, n)
}

func (s *EngineStack) Pop() *ast.LiteralNode {
	x := s.stack[len(s.stack)-1]
	s.stack[len(s.stack)-1] = nil
	s.stack = s.stack[:len(s.stack)-1]
	return x
}

func (s *EngineStack) Reset() {
	s.stack = nil
}

// Scope represents a lookup scope for execution.
type Scope struct {
	// VarMap and FuncMap are the mappings of identifiers to functions
	// and variable values.
	VarMap  map[string]Variable
	FuncMap map[string]Function
}

// Variable is a variable value for execution given as input to the engine.
// It records the value of a variables along with their type.
type Variable struct {
	Value interface{}
	Type  ast.Type
}

// Function defines a function that can be executed by the engine.
// The type checker will validate that the proper types will be called
// to the callback.
type Function struct {
	ArgTypes   []ast.Type
	ReturnType ast.Type
	Callback   func([]interface{}) (interface{}, error)
}

// LookupFunc will look up a variable by name.
// TODO test
func (s *Scope) LookupFunc(n string) (Function, bool) {
	if s == nil {
		return Function{}, false
	}

	v, ok := s.FuncMap[n]
	return v, ok
}

// LookupVar will look up a variable by name.
// TODO test
func (s *Scope) LookupVar(n string) (Variable, bool) {
	if s == nil {
		return Variable{}, false
	}

	v, ok := s.VarMap[n]
	return v, ok
}
<|endoftext|>


<commit_before>package model

// Dependency has project dependency and source distination.
type Dependency struct {
	ProjectUUID          string `json:"project_uuid"`
	DependentProjectUUID string `json:"dependent_project_uuid"`
	SourceCD             string `json:"source_cd"`
}
<commit_msg>Insert description for source cd<commit_after>package model

// SourceCD
// 0 = unknown
// 1 = carthage
// 2 = cocoapods
// 3 = submodule

// Dependency has project dependency and source distination.
type Dependency struct {
	ProjectUUID          string `json:"project_uuid"`
	DependentProjectUUID string `json:"dependent_project_uuid"`
	SourceCD             int    `json:"source_cd"`
}
<|endoftext|>


<commit_before>package TF2RconWrapper

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"strconv"
	"sync"
	"time"
)

type RawMessage struct {
	data []byte
	n    int //length
}

func (r RawMessage) String() string {
	return string(r.data)
}

// RconChatListener maintains an UDP server that receives redirected chat messages from TF2 servers
type RconChatListener struct {
	conn        *net.UDPConn
	servers     map[string]*ServerListener
	serversLock *sync.RWMutex
	wait        *sync.WaitGroup
	addr        *net.UDPAddr
	localip     string
	port        string
	rng         *rand.Rand
}

// NewRconChatListener builds a new RconChatListener. Its arguments are localip (the ip of this server) and
// port (the port the listener will use)
func NewRconChatListener(localip, port string) (*RconChatListener, error) {
	addr, err := net.ResolveUDPAddr("udp4", ":"+port)
	if err != nil {
		return nil, err
	}

	servers := make(map[string]*ServerListener)

	rng := rand.New(rand.NewSource(time.Now().Unix()))

	listener := &RconChatListener{nil, servers, new(sync.RWMutex), new(sync.WaitGroup), addr, localip, port, rng}
	listener.startListening()
	return listener, nil
}

func (r *RconChatListener) startListening() {
	conn, err := net.ListenUDP("udp", r.addr)
	r.conn = conn
	if err != nil {
		log.Fatal(err)
		return
	}
	conn.SetReadBuffer(1048576)
	go r.readStrings()
}

func (r *RconChatListener) readStrings() {
	rawMessageC := make(chan RawMessage, 500)

	go func() {
		for {
			raw := <-rawMessageC
			r.wait.Wait()
			message := raw.data[0:raw.n]
			secret, err := getSecret(message)
			if err != nil {
				continue
			}

			r.serversLock.RLock()
			s, ok := r.servers[secret]
			r.serversLock.RUnlock()
			if ok {
				s.RawMessages <- raw
			}
		}
	}()

	for {
		buff := make([]byte, 65000)
		n, err := r.conn.Read(buff)
		//log.Println(string(buff[0:n]))

		if err != nil {
			fmt.Println("Error receiving server chat data: ", err)
			continue
		}

		rawMessageC <- RawMessage{buff, n}
	}
}

// Close stops the RconChatListener
func (s *ServerListener) Close(m *TF2RconConnection) {
	s.listener.wait.Add(1)
	delete(s.listener.servers, s.secret)
	s.listener.wait.Done()

	m.StopLogRedirection(s.listener.localip, s.listener.port)
}

// CreateServerListener creates a ServerListener that receives chat messages from a
// particular TF2 server
func (r *RconChatListener) CreateServerListener(m *TF2RconConnection) *ServerListener {

	secret := strconv.Itoa(r.rng.Intn(999998) + 1)

	r.serversLock.RLock()
	_, ok := r.servers[secret]
	for ok {
		secret = strconv.Itoa(r.rng.Intn(999998) + 1)
		_, ok = r.servers[secret]
	}
	r.serversLock.RUnlock()

	return r.CreateListenerWithSecret(m, secret)
}

func (r *RconChatListener) CreateListenerWithSecret(m *TF2RconConnection, secret string) *ServerListener {
	s := &ServerListener{make(chan RawMessage, 10), m.host, secret, r}

	r.serversLock.Lock()
	r.servers[secret] = s
	//log.Println(r.servers)
	r.serversLock.Unlock()

	m.Query("sv_logsecret " + secret)
	m.RedirectLogs(r.localip, r.port)
	return s
}

// ServerListener represents a listener that receives chat messages from a particular
// TF2 server. It's built and managed by an RconChatListener instance.
type ServerListener struct {
	RawMessages chan RawMessage

	host     string
	secret   string
	listener *RconChatListener
}
<commit_msg>Make ServerListener.Server public<commit_after>package TF2RconWrapper

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"strconv"
	"sync"
	"time"
)

type RawMessage struct {
	data []byte
	n    int //length
}

func (r RawMessage) String() string {
	return string(r.data)
}

// RconChatListener maintains an UDP server that receives redirected chat messages from TF2 servers
type RconChatListener struct {
	conn        *net.UDPConn
	servers     map[string]*ServerListener
	serversLock *sync.RWMutex
	wait        *sync.WaitGroup
	addr        *net.UDPAddr
	localip     string
	port        string
	rng         *rand.Rand
}

// NewRconChatListener builds a new RconChatListener. Its arguments are localip (the ip of this server) and
// port (the port the listener will use)
func NewRconChatListener(localip, port string) (*RconChatListener, error) {
	addr, err := net.ResolveUDPAddr("udp4", ":"+port)
	if err != nil {
		return nil, err
	}

	servers := make(map[string]*ServerListener)

	rng := rand.New(rand.NewSource(time.Now().Unix()))

	listener := &RconChatListener{nil, servers, new(sync.RWMutex), new(sync.WaitGroup), addr, localip, port, rng}
	listener.startListening()
	return listener, nil
}

func (r *RconChatListener) startListening() {
	conn, err := net.ListenUDP("udp", r.addr)
	r.conn = conn
	if err != nil {
		log.Fatal(err)
		return
	}
	conn.SetReadBuffer(1048576)
	go r.readStrings()
}

func (r *RconChatListener) readStrings() {
	rawMessageC := make(chan RawMessage, 500)

	go func() {
		for {
			raw := <-rawMessageC
			r.wait.Wait()
			message := raw.data[0:raw.n]
			secret, err := getSecret(message)
			if err != nil {
				continue
			}

			r.serversLock.RLock()
			s, ok := r.servers[secret]
			r.serversLock.RUnlock()
			if ok {
				s.RawMessages <- raw
			}
		}
	}()

	for {
		buff := make([]byte, 65000)
		n, err := r.conn.Read(buff)
		//log.Println(string(buff[0:n]))

		if err != nil {
			fmt.Println("Error receiving server chat data: ", err)
			continue
		}

		rawMessageC <- RawMessage{buff, n}
	}
}

// Close stops the RconChatListener
func (s *ServerListener) Close(m *TF2RconConnection) {
	s.listener.wait.Add(1)
	delete(s.listener.servers, s.Secret)
	s.listener.wait.Done()

	m.StopLogRedirection(s.listener.localip, s.listener.port)
}

// CreateServerListener creates a ServerListener that receives chat messages from a
// particular TF2 server
func (r *RconChatListener) CreateServerListener(m *TF2RconConnection) *ServerListener {

	secret := strconv.Itoa(r.rng.Intn(999998) + 1)

	r.serversLock.RLock()
	_, ok := r.servers[secret]
	for ok {
		secret = strconv.Itoa(r.rng.Intn(999998) + 1)
		_, ok = r.servers[secret]
	}
	r.serversLock.RUnlock()

	return r.CreateListenerWithSecret(m, secret)
}

func (r *RconChatListener) CreateListenerWithSecret(m *TF2RconConnection, secret string) *ServerListener {
	s := &ServerListener{make(chan RawMessage, 10), m.host, secret, r}

	r.serversLock.Lock()
	r.servers[secret] = s
	//log.Println(r.servers)
	r.serversLock.Unlock()

	m.Query("sv_logsecret " + secret)
	m.RedirectLogs(r.localip, r.port)
	return s
}

// ServerListener represents a listener that receives chat messages from a particular
// TF2 server. It's built and managed by an RconChatListener instance.
type ServerListener struct {
	RawMessages chan RawMessage

	host     string
	Secret   string
	listener *RconChatListener
}
<|endoftext|>


<commit_before>package simplestreams

import (
	"fmt"
	"strings"
	"time"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/osarch"
)

// Products represents the base of download.json
type Products struct {
	ContentID string             `json:"content_id"`
	DataType  string             `json:"datatype"`
	Format    string             `json:"format"`
	License   string             `json:"license,omitempty"`
	Products  map[string]Product `json:"products"`
	Updated   string             `json:"updated,omitempty"`
}

// Product represents a single product inside download.json
type Product struct {
	Aliases         string                    `json:"aliases"`
	Architecture    string                    `json:"arch"`
	OperatingSystem string                    `json:"os"`
	Release         string                    `json:"release"`
	ReleaseCodename string                    `json:"release_codename,omitempty"`
	ReleaseTitle    string                    `json:"release_title"`
	Supported       bool                      `json:"supported,omitempty"`
	SupportedEOL    string                    `json:"support_eol,omitempty"`
	Version         string                    `json:"version,omitempty"`
	Versions        map[string]ProductVersion `json:"versions"`
}

// ProductVersion represents a particular version of a product
type ProductVersion struct {
	Items      map[string]ProductVersionItem `json:"items"`
	Label      string                        `json:"label,omitempty"`
	PublicName string                        `json:"pubname,omitempty"`
}

// ProductVersionItem represents a file/item of a particular ProductVersion
type ProductVersionItem struct {
	LXDHashSha256DiskImg    string `json:"combined_disk1-img_sha256,omitempty"`
	LXDHashSha256DiskKvmImg string `json:"combined_disk-kvm-img_sha256,omitempty"`
	LXDHashSha256RootXz     string `json:"combined_rootxz_sha256,omitempty"`
	LXDHashSha256           string `json:"combined_sha256,omitempty"`
	LXDHashSha256SquashFs   string `json:"combined_squashfs_sha256,omitempty"`
	FileType                string `json:"ftype"`
	HashMd5                 string `json:"md5,omitempty"`
	Path                    string `json:"path"`
	HashSha256              string `json:"sha256,omitempty"`
	Size                    int64  `json:"size"`
	DeltaBase               string `json:"delta_base,omitempty"`
}

// ToLXD converts the products data into a list of LXD images and associated downloadable files
func (s *Products) ToLXD() ([]api.Image, map[string][][]string) {
	downloads := map[string][][]string{}

	images := []api.Image{}
	nameLayout := "20060102"
	eolLayout := "2006-01-02"

	for _, product := range s.Products {
		// Skip unsupported architectures
		architecture, err := osarch.ArchitectureId(product.Architecture)
		if err != nil {
			continue
		}

		architectureName, err := osarch.ArchitectureName(architecture)
		if err != nil {
			continue
		}

		for name, version := range product.Versions {
			// Short of anything better, use the name as date (see format above)
			if len(name) < 8 {
				continue
			}

			creationDate, err := time.Parse(nameLayout, name[0:8])
			if err != nil {
				continue
			}

			// Image processing function
			addImage := func(meta *ProductVersionItem, root *ProductVersionItem) error {
				// Look for deltas (only on squashfs)
				deltas := []ProductVersionItem{}
				if root != nil && root.FileType == "squashfs" {
					for _, item := range version.Items {
						if item.FileType == "squashfs.vcdiff" {
							deltas = append(deltas, item)
						}
					}
				}

				// Figure out the fingerprint
				fingerprint := ""
				if root != nil {
					if root.FileType == "root.tar.xz" {
						if meta.LXDHashSha256RootXz != "" {
							fingerprint = meta.LXDHashSha256RootXz
						} else {
							fingerprint = meta.LXDHashSha256
						}
					} else if root.FileType == "squashfs" {
						fingerprint = meta.LXDHashSha256SquashFs
					} else if root.FileType == "disk-kvm.img" {
						fingerprint = meta.LXDHashSha256DiskKvmImg
					} else if root.FileType == "disk1.img" {
						fingerprint = meta.LXDHashSha256DiskImg
					}
				} else {
					fingerprint = meta.HashSha256
				}

				if fingerprint == "" {
					return fmt.Errorf("No LXD image fingerprint found")
				}

				// Figure out the size
				size := meta.Size
				if root != nil {
					size += root.Size
				}

				// Determine filename
				if meta.Path == "" {
					return fmt.Errorf("Missing path field on metadata entry")
				}

				fields := strings.Split(meta.Path, "/")
				filename := fields[len(fields)-1]

				// Generate the actual image entry
				description := fmt.Sprintf("%s %s %s", product.OperatingSystem, product.ReleaseTitle, product.Architecture)
				if version.Label != "" {
					description = fmt.Sprintf("%s (%s)", description, version.Label)
				}
				description = fmt.Sprintf("%s (%s)", description, name)

				image := api.Image{}
				image.Architecture = architectureName
				image.Public = true
				image.Size = size
				image.CreatedAt = creationDate
				image.UploadedAt = creationDate
				image.Filename = filename
				image.Fingerprint = fingerprint
				image.Properties = map[string]string{
					"os":           product.OperatingSystem,
					"release":      product.Release,
					"version":      product.Version,
					"architecture": product.Architecture,
					"label":        version.Label,
					"serial":       name,
					"description":  description,
				}
				image.Type = "container"

				if root != nil {
					image.Properties["type"] = root.FileType
					if root.FileType == "disk1.img" || root.FileType == "disk-kvm.img" {
						image.Type = "virtual-machine"
					}
				} else {
					image.Properties["type"] = "tar.gz"
				}

				// Clear unset properties
				for k, v := range image.Properties {
					if v == "" {
						delete(image.Properties, k)
					}
				}

				// Add the provided aliases
				if product.Aliases != "" {
					image.Aliases = []api.ImageAlias{}
					for _, entry := range strings.Split(product.Aliases, ",") {
						image.Aliases = append(image.Aliases, api.ImageAlias{Name: entry})
					}
				}

				// Attempt to parse the EOL
				image.ExpiresAt = time.Unix(0, 0).UTC()
				if product.SupportedEOL != "" {
					eolDate, err := time.Parse(eolLayout, product.SupportedEOL)
					if err == nil {
						image.ExpiresAt = eolDate
					}
				}

				// Set the file list
				var imgDownloads [][]string
				if root == nil {
					imgDownloads = [][]string{{meta.Path, meta.HashSha256, "meta", fmt.Sprintf("%d", meta.Size)}}
				} else {
					imgDownloads = [][]string{
						{meta.Path, meta.HashSha256, "meta", fmt.Sprintf("%d", meta.Size)},
						{root.Path, root.HashSha256, "root", fmt.Sprintf("%d", root.Size)}}
				}

				// Add the deltas
				for _, delta := range deltas {
					srcImage, ok := product.Versions[delta.DeltaBase]
					if !ok {
						// Delta for a since expired image
						continue
					}

					// Locate source image fingerprint
					var srcFingerprint string
					for _, item := range srcImage.Items {
						if item.FileType != "lxd.tar.xz" {
							continue
						}

						srcFingerprint = item.LXDHashSha256SquashFs
						break
					}

					if srcFingerprint == "" {
						// Couldn't find the image
						continue
					}

					// Add the delta
					imgDownloads = append(imgDownloads, []string{
						delta.Path,
						delta.HashSha256,
						fmt.Sprintf("root.delta-%s", srcFingerprint),
						fmt.Sprintf("%d", delta.Size)})
				}

				// Add the image
				downloads[fingerprint] = imgDownloads
				images = append(images, image)

				return nil
			}

			// Locate a valid LXD image
			for _, item := range version.Items {
				if item.FileType == "lxd_combined.tar.gz" {
					err := addImage(&item, nil)
					if err != nil {
						continue
					}
				}

				if item.FileType == "lxd.tar.xz" {
					// Locate the root files
					for _, subItem := range version.Items {
						if shared.StringInSlice(subItem.FileType, []string{"disk1.img", "disk-kvm.img", "root.tar.xz", "squashfs"}) {
							err := addImage(&item, &subItem)
							if err != nil {
								continue
							}
						}
					}
				}
			}
		}
	}

	return images, downloads
}
<commit_msg>shared/simplestreams: Support uefi1.img<commit_after>package simplestreams

import (
	"fmt"
	"strings"
	"time"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/osarch"
)

// Products represents the base of download.json
type Products struct {
	ContentID string             `json:"content_id"`
	DataType  string             `json:"datatype"`
	Format    string             `json:"format"`
	License   string             `json:"license,omitempty"`
	Products  map[string]Product `json:"products"`
	Updated   string             `json:"updated,omitempty"`
}

// Product represents a single product inside download.json
type Product struct {
	Aliases         string                    `json:"aliases"`
	Architecture    string                    `json:"arch"`
	OperatingSystem string                    `json:"os"`
	Release         string                    `json:"release"`
	ReleaseCodename string                    `json:"release_codename,omitempty"`
	ReleaseTitle    string                    `json:"release_title"`
	Supported       bool                      `json:"supported,omitempty"`
	SupportedEOL    string                    `json:"support_eol,omitempty"`
	Version         string                    `json:"version,omitempty"`
	Versions        map[string]ProductVersion `json:"versions"`
}

// ProductVersion represents a particular version of a product
type ProductVersion struct {
	Items      map[string]ProductVersionItem `json:"items"`
	Label      string                        `json:"label,omitempty"`
	PublicName string                        `json:"pubname,omitempty"`
}

// ProductVersionItem represents a file/item of a particular ProductVersion
type ProductVersionItem struct {
	LXDHashSha256DiskImg     string `json:"combined_disk1-img_sha256,omitempty"`
	LXDHashSha256DiskKvmImg  string `json:"combined_disk-kvm-img_sha256,omitempty"`
	LXDHashSha256DiskUefiImg string `json:"combined_uefi1-img_sha256,omitempty"`
	LXDHashSha256RootXz      string `json:"combined_rootxz_sha256,omitempty"`
	LXDHashSha256            string `json:"combined_sha256,omitempty"`
	LXDHashSha256SquashFs    string `json:"combined_squashfs_sha256,omitempty"`
	FileType                 string `json:"ftype"`
	HashMd5                  string `json:"md5,omitempty"`
	Path                     string `json:"path"`
	HashSha256               string `json:"sha256,omitempty"`
	Size                     int64  `json:"size"`
	DeltaBase                string `json:"delta_base,omitempty"`
}

// ToLXD converts the products data into a list of LXD images and associated downloadable files
func (s *Products) ToLXD() ([]api.Image, map[string][][]string) {
	downloads := map[string][][]string{}

	images := []api.Image{}
	nameLayout := "20060102"
	eolLayout := "2006-01-02"

	for _, product := range s.Products {
		// Skip unsupported architectures
		architecture, err := osarch.ArchitectureId(product.Architecture)
		if err != nil {
			continue
		}

		architectureName, err := osarch.ArchitectureName(architecture)
		if err != nil {
			continue
		}

		for name, version := range product.Versions {
			// Short of anything better, use the name as date (see format above)
			if len(name) < 8 {
				continue
			}

			creationDate, err := time.Parse(nameLayout, name[0:8])
			if err != nil {
				continue
			}

			// Image processing function
			addImage := func(meta *ProductVersionItem, root *ProductVersionItem) error {
				// Look for deltas (only on squashfs)
				deltas := []ProductVersionItem{}
				if root != nil && root.FileType == "squashfs" {
					for _, item := range version.Items {
						if item.FileType == "squashfs.vcdiff" {
							deltas = append(deltas, item)
						}
					}
				}

				// Figure out the fingerprint
				fingerprint := ""
				if root != nil {
					if root.FileType == "root.tar.xz" {
						if meta.LXDHashSha256RootXz != "" {
							fingerprint = meta.LXDHashSha256RootXz
						} else {
							fingerprint = meta.LXDHashSha256
						}
					} else if root.FileType == "squashfs" {
						fingerprint = meta.LXDHashSha256SquashFs
					} else if root.FileType == "disk-kvm.img" {
						fingerprint = meta.LXDHashSha256DiskKvmImg
					} else if root.FileType == "disk1.img" {
						fingerprint = meta.LXDHashSha256DiskImg
					} else if root.FileType == "uefi1.img" {
						fingerprint = meta.LXDHashSha256DiskUefiImg
					}
				} else {
					fingerprint = meta.HashSha256
				}

				if fingerprint == "" {
					return fmt.Errorf("No LXD image fingerprint found")
				}

				// Figure out the size
				size := meta.Size
				if root != nil {
					size += root.Size
				}

				// Determine filename
				if meta.Path == "" {
					return fmt.Errorf("Missing path field on metadata entry")
				}

				fields := strings.Split(meta.Path, "/")
				filename := fields[len(fields)-1]

				// Generate the actual image entry
				description := fmt.Sprintf("%s %s %s", product.OperatingSystem, product.ReleaseTitle, product.Architecture)
				if version.Label != "" {
					description = fmt.Sprintf("%s (%s)", description, version.Label)
				}
				description = fmt.Sprintf("%s (%s)", description, name)

				image := api.Image{}
				image.Architecture = architectureName
				image.Public = true
				image.Size = size
				image.CreatedAt = creationDate
				image.UploadedAt = creationDate
				image.Filename = filename
				image.Fingerprint = fingerprint
				image.Properties = map[string]string{
					"os":           product.OperatingSystem,
					"release":      product.Release,
					"version":      product.Version,
					"architecture": product.Architecture,
					"label":        version.Label,
					"serial":       name,
					"description":  description,
				}
				image.Type = "container"

				if root != nil {
					image.Properties["type"] = root.FileType
					if root.FileType == "disk1.img" || root.FileType == "disk-kvm.img" || root.FileType == "uefi1.img" {
						image.Type = "virtual-machine"
					}
				} else {
					image.Properties["type"] = "tar.gz"
				}

				// Clear unset properties
				for k, v := range image.Properties {
					if v == "" {
						delete(image.Properties, k)
					}
				}

				// Add the provided aliases
				if product.Aliases != "" {
					image.Aliases = []api.ImageAlias{}
					for _, entry := range strings.Split(product.Aliases, ",") {
						image.Aliases = append(image.Aliases, api.ImageAlias{Name: entry})
					}
				}

				// Attempt to parse the EOL
				image.ExpiresAt = time.Unix(0, 0).UTC()
				if product.SupportedEOL != "" {
					eolDate, err := time.Parse(eolLayout, product.SupportedEOL)
					if err == nil {
						image.ExpiresAt = eolDate
					}
				}

				// Set the file list
				var imgDownloads [][]string
				if root == nil {
					imgDownloads = [][]string{{meta.Path, meta.HashSha256, "meta", fmt.Sprintf("%d", meta.Size)}}
				} else {
					imgDownloads = [][]string{
						{meta.Path, meta.HashSha256, "meta", fmt.Sprintf("%d", meta.Size)},
						{root.Path, root.HashSha256, "root", fmt.Sprintf("%d", root.Size)}}
				}

				// Add the deltas
				for _, delta := range deltas {
					srcImage, ok := product.Versions[delta.DeltaBase]
					if !ok {
						// Delta for a since expired image
						continue
					}

					// Locate source image fingerprint
					var srcFingerprint string
					for _, item := range srcImage.Items {
						if item.FileType != "lxd.tar.xz" {
							continue
						}

						srcFingerprint = item.LXDHashSha256SquashFs
						break
					}

					if srcFingerprint == "" {
						// Couldn't find the image
						continue
					}

					// Add the delta
					imgDownloads = append(imgDownloads, []string{
						delta.Path,
						delta.HashSha256,
						fmt.Sprintf("root.delta-%s", srcFingerprint),
						fmt.Sprintf("%d", delta.Size)})
				}

				// Add the image
				downloads[fingerprint] = imgDownloads
				images = append(images, image)

				return nil
			}

			// Locate a valid LXD image
			for _, item := range version.Items {
				if item.FileType == "lxd_combined.tar.gz" {
					err := addImage(&item, nil)
					if err != nil {
						continue
					}
				}

				if item.FileType == "lxd.tar.xz" {
					// Locate the root files
					for _, subItem := range version.Items {
						if shared.StringInSlice(subItem.FileType, []string{"disk1.img", "disk-kvm.img", "uefi1.img", "root.tar.xz", "squashfs"}) {
							err := addImage(&item, &subItem)
							if err != nil {
								continue
							}
						}
					}
				}
			}
		}
	}

	return images, downloads
}
<|endoftext|>