<commit_before>package sphere

// IChannels is the interface for ChannelModel
type IChannels interface {
	Name() string
	Subscribe(uri string, connection *Connection) bool
	Disconnect(uri string, connection *Connection)
	Receive(event string, message string) error
}

// ExtendChannelModel lets developer create a IChannals compatible struct
func ExtendChannelModel(name string) *ChannelModel {
	return &ChannelModel{name}
}

// ChannelModel is for user to define channel events and actions
type ChannelModel struct {
	name string
}

// Name to return name of the channel
func (m *ChannelModel) Name() string {
	return m.name
}

// Subscribe decides whether accept the connection into channel or not, return true => accept, false => reject
func (m *ChannelModel) Subscribe(uri string, connection *Connection) bool {
	return false
}

// Disconnect defines the action when user disconnect from channel
func (m *ChannelModel) Disconnect(uri string, connection *Connection) {
}

// Receive defines the action when websocket server receive message from user in this channel
func (m *ChannelModel) Receive(event string, message string) error {
	return nil
}
<commit_msg>rename Name to Namespace<commit_after>package sphere

// IChannels is the interface for ChannelModel
type IChannels interface {
	Namespace() string
	Subscribe(uri string, connection *Connection) bool
	Disconnect(uri string, connection *Connection)
	Receive(event string, message string) error
}

// ExtendChannelModel lets developer create a IChannals compatible struct
func ExtendChannelModel(namespace string) *ChannelModel {
	return &ChannelModel{namespace}
}

// ChannelModel is for user to define channel events and actions
type ChannelModel struct {
	namespace string
}

// Namespace to return name of the channel
func (m *ChannelModel) Namespace() string {
	return m.namespace
}

// Subscribe decides whether accept the connection into channel or not, return true => accept, false => reject
func (m *ChannelModel) Subscribe(uri string, connection *Connection) bool {
	return false
}

// Disconnect defines the action when user disconnect from channel
func (m *ChannelModel) Disconnect(uri string, connection *Connection) {
}

// Receive defines the action when websocket server receive message from user in this channel
func (m *ChannelModel) Receive(event string, message string) error {
	return nil
}
<|endoftext|>


<commit_before>package clocky

import (
	"fmt"
	"http"
	"template"
	"time"

	"solar"
)

const Lat, Lng = 37.79, -122.42

var tmpl = template.Must(template.New("page").Parse(page))

// Pacify converts utc to US Pacific time (2007 rules).  We have to do
// this by hand because Go r60 doesn't have any real time zone
// support.  Things are better in Go 1.
func Pacify(utc *time.Time) *time.Time {
	// Find the second Sunday in March and the first Sunday in
	// November.  The second Sunday in March is the first Sunday
	// that is or follows March 8.
	mar8, _ := time.Parse("2006-01-02 15", fmt.Sprintf("%d-03-08 10", utc.Year))
	dstStart := mar8.Seconds() + int64((7-weekday(mar8))%7*86400)
	nov1, _ := time.Parse("2006-01-02 15", fmt.Sprintf("%d-11-01 09", utc.Year))
	dstEnd := nov1.Seconds() + int64((7-weekday(nov1))%7*86400)

	offset, zone := -8*3600, "PST"
	if utc.Seconds() >= dstStart && utc.Seconds() < dstEnd {
		offset, zone = -7*3600, "PDT"
	}
	local := time.SecondsToUTC(utc.Seconds() + int64(offset))
	local.ZoneOffset = offset
	local.Zone = zone
	return local
}

var sakamotoTable = []int{0, 0, 3, 2, 5, 0, 3, 5, 1, 4, 6, 2, 4}

// weekday calculates the day of a week using Sakamoto's method.
func weekday(t *time.Time) int {
	y := int(t.Year) // This algorithm won't work for years >= 2**31 anyway.
	if t.Month < 3 {
		y--
	}
	return (y + y/4 - y/100 + y/400 + sakamotoTable[t.Month] + t.Day) % 7
}

func handler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}

	// TODO: Refresh minutely; use JS to fake things.
	// TODO: Have browser refresh; safer since error pages will get retried.
	w.Header().Set("Refresh", "2")

	now := Pacify(time.UTC())

	sunrise := Pacify(solar.Rise(now, Lat, Lng))
	sunset := Pacify(solar.Set(now, Lat, Lng))
	sun1 := "sunrise " + sunrise.Format("3:04 pm")
	sun2 := "sunset " + sunset.Format("3:04 pm")
	if sunrise.Seconds() > sunset.Seconds() {
		sun1, sun2 = sun2, sun1
	}

	d := map[string]map[string]string{
		"Time": map[string]string{
			"Big":   now.Format("3:04"),
			"Small": now.Format(":05 pm"),
			"Date":  now.Format("Monday, January 2"),
			"Sun1":  sun1,
			"Sun2":  sun2,
		},
		"Weather": map[string]string{
			"Temp":     "12°",
			"Forecast": dummyForecast,
		},
	}
	tmpl.Execute(w, d)
}

func init() {
	http.HandleFunc("/", handler)
}

const page = `<!DOCTYPE html>

    <title>Clocky</title>

{{/* Everything temporarily down 50px to avoid blinking status bar in Boat. */}}
{{with .Time}}

{{.Big}}{{.Small}}

{{.Date}}

{{.Sun1}}, {{.Sun2}}

{{end}}

{{with .Weather}}

{{.Temp}} calm, 96%

{{.Forecast}}

{{end}}

47 outbound

11, 30, 50, 68 minutes

49 outbound

½, 19, 39, 59 minutes

10, 12 outbound

18 minutes

27 outbound

Probably never

1 inbound

6½, 31, 51, 69 minutes

1 outbound

Now, 41, 59, 79 minutes

`

// km/h, am, pm after number: convert no space or ASCII space to
// line-ending number: change ASCII space to
const dummyForecast = `

Tonight: Patchy fog after
10 pm. Otherwise, mostly cloudy, with a low
around 9. Northwest wind around 10 km/h becoming
calm.

Saturday:
Patchy fog before 10 am. Otherwise, mostly sunny, with a high
near 16. North northeast wind between 10 and 13 km/h
becoming calm.
`
<commit_msg>Lowercase "now".<commit_after>package clocky

import (
	"fmt"
	"http"
	"template"
	"time"

	"solar"
)

const Lat, Lng = 37.79, -122.42

var tmpl = template.Must(template.New("page").Parse(page))

// Pacify converts utc to US Pacific time (2007 rules).  We have to do
// this by hand because Go r60 doesn't have any real time zone
// support.  Things are better in Go 1.
func Pacify(utc *time.Time) *time.Time {
	// Find the second Sunday in March and the first Sunday in
	// November.  The second Sunday in March is the first Sunday
	// that is or follows March 8.
	mar8, _ := time.Parse("2006-01-02 15", fmt.Sprintf("%d-03-08 10", utc.Year))
	dstStart := mar8.Seconds() + int64((7-weekday(mar8))%7*86400)
	nov1, _ := time.Parse("2006-01-02 15", fmt.Sprintf("%d-11-01 09", utc.Year))
	dstEnd := nov1.Seconds() + int64((7-weekday(nov1))%7*86400)

	offset, zone := -8*3600, "PST"
	if utc.Seconds() >= dstStart && utc.Seconds() < dstEnd {
		offset, zone = -7*3600, "PDT"
	}
	local := time.SecondsToUTC(utc.Seconds() + int64(offset))
	local.ZoneOffset = offset
	local.Zone = zone
	return local
}

var sakamotoTable = []int{0, 0, 3, 2, 5, 0, 3, 5, 1, 4, 6, 2, 4}

// weekday calculates the day of a week using Sakamoto's method.
func weekday(t *time.Time) int {
	y := int(t.Year) // This algorithm won't work for years >= 2**31 anyway.
	if t.Month < 3 {
		y--
	}
	return (y + y/4 - y/100 + y/400 + sakamotoTable[t.Month] + t.Day) % 7
}

func handler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}

	// TODO: Refresh minutely; use JS to fake things.
	// TODO: Have browser refresh; safer since error pages will get retried.
	w.Header().Set("Refresh", "2")

	now := Pacify(time.UTC())

	sunrise := Pacify(solar.Rise(now, Lat, Lng))
	sunset := Pacify(solar.Set(now, Lat, Lng))
	sun1 := "sunrise " + sunrise.Format("3:04 pm")
	sun2 := "sunset " + sunset.Format("3:04 pm")
	if sunrise.Seconds() > sunset.Seconds() {
		sun1, sun2 = sun2, sun1
	}

	d := map[string]map[string]string{
		"Time": map[string]string{
			"Big":   now.Format("3:04"),
			"Small": now.Format(":05 pm"),
			"Date":  now.Format("Monday, January 2"),
			"Sun1":  sun1,
			"Sun2":  sun2,
		},
		"Weather": map[string]string{
			"Temp":     "12°",
			"Forecast": dummyForecast,
		},
	}
	tmpl.Execute(w, d)
}

func init() {
	http.HandleFunc("/", handler)
}

const page = `<!DOCTYPE html>

    <title>Clocky</title>

{{/* Everything temporarily down 50px to avoid blinking status bar in Boat. */}}
{{with .Time}}

{{.Big}}{{.Small}}

{{.Date}}

{{.Sun1}}, {{.Sun2}}

{{end}}

{{with .Weather}}

{{.Temp}} calm, 96%

{{.Forecast}}

{{end}}

47 outbound

11, 30, 50, 68 minutes

49 outbound

½, 19, 39, 59 minutes

10, 12 outbound

18 minutes

27 outbound

Probably never

1 inbound

6½, 31, 51, 69 minutes

1 outbound

now, 41, 59, 79 minutes

`

// km/h, am, pm after number: convert no space or ASCII space to
// line-ending number: change ASCII space to
const dummyForecast = `

Tonight: Patchy fog after
10 pm. Otherwise, mostly cloudy, with a low
around 9. Northwest wind around 10 km/h becoming
calm.

Saturday:
Patchy fog before 10 am. Otherwise, mostly sunny, with a high
near 16. North northeast wind between 10 and 13 km/h
becoming calm.
`
<|endoftext|>


<commit_before>package multicast

import (
	"bytes"
	"context"
	"sync"
	"testing"
	"time"
)

var testMulticastAddress = "[ff12::9000]:30001"

func TestListen(t *testing.T) {
	listener := &Multicast{Address: testMulticastAddress}
	respCh, err := listener.Listen(context.TODO())
	if err != nil {
		t.Fatal(err)
	}

	sender := &Multicast{Address: testMulticastAddress}
	if err := sender.Send(context.TODO(), []byte{}); err != nil {
		t.Fatal("Multicast Send should fail", err)
	}

	for {
		select {
		case <-respCh:
			listener.Done()
			sender.Done()
			return
		}
	}

}

