<commit_before>package jujutest

import (
	"fmt"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/version"
	"time"
)

// TestStartStop is similar to Tests.TestStartStop except
// that it does not assume a pristine environment.
func (t *LiveTests) TestStartStop(c *C) {
	insts, err := t.Env.Instances(nil)
	c.Assert(err, IsNil)
	c.Check(insts, HasLen, 0)

	inst, err := t.Env.StartInstance(0, InvalidStateInfo, nil)
	c.Assert(err, IsNil)
	c.Assert(inst, NotNil)
	id0 := inst.Id()

	insts, err = t.Env.Instances([]string{id0, id0})
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 2)
	c.Assert(insts[0].Id(), Equals, id0)
	c.Assert(insts[1].Id(), Equals, id0)

	insts, err = t.Env.AllInstances()
	c.Assert(err, IsNil)
	// differs from the check above because AllInstances returns
	// a set (without duplicates) containing only one instance.
	c.Assert(insts, HasLen, 1)
	c.Assert(insts[0].Id(), Equals, id0)

	dns, err := inst.WaitDNSName()
	c.Assert(err, IsNil)
	c.Assert(dns, Not(Equals), "")

	insts, err = t.Env.Instances([]string{id0, ""})
	c.Assert(err, Equals, environs.ErrPartialInstances)
	c.Assert(insts, HasLen, 2)
	c.Check(insts[0].Id(), Equals, id0)
	c.Check(insts[1], IsNil)

	err = t.Env.StopInstances([]environs.Instance{inst})
	c.Assert(err, IsNil)

	// Stopping may not be noticed at first due to eventual
	// consistency. Repeat a few times to ensure we get the error.
	for i := 0; i < 20; i++ {
		insts, err = t.Env.Instances([]string{id0})
		if err != nil {
			break
		}
		time.Sleep(0.25e9)
	}
	c.Assert(err, Equals, environs.ErrNoInstances)
	c.Assert(insts, HasLen, 0)
}

func (t *LiveTests) TestPorts(c *C) {
	inst1, err := t.Env.StartInstance(1, InvalidStateInfo, nil)
	c.Assert(err, IsNil)
	c.Assert(inst1, NotNil)
	defer t.Env.StopInstances([]environs.Instance{inst1})

	ports, err := inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, HasLen, 0)

	inst2, err := t.Env.StartInstance(2, InvalidStateInfo, nil)
	c.Assert(err, IsNil)
	c.Assert(inst2, NotNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, HasLen, 0)
	defer t.Env.StopInstances([]environs.Instance{inst2})

	// Open some ports and check they're there.
	err = inst1.OpenPorts(1, []state.Port{{"udp", 67}, {"tcp", 45}})
	c.Assert(err, IsNil)
	ports, err = inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"udp", 67}})
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, HasLen, 0)

	// Check there's no crosstalk to another machine
	err = inst2.OpenPorts(2, []state.Port{{"tcp", 89}, {"tcp", 45}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"tcp", 89}})
	ports, err = inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"udp", 67}})

	// Check that opening the same port again is ok.
	err = inst2.OpenPorts(2, []state.Port{{"tcp", 45}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"tcp", 89}})

	// Check that opening the same port again and another port is ok.
	err = inst2.OpenPorts(2, []state.Port{{"tcp", 45}, {"tcp", 99}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"tcp", 89}, {"tcp", 99}})

	// Check that we can close ports and that there's no crosstalk.
	err = inst2.ClosePorts(2, []state.Port{{"tcp", 45}, {"tcp", 99}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 89}})
	ports, err = inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"udp", 67}})

	// Check that we can close multiple ports.
	err = inst1.ClosePorts(1, []state.Port{{"tcp", 45}, {"udp", 67}})
	c.Assert(err, IsNil)
	ports, err = inst1.Ports(1)
	c.Assert(ports, HasLen, 0)

	// Check that we can close ports that aren't there.
	err = inst2.ClosePorts(2, []state.Port{{"tcp", 111}, {"udp", 222}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 89}})
}

func (t *LiveTests) TestBootstrap(c *C) {
	t.BootstrapOnce(c)

	// Wait for a while to let eventual consistency catch up, hopefully.
	time.Sleep(t.ConsistencyDelay)
	err := t.Env.Bootstrap(false)
	c.Assert(err, ErrorMatches, "environment is already bootstrapped")

	info, err := t.Env.StateInfo()
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)
	c.Assert(info.Addrs, Not(HasLen), 0)

	if t.CanOpenState {
		st, err := state.Open(info)
		c.Assert(err, IsNil)
		err = st.Close()
		c.Assert(err, IsNil)
	}

	c.Logf("destroy env")
	t.Destroy(c)

	// check that we can bootstrap after destroy
	t.BootstrapOnce(c)
}

func (t *LiveTests) TestBootstrapProvisioner(c *C) {
	if !t.CanOpenState || !t.HasProvisioner {
		c.Skip(fmt.Sprintf("skipping provisioner test, CanOpenState: %v, HasProvisioner: %v", t.CanOpenState, t.HasProvisioner))
	}
	t.BootstrapOnce(c)

	// TODO(dfc) constructing a juju.Conn by hand is a code smell.
	conn, err := juju.NewConnFromAttrs(t.Env.Config().AllAttrs())
	c.Assert(err, IsNil)

	st, err := conn.State()
	c.Assert(err, IsNil)

	// place a new machine into the state
	m, err := st.AddMachine()
	c.Assert(err, IsNil)

	t.assertStartInstance(c, m)

	// now remove it
	c.Assert(st.RemoveMachine(m.Id()), IsNil)

	// watch the PA remove it
	t.assertStopInstance(c, m)
	assertInstanceId(c, m, nil)

	err = st.Close()
	c.Assert(err, IsNil)
}

var waitAgent = environs.AttemptStrategy{
	Total: 30 * time.Second,
	Delay: 1 * time.Second,
}

func (t *LiveTests) assertStartInstance(c *C, m *state.Machine) {
	// Wait for machine to get an instance id.
	for a := waitAgent.Start(); a.Next(); {
		instId, err := m.InstanceId()
		if _, ok := err.(*state.NotFoundError); ok {
			continue
		}
		c.Assert(err, IsNil)
		_, err = t.Env.Instances([]string{instId})
		c.Assert(err, IsNil)
		return
	}
	c.Fatalf("provisioner failed to start machine after %v", waitAgent.Total)
}

func (t *LiveTests) assertStopInstance(c *C, m *state.Machine) {
	// Wait for machine id to be cleared.
	for a := waitAgent.Start(); a.Next(); {
		if instId, err := m.InstanceId(); instId == "" {
			c.Assert(err, FitsTypeOf, &state.NotFoundError{})
			return
		}
	}
	c.Fatalf("provisioner failed to stop machine after %v", waitAgent.Total)
}

// assertInstanceId asserts that the machine has an instance id
// that matches that of the given instance. If the instance is nil,
// It asserts that the instance id is unset.
func assertInstanceId(c *C, m *state.Machine, inst environs.Instance) {
	// TODO(dfc) add machine.WatchConfig() to avoid having to poll.
	var instId, id string
	var err error
	if inst != nil {
		instId = inst.Id()
	}
	for a := waitAgent.Start(); a.Next(); {
		id, err = m.InstanceId()
		_, notset := err.(*state.NotFoundError)
		if notset {
			if inst == nil {
				return
			}
			continue
		}
		c.Assert(err, IsNil)
		break
	}
	c.Assert(err, IsNil)
	c.Assert(id, Equals, instId)
}

// TODO check that binary data works ok?
var contents = []byte("hello\n")
var contents2 = []byte("goodbye\n\n")

func (t *LiveTests) TestFile(c *C) {
	name := fmt.Sprint("testfile", time.Now().UnixNano())
	storage := t.Env.Storage()

	checkFileDoesNotExist(c, storage, name)
	checkPutFile(c, storage, name, contents)
	checkFileHasContents(c, storage, name, contents)
	checkPutFile(c, storage, name, contents2) // check that we can overwrite the file
	checkFileHasContents(c, storage, name, contents2)

	// check that the listed contents include the
	// expected name.
	names, err := storage.List("")
	c.Assert(err, IsNil)
	found := false
	for _, lname := range names {
		if lname == name {
			found = true
			break
		}
	}
	if !found {
		c.Errorf("file name %q not found in file list %q", name, names)
	}

	err = storage.Remove(name)
	c.Check(err, IsNil)
	checkFileDoesNotExist(c, storage, name)
	// removing a file that does not exist should not be an error.
	err = storage.Remove(name)
	c.Check(err, IsNil)
}

// Check that we can't start an instance running tools
// that correspond with no available platform.
func (t *LiveTests) TestStartInstanceOnUnknownPlatform(c *C) {
	vers := version.Current
	// Note that we want this test to function correctly in the
	// dummy environment, so to avoid enumerating all possible
	// platforms in the dummy provider, it treats only series and/or
	// architectures with the "unknown" prefix as invalid.
	vers.Series = "unknownseries"
	vers.Arch = "unknownarch"
	name := environs.ToolsPath(vers)
	storage := t.Env.Storage()
	checkPutFile(c, storage, name, []byte("fake tools on invalid series"))
	defer storage.Remove(name)

	url, err := storage.URL(name)
	c.Assert(err, IsNil)
	tools := &state.Tools{
		Binary: vers,
		URL:    url,
	}

	inst, err := t.Env.StartInstance(4, InvalidStateInfo, tools)
	if inst != nil {
		err := t.Env.StopInstances([]environs.Instance{inst})
		c.Check(err, IsNil)
	}
	c.Assert(inst, IsNil)
	c.Assert(err, ErrorMatches, "cannot find image.*")
}
<commit_msg>jujutest: speed up TestBootstrap<commit_after>package jujutest

import (
	"fmt"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/version"
	"time"
)

// TestStartStop is similar to Tests.TestStartStop except
// that it does not assume a pristine environment.
func (t *LiveTests) TestStartStop(c *C) {
	insts, err := t.Env.Instances(nil)
	c.Assert(err, IsNil)
	c.Check(insts, HasLen, 0)

	inst, err := t.Env.StartInstance(0, InvalidStateInfo, nil)
	c.Assert(err, IsNil)
	c.Assert(inst, NotNil)
	id0 := inst.Id()

	insts, err = t.Env.Instances([]string{id0, id0})
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 2)
	c.Assert(insts[0].Id(), Equals, id0)
	c.Assert(insts[1].Id(), Equals, id0)

	insts, err = t.Env.AllInstances()
	c.Assert(err, IsNil)
	// differs from the check above because AllInstances returns
	// a set (without duplicates) containing only one instance.
	c.Assert(insts, HasLen, 1)
	c.Assert(insts[0].Id(), Equals, id0)

	dns, err := inst.WaitDNSName()
	c.Assert(err, IsNil)
	c.Assert(dns, Not(Equals), "")

	insts, err = t.Env.Instances([]string{id0, ""})
	c.Assert(err, Equals, environs.ErrPartialInstances)
	c.Assert(insts, HasLen, 2)
	c.Check(insts[0].Id(), Equals, id0)
	c.Check(insts[1], IsNil)

	err = t.Env.StopInstances([]environs.Instance{inst})
	c.Assert(err, IsNil)

	// Stopping may not be noticed at first due to eventual
	// consistency. Repeat a few times to ensure we get the error.
	for i := 0; i < 20; i++ {
		insts, err = t.Env.Instances([]string{id0})
		if err != nil {
			break
		}
		time.Sleep(0.25e9)
	}
	c.Assert(err, Equals, environs.ErrNoInstances)
	c.Assert(insts, HasLen, 0)
}

