// Produce is used by the source to send messages to the backend func (f *RBForwarder) Produce(data []byte, opts map[string]interface{}, opaque interface{}) error { if atomic.LoadUint32(&f.working) == 0 { return errors.New("Forwarder has been closed") } seq := f.currentProducedID f.currentProducedID++ m := utils.NewMessage() r := Report{ seq: seq, Opaque: opaque, } m.PushPayload(data) m.Opts.MSet(opts) m.Reports.Push(r) f.p.input <- m return nil }
func TestBatcher(t *testing.T) { Convey("Given a batcher", t, func() { batcher := &Batcher{ Config: Config{ Workers: 1, TimeoutMillis: 1000, Limit: 10, MaxPendingBatches: 10, Deflate: false, }, } b := batcher.Spawn(0).(*Batcher) b.clk = clock.NewMock() Convey("When the number of workers is requested", func() { workers := batcher.Workers() Convey("Then the number of workers should be correct", func() { So(workers, ShouldEqual, 1) }) }) Convey("When a message is received with no batch group", func() { m := utils.NewMessage() m.PushPayload([]byte("Hello World")) d := new(Doner) d.doneCalled = make(chan *utils.Message, 1) d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "").Times(1) b.OnMessage(m, d.Done) result := <-d.doneCalled Convey("Then the message should be sent as is", func() { So(b.batches.Count(), ShouldEqual, 0) payload, err := result.PopPayload() So(err, ShouldBeNil) So(string(payload), ShouldEqual, "Hello World") d.AssertExpectations(t) }) }) Convey("When a message is received by the worker, but not yet sent", func() { m := utils.NewMessage() m.PushPayload([]byte("Hello World")) m.Opts = cmap.New() m.Opts.Set("batch_group", "group1") b.OnMessage(m, nil) Convey("Message should be present on the batch", func() { tmp, exists := b.batches.Get("group1") So(exists, ShouldBeTrue) batch := tmp.(*Batch) batch.Writer.(*bufio.Writer).Flush() data := batch.Buf.Bytes() So(string(data), ShouldEqual, "Hello World") opts := batch.Message.Opts group, ok := opts.Get("batch_group") So(ok, ShouldBeTrue) So(group.(string), ShouldEqual, "group1") So(b.batches.Count(), ShouldEqual, 1) }) }) Convey("When the max number of messages is reached", func() { var messages []*utils.Message for i := 0; i < int(b.Config.Limit); i++ { m := utils.NewMessage() m.PushPayload([]byte("ABC")) m.Opts = cmap.New() m.Opts.Set("batch_group", "group1") messages = append(messages, m) } d := new(Doner) d.doneCalled = make(chan *utils.Message, 1) d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "limit").Times(1) for i := 0; i < int(b.Config.Limit); i++ { b.OnMessage(messages[i], d.Done) } Convey("Then the batch should be sent", func() { m := <-d.doneCalled data, err := m.PopPayload() So(err, ShouldBeNil) So(string(data), ShouldEqual, "ABCABCABCABCABCABCABCABCABCABC") group1, _ := b.batches.Get("group1") So(group1, ShouldBeNil) So(b.batches.Count(), ShouldEqual, 0) d.AssertExpectations(t) }) }) Convey("When the timeout expires", func() { var messages []*utils.Message for i := 0; i < 5; i++ { m := utils.NewMessage() m.PushPayload([]byte("Hello World")) m.Opts = cmap.New() m.Opts.Set("batch_group", "group1") messages = append(messages, m) } d := new(Doner) d.doneCalled = make(chan *utils.Message, 1) d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "timeout").Times(1) for i := 0; i < 5; i++ { b.OnMessage(messages[i], d.Done) } clk := b.clk.(*clock.Mock) Convey("The batch should be sent", func() { clk.Add(500 * time.Millisecond) group1, _ := b.batches.Get("group1") So(group1.(*Batch), ShouldNotBeNil) clk.Add(500 * time.Millisecond) <-d.doneCalled group1, _ = b.batches.Get("group1") So(group1, ShouldBeNil) So(b.batches.Count(), ShouldEqual, 0) d.AssertExpectations(t) }) }) Convey("When multiple messages are received with differents groups", func() { m1 := utils.NewMessage() m1.PushPayload([]byte("MESSAGE 1")) m1.Reports.Push("Report 1") m1.Opts = cmap.New() m1.Opts.Set("batch_group", "group1") m2 := utils.NewMessage() m2.PushPayload([]byte("MESSAGE 2")) m2.Reports.Push("Report 2") m2.Opts = cmap.New() m2.Opts.Set("batch_group", "group2") m3 := utils.NewMessage() m3.PushPayload([]byte("MESSAGE 3")) m3.Reports.Push("Report 3") m3.Opts = cmap.New() m3.Opts.Set("batch_group", "group2") d := new(Doner) d.doneCalled = make(chan *utils.Message, 2) d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "timeout").Times(2) b.OnMessage(m1, d.Done) b.OnMessage(m2, d.Done) b.OnMessage(m3, d.Done) Convey("Then each message should be in its group", func() { tmp, _ := b.batches.Get("group1") batch := tmp.(*Batch) batch.Writer.(*bufio.Writer).Flush() group1 := batch.Buf.Bytes() So(string(group1), ShouldEqual, "MESSAGE 1") tmp, _ = b.batches.Get("group2") batch = tmp.(*Batch) batch.Writer.(*bufio.Writer).Flush() group2 := batch.Buf.Bytes() So(string(group2), ShouldEqual, "MESSAGE 2MESSAGE 3") So(b.batches.Count(), ShouldEqual, 2) }) Convey("Then after a timeout the messages should be sent", func() { clk := b.clk.(*clock.Mock) So(b.batches.Count(), ShouldEqual, 2) clk.Add(time.Duration(b.Config.TimeoutMillis) * time.Millisecond) group1 := <-d.doneCalled group1Data, err := group1.PopPayload() report1 := group1.Reports.Pop().(string) So(err, ShouldBeNil) So(report1, ShouldEqual, "Report 1") group2 := <-d.doneCalled group2Data, err := group2.PopPayload() So(err, ShouldBeNil) report3 := group2.Reports.Pop().(string) So(report3, ShouldEqual, "Report 3") report2 := group2.Reports.Pop().(string) So(report2, ShouldEqual, "Report 2") So(string(group1Data), ShouldEqual, "MESSAGE 1") So(string(group2Data), ShouldEqual, "MESSAGE 2MESSAGE 3") batch, _ := b.batches.Get("group1") So(batch, ShouldBeNil) batch, _ = b.batches.Get("group2") So(batch, ShouldBeNil) So(b.batches.Count(), ShouldEqual, 0) d.AssertExpectations(t) }) }) }) Convey("Given a batcher with compression", t, func() { batcher := &Batcher{ Config: Config{ TimeoutMillis: 1000, Limit: 10, MaxPendingBatches: 10, Deflate: true, }, } b := batcher.Spawn(0).(*Batcher) b.clk = clock.NewMock() Convey("When the max number of messages is reached", func() { var messages []*utils.Message for i := 0; i < int(b.Config.Limit); i++ { m := utils.NewMessage() m.PushPayload([]byte("ABC")) m.Opts = cmap.New() m.Opts.Set("batch_group", "group1") m.Reports.Push("Report") messages = append(messages, m) } d := new(Doner) d.doneCalled = make(chan *utils.Message, 1) d.On("Done", mock.AnythingOfType("*utils.Message"), 0, "limit").Times(1) for i := 0; i < int(b.Config.Limit); i++ { b.OnMessage(messages[i], d.Done) } Convey("The batch should be sent compressed", func() { m := <-d.doneCalled decompressed := make([]byte, 30) data, err := m.PopPayload() buf := bytes.NewBuffer(data) r, err := zlib.NewReader(buf) r.Read(decompressed) r.Close() So(err, ShouldBeNil) So(string(decompressed), ShouldEqual, "ABCABCABCABCABCABCABCABCABCABC") So(m.Reports.Size(), ShouldEqual, b.Config.Limit) batch, _ := b.batches.Get("group1") So(batch, ShouldBeNil) So(b.batches.Count(), ShouldEqual, 0) d.AssertExpectations(t) }) }) }) }
func TestHTTPSender(t *testing.T) { Convey("Given an Limiter with 100 messages per second without burst", t, func() { limiter := &Limiter{ Config: Config{ MessageLimit: 100, Burst: 1, }, clk: clock.NewMock(), } limiter.Spawn(0) Convey("When the numer of worker is requested", func() { workers := limiter.Workers() Convey("Should be only one", func() { So(workers, ShouldEqual, 1) }) }) Convey("When the limit number of messages are reached", func() { clk := limiter.clk.(*clock.Mock) n := Nexter{ nextCalled: make(chan *utils.Message, limiter.Config.MessageLimit*2), } n.On("Done", mock.AnythingOfType("*utils.Message"), 0, "") for i := uint64(0); i < limiter.Config.MessageLimit; i++ { limiter.OnMessage(nil, n.Done) } Convey("Then the limiter should be paused", func() { So(limiter.currentMessages, ShouldEqual, limiter.Config.MessageLimit) So(limiter.paused, ShouldBeTrue) }) Convey("Then after 1 second the limiter should be ready again", func() { clk.Add(1 * time.Second) limiter.OnMessage(nil, n.Done) So(limiter.currentMessages, ShouldEqual, 1) So(limiter.paused, ShouldBeFalse) }) }) }) Convey("Given an Limiter with 1000 bytes per second without burst", t, func() { limiter := &Limiter{ Config: Config{ BytesLimit: 1000, Burst: 1, }, clk: clock.NewMock(), } limiter.Spawn(0) Convey("When messages are sent", func() { n := Nexter{ nextCalled: make(chan *utils.Message, 100), } n.On("Done", mock.AnythingOfType("*utils.Message"), 0, "") Convey("Then the limiter should not be paused after 750 bytes", func() { for i := uint64(0); i < 3; i++ { m := utils.NewMessage() payload := make([]byte, 250) m.PushPayload(payload) limiter.OnMessage(m, n.Done) } So(limiter.currentBytes, ShouldEqual, 750) So(limiter.paused, ShouldBeFalse) }) Convey("Then the limiter should be paused after 1000 bytes", func() { for i := uint64(0); i < 4; i++ { m := utils.NewMessage() payload := make([]byte, 250) m.PushPayload(payload) limiter.OnMessage(m, n.Done) } So(limiter.currentBytes, ShouldEqual, 1000) So(limiter.paused, ShouldBeTrue) }) Convey("Then after 1 second the limiter should be ready again", func() { clk := limiter.clk.(*clock.Mock) clk.Add(1 * time.Second) m := utils.NewMessage() payload := make([]byte, 250) m.PushPayload(payload) limiter.OnMessage(m, n.Done) So(limiter.currentBytes, ShouldEqual, 250) So(limiter.paused, ShouldBeFalse) }) }) }) Convey("Given a limiter with burst", t, func() { limiter := &Limiter{ Config: Config{ MessageLimit: 100, Burst: 2, }, clk: clock.NewMock(), } limiter.Spawn(0) clk := limiter.clk.(*clock.Mock) clk.Add(0) clk.Add(2 * time.Second) Convey("When the limit number of messages are reached", func() { n := Nexter{ nextCalled: make(chan *utils.Message, limiter.Config.MessageLimit*2), } n.On("Done", mock.AnythingOfType("*utils.Message"), 0, "") for i := uint64(0); i < limiter.Config.MessageLimit; i++ { limiter.OnMessage(nil, n.Done) } Convey("Then should be 2 burst available", func() { So(len(limiter.keepSending), ShouldEqual, 2) }) Convey("Then messages are not blocked after the limit", func() { for i := uint64(0); i < limiter.Config.MessageLimit; i++ { limiter.OnMessage(nil, n.Done) } So(limiter.currentMessages, ShouldEqual, 100) }) Convey("Then the limiter blocks again after reaching limit a second time", func() { So(limiter.paused, ShouldBeTrue) }) }) }) }
func TestHTTPSender(t *testing.T) { Convey("Given an HTTP sender with defined URL", t, func() { sender := &HTTPSender{ Config: Config{ Workers: 1, URL: "http://example.com", Insecure: true, }, } Convey("When it is initialized", func() { s := sender.Spawn(0).(*HTTPSender) Convey("Then the config should be ok", func() { So(s.Client, ShouldNotBeNil) So(s.Workers(), ShouldEqual, 1) }) }) Convey("When a message is sent and the response code is >= 400", func() { var url string s := sender.Spawn(0).(*HTTPSender) m := utils.NewMessage() m.PushPayload([]byte("Hello World")) s.Client = NewTestClient(401, func(req *http.Request) { url = req.URL.String() }) d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then the reporth should contain info about the error", func() { result := <-d.doneCalled So(result.status, ShouldEqual, "HTTPSender error: 401 Unauthorized") So(result.code, ShouldEqual, 2) So(url, ShouldEqual, "http://example.com/") d.AssertExpectations(t) }) }) Convey("When a message is received without endpoint option", func() { var url string s := sender.Spawn(0).(*HTTPSender) m := utils.NewMessage() m.PushPayload([]byte("Hello World")) s.Client = NewTestClient(200, func(req *http.Request) { url = req.URL.String() }) d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then the message should be sent via HTTP to the URL", func() { result := <-d.doneCalled So(result.status, ShouldEqual, "200 OK") So(result.code, ShouldBeZeroValue) So(url, ShouldEqual, "http://example.com/") d.AssertExpectations(t) }) }) Convey("When a message is received with endpoint option", func() { var url string s := sender.Spawn(0).(*HTTPSender) m := utils.NewMessage() m.PushPayload([]byte("Hello World")) m.Opts.Set("http_endpoint", "endpoint1") s.Client = NewTestClient(200, func(req *http.Request) { url = req.URL.String() }) d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then the message should be sent to the URL with endpoint as suffix", func() { result := <-d.doneCalled So(result.status, ShouldEqual, "200 OK") So(result.code, ShouldBeZeroValue) So(url, ShouldEqual, "http://example.com/endpoint1") d.AssertExpectations(t) }) }) Convey("When a message is received with headers", func() { var url string var headerValue string s := sender.Spawn(0).(*HTTPSender) m := utils.NewMessage() m.PushPayload([]byte("Hello World")) m.Opts.Set("http_headers", map[string]string{ "key": "value", }) s.Client = NewTestClient(200, func(req *http.Request) { url = req.URL.String() headerValue = req.Header.Get("key") }) d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then the message should be sent with headers set", func() { result := <-d.doneCalled So(result.status, ShouldEqual, "200 OK") So(result.code, ShouldBeZeroValue) So(url, ShouldEqual, "http://example.com/") So(headerValue, ShouldEqual, "value") d.AssertExpectations(t) }) }) Convey("When a message without payload is received", func() { var url string s := sender.Spawn(0).(*HTTPSender) m := utils.NewMessage() s.Client = NewTestClient(200, func(req *http.Request) { url = req.URL.String() }) d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then the message should not be sent", func() { result := <-d.doneCalled So(result.status, ShouldEqual, "Can't get payload of message: No payload available") So(result.code, ShouldBeGreaterThan, 0) So(url, ShouldBeEmpty) d.AssertExpectations(t) }) }) Convey("When a the HTTP client fails", func() { s := sender.Spawn(0).(*HTTPSender) m := utils.NewMessage() m.PushPayload(nil) s.Client = NewTestClient(200, func(req *http.Request) { req.Write(nil) }) d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then the message should not be sent", func() { result := <-d.doneCalled So(result.status, ShouldEqual, "HTTPSender error: Post http://example.com: EOF") So(result.code, ShouldBeGreaterThan, 0) d.AssertExpectations(t) }) }) }) Convey("Given an HTTP sender with invalid URL", t, func() { sender := &HTTPSender{ Config: Config{Insecure: true}, } s := sender.Spawn(0).(*HTTPSender) Convey("When try to send messages", func() { m := utils.NewMessage() m.PushPayload([]byte("Hello World")) m.Opts.Set("http_endpoint", "endpoint1") d := &Doner{ doneCalled: make(chan struct { code int status string }, 1), } d.On("Done", mock.AnythingOfType("*utils.Message"), mock.AnythingOfType("int"), mock.AnythingOfType("string")) s.OnMessage(m, d.Done) Convey("Then should fail to send messages", func() { So(s.err, ShouldNotBeNil) result := <-d.doneCalled So(result.status, ShouldEqual, "Invalid URL") So(result.code, ShouldBeGreaterThan, 0) }) }) }) }