<commit_before>package mozillaGradingWorker

import (
	"encoding/json"
	"fmt"

	"github.com/mozilla/tls-observatory/connection"
	"github.com/mozilla/tls-observatory/logger"
	"github.com/mozilla/tls-observatory/worker"
)

var workerName = "mozillaGradingWorker"
var workerDesc = `The grading worker provides an SSLabs-like grade for the TLS configuration of
the audited target`

// EvaluationResults contains the results of the mozillaEvaluationWorker
type EvaluationResults struct {
	Grade       float64  `json:"grade"`
	LetterGrade string   `json:"lettergrade"`
	Failures    []string `json:"failures"`
}

type categoryResults struct {
	Grade          int
	MaximumAllowed int
	Remarks        []string
}

// CipherSuite represent a ciphersuite generated and recognised by OpenSSL
type CipherSuite struct {
	Proto string     `json:"proto"`
	Kx    string     `json:"kx"`
	Au    string     `json:"au"`
	Enc   Encryption `json:"encryption"`
	Mac   string     `json:"mac"`
}

//Encryption represents the encryption aspects of a Ciphersuite
type Encryption struct {
	Cipher string `json:"cipher"`
	Bits   int    `json:"key"`
}

type eval struct {
}

var opensslciphersuites = make(map[string]CipherSuite)
var log = logger.GetLogger()

func init() {
	log.Debug("Registering Grading...")
	err := json.Unmarshal([]byte(OpenSSLCiphersuites), &opensslciphersuites)
	if err != nil {
		log.Error(err)
		log.Error("Could not load OpenSSL ciphersuites. Grading Worker not available")
		return
	}
	worker.RegisterWorker(workerName, worker.Info{Runner: new(eval), Description: workerDesc})
}

//go:generate  go run ../../tools/ciphers.go
// Run implements the worker interface.It is called to get the worker results.
func (e eval) Run(in worker.Input, resChan chan worker.Result) {

	res := worker.Result{WorkerName: workerName}

	b, err := Evaluate(in.Connection)
	if err != nil {
		res.Success = false
		res.Errors = append(res.Errors, err.Error())
	} else {
		res.Result = b
		res.Success = true
	}

	resChan <- res
}

// Evaluate runs compliance checks of the provided json Stored connection and returns the results
func Evaluate(connInfo connection.Stored) ([]byte, error) {
	protores, err := gradeProtocol(connInfo)
	if err != nil {
		return nil, err
	}
	cipherres, err := gradeCiphers(connInfo)
	if err != nil {
		return nil, err
	}
	keyxres, err := gradeKeyX(connInfo)
	if err != nil {
		return nil, err
	}

	var score float64
	score = float64(protores.Grade)*0.3 + float64(cipherres.Grade)*0.4 + float64(keyxres.Grade)*0.3

	// fmt.Printf("proto : %d , cipher : %d , keyx: %d\n", int(protores.Grade), int(cipherres.Grade), int(keyxres.Grade))

	er := EvaluationResults{Grade: score, LetterGrade: getLetterfromGrade(score)}
	return json.Marshal(&er)
}

func getLetterfromGrade(grade float64) string {
	if grade < 20 {
		return "F"
	} else if grade < 35 {
		return "E"
	} else if grade < 50 {
		return "D"
	} else if grade < 65 {
		return "C"
	} else if grade < 80 {
		return "B"
	}

	return "A"
}

func (e eval) AnalysisPrinter(r []byte, targetLevel interface{}) (results []string, err error) {
	var eval EvaluationResults
	err = json.Unmarshal(r, &eval)
	if err != nil {
		err = fmt.Errorf("Mozilla grading worker: failed to parse results: %v", err)
		return
	}
	results = append(results, fmt.Sprintf("* Grade: %f", eval.Grade))
	for _, e := range eval.Failures {
		results = append(results, fmt.Sprintf("  - %s", e))
	}
	return
}

func (e eval) Assertor(evresults, assertresults []byte) (pass bool, body []byte, err error) {
	var evres, assertres EvaluationResults
	err = json.Unmarshal(evresults, &evres)
	if err != nil {
		return
	}
	err = json.Unmarshal(assertresults, &assertres)
	if err != nil {
		return
	}
	if evres.Grade != assertres.Grade {
		body = []byte(fmt.Sprintf(`Assertion mozillaGradingWorker. The domain scored %f instead of expected %f`,
			assertres.Grade, evres.Grade))
		pass = false
	} else {
		pass = true
	}
	return
}

// contains checks if an entry exists in a slice and returns
// a booleans.
func contains(slice []string, entry string) bool {
	for _, element := range slice {
		if element == entry {
			return true
		}
	}
	return false
}
<commit_msg>change go:generate command to support running from dockerfile<commit_after>package mozillaGradingWorker

import (
	"encoding/json"
	"fmt"

	"github.com/mozilla/tls-observatory/connection"
	"github.com/mozilla/tls-observatory/logger"
	"github.com/mozilla/tls-observatory/worker"
)

var workerName = "mozillaGradingWorker"
var workerDesc = `The grading worker provides an SSLabs-like grade for the TLS configuration of
the audited target`

// EvaluationResults contains the results of the mozillaEvaluationWorker
type EvaluationResults struct {
	Grade       float64  `json:"grade"`
	LetterGrade string   `json:"lettergrade"`
	Failures    []string `json:"failures"`
}

type categoryResults struct {
	Grade          int
	MaximumAllowed int
	Remarks        []string
}

// CipherSuite represent a ciphersuite generated and recognised by OpenSSL
type CipherSuite struct {
	Proto string     `json:"proto"`
	Kx    string     `json:"kx"`
	Au    string     `json:"au"`
	Enc   Encryption `json:"encryption"`
	Mac   string     `json:"mac"`
}

//Encryption represents the encryption aspects of a Ciphersuite
type Encryption struct {
	Cipher string `json:"cipher"`
	Bits   int    `json:"key"`
}

type eval struct {
}

var opensslciphersuites = make(map[string]CipherSuite)
var log = logger.GetLogger()

func init() {
	err := json.Unmarshal([]byte(OpenSSLCiphersuites), &opensslciphersuites)
	if err != nil {
		log.Error(err)
		log.Error("Could not load OpenSSL ciphersuites. Grading Worker not available")
		return
	}
	worker.RegisterWorker(workerName, worker.Info{Runner: new(eval), Description: workerDesc})
}

//go:generate  go run /go/src/$PROJECT/tools/extractCiphersuites.go
// Run implements the worker interface.It is called to get the worker results.
func (e eval) Run(in worker.Input, resChan chan worker.Result) {

	res := worker.Result{WorkerName: workerName}

	b, err := Evaluate(in.Connection)
	if err != nil {
		res.Success = false
		res.Errors = append(res.Errors, err.Error())
	} else {
		res.Result = b
		res.Success = true
	}

	resChan <- res
}

// Evaluate runs compliance checks of the provided json Stored connection and returns the results
func Evaluate(connInfo connection.Stored) ([]byte, error) {
	protores, err := gradeProtocol(connInfo)
	if err != nil {
		return nil, err
	}
	cipherres, err := gradeCiphers(connInfo)
	if err != nil {
		return nil, err
	}
	keyxres, err := gradeKeyX(connInfo)
	if err != nil {
		return nil, err
	}

	var score float64
	score = float64(protores.Grade)*0.3 + float64(cipherres.Grade)*0.4 + float64(keyxres.Grade)*0.3

	// fmt.Printf("proto : %d , cipher : %d , keyx: %d\n", int(protores.Grade), int(cipherres.Grade), int(keyxres.Grade))

	er := EvaluationResults{Grade: score, LetterGrade: getLetterfromGrade(score)}
	return json.Marshal(&er)
}

func getLetterfromGrade(grade float64) string {
	if grade < 20 {
		return "F"
	} else if grade < 35 {
		return "E"
	} else if grade < 50 {
		return "D"
	} else if grade < 65 {
		return "C"
	} else if grade < 80 {
		return "B"
	}

	return "A"
}

func (e eval) AnalysisPrinter(r []byte, targetLevel interface{}) (results []string, err error) {
	var eval EvaluationResults
	err = json.Unmarshal(r, &eval)
	if err != nil {
		err = fmt.Errorf("Mozilla grading worker: failed to parse results: %v", err)
		return
	}
	results = append(results, fmt.Sprintf("* Grade: %f", eval.Grade))
	for _, e := range eval.Failures {
		results = append(results, fmt.Sprintf("  - %s", e))
	}
	return
}