func TestSend(t *testing.T) {
	var tests = []struct {
		m         *Multicast
		payload   []byte
		shouldErr bool
	}{
		{&Multicast{}, nil, true},
		{&Multicast{Address: testMulticastAddress}, []byte("Hello TestSendAndListen"), false},
		{&Multicast{Address: testMulticastAddress}, []byte("Say hello again"), false},
		{&Multicast{Address: testMulticastAddress}, []byte("123412341234"), false},
	}

	var found = []struct {
		m         *Multicast
		payload   []byte
		shouldErr bool
	}{}

	ctx, cancel := context.WithCancel(context.Background())
	wg := &sync.WaitGroup{}
	errCh := make(chan error, 1)

	// Don't block for test
	go func() {
		listener := &Multicast{Address: testMulticastAddress}

		// Create a listener
		respCh, err := listener.Listen(ctx)
		if err != nil {
			errCh <- err
		}

		// Check if the resp sent over the response channel matches the test's payload
		go func() {
			for {
				select {
				case resp := <-respCh:
					for _, test := range tests {
						skip := false
						// Check the slice of test to see if the response equals what is expected
						if bytes.Equal(resp.Payload, test.payload) {
							// Skip if it's already been found
							for _, f := range found {
								if bytes.Equal(resp.Payload, f.payload) {
									skip = true
								}
							}

							if skip {
								continue
							}

							found = append(found, test)
							test.m.Stop()
							wg.Done()
						}
					}
				case <-ctx.Done():
					return
				case <-listener.Done():
					return
				}
			}
		}()

		for _, test := range tests {
			if !test.shouldErr {
				wg.Add(1)
				if err := test.m.Send(ctx, test.payload); err != nil {
					t.Fatal("Multicast Send failed", err)
				}

			} else {
				if err := test.m.Send(ctx, test.payload); err == nil {
					t.Fatal("Multicast Send should fail", err)
				}
			}
		}

		wg.Wait() // Block until all test multicasts are stopped
		listener.Stop()
		cancel()
	}()

	// Block until the ctx is canceled or we receive an error, such as a timeout
	for {
		select {
		case <-ctx.Done():
			return
		case err := <-errCh:
			if err != nil {
				t.Fatal("err during Send()", err)
			}
			return
		case <-time.After(5000 * time.Millisecond):
			t.Fatal("	 timed out")
			return
		}
	}
}

func TestListenCtxDone(t *testing.T) {
	// Create a listener
	ctx, cancel := context.WithCancel(context.Background())
	listener := &Multicast{Address: testMulticastAddress}
	_, err := listener.Listen(ctx)
	if err != nil {
		t.Fatal(err)
	}

	// Test context Done()
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			}
		}
	}()
	cancel()
}

func TestListenStop(t *testing.T) {
	// Create a listener
	listener := &Multicast{Address: testMulticastAddress}
	_, err := listener.Listen(context.TODO())
	if err != nil {
		t.Fatal(err)
	}

	// Test listener Done()
	go func() {
		for {
			select {
			case <-listener.Done():
				return
			}
		}
	}()
	listener.Stop()
}

func TestCtxCancelFunc(t *testing.T) {
	m := &Multicast{Address: testMulticastAddress}

	payload := []byte("Hello TestCtxCancelFunc")
	ctx, cancel := context.WithCancel(context.Background())
	if err := m.Send(ctx, payload); err != nil {
		t.Fatal("Send error", err)
	}
	time.AfterFunc(100*time.Millisecond, func() { cancel() })
	timeout := time.AfterFunc(200*time.Millisecond, func() { t.Fatal("TestStopChan timedout") })

	// Block until the stopCh is closed
	for {
		select {
		case <-ctx.Done():
			timeout.Stop() // cancel the timeout
			return
		}
	}
}
<commit_msg>Updating multicast_test to avoid panic nil byte slice in ipv6.NewPacketConn WriteTo causes panic bug<commit_after>package multicast

import (
	"bytes"
	"context"
	"sync"
	"testing"
	"time"
)

var testMulticastAddress = "[ff12::9000]:30001"

func TestListen(t *testing.T) {
	listener := &Multicast{Address: testMulticastAddress}
	respCh, err := listener.Listen(context.TODO())
	if err != nil {
		t.Fatal(err)
	}

	sender := &Multicast{Address: testMulticastAddress}
	if err := sender.Send(context.TODO(), []byte{0}); err != nil {
		t.Fatal("Multicast Send should fail", err)
	}

	for {
		select {
		case <-respCh:
			listener.Done()
			sender.Done()
			return
		}
	}

}

func TestSend(t *testing.T) {
	var tests = []struct {
		m         *Multicast
		payload   []byte
		shouldErr bool
	}{
		{&Multicast{}, nil, true},
		{&Multicast{Address: testMulticastAddress}, []byte("Hello TestSendAndListen"), false},
		{&Multicast{Address: testMulticastAddress}, []byte("Say hello again"), false},
		{&Multicast{Address: testMulticastAddress}, []byte("123412341234"), false},
	}

	var found = []struct {
		m         *Multicast
		payload   []byte
		shouldErr bool
	}{}

	ctx, cancel := context.WithCancel(context.Background())
	wg := &sync.WaitGroup{}
	errCh := make(chan error, 1)

	// Don't block for test
	go func() {
		listener := &Multicast{Address: testMulticastAddress}

		// Create a listener
		respCh, err := listener.Listen(ctx)
		if err != nil {
			errCh <- err
		}

		// Check if the resp sent over the response channel matches the test's payload
		go func() {
			for {
				select {
				case resp := <-respCh:
					for _, test := range tests {
						skip := false
						// Check the slice of test to see if the response equals what is expected
						if bytes.Equal(resp.Payload, test.payload) {
							// Skip if it's already been found
							for _, f := range found {
								if bytes.Equal(resp.Payload, f.payload) {
									skip = true
								}
							}

							if skip {
								continue
							}

							found = append(found, test)
							test.m.Stop()
							wg.Done()
						}
					}
				case <-ctx.Done():
					return
				case <-listener.Done():
					return
				}
			}
		}()

		for _, test := range tests {
			if !test.shouldErr {
				wg.Add(1)
				if err := test.m.Send(ctx, test.payload); err != nil {
					t.Fatal("Multicast Send failed", err)
				}

			} else {
				if err := test.m.Send(ctx, test.payload); err == nil {
					t.Fatal("Multicast Send should fail", err)
				}
			}
		}

		wg.Wait() // Block until all test multicasts are stopped
		listener.Stop()
		cancel()
	}()

	// Block until the ctx is canceled or we receive an error, such as a timeout
	for {
		select {
		case <-ctx.Done():
			return
		case err := <-errCh:
			if err != nil {
				t.Fatal("err during Send()", err)
			}
			return
		case <-time.After(5000 * time.Millisecond):
			t.Fatal("	 timed out")
			return
		}
	}
}

func TestListenCtxDone(t *testing.T) {
	// Create a listener
	ctx, cancel := context.WithCancel(context.Background())
	listener := &Multicast{Address: testMulticastAddress}
	_, err := listener.Listen(ctx)
	if err != nil {
		t.Fatal(err)
	}

	// Test context Done()
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			}
		}
	}()
	cancel()
}

func TestListenStop(t *testing.T) {
	// Create a listener
	listener := &Multicast{Address: testMulticastAddress}
	_, err := listener.Listen(context.TODO())
	if err != nil {
		t.Fatal(err)
	}

	// Test listener Done()
	go func() {
		for {
			select {
			case <-listener.Done():
				return
			}
		}
	}()
	listener.Stop()
}

func TestCtxCancelFunc(t *testing.T) {
	m := &Multicast{Address: testMulticastAddress}

	payload := []byte("Hello TestCtxCancelFunc")
	ctx, cancel := context.WithCancel(context.Background())
	if err := m.Send(ctx, payload); err != nil {
		t.Fatal("Send error", err)
	}
	time.AfterFunc(100*time.Millisecond, func() { cancel() })
	timeout := time.AfterFunc(200*time.Millisecond, func() { t.Fatal("TestStopChan timedout") })

	// Block until the stopCh is closed
	for {
		select {
		case <-ctx.Done():
			timeout.Stop() // cancel the timeout
			return
		}
	}
}
<|endoftext|>


<commit_before>package cg

import (
	"context"
	"fmt"
	"time"

	"github.com/Shopify/sarama"
)

const (
	consumerProtocolType       = "consumer"
	consumerGroupMemberVersion = 1
)

// ProtocolKey is an implementation of a protocol that will be announced to
// the Kafka cluster under the given Key. All nodes in the group must agree
// on a common protocol.
type ProtocolKey struct {
	Protocol Protocol
	Key      string
}

// Consume is a function the Coordinator will call when it becomes responsible
// for a new topic-partition. The provided context will be canceled when the
// Coordinator is no longer responsible for the topic-partition. This function
// is expected not to block.
type Consume func(ctx context.Context, topic string, partition int32)

// Config is used to create a new Coordinator.
type Config struct {
	Client         sarama.Client
	GroupID        string
	Protocols      []ProtocolKey
	SessionTimeout time.Duration
	Heartbeat      time.Duration
	Topics         []string
	Consume        Consume
}

// Coordinator implements a Kafka GroupConsumer with semantics available
// after Kafka 0.9.
type Coordinator struct {
	cancels map[string]map[int32]func()
	client  sarama.Client
	cfg     *Config
	gid     int32
	mid     string
}

// NewCoordinator creates a Kafka GroupConsumer.
func NewCoordinator(cfg *Config) *Coordinator {
	c := &Coordinator{
		cancels: map[string]map[int32]func(){},
		client:  cfg.Client,
		cfg:     cfg,
	}
	return c
}

// Run executes the Coordinator until an error or the provided context
// is done.
func (c *Coordinator) Run(outerCtx context.Context) error {
	ctx, cancel := context.WithCancel(outerCtx)
	// ensure that ctx is canceled so that it propagates to all the Consume functions we've called.
	defer cancel()
	runErr := c.run(ctx)
	leaveErr := c.leaveGroup()
	if runErr != nil {
		return runErr
	}
	if leaveErr != nil {
		return leaveErr
	}
	return nil
}

// groupAssignments is only called by the leader and is responsible for assigning
// topic-partitions to the members in the group via a protocol common to the group.
func (c *Coordinator) groupAssignments(resp *sarama.JoinGroupResponse) (map[string]*sarama.ConsumerGroupMemberAssignment, error) {
	// build Candidates.
	members, err := resp.GetMembers()
	if err != nil {
		return map[string]*sarama.ConsumerGroupMemberAssignment{}, err
	}
	cand := Candidates{
		MemberIDs:       make([]string, 0, len(members)),
		TopicPartitions: make(map[string][]int32),
	}
	for memberID, meta := range members {
		cand.MemberIDs = append(cand.MemberIDs, memberID)
		for _, topic := range meta.Topics {
			partitions, err := c.client.Partitions(topic)
			if err != nil {
				return map[string]*sarama.ConsumerGroupMemberAssignment{}, err
			}
			cand.TopicPartitions[topic] = partitions
		}
	}
	// determine which protocol to use to assign topic-partitions.
	for _, p := range c.cfg.Protocols {
		if p.Key == resp.GroupProtocol {
			assignments := p.Protocol.Assign(cand)
			result := make(map[string]*sarama.ConsumerGroupMemberAssignment, len(assignments))
			for memberID, tp := range assignments {
				result[memberID] = &sarama.ConsumerGroupMemberAssignment{
					Version: consumerGroupMemberVersion,
					Topics:  tp,
				}
			}
			return result, nil
		}
	}
	return map[string]*sarama.ConsumerGroupMemberAssignment{}, fmt.Errorf("unhandled protocol")
}

func (c *Coordinator) heartbeat(ctx context.Context) error {
	t := time.NewTicker(c.cfg.Heartbeat)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-t.C:
			b, err := c.client.Coordinator(c.cfg.GroupID)
			if err != nil {
				return err
			}
			resp, err := b.Heartbeat(&sarama.HeartbeatRequest{
				GroupId:      c.cfg.GroupID,
				MemberId:     c.mid,
				GenerationId: c.gid,
			})
			if err != nil {
				return err
			}
			switch resp.Err {
			case sarama.ErrNoError:
				continue // keep heartbeating.
			case sarama.ErrNotCoordinatorForConsumer, sarama.ErrRebalanceInProgress:
				return nil // exit heartbeat to rebalance.
			default:
				return resp.Err
			}
		}
	}
}

func (c *Coordinator) leaveGroup() error {
	// if we were never part of a group.
	if c.mid == "" {
		return nil
	}
	b, err := c.client.Coordinator(c.cfg.GroupID)
	if err != nil {
		return err
	}
	resp, err := b.LeaveGroup(&sarama.LeaveGroupRequest{
		GroupId:  c.cfg.GroupID,
		MemberId: c.mid,
	})
	if err != nil {
		return err
	}
	if resp.Err != sarama.ErrNoError {
		return resp.Err
	}
	return nil
}

