func (c *MockMiddleComponent) OnMessage(m *utils.Message, done utils.Done) {
c.Called(m)
if data, err := m.PopPayload(); err == nil {
processedData := "-> [" + string(data) + "] <-"
m.PushPayload([]byte(processedData))
}
done(m, 0, "")
}
// Add merges a new message in the buffer
func (b *Batch) Add(m *utils.Message) {
newReport := m.Reports.Pop()
b.Message.Reports.Push(newReport)
newPayload, _ := m.PopPayload()
b.Writer.Write(newPayload)
atomic.AddUint64(&b.MessageCount, 1)
}
func (c *MockComponent) OnMessage(m *utils.Message, done utils.Done) {
c.Called(m)
if data, err := m.PopPayload(); err == nil {
c.channel <- string(data)
} else {
c.channel <- err.Error()
}
done(m, c.statusCode, c.status)
}
// OnMessage will block the pipeline when the message rate is too high
func (l *Limiter) OnMessage(m *utils.Message, done utils.Done) {
if l.paused {
<-l.keepSending
l.currentMessages = 0
l.currentBytes = 0
l.paused = false
}
l.currentMessages++
if l.Config.BytesLimit > 0 {
if payload, err := m.PopPayload(); err == nil {
l.currentBytes += uint64(len(payload))
m.PushPayload(payload)
}
}
done(m, 0, "")
if l.Config.MessageLimit > 0 && l.currentMessages >= l.Config.MessageLimit ||
l.Config.BytesLimit > 0 && l.currentBytes >= l.Config.BytesLimit {
l.paused = true
}
}
// NewBatch creates a new instance of Batch
func NewBatch(m *utils.Message, group string, deflate bool, done utils.Done,
clk clock.Clock, timeoutMillis uint, ready chan *Batch) *Batch {
var wg sync.WaitGroup
payload, _ := m.PopPayload()
b := &Batch{
Group: group,
Deflate: deflate,
Done: done,
Message: m,
MessageCount: 1,
Buf: new(bytes.Buffer),
Sent: false,
}
if b.Deflate {
b.Writer = zlib.NewWriter(b.Buf)
} else {
b.Writer = bufio.NewWriter(b.Buf)
}
b.Writer.Write(payload)
if timeoutMillis > 0 {
b.Timer = clk.Timer(time.Duration(timeoutMillis) * time.Millisecond)
wg.Add(1)
go func() {
wg.Done()
<-b.Timer.C
if atomic.LoadUint64(&b.MessageCount) > 0 {
ready <- b
}
}()
}
wg.Wait()
return b
}
// OnMessage is called when a new message should be sent via HTTP
func (httpsender *HTTPSender) OnMessage(m *utils.Message, done utils.Done) {
var u string
var headers map[string]string
var code int
var status string
if httpsender.err != nil {
httpsender.logger.Debugf("Could not send message: %v", httpsender.err)
done(m, 100, httpsender.err.Error())
return
}
data, err := m.PopPayload()
defer func() {
m.PushPayload(data)
done(m, code, status)
}()
if err != nil {
httpsender.logger.Debugf("Could not send message: %v", err)
code = 101
status = "Can't get payload of message: " + err.Error()
return
}
if endpoint, exists := m.Opts.Get("http_endpoint"); exists {
u = httpsender.URL + "/" + endpoint.(string)
} else {
u = httpsender.URL
}
buf := bytes.NewBuffer(data)
req, _ := http.NewRequest("POST", u, buf)
if h, exists := m.Opts.Get("http_headers"); exists {
headers = h.(map[string]string)
for k, v := range headers {
req.Header.Add(k, v)
}
}
res, err := httpsender.Client.Do(req)
if err != nil {
httpsender.logger.Warnf(err.Error())
code = 1
status = "HTTPSender error: " + err.Error()
return
}
io.Copy(ioutil.Discard, res.Body)
res.Body.Close()
if res.StatusCode >= 400 {
httpsender.logger.Warnf("Got status: %v", res.Status)
code = 2
status = "HTTPSender error: " + res.Status
return
}
code = 0
status = res.Status
httpsender.logger.Debugf("Sent message: %v", string(buf.Bytes()))
}