func (t *LiveTests) TestPorts(c *C) {
	inst1, err := t.Env.StartInstance(1, InvalidStateInfo, nil)
	c.Assert(err, IsNil)
	c.Assert(inst1, NotNil)
	defer t.Env.StopInstances([]environs.Instance{inst1})

	ports, err := inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, HasLen, 0)

	inst2, err := t.Env.StartInstance(2, InvalidStateInfo, nil)
	c.Assert(err, IsNil)
	c.Assert(inst2, NotNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, HasLen, 0)
	defer t.Env.StopInstances([]environs.Instance{inst2})

	// Open some ports and check they're there.
	err = inst1.OpenPorts(1, []state.Port{{"udp", 67}, {"tcp", 45}})
	c.Assert(err, IsNil)
	ports, err = inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"udp", 67}})
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, HasLen, 0)

	// Check there's no crosstalk to another machine
	err = inst2.OpenPorts(2, []state.Port{{"tcp", 89}, {"tcp", 45}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"tcp", 89}})
	ports, err = inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"udp", 67}})

	// Check that opening the same port again is ok.
	err = inst2.OpenPorts(2, []state.Port{{"tcp", 45}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"tcp", 89}})

	// Check that opening the same port again and another port is ok.
	err = inst2.OpenPorts(2, []state.Port{{"tcp", 45}, {"tcp", 99}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"tcp", 89}, {"tcp", 99}})

	// Check that we can close ports and that there's no crosstalk.
	err = inst2.ClosePorts(2, []state.Port{{"tcp", 45}, {"tcp", 99}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 89}})
	ports, err = inst1.Ports(1)
	c.Assert(err, IsNil)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 45}, {"udp", 67}})

	// Check that we can close multiple ports.
	err = inst1.ClosePorts(1, []state.Port{{"tcp", 45}, {"udp", 67}})
	c.Assert(err, IsNil)
	ports, err = inst1.Ports(1)
	c.Assert(ports, HasLen, 0)

	// Check that we can close ports that aren't there.
	err = inst2.ClosePorts(2, []state.Port{{"tcp", 111}, {"udp", 222}})
	c.Assert(err, IsNil)
	ports, err = inst2.Ports(2)
	c.Assert(ports, DeepEquals, []state.Port{{"tcp", 89}})
}

func (t *LiveTests) TestBootstrapMultiple(c *C) {
	t.BootstrapOnce(c)

	// Wait for a while to let eventual consistency catch up, hopefully.
	time.Sleep(t.ConsistencyDelay)
	err := t.Env.Bootstrap(false)
	c.Assert(err, ErrorMatches, "environment is already bootstrapped")

	c.Logf("destroy env")
	t.Destroy(c)

	// check that we can bootstrap after destroy
	t.BootstrapOnce(c)
}

func (t *LiveTests) TestBootstrapProvisioner(c *C) {
	if !t.CanOpenState || !t.HasProvisioner {
		c.Skip(fmt.Sprintf("skipping provisioner test, CanOpenState: %v, HasProvisioner: %v", t.CanOpenState, t.HasProvisioner))
	}
	t.BootstrapOnce(c)

	// TODO(dfc) constructing a juju.Conn by hand is a code smell.
	conn, err := juju.NewConnFromAttrs(t.Env.Config().AllAttrs())
	c.Assert(err, IsNil)

	st, err := conn.State()
	c.Assert(err, IsNil)

	// place a new machine into the state
	m, err := st.AddMachine()
	c.Assert(err, IsNil)

	t.assertStartInstance(c, m)

	// now remove it
	c.Assert(st.RemoveMachine(m.Id()), IsNil)

	// watch the PA remove it
	t.assertStopInstance(c, m)
	assertInstanceId(c, m, nil)

	err = st.Close()
	c.Assert(err, IsNil)
}

var waitAgent = environs.AttemptStrategy{
	Total: 30 * time.Second,
	Delay: 1 * time.Second,
}

func (t *LiveTests) assertStartInstance(c *C, m *state.Machine) {
	// Wait for machine to get an instance id.
	for a := waitAgent.Start(); a.Next(); {
		instId, err := m.InstanceId()
		if _, ok := err.(*state.NotFoundError); ok {
			continue
		}
		c.Assert(err, IsNil)
		_, err = t.Env.Instances([]string{instId})
		c.Assert(err, IsNil)
		return
	}
	c.Fatalf("provisioner failed to start machine after %v", waitAgent.Total)
}

func (t *LiveTests) assertStopInstance(c *C, m *state.Machine) {
	// Wait for machine id to be cleared.
	for a := waitAgent.Start(); a.Next(); {
		if instId, err := m.InstanceId(); instId == "" {
			c.Assert(err, FitsTypeOf, &state.NotFoundError{})
			return
		}
	}
	c.Fatalf("provisioner failed to stop machine after %v", waitAgent.Total)
}

// assertInstanceId asserts that the machine has an instance id
// that matches that of the given instance. If the instance is nil,
// It asserts that the instance id is unset.
func assertInstanceId(c *C, m *state.Machine, inst environs.Instance) {
	// TODO(dfc) add machine.WatchConfig() to avoid having to poll.
	var instId, id string
	var err error
	if inst != nil {
		instId = inst.Id()
	}
	for a := waitAgent.Start(); a.Next(); {
		id, err = m.InstanceId()
		_, notset := err.(*state.NotFoundError)
		if notset {
			if inst == nil {
				return
			}
			continue
		}
		c.Assert(err, IsNil)
		break
	}
	c.Assert(err, IsNil)
	c.Assert(id, Equals, instId)
}

// TODO check that binary data works ok?
var contents = []byte("hello\n")
var contents2 = []byte("goodbye\n\n")

func (t *LiveTests) TestFile(c *C) {
	name := fmt.Sprint("testfile", time.Now().UnixNano())
	storage := t.Env.Storage()

	checkFileDoesNotExist(c, storage, name)
	checkPutFile(c, storage, name, contents)
	checkFileHasContents(c, storage, name, contents)
	checkPutFile(c, storage, name, contents2) // check that we can overwrite the file
	checkFileHasContents(c, storage, name, contents2)

	// check that the listed contents include the
	// expected name.
	names, err := storage.List("")
	c.Assert(err, IsNil)
	found := false
	for _, lname := range names {
		if lname == name {
			found = true
			break
		}
	}
	if !found {
		c.Errorf("file name %q not found in file list %q", name, names)
	}

	err = storage.Remove(name)
	c.Check(err, IsNil)
	checkFileDoesNotExist(c, storage, name)
	// removing a file that does not exist should not be an error.
	err = storage.Remove(name)
	c.Check(err, IsNil)
}

// Check that we can't start an instance running tools
// that correspond with no available platform.
func (t *LiveTests) TestStartInstanceOnUnknownPlatform(c *C) {
	vers := version.Current
	// Note that we want this test to function correctly in the
	// dummy environment, so to avoid enumerating all possible
	// platforms in the dummy provider, it treats only series and/or
	// architectures with the "unknown" prefix as invalid.
	vers.Series = "unknownseries"
	vers.Arch = "unknownarch"
	name := environs.ToolsPath(vers)
	storage := t.Env.Storage()
	checkPutFile(c, storage, name, []byte("fake tools on invalid series"))
	defer storage.Remove(name)

	url, err := storage.URL(name)
	c.Assert(err, IsNil)
	tools := &state.Tools{
		Binary: vers,
		URL:    url,
	}

	inst, err := t.Env.StartInstance(4, InvalidStateInfo, tools)
	if inst != nil {
		err := t.Env.StopInstances([]environs.Instance{inst})
		c.Check(err, IsNil)
	}
	c.Assert(inst, IsNil)
	c.Assert(err, ErrorMatches, "cannot find image.*")
}
<|endoftext|>


<commit_before>// Package consistent provides a consistent hashing function.
//
// Consistent hashing is often used to distribute requests to a changing set of servers.  For example,
// say you have some cache servers cacheA, cacheB, and cacheC.  You want to decide which cache server
// to use to look up information on a user.
//
// You could use a typical hash table and hash the user id
// to one of cacheA, cacheB, or cacheC.  But with a typical hash table, if you add or remove a server,
// almost all keys will get remapped to different results, which basically could bring your service
// to a grinding halt while the caches get rebuilt.
//
// With a consistent hash, adding or removing a server drastically reduces the number of keys that
// get remapped.
//
// Read more about consistent hashing on wikipedia:  http://en.wikipedia.org/wiki/Consistent_hashing
//
package consistent

import (
	"errors"
	"hash/crc32"
	"sort"
	"strconv"
)

type uints []uint32

// Len returns the length of the uints array.
func (x uints) Len() int { return len(x) }

// Less returns true if element i is less than element j.
func (x uints) Less(i, j int) bool { return x[i] < x[j] }

// Swap exchanges elements i and j.
func (x uints) Swap(i, j int) { x[i], x[j] = x[j], x[i] }

// ErrEmptyCircle is the error returned when trying to get an element when nothing has been added to hash.
var ErrEmptyCircle = errors.New("empty circle")

// Consistent holds the information about the members of the consistent hash circle.
type Consistent struct {
	circle           map[uint32]string
	sortedHashes     uints
	NumberOfReplicas int
	count            int64
	scratch          [64]byte
}

// New creates a new Consistent object with a default setting of 20 replicas for each entry.
//
// To change the number of replicas, set NumberOfReplicas before adding entries.
func New() *Consistent {
	c := new(Consistent)
	c.NumberOfReplicas = 20
	c.circle = make(map[uint32]string)
	return c
}

// eltKey generates a string key for an element with an index.
func (c *Consistent) eltKey(elt string, idx int) string {
	return elt + "|" + strconv.Itoa(idx)
}

// Add inserts a string element in the consistent hash.
func (c *Consistent) Add(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		c.circle[c.hashKey(c.eltKey(elt, i))] = elt
	}
	c.updateSortedHashes()
	c.count++
}

// Remove removes an element from the hash.
func (c *Consistent) Remove(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		delete(c.circle, c.hashKey(c.eltKey(elt, i)))
	}
	c.updateSortedHashes()
	c.count--
}

// Get returns an element close to where name hashes to in the circle.
func (c *Consistent) Get(name string) (string, error) {
	if len(c.circle) == 0 {
		return "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i := sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	return c.circle[c.sortedHashes[i]], nil
}

// GetTwo returns the two closest distinct elements to the name input in the circle.
func (c *Consistent) GetTwo(name string) (string, string, error) {
	if len(c.circle) == 0 {
		return "", "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i := sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	a := c.circle[c.sortedHashes[i]]

	if c.count == 1 {
		return a, "", nil
	}

	start := i
	i++
	if i >= len(c.sortedHashes) {
		i = 0
	}
	var b string
	for i != start {
		b = c.circle[c.sortedHashes[i]]
		if b != a {
			break
		}
		i++
		if i >= len(c.sortedHashes) {
			i = 0
		}
	}
	return a, b, nil
}

func (c *Consistent) hashKey(key string) uint32 {
	if len(key) < 64 {
		copy(c.scratch[:], key)
		return crc32.ChecksumIEEE(c.scratch[:len(key)])
	}
	return crc32.ChecksumIEEE([]byte(key))
}

func (c *Consistent) updateSortedHashes() {
	hashes := c.sortedHashes[:0]
	//reallocate if we're holding on to too much (1/4th)
	if cap(c.sortedHashes)/(c.NumberOfReplicas*4) > len(c.circle) {
		hashes = nil
	}
	for k := range c.circle {
		hashes = append(hashes, k)
	}
	sort.Sort(hashes)
	c.sortedHashes = hashes
}
<commit_msg>Some code clarity changes<commit_after>// Package consistent provides a consistent hashing function.
//
// Consistent hashing is often used to distribute requests to a changing set of servers.  For example,
// say you have some cache servers cacheA, cacheB, and cacheC.  You want to decide which cache server
// to use to look up information on a user.
//
// You could use a typical hash table and hash the user id
// to one of cacheA, cacheB, or cacheC.  But with a typical hash table, if you add or remove a server,
// almost all keys will get remapped to different results, which basically could bring your service
// to a grinding halt while the caches get rebuilt.
//
// With a consistent hash, adding or removing a server drastically reduces the number of keys that
// get remapped.
//
// Read more about consistent hashing on wikipedia:  http://en.wikipedia.org/wiki/Consistent_hashing
//
package consistent

import (
	"errors"
	"hash/crc32"
	"sort"
	"strconv"
)

type uints []uint32

// Len returns the length of the uints array.
func (x uints) Len() int { return len(x) }

// Less returns true if element i is less than element j.
func (x uints) Less(i, j int) bool { return x[i] < x[j] }

// Swap exchanges elements i and j.
func (x uints) Swap(i, j int) { x[i], x[j] = x[j], x[i] }

// ErrEmptyCircle is the error returned when trying to get an element when nothing has been added to hash.
var ErrEmptyCircle = errors.New("empty circle")

// Consistent holds the information about the members of the consistent hash circle.
type Consistent struct {
	circle           map[uint32]string
	sortedHashes     uints
	NumberOfReplicas int
	count            int64
	scratch          [64]byte
}

// New creates a new Consistent object with a default setting of 20 replicas for each entry.
//
// To change the number of replicas, set NumberOfReplicas before adding entries.
func New() *Consistent {
	c := new(Consistent)
	c.NumberOfReplicas = 20
	c.circle = make(map[uint32]string)
	return c
}

// eltKey generates a string key for an element with an index.
func (c *Consistent) eltKey(elt string, idx int) string {
	return elt + "|" + strconv.Itoa(idx)
}

// Add inserts a string element in the consistent hash.
func (c *Consistent) Add(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		c.circle[c.hashKey(c.eltKey(elt, i))] = elt
	}
	c.updateSortedHashes()
	c.count++
}

// Remove removes an element from the hash.
func (c *Consistent) Remove(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		delete(c.circle, c.hashKey(c.eltKey(elt, i)))
	}
	c.updateSortedHashes()
	c.count--
}