func (c *Coordinator) join() (*sarama.SyncGroupRequest, error) {
	jgr, err := c.joinGroupRequest()
	if err != nil {
		return nil, err
	}
	b, err := c.client.Coordinator(c.cfg.GroupID)
	if err != nil {
		return nil, err
	}
	resp, err := b.JoinGroup(jgr)
	if err != nil {
		return nil, err
	}
	// reset memberId if kafka has forgotten us.
	if resp.Err == sarama.ErrUnknownMemberId {
		c.mid = ""
	}
	if resp.Err != sarama.ErrNoError {
		return nil, resp.Err
	}
	c.mid = resp.MemberId
	c.gid = resp.GenerationId
	sgr := &sarama.SyncGroupRequest{
		GroupId:      c.cfg.GroupID,
		GenerationId: c.gid,
		MemberId:     c.mid,
	}
	if c.mid == resp.LeaderId {
		gas, err := c.groupAssignments(resp)
		for memberID, ga := range gas {
			sgr.AddGroupAssignmentMember(memberID, ga)
		}
		if err != nil {
			return nil, err
		}
	}
	return sgr, nil
}

func (c *Coordinator) joinGroupRequest() (*sarama.JoinGroupRequest, error) {
	req := &sarama.JoinGroupRequest{
		GroupId:        c.cfg.GroupID,
		SessionTimeout: int32(c.cfg.SessionTimeout / time.Millisecond),
		// When a member first joins the group, the memberId will be empty (i.e. "").
		MemberId:     c.mid,
		ProtocolType: consumerProtocolType,
	}
	meta := &sarama.ConsumerGroupMemberMetadata{
		Version: consumerGroupMemberVersion,
		Topics:  c.cfg.Topics,
	}
	for _, p := range c.cfg.Protocols {
		err := req.AddGroupProtocolMetadata(p.Key, meta)
		if err != nil {
			return nil, err
		}
	}
	return req, nil
}

func (c *Coordinator) run(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}
		sgr, err := c.join()
		if err != nil {
			return err
		}
		myAssignments, err := c.sync(sgr)
		if err != nil {
			return err
		}
		err = c.set(ctx, myAssignments)
		if err != nil {
			return err
		}
		err = c.heartbeat(ctx)
		if err != nil {
			return err
		}
	}
}

func (c *Coordinator) set(ctx context.Context, assignments *sarama.ConsumerGroupMemberAssignment) error {
	for topic, partitions := range assignments.Topics {
		// ensure topic entry is created.
		if _, ok := c.cancels[topic]; !ok {
			c.cancels[topic] = map[int32]func(){}
		}
		// remove consumers we already created, but are no longer responsible for.
	MyNextPartition:
		for myPartition, cancel := range c.cancels[topic] {
			for _, targetPartition := range partitions {
				if myPartition == targetPartition {
					continue MyNextPartition
				}
			}
			cancel()
			delete(c.cancels[topic], myPartition)
		}
		for _, partition := range partitions {
			// continue if we're already handling the topic-partition.
			if _, ok := c.cancels[topic][partition]; ok {
				continue
			}
			// start handling new topic-partition.
			ctx, cancel := context.WithCancel(ctx)
			c.cancels[topic][partition] = cancel
			go c.cfg.Consume(ctx, topic, partition)
		}
	}
	return nil
}

func (c *Coordinator) sync(req *sarama.SyncGroupRequest) (*sarama.ConsumerGroupMemberAssignment, error) {
	b, err := c.client.Coordinator(c.cfg.GroupID)
	if err != nil {
		return nil, err
	}
	resp, err := b.SyncGroup(req)
	if err != nil {
		return nil, err
	}
	return resp.GetMemberAssignment()
}
<commit_msg>rework where context enters; better comments<commit_after>package cg

import (
	"context"
	"fmt"
	"time"

	"github.com/Shopify/sarama"
)

const (
	consumerProtocolType       = "consumer"
	consumerGroupMemberVersion = 1
)

// ProtocolKey is an implementation of a protocol that will be announced to
// the Kafka cluster under the given Key. All nodes in the group must agree
// on a common protocol.
type ProtocolKey struct {
	Protocol Protocol
	Key      string
}

// Consume is a function the Coordinator will call when it becomes responsible
// for a new topic-partition. The provided context will be canceled when the
// Coordinator is no longer responsible for the topic-partition. This function
// is expected not to block. If you encounter an error while reading the topic-
// partition, you are still seen as responsible for that topic-partition in
// the kafka consumer group, so you must either recover or stop the Coordinator
// to remove yourself from the consumer group.
type Consume func(ctx context.Context, topic string, partition int32)

// Config is used to create a new Coordinator.
type Config struct {
	Client         sarama.Client
	Context        context.Context
	GroupID        string
	Protocols      []ProtocolKey
	SessionTimeout time.Duration
	Heartbeat      time.Duration
	Topics         []string
	Consume        Consume
}

// Coordinator implements a Kafka GroupConsumer with semantics available
// after Kafka 0.9.
type Coordinator struct {
	cancels map[string]map[int32]func()
	client  sarama.Client
	cfg     *Config
	ctx     context.Context
	gid     int32
	mid     string
}

// NewCoordinator creates a Kafka GroupConsumer.
func NewCoordinator(cfg *Config) *Coordinator {
	c := &Coordinator{
		cancels: map[string]map[int32]func(){},
		client:  cfg.Client,
		cfg:     cfg,
		ctx:     cfg.Context,
	}
	return c
}

// Run executes the Coordinator until an error or the context provided at
// create time closes.
func (c *Coordinator) Run() error {
	ctx, cancel := context.WithCancel(c.ctx)
	// ensure that ctx is canceled so that it propagates to all the
	// Consume functions we've called.
	defer cancel()
	// we run c.leaveGroup() notwithstanding the result of c.run and
	// return the first error between them both (if any)
	runErr := c.run(ctx)
	leaveErr := c.leaveGroup()
	if runErr != nil {
		return runErr
	}
	if leaveErr != nil {
		return leaveErr
	}
	return nil
}

// groupAssignments is only called by the leader and is responsible for assigning
// topic-partitions to the members in the group via a protocol common to the group.
func (c *Coordinator) groupAssignments(resp *sarama.JoinGroupResponse) (map[string]*sarama.ConsumerGroupMemberAssignment, error) {
	// build Candidates.
	members, err := resp.GetMembers()
	if err != nil {
		return map[string]*sarama.ConsumerGroupMemberAssignment{}, err
	}
	cand := Candidates{
		MemberIDs:       make([]string, 0, len(members)),
		TopicPartitions: make(map[string][]int32),
	}
	for memberID, meta := range members {
		cand.MemberIDs = append(cand.MemberIDs, memberID)
		for _, topic := range meta.Topics {
			partitions, err := c.client.Partitions(topic)
			if err != nil {
				return map[string]*sarama.ConsumerGroupMemberAssignment{}, err
			}
			cand.TopicPartitions[topic] = partitions
		}
	}
	// determine which protocol to use to assign topic-partitions.
	for _, p := range c.cfg.Protocols {
		if p.Key == resp.GroupProtocol {
			assignments := p.Protocol.Assign(cand)
			result := make(map[string]*sarama.ConsumerGroupMemberAssignment, len(assignments))
			for memberID, tp := range assignments {
				result[memberID] = &sarama.ConsumerGroupMemberAssignment{
					Version: consumerGroupMemberVersion,
					Topics:  tp,
				}
			}
			return result, nil
		}
	}
	return map[string]*sarama.ConsumerGroupMemberAssignment{}, fmt.Errorf("unhandled protocol")
}

func (c *Coordinator) heartbeat(ctx context.Context) error {
	t := time.NewTicker(c.cfg.Heartbeat)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-t.C:
			b, err := c.client.Coordinator(c.cfg.GroupID)
			if err != nil {
				return err
			}
			resp, err := b.Heartbeat(&sarama.HeartbeatRequest{
				GroupId:      c.cfg.GroupID,
				MemberId:     c.mid,
				GenerationId: c.gid,
			})
			if err != nil {
				return err
			}
			switch resp.Err {
			case sarama.ErrNoError:
				continue // keep heartbeating.
			case sarama.ErrNotCoordinatorForConsumer, sarama.ErrRebalanceInProgress:
				return nil // exit heartbeat to rebalance.
			default:
				return resp.Err
			}
		}
	}
}

func (c *Coordinator) leaveGroup() error {
	// if we were never part of a group.
	if c.mid == "" {
		return nil
	}
	b, err := c.client.Coordinator(c.cfg.GroupID)
	if err != nil {
		return err
	}
	resp, err := b.LeaveGroup(&sarama.LeaveGroupRequest{
		GroupId:  c.cfg.GroupID,
		MemberId: c.mid,
	})
	if err != nil {
		return err
	}
	if resp.Err != sarama.ErrNoError {
		return resp.Err
	}
	return nil
}

func (c *Coordinator) join() (*sarama.SyncGroupRequest, error) {
	jgr, err := c.joinGroupRequest()
	if err != nil {
		return nil, err
	}
	b, err := c.client.Coordinator(c.cfg.GroupID)
	if err != nil {
		return nil, err
	}
	resp, err := b.JoinGroup(jgr)
	if err != nil {
		return nil, err
	}
	// reset memberId if kafka has forgotten us.
	if resp.Err == sarama.ErrUnknownMemberId {
		c.mid = ""
	}
	if resp.Err != sarama.ErrNoError {
		return nil, resp.Err
	}
	c.mid = resp.MemberId
	c.gid = resp.GenerationId
	sgr := &sarama.SyncGroupRequest{
		GroupId:      c.cfg.GroupID,
		GenerationId: c.gid,
		MemberId:     c.mid,
	}
	if c.mid == resp.LeaderId {
		gas, err := c.groupAssignments(resp)
		for memberID, ga := range gas {
			sgr.AddGroupAssignmentMember(memberID, ga)
		}
		if err != nil {
			return nil, err
		}
	}
	return sgr, nil
}

func (c *Coordinator) joinGroupRequest() (*sarama.JoinGroupRequest, error) {
	req := &sarama.JoinGroupRequest{
		GroupId:        c.cfg.GroupID,
		SessionTimeout: int32(c.cfg.SessionTimeout / time.Millisecond),
		// When a member first joins the group, the memberId will be empty (i.e. "").
		MemberId:     c.mid,
		ProtocolType: consumerProtocolType,
	}
	meta := &sarama.ConsumerGroupMemberMetadata{
		Version: consumerGroupMemberVersion,
		Topics:  c.cfg.Topics,
	}
	for _, p := range c.cfg.Protocols {
		err := req.AddGroupProtocolMetadata(p.Key, meta)
		if err != nil {
			return nil, err
		}
	}
	return req, nil
}

func (c *Coordinator) run(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}
		sgr, err := c.join()
		if err != nil {
			return err
		}
		myAssignments, err := c.sync(sgr)
		if err != nil {
			return err
		}
		err = c.set(ctx, myAssignments)
		if err != nil {
			return err
		}
		err = c.heartbeat(ctx)
		if err != nil {
			return err
		}
	}
}