func (e eval) Assertor(evresults, assertresults []byte) (pass bool, body []byte, err error) {
	var evres, assertres EvaluationResults
	err = json.Unmarshal(evresults, &evres)
	if err != nil {
		return
	}
	err = json.Unmarshal(assertresults, &assertres)
	if err != nil {
		return
	}
	if evres.Grade != assertres.Grade {
		body = []byte(fmt.Sprintf(`Assertion mozillaGradingWorker. The domain scored %f instead of expected %f`,
			assertres.Grade, evres.Grade))
		pass = false
	} else {
		pass = true
	}
	return
}

// contains checks if an entry exists in a slice and returns
// a booleans.
func contains(slice []string, entry string) bool {
	for _, element := range slice {
		if element == entry {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package memosort

import "testing"

func TestStableRundom(t *testing.T) {
}
<commit_msg>add initial slice for memosort tests<commit_after>package memosort

import (
	"testing"
	"time"
)

type Track struct {
	Title  string
	Artist string
	Album  string
	Year   int
	Length time.Duration
}

var tracks = []*Track{
	{"Go", "Delilah", "From the Roots Up", 2012, length("3m38s")},
	{"Go", "Moby", "Moby", 1992, length("3m37s")},
	{"Go Ahead", "Alicia Keys", "As I Am", 2007, length("4m36s")},
	{"Ready 2 Go", "Martin Solveig", "Smash", 2011, length("4m24s")},
	{"Bohemian Rhapsody", "Queen", "A Night at the Opera", 1975, length("6m6s")},
	{"Smells Like Teen Spirit", "Nirvana", "Nevermind", 1972, length("4m37s")},
	{"Imagine", "John Lennon", "Imagine", 1971, length("3m54s")},
	{"Hotel California", "Eagles", "Hotel California", 1976, length("6m40s")},
	{"One", "Metallica", "And Justice for All", 1988, length("7m23s")},
	{"Comfortably Numb", "Pink Floyd", "The Wall", 1979, length("6m53s")},
	{"Like a Rolling Stone", "Bob Dylan", "Highway 61 Revisited", 1965, length("6m20s")},
}

func length(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(s)
	}
	return d
}

func TestStableSimple(t *testing.T) {

}
<|endoftext|>


<commit_before>package memqueue_test

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-redis/redis"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	uuid "github.com/satori/go.uuid"

	"github.com/vmihailenco/taskq/v2"
	"github.com/vmihailenco/taskq/v2/memqueue"
)

func TestMemqueue(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "memqueue")
}

var _ = BeforeSuite(func() {
	taskq.SetLogger(log.New(ioutil.Discard, "", 0))
})

var _ = BeforeEach(func() {
	taskq.Tasks.Reset()
})

var _ = Describe("message with args", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(s string, i int) {
				Expect(s).To(Equal("string"))
				Expect(i).To(Equal(42))
				ch <- true
			},
		})
		err := q.Add(task.WithArgs(ctx, "string", 42))
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is called with args", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("context.Context", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(c context.Context, s string, i int) {
				Expect(s).To(Equal("string"))
				Expect(i).To(Equal(42))
				ch <- true
			},
		})
		err := q.Add(task.WithArgs(ctx, "string", 42))
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is called with args", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("message with invalid number of args", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(s string) {
				ch <- true
			},
			RetryLimit: 1,
		})
		q.Consumer().Stop()

		err := q.Add(task.WithArgs(ctx))
		Expect(err).NotTo(HaveOccurred())

		err = q.Consumer().ProcessOne(ctx)
		Expect(err).To(MatchError("taskq: got 0 args, wanted 1"))

		err = q.Consumer().ProcessAll(ctx)
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is not called", func() {
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("HandlerFunc", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(msg *taskq.Message) error {
				Expect(msg.Args).To(Equal([]interface{}{"string", 42}))
				ch <- true
				return nil
			},
		})

		err := q.Add(task.WithArgs(ctx, "string", 42))
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("is called with Message", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("message retry timing", func() {
	ctx := context.Background()
	var q *memqueue.Queue
	var task *taskq.Task
	backoff := 100 * time.Millisecond
	var count int
	var ch chan time.Time

	BeforeEach(func() {
		count = 0
		ch = make(chan time.Time, 10)
		q = memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task = taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() error {
				ch <- time.Now()
				count++
				return fmt.Errorf("fake error #%d", count)
			},
			RetryLimit: 3,
			MinBackoff: backoff,
		})
	})

	Context("without delay", func() {
		var now time.Time

		BeforeEach(func() {
			now = time.Now()
			_ = q.Add(task.WithArgs(ctx))

			err := q.Close()
			Expect(err).NotTo(HaveOccurred())
		})

		It("is retried in time", func() {
			Expect(ch).To(Receive(BeTemporally("~", now, backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(backoff), backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(3*backoff), backoff/10)))
			Expect(ch).NotTo(Receive())
		})
	})

	Context("message with delay", func() {
		var now time.Time

		BeforeEach(func() {
			msg := task.WithArgs(ctx)
			msg.Delay = 5 * backoff
			now = time.Now().Add(msg.Delay)

			q.Add(msg)

			err := q.Close()
			Expect(err).NotTo(HaveOccurred())
		})

		It("is retried in time", func() {
			Expect(ch).To(Receive(BeTemporally("~", now, backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(backoff), backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(3*backoff), backoff/10)))
			Expect(ch).NotTo(Receive())
		})
	})
})

var _ = Describe("failing queue with error handler", func() {
	ctx := context.Background()
	var q *memqueue.Queue
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q = memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() error {
				return errors.New("fake error")
			},
			FallbackHandler: func() {
				ch <- true
			},
			RetryLimit: 1,
		})
		q.Add(task.WithArgs(ctx))

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("error handler is called when handler fails", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("named message", func() {
	ctx := context.Background()
	var count int64

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name:  "test",
			Redis: redisRing(),
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() {
				atomic.AddInt64(&count, 1)
			},
		})

		name := uuid.NewV4().String()

		var wg sync.WaitGroup
		for i := 0; i < 100; i++ {
			wg.Add(1)
			go func() {
				defer GinkgoRecover()
				defer wg.Done()
				msg := task.WithArgs(ctx)
				msg.Name = name
				q.Add(msg)
			}()
		}
		wg.Wait()

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("is processed once", func() {
		n := atomic.LoadInt64(&count)
		Expect(n).To(Equal(int64(1)))
	})
})

var _ = Describe("CallOnce", func() {
	ctx := context.Background()
	var now time.Time
	delay := 3 * time.Second
	ch := make(chan time.Time, 10)

	BeforeEach(func() {
		now = time.Now()

		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name:  "test",
			Redis: redisRing(),
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(slot int64) error {
				ch <- time.Now()
				return nil
			},
		})

		var wg sync.WaitGroup
		for i := 0; i < 10; i++ {
			wg.Add(1)
			go func() {
				defer GinkgoRecover()
				defer wg.Done()

				q.Add(task.OnceWithArgs(ctx, delay, slot(delay)))
			}()
		}
		wg.Wait()

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("processes message once with delay", func() {
		Expect(ch).To(Receive(BeTemporally(">", now.Add(delay), time.Second)))
		Consistently(ch).ShouldNot(Receive())
	})
})

var _ = Describe("stress testing", func() {
	const n = 10000
	ctx := context.Background()
	var count int64

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() {
				atomic.AddInt64(&count, 1)
			},
		})

		for i := 0; i < n; i++ {
			q.Add(task.WithArgs(ctx))
		}

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is called for all messages", func() {
		nn := atomic.LoadInt64(&count)
		Expect(nn).To(Equal(int64(n)))
	})
})

var _ = Describe("stress testing failing queue", func() {
	const n = 100000
	ctx := context.Background()
	var errorCount int64

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name:                 "test",
			PauseErrorsThreshold: -1,
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() error {
				return errors.New("fake error")
			},
			FallbackHandler: func() {
				atomic.AddInt64(&errorCount, 1)
			},
			RetryLimit: 1,
		})

		for i := 0; i < n; i++ {
			q.Add(task.WithArgs(ctx))
		}

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("error handler is called for all messages", func() {
		nn := atomic.LoadInt64(&errorCount)
		Expect(nn).To(Equal(int64(n)))
	})
})

