<commit_before>// Package gpool is a utility for executing jobs in a pool of workers
package gpool

import (
	"errors"
	"sync"
	"time"
)

// ErrClosedPool indicates that a send was attempted on a pool which has already been closed.
var ErrClosedPool = errors.New("send on closed pool")

// ErrKilled indicates that the pool was killed by a call to Kill()
var ErrKilled = errors.New("pool killed by signal")

// Pool is the main pool struct containing a bus and workers.
// Pool should always be invoked via NewPool().
type Pool struct {
	// Workers
	wQ chan jobRequest // Queue
	wR chan JobResult  // Return
	wC chan struct{}   // Cancel
	wD chan bool       // Done

	fJ []JobResult
	iJ int // Job id

	tQ chan ticket // Send ticket requests

	s   *stateManager
	err error

	wg *sync.WaitGroup
	// Hooks are optional functions that are executed during different stages of a Job.
	// They are invoked by the worker and thus are called concurrently.
	// The implementor should consider any race conditions for hook callbacks.
	Hook struct {
		Start HookStart
		Stop  HookStop
	}

	closed bool
	killed bool
	done   bool
}

// NewPool creates a new Pool with the given worker count.
// Workers in the pool are started automatically.
func NewPool(Workers int) *Pool {
	p := &Pool{
		tQ: make(chan ticket),
		wQ: make(chan jobRequest),
		wC: make(chan struct{}),
		wR: make(chan JobResult),
		wD: make(chan bool),
		wg: &sync.WaitGroup{},
		s:  newStateManager(Workers),
	}
	p.start()
	return p
}

// start starts Pool workers and Pool bus
func (p *Pool) start() {
	for range make([]int, p.s.GetTarget()) {
		p.startWorker()
	}
	go func() {
		p.wg.Wait()
		p.wD <- true
	}()
	go p.bus()
}

// ticket request types
const (
	tReqJob int = 1 << iota
	tReqClose
	tReqKill
	tReqWait
	tReqGrow
	tReqHealthy
	tReqGetError
)

// A ticket is a request for input in the queue.
// This prevents direct access to queue channels which reduces the risk of bad things happening.
type ticket struct {
	t    int         // Ticket type
	data interface{} // Ticket request data
	r    chan error  // Return channel
}