// Get returns an element close to where name hashes to in the circle.
func (c *Consistent) Get(name string) (string, error) {
	if len(c.circle) == 0 {
		return "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	i := c.search(key)
	return c.circle[c.sortedHashes[i]], nil
}

func (c *Consistent) search(key uint32) (i int) {
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i = sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	return
}

// GetTwo returns the two closest distinct elements to the name input in the circle.
func (c *Consistent) GetTwo(name string) (string, string, error) {
	if len(c.circle) == 0 {
		return "", "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	i := c.search(key)
	a := c.circle[c.sortedHashes[i]]

	if c.count == 1 {
		return a, "", nil
	}

	start := i
	var b string
	for i = start + 1; i != start; i++ {
		if i >= len(c.sortedHashes) {
			i = 0
		}
		b = c.circle[c.sortedHashes[i]]
		if b != a {
			break
		}
	}
	return a, b, nil
}

func (c *Consistent) hashKey(key string) uint32 {
	if len(key) < 64 {
		copy(c.scratch[:], key)
		return crc32.ChecksumIEEE(c.scratch[:len(key)])
	}
	return crc32.ChecksumIEEE([]byte(key))
}

func (c *Consistent) updateSortedHashes() {
	hashes := c.sortedHashes[:0]
	//reallocate if we're holding on to too much (1/4th)
	if cap(c.sortedHashes)/(c.NumberOfReplicas*4) > len(c.circle) {
		hashes = nil
	}
	for k := range c.circle {
		hashes = append(hashes, k)
	}
	sort.Sort(hashes)
	c.sortedHashes = hashes
}
<|endoftext|>


<commit_before>package speedtest

import (
	"context"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"golang.org/x/sync/errgroup"
)

type downloadWarmUpFunc func(context.Context, string) error
type downloadFunc func(context.Context, string, int) error
type uploadWarmUpFunc func(context.Context, string) error
type uploadFunc func(context.Context, string, int) error

var dlSizes = [...]int{350, 500, 750, 1000, 1500, 2000, 2500, 3000, 3500, 4000}
var ulSizes = [...]int{100, 300, 500, 800, 1000, 1500, 2500, 3000, 3500, 4000} //kB
var client = http.Client{}

// DownloadTest executes the test to measure download speed
func (s *Server) DownloadTest(savingMode bool) error {
	return s.DownloadTestContext(context.Background(), savingMode, dlWarmUp, downloadRequest)
}

// DownloadTestContext executes the test to measure download speed, observing the given context.
func (s *Server) DownloadTestContext(
	ctx context.Context,
	savingMode bool,
	dlWarmUp downloadWarmUpFunc,
	downloadRequest downloadFunc,
) error {
	dlURL := strings.Split(s.URL, "/upload.php")[0]
	eg := errgroup.Group{}

	// Warming up
	sTime := time.Now()
	for i := 0; i < 2; i++ {
		eg.Go(func() error {
			return dlWarmUp(ctx, dlURL)
		})
	}
	if err := eg.Wait(); err != nil {
		return err
	}
	fTime := time.Now()
	// 1.125MB for each request (750 * 750 * 2)
	wuSpeed := 1.125 * 8 * 2 / fTime.Sub(sTime.Add(s.Latency)).Seconds()

	// Decide workload by warm up speed
	workload := 0
	weight := 0
	skip := false
	if savingMode {
		workload = 6
		weight = 3
	} else if 10.0 < wuSpeed {
		workload = 16
		weight = 4
	} else if 4.0 < wuSpeed {
		workload = 8
		weight = 4
	} else if 2.5 < wuSpeed {
		workload = 4
		weight = 4
	} else {
		skip = true
	}

	// Main speedtest
	dlSpeed := wuSpeed
	if !skip {
		sTime = time.Now()
		for i := 0; i < workload; i++ {
			eg.Go(func() error {
				return downloadRequest(ctx, dlURL, weight)
			})
		}
		if err := eg.Wait(); err != nil {
			return err
		}
		fTime = time.Now()

		reqMB := dlSizes[weight] * dlSizes[weight] * 2 / 1000 / 1000
		dlSpeed = float64(reqMB) * 8 * float64(workload) / fTime.Sub(sTime).Seconds()
	}

	s.DLSpeed = dlSpeed
	return nil
}

// UploadTest executes the test to measure upload speed
func (s *Server) UploadTest(savingMode bool) error {
	return s.UploadTestContext(context.Background(), savingMode, ulWarmUp, uploadRequest)
}

// UploadTestContext executes the test to measure upload speed, observing the given context.
func (s *Server) UploadTestContext(
	ctx context.Context,
	savingMode bool,
	ulWarmUp uploadWarmUpFunc,
	uploadRequest uploadFunc,
) error {
	// Warm up
	sTime := time.Now()
	eg := errgroup.Group{}
	for i := 0; i < 2; i++ {
		eg.Go(func() error {
			return ulWarmUp(ctx, s.URL)
		})
	}
	if err := eg.Wait(); err != nil {
		return err
	}
	fTime := time.Now()
	// 1.0 MB for each request
	wuSpeed := 1.0 * 8 * 2 / fTime.Sub(sTime.Add(s.Latency)).Seconds()

	// Decide workload by warm up speed
	workload := 0
	weight := 0
	skip := false
	if savingMode {
		workload = 1
		weight = 7
	} else if 10.0 < wuSpeed {
		workload = 16
		weight = 9
	} else if 4.0 < wuSpeed {
		workload = 8
		weight = 9
	} else if 2.5 < wuSpeed {
		workload = 4
		weight = 5
	} else {
		skip = true
	}

	// Main speedtest
	ulSpeed := wuSpeed
	if !skip {
		sTime = time.Now()
		for i := 0; i < workload; i++ {
			eg.Go(func() error {
				return uploadRequest(ctx, s.URL, weight)
			})
		}
		if err := eg.Wait(); err != nil {
			return err
		}
		fTime = time.Now()

		reqMB := float64(ulSizes[weight]) / 1000
		ulSpeed = reqMB * 8 * float64(workload) / fTime.Sub(sTime).Seconds()
	}

	s.ULSpeed = ulSpeed

	return nil
}

func dlWarmUp(ctx context.Context, dlURL string) error {
	size := dlSizes[2]
	xdlURL := dlURL + "/random" + strconv.Itoa(size) + "x" + strconv.Itoa(size) + ".jpg"

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, xdlURL, nil)
	if err != nil {
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

func ulWarmUp(ctx context.Context, ulURL string) error {
	size := ulSizes[4]
	v := url.Values{}
	v.Add("content", strings.Repeat("0123456789", size*100-51))

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, ulURL, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

func downloadRequest(ctx context.Context, dlURL string, w int) error {
	size := dlSizes[w]
	xdlURL := dlURL + "/random" + strconv.Itoa(size) + "x" + strconv.Itoa(size) + ".jpg"

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, xdlURL, nil)
	if err != nil {
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

func uploadRequest(ctx context.Context, ulURL string, w int) error {
	size := ulSizes[w]
	v := url.Values{}
	v.Add("content", strings.Repeat("0123456789", size*100-51))

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, ulURL, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

// PingTest executes test to measure latency
func (s *Server) PingTest() error {
	return s.PingTestContext(context.Background())
}

// PingTestContext executes test to measure latency, observing the given context.
func (s *Server) PingTestContext(ctx context.Context) error {
	pingURL := strings.Split(s.URL, "/upload.php")[0] + "/latency.txt"

	l := time.Duration(100000000000) // 10sec
	for i := 0; i < 3; i++ {
		sTime := time.Now()

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, pingURL, nil)
		if err != nil {
			return err
		}

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}

		fTime := time.Now()
		if fTime.Sub(sTime) < l {
			l = fTime.Sub(sTime)
		}

		resp.Body.Close()
	}

	s.Latency = time.Duration(int64(l.Nanoseconds() / 2))

	return nil
}
<commit_msg>[add] workload option for high speed env<commit_after>package speedtest

import (
	"context"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"golang.org/x/sync/errgroup"
)

type downloadWarmUpFunc func(context.Context, string) error
type downloadFunc func(context.Context, string, int) error
type uploadWarmUpFunc func(context.Context, string) error
type uploadFunc func(context.Context, string, int) error

var dlSizes = [...]int{350, 500, 750, 1000, 1500, 2000, 2500, 3000, 3500, 4000}
var ulSizes = [...]int{100, 300, 500, 800, 1000, 1500, 2500, 3000, 3500, 4000} //kB
var client = http.Client{}

// DownloadTest executes the test to measure download speed
func (s *Server) DownloadTest(savingMode bool) error {
	return s.DownloadTestContext(context.Background(), savingMode, dlWarmUp, downloadRequest)
}

// DownloadTestContext executes the test to measure download speed, observing the given context.
func (s *Server) DownloadTestContext(
	ctx context.Context,
	savingMode bool,
	dlWarmUp downloadWarmUpFunc,
	downloadRequest downloadFunc,
) error {
	dlURL := strings.Split(s.URL, "/upload.php")[0]
	eg := errgroup.Group{}

	// Warming up
	sTime := time.Now()
	for i := 0; i < 2; i++ {
		eg.Go(func() error {
			return dlWarmUp(ctx, dlURL)
		})
	}
	if err := eg.Wait(); err != nil {
		return err
	}
	fTime := time.Now()
	// 1.125MB for each request (750 * 750 * 2)
	wuSpeed := 1.125 * 8 * 2 / fTime.Sub(sTime.Add(s.Latency)).Seconds()

	// Decide workload by warm up speed
	workload := 0
	weight := 0
	skip := false
	if savingMode {
		workload = 6
		weight = 3
	} else if 50.0 < wuSpeed {
		workload = 32
		weight = 6
	} else if 10.0 < wuSpeed {
		workload = 16
		weight = 4
	} else if 4.0 < wuSpeed {
		workload = 8
		weight = 4
	} else if 2.5 < wuSpeed {
		workload = 4
		weight = 4
	} else {
		skip = true
	}

	// Main speedtest
	dlSpeed := wuSpeed
	if !skip {
		sTime = time.Now()
		for i := 0; i < workload; i++ {
			eg.Go(func() error {
				return downloadRequest(ctx, dlURL, weight)
			})
		}
		if err := eg.Wait(); err != nil {
			return err
		}
		fTime = time.Now()

		reqMB := dlSizes[weight] * dlSizes[weight] * 2 / 1000 / 1000
		dlSpeed = float64(reqMB) * 8 * float64(workload) / fTime.Sub(sTime).Seconds()
	}

	s.DLSpeed = dlSpeed
	return nil
}

// UploadTest executes the test to measure upload speed
func (s *Server) UploadTest(savingMode bool) error {
	return s.UploadTestContext(context.Background(), savingMode, ulWarmUp, uploadRequest)
}

// UploadTestContext executes the test to measure upload speed, observing the given context.
func (s *Server) UploadTestContext(
	ctx context.Context,
	savingMode bool,
	ulWarmUp uploadWarmUpFunc,
	uploadRequest uploadFunc,
) error {
	// Warm up
	sTime := time.Now()
	eg := errgroup.Group{}
	for i := 0; i < 2; i++ {
		eg.Go(func() error {
			return ulWarmUp(ctx, s.URL)
		})
	}
	if err := eg.Wait(); err != nil {
		return err
	}
	fTime := time.Now()
	// 1.0 MB for each request
	wuSpeed := 1.0 * 8 * 2 / fTime.Sub(sTime.Add(s.Latency)).Seconds()

	// Decide workload by warm up speed
	workload := 0
	weight := 0
	skip := false
	if savingMode {
		workload = 1
		weight = 7
	} else if 50.0 < wuSpeed {
		workload = 40
		weight = 9
	} else if 10.0 < wuSpeed {
		workload = 16
		weight = 9
	} else if 4.0 < wuSpeed {
		workload = 8
		weight = 9
	} else if 2.5 < wuSpeed {
		workload = 4
		weight = 5
	} else {
		skip = true
	}

	// Main speedtest
	ulSpeed := wuSpeed
	if !skip {
		sTime = time.Now()
		for i := 0; i < workload; i++ {
			eg.Go(func() error {
				return uploadRequest(ctx, s.URL, weight)
			})
		}
		if err := eg.Wait(); err != nil {
			return err
		}
		fTime = time.Now()

		reqMB := float64(ulSizes[weight]) / 1000
		ulSpeed = reqMB * 8 * float64(workload) / fTime.Sub(sTime).Seconds()
	}

	s.ULSpeed = ulSpeed

	return nil
}

func dlWarmUp(ctx context.Context, dlURL string) error {
	size := dlSizes[2]
	xdlURL := dlURL + "/random" + strconv.Itoa(size) + "x" + strconv.Itoa(size) + ".jpg"

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, xdlURL, nil)
	if err != nil {
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

func ulWarmUp(ctx context.Context, ulURL string) error {
	size := ulSizes[4]
	v := url.Values{}
	v.Add("content", strings.Repeat("0123456789", size*100-51))

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, ulURL, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

func downloadRequest(ctx context.Context, dlURL string, w int) error {
	size := dlSizes[w]
	xdlURL := dlURL + "/random" + strconv.Itoa(size) + "x" + strconv.Itoa(size) + ".jpg"

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, xdlURL, nil)
	if err != nil {
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

func uploadRequest(ctx context.Context, ulURL string, w int) error {
	size := ulSizes[w]
	v := url.Values{}
	v.Add("content", strings.Repeat("0123456789", size*100-51))

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, ulURL, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	_, err = io.Copy(ioutil.Discard, resp.Body)
	return err
}

// PingTest executes test to measure latency
func (s *Server) PingTest() error {
	return s.PingTestContext(context.Background())
}

// PingTestContext executes test to measure latency, observing the given context.
func (s *Server) PingTestContext(ctx context.Context) error {
	pingURL := strings.Split(s.URL, "/upload.php")[0] + "/latency.txt"

	l := time.Duration(100000000000) // 10sec
	for i := 0; i < 3; i++ {
		sTime := time.Now()

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, pingURL, nil)
		if err != nil {
			return err
		}

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}

		fTime := time.Now()
		if fTime.Sub(sTime) < l {
			l = fTime.Sub(sTime)
		}

		resp.Body.Close()
	}

	s.Latency = time.Duration(int64(l.Nanoseconds() / 2))

	return nil
}
<|endoftext|>


<commit_before>package u1

import (
	"bytes"
	"io"
	"net/http"

	"github.com/russross/blackfriday"
)

// GitHub Flavored Markdown-like extensions.
var MarkdownGfmExtensions = 0 |
	blackfriday.EXTENSION_NO_INTRA_EMPHASIS |
	//blackfriday.EXTENSION_TABLES | // TODO: Implement. Maybe.
	blackfriday.EXTENSION_FENCED_CODE |
	blackfriday.EXTENSION_AUTOLINK |
	blackfriday.EXTENSION_STRIKETHROUGH |
	blackfriday.EXTENSION_SPACE_HEADERS
	//blackfriday.EXTENSION_HARD_LINE_BREAK

// Best effort at generating GitHub Flavored Markdown-like HTML output locally.
func MarkdownGfm(input []byte) []byte {
	htmlFlags := 0 |
		blackfriday.HTML_SANITIZE_OUTPUT |
		blackfriday.HTML_GITHUB_BLOCKCODE

	renderer := blackfriday.HtmlRenderer(htmlFlags, "", "")

	return blackfriday.Markdown(input, renderer, MarkdownGfmExtensions)
}

// ---

// Convert GitHub Flavored Markdown to full HTML page and write to w.
// TODO: Do this locally via a native Go library... That's not too much to ask for, is it?
func WriteMarkdownGfmAsHtmlPage(w io.Writer, markdown []byte) {
	// TODO: Do GitHub, fallback to local if it fails.
	writeGitHubFlavoredMarkdownViaGitHub(w, markdown)
	//writeGitHubFlavoredMarkdownViaLocal(w, markdown)
}

func writeGitHubFlavoredMarkdownViaLocal(w io.Writer, markdown []byte) {
	// TODO: Don't hotlink the css file from github.com, serve it locally (it's needed for the GFM html to appear properly)
	// TODO: Use github.com/sourcegraph/syntaxhighlight to add missing syntax highlighting.
	io.WriteString(w, `<meta charset="utf-8"><link href="https://github.com/assets/github.css" media="all" rel="stylesheet" type="text/css" />
`)
	w.Write(MarkdownGfm(markdown))
	io.WriteString(w, `</article></body></html>`)
}

// TODO: Remove once local version gives matching results.
func writeGitHubFlavoredMarkdownViaGitHub(w io.Writer, markdown []byte) {
	// TODO: Don't hotlink the css file from github.com, serve it locally (it's needed for the GFM html to appear properly)
	io.WriteString(w, `<meta charset="utf-8"><link href="https://github.com/assets/github.css" media="all" rel="stylesheet" type="text/css" />
`)

	// Convert GitHub-Flavored-Markdown to HTML (includes syntax highlighting for diff, Go, etc.)
	resp, err := http.Post("https://api.github.com/markdown/raw", "text/x-markdown", bytes.NewReader(markdown))
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		panic(err)
	}

	io.WriteString(w, `</article></body></html>`)
}
<commit_msg>Export WriteGitHubFlavoredMarkdownViaLocal().<commit_after>package u1

import (
	"bytes"
	"io"
	"net/http"

	"github.com/russross/blackfriday"
)

// GitHub Flavored Markdown-like extensions.
var MarkdownGfmExtensions = 0 |
	blackfriday.EXTENSION_NO_INTRA_EMPHASIS |
	//blackfriday.EXTENSION_TABLES | // TODO: Implement. Maybe.
	blackfriday.EXTENSION_FENCED_CODE |
	blackfriday.EXTENSION_AUTOLINK |
	blackfriday.EXTENSION_STRIKETHROUGH |
	blackfriday.EXTENSION_SPACE_HEADERS
	//blackfriday.EXTENSION_HARD_LINE_BREAK

// Best effort at generating GitHub Flavored Markdown-like HTML output locally.
func MarkdownGfm(input []byte) []byte {
	htmlFlags := 0 |
		blackfriday.HTML_SANITIZE_OUTPUT |
		blackfriday.HTML_GITHUB_BLOCKCODE

	renderer := blackfriday.HtmlRenderer(htmlFlags, "", "")

	return blackfriday.Markdown(input, renderer, MarkdownGfmExtensions)
}

// ---

// Convert GitHub Flavored Markdown to full HTML page and write to w.
// TODO: Do this locally via a native Go library... That's not too much to ask for, is it?
func WriteMarkdownGfmAsHtmlPage(w io.Writer, markdown []byte) {
	// TODO: Do GitHub, fallback to local if it fails.
	writeGitHubFlavoredMarkdownViaGitHub(w, markdown)
	//WriteGitHubFlavoredMarkdownViaLocal(w, markdown)
}

func WriteGitHubFlavoredMarkdownViaLocal(w io.Writer, markdown []byte) {
	// TODO: Don't hotlink the css file from github.com, serve it locally (it's needed for the GFM html to appear properly)
	// TODO: Use github.com/sourcegraph/syntaxhighlight to add missing syntax highlighting.
	io.WriteString(w, `<meta charset="utf-8"><link href="https://github.com/assets/github.css" media="all" rel="stylesheet" type="text/css" />
`)
	w.Write(MarkdownGfm(markdown))
	io.WriteString(w, `</article></body></html>`)
}

// TODO: Remove once local version gives matching results.
func writeGitHubFlavoredMarkdownViaGitHub(w io.Writer, markdown []byte) {
	// TODO: Don't hotlink the css file from github.com, serve it locally (it's needed for the GFM html to appear properly)
	io.WriteString(w, `<meta charset="utf-8"><link href="https://github.com/assets/github.css" media="all" rel="stylesheet" type="text/css" />
`)

	// Convert GitHub-Flavored-Markdown to HTML (includes syntax highlighting for diff, Go, etc.)
	resp, err := http.Post("https://api.github.com/markdown/raw", "text/x-markdown", bytes.NewReader(markdown))
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		panic(err)
	}

	io.WriteString(w, `</article></body></html>`)
}
<|endoftext|>


<commit_before>package widget

import (
	"net/http"

	"github.com/qor/admin"
	"github.com/qor/responder"
	"github.com/qor/serializable_meta"
)

type widgetController struct {
	Widgets *Widgets
}

func (wc widgetController) Index(context *admin.Context) {
	context = context.NewResourceContext(wc.Widgets.WidgetSettingResource)
	result, _, err := wc.getWidget(context)
	context.AddError(err)

	if context.HasError() {
		http.NotFound(context.Writer, context.Request)
	} else {
		responder.With("html", func() {
			context.Execute("index", result)
		}).With("json", func() {
			context.JSON("index", result)
		}).Respond(context.Request)
	}
}

func (wc widgetController) Edit(context *admin.Context) {
	context.Resource = wc.Widgets.WidgetSettingResource
	widgetSetting, scopes, err := wc.getWidget(context)
	context.AddError(err)
	context.Execute("edit", map[string]interface{}{"Scopes": scopes, "Widget": widgetSetting})
}

func (wc widgetController) Update(context *admin.Context) {
	context.Resource = wc.Widgets.WidgetSettingResource
	widgetSetting, scopes, err := wc.getWidget(context)
	context.AddError(err)

	if context.AddError(context.Resource.Decode(context.Context, widgetSetting)); !context.HasError() {
		context.AddError(context.Resource.CallSave(widgetSetting, context.Context))
	}

	if context.HasError() {
		context.Execute("edit", map[string]interface{}{"Scopes": scopes, "Widget": widgetSetting})
	} else {
		http.Redirect(context.Writer, context.Request, context.Request.URL.Path, http.StatusFound)
	}
}

func (wc widgetController) InlineEdit(context *admin.Context) {
	context.Writer.Write([]byte(context.Render("inline_edit")))
}

func (wc widgetController) getWidget(context *admin.Context) (interface{}, []string, error) {
	if context.ResourceID == "" {
		// index page
		context.SetDB(context.GetDB().Where("scope = ?", "default"))
		results, err := context.FindMany()
		return results, []string{}, err
	}

	// show page
	result := wc.Widgets.WidgetSettingResource.NewStruct()
	scope := context.Request.URL.Query().Get("widget_scope")

	var scopes []string
	context.GetDB().Model(result).Where("name = ?", context.ResourceID).Pluck("scope", &scopes)

	var hasScope bool

	for _, s := range scopes {
		if scope == s {
			hasScope = true
			break
		}
	}

	if !hasScope {
		scope = "default"
	}

	err := context.GetDB().First(result, "name = ? AND scope = ?", context.ResourceID, scope).Error

	if widgetType := context.Request.URL.Query().Get("widget_type"); widgetType != "" {
		if serializableMeta, ok := result.(serializable_meta.SerializableMetaInterface); ok {
			serializableMeta.SetSerializableArgumentKind(widgetType)
			serializableMeta.SetSerializableArgumentValue(nil)
		}
	}
	return result, scopes, err
}
<commit_msg>Add HTTPUnprocessableEntity code for widget update<commit_after>package widget

import (
	"net/http"

	"github.com/qor/admin"
	"github.com/qor/responder"
	"github.com/qor/serializable_meta"
)

type widgetController struct {
	Widgets *Widgets
}

func (wc widgetController) Index(context *admin.Context) {
	context = context.NewResourceContext(wc.Widgets.WidgetSettingResource)
	result, _, err := wc.getWidget(context)
	context.AddError(err)

	if context.HasError() {
		http.NotFound(context.Writer, context.Request)
	} else {
		responder.With("html", func() {
			context.Execute("index", result)
		}).With("json", func() {
			context.JSON("index", result)
		}).Respond(context.Request)
	}
}

func (wc widgetController) Edit(context *admin.Context) {
	context.Resource = wc.Widgets.WidgetSettingResource
	widgetSetting, scopes, err := wc.getWidget(context)
	context.AddError(err)
	context.Execute("edit", map[string]interface{}{"Scopes": scopes, "Widget": widgetSetting})
}

func (wc widgetController) Update(context *admin.Context) {
	context.Resource = wc.Widgets.WidgetSettingResource
	widgetSetting, scopes, err := wc.getWidget(context)
	context.AddError(err)

	if context.AddError(context.Resource.Decode(context.Context, widgetSetting)); !context.HasError() {
		context.AddError(context.Resource.CallSave(widgetSetting, context.Context))
	}

	if context.HasError() {
		context.Writer.WriteHeader(admin.HTTPUnprocessableEntity)
		context.Execute("edit", map[string]interface{}{"Scopes": scopes, "Widget": widgetSetting})
	} else {
		http.Redirect(context.Writer, context.Request, context.Request.URL.Path, http.StatusFound)
	}
}

func (wc widgetController) InlineEdit(context *admin.Context) {
	context.Writer.Write([]byte(context.Render("inline_edit")))
}

func (wc widgetController) getWidget(context *admin.Context) (interface{}, []string, error) {
	if context.ResourceID == "" {
		// index page
		context.SetDB(context.GetDB().Where("scope = ?", "default"))
		results, err := context.FindMany()
		return results, []string{}, err
	}

	// show page
	result := wc.Widgets.WidgetSettingResource.NewStruct()
	scope := context.Request.URL.Query().Get("widget_scope")

	var scopes []string
	context.GetDB().Model(result).Where("name = ?", context.ResourceID).Pluck("scope", &scopes)

	var hasScope bool

	for _, s := range scopes {
		if scope == s {
			hasScope = true
			break
		}
	}

	if !hasScope {
		scope = "default"
	}

	err := context.GetDB().First(result, "name = ? AND scope = ?", context.ResourceID, scope).Error

	if widgetType := context.Request.URL.Query().Get("widget_type"); widgetType != "" {
		if serializableMeta, ok := result.(serializable_meta.SerializableMetaInterface); ok {
			serializableMeta.SetSerializableArgumentKind(widgetType)
			serializableMeta.SetSerializableArgumentValue(nil)
		}
	}
	return result, scopes, err
}
<|endoftext|>


<commit_before>package handlers

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/client"
	"github.com/openfaas/faas/gateway/metrics"
	"github.com/openfaas/faas/gateway/requests"
)

// MakeUpdateFunctionHandler request to update an existing function with new configuration such as image, envvars etc.
func MakeUpdateFunctionHandler(metricsOptions metrics.MetricOptions, c *client.Client, maxRestarts uint64, restartDelay time.Duration) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := context.Background()

		defer r.Body.Close()
		body, _ := ioutil.ReadAll(r.Body)

		request := requests.CreateFunctionRequest{}
		err := json.Unmarshal(body, &request)
		if err != nil {
			log.Println("Error parsing request:", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
			return
		}

		serviceInspectopts := types.ServiceInspectOptions{
			InsertDefaults: true,
		}

		service, _, err := c.ServiceInspectWithRaw(ctx, request.Service, serviceInspectopts)
		if err != nil {
			log.Println("Error inspecting service", err)
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte(err.Error()))
			return
		}

		updateSpec(&request, &service.Spec, maxRestarts, restartDelay)

		updateOpts := types.ServiceUpdateOptions{}
		updateOpts.RegistryAuthFrom = types.RegistryAuthFromSpec

		if len(request.RegistryAuth) > 0 {
			auth, err := BuildEncodedAuthConfig(request.RegistryAuth, request.Image)
			if err != nil {
				log.Println("Error building registry auth configuration:", err)
				w.WriteHeader(http.StatusBadRequest)
				w.Write([]byte("Invalid registry auth"))
				return
			}
			updateOpts.EncodedRegistryAuth = auth
		}

		response, err := c.ServiceUpdate(ctx, service.ID, service.Version, service.Spec, updateOpts)
		if err != nil {
			log.Println("Error updating service:", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Update error: " + err.Error()))
			return
		}
		log.Println(response.Warnings)
	}
}

func updateSpec(request *requests.CreateFunctionRequest, spec *swarm.ServiceSpec, maxRestarts uint64, restartDelay time.Duration) {

	constraints := []string{}
	if request.Constraints != nil && len(request.Constraints) > 0 {
		constraints = request.Constraints
	} else {
		constraints = linuxOnlyConstraints
	}

	spec.TaskTemplate.RestartPolicy.MaxAttempts = &maxRestarts
	spec.TaskTemplate.RestartPolicy.Condition = swarm.RestartPolicyConditionAny
	spec.TaskTemplate.RestartPolicy.Delay = &restartDelay
	spec.TaskTemplate.ContainerSpec.Image = request.Image
	spec.TaskTemplate.ContainerSpec.Labels = map[string]string{
		"function":              "true",
		"com.openfaas.function": request.Service,
		"com.openfaas.uid":      fmt.Sprintf("%d", time.Now().Nanosecond()),
	}

	if request.Labels != nil {
		for k, v := range *request.Labels {
			spec.TaskTemplate.ContainerSpec.Labels[k] = v
			spec.Annotations.Labels[k] = v
		}
	}

	spec.TaskTemplate.Networks = []swarm.NetworkAttachmentConfig{
		{
			Target: request.Network,
		},
	}

	spec.TaskTemplate.Resources = buildResources(request)

	spec.TaskTemplate.Placement = &swarm.Placement{
		Constraints: constraints,
	}

	spec.Annotations.Name = request.Service

	spec.RollbackConfig = &swarm.UpdateConfig{
		FailureAction: "pause",
	}

	spec.UpdateConfig = &swarm.UpdateConfig{
		Parallelism:   1,
		FailureAction: "rollback",
	}

	env := buildEnv(request.EnvProcess, request.EnvVars)

	if len(env) > 0 {
		spec.TaskTemplate.ContainerSpec.Env = env
	}

	if spec.Mode.Replicated != nil {
		spec.Mode.Replicated.Replicas = getMinReplicas(request)
	}
}
<commit_msg>Add secret management to the update handler<commit_after>package handlers

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/client"
	"github.com/openfaas/faas/gateway/metrics"
	"github.com/openfaas/faas/gateway/requests"
)