var _ = Describe("empty queue", func() {
	ctx := context.Background()
	var q *memqueue.Queue
	var task *taskq.Task
	var processed uint32

	BeforeEach(func() {
		processed = 0
		q = memqueue.NewQueue(&taskq.QueueOptions{
			Name:  "test",
			Redis: redisRing(),
		})
		task = taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() {
				atomic.AddUint32(&processed, 1)
			},
		})
	})

	AfterEach(func() {
		_ = q.Close()
	})

	It("can be closed", func() {
		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("stops processor", func() {
		err := q.Consumer().Stop()
		Expect(err).NotTo(HaveOccurred())
	})

	testEmptyQueue := func() {
		It("processes all messages", func() {
			err := q.Consumer().ProcessAll(ctx)
			Expect(err).NotTo(HaveOccurred())
		})

		It("processes one message", func() {
			err := q.Consumer().ProcessOne(ctx)
			Expect(err).To(MatchError("taskq: queue is empty"))

			err = q.Consumer().ProcessAll(ctx)
			Expect(err).NotTo(HaveOccurred())
		})
	}

	Context("when processor is stopped", func() {
		BeforeEach(func() {
			err := q.Consumer().Stop()
			Expect(err).NotTo(HaveOccurred())
		})

		testEmptyQueue()

		Context("when there are messages in the queue", func() {
			BeforeEach(func() {
				for i := 0; i < 3; i++ {
					err := q.Add(task.WithArgs(ctx))
					Expect(err).NotTo(HaveOccurred())
				}
			})

			It("processes all messages", func() {
				p := q.Consumer()

				err := p.ProcessAll(ctx)
				Expect(err).NotTo(HaveOccurred())

				n := atomic.LoadUint32(&processed)
				Expect(n).To(Equal(uint32(3)))
			})

			It("processes one message", func() {
				p := q.Consumer()

				err := p.ProcessOne(ctx)
				Expect(err).NotTo(HaveOccurred())

				n := atomic.LoadUint32(&processed)
				Expect(n).To(Equal(uint32(1)))

				err = p.ProcessAll(ctx)
				Expect(err).NotTo(HaveOccurred())

				n = atomic.LoadUint32(&processed)
				Expect(n).To(Equal(uint32(3)))
			})
		})
	})
})

// slot splits time into equal periods (called slots) and returns
// slot number for provided time.
func slot(period time.Duration) int64 {
	tm := time.Now()
	periodSec := int64(period / time.Second)
	if periodSec == 0 {
		return tm.Unix()
	}
	return tm.Unix() / periodSec
}

var (
	ringOnce sync.Once
	ring     *redis.Ring
)

func redisRing() *redis.Ring {
	ringOnce.Do(func() {
		ring = redis.NewRing(&redis.RingOptions{
			Addrs: map[string]string{"0": ":6379"},
		})
	})
	return ring
}
<commit_msg>Fix build<commit_after>package memqueue_test

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-redis/redis"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	uuid "github.com/satori/go.uuid"

	"github.com/vmihailenco/taskq/v2"
	"github.com/vmihailenco/taskq/v2/memqueue"
)

func TestMemqueue(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "memqueue")
}

var _ = BeforeSuite(func() {
	taskq.SetLogger(log.New(ioutil.Discard, "", 0))
})

var _ = BeforeEach(func() {
	taskq.Tasks.Reset()
})

var _ = Describe("message with args", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(s string, i int) {
				Expect(s).To(Equal("string"))
				Expect(i).To(Equal(42))
				ch <- true
			},
		})
		err := q.Add(task.WithArgs(ctx, "string", 42))
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is called with args", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("context.Context", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(c context.Context, s string, i int) {
				Expect(s).To(Equal("string"))
				Expect(i).To(Equal(42))
				ch <- true
			},
		})
		err := q.Add(task.WithArgs(ctx, "string", 42))
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is called with args", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("message with invalid number of args", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(s string) {
				ch <- true
			},
			RetryLimit: 1,
		})
		q.Consumer().Stop()

		err := q.Add(task.WithArgs(ctx))
		Expect(err).NotTo(HaveOccurred())

		err = q.Consumer().ProcessOne(ctx)
		Expect(err).To(MatchError("taskq: got 0 args, wanted 1"))

		err = q.Consumer().ProcessAll(ctx)
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is not called", func() {
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("HandlerFunc", func() {
	ctx := context.Background()
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(msg *taskq.Message) error {
				Expect(msg.Args).To(Equal([]interface{}{"string", 42}))
				ch <- true
				return nil
			},
		})

		err := q.Add(task.WithArgs(ctx, "string", 42))
		Expect(err).NotTo(HaveOccurred())

		err = q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("is called with Message", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("message retry timing", func() {
	ctx := context.Background()
	var q *memqueue.Queue
	var task *taskq.Task
	backoff := 100 * time.Millisecond
	var count int
	var ch chan time.Time

	BeforeEach(func() {
		count = 0
		ch = make(chan time.Time, 10)
		q = memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task = taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() error {
				ch <- time.Now()
				count++
				return fmt.Errorf("fake error #%d", count)
			},
			RetryLimit: 3,
			MinBackoff: backoff,
		})
	})

	Context("without delay", func() {
		var now time.Time

		BeforeEach(func() {
			now = time.Now()
			_ = q.Add(task.WithArgs(ctx))

			err := q.Close()
			Expect(err).NotTo(HaveOccurred())
		})

		It("is retried in time", func() {
			Expect(ch).To(Receive(BeTemporally("~", now, backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(backoff), backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(3*backoff), backoff/10)))
			Expect(ch).NotTo(Receive())
		})
	})

	Context("message with delay", func() {
		var now time.Time

		BeforeEach(func() {
			msg := task.WithArgs(ctx)
			msg.Delay = 5 * backoff
			now = time.Now().Add(msg.Delay)

			q.Add(msg)

			err := q.Close()
			Expect(err).NotTo(HaveOccurred())
		})

		It("is retried in time", func() {
			Expect(ch).To(Receive(BeTemporally("~", now, backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(backoff), backoff/10)))
			Expect(ch).To(Receive(BeTemporally("~", now.Add(3*backoff), backoff/10)))
			Expect(ch).NotTo(Receive())
		})
	})
})

var _ = Describe("failing queue with error handler", func() {
	ctx := context.Background()
	var q *memqueue.Queue
	ch := make(chan bool, 10)

	BeforeEach(func() {
		q = memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() error {
				return errors.New("fake error")
			},
			FallbackHandler: func() {
				ch <- true
			},
			RetryLimit: 1,
		})
		q.Add(task.WithArgs(ctx))

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("error handler is called when handler fails", func() {
		Expect(ch).To(Receive())
		Expect(ch).NotTo(Receive())
	})
})

var _ = Describe("named message", func() {
	ctx := context.Background()
	var count int64

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name:  "test",
			Redis: redisRing(),
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() {
				atomic.AddInt64(&count, 1)
			},
		})

		name := uuid.NewV4().String()

		var wg sync.WaitGroup
		for i := 0; i < 100; i++ {
			wg.Add(1)
			go func() {
				defer GinkgoRecover()
				defer wg.Done()
				msg := task.WithArgs(ctx)
				msg.Name = name
				q.Add(msg)
			}()
		}
		wg.Wait()

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("is processed once", func() {
		n := atomic.LoadInt64(&count)
		Expect(n).To(Equal(int64(1)))
	})
})

var _ = Describe("CallOnce", func() {
	ctx := context.Background()
	var now time.Time
	delay := 3 * time.Second
	ch := make(chan time.Time, 10)

	BeforeEach(func() {
		now = time.Now()

		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name:  "test",
			Redis: redisRing(),
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func(slot int64) error {
				ch <- time.Now()
				return nil
			},
		})

		var wg sync.WaitGroup
		for i := 0; i < 10; i++ {
			wg.Add(1)
			go func() {
				defer GinkgoRecover()
				defer wg.Done()

				q.Add(task.WithArgs(ctx, slot(delay)).OnceInPeriod(delay))
			}()
		}
		wg.Wait()

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("processes message once with delay", func() {
		Expect(ch).To(Receive(BeTemporally(">", now.Add(delay), time.Second)))
		Consistently(ch).ShouldNot(Receive())
	})
})

var _ = Describe("stress testing", func() {
	const n = 10000
	ctx := context.Background()
	var count int64

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name: "test",
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() {
				atomic.AddInt64(&count, 1)
			},
		})

		for i := 0; i < n; i++ {
			q.Add(task.WithArgs(ctx))
		}

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("handler is called for all messages", func() {
		nn := atomic.LoadInt64(&count)
		Expect(nn).To(Equal(int64(n)))
	})
})

var _ = Describe("stress testing failing queue", func() {
	const n = 100000
	ctx := context.Background()
	var errorCount int64

	BeforeEach(func() {
		q := memqueue.NewQueue(&taskq.QueueOptions{
			Name:                 "test",
			PauseErrorsThreshold: -1,
		})
		task := taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() error {
				return errors.New("fake error")
			},
			FallbackHandler: func() {
				atomic.AddInt64(&errorCount, 1)
			},
			RetryLimit: 1,
		})

		for i := 0; i < n; i++ {
			q.Add(task.WithArgs(ctx))
		}

		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("error handler is called for all messages", func() {
		nn := atomic.LoadInt64(&errorCount)
		Expect(nn).To(Equal(int64(n)))
	})
})