// Kill sends a kill request to the pool bus.
// When sent, any currently running jobs have Cancel() called.
// If the pool has already been killed or closed ErrClosedPool is returned.
func (p *Pool) Kill() error {
	t := ticket{
		tReqKill, nil, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Close sends a graceful close request to the pool bus.
// Workers will finish after the last submitted job is complete.
// Close does not return an error if the pool is already closed.
func (p *Pool) Close() error {
	t := ticket{
		tReqClose, nil, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Wait waits for the pool worker group to finish and then returns all jobs completed during execution.
// If the pool has an error it is returned here.
// Wait will block until the pool is marked as done (via call to Pool.Close) or has an error.
// As such it is important that the caller either implements a timeout around Wait,
// or ensures a call to Pool.Close will be made.
// Wait requests are stacked until they can be resolved in the next bus cycle.
func (p *Pool) Wait() ([]JobResult, error) {
	t := ticket{
		tReqWait, nil, make(chan error),
	}
	p.tQ <- t
	return p.fJ, <-t.r
}

// Send sends the given PoolJob as a request to the pool bus.
// If the pool is closed the error ErrClosedPool is returned.
// No error is returned if the Send() was successful.
// A call to Send is blocked until a worker accepts the Job.
func (p *Pool) Send(job Job) error {
	if job == nil {
		panic("send of nil job")
	}
	t := ticket{
		tReqJob, job, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Healthy returns true if the pool is healthy and able to receive further jobs.
func (p *Pool) Healthy() bool {
	t := ticket{
		tReqHealthy, nil, make(chan error),
	}
	p.tQ <- t
	e := <-t.r
	if e == nil {
		return true
	}
	return false
}

// Error returns the current error present in the pool.
func (p *Pool) Error() error {
	t := ticket{
		tReqGetError, nil, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Grow grows the amount of workers running in the pool by the requested amount.
// If the pool is closed ErrClosedPool is return.
func (p *Pool) Grow(Req int) error {
	if Req == 0 {
		return nil
	}
	t := ticket{
		tReqGrow, Req, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// WorkerState returns the current amount of running workers and the amount of requested workers.
func (p *Pool) WorkerState() (c int, t int) {
	return p.s.WorkerState()
}

// JobState returns the amount of currently executing jobs in the pool
// and the amount of executed jobs in the pool (including failed).
func (p *Pool) JobState() (c int, f int) {
	return p.s.JobState()
}

type jobRequest struct {
	Job Job
	ID  int
}

// close closes the pool if not already closed.
// Returns true if the pool was closed.
func (p *Pool) close(kill bool) bool {
	if !p.killed && kill {
		p.killed = true
		if p.err == nil {
			p.err = ErrKilled
		}
		close(p.wC)
	}
	if !p.closed {
		p.closed = true
		close(p.wQ)
		return true
	}
	return false
}

// res receives a JobResult and processes it.
func (p *Pool) res(resp JobResult) {
	if resp.Error != nil {
		p.err = resp.Error
		p.close(true)
		return
	}
	p.fJ = append(p.fJ, resp)
}

// doSend tries to send a jobRequest to the worker pool.
// If the pool is full a deadlock is possible unless we also collect the return message from the workers
func (p *Pool) doSend(r jobRequest) {
	for {
		select {
		case p.wQ <- r:
			// Request successfully sent
			return
		case resp := <-p.wR:
			p.res(resp)
		}
	}
}

func (p *Pool) unhealthy() bool {
	return p.done || p.err != nil || p.closed || p.killed
}

// bus is the central communication bus for the pool.
// All pool inputs are collected here as tickets and then actioned on depending on the ticket type.
// If the ticket is a tReqWait request they are added to the slice of pending tickets and upon pool closure.
// A central bus is used to prevent concurrent access to any property of the pool.
func (p *Pool) bus() {
	// Pending tReqWaits
	pendWait := []ticket{}
	for {
		select {
		case <-p.wD:
			p.close(false)
			p.done = true
		// Receive worker response
		case resp := <-p.wR:
			p.res(resp)
		// Receive ticket request
		case t := <-p.tQ:
			switch t.t {
			// New Job request
			case tReqJob:
				if p.unhealthy() {
					t.r <- ErrClosedPool
					continue
				}
				p.iJ++
				p.doSend(jobRequest{
					ID:  p.iJ,
					Job: t.data.(Job),
				})
				t.r <- nil
			// Pool close request
			case tReqClose:
				if p.done {
					t.r <- ErrClosedPool
					continue
				}
				p.close(false)
				t.r <- nil
			// Pool kill request
			case tReqKill:
				if p.done {
					t.r <- ErrClosedPool
					continue
				}
				if ok := p.close(true); ok {
					t.r <- nil
				} else {
					t.r <- ErrClosedPool
				}
			// Pool wait request
			case tReqWait:
				pendWait = append(pendWait, t)
			case tReqGrow:
				if p.unhealthy() {
					t.r <- ErrClosedPool
				}
				for range make([]int, t.data.(int)) {
					p.s.IncrTarget()
					p.startWorker()
				}
				t.r <- nil
			// Pool is open request
			case tReqHealthy:
				if p.unhealthy() {
					t.r <- ErrClosedPool
					continue
				}
				t.r <- nil
			case tReqGetError:
				t.r <- p.err
			default:
				panic("unknown ticket type")
			}
		// Default case resolves any pending tickets
		default:
			if p.done && len(pendWait) != 0 {
				// Pop every item in pending and send return signal to ticket
				for {
					if len(pendWait) == 0 {
						// No more pending message, continue with next cycle
						break
					}
					var t ticket
					t, pendWait = pendWait[len(pendWait)-1], pendWait[:len(pendWait)-1]
					t.r <- p.err
				}
			}

		}

	}
}

// startWorker starts a worker and waits for it to complete starting before return.
// This prevents a race condition on call to Pool.Running() where the worker may not have completed it's state registration.
func (p *Pool) startWorker() {
	ok := make(chan bool)
	p.wg.Add(1)
	go p.worker(ok)
	<-ok
}

// worker is the Pool worker routine that receives jobRequests from the pool worker queue until the queue is closed.
// It executes the job and returns the result to the worker return channel as a JobResult.
func (p *Pool) worker(started chan bool) {
	defer p.wg.Done()
	p.s.AddWorker()
	defer p.s.RemoveWorker()
	// Acknowledge start
	close(started)
	for {
		select {
		case req, ok := <-p.wQ:
			// Worker queue has been closed
			if !ok {
				return
			}

			p.s.AddJob()

			if p.Hook.Start != nil {
				p.Hook.Start(req.ID, req.Job)
			}

			// Cancel wrapper
			d := make(chan struct{})
			go func() {
				select {
				// Job done
				case <-d:
					return
				// Worker cancel signal received
				case <-p.wC:
					req.Job.Cancel()
				}
			}()

			s := time.Now()

			// Execute Job
			e := req.Job.Run()

			// Close cancel wrapping
			close(d)

			if e != nil {
				e = newPoolError(&req, e)
			}
			j := JobResult{
				Job:      req.Job,
				ID:       req.ID,
				Duration: time.Since(s),
				Error:    e,
			}
			// Send Job result to return queue
			p.wR <- j

			if p.Hook.Stop != nil {
				p.Hook.Stop(req.ID, j)
			}

			p.s.RemoveJob()

		case <-p.wC:
			return
		}
	}
}
<commit_msg>ticket request type as type<commit_after>// Package gpool is a utility for executing jobs in a pool of workers
package gpool

import (
	"errors"
	"sync"
	"time"
)

// ErrClosedPool indicates that a send was attempted on a pool which has already been closed.
var ErrClosedPool = errors.New("send on closed pool")

// ErrKilled indicates that the pool was killed by a call to Kill()
var ErrKilled = errors.New("pool killed by signal")

// Pool is the main pool struct containing a bus and workers.
// Pool should always be invoked via NewPool().
type Pool struct {
	// Workers
	wQ chan jobRequest // Queue
	wR chan JobResult  // Return
	wC chan struct{}   // Cancel
	wD chan bool       // Done

	fJ []JobResult
	iJ int // Job id

	tQ chan ticket // Send ticket requests

	s   *stateManager
	err error

	wg *sync.WaitGroup
	// Hooks are optional functions that are executed during different stages of a Job.
	// They are invoked by the worker and thus are called concurrently.
	// The implementor should consider any race conditions for hook callbacks.
	Hook struct {
		Start HookStart
		Stop  HookStop
	}

	closed bool
	killed bool
	done   bool
}

// NewPool creates a new Pool with the given worker count.
// Workers in the pool are started automatically.
func NewPool(Workers int) *Pool {
	p := &Pool{
		tQ: make(chan ticket),
		wQ: make(chan jobRequest),
		wC: make(chan struct{}),
		wR: make(chan JobResult),
		wD: make(chan bool),
		wg: &sync.WaitGroup{},
		s:  newStateManager(Workers),
	}
	p.start()
	return p
}

// start starts Pool workers and Pool bus
func (p *Pool) start() {
	for range make([]int, p.s.GetTarget()) {
		p.startWorker()
	}
	go func() {
		p.wg.Wait()
		p.wD <- true
	}()
	go p.bus()
}

// ticket request types
type tReq int

const (
	tReqJob tReq = 1 << iota
	tReqClose
	tReqKill
	tReqWait
	tReqGrow
	tReqHealthy
	tReqGetError
)

// A ticket is a request for input in the queue.
// This prevents direct access to queue channels which reduces the risk of bad things happening.
type ticket struct {
	t    tReq        // Ticket type
	data interface{} // Ticket request data
	r    chan error  // Return channel
}

// Kill sends a kill request to the pool bus.
// When sent, any currently running jobs have Cancel() called.
// If the pool has already been killed or closed ErrClosedPool is returned.
func (p *Pool) Kill() error {
	t := ticket{
		tReqKill, nil, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Close sends a graceful close request to the pool bus.
// Workers will finish after the last submitted job is complete.
// Close does not return an error if the pool is already closed.
func (p *Pool) Close() error {
	t := ticket{
		tReqClose, nil, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Wait waits for the pool worker group to finish and then returns all jobs completed during execution.
// If the pool has an error it is returned here.
// Wait will block until the pool is marked as done (via call to Pool.Close) or has an error.
// As such it is important that the caller either implements a timeout around Wait,
// or ensures a call to Pool.Close will be made.
// Wait requests are stacked until they can be resolved in the next bus cycle.
func (p *Pool) Wait() ([]JobResult, error) {
	t := ticket{
		tReqWait, nil, make(chan error),
	}
	p.tQ <- t
	return p.fJ, <-t.r
}

// Send sends the given PoolJob as a request to the pool bus.
// If the pool is closed the error ErrClosedPool is returned.
// No error is returned if the Send() was successful.
// A call to Send is blocked until a worker accepts the Job.
func (p *Pool) Send(job Job) error {
	if job == nil {
		panic("send of nil job")
	}
	t := ticket{
		tReqJob, job, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Healthy returns true if the pool is healthy and able to receive further jobs.
func (p *Pool) Healthy() bool {
	t := ticket{
		tReqHealthy, nil, make(chan error),
	}
	p.tQ <- t
	if <-t.r == nil {
		return true
	}
	return false
}

// Error returns the current error present in the pool.
func (p *Pool) Error() error {
	t := ticket{
		tReqGetError, nil, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// Grow grows the amount of workers running in the pool by the requested amount.
// If the pool is closed ErrClosedPool is return.
func (p *Pool) Grow(Req int) error {
	if Req == 0 {
		return nil
	}
	t := ticket{
		tReqGrow, Req, make(chan error),
	}
	p.tQ <- t
	return <-t.r
}

// WorkerState returns the current amount of running workers and the amount of requested workers.
func (p *Pool) WorkerState() (c int, t int) {
	return p.s.WorkerState()
}

// JobState returns the amount of currently executing jobs in the pool
// and the amount of executed jobs in the pool (including failed).
func (p *Pool) JobState() (c int, f int) {
	return p.s.JobState()
}

type jobRequest struct {
	Job Job
	ID  int
}

// close closes the pool if not already closed.
// Returns true if the pool was closed.
func (p *Pool) close(kill bool) bool {
	if !p.killed && kill {
		p.killed = true
		if p.err == nil {
			p.err = ErrKilled
		}
		close(p.wC)
	}
	if !p.closed {
		p.closed = true
		close(p.wQ)
		return true
	}
	return false
}

// res receives a JobResult and processes it.
func (p *Pool) res(resp JobResult) {
	if resp.Error != nil {
		p.err = resp.Error
		p.close(true)
		return
	}
	p.fJ = append(p.fJ, resp)
}

// doSend tries to send a jobRequest to the worker pool.
// If the pool is full a deadlock is possible unless we also collect the return message from the workers
func (p *Pool) doSend(r jobRequest) {
	for {
		select {
		case p.wQ <- r:
			// Request successfully sent
			return
		case resp := <-p.wR:
			p.res(resp)
		}
	}
}

func (p *Pool) unhealthy() bool {
	return p.done || p.err != nil || p.closed || p.killed
}

// bus is the central communication bus for the pool.
// All pool inputs are collected here as tickets and then actioned on depending on the ticket type.
// If the ticket is a tReqWait request they are added to the slice of pending tickets and upon pool closure.
// A central bus is used to prevent concurrent access to any property of the pool.
func (p *Pool) bus() {
	// Pending tReqWaits
	pendWait := []ticket{}
	for {
		select {
		case <-p.wD:
			p.close(false)
			p.done = true
		// Receive worker response
		case resp := <-p.wR:
			p.res(resp)
		// Receive ticket request
		case t := <-p.tQ:
			switch t.t {
			// New Job request
			case tReqJob:
				if p.unhealthy() {
					t.r <- ErrClosedPool
					continue
				}
				p.iJ++
				p.doSend(jobRequest{
					ID:  p.iJ,
					Job: t.data.(Job),
				})
				t.r <- nil
			// Pool close request
			case tReqClose:
				if p.done {
					t.r <- ErrClosedPool
					continue
				}
				p.close(false)
				t.r <- nil
			// Pool kill request
			case tReqKill:
				if p.done {
					t.r <- ErrClosedPool
					continue
				}
				if ok := p.close(true); ok {
					t.r <- nil
				} else {
					t.r <- ErrClosedPool
				}
			// Pool wait request
			case tReqWait:
				pendWait = append(pendWait, t)
			case tReqGrow:
				if p.unhealthy() {
					t.r <- ErrClosedPool
				}
				for range make([]int, t.data.(int)) {
					p.s.IncrTarget()
					p.startWorker()
				}
				t.r <- nil
			// Pool is open request
			case tReqHealthy:
				if p.unhealthy() {
					t.r <- ErrClosedPool
					continue
				}
				t.r <- nil
			case tReqGetError:
				t.r <- p.err
			default:
				panic("unknown ticket type")
			}
		// Default case resolves any pending tickets
		default:
			if p.done && len(pendWait) != 0 {
				// Pop every item in pending and send return signal to ticket
				for {
					if len(pendWait) == 0 {
						// No more pending message, continue with next cycle
						break
					}
					var t ticket
					t, pendWait = pendWait[len(pendWait)-1], pendWait[:len(pendWait)-1]
					t.r <- p.err
				}
			}

		}

	}
}

// startWorker starts a worker and waits for it to complete starting before return.
// This prevents a race condition on call to Pool.Running() where the worker may not have completed it's state registration.
func (p *Pool) startWorker() {
	ok := make(chan bool)
	p.wg.Add(1)
	go p.worker(ok)
	<-ok
}

// worker is the Pool worker routine that receives jobRequests from the pool worker queue until the queue is closed.
// It executes the job and returns the result to the worker return channel as a JobResult.
func (p *Pool) worker(started chan bool) {
	defer p.wg.Done()
	p.s.AddWorker()
	defer p.s.RemoveWorker()
	// Acknowledge start
	close(started)
	for {
		select {
		case req, ok := <-p.wQ:
			// Worker queue has been closed
			if !ok {
				return
			}

			p.s.AddJob()

			if p.Hook.Start != nil {
				p.Hook.Start(req.ID, req.Job)
			}

			// Cancel wrapper
			d := make(chan struct{})
			go func() {
				select {
				// Job done
				case <-d:
					return
				// Worker cancel signal received
				case <-p.wC:
					req.Job.Cancel()
				}
			}()

			s := time.Now()

			// Execute Job
			e := req.Job.Run()

			// Close cancel wrapping
			close(d)

			if e != nil {
				e = newPoolError(&req, e)
			}
			j := JobResult{
				Job:      req.Job,
				ID:       req.ID,
				Duration: time.Since(s),
				Error:    e,
			}
			// Send Job result to return queue
			p.wR <- j

			if p.Hook.Stop != nil {
				p.Hook.Stop(req.ID, j)
			}

			p.s.RemoveJob()

		case <-p.wC:
			return
		}
	}
}
<|endoftext|>


<commit_before>package router

import (
	"encoding/json"
	. "launchpad.net/gocheck"
	"router/config"
	"time"
)

type RegistrySuite struct {
	*Registry
}

var _ = Suite(&RegistrySuite{})

var fooReg = &registryMessage{
	Host: "192.168.1.1",
	Port: 1234,
	Uris: []Uri{"foo.vcap.me", "fooo.vcap.me"},
	Tags: map[string]string{
		"runtime":   "ruby18",
		"framework": "sinatra",
	},
	App: "12345",
}

var barReg = &registryMessage{
	Host: "192.168.1.2",
	Port: 4321,
	Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	Tags: map[string]string{
		"runtime":   "javascript",
		"framework": "node",
	},
	App: "54321",
}

var bar2Reg = &registryMessage{
	Host: "192.168.1.3",
	Port: 1234,
	Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	Tags: map[string]string{
		"runtime":   "javascript",
		"framework": "node",
	},
	App: "54321",
}

func (s *RegistrySuite) SetUpTest(c *C) {
	var x *config.Config

	x = config.DefaultConfig()
	x.DropletStaleThreshold = 1

	s.Registry = NewRegistry(x)
}

func (s *RegistrySuite) TestRegister(c *C) {
	s.Register(fooReg)
	c.Check(s.NumUris(), Equals, 2)

	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 4)

	c.Assert(s.staleTracker.Len(), Equals, 2)
}

func (s *RegistrySuite) TestRegisterIgnoreEmpty(c *C) {
	s.Register(&registryMessage{})
	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestRegisterIgnoreDuplicates(c *C) {
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Unregister(barReg)
	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestRegisterUppercase(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1235,
		Uris: []Uri{"FOO.VCAP.ME"},
	}

	s.Register(m1)
	s.Register(m2)

	c.Check(s.NumUris(), Equals, 1)
}

func (s *RegistrySuite) TestRegisterDoesntReplace(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"bar.vcap.me"},
	}

	s.Register(m1)
	s.Register(m2)

	c.Check(s.NumUris(), Equals, 2)
}

func (s *RegistrySuite) TestRegisterWithoutUris(c *C) {
	m := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{},
	}

	s.Register(m)

	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestUnregister(c *C) {
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Register(bar2Reg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 2)

	s.Unregister(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Unregister(bar2Reg)
	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestUnregisterUppercase(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"FOO.VCAP.ME"},
	}

	s.Register(m1)
	s.Unregister(m2)

	c.Check(s.NumUris(), Equals, 0)
}

func (s *RegistrySuite) TestUnregisterDoesntDemolish(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me", "bar.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	s.Register(m1)
	s.Unregister(m2)

	c.Check(s.NumUris(), Equals, 1)
}

func (s *RegistrySuite) TestLookup(c *C) {
	m := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	s.Register(m)

	var b *Backend
	var ok bool

	b, ok = s.Lookup("foo.vcap.me")
	c.Assert(ok, Equals, true)
	c.Check(b.BackendId, Equals, BackendId("192.168.1.1:1234"))

	b, ok = s.Lookup("FOO.VCAP.ME")
	c.Assert(ok, Equals, true)
	c.Check(b.BackendId, Equals, BackendId("192.168.1.1:1234"))
}

func (s *RegistrySuite) TestLookupDoubleRegister(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.2",
		Port: 1234,
		Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.2",
		Port: 1235,
		Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	}

	s.Register(m1)
	s.Register(m2)

	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 2)
}

func (s *RegistrySuite) TestTracker(c *C) {
	s.Register(fooReg)
	s.Register(barReg)
	c.Assert(s.staleTracker.Len(), Equals, 2)

	s.Unregister(fooReg)
	s.Unregister(barReg)
	c.Assert(s.staleTracker.Len(), Equals, 0)
}

func (s *RegistrySuite) TestPruneStaleApps(c *C) {
	s.Register(fooReg)
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 4)
	c.Check(s.NumBackends(), Equals, 2)
	c.Assert(s.staleTracker.Len(), Equals, 2)

	time.Sleep(s.dropletStaleThreshold + 1*time.Millisecond)
	s.PruneStaleDroplets()

	s.Register(bar2Reg)

	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)
	c.Assert(s.staleTracker.Len(), Equals, 1)
}

func (s *RegistrySuite) TestInfoMarshalling(c *C) {
	m := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	s.Register(m)
	marshalled, err := json.Marshal(s)
	c.Check(err, IsNil)
	c.Check(string(marshalled), Equals, "{\"foo.vcap.me\":[\"192.168.1.1:1234\"]}")
}
<commit_msg>add test for stale tracker when state is stale<commit_after>package router

import (
	"encoding/json"
	. "launchpad.net/gocheck"
	"router/config"
	"time"
)

type RegistrySuite struct {
	*Registry
}

var _ = Suite(&RegistrySuite{})

var fooReg = &registryMessage{
	Host: "192.168.1.1",
	Port: 1234,
	Uris: []Uri{"foo.vcap.me", "fooo.vcap.me"},
	Tags: map[string]string{
		"runtime":   "ruby18",
		"framework": "sinatra",
	},
	App: "12345",
}

var barReg = &registryMessage{
	Host: "192.168.1.2",
	Port: 4321,
	Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	Tags: map[string]string{
		"runtime":   "javascript",
		"framework": "node",
	},
	App: "54321",
}

var bar2Reg = &registryMessage{
	Host: "192.168.1.3",
	Port: 1234,
	Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	Tags: map[string]string{
		"runtime":   "javascript",
		"framework": "node",
	},
	App: "54321",
}

func (s *RegistrySuite) SetUpTest(c *C) {
	var x *config.Config

	x = config.DefaultConfig()
	x.DropletStaleThreshold = 1

	s.Registry = NewRegistry(x)
}

func (s *RegistrySuite) TestRegister(c *C) {
	s.Register(fooReg)
	c.Check(s.NumUris(), Equals, 2)

	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 4)

	c.Assert(s.staleTracker.Len(), Equals, 2)
}

func (s *RegistrySuite) TestRegisterIgnoreEmpty(c *C) {
	s.Register(&registryMessage{})
	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestRegisterIgnoreDuplicates(c *C) {
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Unregister(barReg)
	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestRegisterUppercase(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1235,
		Uris: []Uri{"FOO.VCAP.ME"},
	}

	s.Register(m1)
	s.Register(m2)

	c.Check(s.NumUris(), Equals, 1)
}

func (s *RegistrySuite) TestRegisterDoesntReplace(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"bar.vcap.me"},
	}

	s.Register(m1)
	s.Register(m2)

	c.Check(s.NumUris(), Equals, 2)
}

func (s *RegistrySuite) TestRegisterWithoutUris(c *C) {
	m := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{},
	}

	s.Register(m)

	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestUnregister(c *C) {
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Register(bar2Reg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 2)

	s.Unregister(barReg)
	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)

	s.Unregister(bar2Reg)
	c.Check(s.NumUris(), Equals, 0)
	c.Check(s.NumBackends(), Equals, 0)
}

func (s *RegistrySuite) TestUnregisterUppercase(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"FOO.VCAP.ME"},
	}

	s.Register(m1)
	s.Unregister(m2)

	c.Check(s.NumUris(), Equals, 0)
}

func (s *RegistrySuite) TestUnregisterDoesntDemolish(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me", "bar.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	s.Register(m1)
	s.Unregister(m2)

	c.Check(s.NumUris(), Equals, 1)
}

func (s *RegistrySuite) TestLookup(c *C) {
	m := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	s.Register(m)

	var b *Backend
	var ok bool

	b, ok = s.Lookup("foo.vcap.me")
	c.Assert(ok, Equals, true)
	c.Check(b.BackendId, Equals, BackendId("192.168.1.1:1234"))

	b, ok = s.Lookup("FOO.VCAP.ME")
	c.Assert(ok, Equals, true)
	c.Check(b.BackendId, Equals, BackendId("192.168.1.1:1234"))
}

func (s *RegistrySuite) TestLookupDoubleRegister(c *C) {
	m1 := &registryMessage{
		Host: "192.168.1.2",
		Port: 1234,
		Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	}

	m2 := &registryMessage{
		Host: "192.168.1.2",
		Port: 1235,
		Uris: []Uri{"bar.vcap.me", "barr.vcap.me"},
	}

	s.Register(m1)
	s.Register(m2)

	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 2)
}

func (s *RegistrySuite) TestTracker(c *C) {
	s.Register(fooReg)
	s.Register(barReg)
	c.Assert(s.staleTracker.Len(), Equals, 2)

	s.Unregister(fooReg)
	s.Unregister(barReg)
	c.Assert(s.staleTracker.Len(), Equals, 0)
}

func (s *RegistrySuite) TestPruneStaleApps(c *C) {
	s.Register(fooReg)
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 4)
	c.Check(s.NumBackends(), Equals, 2)
	c.Assert(s.staleTracker.Len(), Equals, 2)

	time.Sleep(s.dropletStaleThreshold + 1*time.Millisecond)
	s.PruneStaleDroplets()

	s.Register(bar2Reg)

	c.Check(s.NumUris(), Equals, 2)
	c.Check(s.NumBackends(), Equals, 1)
	c.Assert(s.staleTracker.Len(), Equals, 1)
}

func (s *RegistrySuite) TestPruneStaleAppsWhenStateStale(c *C) {
	s.Register(fooReg)
	s.Register(barReg)
	c.Check(s.NumUris(), Equals, 4)
	c.Check(s.NumBackends(), Equals, 2)
	c.Assert(s.staleTracker.Len(), Equals, 2)

	time.Sleep(s.dropletStaleThreshold + 1*time.Millisecond)
	s.isStateStale = func() bool { return true }
	s.PruneStaleDroplets()

	c.Check(s.NumUris(), Equals, 4)
	c.Check(s.NumBackends(), Equals, 2)
	c.Assert(s.staleTracker.Len(), Equals, 0)
}

func (s *RegistrySuite) TestInfoMarshalling(c *C) {
	m := &registryMessage{
		Host: "192.168.1.1",
		Port: 1234,
		Uris: []Uri{"foo.vcap.me"},
	}

	s.Register(m)
	marshalled, err := json.Marshal(s)
	c.Check(err, IsNil)
	c.Check(string(marshalled), Equals, "{\"foo.vcap.me\":[\"192.168.1.1:1234\"]}")
}
<|endoftext|>


<commit_before>package goprop

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
)

type Prop struct {
	m map[string]string
}

func NewProp() *Prop {
	return &Prop{}
}

func (p *Prop) Get(key string) string {
	if key == "" {
		return ""
	}
	return p.m[key]
}

func (p *Prop) Read(propFileName string) error {
	file, err := os.OpenFile("prop.properties", os.O_RDONLY, 0660)
	if err != nil {
		log.Println(err)
		return err
	}
	defer file.Close()
	reader := bufio.NewReader(file)
	p.m = make(map[string]string)
	for {
		line, err := reader.ReadString('\n')
		props := strings.SplitN(line, "=", 2)
		p.m[strings.TrimSpace(props[0])] = strings.TrimSpace(props[1])
		if err != nil || io.EOF == err {
			break
		}
	}
	return nil
}
<commit_msg>Update prop.go<commit_after>package goprop

import (
	"bufio"
	"io"
	"log"
	"os"
	"strings"
)

type Prop struct {
	m map[string]string
}

func NewProp() *Prop {
	return &Prop{}
}

func (p *Prop) Get(key string) string {
	if key == "" {
		return ""
	}
	return p.m[key]
}

func (p *Prop) Read(propFileName string) error {
	file, err := os.OpenFile("prop.properties", os.O_RDONLY, 0660)
	if err != nil {
		log.Println(err)
		return err
	}
	defer file.Close()
	reader := bufio.NewReader(file)
	p.m = make(map[string]string)
	for {
		line, err := reader.ReadString('\n')
		props := strings.SplitN(line, "=", 2)
		p.m[strings.TrimSpace(props[0])] = strings.TrimSpace(props[1])
		if err != nil || io.EOF == err {
			break
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// The pubsub package provides a useful stable Redis pubsub connection.
// After opening a connection, it allows you to subscribe to and receieve
// events even in the case of network failures - you don't have to deal
// with that in your code!
//
// Basic usage, prints any messages it gets in the "foobar" channel:
//
// 	client := NewPubsub("127.0.0.1:6379")
// 	defer client.TearDown()
// 	go client.Connect()
//
// 	listener := client.Listener(Channel, "foobar")
// 	for {
// 		fmt.Println(<-listener.Channel)
// 	}
//
// Events are emitted down the client's "Event" channel. If you wanted to
// wait until the client was open (not necessary, but may be useful):
//
// 	client := NewPubsub("127.0.0.1:6379")
// 	go client.Connect()
// 	client.WaitFor(ConnectedEvent)
//
// You can also subscribe to patterns and unsubscribe, of course:
//
// 	listener := client.Listener(Pattern, "foo:*:bar")
// 	doStuff()
// 	listener.Unsubscribe()
//
package pubsub

import (
	"github.com/garyburd/redigo/redis"
	"math"
	"sync"
	"time"
)

// The Client is responsible for maintaining a subscribed redis client,
// reconnecting and resubscribing if it drops.
type Client struct {
	// Event channel you can pull from to know what's happening in the app.
	// By default this is a buffered channel with a size of 1.
	Events chan Event
	// The current state that the client is in.
	State State
	// Lock used to force thread safety on the state.
	stateLock *sync.Mutex
	// Lock used to force thread safety on the pubsub client.
	pubsubLock *sync.Mutex
	// The address we're connected to
	address string
	// The subscription client we're currently using.
	pubsub *redis.PubSubConn
	// How many times we've tried to reconnect.
	retries uint
	// A list of events that we're subscribed to. If the connection closes,
	// we'll reestablish it and resubscribe to events.
	subscribed map[string][]*Listener
	// A queue so that we can asynchronously send "todos" to the main
	// subscribe routine to listen and unsubscribe from events.
	actionQueue chan task
}

// Used to denote the type of listener - channel or pattern.
type ListenerType uint8

const (
	Channel ListenerType = iota
	Pattern
)

// The event is sent down the client's Events channel when something happens!
type Event struct {
	Type   EventType
	Packet interface{}
}

// Events that are sent down the "Events" channel.
type EventType uint8

const (
	ConnectedEvent EventType = iota
	DisconnectedEvent
	MessageEvent
	SubscribeEvent
	UnsubscribeEvent
)

// Tasks we send to the main pubsub thread to subscribe/unsubscribe.
type task struct {
	Action   action
	Listener *Listener
	// If true, the *action* (subscribed/unsubscribed) will be undertaken
	// even if we think we might have already done it.
	Force bool
}

// actions are "things" we can do in tasks
type action uint8

// List of actions we can use in tasks. Internal use
const (
	subscribeAction action = iota
	unsubscribeAction
)

// Client state information
type State uint8

const (
	// Not currently connected to a server.
	DisconnectedState State = iota
	// We were connected, but closed gracefully.
	ClosedState
	// Currently in the process of trying to connect to a server.
	ConnectingState
	// Connected to a server, but not yet linked as a pubsub client.
	ConnectedState
	// Fully connected and listening for events.
	BootedState
)

// Creates a new Radix client and subscribes to it.
func New(address string) *Client {
	client := &Client{
		Events:      make(chan Event),
		State:       DisconnectedState,
		address:     address,
		subscribed:  map[string][]*Listener{},
		actionQueue: make(chan task, 10),
		stateLock:   new(sync.Mutex),
		pubsubLock:  new(sync.Mutex),
	}

	return client
}

// Convenience function to create a new listener for an event.
func (c *Client) Listener(kind ListenerType, event string) *Listener {
	listener := &Listener{
		Type:      kind,
		Event:     event,
		Messages:  make(chan *redis.Message),
		PMessages: make(chan *redis.PMessage),
		Client:    c,
	}
	c.Subscribe(listener)

	return listener
}

// Tears down the client - closes the connection and stops listening for
// connections.
func (c *Client) TearDown() {
	c.setState(ClosedState)
	c.pubsub.Close()
}

// Gets the current client state.
func (c *Client) GetState() State {
	c.stateLock.Lock()
	defer c.stateLock.Unlock()

	return c.State
}

// Sets the client state in a thread-safe manner.
func (c *Client) setState(s State) {
	c.stateLock.Lock()
	c.State = s
	c.stateLock.Unlock()
}

// Blocks until a certain event type is emitted on the client.
func (c *Client) WaitFor(event EventType) Event {
	for {
		if ev := <-c.Events; ev.Type == event {
			return ev
		}
	}
}

// Creates a connection for the client, based on the address. Returns
// an error if there was any issues.
func (c *Client) setupConnection() error {
	cnx, err := redis.Dial("tcp", c.address)
	if err != nil {
		return err
	}

	c.pubsubLock.Lock()
	c.pubsub = &redis.PubSubConn{Conn: cnx}
	c.pubsubLock.Unlock()

	c.retries = 0
	c.setState(ConnectedState)
	c.emit(Event{Type: ConnectedEvent})

	return nil
}

// Should be called when the client is opened. Recieves events, reconnects,
// and dispatches incoming messages.
func (c *Client) boot() {
	c.setState(BootedState)

	halt := make(chan bool, 1)
	ch := make(chan interface{})
	go func() {

		for {
			c.pubsubLock.Lock()
			e := c.pubsub.Receive()
			c.pubsubLock.Unlock()

			if !<-halt {
				ch <- e
			} else {
				return
			}
		}
	}()

	halt <- false
	for c.GetState() == BootedState {

		// Listen for both pubsub receives and new things we have to
		// subscribe to...
		select {
		case e := <-c.actionQueue:
			c.workOnTask(e)
		case e := <-ch:
			halt <- c.dispatchReply(e)
		}
	}
}

// Takes a task, modifies redis and the internal subscribed registery.
// This is done here (called in boot()) since Go's maps are not thread safe.
func (c *Client) workOnTask(t task) {
	event := t.Listener.Event

	// If the action is subscribe...
	if t.Action == subscribeAction {

		// Check to see if it already exists in the subscribed map.
		// If not, then we need to create a new listener list and
		// start listening on the pubsub client. Otherwise just
		// append it.
		if _, exists := c.subscribed[event]; !exists || t.Force {
			c.subscribed[event] = []*Listener{t.Listener}

			switch t.Listener.Type {
			case Channel:
				c.pubsub.Subscribe(event)
			case Pattern:
				c.pubsub.PSubscribe(event)
			}
		} else {
			c.subscribed[event] = append(c.subscribed[event], t.Listener)
		}
	} else {
		// Look for the listener in the subscribed list and,
		// once found, remove it.
		list := c.subscribed[event]
		for i, l := range list {
			if l == t.Listener {
				c.subscribed[event] = append(list[:i], list[i+1:]...)
				break
			}
		}

		// If the list is now empty, we can unsubscribe on Redis and
		// remove it from our registery
		if len(c.subscribed[event]) == 0 || t.Force {
			switch t.Listener.Type {
			case Channel:
				c.pubsub.Unsubscribe(event)
			case Pattern:
				c.pubsub.PUnsubscribe(event)
			}
			delete(c.subscribed, event)
		}
	}
}

// Takes a reply from Redis' recieve and dispatches the
// appropriate message. Returns whether the event loop should be halted.
func (c *Client) dispatchReply(r interface{}) bool {
	switch reply := r.(type) {
	case redis.Message:
		c.dispatchMessage(&reply)
	case redis.PMessage:
		c.dispatchPMessage(&reply)
	case redis.Subscription:
		switch reply.Kind {
		case "subscribe", "psubscribe":
			c.emit(Event{Type: SubscribeEvent, Packet: &reply})
		case "unsubscribe", "punsubscribe":
			c.emit(Event{Type: UnsubscribeEvent, Packet: &reply})
		}
	case error:
		c.handleError(reply)
		return true
	}

	return false
}

// Emits a non-blocking event on the client.
func (c *Client) emit(event Event) {
	select {
	case c.Events <- event:
	default:
	}
}

// Returns the delay time before we should try to reconnect. Increases
// in log time.
func (c *Client) getNextDelay() time.Duration {
	c.retries += 1
	return time.Duration(math.Log(float64(c.retries)) * float64(time.Millisecond))
}

// Returns whether the pubsub client is in a state that you can connect to.
func (c *Client) clientAvailable() bool {
	s := c.GetState()
	return s == ConnectedState || s == BootedState
}

// Tries to reconnect to pubsub, looping until we're
// able to do so successfully.
func (c *Client) Connect() {
	c.setState(ConnectingState)

	for !c.clientAvailable() {
		time.Sleep(c.getNextDelay())
		if c.GetState() == ClosedState {
			return
		}

		c.setupConnection()
	}

	c.resubscribe()
	c.boot()
}

// Resubscribes to all Redis events. Good to do after a disconnection.
func (c *Client) resubscribe() {
	for _, events := range c.subscribed {
		// We only need to subscribe on one event, so that the Redis
		// connection gets registered. We don't care about the others.
		c.actionQueue <- task{Listener: events[0], Action: subscribeAction, Force: true}
	}
}

// Handles error that occur in redis pubsub. Reconnects if necessary.
func (c *Client) handleError(err error) {
	c.emit(Event{Type: DisconnectedEvent, Packet: err})

	if c.GetState() != DisconnectedState {
		c.setState(DisconnectedState)
		c.pubsub.Close()
		c.Connect()
	}
}

// Takes in a Redis message and sends it out to any "listening" channels.
func (c *Client) dispatchMessage(message *redis.Message) {
	if listeners, exists := c.subscribed[message.Channel]; exists {
		for _, l := range listeners {
			l.Messages <- message
		}
	}
}

// Takes in a Redis pmessage and sends it out to any "listening" channels.
func (c *Client) dispatchPMessage(message *redis.PMessage) {
	if listeners, exists := c.subscribed[message.Pattern]; exists {
		for _, l := range listeners {
			l.PMessages <- message
		}
	}
}

// Subscribes to a Redis event. Strings are sent back down the listener
// channel when they come in, and
func (c *Client) Subscribe(listener *Listener) {
	listener.Active = true
	c.actionQueue <- task{Listener: listener, Action: subscribeAction}
}

// Removes the listener from the list of subscribers. If it's the last one
// listening to that Redis event, we unsubscribe entirely.
func (c *Client) Unsubscribe(listener *Listener) {
	listener.Active = false
	c.actionQueue <- task{Listener: listener, Action: unsubscribeAction}
}
<commit_msg>Prevent panics when tearing down a client that was not set up<commit_after>// The pubsub package provides a useful stable Redis pubsub connection.
// After opening a connection, it allows you to subscribe to and receieve
// events even in the case of network failures - you don't have to deal
// with that in your code!
//
// Basic usage, prints any messages it gets in the "foobar" channel:
//
// 	client := NewPubsub("127.0.0.1:6379")
// 	defer client.TearDown()
// 	go client.Connect()
//
// 	listener := client.Listener(Channel, "foobar")
// 	for {
// 		fmt.Println(<-listener.Channel)
// 	}
//
// Events are emitted down the client's "Event" channel. If you wanted to
// wait until the client was open (not necessary, but may be useful):
//
// 	client := NewPubsub("127.0.0.1:6379")
// 	go client.Connect()
// 	client.WaitFor(ConnectedEvent)
//
// You can also subscribe to patterns and unsubscribe, of course:
//
// 	listener := client.Listener(Pattern, "foo:*:bar")
// 	doStuff()
// 	listener.Unsubscribe()
//
package pubsub

import (
	"github.com/garyburd/redigo/redis"
	"math"
	"sync"
	"time"
)

// The Client is responsible for maintaining a subscribed redis client,
// reconnecting and resubscribing if it drops.
type Client struct {
	// Event channel you can pull from to know what's happening in the app.
	// By default this is a buffered channel with a size of 1.
	Events chan Event
	// The current state that the client is in.
	State State
	// Lock used to force thread safety on the state.
	stateLock *sync.Mutex
	// Lock used to force thread safety on the pubsub client.
	pubsubLock *sync.Mutex
	// The address we're connected to
	address string
	// The subscription client we're currently using.
	pubsub *redis.PubSubConn
	// How many times we've tried to reconnect.
	retries uint
	// A list of events that we're subscribed to. If the connection closes,
	// we'll reestablish it and resubscribe to events.
	subscribed map[string][]*Listener
	// A queue so that we can asynchronously send "todos" to the main
	// subscribe routine to listen and unsubscribe from events.
	actionQueue chan task
}

// Used to denote the type of listener - channel or pattern.
type ListenerType uint8

const (
	Channel ListenerType = iota
	Pattern
)

// The event is sent down the client's Events channel when something happens!
type Event struct {
	Type   EventType
	Packet interface{}
}

// Events that are sent down the "Events" channel.
type EventType uint8

const (
	ConnectedEvent EventType = iota
	DisconnectedEvent
	MessageEvent
	SubscribeEvent
	UnsubscribeEvent
)

// Tasks we send to the main pubsub thread to subscribe/unsubscribe.
type task struct {
	Action   action
	Listener *Listener
	// If true, the *action* (subscribed/unsubscribed) will be undertaken
	// even if we think we might have already done it.
	Force bool
}

// actions are "things" we can do in tasks
type action uint8

// List of actions we can use in tasks. Internal use
const (
	subscribeAction action = iota
	unsubscribeAction
)

// Client state information
type State uint8

const (
	// Not currently connected to a server.
	DisconnectedState State = iota
	// We were connected, but closed gracefully.
	ClosedState
	// Currently in the process of trying to connect to a server.
	ConnectingState
	// Connected to a server, but not yet linked as a pubsub client.
	ConnectedState
	// Fully connected and listening for events.
	BootedState
)

// Creates a new Radix client and subscribes to it.
func New(address string) *Client {
	client := &Client{
		Events:      make(chan Event),
		State:       DisconnectedState,
		address:     address,
		subscribed:  map[string][]*Listener{},
		actionQueue: make(chan task, 10),
		stateLock:   new(sync.Mutex),
		pubsubLock:  new(sync.Mutex),
	}

	return client
}

// Convenience function to create a new listener for an event.
func (c *Client) Listener(kind ListenerType, event string) *Listener {
	listener := &Listener{
		Type:      kind,
		Event:     event,
		Messages:  make(chan *redis.Message),
		PMessages: make(chan *redis.PMessage),
		Client:    c,
	}
	c.Subscribe(listener)

	return listener
}

// Tears down the client - closes the connection and stops listening for
// connections.
func (c *Client) TearDown() {
	c.setState(ClosedState)

	if c.pubsub != nil {
		c.pubsub.Close()
	}
}

// Gets the current client state.
func (c *Client) GetState() State {
	c.stateLock.Lock()
	defer c.stateLock.Unlock()

	return c.State
}

// Sets the client state in a thread-safe manner.
func (c *Client) setState(s State) {
	c.stateLock.Lock()
	c.State = s
	c.stateLock.Unlock()
}

// Blocks until a certain event type is emitted on the client.
func (c *Client) WaitFor(event EventType) Event {
	for {
		if ev := <-c.Events; ev.Type == event {
			return ev
		}
	}
}

// Creates a connection for the client, based on the address. Returns
// an error if there was any issues.
func (c *Client) setupConnection() error {
	cnx, err := redis.Dial("tcp", c.address)
	if err != nil {
		return err
	}

	c.pubsubLock.Lock()
	c.pubsub = &redis.PubSubConn{Conn: cnx}
	c.pubsubLock.Unlock()

	c.retries = 0
	c.setState(ConnectedState)
	c.emit(Event{Type: ConnectedEvent})

	return nil
}

// Should be called when the client is opened. Recieves events, reconnects,
// and dispatches incoming messages.
func (c *Client) boot() {
	c.setState(BootedState)

	halt := make(chan bool, 1)
	ch := make(chan interface{})
	go func() {

		for {
			c.pubsubLock.Lock()
			e := c.pubsub.Receive()
			c.pubsubLock.Unlock()

			if !<-halt {
				ch <- e
			} else {
				return
			}
		}
	}()

	halt <- false
	for c.GetState() == BootedState {

		// Listen for both pubsub receives and new things we have to
		// subscribe to...
		select {
		case e := <-c.actionQueue:
			c.workOnTask(e)
		case e := <-ch:
			halt <- c.dispatchReply(e)
		}
	}
}

// Takes a task, modifies redis and the internal subscribed registery.
// This is done here (called in boot()) since Go's maps are not thread safe.
func (c *Client) workOnTask(t task) {
	event := t.Listener.Event

	// If the action is subscribe...
	if t.Action == subscribeAction {

		// Check to see if it already exists in the subscribed map.
		// If not, then we need to create a new listener list and
		// start listening on the pubsub client. Otherwise just
		// append it.
		if _, exists := c.subscribed[event]; !exists || t.Force {
			c.subscribed[event] = []*Listener{t.Listener}

			switch t.Listener.Type {
			case Channel:
				c.pubsub.Subscribe(event)
			case Pattern:
				c.pubsub.PSubscribe(event)
			}
		} else {
			c.subscribed[event] = append(c.subscribed[event], t.Listener)
		}
	} else {
		// Look for the listener in the subscribed list and,
		// once found, remove it.
		list := c.subscribed[event]
		for i, l := range list {
			if l == t.Listener {
				c.subscribed[event] = append(list[:i], list[i+1:]...)
				break
			}
		}

		// If the list is now empty, we can unsubscribe on Redis and
		// remove it from our registery
		if len(c.subscribed[event]) == 0 || t.Force {
			switch t.Listener.Type {
			case Channel:
				c.pubsub.Unsubscribe(event)
			case Pattern:
				c.pubsub.PUnsubscribe(event)
			}
			delete(c.subscribed, event)
		}
	}
}

// Takes a reply from Redis' recieve and dispatches the
// appropriate message. Returns whether the event loop should be halted.
func (c *Client) dispatchReply(r interface{}) bool {
	switch reply := r.(type) {
	case redis.Message:
		c.dispatchMessage(&reply)
	case redis.PMessage:
		c.dispatchPMessage(&reply)
	case redis.Subscription:
		switch reply.Kind {
		case "subscribe", "psubscribe":
			c.emit(Event{Type: SubscribeEvent, Packet: &reply})
		case "unsubscribe", "punsubscribe":
			c.emit(Event{Type: UnsubscribeEvent, Packet: &reply})
		}
	case error:
		c.handleError(reply)
		return true
	}

	return false
}

// Emits a non-blocking event on the client.
func (c *Client) emit(event Event) {
	select {
	case c.Events <- event:
	default:
	}
}

// Returns the delay time before we should try to reconnect. Increases
// in log time.
func (c *Client) getNextDelay() time.Duration {
	c.retries += 1
	return time.Duration(math.Log(float64(c.retries)) * float64(time.Millisecond))
}

// Returns whether the pubsub client is in a state that you can connect to.
func (c *Client) clientAvailable() bool {
	s := c.GetState()
	return s == ConnectedState || s == BootedState
}

// Tries to reconnect to pubsub, looping until we're
// able to do so successfully.
func (c *Client) Connect() {
	c.setState(ConnectingState)

	for !c.clientAvailable() {
		time.Sleep(c.getNextDelay())
		if c.GetState() == ClosedState {
			return
		}

		c.setupConnection()
	}

	c.resubscribe()
	c.boot()
}

// Resubscribes to all Redis events. Good to do after a disconnection.
func (c *Client) resubscribe() {
	for _, events := range c.subscribed {
		// We only need to subscribe on one event, so that the Redis
		// connection gets registered. We don't care about the others.
		c.actionQueue <- task{Listener: events[0], Action: subscribeAction, Force: true}
	}
}

// Handles error that occur in redis pubsub. Reconnects if necessary.
func (c *Client) handleError(err error) {
	c.emit(Event{Type: DisconnectedEvent, Packet: err})

	if c.GetState() != DisconnectedState {
		c.setState(DisconnectedState)
		c.pubsub.Close()
		c.Connect()
	}
}

// Takes in a Redis message and sends it out to any "listening" channels.
func (c *Client) dispatchMessage(message *redis.Message) {
	if listeners, exists := c.subscribed[message.Channel]; exists {
		for _, l := range listeners {
			l.Messages <- message
		}
	}
}

// Takes in a Redis pmessage and sends it out to any "listening" channels.
func (c *Client) dispatchPMessage(message *redis.PMessage) {
	if listeners, exists := c.subscribed[message.Pattern]; exists {
		for _, l := range listeners {
			l.PMessages <- message
		}
	}
}

// Subscribes to a Redis event. Strings are sent back down the listener
// channel when they come in, and
func (c *Client) Subscribe(listener *Listener) {
	listener.Active = true
	c.actionQueue <- task{Listener: listener, Action: subscribeAction}
}

// Removes the listener from the list of subscribers. If it's the last one
// listening to that Redis event, we unsubscribe entirely.
func (c *Client) Unsubscribe(listener *Listener) {
	listener.Active = false
	c.actionQueue <- task{Listener: listener, Action: unsubscribeAction}
}
<|endoftext|>


<commit_before>package sms

import (
	"encoding/xml"
	"fmt"
	"log"
	"os"

	xco "github.com/mndrix/go-xco"
)

// xmppProcess is the piece which interacts with the XMPP network and
// converts those communications into values which the rest of
// sms-over-xmpp can understand.
type xmppProcess struct {
	// where to connect to the XMPP server
	host string
	port int

	// credentials for XMPP auth
	name   string
	secret string

	// channels for communicating with the Gateway process
	gatewayRx chan<- *xco.Message
	gatewayTx <-chan *xco.Message
}

// runXmppComponent creates a goroutine for sending and receiving XMPP
// stanzas.  it returns a channel for monitoring the goroutine's
// health.  if that channel closes, the XMPP process has died.
func (x *xmppProcess) run() <-chan struct{} {
	opts := xco.Options{
		Name:         x.name,
		SharedSecret: x.secret,
		Address:      fmt.Sprintf("%s:%d", x.host, x.port),
		Logger:       log.New(os.Stderr, "", log.LstdFlags),
	}

	healthCh := make(chan struct{})
	go x.loop(opts, healthCh)
	return healthCh
}

func (x *xmppProcess) loop(opts xco.Options, healthCh chan<- struct{}) {
	defer func() { close(healthCh) }()

	c, err := xco.NewComponent(opts)
	if err != nil {
		log.Printf("can't create internal XMPP component: %s", err)
		return
	}

	tx, rx, errx := c.RunAsync()
	for {
		select {
		case st := <-rx:
			switch stanza := st.(type) {
			case *xco.Message:
				log.Printf("Message: %+v", stanza)
				if stanza.Body == "" {
					log.Printf("  ignoring message with empty body")
					break
				}
				go func() { x.gatewayRx <- stanza }()
			case *xco.Iq:
				if stanza.IsDiscoInfo() {
					log.Printf("Disco: %+v", stanza)
					ids, features := x.describeService()
					stanza, err = stanza.DiscoInfoReply(ids, features)
					if err == nil {
						go func() { tx <- stanza }()
					}
				} else {
					log.Printf("Iq: %+v", stanza)
				}
			case *xco.Presence:
				log.Printf("Presence: %+v", stanza)
				switch stanza.Type {
				case "probe":
					stanza, err = x.presenceAvailable(stanza)
					if err == nil {
						go func() { tx <- stanza }()
					}
				case "subscribe", "unsubscribe":
					stanza, err = x.handleSubscription(stanza)
					if err == nil {
						go func() { tx <- stanza }()
					}
				}
			case *xml.StartElement:
				log.Printf("Unknown: %+v", stanza)
			default:
				panic(fmt.Sprintf("Unexpected stanza type: %#v", st))
			}
		case stanza := <-x.gatewayTx:
			go func() { tx <- stanza }()
		case err = <-errx:
		}

		if err != nil {
			break
		}
	}

	log.Printf("lost XMPP connection: %s", err)
}

func (x *xmppProcess) describeService() ([]xco.DiscoIdentity, []xco.DiscoFeature) {
	ids := []xco.DiscoIdentity{
		{
			Category: "gateway",
			Type:     "sms",
			Name:     "SMS over XMPP",
		},
	}
	features := []xco.DiscoFeature{
		{
			Var: "urn:xmpp:receipts",
		},
	}
	return ids, features
}

func (x *xmppProcess) presenceAvailable(p *xco.Presence) (*xco.Presence, error) {
	stanza := &xco.Presence{
		Header: xco.Header{
			From: p.Header.To,
			To:   p.Header.From,
			ID:   NewId(),
		},
	}
	return stanza, nil
}

func (x *xmppProcess) handleSubscription(p *xco.Presence) (*xco.Presence, error) {
	// RFC says to use full JIDs
	p.Header.To.ResourcePart = ""
	p.Header.From.ResourcePart = ""

	stanza := &xco.Presence{
		Header: xco.Header{
			From: p.Header.To,
			To:   p.Header.From,
			ID:   NewId(),
		},
	}
	switch p.Type {
	case "subscribe":
		stanza.Type = "subscribed"
	case "unsubscribe":
		stanza.Type = "unsubscribed"
	}
	return stanza, nil
}
<commit_msg>Reply to "unsubscribe" with "unavailable"<commit_after>package sms

import (
	"encoding/xml"
	"fmt"
	"log"
	"os"

	xco "github.com/mndrix/go-xco"
)

// xmppProcess is the piece which interacts with the XMPP network and
// converts those communications into values which the rest of
// sms-over-xmpp can understand.
type xmppProcess struct {
	// where to connect to the XMPP server
	host string
	port int

	// credentials for XMPP auth
	name   string
	secret string

	// channels for communicating with the Gateway process
	gatewayRx chan<- *xco.Message
	gatewayTx <-chan *xco.Message
}

// runXmppComponent creates a goroutine for sending and receiving XMPP
// stanzas.  it returns a channel for monitoring the goroutine's
// health.  if that channel closes, the XMPP process has died.
func (x *xmppProcess) run() <-chan struct{} {
	opts := xco.Options{
		Name:         x.name,
		SharedSecret: x.secret,
		Address:      fmt.Sprintf("%s:%d", x.host, x.port),
		Logger:       log.New(os.Stderr, "", log.LstdFlags),
	}

	healthCh := make(chan struct{})
	go x.loop(opts, healthCh)
	return healthCh
}

func (x *xmppProcess) loop(opts xco.Options, healthCh chan<- struct{}) {
	defer func() { close(healthCh) }()

	c, err := xco.NewComponent(opts)
	if err != nil {
		log.Printf("can't create internal XMPP component: %s", err)
		return
	}

	tx, rx, errx := c.RunAsync()
	for {
		select {
		case st := <-rx:
			switch stanza := st.(type) {
			case *xco.Message:
				log.Printf("Message: %+v", stanza)
				if stanza.Body == "" {
					log.Printf("  ignoring message with empty body")
					break
				}
				go func() { x.gatewayRx <- stanza }()
			case *xco.Iq:
				if stanza.IsDiscoInfo() {
					log.Printf("Disco: %+v", stanza)
					ids, features := x.describeService()
					stanza, err = stanza.DiscoInfoReply(ids, features)
					if err == nil {
						go func() { tx <- stanza }()
					}
				} else {
					log.Printf("Iq: %+v", stanza)
				}
			case *xco.Presence:
				log.Printf("Presence: %+v", stanza)
				switch stanza.Type {
				case "probe":
					stanza, err = x.presenceAvailable(stanza)
					if err == nil {
						go func() { tx <- stanza }()
					}
				case "subscribe", "unsubscribe":
					stanza, err = x.handleSubscription(stanza)
					if err == nil {
						go func() { tx <- stanza }()
					}
				}
			case *xml.StartElement:
				log.Printf("Unknown: %+v", stanza)
			default:
				panic(fmt.Sprintf("Unexpected stanza type: %#v", st))
			}
		case stanza := <-x.gatewayTx:
			go func() { tx <- stanza }()
		case err = <-errx:
		}

		if err != nil {
			break
		}
	}

	log.Printf("lost XMPP connection: %s", err)
}

func (x *xmppProcess) describeService() ([]xco.DiscoIdentity, []xco.DiscoFeature) {
	ids := []xco.DiscoIdentity{
		{
			Category: "gateway",
			Type:     "sms",
			Name:     "SMS over XMPP",
		},
	}
	features := []xco.DiscoFeature{
		{
			Var: "urn:xmpp:receipts",
		},
	}
	return ids, features
}

func (x *xmppProcess) presenceAvailable(p *xco.Presence) (*xco.Presence, error) {
	stanza := &xco.Presence{
		Header: xco.Header{
			From: p.Header.To,
			To:   p.Header.From,
			ID:   NewId(),
		},
	}
	return stanza, nil
}

func (x *xmppProcess) handleSubscription(p *xco.Presence) (*xco.Presence, error) {
	// RFC says to use full JIDs
	p.Header.To.ResourcePart = ""
	p.Header.From.ResourcePart = ""

	stanza := &xco.Presence{
		Header: xco.Header{
			From: p.Header.To,
			To:   p.Header.From,
			ID:   NewId(),
		},
	}
	switch p.Type {
	case "subscribe":
		stanza.Type = "subscribed"
	case "unsubscribe":
		stanza.Type = "unavailable"
	}
	return stanza, nil
}
<|endoftext|>


<commit_before>package rtime

import (
	"strings"
	"time"
)

// Returns the minute offset within the hour specified by t, in the range [0, 59].
func (t Time) Minute() int {
	return time.Time(t).Minute()
}

// Returns the second offset within the minute specified by t, in the range [0, 59].
func (t Time) Second() int {
	return time.Time(t).Second()
}

// Returns the hour within the day specified by t, in the range [0, 23].
func (t Time) Hour() int {
	return time.Time(t).Hour()
}

// Returns string representation of Hour
func (t Time) HourString() (result string) {
	result = numberInString(t.Hour(), false)
	switch t.Hour() {
	case 1:
		result += " час"
	case 2, 3, 4:
		result += " часа"
	default:
		result += " часов"
	}
	return
}

// Returns string representation of Minute
func (t Time) MinuteString() (result string) {
	result = numberInString(t.Minute(), true)
	switch t.Minute() {
	case 1:
		result += " минута"
	case 2, 3, 4:
		result += " минуты"
	default:
		result += " минут"
	}
	return
}

// Returns string representation of Second
func (t Time) SecondString() (result string) {
	result = numberInString(t.Second(), true)
	switch t.Second() {
	case 1:
		result += " секунда"
	case 2, 3, 4:
		result += " секунды"
	default:
		result += " секунд"
	}
	return
}

func numberInString(num int, female bool) (result string) {
	switch num {
	case 0:
		result = "Ноль"
	case 1:
		if female {
			result = "Одна"
		} else {
			result = "Один"
		}
	case 2:
		if female {
			result = "Две"
		} else {
			result = "Две"
		}
	case 3:
		result = "Три"
	case 4:
		result = "Четыре"
	case 5:
		result = "Пять"
	case 6:
		result = "Шесть"
	case 7:
		result = "Семь"
	case 8:
		result = "Восемь"
	case 9:
		result = "Девять"
	case 10:
		result = "Десять"
	case 11:
		result = "Одиннадцать"
	case 12:
		result = "Двенадцать"
	case 13:
		result = "Тринадцать"
	case 14:
		result = "Четырнадцать"
	case 15:
		result = "Пятнадцать"
	case 16:
		result = "Шестнадцать"
	case 17:
		result = "Семнадцать"
	case 18:
		result = "Восемнадцать"
	case 19:
		result = "Девятнадцать"
	case 20:
		result = "Двадцать"
	case 30:
		result = "Тридцать"
	case 40:
		result = "Сорок"
	case 50:
		result = "Пятьдесят"
	case 60:
		result = "Шестьдесят"
	case 70:
		result = "Семьдесят"
	case 80:
		result = "Восемьдесят"
	case 90:
		result = "Девяносто"
	default:
		lastDigit := num % 10
		result = numberInString(num-lastDigit, female)
		result += " " + strings.ToLower(numberInString(lastDigit, female))
	}
	return
}
<commit_msg>Little fix<commit_after>package rtime

import (
	"strings"
	"time"
)

// Returns the minute offset within the hour specified by t, in the range [0, 59].
func (t Time) Minute() int {
	return time.Time(t).Minute()
}

// Returns the second offset within the minute specified by t, in the range [0, 59].
func (t Time) Second() int {
	return time.Time(t).Second()
}

// Returns the hour within the day specified by t, in the range [0, 23].
func (t Time) Hour() int {
	return time.Time(t).Hour()
}

// Returns string representation of Hour
func (t Time) HourString() (result string) {
	result = numberInString(t.Hour(), false)
	switch t.Hour() {
	case 1:
		result += " час"
	case 2, 3, 4:
		result += " часа"
	default:
		result += " часов"
	}
	return
}

// Returns string representation of Minute
func (t Time) MinuteString() (result string) {
	result = numberInString(t.Minute(), true)
	switch t.Minute() % 10 {
	case 1:
		result += " минута"
	case 2, 3, 4:
		result += " минуты"
	default:
		result += " минут"
	}
	return
}

// Returns string representation of Second
func (t Time) SecondString() (result string) {
	result = numberInString(t.Second(), true)
	switch t.Second() % 10 {
	case 1:
		result += " секунда"
	case 2, 3, 4:
		result += " секунды"
	default:
		result += " секунд"
	}
	return
}

func numberInString(num int, female bool) (result string) {
	switch num {
	case 0:
		result = "Ноль"
	case 1:
		if female {
			result = "Одна"
		} else {
			result = "Один"
		}
	case 2:
		if female {
			result = "Две"
		} else {
			result = "Две"
		}
	case 3:
		result = "Три"
	case 4:
		result = "Четыре"
	case 5:
		result = "Пять"
	case 6:
		result = "Шесть"
	case 7:
		result = "Семь"
	case 8:
		result = "Восемь"
	case 9:
		result = "Девять"
	case 10:
		result = "Десять"
	case 11:
		result = "Одиннадцать"
	case 12:
		result = "Двенадцать"
	case 13:
		result = "Тринадцать"
	case 14:
		result = "Четырнадцать"
	case 15:
		result = "Пятнадцать"
	case 16:
		result = "Шестнадцать"
	case 17:
		result = "Семнадцать"
	case 18:
		result = "Восемнадцать"
	case 19:
		result = "Девятнадцать"
	case 20:
		result = "Двадцать"
	case 30:
		result = "Тридцать"
	case 40:
		result = "Сорок"
	case 50:
		result = "Пятьдесят"
	case 60:
		result = "Шестьдесят"
	case 70:
		result = "Семьдесят"
	case 80:
		result = "Восемьдесят"
	case 90:
		result = "Девяносто"
	default:
		lastDigit := num % 10
		result = numberInString(num-lastDigit, female)
		result += " " + strings.ToLower(numberInString(lastDigit, female))
	}
	return
}
<|endoftext|>


<commit_before>package tlsh

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"io"
	"math"
	"os"
)

const (
	log1_5       = 0.4054651
	log1_3       = 0.26236426
	log1_1       = 0.095310180
	codeSize     = 32
	windowLength = 5
	effBuckets   = 128
	numBuckets   = 256
)

// Tlsh holds hash components
type Tlsh struct {
	checksum byte
	lValue   byte
	q1Ratio  byte
	q2Ratio  byte
	qRatio   byte
	code     [codeSize]byte
}

// New represents type factory for Tlsh
func New(checksum, lValue, q1Ratio, q2Ratio, qRatio byte, code [codeSize]byte) *Tlsh {
	return &Tlsh{
		checksum: checksum,
		lValue:   lValue,
		q1Ratio:  q1Ratio,
		q2Ratio:  q2Ratio,
		qRatio:   qRatio,
		code:     code,
	}
}

// Binary returns the binary representation of the hash
func (t *Tlsh) Binary() []byte {
	return append([]byte{swapByte(t.checksum), swapByte(t.lValue), t.qRatio}, t.code[:]...)
}

// String returns the string representation of the hash`
func (t *Tlsh) String() string {
	return hex.EncodeToString(t.Binary())
}

func quartilePoints(buckets [numBuckets]uint) (q1, q2, q3 uint) {
	var spl, spr uint
	p1 := uint(effBuckets/4 - 1)
	p2 := uint(effBuckets/2 - 1)
	p3 := uint(effBuckets - effBuckets/4 - 1)
	end := uint(effBuckets - 1)

	bucketCopy := make([]uint, effBuckets)
	copy(bucketCopy, buckets[:effBuckets])

	shortCutLeft := make([]uint, effBuckets)
	shortCutRight := make([]uint, effBuckets)

	for l, r := uint(0), end; ; {
		ret := partition(&bucketCopy, l, r)
		if ret > p2 {
			r = ret - 1
			shortCutRight[spr] = ret
			spr++
		} else if ret < p2 {
			l = ret + 1
			shortCutLeft[spl] = ret
			spl++
		} else {
			q2 = bucketCopy[p2]
			break
		}
	}

	shortCutLeft[spl] = p2 - 1
	shortCutRight[spr] = p2 + 1

	for i, l := uint(0), uint(0); i <= spl; i++ {
		r := shortCutLeft[i]
		if r > p1 {
			for {
				ret := partition(&bucketCopy, l, r)
				if ret > p1 {
					r = ret - 1
				} else if ret < p1 {
					l = ret + 1
				} else {
					q1 = bucketCopy[p1]
					break
				}
			}
			break
		} else if r < p1 {
			l = r
		} else {
			q1 = bucketCopy[p1]
			break
		}
	}

	for i, r := uint(0), end; i <= spr; i++ {
		l := shortCutRight[i]
		if l < p3 {
			for {
				ret := partition(&bucketCopy, l, r)
				if ret > p3 {
					r = ret - 1
				} else if ret < p3 {
					l = ret + 1
				} else {
					q3 = bucketCopy[p3]
					break
				}
			}
			break
		} else if l > p3 {
			r = l
		} else {
			q3 = bucketCopy[p3]
			break
		}
	}

	return q1, q2, q3
}

func partition(pbuf *[]uint, left, right uint) uint {
	if left == right {
		return left
	}

	rbuf := (*pbuf)

	if left+1 == right {
		if rbuf[left] > rbuf[right] {
			rbuf[right], rbuf[left] = rbuf[left], rbuf[right]
		}
		return left
	}

	ret := left
	pivot := (left + right) >> 1
	val := rbuf[pivot]

	rbuf[pivot] = rbuf[right]
	rbuf[right] = val

	for i := left; i < right; i++ {
		if rbuf[i] < val {
			rbuf[i], rbuf[ret] = rbuf[ret], rbuf[i]
			ret++
		}
	}

	rbuf[right] = rbuf[ret]
	rbuf[ret] = val

	return ret
}

func lValue(length int) byte {
	var l byte

	if length <= 656 {
		l = byte(math.Floor(math.Log(float64(length)) / log1_5))
	} else if length <= 3199 {
		l = byte(math.Floor(math.Log(float64(length))/log1_3 - 8.72777))
	} else {
		l = byte(math.Floor(math.Log(float64(length))/log1_1 - 62.5472))
	}

	return l % 255
}

func swapByte(in byte) byte {
	var out byte

	out = ((in & 0xF0) >> 4) & 0x0F
	out |= ((in & 0x0F) << 4) & 0xF0

	return out
}

func bucketsBinaryRepresentation(buckets [numBuckets]uint, q1, q2, q3 uint) [codeSize]byte {
	var biHash [codeSize]byte

	for i := 0; i < codeSize; i++ {
		var h byte
		for j := 0; j < 4; j++ {
			k := buckets[4*i+j]
			if q3 < k {
				h += 3 << (byte(j) * 2)
			} else if q2 < k {
				h += 2 << (byte(j) * 2)
			} else if q1 < k {
				h += 1 << (byte(j) * 2)
			}
		}
		// Prepend the new h to the hash
		biHash[(codeSize-1)-i] = h
	}
	return biHash
}

func reverse(s [5]byte) [5]byte {
	for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func fillBuckets(r FuzzyReader) ([numBuckets]uint, byte, int, error) {
	buckets := [numBuckets]uint{}
	chunkSlice := make([]byte, windowLength)
	chunk := [windowLength]byte{}
	salt := [6]byte{2, 3, 5, 7, 11, 13}
	fileSize := 0
	checksum := byte(0)

	n, err := r.Read(chunkSlice)
	if err != nil {
		return [numBuckets]uint{}, 0, 0, err
	}
	copy(chunk[:], chunkSlice[0:5])
	chunk = reverse(chunk)
	fileSize += n

	chunk3 := &[3]byte{}

	for {
		chunk3[0] = chunk[0]
		chunk3[1] = chunk[1]
		chunk3[2] = checksum
		checksum = pearsonHash(0, chunk3)

		chunk3[2] = chunk[2]
		buckets[pearsonHash(salt[0], chunk3)]++

		chunk3[2] = chunk[3]
		buckets[pearsonHash(salt[1], chunk3)]++

		chunk3[1] = chunk[2]
		buckets[pearsonHash(salt[2], chunk3)]++

		chunk3[2] = chunk[4]
		buckets[pearsonHash(salt[3], chunk3)]++

		chunk3[1] = chunk[1]
		buckets[pearsonHash(salt[4], chunk3)]++

		chunk3[1] = chunk[3]
		buckets[pearsonHash(salt[5], chunk3)]++

		copy(chunk[1:], chunk[0:4])
		chunk[0], err = r.ReadByte()
		if err != nil {
			if err != io.EOF {
				return [numBuckets]uint{}, 0, 0, err
			}
			break
		}
		fileSize++
	}
	return buckets, checksum, fileSize, nil
}

// hashCalculate calculate TLSH
func hashCalculate(r FuzzyReader) (*Tlsh, error) {
	buckets, checksum, fileSize, err := fillBuckets(r)
	if err != nil {
		return &Tlsh{}, err
	}

	q1, q2, q3 := quartilePoints(buckets)
	q1Ratio := byte(float32(q1)*100/float32(q3)) % 16
	q2Ratio := byte(float32(q2)*100/float32(q3)) % 16
	qRatio := ((q1Ratio & 0xF) << 4) | (q2Ratio & 0xF)

	biHash := bucketsBinaryRepresentation(buckets, q1, q2, q3)

	return New(checksum, lValue(fileSize), q1Ratio, q2Ratio, qRatio, biHash), nil
}

// FuzzyReader interface
type FuzzyReader interface {
	Read([]byte) (int, error)
	ReadByte() (byte, error)
}

//HashReader calculates the TLSH for the input reader
func HashReader(r FuzzyReader) (string, error) {
	tlsh, err := hashCalculate(r)
	if err != nil {
		return "", err
	}

	return tlsh.String(), err
}

//HashBytes calculates the TLSH for the input byte slice
func HashBytes(blob []byte) (hash string, err error) {
	r := bytes.NewReader(blob)
	return HashReader(r)
}

//HashFilename calculates the TLSH for the input file
func HashFilename(filename string) (tlsh *Tlsh, err error) {
	f, err := os.Open(filename)
	defer f.Close()
	if err != nil {
		return &Tlsh{}, err
	}

	r := bufio.NewReader(f)
	tlsh, err = hashCalculate(r)
	if err != nil {
		return &Tlsh{}, err
	}
	return tlsh, nil
}

// Diff current hash with other hash
func (t *Tlsh) Diff(t2 *Tlsh) int {
	return diffTotal(t, t2, true)
}

// DiffFilenames calculate distance between two files
func DiffFilenames(filenameA, filenameB string) (int, error) {
	f, err := os.Open(filenameA)
	defer f.Close()
	if err != nil {
		return -1, err
	}
	r := bufio.NewReader(f)
	tlshA, err := hashCalculate(r)

	f, err = os.Open(filenameB)
	defer f.Close()
	if err != nil {
		return -1, err
	}
	r = bufio.NewReader(f)
	tlshB, err := hashCalculate(r)
	if err != nil {
		return -1, err
	}

	return tlshA.Diff(tlshB), nil
}
<commit_msg>Handle error correctly<commit_after>package tlsh

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"io"
	"math"
	"os"
)

const (
	log1_5       = 0.4054651
	log1_3       = 0.26236426
	log1_1       = 0.095310180
	codeSize     = 32
	windowLength = 5
	effBuckets   = 128
	numBuckets   = 256
)

// Tlsh holds hash components
type Tlsh struct {
	checksum byte
	lValue   byte
	q1Ratio  byte
	q2Ratio  byte
	qRatio   byte
	code     [codeSize]byte
}

// New represents type factory for Tlsh
func New(checksum, lValue, q1Ratio, q2Ratio, qRatio byte, code [codeSize]byte) *Tlsh {
	return &Tlsh{
		checksum: checksum,
		lValue:   lValue,
		q1Ratio:  q1Ratio,
		q2Ratio:  q2Ratio,
		qRatio:   qRatio,
		code:     code,
	}
}

// Binary returns the binary representation of the hash
func (t *Tlsh) Binary() []byte {
	return append([]byte{swapByte(t.checksum), swapByte(t.lValue), t.qRatio}, t.code[:]...)
}

// String returns the string representation of the hash`
func (t *Tlsh) String() string {
	return hex.EncodeToString(t.Binary())
}

func quartilePoints(buckets [numBuckets]uint) (q1, q2, q3 uint) {
	var spl, spr uint
	p1 := uint(effBuckets/4 - 1)
	p2 := uint(effBuckets/2 - 1)
	p3 := uint(effBuckets - effBuckets/4 - 1)
	end := uint(effBuckets - 1)

	bucketCopy := make([]uint, effBuckets)
	copy(bucketCopy, buckets[:effBuckets])

	shortCutLeft := make([]uint, effBuckets)
	shortCutRight := make([]uint, effBuckets)

	for l, r := uint(0), end; ; {
		ret := partition(&bucketCopy, l, r)
		if ret > p2 {
			r = ret - 1
			shortCutRight[spr] = ret
			spr++
		} else if ret < p2 {
			l = ret + 1
			shortCutLeft[spl] = ret
			spl++
		} else {
			q2 = bucketCopy[p2]
			break
		}
	}

	shortCutLeft[spl] = p2 - 1
	shortCutRight[spr] = p2 + 1

	for i, l := uint(0), uint(0); i <= spl; i++ {
		r := shortCutLeft[i]
		if r > p1 {
			for {
				ret := partition(&bucketCopy, l, r)
				if ret > p1 {
					r = ret - 1
				} else if ret < p1 {
					l = ret + 1
				} else {
					q1 = bucketCopy[p1]
					break
				}
			}
			break
		} else if r < p1 {
			l = r
		} else {
			q1 = bucketCopy[p1]
			break
		}
	}

	for i, r := uint(0), end; i <= spr; i++ {
		l := shortCutRight[i]
		if l < p3 {
			for {
				ret := partition(&bucketCopy, l, r)
				if ret > p3 {
					r = ret - 1
				} else if ret < p3 {
					l = ret + 1
				} else {
					q3 = bucketCopy[p3]
					break
				}
			}
			break
		} else if l > p3 {
			r = l
		} else {
			q3 = bucketCopy[p3]
			break
		}
	}

	return q1, q2, q3
}

func partition(pbuf *[]uint, left, right uint) uint {
	if left == right {
		return left
	}

	rbuf := (*pbuf)

	if left+1 == right {
		if rbuf[left] > rbuf[right] {
			rbuf[right], rbuf[left] = rbuf[left], rbuf[right]
		}
		return left
	}

	ret := left
	pivot := (left + right) >> 1
	val := rbuf[pivot]

	rbuf[pivot] = rbuf[right]
	rbuf[right] = val

	for i := left; i < right; i++ {
		if rbuf[i] < val {
			rbuf[i], rbuf[ret] = rbuf[ret], rbuf[i]
			ret++
		}
	}

	rbuf[right] = rbuf[ret]
	rbuf[ret] = val

	return ret
}

func lValue(length int) byte {
	var l byte

	if length <= 656 {
		l = byte(math.Floor(math.Log(float64(length)) / log1_5))
	} else if length <= 3199 {
		l = byte(math.Floor(math.Log(float64(length))/log1_3 - 8.72777))
	} else {
		l = byte(math.Floor(math.Log(float64(length))/log1_1 - 62.5472))
	}

	return l % 255
}

func swapByte(in byte) byte {
	var out byte

	out = ((in & 0xF0) >> 4) & 0x0F
	out |= ((in & 0x0F) << 4) & 0xF0

	return out
}

func bucketsBinaryRepresentation(buckets [numBuckets]uint, q1, q2, q3 uint) [codeSize]byte {
	var biHash [codeSize]byte

	for i := 0; i < codeSize; i++ {
		var h byte
		for j := 0; j < 4; j++ {
			k := buckets[4*i+j]
			if q3 < k {
				h += 3 << (byte(j) * 2)
			} else if q2 < k {
				h += 2 << (byte(j) * 2)
			} else if q1 < k {
				h += 1 << (byte(j) * 2)
			}
		}
		// Prepend the new h to the hash
		biHash[(codeSize-1)-i] = h
	}
	return biHash
}

func reverse(s [5]byte) [5]byte {
	for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func fillBuckets(r FuzzyReader) ([numBuckets]uint, byte, int, error) {
	buckets := [numBuckets]uint{}
	chunkSlice := make([]byte, windowLength)
	chunk := [windowLength]byte{}
	salt := [6]byte{2, 3, 5, 7, 11, 13}
	fileSize := 0
	checksum := byte(0)

	n, err := r.Read(chunkSlice)
	if err != nil {
		return [numBuckets]uint{}, 0, 0, err
	}
	copy(chunk[:], chunkSlice[0:5])
	chunk = reverse(chunk)
	fileSize += n

	chunk3 := &[3]byte{}

	for {
		chunk3[0] = chunk[0]
		chunk3[1] = chunk[1]
		chunk3[2] = checksum
		checksum = pearsonHash(0, chunk3)

		chunk3[2] = chunk[2]
		buckets[pearsonHash(salt[0], chunk3)]++

		chunk3[2] = chunk[3]
		buckets[pearsonHash(salt[1], chunk3)]++

		chunk3[1] = chunk[2]
		buckets[pearsonHash(salt[2], chunk3)]++

		chunk3[2] = chunk[4]
		buckets[pearsonHash(salt[3], chunk3)]++

		chunk3[1] = chunk[1]
		buckets[pearsonHash(salt[4], chunk3)]++

		chunk3[1] = chunk[3]
		buckets[pearsonHash(salt[5], chunk3)]++

		copy(chunk[1:], chunk[0:4])
		chunk[0], err = r.ReadByte()
		if err != nil {
			if err != io.EOF {
				return [numBuckets]uint{}, 0, 0, err
			}
			break
		}
		fileSize++
	}
	return buckets, checksum, fileSize, nil
}

// hashCalculate calculate TLSH
func hashCalculate(r FuzzyReader) (*Tlsh, error) {
	buckets, checksum, fileSize, err := fillBuckets(r)
	if err != nil {
		return &Tlsh{}, err
	}

	q1, q2, q3 := quartilePoints(buckets)
	q1Ratio := byte(float32(q1)*100/float32(q3)) % 16
	q2Ratio := byte(float32(q2)*100/float32(q3)) % 16
	qRatio := ((q1Ratio & 0xF) << 4) | (q2Ratio & 0xF)

	biHash := bucketsBinaryRepresentation(buckets, q1, q2, q3)

	return New(checksum, lValue(fileSize), q1Ratio, q2Ratio, qRatio, biHash), nil
}

// FuzzyReader interface
type FuzzyReader interface {
	Read([]byte) (int, error)
	ReadByte() (byte, error)
}

//HashReader calculates the TLSH for the input reader
func HashReader(r FuzzyReader) (string, error) {
	tlsh, err := hashCalculate(r)
	if err != nil {
		return "", err
	}

	return tlsh.String(), err
}

//HashBytes calculates the TLSH for the input byte slice
func HashBytes(blob []byte) (hash string, err error) {
	r := bytes.NewReader(blob)
	return HashReader(r)
}

//HashFilename calculates the TLSH for the input file
func HashFilename(filename string) (tlsh *Tlsh, err error) {
	f, err := os.Open(filename)
	defer f.Close()
	if err != nil {
		return &Tlsh{}, err
	}

	r := bufio.NewReader(f)
	tlsh, err = hashCalculate(r)
	if err != nil {
		return &Tlsh{}, err
	}
	return tlsh, nil
}

// Diff current hash with other hash
func (t *Tlsh) Diff(t2 *Tlsh) int {
	return diffTotal(t, t2, true)
}

// DiffFilenames calculate distance between two files
func DiffFilenames(filenameA, filenameB string) (int, error) {
	f, err := os.Open(filenameA)
	defer f.Close()
	if err != nil {
		return -1, err
	}
	r := bufio.NewReader(f)
	tlshA, err := hashCalculate(r)
	if err != nil {
		return -1, err
	}

	f, err = os.Open(filenameB)
	defer f.Close()
	if err != nil {
		return -1, err
	}
	r = bufio.NewReader(f)
	tlshB, err := hashCalculate(r)
	if err != nil {
		return -1, err
	}

	return tlshA.Diff(tlshB), nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// globals are eeeeevil
var paused bool
var signalLock = sync.RWMutex{}
var maintModeLock = sync.RWMutex{}

// we wrap access to `paused` in a RLock so that if we're in the middle of
// marking services for maintenance we don't get stale reads
func inMaintenanceMode() bool {
	maintModeLock.RLock()
	defer maintModeLock.RUnlock()
	return paused
}

func toggleMaintenanceMode(config *Config) {
	maintModeLock.RLock()
	signalLock.Lock()
	defer signalLock.Unlock()
	defer maintModeLock.RUnlock()
	paused = !paused
	if paused {
		forAllServices(config, func(service *ServiceConfig) {
			log.Printf("Marking for maintenance: %s\n", service.Name)
			service.MarkForMaintenance()
		})
	}
}

func terminate(config *Config) {
	signalLock.Lock()
	defer signalLock.Unlock()
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	// Run and wait for preStop command to exit
	run(config.preStopCmd)

	cmd := config.Command
	if cmd == nil || cmd.Process == nil {
		// Not managing the process, so don't do anything
		return
	}
	if config.StopTimeout > 0 {
		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
			log.Printf("Error sending SIGTERM to application: %s\n", err)
		} else {
			time.AfterFunc(time.Duration(config.StopTimeout)*time.Second, func() {
				log.Printf("Killing Process %#v\n", cmd.Process)
				cmd.Process.Kill()
			})
			return
		}
	}
	log.Printf("Killing Process %#v\n", config.Command.Process)
	cmd.Process.Kill()
}

func reloadConfig(config *Config) *Config {
	signalLock.Lock()
	defer signalLock.Unlock()
	log.Printf("Reloading configuration.\n")
	newConfig, err := loadConfig()
	if err != nil {
		log.Printf("Could not reload config: %v\n", err)
		return nil
	}
	// stop advertising the existing services so that we can
	// make sure we update them if ports, etc. change.
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	signal.Reset()
	handleSignals(newConfig)
	handlePolling(newConfig)

	return newConfig // return for debuggability
}

func stopPolling(config *Config) {
	for _, quit := range config.QuitChannels {
		quit <- true
	}
}

type serviceFunc func(service *ServiceConfig)

func forAllServices(config *Config, fn serviceFunc) {
	for _, service := range config.Services {
		fn(service)
	}
}

// Listen for and capture signals used for orchestration
func handleSignals(config *Config) {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGUSR1, syscall.SIGTERM, syscall.SIGHUP)
	go func() {
		for signal := range sig {
			switch signal {
			case syscall.SIGUSR1:
				toggleMaintenanceMode(config)
			case syscall.SIGTERM:
				terminate(config)
			case syscall.SIGHUP:
				reloadConfig(config)
			}
		}
	}()
}

// on SIGCHLD send wait4() (ref http://linux.die.net/man/2/waitpid)
// to clean up any potential zombies
func reapChildren() {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGCHLD)
	go func() {
		<-sig // block here
		for {
			// only 1 SIGCHLD can be handled at a time from the channel,
			// so we need to allow for the possibility that multiple child
			// processes have terminated while one is already being reaped.
			var wstatus syscall.WaitStatus
			if _, err := syscall.Wait4(-1, &wstatus,
				syscall.WNOHANG|syscall.WUNTRACED|syscall.WCONTINUED,
				nil); err == syscall.EINTR {
				continue
			}
			break
		}
	}()
}
<commit_msg>Prevent exit from reapChildren goroutine<commit_after>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// globals are eeeeevil
var paused bool
var signalLock = sync.RWMutex{}
var maintModeLock = sync.RWMutex{}

// we wrap access to `paused` in a RLock so that if we're in the middle of
// marking services for maintenance we don't get stale reads
func inMaintenanceMode() bool {
	maintModeLock.RLock()
	defer maintModeLock.RUnlock()
	return paused
}

func toggleMaintenanceMode(config *Config) {
	maintModeLock.RLock()
	signalLock.Lock()
	defer signalLock.Unlock()
	defer maintModeLock.RUnlock()
	paused = !paused
	if paused {
		forAllServices(config, func(service *ServiceConfig) {
			log.Printf("Marking for maintenance: %s\n", service.Name)
			service.MarkForMaintenance()
		})
	}
}

func terminate(config *Config) {
	signalLock.Lock()
	defer signalLock.Unlock()
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	// Run and wait for preStop command to exit
	run(config.preStopCmd)

	cmd := config.Command
	if cmd == nil || cmd.Process == nil {
		// Not managing the process, so don't do anything
		return
	}
	if config.StopTimeout > 0 {
		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
			log.Printf("Error sending SIGTERM to application: %s\n", err)
		} else {
			time.AfterFunc(time.Duration(config.StopTimeout)*time.Second, func() {
				log.Printf("Killing Process %#v\n", cmd.Process)
				cmd.Process.Kill()
			})
			return
		}
	}
	log.Printf("Killing Process %#v\n", config.Command.Process)
	cmd.Process.Kill()
}

func reloadConfig(config *Config) *Config {
	signalLock.Lock()
	defer signalLock.Unlock()
	log.Printf("Reloading configuration.\n")
	newConfig, err := loadConfig()
	if err != nil {
		log.Printf("Could not reload config: %v\n", err)
		return nil
	}
	// stop advertising the existing services so that we can
	// make sure we update them if ports, etc. change.
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	signal.Reset()
	handleSignals(newConfig)
	handlePolling(newConfig)

	return newConfig // return for debuggability
}

func stopPolling(config *Config) {
	for _, quit := range config.QuitChannels {
		quit <- true
	}
}

type serviceFunc func(service *ServiceConfig)

func forAllServices(config *Config, fn serviceFunc) {
	for _, service := range config.Services {
		fn(service)
	}
}

// Listen for and capture signals used for orchestration
func handleSignals(config *Config) {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGUSR1, syscall.SIGTERM, syscall.SIGHUP)
	go func() {
		for signal := range sig {
			switch signal {
			case syscall.SIGUSR1:
				toggleMaintenanceMode(config)
			case syscall.SIGTERM:
				terminate(config)
			case syscall.SIGHUP:
				reloadConfig(config)
			}
		}
	}()
}

// on SIGCHLD send wait4() (ref http://linux.die.net/man/2/waitpid)
// to clean up any potential zombies
func reapChildren() {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGCHLD)
	go func() {
		// wait for signals on the channel until it closes
		for _ = range sig {
			for {
				// only 1 SIGCHLD can be handled at a time from the channel,
				// so we need to allow for the possibility that multiple child
				// processes have terminated while one is already being reaped.
				var wstatus syscall.WaitStatus
				if _, err := syscall.Wait4(-1, &wstatus,
					syscall.WNOHANG|syscall.WUNTRACED|syscall.WCONTINUED,
					nil); err == syscall.EINTR {
					continue
				}
				// return to the outer loop and wait for another signal
				break
			}
		}
	}()
}
<|endoftext|>


<commit_before>package random

import (
	"math/rand"
	"time"
)

const (
	UPPER_LETTER     = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	UPPER_LETTER_LEN = len(UPPER_LETTER)
	LOWER_LETTER     = "abcdefghijklmnopqrstuvwxyz"
	LOWER_LETTER_LEN = len(LOWER_LETTER)
	NUMBER           = "0123456789"
	LETTER           = UPPER_LETTER + LOWER_LETTER
	LETTER_LEN       = len(LETTER)
	CHAR             = NUMBER + UPPER_LETTER + LOWER_LETTER
	CHAR_LEN         = len(CHAR)
)

func Int(min, max int) int {
	source := rand.NewSource(time.Now().UnixNano())
	return min + rand.New(source).Intn(max-min)
}

func Char() string {
	rint := Int(0, CHAR_LEN)
	return CHAR[rint : rint+1]
}

func String(length int) string {
	return strs(Char, length)
}

func Letter() string {
	rint := Int(0, LETTER_LEN)
	return LETTER[rint : rint+1]
}

func Letters(length int) string {
	return strs(Letter, length)
}

func strs(fn func() string, length int) string {
	if 0 >= length {
		panic("lenght must be greater than zero")
	}

	str := ""
	for i := 0; i < length; i++ {
		str = str + fn()
	}

	return str
}
<commit_msg>improve the String method performance<commit_after>package random

import (
	"crypto/md5"
	"fmt"
	"math/rand"
	"strconv"
	"time"
)

const (
	UPPER_LETTER     = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	UPPER_LETTER_LEN = len(UPPER_LETTER)
	LOWER_LETTER     = "abcdefghijklmnopqrstuvwxyz"
	LOWER_LETTER_LEN = len(LOWER_LETTER)
	NUMBER           = "0123456789"
	LETTER           = UPPER_LETTER + LOWER_LETTER
	LETTER_LEN       = len(LETTER)
	CHAR             = NUMBER + UPPER_LETTER + LOWER_LETTER
	CHAR_LEN         = len(CHAR)
)

func Md5(str string) string {
	h := md5.New()
	h.Write([]byte(str))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func Int(min, max int) int {
	source := rand.NewSource(time.Now().UnixNano())
	return min + rand.New(source).Intn(max-min)
}

func Char() string {
	rint := Int(0, CHAR_LEN)
	return CHAR[rint : rint+1]
}

func String(length int) string {
	s := ""
	if length > 32 {
		var ts int = length / 32
		for i := 0; i < ts; i++ {
			s += String(32)
		}

		if tr := length % 32; tr > 0 {
			s += String(tr)
		}

		return s
	}

	s = Md5(strconv.Itoa(time.Now().Nanosecond()))
	if length == 32 {
		return s
	}

	return s[0:length]

}

func Letter() string {
	rint := Int(0, LETTER_LEN)
	return LETTER[rint : rint+1]
}

func Letters(length int) string {
	return strs(Letter, length)
}

func strs(fn func() string, length int) string {
	if 0 >= length {
		panic("lenght must be greater than zero")
	}

	str := ""
	for i := 0; i < length; i++ {
		str = str + fn()
	}

	return str
}
<|endoftext|>


<commit_before>package rdigo

import (
	"encoding/json"
	"github.com/mrjones/oauth"
	"io/ioutil"
	"log"
)

const baseUrl = "http://api.rdio.com/1/"

type Rdio struct {
	Consumer    *oauth.Consumer
	AccessToken *oauth.AccessToken
}

func NewClient(consumerKey, consumerSecret string) Rdio {
	c := oauth.NewConsumer(consumerKey, consumerSecret,
		oauth.ServiceProvider{
			RequestTokenUrl:   "http://api.rdio.com/oauth/request_token",
			AuthorizeTokenUrl: "https://www.rdio.com/oauth/authorize",
			AccessTokenUrl:    "http://api.rdio.com/oauth/access_token",
		})
	return Rdio{Consumer: c}
}

func (r *Rdio) BeginAuthentication(callbackUrl string) (*oauth.RequestToken, string, error) {
	return r.Consumer.GetRequestTokenAndUrl(callbackUrl)
}

func (r *Rdio) CompleteAuthentication(rToken *oauth.RequestToken, verifier string) error {
	token, err := r.Consumer.AuthorizeToken(rToken, verifier)
	r.AccessToken = token
	return err
}

// Authenticated only for now
func (r *Rdio) Call(method string, query map[string]string) (interface{}, error) {
	query["method"] = method
	resp, err := r.Consumer.Post(baseUrl, query, r.AccessToken)
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("Rdio Call Error:", err.Error())
		return nil, err
	}
	var rdioResp interface{}
	err = json.Unmarshal(body, &rdioResp)
	if err != nil {
		log.Println("Rdio Call Error:", err.Error())
		return nil, err
	}
	return rdioResp, err
}
<commit_msg>Modified CompleteAuthentication to take in token and secret as strings so they can be passed directly from cookies.<commit_after>package rdigo

import (
	"encoding/json"
	"github.com/mrjones/oauth"
	"io/ioutil"
	"log"
)

const baseUrl = "http://api.rdio.com/1/"

type Rdio struct {
	Consumer    *oauth.Consumer
	AccessToken *oauth.AccessToken
}

func NewClient(consumerKey, consumerSecret string) Rdio {
	c := oauth.NewConsumer(consumerKey, consumerSecret,
		oauth.ServiceProvider{
			RequestTokenUrl:   "http://api.rdio.com/oauth/request_token",
			AuthorizeTokenUrl: "https://www.rdio.com/oauth/authorize",
			AccessTokenUrl:    "http://api.rdio.com/oauth/access_token",
		})
	return Rdio{Consumer: c}
}

func (r *Rdio) BeginAuthentication(callbackUrl string) (*oauth.RequestToken, string, error) {
	return r.Consumer.GetRequestTokenAndUrl(callbackUrl)
}

func (r *Rdio) CompleteAuthentication(requestToken, requestTokenSecret, verifier string) error {
	rToken := oauth.RequestToken{Token: requestToken, Secret: requestTokenSecret}
	token, err := r.Consumer.AuthorizeToken(&rToken, verifier)
	r.AccessToken = token
	return err
}

// Authenticated only for now
func (r *Rdio) Call(method string, query map[string]string) (interface{}, error) {
	query["method"] = method
	resp, err := r.Consumer.Post(baseUrl, query, r.AccessToken)
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("Rdio Call Error:", err.Error())
		return nil, err
	}
	var rdioResp interface{}
	err = json.Unmarshal(body, &rdioResp)
	if err != nil {
		log.Println("Rdio Call Error:", err.Error())
		return nil, err
	}
	return rdioResp, err
}
<|endoftext|>