func (c *Coordinator) set(ctx context.Context, assignments *sarama.ConsumerGroupMemberAssignment) error {
	for topic, partitions := range assignments.Topics {
		// ensure topic entry is created.
		if _, ok := c.cancels[topic]; !ok {
			c.cancels[topic] = map[int32]func(){}
		}
		// remove consumers we already created, but are no longer responsible for.
	MyNextPartition:
		for myPartition, cancel := range c.cancels[topic] {
			for _, targetPartition := range partitions {
				if myPartition == targetPartition {
					continue MyNextPartition
				}
			}
			cancel()
			delete(c.cancels[topic], myPartition)
		}
		for _, partition := range partitions {
			// continue if we're already handling the topic-partition.
			if _, ok := c.cancels[topic][partition]; ok {
				continue
			}
			// start handling new topic-partition.
			ctx, cancel := context.WithCancel(ctx)
			c.cancels[topic][partition] = cancel
			go c.cfg.Consume(ctx, topic, partition)
		}
	}
	return nil
}

func (c *Coordinator) sync(req *sarama.SyncGroupRequest) (*sarama.ConsumerGroupMemberAssignment, error) {
	b, err := c.client.Coordinator(c.cfg.GroupID)
	if err != nil {
		return nil, err
	}
	resp, err := b.SyncGroup(req)
	if err != nil {
		return nil, err
	}
	return resp.GetMemberAssignment()
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"sync"

	"errors"

	"github.com/TykTechnologies/tyk/config"
)

var orgChanMap = make(map[string]chan bool)

type orgActiveMapMu struct {
	sync.RWMutex
	OrgMap map[string]bool
}

var orgActiveMap = orgActiveMapMu{
	OrgMap: map[string]bool{},
}

// RateLimitAndQuotaCheck will check the incomming request and key whether it is within it's quota and
// within it's rate limit, it makes use of the SessionLimiter object to do this
type OrganizationMonitor struct {
	BaseMiddleware
	sessionlimiter SessionLimiter
	mon            Monitor
}

func (k *OrganizationMonitor) Name() string {
	return "OrganizationMonitor"
}

func (k *OrganizationMonitor) EnabledForSpec() bool {
	// If false, we aren't enforcing quotas so skip this mw
	// altogether
	return config.Global.EnforceOrgQuotas
}

func (k *OrganizationMonitor) ProcessRequest(w http.ResponseWriter, r *http.Request, conf interface{}) (error, int) {
	if config.Global.ExperimentalProcessOrgOffThread {
		return k.ProcessRequestOffThread(r)
	}
	return k.ProcessRequestLive(r)
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (k *OrganizationMonitor) ProcessRequestLive(r *http.Request) (error, int) {

	session, found := k.OrgSession(k.Spec.OrgID)
	if !found {
		// No organisation session has been created, should not be a pre-requisite in site setups, so we pass the request on
		return nil, 200
	}

	// Is it active?
	logEntry := getLogEntryForRequest(r, k.Spec.OrgID, nil)
	if session.IsInactive {
		logEntry.Warning("Organisation access is disabled.")

		return errors.New("this organisation access has been disabled, please contact your API administrator"), 403
	}

	// We found a session, apply the quota limiter
	reason := k.sessionlimiter.ForwardMessage(&session,
		k.Spec.OrgID,
		k.Spec.OrgSessionManager.Store(), false, false)

	k.Spec.OrgSessionManager.UpdateSession(k.Spec.OrgID, &session, session.Lifetime(k.Spec.SessionLifetime))

	// org limits apply only to quotas, so we don't care about
	// sessionFailRateLimit.
	switch reason {
	case sessionFailNone:
	case sessionFailQuota:
		logEntry.Warning("Organisation quota has been exceeded.")

		// Fire a quota exceeded event
		k.FireEvent(EventOrgQuotaExceeded, EventKeyFailureMeta{
			EventMetaDefault: EventMetaDefault{Message: "Organisation quota has been exceeded", OriginatingRequest: EncodeRequestToEvent(r)},
			Path:             r.URL.Path,
			Origin:           requestIP(r),
			Key:              k.Spec.OrgID,
		})

		return errors.New("This organisation quota has been exceeded, please contact your API administrator"), 403
	}

	if config.Global.Monitor.MonitorOrgKeys {
		// Run the trigger monitor
		k.mon.Check(&session, "")
	}

	// Lets keep a reference of the org
	setCtxValue(r, OrgSessionContext, session)

	// Request is valid, carry on
	return nil, 200
}

func (k *OrganizationMonitor) SetOrgSentinel(orgChan chan bool, orgId string) {
	for isActive := range orgChan {
		log.Debug("Chan got:", isActive)
		orgActiveMap.Lock()
		orgActiveMap.OrgMap[orgId] = isActive
		orgActiveMap.Unlock()
	}
}

func (k *OrganizationMonitor) ProcessRequestOffThread(r *http.Request) (error, int) {

	orgActiveMap.Lock()
	orgChan, ok := orgChanMap[k.Spec.OrgID]
	if !ok {
		orgChanMap[k.Spec.OrgID] = make(chan bool)
		orgChan = orgChanMap[k.Spec.OrgID]
		go k.SetOrgSentinel(orgChan, k.Spec.OrgID)
	}
	active, found := orgActiveMap.OrgMap[k.Spec.OrgID]
	orgActiveMap.Unlock()

	requestCopy := copyRequest(r)
	go k.AllowAccessNext(orgChan, requestCopy)

	if found && !active {
		log.Debug("Is not active")
		return errors.New("This organisation access has been disabled or quota is exceeded, please contact your API administrator"), 403
	}

	log.Debug("Key not found")

	// Request is valid, carry on
	return nil, 200
}

func (k *OrganizationMonitor) AllowAccessNext(orgChan chan bool, r *http.Request) {

	session, found := k.OrgSession(k.Spec.OrgID)

	if !found {
		// No organisation session has been created, should not be a pre-requisite in site setups, so we pass the request on
		log.Debug("No session for org, skipping")
		return
	}

	// Is it active?
	logEntry := getLogEntryForRequest(r, k.Spec.OrgID, nil)
	if session.IsInactive {
		logEntry.Warning("Organisation access is disabled.")

		//return errors.New("This organisation access has been disabled, please contact your API administrator."), 403
		orgChan <- false
		return
	}

	// We found a session, apply the quota limiter
	isQuotaExceeded := k.sessionlimiter.RedisQuotaExceeded(&session, k.Spec.OrgID, k.Spec.OrgSessionManager.Store())

	k.Spec.OrgSessionManager.UpdateSession(k.Spec.OrgID, &session, session.Lifetime(k.Spec.SessionLifetime))

	if isQuotaExceeded {
		logEntry.Warning("Organisation quota has been exceeded.")

		// Fire a quota exceeded event
		k.FireEvent(EventOrgQuotaExceeded, EventKeyFailureMeta{
			EventMetaDefault: EventMetaDefault{Message: "Organisation quota has been exceeded", OriginatingRequest: EncodeRequestToEvent(r)},
			Path:             r.URL.Path,
			Origin:           requestIP(r),
			Key:              k.Spec.OrgID,
		})

		//return errors.New("This organisation quota has been exceeded, please contact your API administrator"), 403
		orgChan <- false

		if config.Global.Monitor.MonitorOrgKeys {
			// Run the trigger monitor
			k.mon.Check(&session, "")
		}

		return
	}

	if config.Global.Monitor.MonitorOrgKeys {
		// Run the trigger monitor
		k.mon.Check(&session, "")
	}

	// Lets keep a reference of the org
	setCtxValue(r, OrgSessionContext, session)

	orgChan <- true
}
<commit_msg>Changed to a sync.Map instead of a mutex<commit_after>package main

import (
	"net/http"
	"sync"

	"errors"

	"github.com/TykTechnologies/tyk/config"
)

var orgChanMap = make(map[string]chan bool)

var orgActiveMap sync.Map

// RateLimitAndQuotaCheck will check the incomming request and key whether it is within it's quota and
// within it's rate limit, it makes use of the SessionLimiter object to do this
type OrganizationMonitor struct {
	BaseMiddleware
	sessionlimiter SessionLimiter
	mon            Monitor
}

func (k *OrganizationMonitor) Name() string {
	return "OrganizationMonitor"
}

func (k *OrganizationMonitor) EnabledForSpec() bool {
	// If false, we aren't enforcing quotas so skip this mw
	// altogether
	return config.Global.EnforceOrgQuotas
}

func (k *OrganizationMonitor) ProcessRequest(w http.ResponseWriter, r *http.Request, conf interface{}) (error, int) {
	if config.Global.ExperimentalProcessOrgOffThread {
		return k.ProcessRequestOffThread(r)
	}
	return k.ProcessRequestLive(r)
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (k *OrganizationMonitor) ProcessRequestLive(r *http.Request) (error, int) {

	session, found := k.OrgSession(k.Spec.OrgID)
	if !found {
		// No organisation session has been created, should not be a pre-requisite in site setups, so we pass the request on
		return nil, 200
	}

	// Is it active?
	logEntry := getLogEntryForRequest(r, k.Spec.OrgID, nil)
	if session.IsInactive {
		logEntry.Warning("Organisation access is disabled.")

		return errors.New("this organisation access has been disabled, please contact your API administrator"), 403
	}

	// We found a session, apply the quota limiter
	reason := k.sessionlimiter.ForwardMessage(&session,
		k.Spec.OrgID,
		k.Spec.OrgSessionManager.Store(), false, false)

	k.Spec.OrgSessionManager.UpdateSession(k.Spec.OrgID, &session, session.Lifetime(k.Spec.SessionLifetime))

	// org limits apply only to quotas, so we don't care about
	// sessionFailRateLimit.
	switch reason {
	case sessionFailNone:
	case sessionFailQuota:
		logEntry.Warning("Organisation quota has been exceeded.")

		// Fire a quota exceeded event
		k.FireEvent(EventOrgQuotaExceeded, EventKeyFailureMeta{
			EventMetaDefault: EventMetaDefault{Message: "Organisation quota has been exceeded", OriginatingRequest: EncodeRequestToEvent(r)},
			Path:             r.URL.Path,
			Origin:           requestIP(r),
			Key:              k.Spec.OrgID,
		})

		return errors.New("This organisation quota has been exceeded, please contact your API administrator"), 403
	}

	if config.Global.Monitor.MonitorOrgKeys {
		// Run the trigger monitor
		k.mon.Check(&session, "")
	}

	// Lets keep a reference of the org
	setCtxValue(r, OrgSessionContext, session)

	// Request is valid, carry on
	return nil, 200
}

func (k *OrganizationMonitor) SetOrgSentinel(orgChan chan bool, orgId string) {
	for isActive := range orgChan {
		log.Debug("Chan got:", isActive)
		o, ok := orgActiveMap.Load(orgId)
		if ok {
			o = isActive
			orgActiveMap.Store(orgId, o)
		}
	}
}

func (k *OrganizationMonitor) ProcessRequestOffThread(r *http.Request) (error, int) {

	orgChan, ok := orgChanMap[k.Spec.OrgID]
	if !ok {
		orgChanMap[k.Spec.OrgID] = make(chan bool)
		orgChan = orgChanMap[k.Spec.OrgID]
		go k.SetOrgSentinel(orgChan, k.Spec.OrgID)
	}
	active, found := orgActiveMap.Load(k.Spec.OrgID)

	requestCopy := copyRequest(r)
	go k.AllowAccessNext(orgChan, requestCopy)

	if found && !active.(bool) {
		log.Debug("Is not active")
		return errors.New("This organisation access has been disabled or quota is exceeded, please contact your API administrator"), 403
	}

	log.Debug("Key not found")

	// Request is valid, carry on
	return nil, 200
}

func (k *OrganizationMonitor) AllowAccessNext(orgChan chan bool, r *http.Request) {

	session, found := k.OrgSession(k.Spec.OrgID)

	if !found {
		// No organisation session has been created, should not be a pre-requisite in site setups, so we pass the request on
		log.Debug("No session for org, skipping")
		return
	}

	// Is it active?
	logEntry := getLogEntryForRequest(r, k.Spec.OrgID, nil)
	if session.IsInactive {
		logEntry.Warning("Organisation access is disabled.")

		//return errors.New("This organisation access has been disabled, please contact your API administrator."), 403
		orgChan <- false
		return
	}

	// We found a session, apply the quota limiter
	isQuotaExceeded := k.sessionlimiter.RedisQuotaExceeded(&session, k.Spec.OrgID, k.Spec.OrgSessionManager.Store())

	k.Spec.OrgSessionManager.UpdateSession(k.Spec.OrgID, &session, session.Lifetime(k.Spec.SessionLifetime))

	if isQuotaExceeded {
		logEntry.Warning("Organisation quota has been exceeded.")

		// Fire a quota exceeded event
		k.FireEvent(EventOrgQuotaExceeded, EventKeyFailureMeta{
			EventMetaDefault: EventMetaDefault{Message: "Organisation quota has been exceeded", OriginatingRequest: EncodeRequestToEvent(r)},
			Path:             r.URL.Path,
			Origin:           requestIP(r),
			Key:              k.Spec.OrgID,
		})

		//return errors.New("This organisation quota has been exceeded, please contact your API administrator"), 403
		orgChan <- false

		if config.Global.Monitor.MonitorOrgKeys {
			// Run the trigger monitor
			k.mon.Check(&session, "")
		}

		return
	}

	if config.Global.Monitor.MonitorOrgKeys {
		// Run the trigger monitor
		k.mon.Check(&session, "")
	}

	// Lets keep a reference of the org
	setCtxValue(r, OrgSessionContext, session)

	orgChan <- true
}
<|endoftext|>


<commit_before>package nsq

import (
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

// ProducerConfig carries the different variables to tune a newly started
// producer.
type ProducerConfig struct {
	Address         string
	Topic           string
	MaxConcurrency  int
	ForceConnect    bool
	DialTimeout     time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	MaxRetryTimeout time.Duration
	MinRetryTimeout time.Duration
}

// Producer provide an abstraction around using direct connections to nsqd
// nodes to send messages.
type Producer struct {
	// Communication channels of the producer.
	reqs chan ProducerRequest
	done chan struct{}
	once sync.Once
	join sync.WaitGroup
	ok   uint32

	// Immutable state of the producer.
	address         string
	topic           string
	forceConnect    bool
	dialTimeout     time.Duration
	readTimeout     time.Duration
	writeTimeout    time.Duration
	maxRetryTimeout time.Duration
	minRetryTimeout time.Duration
}

// ProducerRequest are used to represent operations that are submitted to
// producers.
type ProducerRequest struct {
	Topic    string
	Message  []byte
	Response chan<- error
	Deadline time.Time
}

// StartProducer starts and returns a new producer p, configured with the
// variables from the config parameter, or returning an non-nil error if
// some of the configuration variables were invalid.
func StartProducer(config ProducerConfig) (p *Producer, err error) {
	if len(config.Address) == 0 {
		config.Address = "localhost:4151"
	}

	if config.MaxConcurrency == 0 {
		config.MaxConcurrency = DefaultMaxConcurrency
	}

	if config.DialTimeout == 0 {
		config.DialTimeout = DefaultDialTimeout
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = DefaultReadTimeout
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = DefaultWriteTimeout
	}

	if config.MaxRetryTimeout == 0 {
		config.MaxRetryTimeout = DefaultMaxRetryTimeout
	}

	if config.MinRetryTimeout == 0 {
		config.MinRetryTimeout = DefaultMinRetryTimeout
	}

	p = &Producer{
		reqs:            make(chan ProducerRequest, config.MaxConcurrency),
		done:            make(chan struct{}),
		address:         config.Address,
		topic:           config.Topic,
		forceConnect:    config.ForceConnect,
		dialTimeout:     config.DialTimeout,
		readTimeout:     config.ReadTimeout,
		writeTimeout:    config.WriteTimeout,
		maxRetryTimeout: config.MaxRetryTimeout,
		minRetryTimeout: config.MinRetryTimeout,
	}
	p.join.Add(config.MaxConcurrency)

	for i := 0; i != config.MaxConcurrency; i++ {
		go p.run()
	}

	return
}

// Stop gracefully shutsdown the producer, cancelling all inflight requests and
// waiting for all backend connections to be closed.
//
// It is safe to call the method multiple times and from multiple goroutines,
// they will all block until the producer has been completely shutdown.
func (p *Producer) Stop() {
	p.once.Do(p.stop)
	err := errors.New("publishing to a producer that was already stopped")

	for req := range p.reqs {
		req.complete(err)
	}

	p.join.Wait()
}

// Publish sends a message using the producer p, returning an error if it was
// already closed or if an error occurred while publishing the message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) Publish(message []byte) (err error) {
	return p.PublishTo(p.topic, message)
}

// PublishTo sends a message to a specific topic using the producer p, returning
// an error if it was already closed or if an error occurred while publishing the
// message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) PublishTo(topic string, message []byte) (err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("publishing to a producer that was already stopped")
		}
	}()

	if len(topic) == 0 {
		return errors.New("no topic set for publishing message")
	}

	response := make(chan error, 1)
	deadline := time.Now().Add(p.dialTimeout + p.readTimeout + p.writeTimeout)

	// Attempts to queue the request so one of the active connections can pick
	// it up.
	p.reqs <- ProducerRequest{
		Topic:    topic,
		Message:  message,
		Response: response,
		Deadline: deadline,
	}

	// This will always trigger, either if the connection was lost or if a
	// response was successfully sent.
	err = <-response
	return
}