var _ = Describe("empty queue", func() {
	ctx := context.Background()
	var q *memqueue.Queue
	var task *taskq.Task
	var processed uint32

	BeforeEach(func() {
		processed = 0
		q = memqueue.NewQueue(&taskq.QueueOptions{
			Name:  "test",
			Redis: redisRing(),
		})
		task = taskq.RegisterTask(&taskq.TaskOptions{
			Name: "test",
			Handler: func() {
				atomic.AddUint32(&processed, 1)
			},
		})
	})

	AfterEach(func() {
		_ = q.Close()
	})

	It("can be closed", func() {
		err := q.Close()
		Expect(err).NotTo(HaveOccurred())
	})

	It("stops processor", func() {
		err := q.Consumer().Stop()
		Expect(err).NotTo(HaveOccurred())
	})

	testEmptyQueue := func() {
		It("processes all messages", func() {
			err := q.Consumer().ProcessAll(ctx)
			Expect(err).NotTo(HaveOccurred())
		})

		It("processes one message", func() {
			err := q.Consumer().ProcessOne(ctx)
			Expect(err).To(MatchError("taskq: queue is empty"))

			err = q.Consumer().ProcessAll(ctx)
			Expect(err).NotTo(HaveOccurred())
		})
	}

	Context("when processor is stopped", func() {
		BeforeEach(func() {
			err := q.Consumer().Stop()
			Expect(err).NotTo(HaveOccurred())
		})

		testEmptyQueue()

		Context("when there are messages in the queue", func() {
			BeforeEach(func() {
				for i := 0; i < 3; i++ {
					err := q.Add(task.WithArgs(ctx))
					Expect(err).NotTo(HaveOccurred())
				}
			})

			It("processes all messages", func() {
				p := q.Consumer()

				err := p.ProcessAll(ctx)
				Expect(err).NotTo(HaveOccurred())

				n := atomic.LoadUint32(&processed)
				Expect(n).To(Equal(uint32(3)))
			})

			It("processes one message", func() {
				p := q.Consumer()

				err := p.ProcessOne(ctx)
				Expect(err).NotTo(HaveOccurred())

				n := atomic.LoadUint32(&processed)
				Expect(n).To(Equal(uint32(1)))

				err = p.ProcessAll(ctx)
				Expect(err).NotTo(HaveOccurred())

				n = atomic.LoadUint32(&processed)
				Expect(n).To(Equal(uint32(3)))
			})
		})
	})
})

// slot splits time into equal periods (called slots) and returns
// slot number for provided time.
func slot(period time.Duration) int64 {
	tm := time.Now()
	periodSec := int64(period / time.Second)
	if periodSec == 0 {
		return tm.Unix()
	}
	return tm.Unix() / periodSec
}

var (
	ringOnce sync.Once
	ring     *redis.Ring
)

func redisRing() *redis.Ring {
	ringOnce.Do(func() {
		ring = redis.NewRing(&redis.RingOptions{
			Addrs: map[string]string{"0": ":6379"},
		})
	})
	return ring
}
<|endoftext|>


<commit_before>// Copyright 2014 Apcera Inc. All rights reserved.

package server

import (
	"strconv"
	"sync"
	"testing"
)

func BenchmarkParseInt(b *testing.B) {
	b.SetBytes(1)
	n := "12345678"
	for i := 0; i < b.N; i++ {
		strconv.ParseInt(n, 10, 0)
	}
}

func BenchmarkParseSize(b *testing.B) {
	b.SetBytes(1)
	n := []byte("12345678")
	for i := 0; i < b.N; i++ {
		parseSize(n)
	}
}

func deferUnlock(mu sync.Mutex) {
	mu.Lock()
	defer mu.Unlock()
}

func BenchmarkDeferMutex(b *testing.B) {
	var mu sync.Mutex
	b.SetBytes(1)
	for i := 0; i < b.N; i++ {
		deferUnlock(mu)
	}
}

func noDeferUnlock(mu sync.Mutex) {
	mu.Lock()
	mu.Unlock()
}

func BenchmarkNoDeferMutex(b *testing.B) {
	var mu sync.Mutex
	b.SetBytes(1)
	for i := 0; i < b.N; i++ {
		noDeferUnlock(mu)
	}
}
<commit_msg>pass by address<commit_after>// Copyright 2014 Apcera Inc. All rights reserved.

package server

import (
	"strconv"
	"sync"
	"testing"
)

func BenchmarkParseInt(b *testing.B) {
	b.SetBytes(1)
	n := "12345678"
	for i := 0; i < b.N; i++ {
		strconv.ParseInt(n, 10, 0)
	}
}

func BenchmarkParseSize(b *testing.B) {
	b.SetBytes(1)
	n := []byte("12345678")
	for i := 0; i < b.N; i++ {
		parseSize(n)
	}
}

func deferUnlock(mu *sync.Mutex) {
	mu.Lock()
	defer mu.Unlock()
}

func BenchmarkDeferMutex(b *testing.B) {
	var mu sync.Mutex
	b.SetBytes(1)
	for i := 0; i < b.N; i++ {
		deferUnlock(&mu)
	}
}

func noDeferUnlock(mu *sync.Mutex) {
	mu.Lock()
	mu.Unlock()
}

func BenchmarkNoDeferMutex(b *testing.B) {
	var mu sync.Mutex
	b.SetBytes(1)
	for i := 0; i < b.N; i++ {
		noDeferUnlock(&mu)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 512
)

type socket struct {
	ws   *websocket.Conn
	send chan []byte
}

func (c *socket) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

func (s *Server) websocketBroadcast(v interface{}) {
	out, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	s.broadcast <- out
}

func (s *Server) websocketRouter() {
	hub := make(map[*socket]bool)
	for {
		select {
		case c := <-s.addSocket:
			hub[c] = true
		case c := <-s.delSocket:
			delete(hub, c)
		case m := <-s.broadcast:
			for c := range hub {
				c.write(websocket.TextMessage, m)
			}
		}
	}
}
func (s *Server) websocketReadPump(c *socket) {
	defer func() {
		s.delSocket <- c
		c.ws.Close()
	}()
	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.ws.ReadMessage()
		if string(message) == "list" {
			s.Lock()
			blocks := s.ListBlocks()
			groups := s.ListGroups()
			sources := s.ListSources()
			connections := s.ListConnections()
			links := s.listLinks()

			cache := make(map[int]Node)

			for i, _ := range blocks {
				cache[blocks[i].Id] = &blocks[i]
			}

			for i, _ := range groups {
				cache[groups[i].Id] = &groups[i]
			}

			for i, _ := range sources {
				cache[sources[i].Id] = &sources[i]
			}

			var recurseGroups func(int)
			recurseGroups = func(id int) {
				g := cache[id].(*Group)
				for _, child := range g.Children {
					switch n := cache[child].(type) {
					case *BlockLedger:
						o, _ := json.Marshal(Update{Action: CREATE, Type: BLOCK, Data: wsBlock{n}})
						c.write(websocket.TextMessage, o)

						o, _ = json.Marshal(Update{Action: CREATE, Type: CHILD, Data: wsGroupChild{
							Group: wsId{id},
							Child: wsId{child},
						}})
						c.write(websocket.TextMessage, o)
					case *Group:
						// we have to set empty out the children because the children don't exist
						// in the client at the time that this message is posted.
						// this is suboptimal, but provides a cleaner representation of state to
						// the client.
						ng := *n
						ng.Children = []int{}
						o, _ := json.Marshal(Update{Action: CREATE, Type: GROUP, Data: wsGroup{ng}})
						c.write(websocket.TextMessage, o)
						o, _ = json.Marshal(Update{Action: CREATE, Type: CHILD, Data: wsGroupChild{
							Group: wsId{id},
							Child: wsId{child},
						}})
						c.write(websocket.TextMessage, o)
						recurseGroups(child)
					case *SourceLedger:
						o, _ := json.Marshal(Update{Action: CREATE, Type: SOURCE, Data: wsSource{n}})
						c.write(websocket.TextMessage, o)
						o, _ = json.Marshal(Update{Action: CREATE, Type: CHILD, Data: wsGroupChild{
							Group: wsId{id},
							Child: wsId{child},
						}})
						c.write(websocket.TextMessage, o)
					}
				}
			}

			// group 0 is the root, and does not require to be added to a group.
			// TODO: in the future, we may consider moving to a non-root tree structure.
			o, _ := json.Marshal(Update{Action: CREATE, Type: GROUP, Data: wsGroup{cache[0]}})
			c.write(websocket.TextMessage, o)
			recurseGroups(0)

			for _, connection := range connections {
				o, _ := json.Marshal(Update{Action: CREATE, Type: CONNECTION, Data: wsConnection{connection}})
				c.write(websocket.TextMessage, o)

			}

			for _, l := range links {
				o, _ := json.Marshal(Update{Action: CREATE, Type: LINK, Data: wsLink{l}})
				c.write(websocket.TextMessage, o)
			}

			s.Unlock()
			// we want to lock for this entire time, so that nothing can interfere
			// with our state as we are dumping it

		}

		if err != nil {
			break
		}
		s.emitChan <- message
	}
}