// MakeUpdateFunctionHandler request to update an existing function with new configuration such as image, envvars etc.
func MakeUpdateFunctionHandler(metricsOptions metrics.MetricOptions, c *client.Client, maxRestarts uint64, restartDelay time.Duration) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := context.Background()

		defer r.Body.Close()
		body, _ := ioutil.ReadAll(r.Body)

		request := requests.CreateFunctionRequest{}
		err := json.Unmarshal(body, &request)
		if err != nil {
			log.Println("Error parsing request:", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
			return
		}

		serviceInspectopts := types.ServiceInspectOptions{
			InsertDefaults: true,
		}

		service, _, err := c.ServiceInspectWithRaw(ctx, request.Service, serviceInspectopts)
		if err != nil {
			log.Println("Error inspecting service", err)
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte(err.Error()))
			return
		}

		secrets, err := makeSecretsArray(c, request.Secrets)
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Deployment error: " + err.Error()))
			return
		}
		updateSpec(&request, &service.Spec, maxRestarts, restartDelay, secrets)

		updateOpts := types.ServiceUpdateOptions{}
		updateOpts.RegistryAuthFrom = types.RegistryAuthFromSpec

		if len(request.RegistryAuth) > 0 {
			auth, err := BuildEncodedAuthConfig(request.RegistryAuth, request.Image)
			if err != nil {
				log.Println("Error building registry auth configuration:", err)
				w.WriteHeader(http.StatusBadRequest)
				w.Write([]byte("Invalid registry auth"))
				return
			}
			updateOpts.EncodedRegistryAuth = auth
		}

		response, err := c.ServiceUpdate(ctx, service.ID, service.Version, service.Spec, updateOpts)
		if err != nil {
			log.Println("Error updating service:", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Update error: " + err.Error()))
			return
		}
		log.Println(response.Warnings)
	}
}

