4e909185d1
* Add GetRecords function to Kinesis Utils * Add sample .env to showcase ENV vars needed * Fix RedisCheckpoint to look for empty string (default from Redis) * Extract example code into its own repository
62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
package buffers
|
|
|
|
import "bytes"
|
|
|
|
type MsgBuffer struct {
|
|
buffer bytes.Buffer
|
|
firstSequenceNumber string
|
|
lastSequenceNumber string
|
|
NumMessagesToBuffer int
|
|
sequencesInBuffer []string
|
|
}
|
|
|
|
func NewMessageBuffer(n int) *MsgBuffer {
|
|
return &MsgBuffer{NumMessagesToBuffer: n}
|
|
}
|
|
|
|
func (b *MsgBuffer) ConsumeRecord(data []byte, sequenceNumber string) {
|
|
if len(b.sequencesInBuffer) == 0 {
|
|
b.firstSequenceNumber = sequenceNumber
|
|
}
|
|
|
|
b.lastSequenceNumber = sequenceNumber
|
|
|
|
if !b.SequenceExists(sequenceNumber) {
|
|
b.buffer.Write(data)
|
|
b.sequencesInBuffer = append(b.sequencesInBuffer, sequenceNumber)
|
|
}
|
|
}
|
|
|
|
func (b MsgBuffer) SequenceExists(sequenceNumber string) bool {
|
|
for _, v := range b.sequencesInBuffer {
|
|
if v == sequenceNumber {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (b MsgBuffer) Data() []byte {
|
|
return b.buffer.Bytes()
|
|
}
|
|
|
|
func (b MsgBuffer) NumMessagesInBuffer() int {
|
|
return len(b.sequencesInBuffer)
|
|
}
|
|
|
|
func (b *MsgBuffer) Flush() {
|
|
b.buffer.Reset()
|
|
b.sequencesInBuffer = b.sequencesInBuffer[:0]
|
|
}
|
|
|
|
func (b MsgBuffer) ShouldFlush() bool {
|
|
return len(b.sequencesInBuffer) >= b.NumMessagesToBuffer
|
|
}
|
|
|
|
func (b MsgBuffer) FirstSequenceNumber() string {
|
|
return b.firstSequenceNumber
|
|
}
|
|
|
|
func (b MsgBuffer) LastSequenceNumber() string {
|
|
return b.lastSequenceNumber
|
|
}
|