snapshot of https://github.com/ThreeDotsLabs/watermill.git (master, 2026-09-15; blobs > 131072B omitted)

message/subscriber/read.go

@@ -0,0 +1,53 @@
+package subscriber
+
+import (
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// BulkRead reads provided amount of messages from the provided channel, until a timeout occurs or the limit is reached.
+func BulkRead(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool) {
+MessagesLoop:
+	for len(receivedMessages) < limit {
+		select {
+		case msg, ok := <-messagesCh:
+			if !ok {
+				break MessagesLoop
+			}
+
+			receivedMessages = append(receivedMessages, msg)
+			msg.Ack()
+		case <-time.After(timeout):
+			break MessagesLoop
+		}
+	}
+
+	return receivedMessages, len(receivedMessages) == limit
+}
+
+// BulkReadWithDeduplication reads provided number of messages from the provided channel, ignoring duplicates,
+// until a timeout occurs or the limit is reached.
+func BulkReadWithDeduplication(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool) {
+	receivedIDs := map[string]struct{}{}
+
+MessagesLoop:
+	for len(receivedMessages) < limit {
+		select {
+		case msg, ok := <-messagesCh:
+			if !ok {
+				break MessagesLoop
+			}
+
+			if _, ok := receivedIDs[msg.UUID]; !ok {
+				receivedIDs[msg.UUID] = struct{}{}
+				receivedMessages = append(receivedMessages, msg)
+			}
+			msg.Ack()
+		case <-time.After(timeout):
+			break MessagesLoop
+		}
+	}
+
+	return receivedMessages, len(receivedMessages) == limit
+}

message/subscriber/read_test.go

@@ -0,0 +1,185 @@
+package subscriber_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/subscriber"
+	"github.com/stretchr/testify/assert"
+)
+
+type bulkReadFunc func(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool)
+
+func TestBulkRead(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 100
+
+			var messages []*message.Message
+			messagesCh := make(chan *message.Message, messagesCount)
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				messages = append(messages, msg)
+				messagesCh <- msg
+			}
+
+			readMessages, all := subscriber.BulkRead(messagesCh, messagesCount, time.Second)
+			assert.True(t, all)
+
+			tests.AssertAllMessagesReceived(t, messages, readMessages)
+		})
+	}
+}
+
+func TestBulkRead_timeout(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 100
+			sendLimit := 90
+
+			messagesCh := make(chan *message.Message, messagesCount)
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				if i < sendLimit {
+					messagesCh <- msg
+				}
+			}
+
+			bulkReadStart := time.Now()
+			readMessages, all := subscriber.BulkRead(messagesCh, messagesCount, time.Millisecond)
+
+			assert.WithinDuration(t, bulkReadStart, time.Now(), time.Millisecond*100)
+			assert.False(t, all)
+			assert.Equal(t, sendLimit, len(readMessages))
+		})
+	}
+}
+
+func TestBulkRead_with_limit(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 110
+			limit := 100
+
+			messagesCh := make(chan *message.Message, messagesCount)
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				messagesCh <- msg
+			}
+
+			readMessages, all := subscriber.BulkRead(messagesCh, limit, time.Second)
+			assert.True(t, all)
+			assert.Equal(t, limit, len(readMessages))
+		})
+	}
+}
+
+func TestBulkRead_return_on_channel_close(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		BulkReadFunc bulkReadFunc
+	}{
+		{
+			Name:         "BulkRead",
+			BulkReadFunc: subscriber.BulkRead,
+		},
+		{
+			Name:         "BulkReadWithDeduplication",
+			BulkReadFunc: subscriber.BulkReadWithDeduplication,
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			messagesCount := 100
+			sendLimit := 90
+
+			messagesCh := make(chan *message.Message, messagesCount)
+			messagesChClosed := false
+
+			for i := 0; i < messagesCount; i++ {
+				msg := message.NewMessage(watermill.NewUUID(), nil)
+
+				if i < sendLimit {
+					messagesCh <- msg
+				} else if !messagesChClosed {
+					close(messagesCh)
+					messagesChClosed = true
+				}
+			}
+
+			bulkReadStart := time.Now()
+			_, all := subscriber.BulkRead(messagesCh, messagesCount, time.Second)
+
+			assert.WithinDuration(t, bulkReadStart, time.Now(), time.Millisecond*100)
+			assert.False(t, all)
+		})
+	}
+}
+
+func TestBulkReadWithDeduplication(t *testing.T) {
+	messagesCh := make(chan *message.Message, 3)
+
+	msg1 := message.NewMessage(watermill.NewUUID(), nil)
+	msg2 := message.NewMessage(watermill.NewUUID(), nil)
+	messagesCh <- msg1
+	messagesCh <- msg1
+	messagesCh <- msg2
+
+	readMessages, all := subscriber.BulkReadWithDeduplication(messagesCh, 2, time.Second)
+	assert.True(t, all)
+
+	assert.Equal(t, []string{msg1.UUID, msg2.UUID}, readMessages.IDs())
+}