func (s *Server) websocketWritePump(c *socket) {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.write(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (s *Server) UpdateSocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	c := &socket{send: make(chan []byte, 256), ws: ws}
	s.addSocket <- c
	go s.websocketWritePump(c)
	go s.websocketReadPump(c)
}
<commit_msg>root needs to report 0 children as well<commit_after>package server

import (
	"encoding/json"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 512
)

type socket struct {
	ws   *websocket.Conn
	send chan []byte
}

func (c *socket) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

func (s *Server) websocketBroadcast(v interface{}) {
	out, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	s.broadcast <- out
}

func (s *Server) websocketRouter() {
	hub := make(map[*socket]bool)
	for {
		select {
		case c := <-s.addSocket:
			hub[c] = true
		case c := <-s.delSocket:
			delete(hub, c)
		case m := <-s.broadcast:
			for c := range hub {
				c.write(websocket.TextMessage, m)
			}
		}
	}
}
func (s *Server) websocketReadPump(c *socket) {
	defer func() {
		s.delSocket <- c
		c.ws.Close()
	}()
	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.ws.ReadMessage()
		if string(message) == "list" {
			s.Lock()
			blocks := s.ListBlocks()
			groups := s.ListGroups()
			sources := s.ListSources()
			connections := s.ListConnections()
			links := s.listLinks()

			cache := make(map[int]Node)

			for i, _ := range blocks {
				cache[blocks[i].Id] = &blocks[i]
			}

			for i, _ := range groups {
				cache[groups[i].Id] = &groups[i]
			}

			for i, _ := range sources {
				cache[sources[i].Id] = &sources[i]
			}

			var recurseGroups func(int)
			recurseGroups = func(id int) {
				g := cache[id].(*Group)
				for _, child := range g.Children {
					switch n := cache[child].(type) {
					case *BlockLedger:
						o, _ := json.Marshal(Update{Action: CREATE, Type: BLOCK, Data: wsBlock{n}})
						c.write(websocket.TextMessage, o)

						o, _ = json.Marshal(Update{Action: CREATE, Type: CHILD, Data: wsGroupChild{
							Group: wsId{id},
							Child: wsId{child},
						}})
						c.write(websocket.TextMessage, o)
					case *Group:
						// we have to set empty out the children because the children don't exist
						// in the client at the time that this message is posted.
						// this is suboptimal, but provides a cleaner representation of state to
						// the client.
						ng := *n
						ng.Children = []int{}
						o, _ := json.Marshal(Update{Action: CREATE, Type: GROUP, Data: wsGroup{ng}})
						c.write(websocket.TextMessage, o)
						o, _ = json.Marshal(Update{Action: CREATE, Type: CHILD, Data: wsGroupChild{
							Group: wsId{id},
							Child: wsId{child},
						}})
						c.write(websocket.TextMessage, o)
						recurseGroups(child)
					case *SourceLedger:
						o, _ := json.Marshal(Update{Action: CREATE, Type: SOURCE, Data: wsSource{n}})
						c.write(websocket.TextMessage, o)
						o, _ = json.Marshal(Update{Action: CREATE, Type: CHILD, Data: wsGroupChild{
							Group: wsId{id},
							Child: wsId{child},
						}})
						c.write(websocket.TextMessage, o)
					}
				}
			}

			// group 0 is the root, and does not require to be added to a group.
			// TODO: in the future, we may consider moving to a non-root tree structure.
			ng := cache[0].(*Group)
			cg := *ng
			cg.Children = []int{}
			o, _ := json.Marshal(Update{Action: CREATE, Type: GROUP, Data: wsGroup{cg}})
			c.write(websocket.TextMessage, o)
			recurseGroups(0)

			for _, connection := range connections {
				o, _ := json.Marshal(Update{Action: CREATE, Type: CONNECTION, Data: wsConnection{connection}})
				c.write(websocket.TextMessage, o)

			}

			for _, l := range links {
				o, _ := json.Marshal(Update{Action: CREATE, Type: LINK, Data: wsLink{l}})
				c.write(websocket.TextMessage, o)
			}

			s.Unlock()
			// we want to lock for this entire time, so that nothing can interfere
			// with our state as we are dumping it

		}

		if err != nil {
			break
		}
		s.emitChan <- message
	}
}

func (s *Server) websocketWritePump(c *socket) {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.write(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (s *Server) UpdateSocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	c := &socket{send: make(chan []byte, 256), ws: ws}
	s.addSocket <- c
	go s.websocketWritePump(c)
	go s.websocketReadPump(c)
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   []byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}

	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}

	_, err = msg.Decode(b)
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  []byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<commit_msg>修改了buffer.go;sendrecv.go;process.go<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   []byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}

	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  []byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				//return
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("向ringbuffer些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Sendering from panic(sender): %v", this.cid(), r)
			})
		} else {
			Log.Infoc(func() string {
				return fmt.Sprintf("向conn些数据成功！")
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//Log.Debugc(func() string {
	//	return fmt.Sprintf("(%s) Starting sender", this.cid())
	//})

	this.wgStarted.Done()
	//Log.Debugc(func() string {
	//	return fmt.Sprintf("sender_1(%s)", this.cid())
	//})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			//Log.Debugc(func() string {
			//	return fmt.Sprintf("sender_2(%s)", this.cid())
			//})
			_, err := this.out.WriteTo(conn)
			Log.Debugc(func() string {
				return fmt.Sprintf("sender_3(%s)", this.cid())
			})
			if err != nil {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_4(%s)", this.cid())
				})
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("向conn些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
	Log.Debugc(func() string {
		return fmt.Sprintf("sender_6(%s)", this.cid())
	})
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	var b []byte
	var index int64
	var ok bool

	if this.isDone() {
		return nil, err
	}
	b, index, ok = this.in.ReadBuffer()
	defer this.in.ReadCommit(index)
	if !ok {
		return nil, err
	}

	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) error {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	n, err := msg.Encode(b)
	if err != nil {
		return err
	}

	if this.out.WriteBuffer(b) {
		return err
	}

	this.outStat.increment(int64(n))

	return nil
}
<commit_msg>添加测试代码<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				//return
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("向ringbuffer些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Sendering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//Log.Debugc(func() string {
	//	return fmt.Sprintf("(%s) Starting sender", this.cid())
	//})

	this.wgStarted.Done()
	//Log.Debugc(func() string {
	//	return fmt.Sprintf("sender_1(%s)", this.cid())
	//})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			//Log.Debugc(func() string {
			//	return fmt.Sprintf("sender_2(%s)", this.cid())
			//})
			_, err := this.out.WriteTo(conn)
			Log.Debugc(func() string {
				return fmt.Sprintf("sender_3(%s)", this.cid())
			})
			if err != nil {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_4(%s)", this.cid())
				})
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("向conn些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
	Log.Debugc(func() string {
		return fmt.Sprintf("sender_6(%s)", this.cid())
	})
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	var b []byte
	var index int64
	var ok bool

	if this.isDone() {
		return nil, err
	}
	b, index, ok = this.in.ReadBuffer()
	defer this.in.ReadCommit(index)
	if !ok {
		return nil, err
	}

	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) error {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	n, err := msg.Encode(b)
	if err != nil {
		return err
	}

	if this.out.WriteBuffer(b) {
		return err
	}

	this.outStat.increment(int64(n))

	return nil
}
<|endoftext|>


<commit_before>package services

import (
	"errors"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"github.com/fatih/color"
	"github.com/hpcloud/tail"
	"github.com/yext/edward/common"
	"github.com/yext/edward/home"
	"github.com/yext/errgo"
)

type ServiceCommand struct {
	// Parent service config
	Service *ServiceConfig
	// Path to string
	Scripts struct {
		Build  string
		Launch string
		Stop   string
	}
	Pid  int
	Logs struct {
		Build string
		Run   string
		Stop  string
	}
	Logger common.Logger
}

func (c *ServiceCommand) printf(format string, v ...interface{}) {
	if c.Logger == nil {
		return
	}
	c.Logger.Printf(format, v...)
}

func (sc *ServiceCommand) createScript(content string, scriptType string) (*os.File, error) {
	file, err := os.Create(path.Join(home.EdwardConfig.ScriptDir, sc.Service.Name+"-"+scriptType))
	if err != nil {
		return nil, err
	}
	file.WriteString(content)
	file.Close()

	err = os.Chmod(file.Name(), 0777)
	if err != nil {
		return nil, err
	}

	return file, nil
}

func (sc *ServiceCommand) deleteScript(scriptType string) error {
	return os.Remove(path.Join(home.EdwardConfig.ScriptDir, sc.Service.Name+"-"+scriptType))
}

func (sc *ServiceCommand) BuildSync() error {
	printOperation("Building " + sc.Service.Name)
	sc.printf("Building %v\n", sc.Service.Name)

	if sc.Pid != 0 {
		sc.printf("%v is already running\n", sc.Service.Name)
		printResult("Already running", color.FgYellow)
		return nil
	}

	if sc.Scripts.Build == "" {
		printResult("No build", color.FgGreen)
		sc.printf("No build needed for %v\n", sc.Service.Name)
		return nil
	}

	file, err := sc.createScript(sc.Scripts.Build, "Build")
	// Build the project and wait for completion
	if err != nil {
		return err
	}
	defer sc.deleteScript("Build")

	cmd := exec.Command(file.Name())
	err = cmd.Run()
	if err != nil {
		printResult("Failed", color.FgRed)
		printFile(sc.Logs.Build)
		return errgo.Mask(err)
	}

	printResult("OK", color.FgGreen)
	sc.printf("%v build succeeded.\n", sc.Service.Name)

	return nil
}

func (sc *ServiceCommand) waitUntilLive(command *exec.Cmd) error {

	sc.printf("Waiting for %v to start.\n", sc.Service.Name)

	var err error = nil
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		// Read output until we get the success
		var t *tail.Tail
		t, err = tail.TailFile(sc.Logs.Run, tail.Config{Follow: true, Logger: tail.DiscardingLogger})
		for line := range t.Lines {
			if strings.Contains(line.Text, sc.Service.Properties.Started) {
				wg.Done()
				return
			}
		}
	}()

	go func() {
		// Wait until the process exists
		command.Wait()
		err = errors.New("Command failed!")
		wg.Done()
	}()

	wg.Wait()

	return err
}

func (sc *ServiceCommand) StartAsync() error {

	printOperation("Launching " + sc.Service.Name)
	sc.printf("Launching %v\n", sc.Service.Name)

	if sc.Pid != 0 {
		printResult("Already running", color.FgYellow)
		sc.printf("%v is already running.\n", sc.Service.Name)
		return nil
	}
	// Clear logs
	os.Remove(sc.Logs.Run)

	// Start the project and get the PID
	file, err := sc.createScript(sc.Scripts.Launch, "Launch")
	if err != nil {
		return err
	}
	//defer os.Remove(file.Name())

	cmd := exec.Command(file.Name())
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, sc.Service.Env...)
	err = cmd.Start()
	if err != nil {
		printResult("Failed", color.FgRed)
		return errgo.Mask(err)
	}

	pid := cmd.Process.Pid

	sc.printf("%v has PID: %d.\n", sc.Service.Name, pid)

	pidStr := strconv.Itoa(pid)
	f, err := os.Create(sc.getPidPath())
	if err != nil {
		return err
	}
	f.WriteString(pidStr)
	f.Close()

	err = sc.waitUntilLive(cmd)
	if err == nil {
		printResult("OK", color.FgGreen)
		sc.printf("%v start succeeded.\n", sc.Service.Name)
	} else {
		printResult("Failed!", color.FgRed)
		printFile(sc.Logs.Run)
	}
	return errgo.Mask(err)
}