// Requests returns a write-only channel that can be used to submit requests to p.
//
// This method is useful when the publish operation needs to be associated with
// other operations on channels in a select statement for example, or to publish
// in a non-blocking fashion.
func (p *Producer) Requests() chan<- ProducerRequest {
	return p.reqs
}

// Connected returns true if the producer has successfully established a
// connection to nsqd, false otherwise.
func (p *Producer) Connected() bool {
	return atomic.LoadUint32(&p.ok) != 0
}

func (p *Producer) stop() {
	close(p.done)
	close(p.reqs)
}

func (p *Producer) run() {
	var conn *Conn
	var resChan chan Frame
	var pending []ProducerRequest
	var retry time.Duration

	shutdown := func(err error) {
		atomic.StoreUint32(&p.ok, 0)

		if conn != nil {
			close(resChan)
			conn.Close()
			conn = nil
			resChan = nil
			pending = completeAllProducerRequests(pending, err)
		}
	}

	connect := func() (err error) {
		if conn, err = DialTimeout(p.address, p.dialTimeout); err != nil {
			log.Printf("failed to connect to %s, retrying after %s: %s", p.address, retry, err)
			retry = p.sleep(retry)
			return
		}

		retry = 0
		resChan = make(chan Frame, 16)
		go p.flush(conn, resChan)

		atomic.StoreUint32(&p.ok, 1)
		return
	}

	defer p.join.Done()
	defer shutdown(nil)

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	if p.forceConnect {
		connect()
	}

	for {
		select {
		case <-p.done:
			return

		case now := <-ticker.C:
			if p.forceConnect && conn == nil {
				connect()
			}

			if producerRequestsTimedOut(pending, now) {
				shutdown(errors.New("timeout"))
				continue
			}

		case req, ok := <-p.reqs:
			if !ok {
				return
			}

			if conn == nil {
				if err := connect(); err != nil {
					req.complete(err)
					continue
				}
			}

			if err := p.publish(conn, req.Topic, req.Message); err != nil {
				req.complete(err)
				shutdown(err)
				continue
			}

			pending = append(pending, req)

		case frame := <-resChan:
			switch f := frame.(type) {
			case Response:
				switch f {
				case OK:
					pending = completeNextProducerRequest(pending, nil)

				case Heartbeat:
					if err := p.ping(conn); err != nil {
						shutdown(err)
						continue
					}

				case CloseWait:
					return

				default:
					shutdown(errors.Errorf("closing connection after receiving an unexpected response from %s: %s", conn.RemoteAddr(), f))
					continue
				}

			case Error:
				shutdown(errors.Errorf("closing connection after receiving an error from %s: %s", conn.RemoteAddr(), f))
				continue

			case Message:
				shutdown(errors.Errorf("closing connection after receiving an unexpected message from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue

			default:
				shutdown(errors.Errorf("closing connection after receiving an unsupported frame from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue
			}
		}
	}
}

func (p *Producer) flush(conn *Conn, resChan chan<- Frame) {
	defer func() { recover() }() // may happen if the resChan is closed
	defer conn.Close()

	for {
		frame, err := conn.ReadFrame()

		if err != nil {
			resChan <- Error(err.Error())
			return
		}

		resChan <- frame
	}
}

func (p *Producer) write(conn *Conn, cmd Command) (err error) {
	if err = conn.SetDeadline(time.Now().Add(p.writeTimeout)); err == nil {
		err = conn.WriteCommand(cmd)
	}
	return
}

func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
	return p.write(conn, Pub{
		Topic:   topic,
		Message: message,
	})
}

func (p *Producer) ping(conn *Conn) error {
	return p.write(conn, Nop{})
}

func (p *Producer) sleep(d time.Duration) time.Duration {
	if d < p.minRetryTimeout {
		d = p.minRetryTimeout
	}

	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-t.C:
	case <-p.done:
	}

	if d *= 2; d > p.maxRetryTimeout {
		d = p.maxRetryTimeout
	}

	return d
}

func (r ProducerRequest) complete(err error) {
	if r.Response != nil {
		r.Response <- err
	}
}

func completeNextProducerRequest(reqs []ProducerRequest, err error) []ProducerRequest {
	reqs[0].complete(err)
	return reqs[1:]
}

func completeAllProducerRequests(reqs []ProducerRequest, err error) []ProducerRequest {
	for _, req := range reqs {
		req.complete(err)
	}
	return nil
}

func producerRequestsTimedOut(reqs []ProducerRequest, now time.Time) bool {
	return len(reqs) != 0 && now.After(reqs[0].Deadline)
}
<commit_msg>remove ForceConnect<commit_after>package nsq

import (
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

// ProducerConfig carries the different variables to tune a newly started
// producer.
type ProducerConfig struct {
	Address         string
	Topic           string
	MaxConcurrency  int
	DialTimeout     time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	MaxRetryTimeout time.Duration
	MinRetryTimeout time.Duration
}

// Producer provide an abstraction around using direct connections to nsqd
// nodes to send messages.
type Producer struct {
	// Communication channels of the producer.
	reqs chan ProducerRequest
	done chan struct{}
	once sync.Once
	join sync.WaitGroup
	ok   uint32

	// Immutable state of the producer.
	address         string
	topic           string
	dialTimeout     time.Duration
	readTimeout     time.Duration
	writeTimeout    time.Duration
	maxRetryTimeout time.Duration
	minRetryTimeout time.Duration
}

// ProducerRequest are used to represent operations that are submitted to
// producers.
type ProducerRequest struct {
	Topic    string
	Message  []byte
	Response chan<- error
	Deadline time.Time
}

// StartProducer starts and returns a new producer p, configured with the
// variables from the config parameter, or returning an non-nil error if
// some of the configuration variables were invalid.
func StartProducer(config ProducerConfig) (p *Producer, err error) {
	if len(config.Address) == 0 {
		config.Address = "localhost:4151"
	}

	if config.MaxConcurrency == 0 {
		config.MaxConcurrency = DefaultMaxConcurrency
	}

	if config.DialTimeout == 0 {
		config.DialTimeout = DefaultDialTimeout
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = DefaultReadTimeout
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = DefaultWriteTimeout
	}

	if config.MaxRetryTimeout == 0 {
		config.MaxRetryTimeout = DefaultMaxRetryTimeout
	}

	if config.MinRetryTimeout == 0 {
		config.MinRetryTimeout = DefaultMinRetryTimeout
	}

	p = &Producer{
		reqs:            make(chan ProducerRequest, config.MaxConcurrency),
		done:            make(chan struct{}),
		address:         config.Address,
		topic:           config.Topic,
		dialTimeout:     config.DialTimeout,
		readTimeout:     config.ReadTimeout,
		writeTimeout:    config.WriteTimeout,
		maxRetryTimeout: config.MaxRetryTimeout,
		minRetryTimeout: config.MinRetryTimeout,
	}
	p.join.Add(config.MaxConcurrency)

	for i := 0; i != config.MaxConcurrency; i++ {
		go p.run()
	}

	return
}

// Stop gracefully shutsdown the producer, cancelling all inflight requests and
// waiting for all backend connections to be closed.
//
// It is safe to call the method multiple times and from multiple goroutines,
// they will all block until the producer has been completely shutdown.
func (p *Producer) Stop() {
	p.once.Do(p.stop)
	err := errors.New("publishing to a producer that was already stopped")

	for req := range p.reqs {
		req.complete(err)
	}

	p.join.Wait()
}

// Publish sends a message using the producer p, returning an error if it was
// already closed or if an error occurred while publishing the message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) Publish(message []byte) (err error) {
	return p.PublishTo(p.topic, message)
}

// PublishTo sends a message to a specific topic using the producer p, returning
// an error if it was already closed or if an error occurred while publishing the
// message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) PublishTo(topic string, message []byte) (err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("publishing to a producer that was already stopped")
		}
	}()

	if len(topic) == 0 {
		return errors.New("no topic set for publishing message")
	}

	response := make(chan error, 1)
	deadline := time.Now().Add(p.dialTimeout + p.readTimeout + p.writeTimeout)

	// Attempts to queue the request so one of the active connections can pick
	// it up.
	p.reqs <- ProducerRequest{
		Topic:    topic,
		Message:  message,
		Response: response,
		Deadline: deadline,
	}

	// This will always trigger, either if the connection was lost or if a
	// response was successfully sent.
	err = <-response
	return
}

