Browse Source

fix inflight message tracker

pull/5890/head
chrislu 7 months ago
parent
commit
ad2b878130
  1. 17
      weed/mq/broker/broker_grpc_sub.go
  2. 95
      weed/mq/sub_coordinator/inflight_message_tracker.go
  3. 81
      weed/mq/sub_coordinator/inflight_message_tracker_test.go

17
weed/mq/broker/broker_grpc_sub.go

@ -101,11 +101,15 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
}}) }})
break break
} }
glog.V(0).Infof("topic %v partition %v subscriber %s error: %v", t, partition, clientName, err)
glog.V(0).Infof("topic %v partition %v subscriber %s lastOffset %d error: %v", t, partition, clientName, lastOffset, err)
break break
} }
if ack.GetAck().Key == nil {
// skip ack for control messages
continue
}
imt.AcknowledgeMessage(ack.GetAck().Key, ack.GetAck().Sequence) imt.AcknowledgeMessage(ack.GetAck().Key, ack.GetAck().Sequence)
currentLastOffset := imt.GetOldest()
currentLastOffset := imt.GetOldestAckedTimestamp()
fmt.Printf("%+v recv (%s,%d), oldest %d\n", partition, string(ack.GetAck().Key), ack.GetAck().Sequence, currentLastOffset) fmt.Printf("%+v recv (%s,%d), oldest %d\n", partition, string(ack.GetAck().Key), ack.GetAck().Sequence, currentLastOffset)
if subscribeFollowMeStream != nil && currentLastOffset > lastOffset { if subscribeFollowMeStream != nil && currentLastOffset > lastOffset {
if err := subscribeFollowMeStream.Send(&mq_pb.SubscribeFollowMeRequest{ if err := subscribeFollowMeStream.Send(&mq_pb.SubscribeFollowMeRequest{
@ -124,7 +128,8 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
} }
if lastOffset > 0 { if lastOffset > 0 {
if err := b.saveConsumerGroupOffset(t, partition, req.GetInit().ConsumerGroup, lastOffset); err != nil { if err := b.saveConsumerGroupOffset(t, partition, req.GetInit().ConsumerGroup, lastOffset); err != nil {
glog.Errorf("saveConsumerGroupOffset: %v", err)
glog.Errorf("saveConsumerGroupOffset partition %v lastOffset %d: %v", partition, lastOffset, err)
}
} }
if subscribeFollowMeStream != nil { if subscribeFollowMeStream != nil {
if err := subscribeFollowMeStream.Send(&mq_pb.SubscribeFollowMeRequest{ if err := subscribeFollowMeStream.Send(&mq_pb.SubscribeFollowMeRequest{
@ -135,7 +140,6 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
glog.Errorf("Error sending close to follower: %v", err) glog.Errorf("Error sending close to follower: %v", err)
} }
} }
}
}() }()
return localTopicPartition.Subscribe(clientName, startPosition, func() bool { return localTopicPartition.Subscribe(clientName, startPosition, func() bool {
@ -170,8 +174,9 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
for imt.IsInflight(logEntry.Key) { for imt.IsInflight(logEntry.Key) {
time.Sleep(137 * time.Millisecond) time.Sleep(137 * time.Millisecond)
} }
imt.InflightMessage(logEntry.Key, logEntry.TsNs)
if logEntry.Key != nil {
imt.EnflightMessage(logEntry.Key, logEntry.TsNs)
}
if err := stream.Send(&mq_pb.SubscribeMessageResponse{Message: &mq_pb.SubscribeMessageResponse_Data{ if err := stream.Send(&mq_pb.SubscribeMessageResponse{Message: &mq_pb.SubscribeMessageResponse_Data{
Data: &mq_pb.DataMessage{ Data: &mq_pb.DataMessage{

95
weed/mq/sub_coordinator/inflight_message_tracker.go

@ -1,6 +1,7 @@
package sub_coordinator package sub_coordinator
import ( import (
"fmt"
"sort" "sort"
"sync" "sync"
) )
@ -18,13 +19,14 @@ func NewInflightMessageTracker(capacity int) *InflightMessageTracker {
} }
} }
// InflightMessage tracks the message with the key and timestamp.
// EnflightMessage tracks the message with the key and timestamp.
// These messages are sent to the consumer group instances and waiting for ack. // These messages are sent to the consumer group instances and waiting for ack.
func (imt *InflightMessageTracker) InflightMessage(key []byte, tsNs int64) {
func (imt *InflightMessageTracker) EnflightMessage(key []byte, tsNs int64) {
fmt.Printf("EnflightMessage(%s,%d)\n", string(key), tsNs)
imt.mu.Lock() imt.mu.Lock()
defer imt.mu.Unlock() defer imt.mu.Unlock()
imt.messages[string(key)] = tsNs imt.messages[string(key)] = tsNs
imt.timestamps.Add(tsNs)
imt.timestamps.EnflightTimestamp(tsNs)
} }
// IsMessageAcknowledged returns true if the message has been acknowledged. // IsMessageAcknowledged returns true if the message has been acknowledged.
@ -35,7 +37,7 @@ func (imt *InflightMessageTracker) IsMessageAcknowledged(key []byte, tsNs int64)
imt.mu.Lock() imt.mu.Lock()
defer imt.mu.Unlock() defer imt.mu.Unlock()
if tsNs < imt.timestamps.Oldest() {
if tsNs <= imt.timestamps.OldestAckedTimestamp() {
return true return true
} }
if tsNs > imt.timestamps.Latest() { if tsNs > imt.timestamps.Latest() {
@ -51,6 +53,7 @@ func (imt *InflightMessageTracker) IsMessageAcknowledged(key []byte, tsNs int64)
// AcknowledgeMessage acknowledges the message with the key and timestamp. // AcknowledgeMessage acknowledges the message with the key and timestamp.
func (imt *InflightMessageTracker) AcknowledgeMessage(key []byte, tsNs int64) bool { func (imt *InflightMessageTracker) AcknowledgeMessage(key []byte, tsNs int64) bool {
fmt.Printf("AcknowledgeMessage(%s,%d)\n", string(key), tsNs)
imt.mu.Lock() imt.mu.Lock()
defer imt.mu.Unlock() defer imt.mu.Unlock()
timestamp, exists := imt.messages[string(key)] timestamp, exists := imt.messages[string(key)]
@ -59,12 +62,12 @@ func (imt *InflightMessageTracker) AcknowledgeMessage(key []byte, tsNs int64) bo
} }
delete(imt.messages, string(key)) delete(imt.messages, string(key))
// Remove the specific timestamp from the ring buffer. // Remove the specific timestamp from the ring buffer.
imt.timestamps.Remove(tsNs)
imt.timestamps.AckTimestamp(tsNs)
return true return true
} }
func (imt *InflightMessageTracker) GetOldest() int64 {
return imt.timestamps.Oldest()
func (imt *InflightMessageTracker) GetOldestAckedTimestamp() int64 {
return imt.timestamps.OldestAckedTimestamp()
} }
// IsInflight returns true if the message with the key is inflight. // IsInflight returns true if the message with the key is inflight.
@ -75,63 +78,81 @@ func (imt *InflightMessageTracker) IsInflight(key []byte) bool {
return found return found
} }
type TimestampStatus struct {
Timestamp int64
Acked bool
}
// RingBuffer represents a circular buffer to hold timestamps. // RingBuffer represents a circular buffer to hold timestamps.
type RingBuffer struct { type RingBuffer struct {
buffer []int64
buffer []*TimestampStatus
head int head int
size int size int
maxTimestamp int64
minAckedTs int64
} }
// NewRingBuffer creates a new RingBuffer of the given capacity. // NewRingBuffer creates a new RingBuffer of the given capacity.
func NewRingBuffer(capacity int) *RingBuffer { func NewRingBuffer(capacity int) *RingBuffer {
return &RingBuffer{ return &RingBuffer{
buffer: make([]int64, capacity),
buffer: newBuffer(capacity),
} }
} }
// Add adds a new timestamp to the ring buffer.
func (rb *RingBuffer) Add(timestamp int64) {
rb.buffer[rb.head] = timestamp
rb.head = (rb.head + 1) % len(rb.buffer)
func newBuffer(capacity int) []*TimestampStatus {
buffer := make([]*TimestampStatus, capacity)
for i := range buffer {
buffer[i] = &TimestampStatus{}
}
return buffer
}
// EnflightTimestamp adds a new timestamp to the ring buffer.
func (rb *RingBuffer) EnflightTimestamp(timestamp int64) {
if rb.size < len(rb.buffer) { if rb.size < len(rb.buffer) {
rb.size++ rb.size++
} else {
newBuf := newBuffer(2*len(rb.buffer))
for i := 0; i < rb.size; i++ {
newBuf[i] = rb.buffer[(rb.head+len(rb.buffer)-rb.size+i)%len(rb.buffer)]
}
rb.buffer = newBuf
rb.head = rb.size
rb.size++
}
head := rb.buffer[rb.head]
head.Timestamp = timestamp
head.Acked = false
rb.head = (rb.head + 1) % len(rb.buffer)
if timestamp > rb.maxTimestamp {
rb.maxTimestamp = timestamp
} }
} }
// Remove removes the specified timestamp from the ring buffer.
func (rb *RingBuffer) Remove(timestamp int64) {
// AckTimestamp removes the specified timestamp from the ring buffer.
func (rb *RingBuffer) AckTimestamp(timestamp int64) {
// Perform binary search // Perform binary search
index := sort.Search(rb.size, func(i int) bool { index := sort.Search(rb.size, func(i int) bool {
return rb.buffer[(rb.head+len(rb.buffer)-rb.size+i)%len(rb.buffer)] >= timestamp
return rb.buffer[(rb.head+len(rb.buffer)-rb.size+i)%len(rb.buffer)].Timestamp >= timestamp
}) })
actualIndex := (rb.head + len(rb.buffer) - rb.size + index) % len(rb.buffer) actualIndex := (rb.head + len(rb.buffer) - rb.size + index) % len(rb.buffer)
if index < rb.size && rb.buffer[actualIndex] == timestamp {
// Shift elements to maintain the buffer order
for i := index; i < rb.size-1; i++ {
fromIndex := (rb.head + len(rb.buffer) - rb.size + i + 1) % len(rb.buffer)
toIndex := (rb.head + len(rb.buffer) - rb.size + i) % len(rb.buffer)
rb.buffer[toIndex] = rb.buffer[fromIndex]
}
rb.buffer[actualIndex].Acked = true
// Remove all the acknowledged timestamps from the buffer
startPos := (rb.head + len(rb.buffer) - rb.size) % len(rb.buffer)
for i := 0; i < len(rb.buffer) && rb.buffer[(startPos+i)%len(rb.buffer)].Acked; i++ {
rb.size-- rb.size--
rb.buffer[(rb.head+len(rb.buffer)-1)%len(rb.buffer)] = 0 // Clear the last element
rb.minAckedTs = rb.buffer[(startPos+i)%len(rb.buffer)].Timestamp
} }
} }
// Oldest returns the oldest timestamp in the ring buffer.
func (rb *RingBuffer) Oldest() int64 {
if rb.size == 0 {
return 0
}
oldestIndex := (rb.head + len(rb.buffer) - rb.size) % len(rb.buffer)
return rb.buffer[oldestIndex]
// OldestAckedTimestamp returns the oldest that is already acked timestamp in the ring buffer.
func (rb *RingBuffer) OldestAckedTimestamp() int64 {
return rb.minAckedTs
} }
// Latest returns the most recently added timestamp in the ring buffer.
// Latest returns the most recently known timestamp in the ring buffer.
func (rb *RingBuffer) Latest() int64 { func (rb *RingBuffer) Latest() int64 {
if rb.size == 0 {
return 0
}
latestIndex := (rb.head + len(rb.buffer) - 1) % len(rb.buffer)
return rb.buffer[latestIndex]
return rb.maxTimestamp
} }

81
weed/mq/sub_coordinator/inflight_message_tracker_test.go

@ -1,9 +1,8 @@
package sub_coordinator package sub_coordinator
import ( import (
"sort"
"github.com/stretchr/testify/assert"
"testing" "testing"
"time"
) )
func TestRingBuffer(t *testing.T) { func TestRingBuffer(t *testing.T) {
@ -13,7 +12,7 @@ func TestRingBuffer(t *testing.T) {
// Add timestamps to the buffer // Add timestamps to the buffer
timestamps := []int64{100, 200, 300, 400, 500} timestamps := []int64{100, 200, 300, 400, 500}
for _, ts := range timestamps { for _, ts := range timestamps {
rb.Add(ts)
rb.EnflightTimestamp(ts)
} }
// Test Add method and buffer size // Test Add method and buffer size
@ -22,38 +21,25 @@ func TestRingBuffer(t *testing.T) {
t.Errorf("Expected buffer size %d, got %d", expectedSize, rb.size) t.Errorf("Expected buffer size %d, got %d", expectedSize, rb.size)
} }
// Test Oldest and Latest methods
expectedOldest := int64(100)
if oldest := rb.Oldest(); oldest != expectedOldest {
t.Errorf("Expected oldest timestamp %d, got %d", expectedOldest, oldest)
}
expectedLatest := int64(500)
if latest := rb.Latest(); latest != expectedLatest {
t.Errorf("Expected latest timestamp %d, got %d", expectedLatest, latest)
}
assert.Equal(t, int64(0), rb.OldestAckedTimestamp())
assert.Equal(t, int64(500), rb.Latest())
// Test Remove method
rb.Remove(200)
expectedSize--
if rb.size != expectedSize {
t.Errorf("Expected buffer size %d after removal, got %d", expectedSize, rb.size)
}
rb.AckTimestamp(200)
assert.Equal(t, int64(0), rb.OldestAckedTimestamp())
rb.AckTimestamp(100)
assert.Equal(t, int64(200), rb.OldestAckedTimestamp())
// Test removal of non-existent element
rb.Remove(600)
if rb.size != expectedSize {
t.Errorf("Expected buffer size %d after attempting removal of non-existent element, got %d", expectedSize, rb.size)
}
rb.EnflightTimestamp(int64(600))
rb.EnflightTimestamp(int64(700))
// Test binary search correctness
target := int64(300)
index := sort.Search(rb.size, func(i int) bool {
return rb.buffer[(rb.head+len(rb.buffer)-rb.size+i)%len(rb.buffer)] >= target
})
actualIndex := (rb.head + len(rb.buffer) - rb.size + index) % len(rb.buffer)
if rb.buffer[actualIndex] != target {
t.Errorf("Binary search failed to find the correct index for timestamp %d", target)
}
rb.AckTimestamp(500)
assert.Equal(t, int64(200), rb.OldestAckedTimestamp())
rb.AckTimestamp(400)
assert.Equal(t, int64(200), rb.OldestAckedTimestamp())
rb.AckTimestamp(300)
assert.Equal(t, int64(500), rb.OldestAckedTimestamp())
assert.Equal(t, int64(700), rb.Latest())
} }
func TestInflightMessageTracker(t *testing.T) { func TestInflightMessageTracker(t *testing.T) {
@ -61,9 +47,9 @@ func TestInflightMessageTracker(t *testing.T) {
tracker := NewInflightMessageTracker(5) tracker := NewInflightMessageTracker(5)
// Add inflight messages // Add inflight messages
key := []byte("exampleKey")
timestamp := time.Now().UnixNano()
tracker.InflightMessage(key, timestamp)
key := []byte("1")
timestamp := int64(1)
tracker.EnflightMessage(key, timestamp)
// Test IsMessageAcknowledged method // Test IsMessageAcknowledged method
isOld := tracker.IsMessageAcknowledged(key, timestamp-10) isOld := tracker.IsMessageAcknowledged(key, timestamp-10)
@ -82,4 +68,29 @@ func TestInflightMessageTracker(t *testing.T) {
if tracker.timestamps.size != 0 { if tracker.timestamps.size != 0 {
t.Error("Expected buffer size to be 0 after ack") t.Error("Expected buffer size to be 0 after ack")
} }
assert.Equal(t, timestamp, tracker.GetOldestAckedTimestamp())
}
func TestInflightMessageTracker2(t *testing.T) {
// Initialize an InflightMessageTracker with initial capacity 1
tracker := NewInflightMessageTracker(1)
tracker.EnflightMessage([]byte("1"), int64(1))
tracker.EnflightMessage([]byte("2"), int64(2))
tracker.EnflightMessage([]byte("3"), int64(3))
tracker.EnflightMessage([]byte("4"), int64(4))
tracker.EnflightMessage([]byte("5"), int64(5))
assert.True(t, tracker.AcknowledgeMessage([]byte("1"), int64(1)))
assert.Equal(t, int64(1), tracker.GetOldestAckedTimestamp())
// Test IsMessageAcknowledged method
isAcked := tracker.IsMessageAcknowledged([]byte("2"), int64(2))
if isAcked {
t.Error("Expected message to be not acked")
}
// Test AcknowledgeMessage method
assert.True(t, tracker.AcknowledgeMessage([]byte("2"), int64(2)))
assert.Equal(t, int64(2), tracker.GetOldestAckedTimestamp())
} }
Loading…
Cancel
Save