func (sc *ServiceCommand) StopScript() error {

	sc.printf("Running stop script for %v\n", sc.Service.Name)

	// Start the project and get the PID
	file, err := sc.createScript(sc.Scripts.Stop, "Stop")
	if err != nil {
		return err
	}
	defer os.Remove(file.Name())

	cmd := exec.Command(file.Name())
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	err = cmd.Run()
	return err
}

func (sc *ServiceCommand) clearPid() {
	sc.Pid = 0
	os.Remove(sc.getPidPath())
}

func (sc *ServiceCommand) clearState() {
	sc.clearPid()
	sc.deleteScript("Stop")
	sc.deleteScript("Launch")
	sc.deleteScript("Build")
}

func (sc *ServiceCommand) getPidPath() string {
	dir := home.EdwardConfig.PidDir
	return path.Join(dir, sc.Service.Name+".pid")
}
<commit_msg>Allow for services that build but do not launch anything. This can be useful for common pre-build steps for groups.<commit_after>package services

import (
	"errors"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"github.com/fatih/color"
	"github.com/hpcloud/tail"
	"github.com/yext/edward/common"
	"github.com/yext/edward/home"
	"github.com/yext/errgo"
)

type ServiceCommand struct {
	// Parent service config
	Service *ServiceConfig
	// Path to string
	Scripts struct {
		Build  string
		Launch string
		Stop   string
	}
	Pid  int
	Logs struct {
		Build string
		Run   string
		Stop  string
	}
	Logger common.Logger
}

func (c *ServiceCommand) printf(format string, v ...interface{}) {
	if c.Logger == nil {
		return
	}
	c.Logger.Printf(format, v...)
}

func (sc *ServiceCommand) createScript(content string, scriptType string) (*os.File, error) {
	file, err := os.Create(path.Join(home.EdwardConfig.ScriptDir, sc.Service.Name+"-"+scriptType))
	if err != nil {
		return nil, err
	}
	file.WriteString(content)
	file.Close()

	err = os.Chmod(file.Name(), 0777)
	if err != nil {
		return nil, err
	}

	return file, nil
}

func (sc *ServiceCommand) deleteScript(scriptType string) error {
	return os.Remove(path.Join(home.EdwardConfig.ScriptDir, sc.Service.Name+"-"+scriptType))
}

func (sc *ServiceCommand) BuildSync() error {
	printOperation("Building " + sc.Service.Name)
	sc.printf("Building %v\n", sc.Service.Name)

	if sc.Pid != 0 {
		sc.printf("%v is already running\n", sc.Service.Name)
		printResult("Already running", color.FgYellow)
		return nil
	}

	if sc.Scripts.Build == "" {
		printResult("No build", color.FgGreen)
		sc.printf("No build needed for %v\n", sc.Service.Name)
		return nil
	}

	file, err := sc.createScript(sc.Scripts.Build, "Build")
	// Build the project and wait for completion
	if err != nil {
		return err
	}
	defer sc.deleteScript("Build")

	cmd := exec.Command(file.Name())
	err = cmd.Run()
	if err != nil {
		printResult("Failed", color.FgRed)
		printFile(sc.Logs.Build)
		return errgo.Mask(err)
	}

	printResult("OK", color.FgGreen)
	sc.printf("%v build succeeded.\n", sc.Service.Name)

	return nil
}

func (sc *ServiceCommand) waitUntilLive(command *exec.Cmd) error {

	sc.printf("Waiting for %v to start.\n", sc.Service.Name)

	var err error = nil
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		// Read output until we get the success
		var t *tail.Tail
		t, err = tail.TailFile(sc.Logs.Run, tail.Config{Follow: true, Logger: tail.DiscardingLogger})
		for line := range t.Lines {
			if strings.Contains(line.Text, sc.Service.Properties.Started) {
				wg.Done()
				return
			}
		}
	}()

	go func() {
		// Wait until the process exists
		command.Wait()
		err = errors.New("Command failed!")
		wg.Done()
	}()

	wg.Wait()

	return err
}

func (sc *ServiceCommand) StartAsync() error {

	printOperation("Launching " + sc.Service.Name)
	sc.printf("Launching %v\n", sc.Service.Name)

	if sc.Pid != 0 {
		printResult("Already running", color.FgYellow)
		sc.printf("%v is already running.\n", sc.Service.Name)
		return nil
	}

	if sc.Scripts.Launch == "" {
		printResult("No launch", color.FgGreen)
		sc.printf("No launch needed for %v\n", sc.Service.Name)
		return nil
	}

	// Clear logs
	os.Remove(sc.Logs.Run)

	// Start the project and get the PID
	file, err := sc.createScript(sc.Scripts.Launch, "Launch")
	if err != nil {
		return err
	}
	//defer os.Remove(file.Name())

	cmd := exec.Command(file.Name())
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, sc.Service.Env...)
	err = cmd.Start()
	if err != nil {
		printResult("Failed", color.FgRed)
		return errgo.Mask(err)
	}

	pid := cmd.Process.Pid

	sc.printf("%v has PID: %d.\n", sc.Service.Name, pid)

	pidStr := strconv.Itoa(pid)
	f, err := os.Create(sc.getPidPath())
	if err != nil {
		return err
	}
	f.WriteString(pidStr)
	f.Close()

	err = sc.waitUntilLive(cmd)
	if err == nil {
		printResult("OK", color.FgGreen)
		sc.printf("%v start succeeded.\n", sc.Service.Name)
	} else {
		printResult("Failed!", color.FgRed)
		printFile(sc.Logs.Run)
	}
	return errgo.Mask(err)
}