func updateSpec(request *requests.CreateFunctionRequest, spec *swarm.ServiceSpec, maxRestarts uint64, restartDelay time.Duration, secrets []*swarm.SecretReference) {

	constraints := []string{}
	if request.Constraints != nil && len(request.Constraints) > 0 {
		constraints = request.Constraints
	} else {
		constraints = linuxOnlyConstraints
	}

	spec.TaskTemplate.RestartPolicy.MaxAttempts = &maxRestarts
	spec.TaskTemplate.RestartPolicy.Condition = swarm.RestartPolicyConditionAny
	spec.TaskTemplate.RestartPolicy.Delay = &restartDelay
	spec.TaskTemplate.ContainerSpec.Image = request.Image
	spec.TaskTemplate.ContainerSpec.Labels = map[string]string{
		"function":              "true",
		"com.openfaas.function": request.Service,
		"com.openfaas.uid":      fmt.Sprintf("%d", time.Now().Nanosecond()),
	}

	if request.Labels != nil {
		for k, v := range *request.Labels {
			spec.TaskTemplate.ContainerSpec.Labels[k] = v
			spec.Annotations.Labels[k] = v
		}
	}

	spec.TaskTemplate.Networks = []swarm.NetworkAttachmentConfig{
		{
			Target: request.Network,
		},
	}
	spec.TaskTemplate.ContainerSpec.Secrets = secrets

	spec.TaskTemplate.Resources = buildResources(request)

	spec.TaskTemplate.Placement = &swarm.Placement{
		Constraints: constraints,
	}

	spec.Annotations.Name = request.Service

	spec.RollbackConfig = &swarm.UpdateConfig{
		FailureAction: "pause",
	}

	spec.UpdateConfig = &swarm.UpdateConfig{
		Parallelism:   1,
		FailureAction: "rollback",
	}

	env := buildEnv(request.EnvProcess, request.EnvVars)

	if len(env) > 0 {
		spec.TaskTemplate.ContainerSpec.Env = env
	}

	if spec.Mode.Replicated != nil {
		spec.Mode.Replicated.Replicas = getMinReplicas(request)
	}
}
<|endoftext|>


<commit_before>package revok

import (
	"context"

	"cloud.google.com/go/trace"
	"code.cloudfoundry.org/lager"

	"cred-alert/db"
	"cred-alert/notifications"
)

//go:generate counterfeiter . NotificationComposerScanner

type NotificationComposerScanner interface {
	Scan(lager.Logger, string, string, map[string]struct{}, string, string, string) ([]db.Credential, error)
}

type NotificationComposer struct {
	repositoryRepository db.RepositoryRepository
	router               notifications.Router
	scanner              NotificationComposerScanner
}

func NewNotificationComposer(
	repositoryRepository db.RepositoryRepository,
	router notifications.Router,
	scanner NotificationComposerScanner,
) *NotificationComposer {
	return &NotificationComposer{
		repositoryRepository: repositoryRepository,
		router:               router,
		scanner:              scanner,
	}
}

func (n *NotificationComposer) ScanAndNotify(
	ctx context.Context,
	logger lager.Logger,
	owner string,
	repository string,
	scannedOids map[string]struct{},
	branch string,
	startSHA string,
	stopSHA string,
) error {

	dbRepository, err := n.repositoryRepository.MustFind(owner, repository)
	if err != nil {
		logger.Error("failed-to-find-db-repo", err)
		return err
	}

	scanSpan := trace.FromContext(ctx).NewChild("scanning")
	scanSpan.SetLabel("branch", branch)
	credentials, err := n.scanner.Scan(logger, owner, repository, scannedOids, branch, startSHA, stopSHA)
	if err != nil {
		return err
	}
	scanSpan.Finish()

	var batch []notifications.Notification

	for _, credential := range credentials {
		batch = append(batch, notifications.Notification{
			Owner:      credential.Owner,
			Repository: credential.Repository,
			Branch:     branch,
			SHA:        credential.SHA,
			Path:       credential.Path,
			LineNumber: credential.LineNumber,
			Private:    dbRepository.Private,
		})
		logger.Info("notifing-for-credential", lager.Data{
			"credential_owner":       credential.Owner,
			"credential_repository":  credential.Repository,
			"credential_branch":      branch,
			"credential_sha":         credential.SHA,
			"credential_path":        credential.Path,
			"credential_line_number": credential.LineNumber,
			"credential_private":     credential.Private,
		})
	}

	if len(batch) > 0 {
		err = n.router.Deliver(ctx, logger, batch)
		if err != nil {
			logger.Error("failed", err)
			return err
		}
	}

	return nil
}
<commit_msg>Fix typo in log message.<commit_after>package revok

import (
	"context"

	"cloud.google.com/go/trace"
	"code.cloudfoundry.org/lager"

	"cred-alert/db"
	"cred-alert/notifications"
)

//go:generate counterfeiter . NotificationComposerScanner

type NotificationComposerScanner interface {
	Scan(lager.Logger, string, string, map[string]struct{}, string, string, string) ([]db.Credential, error)
}

type NotificationComposer struct {
	repositoryRepository db.RepositoryRepository
	router               notifications.Router
	scanner              NotificationComposerScanner
}