// Requests returns a write-only channel that can be used to submit requests to p.
//
// This method is useful when the publish operation needs to be associated with
// other operations on channels in a select statement for example, or to publish
// in a non-blocking fashion.
func (p *Producer) Requests() chan<- ProducerRequest {
	return p.reqs
}

// Connected returns true if the producer has successfully established a
// connection to nsqd, false otherwise.
func (p *Producer) Connected() bool {
	return atomic.LoadUint32(&p.ok) != 0
}

func (p *Producer) stop() {
	close(p.done)
	close(p.reqs)
}

func (p *Producer) run() {
	var conn *Conn
	var reqChan <-chan ProducerRequest
	var resChan chan Frame
	var pending []ProducerRequest
	var retry time.Duration

	shutdown := func(err error) {
		atomic.StoreUint32(&p.ok, 0)

		if conn != nil {
			close(resChan)
			conn.Close()
			conn = nil
			reqChan = nil
			resChan = nil
			pending = completeAllProducerRequests(pending, err)
		}
	}

	connect := func() (err error) {
		if conn, err = DialTimeout(p.address, p.dialTimeout); err != nil {
			log.Printf("failed to connect to %s, retrying after %s: %s", p.address, retry, err)
			retry = p.sleep(retry)
			return
		}

		retry = 0
		reqChan = p.reqs
		resChan = make(chan Frame, 16)
		go p.flush(conn, resChan)

		atomic.StoreUint32(&p.ok, 1)
		return
	}

	defer p.join.Done()
	defer shutdown(nil)

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	connect()

	for {
		select {
		case <-p.done:
			return

		case now := <-ticker.C:
			if conn == nil {
				connect()
			}

			if producerRequestsTimedOut(pending, now) {
				shutdown(errors.New("timeout"))
				continue
			}

		case req, ok := <-reqChan:
			if !ok {
				return
			}

			if err := p.publish(conn, req.Topic, req.Message); err != nil {
				req.complete(err)
				shutdown(err)
				continue
			}

			pending = append(pending, req)

		case frame := <-resChan:
			switch f := frame.(type) {
			case Response:
				switch f {
				case OK:
					pending = completeNextProducerRequest(pending, nil)

				case Heartbeat:
					if err := p.ping(conn); err != nil {
						shutdown(err)
						continue
					}

				case CloseWait:
					return

				default:
					shutdown(errors.Errorf("closing connection after receiving an unexpected response from %s: %s", conn.RemoteAddr(), f))
					continue
				}

			case Error:
				shutdown(errors.Errorf("closing connection after receiving an error from %s: %s", conn.RemoteAddr(), f))
				continue

			case Message:
				shutdown(errors.Errorf("closing connection after receiving an unexpected message from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue

			default:
				shutdown(errors.Errorf("closing connection after receiving an unsupported frame from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue
			}
		}
	}
}

func (p *Producer) flush(conn *Conn, resChan chan<- Frame) {
	defer func() { recover() }() // may happen if the resChan is closed
	defer conn.Close()

	for {
		frame, err := conn.ReadFrame()

		if err != nil {
			resChan <- Error(err.Error())
			return
		}

		resChan <- frame
	}
}

func (p *Producer) write(conn *Conn, cmd Command) (err error) {
	if err = conn.SetDeadline(time.Now().Add(p.writeTimeout)); err == nil {
		err = conn.WriteCommand(cmd)
	}
	return
}

func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
	return p.write(conn, Pub{
		Topic:   topic,
		Message: message,
	})
}

func (p *Producer) ping(conn *Conn) error {
	return p.write(conn, Nop{})
}

func (p *Producer) sleep(d time.Duration) time.Duration {
	if d < p.minRetryTimeout {
		d = p.minRetryTimeout
	}

	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-t.C:
	case <-p.done:
	}

	if d *= 2; d > p.maxRetryTimeout {
		d = p.maxRetryTimeout
	}

	return d
}

func (r ProducerRequest) complete(err error) {
	if r.Response != nil {
		r.Response <- err
	}
}

func completeNextProducerRequest(reqs []ProducerRequest, err error) []ProducerRequest {
	reqs[0].complete(err)
	return reqs[1:]
}

func completeAllProducerRequests(reqs []ProducerRequest, err error) []ProducerRequest {
	for _, req := range reqs {
		req.complete(err)
	}
	return nil
}

func producerRequestsTimedOut(reqs []ProducerRequest, now time.Time) bool {
	return len(reqs) != 0 && now.After(reqs[0].Deadline)
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"github.com/conformal/btcjson"
	"io/ioutil"
	"log"
	"os"
	//"os/exec"
	"os/signal"
	"path/filepath"
	"time"

	"github.com/conformal/btcrpcclient"
	"github.com/conformal/btcutil"
	//"github.com/davecgh/go-spew/spew"
)

func Serialize(item btcjson.ListTransactionsResult) string {
	amt := fmt.Sprintf("%.8f", item.Amount)
	fee := fmt.Sprintf("%.8f", item.Fee)
	return item.Address + "," + amt + "," + fee
}

func main() {
	//var cmd *exec.Cmd
	connected := make(chan struct{})
	var firstConn bool
	// based off of btcwebsocket example for reference
	ntfnHandlers := btcrpcclient.NotificationHandlers{
		OnBtcdConnected: func(conn bool) {
			if conn && !firstConn {
				firstConn = true
				connected <- struct{}{}
			}
		},
	}

	// Connect to local btcwallet RPC server using websockets.
	certHomeDir := btcutil.AppDataDir("btcwallet", false)
	certs, err := ioutil.ReadFile(filepath.Join(certHomeDir, "rpc.cert"))
	if err != nil {
		log.Fatal(err)
	}

	connCfg := &btcrpcclient.ConnConfig{
		Host:         "localhost:18554",
		Endpoint:     "ws",
		User:         "rpcuser",
		Pass:         "rpcpass",
		Certificates: certs,
	}
	client, err := btcrpcclient.New(connCfg, &ntfnHandlers)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Waiting for BTCD")

	// Unlock wallet.
	log.Println("Unlock Wallet")
	if err := client.WalletPassphrase("walletpass", 3600); err != nil {
		log.Printf("%s: Cannot unlock wallet: %v", "localhost:18556", err)

	}

	// count transactions
	client.NotifyNewTransactions(true)
	log.Println("Monitoring TXns")
	for {

		StartTime := time.Now().Unix()
		data, err := client.ListTransactionsCount("", 100)

		// Client
		if err != nil {
			log.Printf("ListTransactionsCount RPC Error: %s", err)
			break
		} else {
			txn := data[0]
			for {

				TxnArray, err := client.ListTransactionsCount("", 100)
				if err != nil {
					log.Printf("ListTransactionsCount RPC Error: %s", err)
					break
				}

				log.Println("Checking Transactions")
				var TxnFound bool = true

				for i := range TxnArray {
					if Serialize(TxnArray[i]) == Serialize(txn) {
						TxnFound = true
						break
					} else {

						TxnFound = false

					}

				}

				if TxnFound == false {
					EndTime := time.Now().Unix()
					delta := float64((StartTime - EndTime))

					if delta == 0 {
						delta = float64(0.00000001) // at least a satoshi
					}
					tps := 100.0 / delta
					log.Printf("Transactions per second: %f", tps)
					break
				}
			}
		}
	}

	// shutdown with ctrl-c
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		for sig := range c {
			log.Printf("Client shutdown: %v, stopping proc and exiting..", sig)

			time.AfterFunc(time.Second*1, func() {
				log.Println("Client shutting down...")
				client.Shutdown()
				log.Println("Client shutdown complete.")
			})
		}
		os.Exit(1)
	}()

	client.WaitForShutdown()
}
<commit_msg>wait for btcd<commit_after>// Copyright (c) 2014 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"github.com/conformal/btcjson"
	"io/ioutil"
	"log"
	"os"
	//"os/exec"
	"os/signal"
	"path/filepath"
	"time"

	"github.com/conformal/btcrpcclient"
	"github.com/conformal/btcutil"
	//"github.com/davecgh/go-spew/spew"
)

func Serialize(item btcjson.ListTransactionsResult) string {
	amt := fmt.Sprintf("%.8f", item.Amount)
	fee := fmt.Sprintf("%.8f", item.Fee)
	return item.Address + "," + amt + "," + fee
}

func main() {
	//var cmd *exec.Cmd
	connected := make(chan struct{})
	var firstConn bool
	// based off of btcwebsocket example for reference
	ntfnHandlers := btcrpcclient.NotificationHandlers{
		OnBtcdConnected: func(conn bool) {
			if conn && !firstConn {
				firstConn = true
				connected <- struct{}{}
			}
		},
	}

	// Connect to local btcwallet RPC server using websockets.
	certHomeDir := btcutil.AppDataDir("btcwallet", false)
	certs, err := ioutil.ReadFile(filepath.Join(certHomeDir, "rpc.cert"))
	if err != nil {
		log.Fatal(err)
	}

	connCfg := &btcrpcclient.ConnConfig{
		Host:         "localhost:18554",
		Endpoint:     "ws",
		User:         "rpcuser",
		Pass:         "rpcpass",
		Certificates: certs,
	}
	client, err := btcrpcclient.New(connCfg, &ntfnHandlers)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Waiting for BTCD")
	<-connected

	// Unlock wallet.
	log.Println("Unlock Wallet")
	if err := client.WalletPassphrase("walletpass", 3600); err != nil {
		log.Printf("%s: Cannot unlock wallet: %v", "localhost:18556", err)

	}

	// count transactions
	client.NotifyNewTransactions(true)
	log.Println("Monitoring TXns")
	for {

		StartTime := time.Now().Unix()
		data, err := client.ListTransactionsCount("", 100)

		// Client
		if err != nil {
			log.Printf("ListTransactionsCount RPC Error: %s", err)
			break
		} else {
			txn := data[0]
			for {

				TxnArray, err := client.ListTransactionsCount("", 100)
				if err != nil {
					log.Printf("ListTransactionsCount RPC Error: %s", err)
					break
				}

				log.Println("Checking Transactions")
				var TxnFound bool = true

				for i := range TxnArray {
					if Serialize(TxnArray[i]) == Serialize(txn) {
						TxnFound = true
						break
					} else {

						TxnFound = false

					}

				}

				if TxnFound == false {
					EndTime := time.Now().Unix()
					delta := float64((StartTime - EndTime))

					if delta == 0 {
						delta = float64(0.00000001) // at least a satoshi
					}
					tps := 100.0 / delta
					log.Printf("Transactions per second: %f", tps)
					break
				}
			}
		}
	}

	// shutdown with ctrl-c
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		for sig := range c {
			log.Printf("Client shutdown: %v, stopping proc and exiting..", sig)

			time.AfterFunc(time.Second*1, func() {
				log.Println("Client shutting down...")
				client.Shutdown()
				log.Println("Client shutdown complete.")
			})
		}
		os.Exit(1)
	}()

	client.WaitForShutdown()
}
<|endoftext|>


