File size: 1,966 Bytes
d6f631f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | package sink
import (
"fmt"
"sync"
"github.com/confluentinc/confluent-kafka-go/v2/kafka"
sinkmodels "github.com/openmeterio/openmeter/openmeter/sink/models"
)
type SinkBuffer struct {
mu sync.Mutex
data map[string]sinkmodels.SinkMessage
}
func NewSinkBuffer() *SinkBuffer {
return &SinkBuffer{
data: map[string]sinkmodels.SinkMessage{},
}
}
func (b *SinkBuffer) Size() int {
b.mu.Lock()
defer b.mu.Unlock()
return len(b.data)
}
func (b *SinkBuffer) Add(message sinkmodels.SinkMessage) {
b.mu.Lock()
defer b.mu.Unlock()
// Unique identifier for each message (topic + partition + offset)
key := message.KafkaMessage.String()
b.data[key] = message
}
type MessageTransformerFunc func(*sinkmodels.SinkMessage)
func (b *SinkBuffer) Dequeue(transformers ...MessageTransformerFunc) []sinkmodels.SinkMessage {
b.mu.Lock()
defer b.mu.Unlock()
messages := make([]sinkmodels.SinkMessage, 0, len(b.data))
for key, message := range b.data {
for _, transformer := range transformers {
transformer(&message)
}
messages = append(messages, message)
delete(b.data, key)
}
return messages
}
// RemoveByPartitions removes messages from the buffer by partitions
// Useful when partitions are revoked.
func (b *SinkBuffer) RemoveByPartitions(partitions []kafka.TopicPartition) {
b.mu.Lock()
defer b.mu.Unlock()
partitionMap := map[string]bool{}
for _, topicPartition := range partitions {
key := topicPartitionKey(topicPartition)
partitionMap[key] = true
}
for key, message := range b.data {
topicKey := topicPartitionKey(message.KafkaMessage.TopicPartition)
if partitionMap[topicKey] {
delete(b.data, key)
}
}
}
func topicPartitionKey(partition kafka.TopicPartition) string {
var topic string
if partition.Topic != nil {
topic = *partition.Topic
}
return partitionKey(topic, partition.Partition)
}
func partitionKey(topic string, partition int32) string {
return fmt.Sprintf("%s-%d", topic, partition)
}
|