func NewNotificationComposer(
	repositoryRepository db.RepositoryRepository,
	router notifications.Router,
	scanner NotificationComposerScanner,
) *NotificationComposer {
	return &NotificationComposer{
		repositoryRepository: repositoryRepository,
		router:               router,
		scanner:              scanner,
	}
}

func (n *NotificationComposer) ScanAndNotify(
	ctx context.Context,
	logger lager.Logger,
	owner string,
	repository string,
	scannedOids map[string]struct{},
	branch string,
	startSHA string,
	stopSHA string,
) error {

	dbRepository, err := n.repositoryRepository.MustFind(owner, repository)
	if err != nil {
		logger.Error("failed-to-find-db-repo", err)
		return err
	}

	scanSpan := trace.FromContext(ctx).NewChild("scanning")
	scanSpan.SetLabel("branch", branch)
	credentials, err := n.scanner.Scan(logger, owner, repository, scannedOids, branch, startSHA, stopSHA)
	if err != nil {
		return err
	}
	scanSpan.Finish()

	var batch []notifications.Notification

	for _, credential := range credentials {
		batch = append(batch, notifications.Notification{
			Owner:      credential.Owner,
			Repository: credential.Repository,
			Branch:     branch,
			SHA:        credential.SHA,
			Path:       credential.Path,
			LineNumber: credential.LineNumber,
			Private:    dbRepository.Private,
		})
		logger.Info("notifying-for-credential", lager.Data{
			"credential_owner":       credential.Owner,
			"credential_repository":  credential.Repository,
			"credential_branch":      branch,
			"credential_sha":         credential.SHA,
			"credential_path":        credential.Path,
			"credential_line_number": credential.LineNumber,
			"credential_private":     credential.Private,
		})
	}

	if len(batch) > 0 {
		err = n.router.Deliver(ctx, logger, batch)
		if err != nil {
			logger.Error("failed", err)
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package encoding

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"

	"github.com/yarpc/yab/encoding/encodingerror"
	"github.com/yarpc/yab/encoding/inputdecoder"
	"github.com/yarpc/yab/protobuf"
	"github.com/yarpc/yab/transport"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/jhump/protoreflect/desc"
	"github.com/jhump/protoreflect/dynamic"
	"go.uber.org/yarpc/pkg/procedure"
)

type protoSerializer struct {
	serviceName string
	methodName  string
	method      *desc.MethodDescriptor
	anyResolver jsonpb.AnyResolver

	reqDecoder inputdecoder.Decoder
}

// bytesMsg wraps a raw byte slice for serialization purposes. Especially
// useful for any types where we can't find the value in the registry.
type bytesMsg struct {
	V []byte
}

// anyResolver is a custom resolver which will use the descriptor provider
// and fallback to a simple byte array structure in the json encoding.
// This is the needed otherwise the default behaviour of protoreflect
// is to fail with an error when it can't find a type from protobuf.Any.
type anyResolver struct {
	source protobuf.DescriptorProvider
}

func (*bytesMsg) ProtoMessage()             {}
func (*bytesMsg) XXX_WellKnownType() string { return "BytesValue" }
func (m *bytesMsg) Reset()                  { *m = bytesMsg{} }
func (m *bytesMsg) String() string {
	return fmt.Sprintf("%x", m.V) // not compatible w/ pb oct
}
func (m *bytesMsg) Unmarshal(b []byte) error {
	m.V = append([]byte(nil), b...)
	return nil
}

func (r anyResolver) Resolve(typeUrl string) (proto.Message, error) {
	mname := typeUrl
	// This resolver supports types based on the spec at https://developers.google.com/protocol-buffers/docs/proto3#any
	// Example: type.googleapis.com/_packagename_._messagename_
	// so we want to get the fully qualified name of the type that we are dealing with
	// i.e. everything after the last '/'
	if slash := strings.LastIndex(mname, "/"); slash >= 0 {
		mname = mname[slash+1:]
	}

	msgDescriptor, err := r.source.FindMessage(mname)
	if err != nil {
		return nil, err
	}
	if msgDescriptor != nil {
		// We found a registered descriptor for our type, return it for human
		// readable format
		return dynamic.NewMessage(msgDescriptor), nil
	}
	// If me couldn't find the msg descriptor then provide a default implementation which will just
	// output the raw bytes as base64 - it's better than nothing.
	return &bytesMsg{}, nil
}

// NewProtobuf returns a protobuf serializer.
func NewProtobuf(fullMethodName string, source protobuf.DescriptorProvider, r io.Reader) (Serializer, error) {
	serviceName, methodName, err := splitMethod(fullMethodName)
	if err != nil {
		return nil, err
	}

	serviceDescriptor, err := source.FindService(serviceName)
	if err != nil {
		return nil, err
	}

	methodDescriptor, err := findProtoMethodDescriptor(serviceDescriptor, methodName)
	if err != nil {
		return nil, err
	}

	reqDecoder, err := inputdecoder.New(r)
	if err != nil {
		return nil, err
	}
	return &protoSerializer{
		serviceName: serviceName,
		methodName:  methodName,
		method:      methodDescriptor,
		anyResolver: anyResolver{
			source: source,
		},
		reqDecoder: reqDecoder,
	}, nil
}

func (p protoSerializer) Encoding() Encoding {
	return Protobuf
}

func (p protoSerializer) Request() (*transport.Request, error) {
	if p.method.IsClientStreaming() || p.method.IsServerStreaming() {
		return &transport.Request{
			Method: procedure.ToName(p.serviceName, p.methodName),
		}, nil
	}
	bytes, err := p.reqDecoder.Next()
	if err != nil && err != io.EOF {
		return nil, err
	}
	req := dynamic.NewMessage(p.method.GetInputType())
	if err := req.UnmarshalJSON(bytes); err != nil {
		return nil, fmt.Errorf("could not parse given request body as message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
	}
	bytes, err = proto.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("could marshal message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
	}
	return &transport.Request{
		Method: procedure.ToName(p.serviceName, p.methodName),
		Body:   bytes,
	}, nil
}

func (p protoSerializer) Response(body *transport.Response) (interface{}, error) {
	resp := dynamic.NewMessage(p.method.GetOutputType())
	if err := resp.Unmarshal(body.Body); err != nil {
		return nil, fmt.Errorf("could not parse given response body as message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
	}

	marshaler := &jsonpb.Marshaler{
		AnyResolver: p.anyResolver,
	}
	str, err := resp.MarshalJSONPB(marshaler)
	if err != nil {
		return nil, err
	}
	var unmarshaledJSON json.RawMessage
	if err = json.Unmarshal(str, &unmarshaledJSON); err != nil {
		return nil, err
	}
	return unmarshaledJSON, nil
}

func (p protoSerializer) CheckSuccess(body *transport.Response) error {
	_, err := p.Response(body)
	return err
}

func splitMethod(fullMethod string) (svc, method string, err error) {
	parts := strings.Split(fullMethod, "/")
	switch len(parts) {
	case 1:
		return parts[0], "", nil
	case 2:
		return parts[0], parts[1], nil
	default:
		return "", "", fmt.Errorf("invalid proto method %q, expected form package.Service/Method", fullMethod)
	}
}

func findProtoMethodDescriptor(s *desc.ServiceDescriptor, m string) (*desc.MethodDescriptor, error) {
	methodDescriptor := s.FindMethodByName(m)
	if methodDescriptor == nil {
		available := make([]string, len(s.GetMethods()))
		for i, method := range s.GetMethods() {
			available[i] = s.GetFullyQualifiedName() + "/" + method.GetName()
		}

		return nil, encodingerror.NotFound{
			Encoding:   "gRPC",
			SearchType: "method",
			Search:     m,
			LookIn:     fmt.Sprintf("service %q", s.GetFullyQualifiedName()),
			Example:    "--method package.Service/Method",
			Available:  available,
		}
	}
	return methodDescriptor, nil
}
<commit_msg>Remove stream code from this branch<commit_after>package encoding

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"

	"github.com/yarpc/yab/encoding/encodingerror"
	"github.com/yarpc/yab/encoding/inputdecoder"
	"github.com/yarpc/yab/protobuf"
	"github.com/yarpc/yab/transport"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/jhump/protoreflect/desc"
	"github.com/jhump/protoreflect/dynamic"
	"go.uber.org/yarpc/pkg/procedure"
)

type protoSerializer struct {
	serviceName string
	methodName  string
	method      *desc.MethodDescriptor
	anyResolver jsonpb.AnyResolver

	reqDecoder inputdecoder.Decoder
}

// bytesMsg wraps a raw byte slice for serialization purposes. Especially
// useful for any types where we can't find the value in the registry.
type bytesMsg struct {
	V []byte
}

// anyResolver is a custom resolver which will use the descriptor provider
// and fallback to a simple byte array structure in the json encoding.
// This is the needed otherwise the default behaviour of protoreflect
// is to fail with an error when it can't find a type from protobuf.Any.
type anyResolver struct {
	source protobuf.DescriptorProvider
}

func (*bytesMsg) ProtoMessage()             {}
func (*bytesMsg) XXX_WellKnownType() string { return "BytesValue" }
func (m *bytesMsg) Reset()                  { *m = bytesMsg{} }
func (m *bytesMsg) String() string {
	return fmt.Sprintf("%x", m.V) // not compatible w/ pb oct
}
func (m *bytesMsg) Unmarshal(b []byte) error {
	m.V = append([]byte(nil), b...)
	return nil
}

func (r anyResolver) Resolve(typeUrl string) (proto.Message, error) {
	mname := typeUrl
	// This resolver supports types based on the spec at https://developers.google.com/protocol-buffers/docs/proto3#any
	// Example: type.googleapis.com/_packagename_._messagename_
	// so we want to get the fully qualified name of the type that we are dealing with
	// i.e. everything after the last '/'
	if slash := strings.LastIndex(mname, "/"); slash >= 0 {
		mname = mname[slash+1:]
	}

	msgDescriptor, err := r.source.FindMessage(mname)
	if err != nil {
		return nil, err
	}
	if msgDescriptor != nil {
		// We found a registered descriptor for our type, return it for human
		// readable format
		return dynamic.NewMessage(msgDescriptor), nil
	}
	// If me couldn't find the msg descriptor then provide a default implementation which will just
	// output the raw bytes as base64 - it's better than nothing.
	return &bytesMsg{}, nil
}

// NewProtobuf returns a protobuf serializer.
func NewProtobuf(fullMethodName string, source protobuf.DescriptorProvider, r io.Reader) (Serializer, error) {
	serviceName, methodName, err := splitMethod(fullMethodName)
	if err != nil {
		return nil, err
	}

	serviceDescriptor, err := source.FindService(serviceName)
	if err != nil {
		return nil, err
	}

	methodDescriptor, err := findProtoMethodDescriptor(serviceDescriptor, methodName)
	if err != nil {
		return nil, err
	}

	reqDecoder, err := inputdecoder.New(r)
	if err != nil {
		return nil, err
	}
	return &protoSerializer{
		serviceName: serviceName,
		methodName:  methodName,
		method:      methodDescriptor,
		anyResolver: anyResolver{
			source: source,
		},
		reqDecoder: reqDecoder,
	}, nil
}

func (p protoSerializer) Encoding() Encoding {
	return Protobuf
}

func (p protoSerializer) Request() (*transport.Request, error) {
	bytes, err := p.reqDecoder.Next()
	if err != nil && err != io.EOF {
		return nil, err
	}
	req := dynamic.NewMessage(p.method.GetInputType())
	if err := req.UnmarshalJSON(bytes); err != nil {
		return nil, fmt.Errorf("could not parse given request body as message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
	}
	bytes, err = proto.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("could marshal message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
	}
	return &transport.Request{
		Method: procedure.ToName(p.serviceName, p.methodName),
		Body:   bytes,
	}, nil
}

func (p protoSerializer) Response(body *transport.Response) (interface{}, error) {
	resp := dynamic.NewMessage(p.method.GetOutputType())
	if err := resp.Unmarshal(body.Body); err != nil {
		return nil, fmt.Errorf("could not parse given response body as message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
	}

	marshaler := &jsonpb.Marshaler{
		AnyResolver: p.anyResolver,
	}
	str, err := resp.MarshalJSONPB(marshaler)
	if err != nil {
		return nil, err
	}
	var unmarshaledJSON json.RawMessage
	if err = json.Unmarshal(str, &unmarshaledJSON); err != nil {
		return nil, err
	}
	return unmarshaledJSON, nil
}

func (p protoSerializer) CheckSuccess(body *transport.Response) error {
	_, err := p.Response(body)
	return err
}

func splitMethod(fullMethod string) (svc, method string, err error) {
	parts := strings.Split(fullMethod, "/")
	switch len(parts) {
	case 1:
		return parts[0], "", nil
	case 2:
		return parts[0], parts[1], nil
	default:
		return "", "", fmt.Errorf("invalid proto method %q, expected form package.Service/Method", fullMethod)
	}
}

func findProtoMethodDescriptor(s *desc.ServiceDescriptor, m string) (*desc.MethodDescriptor, error) {
	methodDescriptor := s.FindMethodByName(m)
	if methodDescriptor == nil {
		available := make([]string, len(s.GetMethods()))
		for i, method := range s.GetMethods() {
			available[i] = s.GetFullyQualifiedName() + "/" + method.GetName()
		}

		return nil, encodingerror.NotFound{
			Encoding:   "gRPC",
			SearchType: "method",
			Search:     m,
			LookIn:     fmt.Sprintf("service %q", s.GetFullyQualifiedName()),
			Example:    "--method package.Service/Method",
			Available:  available,
		}
	}
	return methodDescriptor, nil
}
<|endoftext|>


<commit_before>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package image

var (
	// Black is an opaque black ColorImage.
	Black = ColorImage{Gray16Color{0}}
	// White is an opaque white ColorImage.
	White = ColorImage{Gray16Color{0xffff}}
	// Transparent is a fully transparent ColorImage.
	Transparent = ColorImage{Alpha16Color{0}}
	// Opaque is a fully opaque ColorImage.
	Opaque = ColorImage{Alpha16Color{0xffff}}
)

// A ColorImage is a practically infinite-sized Image of uniform Color.
// It implements both the Color and Image interfaces.
type ColorImage struct {
	C Color
}

func (c ColorImage) RGBA() (r, g, b, a uint32) {
	return c.C.RGBA()
}

func (c ColorImage) ColorModel() ColorModel {
	return ColorModelFunc(func(Color) Color { return c.C })
}

func (c ColorImage) Bounds() Rectangle { return Rectangle{ZP, Point{1e9, 1e9}} }

func (c ColorImage) At(x, y int) Color { return c.C }

// Opaque scans the entire image and returns whether or not it is fully opaque.
func (c ColorImage) Opaque() bool {
	_, _, _, a := c.C.RGBA()
	return a == 0xffff
}
<commit_msg>image: change a ColorImage's minimum point from (0, 0) to (-1e9, -1e9).<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package image

var (
	// Black is an opaque black ColorImage.
	Black = ColorImage{Gray16Color{0}}
	// White is an opaque white ColorImage.
	White = ColorImage{Gray16Color{0xffff}}
	// Transparent is a fully transparent ColorImage.
	Transparent = ColorImage{Alpha16Color{0}}
	// Opaque is a fully opaque ColorImage.
	Opaque = ColorImage{Alpha16Color{0xffff}}
)

// A ColorImage is a practically infinite-sized Image of uniform Color.
// It implements both the Color and Image interfaces.
type ColorImage struct {
	C Color
}

func (c ColorImage) RGBA() (r, g, b, a uint32) {
	return c.C.RGBA()
}

func (c ColorImage) ColorModel() ColorModel {
	return ColorModelFunc(func(Color) Color { return c.C })
}

func (c ColorImage) Bounds() Rectangle { return Rectangle{Point{-1e9, -1e9}, Point{1e9, 1e9}} }

func (c ColorImage) At(x, y int) Color { return c.C }

// Opaque scans the entire image and returns whether or not it is fully opaque.
func (c ColorImage) Opaque() bool {
	_, _, _, a := c.C.RGBA()
	return a == 0xffff
}
<|endoftext|>


<commit_before>package kd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"koding/newkite/kd/util"
	"koding/newkite/kodingkey"
	"log"
	"net/http"
	"time"
)

const KeyLength = 64

var (
	AuthServer      = "https://koding.com"
	AuthServerLocal = "http://localhost:3020"
)

type Register struct {
	authServer string
}

func NewRegister() *Register {
	return &Register{
		authServer: AuthServer,
	}
}

func (r *Register) Definition() string {
	return "Register this host to Koding"
}

func (r *Register) Exec(args []string) error {
	// change authServer address if debug mode is enabled
	if len(args) == 1 && (args[0] == "--debug" || args[0] == "-d") {
		r.authServer = AuthServerLocal
	}

	hostID, err := util.HostID()
	if err != nil {
		return err
	}

	var key string
	keyExist := false

	key, err = util.GetKey()
	if err != nil {
		k, err := kodingkey.NewKodingKey()
		if err != nil {
			return err
		}

		key = k.String()
	} else {
		fmt.Printf("Found a key under '%s'. Going to use it to register\n", KdPath)
		keyExist = true
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", r.authServer, hostID, key)

	fmt.Printf("Please open the following url for authentication:\n\n")
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	err = r.checker(key)
	if err != nil {
		return err
	}
	fmt.Println("successfully authenticated.")

	if keyExist {
		return nil
	}

	err = util.WriteKey(key)
	if err != nil {
		return err
	}

	return nil
}

// checker checks if the user has browsed the register URL by polling the check URL.
func (r *Register) checker(key string) error {
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", r.authServer, key)

	// check the result every two seconds
	ticker := time.NewTicker(2 * time.Second).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(3 * time.Minute)

	for {
		select {
		case <-ticker:
			err := checkResponse(checkUrl)
			if err != nil {
				// we didn't get OK message, continue until timout
				continue
			}

			return nil
		case <-timeout:
			return errors.New("timeout")
		}
	}
}

func checkResponse(checkUrl string) error {
	resp, err := http.Get(checkUrl)
	if err != nil {
		return err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	resp.Body.Close()

	fmt.Printf(". ") // animation

	if resp.StatusCode != 200 {
		return errors.New("non 200 response")
	}

	type Result struct {
		Result string `json:"result"`
	}

	res := Result{}
	err = json.Unmarshal(bytes.TrimSpace(body), &res)
	if err != nil {
		log.Fatalln(err) // this should not happen, exit here
	}

	return nil
}
<commit_msg>kd/register: fix path<commit_after>package kd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"koding/newkite/kd/util"
	"koding/newkite/kodingkey"
	"log"
	"net/http"
	"time"
)

const KeyLength = 64

var (
	AuthServer      = "https://koding.com"
	AuthServerLocal = "http://localhost:3020"
)

type Register struct {
	authServer string
}

func NewRegister() *Register {
	return &Register{
		authServer: AuthServer,
	}
}

func (r *Register) Definition() string {
	return "Register this host to Koding"
}

func (r *Register) Exec(args []string) error {
	// change authServer address if debug mode is enabled
	if len(args) == 1 && (args[0] == "--debug" || args[0] == "-d") {
		r.authServer = AuthServerLocal
	}

	hostID, err := util.HostID()
	if err != nil {
		return err
	}

	var key string
	keyExist := false

	key, err = util.GetKey()
	if err != nil {
		k, err := kodingkey.NewKodingKey()
		if err != nil {
			return err
		}

		key = k.String()
	} else {
		fmt.Printf("Found a key under '%s'. Going to use it to register\n", util.GetKdPath())
		keyExist = true
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", r.authServer, hostID, key)

	fmt.Printf("Please open the following url for authentication:\n\n")
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	err = r.checker(key)
	if err != nil {
		return err
	}
	fmt.Println("successfully authenticated.")

	if keyExist {
		return nil
	}

	err = util.WriteKey(key)
	if err != nil {
		return err
	}

	return nil
}

// checker checks if the user has browsed the register URL by polling the check URL.
func (r *Register) checker(key string) error {
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", r.authServer, key)

	// check the result every two seconds
	ticker := time.NewTicker(2 * time.Second).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(3 * time.Minute)

	for {
		select {
		case <-ticker:
			err := checkResponse(checkUrl)
			if err != nil {
				// we didn't get OK message, continue until timout
				continue
			}

			return nil
		case <-timeout:
			return errors.New("timeout")
		}
	}
}

func checkResponse(checkUrl string) error {
	resp, err := http.Get(checkUrl)
	if err != nil {
		return err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	resp.Body.Close()

	fmt.Printf(". ") // animation

	if resp.StatusCode != 200 {
		return errors.New("non 200 response")
	}

	type Result struct {
		Result string `json:"result"`
	}

	res := Result{}
	err = json.Unmarshal(bytes.TrimSpace(body), &res)
	if err != nil {
		log.Fatalln(err) // this should not happen, exit here
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"github.com/streadway/amqp"
	"koding/tools/amqputil"
	"log"
)

type Consumer struct {
	conn    *amqp.Connection
	channel *amqp.Channel
	tag     string
}

type Producer struct {
	conn    *amqp.Connection
	channel *amqp.Channel
}

type JoinMsg struct {
	Name               string  `json:"name"`
	BindingExchange    string  `json:"bindingExchange"`
	BindingKey         string  `json:"bindingKey"`
	PublishingExchange *string `json:"publishingExchange"`
	RoutingKey         string  `json:"routingKey"`
	ConsumerTag        string
	Suffix             string `json:"suffix"`
}

type LeaveMsg struct {
	RoutingKey string `json:"routingKey"`
}

var authPairs = make(map[string]JoinMsg)
var exchanges = make(map[string]uint)
var producer *Producer

func main() {
	log.Println("routing worker started")

	var err error
	producer, err = createProducer()
	if err != nil {
		log.Printf("create producer: %v", err)
	}

	startRouting()
}

func startRouting() {
	c := &Consumer{
		conn:    nil,
		channel: nil,
		tag:     "",
	}

	var err error

	log.Printf("creating consumer connections")
	c.conn = amqputil.CreateConnection("routing")
	c.channel = amqputil.CreateChannel(c.conn)

	err = c.channel.ExchangeDeclare("routing-control", "fanout", false, true, false, false, nil)
	if err != nil {
		log.Fatalf("exchange.declare: %s", err)
	}

	if _, err := c.channel.QueueDeclare("", false, true, false, false, nil); err != nil {
		log.Fatalf("queue.declare: %s", err)
	}

	if err := c.channel.QueueBind("", "", "routing-control", false, nil); err != nil {
		log.Fatalf("queue.bind: %s", err)
	}

	authStream, err := c.channel.Consume("", "", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %s", err)
	}

	log.Println("routing started...")
	for msg := range authStream {
		// used for debug
		// log.Printf("got %dB message data: [%v]-[%s] %s",
		// 	len(msg.Body),
		// 	msg.DeliveryTag,
		// 	msg.RoutingKey,
		// 	msg.Body,
		// )

		switch msg.RoutingKey {
		case "auth.join":
			var join JoinMsg
			err := json.Unmarshal(msg.Body, &join)
			if err != nil {
				log.Print("bad json incoming msg: ", err)
			}

			var publishingExchange string
			if join.PublishingExchange != nil {
				publishingExchange = *join.PublishingExchange
			} else {
				publishingExchange = "broker"
			}

			join.ConsumerTag = generateUniqueConsumerTag(join.BindingKey)
			authPairs[join.RoutingKey] = join

			// log.Printf("Auth pairs: %+v\n", authPairs) // this is just for debug

			declareExchange(c, join.BindingExchange)

			errors := make(chan error)

			go consumeAndRepublish(c,
				join.BindingExchange,
				join.BindingKey,
				publishingExchange,
				join.RoutingKey,
				join.Suffix,
				join.ConsumerTag,
				errors,
			)
			// select {
			// case errors <- err:
			// 	log.Printf("Handled an error: %v", err)
			// }
		case "auth.leave":
			var leave LeaveMsg
			err := json.Unmarshal(msg.Body, &leave)
			if err != nil {
				log.Print("bad json incoming msg: ", err)
			}

			// auth.leave could be received before auth.join, thus check it
			msg, ok := authPairs[leave.RoutingKey]
			if !ok {
				log.Printf("no tag available for %s\n", leave.RoutingKey)
				continue // do not close our main for clause 'authStream'
			}

			err = c.channel.Cancel(msg.ConsumerTag, false)
			if err != nil {
				log.Fatalf("basic.cancel: %s", err)
			}

			decrementExchangeCounter(leave.RoutingKey)
		default:
			log.Println("routing key is not defined: ", msg.RoutingKey)
		}
	}
}

func generateUniqueConsumerTag(bindingKey string) string {
	r := make([]byte, 32/8)
	rand.Read(r)
	return bindingKey + "." + base64.StdEncoding.EncodeToString(r)
}

func generateUniqueQueueName() string {
	r := make([]byte, 32/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

func declareExchange(c *Consumer, exchange string) {
	if exchanges[exchange] <= 0 {
		if err := c.channel.ExchangeDeclare(exchange, "topic", false, true, false, false, nil); err != nil {
			log.Fatalf("exchange.declare: %s", err)
		}
		exchanges[exchange] = 0
	}
	exchanges[exchange]++
}

func decrementExchangeCounter(routingKey string) {
	exchange := authPairs[routingKey].BindingExchange
	exchanges[exchange]--
	delete(authPairs, routingKey)
}

func consumeAndRepublish(
	c *Consumer,
	bindingExchange,
	bindingKey,
	publishingExchange,
	routingKey,
	suffix,
	consumerTag string,
	done chan error,
) {
	// used fo debug
	// log.Printf("Consume from:\n bindingExchange %s\n bindingKey %s\n routingKey %s\n consumerTag %s\n",
	// 	bindingExchange, bindingKey, routingKey, consumerTag)

	if len(suffix) > 0 {
		routingKey += suffix
	}

	channel, err := c.conn.Channel()
	if err != nil {
		done <- err
		return
	}
	defer channel.Close()

	uniqueQueueName := generateUniqueQueueName()

	if _, err := channel.QueueDeclare(uniqueQueueName, false, true, true, false, nil); err != nil {
		log.Fatalf("queue.declare: %s", err)
	}

	if err := channel.QueueBind(uniqueQueueName, bindingKey, bindingExchange, false, nil); err != nil {
		log.Fatalf("queue.bind: %s", err)
	}

	messages, err := channel.Consume(uniqueQueueName, consumerTag, true, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %s", err)
	}

	for msg := range messages {
		// used for debug
		// log.Printf("messages stream got %dB message data: [%v] %s",
		// 	len(msg.Body),
		// 	msg.DeliveryTag,
		// 	msg.Body,
		// )

		publishTo(publishingExchange, routingKey, msg.Body)
	}

}

func publishTo(exchange, routingKey string, data []byte) {
	msg := amqp.Publishing{
		Headers:         amqp.Table{},
		ContentType:     "text/plain",
		ContentEncoding: "",
		Body:            data,
		DeliveryMode:    1, // 1=non-persistent, 2=persistent
		Priority:        0, // 0-9
	}

	// used for debug
	// log.Println("publishing data ", string(data), routingKey)
	err := producer.channel.Publish(exchange, routingKey, false, false, msg)
	if err != nil {
		log.Printf("error while publishing proxy message: %s", err)
	}

}

func createProducer() (*Producer, error) {
	p := &Producer{
		conn:    nil,
		channel: nil,
	}

	log.Printf("creating publisher connections")

	p.conn = amqputil.CreateConnection("routing")
	p.channel = amqputil.CreateChannel(p.conn)

	return p, nil
}
<commit_msg>fix channel canceling and closing by storing a reference to each individual channel<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"github.com/streadway/amqp"
	"koding/tools/amqputil"
	"log"
)

type Consumer struct {
	conn    *amqp.Connection
	channel *amqp.Channel
	tag     string
}

type Producer struct {
	conn    *amqp.Connection
	channel *amqp.Channel
}

type JoinMsg struct {
	Name               string  `json:"name"`
	BindingExchange    string  `json:"bindingExchange"`
	BindingKey         string  `json:"bindingKey"`
	PublishingExchange *string `json:"publishingExchange"`
	RoutingKey         string  `json:"routingKey"`
	ConsumerTag        string
	Suffix             string `json:"suffix"`
	Channel            *amqp.Channel
}

type LeaveMsg struct {
	RoutingKey string `json:"routingKey"`
}

var authPairs = make(map[string]JoinMsg)
var exchanges = make(map[string]uint)
var producer *Producer

func main() {
	log.Println("routing worker started")

	var err error
	producer, err = createProducer()
	if err != nil {
		log.Printf("create producer: %v", err)
	}

	startRouting()
}

func startRouting() {
	c := &Consumer{
		conn:    nil,
		channel: nil,
		tag:     "",
	}

	var err error

	log.Printf("creating consumer connections")
	c.conn = amqputil.CreateConnection("routing")
	c.channel = amqputil.CreateChannel(c.conn)

	err = c.channel.ExchangeDeclare("routing-control", "fanout", false, true, false, false, nil)
	if err != nil {
		log.Fatalf("exchange.declare: %s", err)
	}

	if _, err := c.channel.QueueDeclare("", false, true, false, false, nil); err != nil {
		log.Fatalf("queue.declare: %s", err)
	}

	if err := c.channel.QueueBind("", "", "routing-control", false, nil); err != nil {
		log.Fatalf("queue.bind: %s", err)
	}

	authStream, err := c.channel.Consume("", "", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %s", err)
	}

	log.Println("routing started...")
	for msg := range authStream {
		// used for debug
		// log.Printf("got %dB message data: [%v]-[%s] %s",
		// 	len(msg.Body),
		// 	msg.DeliveryTag,
		// 	msg.RoutingKey,
		// 	msg.Body,
		// )

		switch msg.RoutingKey {
		case "auth.join":
			var join JoinMsg
			err := json.Unmarshal(msg.Body, &join)
			if err != nil {
				log.Print("bad json incoming msg: ", err)
			}

			var publishingExchange string
			if join.PublishingExchange != nil {
				publishingExchange = *join.PublishingExchange
			} else {
				publishingExchange = "broker"
			}

			join.ConsumerTag = generateUniqueConsumerTag(join.BindingKey)
			authPairs[join.RoutingKey] = join

			// log.Printf("Auth pairs: %+v\n", authPairs) // this is just for debug

			declareExchange(c, join.BindingExchange)

			errors := make(chan error)

			go consumeAndRepublish(c,
				join.BindingExchange,
				join.BindingKey,
				publishingExchange,
				join.RoutingKey,
				join.Suffix,
				join.ConsumerTag,
				errors,
			)
			// select {
			// case errors <- err:
			// 	log.Printf("Handled an error: %v", err)
			// }
		case "auth.leave":
			var leave LeaveMsg
			err := json.Unmarshal(msg.Body, &leave)
			if err != nil {
				log.Print("bad json incoming msg: ", err)
			}

			// auth.leave could be received before auth.join, thus check it
			msg, ok := authPairs[leave.RoutingKey]
			if !ok {
				log.Printf("no tag available for %s\n", leave.RoutingKey)
				continue // do not close our main for clause 'authStream'
			}

			err = msg.Channel.Cancel(msg.ConsumerTag, false)
			if err != nil {
				log.Fatalf("basic.cancel: %s", err)
			}

			err = msg.Channel.Close()
			if err != nil {
				log.Fatalf("channel.close: %s", err)
			}

			decrementExchangeCounter(leave.RoutingKey)
		default:
			log.Println("routing key is not defined: ", msg.RoutingKey)
		}
	}
}

func generateUniqueConsumerTag(bindingKey string) string {
	r := make([]byte, 32/8)
	rand.Read(r)
	return bindingKey + "." + base64.StdEncoding.EncodeToString(r)
}

func generateUniqueQueueName() string {
	r := make([]byte, 32/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

func declareExchange(c *Consumer, exchange string) {
	if exchanges[exchange] <= 0 {
		if err := c.channel.ExchangeDeclare(exchange, "topic", false, true, false, false, nil); err != nil {
			log.Fatalf("exchange.declare: %s", err)
		}
		exchanges[exchange] = 0
	}
	exchanges[exchange]++
}

func decrementExchangeCounter(routingKey string) {
	exchange := authPairs[routingKey].BindingExchange
	exchanges[exchange]--
	delete(authPairs, routingKey)
}

func consumeAndRepublish(
	c *Consumer,
	bindingExchange,
	bindingKey,
	publishingExchange,
	routingKey,
	suffix,
	consumerTag string,
	done chan error,
) {
	// used fo debug
	// log.Printf("Consume from:\n bindingExchange %s\n bindingKey %s\n routingKey %s\n consumerTag %s\n",
	// 	bindingExchange, bindingKey, routingKey, consumerTag)

	if len(suffix) > 0 {
		routingKey += suffix
	}

	channel, err := c.conn.Channel()
	if err != nil {
		done <- err
		return
	}

	// we need to extract join because GO doesn't assignments in form
	// authPairs[routingKey].Channel = channel
	join := authPairs[routingKey]
	join.Channel = channel
	authPairs[routingKey] = join

	uniqueQueueName := generateUniqueQueueName()

	if _, err := channel.QueueDeclare(uniqueQueueName, false, true, true, false, nil); err != nil {
		log.Fatalf("queue.declare: %s", err)
	}

	if err := channel.QueueBind(uniqueQueueName, bindingKey, bindingExchange, false, nil); err != nil {
		log.Fatalf("queue.bind: %s", err)
	}

	messages, err := channel.Consume(uniqueQueueName, consumerTag, true, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %s", err)
	}

	for msg := range messages {
		// used for debug
		// log.Printf("messages stream got %dB message data: [%v] %s",
		// 	len(msg.Body),
		// 	msg.DeliveryTag,
		// 	msg.Body,
		// )

		publishTo(publishingExchange, routingKey, msg.Body)
	}

}

func publishTo(exchange, routingKey string, data []byte) {
	msg := amqp.Publishing{
		Headers:         amqp.Table{},
		ContentType:     "text/plain",
		ContentEncoding: "",
		Body:            data,
		DeliveryMode:    1, // 1=non-persistent, 2=persistent
		Priority:        0, // 0-9
	}

	// used for debug
	// log.Println("publishing data ", string(data), routingKey)
	err := producer.channel.Publish(exchange, routingKey, false, false, msg)
	if err != nil {
		log.Printf("error while publishing proxy message: %s", err)
	}

}

func createProducer() (*Producer, error) {
	p := &Producer{
		conn:    nil,
		channel: nil,
	}

	log.Printf("creating publisher connections")

	p.conn = amqputil.CreateConnection("routing")
	p.channel = amqputil.CreateChannel(p.conn)

	return p, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package test

import (
	"time"

	"v.io/jiri/collect"
	"v.io/jiri/jiri"
	"v.io/jiri/runutil"
	"v.io/x/devtools/internal/test"
	"v.io/x/lib/envvar"
)

// runMakefileTest is a helper for running tests through make commands.
func runMakefileTest(jirix *jiri.X, testName, testDir, target string, env map[string]string, profiles []string, timeout time.Duration) (_ *test.Result, e error) {
	// Install base profile first, before any test-specific profiles.
	profiles = append([]string{"base"}, profiles...)

	// Initialize the test.
	cleanup, err := initTest(jirix, testName, profiles)
	if err != nil {
		return nil, newInternalError(err, "Init")
	}
	defer collect.Error(func() error { return cleanup() }, &e)

	s := jirix.NewSeq()

	// Set up the environment
	merged := envvar.MergeMaps(jirix.Env(), env)

	// Navigate to project directory, run make clean and make target.
	err = s.Pushd(testDir).
		Run("make", "clean").
		Timeout(timeout).Env(merged).Last("make", target)
	if err != nil {
		if runutil.IsTimeout(err) {
			return &test.Result{
				Status:       test.TimedOut,
				TimeoutValue: timeout,
			}, nil
		} else {
			return nil, newInternalError(err, "Make "+target)
		}
	}

	return &test.Result{Status: test.Passed}, nil
}
<commit_msg>TBR: Always print output from makefile tests.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package test

import (
	"time"

	"v.io/jiri/collect"
	"v.io/jiri/jiri"
	"v.io/jiri/runutil"
	"v.io/x/devtools/internal/test"
	"v.io/x/lib/envvar"
)

// runMakefileTest is a helper for running tests through make commands.
func runMakefileTest(jirix *jiri.X, testName, testDir, target string, env map[string]string, profiles []string, timeout time.Duration) (_ *test.Result, e error) {
	// Install base profile first, before any test-specific profiles.
	profiles = append([]string{"base"}, profiles...)

	// Initialize the test.
	cleanup, err := initTest(jirix, testName, profiles)
	if err != nil {
		return nil, newInternalError(err, "Init")
	}
	defer collect.Error(func() error { return cleanup() }, &e)

	s := jirix.NewSeq()

	// Set up the environment
	merged := envvar.MergeMaps(jirix.Env(), env)

	// Navigate to project directory, run make clean and make target.
	err = s.Pushd(testDir).
		Verbose(true).
		Run("make", "clean").
		Verbose(true).
		Timeout(timeout).Env(merged).Last("make", target)
	if err != nil {
		if runutil.IsTimeout(err) {
			return &test.Result{
				Status:       test.TimedOut,
				TimeoutValue: timeout,
			}, nil
		} else {
			return nil, newInternalError(err, "Make "+target)
		}
	}

	return &test.Result{Status: test.Passed}, nil
}
<|endoftext|>