<commit_before>package analytics

import (
	"net/http"
	"runtime"

	"github.com/getlantern/flashlight/config"
	"github.com/mitchellh/mapstructure"

	"github.com/getlantern/analytics"
	"github.com/getlantern/flashlight/ui"
	"github.com/getlantern/golog"
)

const (
	messageType = `Analytics`
)

var (
	log        = golog.LoggerFor("flashlight.analytics")
	service    *ui.Service
	httpClient *http.Client
	hostName   *string
	stopCh     chan bool
)

func Configure(cfg *config.Config, serverSession bool, newClient *http.Client) {

	httpClient = newClient

	sessionPayload := &analytics.Payload{
		ClientId:      cfg.InstanceId,
		ClientVersion: string(cfg.Version),
		HitType:       analytics.EventType,
		Event: &analytics.Event{
			Category: "Session",
			Action:   "Start",
			Label:    runtime.GOOS,
		},
	}

	if cfg == nil {
		analytics.SessionEvent(httpClient, sessionPayload)
		return
	}

	if cfg.InstanceId != "" {
		sessionPayload.ClientId = cfg.InstanceId
	}
	if cfg.Version != 0 {
		sessionPayload.ClientVersion = string(cfg.Version)
	}

	if serverSession {
		sessionPayload.Hostname = cfg.Server.RegisterAt
	} else {
		sessionPayload.Hostname = "localhost"
	}

	analytics.SessionEvent(httpClient, sessionPayload)

	if !serverSession && cfg.AutoReport != nil && *cfg.AutoReport {
		err := StartService()
		if err != nil {
			log.Errorf("Error starting analytics service: %q", err)
		}
	}
}

// Used with clients to track user interaction with the UI
func StartService() error {

	var err error

	if service != nil {
		return nil
	}

	newMessage := func() interface{} {
		return &analytics.Payload{}
	}

	if service, err = ui.Register(messageType, newMessage, nil); err != nil {
		log.Errorf("Unable to register analytics service: %q", err)
		return err
	}

	stopCh = make(chan bool)

	// process analytics messages
	go read()

	return err
}

func StopService() {
	if service != nil && stopCh != nil {
		ui.Unregister(messageType)
		stopCh <- true
		service = nil
		log.Debug("Successfully stopped analytics service")
	}
}

func read() {

	for {
		select {
		case <-stopCh:
			return
		case msg := <-service.In:
			log.Debugf("New UI analytics message: %q", msg)
			var payload analytics.Payload
			if err := mapstructure.Decode(msg, &payload); err != nil {
				log.Errorf("Could not decode payload: %q", err)
			} else {
				// set to localhost on clients
				payload.Hostname = "localhost"
				payload.HitType = analytics.PageViewType
				// for now, the only analytics messages we are
				// currently receiving from the UI are initial page
				// views which indicate new UI sessions
				analytics.UIEvent(httpClient, &payload)
			}
		}
	}
}
<commit_msg>fix unit tests<commit_after>package analytics

import (
	"net/http"
	"runtime"

	"github.com/getlantern/flashlight/config"
	"github.com/mitchellh/mapstructure"

	"github.com/getlantern/analytics"
	"github.com/getlantern/flashlight/ui"
	"github.com/getlantern/golog"
)

const (
	messageType = `Analytics`
)

var (
	log        = golog.LoggerFor("flashlight.analytics")
	service    *ui.Service
	httpClient *http.Client
	hostName   *string
	stopCh     chan bool
)

func Configure(cfg *config.Config, serverSession bool, newClient *http.Client) {

	httpClient = newClient

	sessionPayload := &analytics.Payload{
		HitType: analytics.EventType,
		Event: &analytics.Event{
			Category: "Session",
			Action:   "Start",
			Label:    runtime.GOOS,
		},
	}

	if cfg == nil {
		analytics.SessionEvent(httpClient, sessionPayload)
		return
	}

	if cfg.InstanceId != "" {
		sessionPayload.ClientId = cfg.InstanceId
	}
	if cfg.Version != 0 {
		sessionPayload.ClientVersion = string(cfg.Version)
	}

	if serverSession {
		sessionPayload.Hostname = cfg.Server.RegisterAt
	} else {
		sessionPayload.Hostname = "localhost"
	}

	analytics.SessionEvent(httpClient, sessionPayload)

	if !serverSession && cfg.AutoReport != nil && *cfg.AutoReport {
		err := StartService()
		if err != nil {
			log.Errorf("Error starting analytics service: %q", err)
		}
	}
}

// Used with clients to track user interaction with the UI
func StartService() error {

	var err error

	if service != nil {
		return nil
	}

	newMessage := func() interface{} {
		return &analytics.Payload{}
	}

	if service, err = ui.Register(messageType, newMessage, nil); err != nil {
		log.Errorf("Unable to register analytics service: %q", err)
		return err
	}

	stopCh = make(chan bool)

	// process analytics messages
	go read()

	return err
}

func StopService() {
	if service != nil && stopCh != nil {
		ui.Unregister(messageType)
		stopCh <- true
		service = nil
		log.Debug("Successfully stopped analytics service")
	}
}

func read() {

	for {
		select {
		case <-stopCh:
			return
		case msg := <-service.In:
			log.Debugf("New UI analytics message: %q", msg)
			var payload analytics.Payload
			if err := mapstructure.Decode(msg, &payload); err != nil {
				log.Errorf("Could not decode payload: %q", err)
			} else {
				// set to localhost on clients
				payload.Hostname = "localhost"
				payload.HitType = analytics.PageViewType
				// for now, the only analytics messages we are
				// currently receiving from the UI are initial page
				// views which indicate new UI sessions
				analytics.UIEvent(httpClient, &payload)
			}
		}
	}
}
<|endoftext|>


<commit_before>package dokugen

import (
	"strings"
	"testing"
)

func TestCellCreation(t *testing.T) {
	data := "1"
	cell := NewCell(nil, 0, 0, data)
	if cell.Number != 1 {
		t.Log("Number came back wrong")
		t.Fail()
	}
	if cell.Row != 0 {
		t.Log("Row came back wrong")
		t.Fail()
	}
	if cell.Col != 0 {
		t.Log("Cell came back wrong")
		t.Fail()
	}
	if cell.DataString() != data {
		t.Log("Cell round-tripped out with different string than data in")
		t.Fail()
	}
	//TODO: test failing for values that are too high.
}

func TestGridCreation(t *testing.T) {

	blockUpperLeftRow := make([]int, DIM)
	blockUpperLeftCol := make([]int, DIM)

	row := 0
	col := -1 * BLOCK_DIM

	for i := 0; i < DIM; i++ {
		col += BLOCK_DIM
		if col >= DIM {
			col = 0
			row += BLOCK_DIM
		}
		blockUpperLeftRow[i] = row
		blockUpperLeftCol[i] = col
	}

	cellData := "1"
	rowData := strings.Join(nCopies(cellData, DIM), COL_SEP)
	data := strings.Join(nCopies(rowData, DIM), ROW_SEP)
	grid := NewGrid(data)
	if len(grid.cells) != DIM*DIM {
		t.Log("Didn't generate enough cells")
		t.Fail()
	}
	if grid.DataString() != data {
		t.Log("The grid round-tripped with different result than data in")
		t.Fail()
	}
	if grid.cells[10].Number != 1 {
		t.Log("A random spot check of a cell had the wrong number: %s", grid.cells[10])
		t.Fail()
	}

	for count := 0; count < DIM; count++ {
		col := grid.Col(count)
		if num := len(col); num != DIM {
			t.Log("We got back a column but it had the wrong amount of items: ", num, "\n")
			t.Fail()
		}
		for i, cell := range col {
			if cell.Col != count {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right column.")
				t.Fail()
			}
			if cell.Row != i {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right row.")
				t.Fail()
			}
		}

		row := grid.Row(count)
		if len(row) != DIM {
			t.Log("We got back a row but it had the wrong number of items.")
			t.Fail()
		}
		for i, cell := range row {
			if cell.Row != count {
				t.Log("One of the cells we got back when asking for row ", count, " was not in the right rows.")
				t.Fail()
			}
			if cell.Col != i {
				t.Log("One of the cells we got back from row ", count, " was not in the right column.")
				t.Fail()
			}
		}

		block := grid.Block(count)
		if len(block) != DIM {
			t.Log("We got back a block but it had the wrong number of items.")
			t.Fail()
		}

		for _, cell := range block {
			if cell.Block != count {
				t.Log("We got a cell back in a block with the wrong block number")
				t.Fail()
			}
		}

		if block[0].Row != blockUpperLeftRow[count] || block[0].Col != blockUpperLeftCol[count] {
			t.Log("We got back the wrong first cell from block ", count, ": ", block[0])
			t.Fail()
		}

		if block[DIM-1].Row != blockUpperLeftRow[count]+BLOCK_DIM-1 || block[DIM-1].Col != blockUpperLeftCol[count]+BLOCK_DIM-1 {
			t.Log("We got back the wrong last cell from block ", count, ": ", block[0])
			t.Fail()
		}

	}

	//TODO: test neighbors list is set correctly.

	cell := grid.Cell(2, 2)

	if cell.Row != 2 || cell.Col != 2 {
		t.Log("We grabbed a cell but what we got back was the wrong row and col.")
		t.Fail()
	}

	neighbors := cell.Neighbors()

	if len(neighbors) != (DIM-1)*3-(BLOCK_DIM-1)*2 {
		t.Log("We got a different number of neighbors than what we were expecting: ", len(neighbors))
		t.Fail()
	}
	for _, neighbor := range neighbors {
		if neighbor == nil {
			t.Log("We found a nil neighbor")
			t.Fail()
		}
		if neighbor.Row != cell.Row && neighbor.Col != cell.Col && neighbor.Block != cell.Block {
			t.Log("We found a neighbor in ourselves that doesn't appear to be related: Neighbor: ", neighbor, " Cell: ", cell)
			t.Fail()
		}
		//TODO: Check to make sure we don't get duplicates in this list.
	}

}

func nCopies(in string, copies int) (result []string) {
	for i := 0; i < copies; i++ {
		result = append(result, in)
	}
	return
}
<commit_msg>Added a test to ensure no duplicates are in the neighbor list.<commit_after>package dokugen

import (
	"strings"
	"testing"
)

func TestCellCreation(t *testing.T) {
	data := "1"
	cell := NewCell(nil, 0, 0, data)
	if cell.Number != 1 {
		t.Log("Number came back wrong")
		t.Fail()
	}
	if cell.Row != 0 {
		t.Log("Row came back wrong")
		t.Fail()
	}
	if cell.Col != 0 {
		t.Log("Cell came back wrong")
		t.Fail()
	}
	if cell.DataString() != data {
		t.Log("Cell round-tripped out with different string than data in")
		t.Fail()
	}
	//TODO: test failing for values that are too high.
}