func (sc *ServiceCommand) StopScript() error {

	sc.printf("Running stop script for %v\n", sc.Service.Name)

	// Start the project and get the PID
	file, err := sc.createScript(sc.Scripts.Stop, "Stop")
	if err != nil {
		return err
	}
	defer os.Remove(file.Name())

	cmd := exec.Command(file.Name())
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	err = cmd.Run()
	return err
}

func (sc *ServiceCommand) clearPid() {
	sc.Pid = 0
	os.Remove(sc.getPidPath())
}

func (sc *ServiceCommand) clearState() {
	sc.clearPid()
	sc.deleteScript("Stop")
	sc.deleteScript("Launch")
	sc.deleteScript("Build")
}

func (sc *ServiceCommand) getPidPath() string {
	dir := home.EdwardConfig.PidDir
	return path.Join(dir, sc.Service.Name+".pid")
}
<|endoftext|>


<commit_before>package services_test

import (
	"net/http"
	"net/http/httptest"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/s3"

	"github.com/rosenhouse/awsfaker"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

type FakeS3Backend struct {
	PutObjectCall struct {
		Receives      *s3.PutObjectInput
		ReturnsResult *s3.PutObjectOutput
		ReturnsError  error
	}
}

func (f *FakeS3Backend) PutObject(input *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
	f.PutObjectCall.Receives = input
	return f.PutObjectCall.ReturnsResult, f.PutObjectCall.ReturnsError
}

var _ = XDescribe("Mocking out the S3 service", func() {
	var (
		fakeBackend *FakeS3Backend
		fakeServer  *httptest.Server
		client      *s3.S3
	)

	BeforeEach(func() {
		fakeBackend = &FakeS3Backend{}
		fakeServer = httptest.NewServer(awsfaker.New(fakeBackend))
		client = s3.New(
			newSession(fakeServer.URL),
			&aws.Config{S3ForcePathStyle: aws.Bool(true)},
		)
	})

	AfterEach(func() {
		if fakeServer != nil {
			fakeServer.Close()
		}
	})

	It("should call the backend method", func() {
		client.PutObject(
			&s3.PutObjectInput{
				Bucket: aws.String("some-bucket"),
				Key:    aws.String("some/object/path"),
				Body:   strings.NewReader("some object data to upload"),
				Metadata: map[string]*string{
					"key1": aws.String("a-value"),
					"key2": aws.String("b-value"),
				},
			})

		Expect(fakeBackend.PutObjectCall.Receives).NotTo(BeNil())
		Expect(fakeBackend.PutObjectCall.Receives.Bucket).To(Equal(aws.String("some-bucket")))
		Expect(fakeBackend.PutObjectCall.Receives.Key).To(Equal(aws.String("some/object/path")))
		Expect(fakeBackend.PutObjectCall.Receives.Body).To(Equal(strings.NewReader("some object data to upload")))
		Expect(fakeBackend.PutObjectCall.Receives.Metadata).To(Equal(map[string]*string{
			"key1": aws.String("a-value"),
			"key2": aws.String("b-value"),
		}))
	})

	Context("when the backend succeeds", func() {
		It("should return the data in a format parsable by the client library", func() {
			fakeBackend.PutObjectCall.ReturnsResult = &s3.PutObjectOutput{
				ETag: aws.String("some-etag"),
			}

			output, err := client.PutObject(
				&s3.PutObjectInput{
					Bucket: aws.String("some-bucket"),
					Key:    aws.String("some/object/path"),
					Body:   strings.NewReader("some object data to upload"),
				})

			Expect(err).NotTo(HaveOccurred())
			Expect(output).To(Equal(&s3.PutObjectOutput{
				ETag: aws.String("some-etag"),
			}))
		})
	})

	Context("when the backend returns an error", func() {
		It("should return the error in a format that is parsable by the client library", func() {
			fakeBackend.PutObjectCall.ReturnsError = &awsfaker.ErrorResponse{
				AWSErrorCode:    "ValidationError",
				AWSErrorMessage: "some error message",
				HTTPStatusCode:  http.StatusBadRequest,
			}

			_, err := client.PutObject(
				&s3.PutObjectInput{
					Bucket: aws.String("some-bucket"),
					Key:    aws.String("some/object/path"),
					Body:   strings.NewReader("some object data to upload"),
				})

			Expect(err).To(HaveOccurred())
			awsErr := err.(awserr.RequestFailure)
			Expect(awsErr.StatusCode()).To(Equal(400))
			Expect(awsErr.Code()).To(Equal("ValidationError"))
			Expect(awsErr.Message()).To(Equal("some error message"))
		})
	})
})
<commit_msg>Remove pending s3 tests for now<commit_after><|endoftext|>


<commit_before>package wikidump

import (
	"os"
	"regexp"
	"sort"
	"strings"
	"testing"
)

func assertStringEq(t *testing.T, a, b string) {
	if a != b {
		t.Errorf("%q != %q", a, b)
	}
}

func TestCleanup(t *testing.T) {
	in := "Hello, table! {|\n|bla\n|bla\n|}"
	assertStringEq(t, strings.TrimSpace(Cleanup(in)), "Hello, table!")
	in = `|}Hello,<ref group="note">1</rf> world{{math|bla{{?}}}}!{{bla`
	assertStringEq(t, Cleanup(in), "Hello, world!")
}

var ws = regexp.MustCompile(`\s+`)

func checkLink(t *testing.T, got Link, target, anchor string) {
	// Don't care about whitespace in the anchor...
	gotAnchor := ws.ReplaceAllString(strings.TrimSpace(got.Anchor), " ")
	if gotAnchor != anchor {
		t.Errorf("expected anchor %q, got %q", anchor, gotAnchor)
	}

	// ... but the target should be normalized.
	if got.Target != target {
		t.Errorf("expected target %q, got %q", target, got.Target)
	}
}

func TestExtractLinks_single(t *testing.T) {
	onlyLink := func(text string) Link {
		links := ExtractLinks(text)
		if len(links) != 1 {
			t.Errorf("expected one link, got at least %d", len(links))
		}
		for link, count := range links {
			if count != 1 {
				t.Errorf("expected one link, got %d", count)
			}
			return link
		}
		panic("no links")
	}

	cases := []struct {
		text, target, anchor string
	}{
		{"[[foo|bar]]", "Foo", "bar"},
		{"[[foo]]", "Foo", "foo"},
		{"[[File:picture!]] [[foo]]", "Foo", "foo"},
		{"[[foo]]bar.", "Foo", "foobar"},
		{"[[baz|foobar]];", "Baz", "foobar"},
		{"[[baz#quux]]", "Baz", "baz#quux"},
		{"[[FOO_BAR|foo bar]]", "FOO BAR", "foo bar"},

		{"[[C. Stephen Evans | Evans, C. Stephen]]",
			"C. Stephen Evans", "Evans, C. Stephen"},

		// Links like these commonly occur in nlwiki (and presumably dewiki
		// and other compounding languages):
		{"foo[[baz|bar]]", "Baz", "foobar"},
		{"before[[_target _page_ #\nsection|inside]]after",
			"Target page", "beforeinsideafter"},

		// MediaWiki only considers alphabetic characters outside [[]] part
		// of the anchor.
		{"foo-[[bar]]", "Bar", "bar"},
		{"[[bar]]/baz", "Bar", "bar"},

		// XXX The following are broken. They do occur in the wild, e.g.,
		// -18[[Celsius|°C]] and 700[[Megabyte|MB]]-cd (found in nlwiki dump).
		//{"[[bar]]0", "Bar", "bar"},
		//{"[[bar]]_", "Bar", "bar"},

		// We're not interested in section links
		{"[[#Some section|elsewhere]] [[other_article]]",
			"Other article", "other_article"},

		// Nor file and category links
		{"[[File:foo.png]] [[foo|see picture]]",
			"Foo", "see picture"},
		{"[[Category:Foos of the world]] [[foo]]", "Foo", "foo"},
	}

	for _, c := range cases {
		checkLink(t, onlyLink(c.text), c.target, c.anchor)
	}
}

type sortByAnchor []Link

func (s sortByAnchor) Len() int { return len(s) }

func (s sortByAnchor) Less(i, j int) bool {
	l := ([]Link)(s)
	return l[i].Anchor < l[j].Anchor
}

func (s sortByAnchor) Swap(i, j int) {
	l := ([]Link)(s)
	l[i], l[j] = l[j], l[i]
}

// Simulate the old API, except for the ordering.
func extractLinks(s string) []Link {
	links := make(sortByAnchor, 0)
	for k, v := range ExtractLinks(s) {
		for i := 0; i < v; i++ {
			links = append(links, k)
		}
	}
	sort.Sort(links)
	return ([]Link)(links)
}

func TestExtractLinks_multiple(t *testing.T) {
	// Expected links have to be sorted by anchor, UTF8-betically.
	cases := [][]string{
		// This construct appears in enwiki for chemical formulae etc.,
		// but also in nlwiki (and dewiki?) for more general compound nouns.
		{"[[Lithium|Li]][[Fluorine|F]]", "Fluorine", "F", "Lithium", "Li"},

		{"[[tera-|tera]][[becquerel]]s",
			"Becquerel", "becquerels", "Tera-", "tera"},

		// Newlines in links.
		{`[[Lord's
          prayer]]
          [[Dismissal
          (cricket)|dismissal]] [[Badass|Chuck
          Norris]]`,
			"Badass", "Chuck Norris",
			"Lord's prayer", "Lord's prayer",
			"Dismissal (cricket)", "dismissal"},
	}

	for _, c := range cases {
		links := extractLinks(c[0])
		if len(links) != (len(c)-1)/2 {
			t.Errorf("Wrong number of links %d in %q", len(links), c[0])
		}
		for i, l := range links {
			checkLink(t, l, c[i*2+1], c[i*2+2])
		}
	}
}

func getPages() []string {
	f, err := os.Open("nlwiki-20140927-sample.xml")
	if err != nil {
		panic(err)
	}
	pc, rc := make(chan *Page), make(chan *Redirect)
	go GetPages(f, pc, rc)
	go func() {
		for _ = range rc {
		}
	}()

	pages := make([]string, 0)
	for p := range pc {
		pages = append(pages, p.Text)
	}
	return pages
}

func BenchmarkCleanup(b *testing.B) {
	b.StopTimer()
	pages := getPages()

	for i := 0; i < b.N; i++ {
		b.StartTimer()
		for _, p := range pages {
			Cleanup(p)
		}
		b.StopTimer()
	}
}

func BenchmarkExtractLinks(b *testing.B) {
	b.StopTimer()
	pages := getPages()

	for i := 0; i < b.N; i++ {
		b.StartTimer()
		for _, p := range pages {
			ExtractLinks(p)
		}
		b.StopTimer()
	}
}
<commit_msg>improve test for wikidump.Cleanup<commit_after>package wikidump

import (
	"os"
	"regexp"
	"sort"
	"strings"
	"testing"
)

func TestCleanup(t *testing.T) {
	for _, c := range []struct {
		in, out string
	}{
		{
			in:  "Hello, table! {|\n|bla\n|bla\n|}",
			out: "Hello, table!",
		},
		{
			in:  `|}Hello,<ref group="note">1</rf> world{{math|bla{{?}}}}!{{bla`,
			out: "Hello, world!",
		},
		{
			// XXX Is this what we want?
			in:  "Text before, <references/> and after",
			out: "Text before,",
		},
	} {
		if out := strings.TrimSpace(Cleanup(c.in)); out != c.out {
			t.Errorf("expected %q for %q, got %q", c.out, c.in, out)
		}
	}
}

var ws = regexp.MustCompile(`\s+`)

func checkLink(t *testing.T, got Link, target, anchor string) {
	// Don't care about whitespace in the anchor...
	gotAnchor := ws.ReplaceAllString(strings.TrimSpace(got.Anchor), " ")
	if gotAnchor != anchor {
		t.Errorf("expected anchor %q, got %q", anchor, gotAnchor)
	}

	// ... but the target should be normalized.
	if got.Target != target {
		t.Errorf("expected target %q, got %q", target, got.Target)
	}
}

func TestExtractLinks_single(t *testing.T) {
	onlyLink := func(text string) Link {
		links := ExtractLinks(text)
		if len(links) != 1 {
			t.Errorf("expected one link, got at least %d", len(links))
		}
		for link, count := range links {
			if count != 1 {
				t.Errorf("expected one link, got %d", count)
			}
			return link
		}
		panic("no links")
	}

	cases := []struct {
		text, target, anchor string
	}{
		{"[[foo|bar]]", "Foo", "bar"},
		{"[[foo]]", "Foo", "foo"},
		{"[[File:picture!]] [[foo]]", "Foo", "foo"},
		{"[[foo]]bar.", "Foo", "foobar"},
		{"[[baz|foobar]];", "Baz", "foobar"},
		{"[[baz#quux]]", "Baz", "baz#quux"},
		{"[[FOO_BAR|foo bar]]", "FOO BAR", "foo bar"},

		{"[[C. Stephen Evans | Evans, C. Stephen]]",
			"C. Stephen Evans", "Evans, C. Stephen"},

		// Links like these commonly occur in nlwiki (and presumably dewiki
		// and other compounding languages):
		{"foo[[baz|bar]]", "Baz", "foobar"},
		{"before[[_target _page_ #\nsection|inside]]after",
			"Target page", "beforeinsideafter"},

		// MediaWiki only considers alphabetic characters outside [[]] part
		// of the anchor.
		{"foo-[[bar]]", "Bar", "bar"},
		{"[[bar]]/baz", "Bar", "bar"},

		// XXX The following are broken. They do occur in the wild, e.g.,
		// -18[[Celsius|°C]] and 700[[Megabyte|MB]]-cd (found in nlwiki dump).
		//{"[[bar]]0", "Bar", "bar"},
		//{"[[bar]]_", "Bar", "bar"},

		// We're not interested in section links
		{"[[#Some section|elsewhere]] [[other_article]]",
			"Other article", "other_article"},

		// Nor file and category links
		{"[[File:foo.png]] [[foo|see picture]]",
			"Foo", "see picture"},
		{"[[Category:Foos of the world]] [[foo]]", "Foo", "foo"},
	}

	for _, c := range cases {
		checkLink(t, onlyLink(c.text), c.target, c.anchor)
	}
}

type sortByAnchor []Link

func (s sortByAnchor) Len() int { return len(s) }

func (s sortByAnchor) Less(i, j int) bool {
	l := ([]Link)(s)
	return l[i].Anchor < l[j].Anchor
}

func (s sortByAnchor) Swap(i, j int) {
	l := ([]Link)(s)
	l[i], l[j] = l[j], l[i]
}

// Simulate the old API, except for the ordering.
func extractLinks(s string) []Link {
	links := make(sortByAnchor, 0)
	for k, v := range ExtractLinks(s) {
		for i := 0; i < v; i++ {
			links = append(links, k)
		}
	}
	sort.Sort(links)
	return ([]Link)(links)
}

func TestExtractLinks_multiple(t *testing.T) {
	// Expected links have to be sorted by anchor, UTF8-betically.
	cases := [][]string{
		// This construct appears in enwiki for chemical formulae etc.,
		// but also in nlwiki (and dewiki?) for more general compound nouns.
		{"[[Lithium|Li]][[Fluorine|F]]", "Fluorine", "F", "Lithium", "Li"},

		{"[[tera-|tera]][[becquerel]]s",
			"Becquerel", "becquerels", "Tera-", "tera"},

		// Newlines in links.
		{`[[Lord's
          prayer]]
          [[Dismissal
          (cricket)|dismissal]] [[Badass|Chuck
          Norris]]`,
			"Badass", "Chuck Norris",
			"Lord's prayer", "Lord's prayer",
			"Dismissal (cricket)", "dismissal"},
	}

	for _, c := range cases {
		links := extractLinks(c[0])
		if len(links) != (len(c)-1)/2 {
			t.Errorf("Wrong number of links %d in %q", len(links), c[0])
		}
		for i, l := range links {
			checkLink(t, l, c[i*2+1], c[i*2+2])
		}
	}
}

func getPages() []string {
	f, err := os.Open("nlwiki-20140927-sample.xml")
	if err != nil {
		panic(err)
	}
	pc, rc := make(chan *Page), make(chan *Redirect)
	go GetPages(f, pc, rc)
	go func() {
		for _ = range rc {
		}
	}()

	pages := make([]string, 0)
	for p := range pc {
		pages = append(pages, p.Text)
	}
	return pages
}

func BenchmarkCleanup(b *testing.B) {
	b.StopTimer()
	pages := getPages()

	for i := 0; i < b.N; i++ {
		b.StartTimer()
		for _, p := range pages {
			Cleanup(p)
		}
		b.StopTimer()
	}
}

func BenchmarkExtractLinks(b *testing.B) {
	b.StopTimer()
	pages := getPages()

	for i := 0; i < b.N; i++ {
		b.StartTimer()
		for _, p := range pages {
			ExtractLinks(p)
		}
		b.StopTimer()
	}
}
<|endoftext|>