Exemple #1
0
// 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)
			})
		})
	})
}