func TestGridCreation(t *testing.T) {

	blockUpperLeftRow := make([]int, DIM)
	blockUpperLeftCol := make([]int, DIM)

	row := 0
	col := -1 * BLOCK_DIM

	for i := 0; i < DIM; i++ {
		col += BLOCK_DIM
		if col >= DIM {
			col = 0
			row += BLOCK_DIM
		}
		blockUpperLeftRow[i] = row
		blockUpperLeftCol[i] = col
	}

	cellData := "1"
	rowData := strings.Join(nCopies(cellData, DIM), COL_SEP)
	data := strings.Join(nCopies(rowData, DIM), ROW_SEP)
	grid := NewGrid(data)
	if len(grid.cells) != DIM*DIM {
		t.Log("Didn't generate enough cells")
		t.Fail()
	}
	if grid.DataString() != data {
		t.Log("The grid round-tripped with different result than data in")
		t.Fail()
	}
	if grid.cells[10].Number != 1 {
		t.Log("A random spot check of a cell had the wrong number: %s", grid.cells[10])
		t.Fail()
	}

	for count := 0; count < DIM; count++ {
		col := grid.Col(count)
		if num := len(col); num != DIM {
			t.Log("We got back a column but it had the wrong amount of items: ", num, "\n")
			t.Fail()
		}
		for i, cell := range col {
			if cell.Col != count {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right column.")
				t.Fail()
			}
			if cell.Row != i {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right row.")
				t.Fail()
			}
		}

		row := grid.Row(count)
		if len(row) != DIM {
			t.Log("We got back a row but it had the wrong number of items.")
			t.Fail()
		}
		for i, cell := range row {
			if cell.Row != count {
				t.Log("One of the cells we got back when asking for row ", count, " was not in the right rows.")
				t.Fail()
			}
			if cell.Col != i {
				t.Log("One of the cells we got back from row ", count, " was not in the right column.")
				t.Fail()
			}
		}

		block := grid.Block(count)
		if len(block) != DIM {
			t.Log("We got back a block but it had the wrong number of items.")
			t.Fail()
		}

		for _, cell := range block {
			if cell.Block != count {
				t.Log("We got a cell back in a block with the wrong block number")
				t.Fail()
			}
		}

		if block[0].Row != blockUpperLeftRow[count] || block[0].Col != blockUpperLeftCol[count] {
			t.Log("We got back the wrong first cell from block ", count, ": ", block[0])
			t.Fail()
		}

		if block[DIM-1].Row != blockUpperLeftRow[count]+BLOCK_DIM-1 || block[DIM-1].Col != blockUpperLeftCol[count]+BLOCK_DIM-1 {
			t.Log("We got back the wrong last cell from block ", count, ": ", block[0])
			t.Fail()
		}

	}

	cell := grid.Cell(2, 2)

	if cell.Row != 2 || cell.Col != 2 {
		t.Log("We grabbed a cell but what we got back was the wrong row and col.")
		t.Fail()
	}

	neighbors := cell.Neighbors()

	if len(neighbors) != (DIM-1)*3-(BLOCK_DIM-1)*2 {
		t.Log("We got a different number of neighbors than what we were expecting: ", len(neighbors))
		t.Fail()
	}
	neighborsMap := make(map[*Cell]bool)
	for _, neighbor := range neighbors {
		if neighbor == nil {
			t.Log("We found a nil neighbor")
			t.Fail()
		}
		if neighbor.Row != cell.Row && neighbor.Col != cell.Col && neighbor.Block != cell.Block {
			t.Log("We found a neighbor in ourselves that doesn't appear to be related: Neighbor: ", neighbor, " Cell: ", cell)
			t.Fail()
		}
		if _, ok := neighborsMap[neighbor]; ok {
			t.Log("We found a duplicate in the neighbors list")
			t.Fail()
		} else {
			neighborsMap[cell] = true
		}
	}

}

func nCopies(in string, copies int) (result []string) {
	for i := 0; i < copies; i++ {
		result = append(result, in)
	}
	return
}
<|endoftext|>


<commit_before>package mpb

import (
	"fmt"
	"io"
	"os"
	"sort"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

type opType uint

const (
	opBarAdd opType = iota
	opBarRemove
)

type SortType uint

const (
	SortNone SortType = iota
	SortTop
	SortBottom
)

const refreshRate = 100

// Progress represents the container that renders Progress bars
type Progress struct {
	out     io.Writer
	width   int
	sort    SortType
	stopped bool

	op             chan *operation
	rrChangeReqCh  chan time.Duration
	outChangeReqCh chan io.Writer

	wg *sync.WaitGroup
}

type operation struct {
	kind   opType
	bar    *Bar
	result chan bool
}

// New returns a new progress bar with defaults
func New() *Progress {
	p := &Progress{
		width:          70,
		op:             make(chan *operation),
		rrChangeReqCh:  make(chan time.Duration),
		outChangeReqCh: make(chan io.Writer),
		wg:             new(sync.WaitGroup),
	}
	go p.server(cwriter.New(os.Stdout))
	return p
}

func (p *Progress) SetWidth(n int) *Progress {
	if n <= 0 {
		return p
	}
	p.width = n
	return p
}

// SetOut sets underlying writer of progress
// default is os.Stdout
func (p *Progress) SetOut(w io.Writer) *Progress {
	if w == nil {
		return p
	}
	p.outChangeReqCh <- w
	return p
}

// RefreshRate overrides default (30ms) refreshRate value
func (p *Progress) RefreshRate(d time.Duration) *Progress {
	p.rrChangeReqCh <- d
	return p
}

func (p *Progress) WithSort(sort SortType) *Progress {
	p.sort = sort
	return p
}

// AddBar creates a new progress bar and adds to the container
func (p *Progress) AddBar(total int) *Bar {
	p.wg.Add(1)
	bar := newBar(total, p.width, p.wg)
	p.op <- &operation{opBarAdd, bar, nil}
	return bar
}

func (p *Progress) RemoveBar(b *Bar) bool {
	result := make(chan bool)
	p.op <- &operation{opBarRemove, b, result}
	return <-result
}

// WaitAndStop stops listening
func (p *Progress) WaitAndStop() {
	if !p.stopped {
		// fmt.Fprintln(os.Stderr, "p.WaitAndStop")
		p.stopped = true
		p.wg.Wait()
		close(p.op)
	}
}

// server monitors underlying channels and renders any progress bars
func (p *Progress) server(cw *cwriter.Writer) {
	t := time.NewTicker(refreshRate * time.Millisecond)
	bars := make([]*Bar, 0, 4)
	for {
		select {
		case w := <-p.outChangeReqCh:
			cw.Flush()
			cw = cwriter.New(w)
		case op, ok := <-p.op:
			if !ok {
				// fmt.Fprintln(os.Stderr, "Sopping bars")
				for _, b := range bars {
					b.Stop()
				}
				t.Stop()
				return
			}
			switch op.kind {
			case opBarAdd:
				bars = append(bars, op.bar)
			case opBarRemove:
				var ok bool
				for i, b := range bars {
					if b == op.bar {
						bars = append(bars[:i], bars[i+1:]...)
						ok = true
						b.Stop()
						break
					}
				}
				op.result <- ok
			}
		case <-t.C:
			switch p.sort {
			case SortTop:
				sort.Sort(sort.Reverse(SortableBarSlice(bars)))
			case SortBottom:
				sort.Sort(SortableBarSlice(bars))
			}
			for _, b := range bars {
				// cannot parallel this, because order matters
				fmt.Fprintln(cw, b)
			}
			cw.Flush()
			for _, b := range bars {
				go func(b *Bar) {
					b.flushedCh <- struct{}{}
				}(b)
			}
		case d := <-p.rrChangeReqCh:
			t.Stop()
			t = time.NewTicker(d)
		}
	}
}
<commit_msg>refactor const RefreshRate<commit_after>package mpb

import (
	"fmt"
	"io"
	"os"
	"sort"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

type opType uint

const (
	opBarAdd opType = iota
	opBarRemove
)

type SortType uint

const (
	SortNone SortType = iota
	SortTop
	SortBottom
)

// RefreshRate
const rr = 100

// Progress represents the container that renders Progress bars
type Progress struct {
	out     io.Writer
	width   int
	sort    SortType
	stopped bool

	op             chan *operation
	rrChangeReqCh  chan time.Duration
	outChangeReqCh chan io.Writer

	wg *sync.WaitGroup
}

type operation struct {
	kind   opType
	bar    *Bar
	result chan bool
}

// New returns a new progress bar with defaults
func New() *Progress {
	p := &Progress{
		width:          70,
		op:             make(chan *operation),
		rrChangeReqCh:  make(chan time.Duration),
		outChangeReqCh: make(chan io.Writer),
		wg:             new(sync.WaitGroup),
	}
	go p.server(cwriter.New(os.Stdout), time.NewTicker(rr*time.Millisecond))
	return p
}

func (p *Progress) SetWidth(n int) *Progress {
	if n <= 0 {
		return p
	}
	p.width = n
	return p
}

// SetOut sets underlying writer of progress
// default is os.Stdout
func (p *Progress) SetOut(w io.Writer) *Progress {
	if w == nil {
		return p
	}
	p.outChangeReqCh <- w
	return p
}

// RefreshRate overrides default (30ms) refreshRate value
func (p *Progress) RefreshRate(d time.Duration) *Progress {
	p.rrChangeReqCh <- d
	return p
}

func (p *Progress) WithSort(sort SortType) *Progress {
	p.sort = sort
	return p
}

// AddBar creates a new progress bar and adds to the container
func (p *Progress) AddBar(total int) *Bar {
	p.wg.Add(1)
	bar := newBar(total, p.width, p.wg)
	p.op <- &operation{opBarAdd, bar, nil}
	return bar
}

func (p *Progress) RemoveBar(b *Bar) bool {
	result := make(chan bool)
	p.op <- &operation{opBarRemove, b, result}
	return <-result
}

// WaitAndStop stops listening
func (p *Progress) WaitAndStop() {
	if !p.stopped {
		// fmt.Fprintln(os.Stderr, "p.WaitAndStop")
		p.stopped = true
		p.wg.Wait()
		close(p.op)
	}
}

// server monitors underlying channels and renders any progress bars
func (p *Progress) server(cw *cwriter.Writer, t *time.Ticker) {
	bars := make([]*Bar, 0, 4)
	for {
		select {
		case w := <-p.outChangeReqCh:
			cw.Flush()
			cw = cwriter.New(w)
		case op, ok := <-p.op:
			if !ok {
				// fmt.Fprintln(os.Stderr, "Sopping bars")
				for _, b := range bars {
					b.Stop()
				}
				t.Stop()
				return
			}
			switch op.kind {
			case opBarAdd:
				bars = append(bars, op.bar)
			case opBarRemove:
				var ok bool
				for i, b := range bars {
					if b == op.bar {
						bars = append(bars[:i], bars[i+1:]...)
						ok = true
						b.Stop()
						break
					}
				}
				op.result <- ok
			}
		case <-t.C:
			switch p.sort {
			case SortTop:
				sort.Sort(sort.Reverse(SortableBarSlice(bars)))
			case SortBottom:
				sort.Sort(SortableBarSlice(bars))
			}
			for _, b := range bars {
				// cannot parallel this, because order matters
				fmt.Fprintln(cw, b)
			}
			cw.Flush()
			for _, b := range bars {
				go func(b *Bar) {
					b.flushedCh <- struct{}{}
				}(b)
			}
		case d := <-p.rrChangeReqCh:
			t.Stop()
			t = time.NewTicker(d)
		}
	}
}
<|endoftext|>