Repository: Azure/azure-service-bus-go
Pull request #80 (2018-12-13 18:06:15)

add prefetch support for queues and subscriptions
- add prefetch support for queues and subscriptions
- add tests

#78 


Patch:
diff --git a/queue.go b/queue.go
index 93eeb3b..4ada70c 100644
--- a/queue.go
+++ b/queue.go
@@ -57,6 +57,7 @@ type (
 		senderMu          sync.Mutex
 		receiveMode       ReceiveMode
 		requiredSessionID *string
+		prefetchCount     *uint32
 	}
 
 	// queueContent is a specialized Queue body for an Atom entry
@@ -126,6 +127,21 @@ func QueueWithReceiveAndDelete() QueueOption {
 	}
 }
 
+// QueueWithPrefetchCount configures the queue to attempt to fetch the number of messages specified by the
+// prefetch count at one time.
+//
+// The default is 1 message at a time.
+//
+// Caution: Using PeekLock, messages have a set lock timeout, which can be renewed. By setting a high prefetch count, a
+// local queue of messages could build up and cause message locks to expire before the message lands in the handler. If
+// this happens, the message disposition will fail and will be re-queued and processed again.
+func QueueWithPrefetchCount(prefetch uint32) QueueOption {
+	return func(q *Queue) error {
+		q.prefetchCount = &prefetch
+		return nil
+	}
+}
+
 // NewQueue creates a new Queue Sender / Receiver
 func (ns *Namespace) NewQueue(name string, opts ...QueueOption) (*Queue, error) {
 	queue := &Queue{
@@ -475,6 +491,11 @@ func (q *Queue) newReceiver(ctx context.Context, opts ...ReceiverOption) (*Recei
 	defer span.Finish()
 
 	opts = append(opts, ReceiverWithReceiveMode(q.receiveMode))
+
+	if q.prefetchCount != nil {
+		opts = append(opts, ReceiverWithPrefetchCount(*q.prefetchCount))
+	}
+
 	return q.namespace.NewReceiver(ctx, q.Name, opts...)
 }
 
diff --git a/receiver.go b/receiver.go
index ae9b40c..b986648 100644
--- a/receiver.go
+++ b/receiver.go
@@ -78,6 +78,21 @@ func ReceiverWithReceiveMode(mode ReceiveMode) ReceiverOption {
 	}
 }
 
+// ReceiverWithPrefetchCount configures the receiver to attempt to fetch the number of messages specified by the prefect
+// at one time.
+//
+// The default is 1 message at a time.
+//
+// Caution: Using PeekLock, messages have a set lock timeout, which can be renewed. By setting a high prefetch count, a
+// local queue of messages could build up and cause message locks to expire before the message lands in the handler. If
+// this happens, the message disposition will fail and will be re-queued and processed again.
+func ReceiverWithPrefetchCount(prefetch uint32) ReceiverOption {
+	return func(receiver *Receiver) error {
+		receiver.prefetch = prefetch
+		return nil
+	}
+}
+
 // NewReceiver creates a new Service Bus message listener given an AMQP client and an entity path
 func (ns *Namespace) NewReceiver(ctx context.Context, entityPath string, opts ...ReceiverOption) (*Receiver, error) {
 	span, ctx := ns.startSpanFromContext(ctx, "sb.Hub.NewReceiver")
diff --git a/subscription.go b/subscription.go
index aab182d..c6e57c5 100644
--- a/subscription.go
+++ b/subscription.go
@@ -43,6 +43,7 @@ type (
 		receiverMu        sync.Mutex
 		receiveMode       ReceiveMode
 		requiredSessionID *string
+		prefetchCount     *uint32
 	}
 
 	// SubscriptionDescription is the content type for Subscription management requests
@@ -80,6 +81,21 @@ func SubscriptionWithReceiveAndDelete() SubscriptionOption {
 	}
 }
 
+// SubscriptionWithPrefetchCount configures the subscription to attempt to fetch the number of messages specified by the
+// prefetch count at one time.
+//
+// The default is 1 message at a time.
+//
+// Caution: Using PeekLock, messages have a set lock timeout, which can be renewed. By setting a high prefetch count, a
+// local queue of messages could build up and cause message locks to expire before the message lands in the handler. If
+// this happens, the message disposition will fail and will be re-queued and processed again.
+func SubscriptionWithPrefetchCount(prefetch uint32) SubscriptionOption {
+	return func(q *Subscription) error {
+		q.prefetchCount = &prefetch
+		return nil
+	}
+}
+
 // NewSubscription creates a new Topic Subscription client
 func (t *Topic) NewSubscription(name string, opts ...SubscriptionOption) (*Subscription, error) {
 	sub := &Subscription{
@@ -185,6 +201,11 @@ func (s *Subscription) NewReceiver(ctx context.Context, opts ...ReceiverOption)
 	defer span.Finish()
 
 	opts = append(opts, ReceiverWithReceiveMode(s.receiveMode))
+
+	if s.prefetchCount != nil {
+		opts = append(opts, ReceiverWithPrefetchCount(*s.prefetchCount))
+	}
+
 	return s.namespace.NewReceiver(ctx, s.Topic.Name+"/Subscriptions/"+s.Name, opts...)
 }
 


Test patch:
diff --git a/queue_test.go b/queue_test.go
index 52c3c7a..9ebd1a6 100644
--- a/queue_test.go
+++ b/queue_test.go
@@ -597,6 +597,31 @@ func testQueueSendAndReceiveWithReceiveAndDelete(ctx context.Context, t *testing
 	}
 }
 
+func (suite *serviceBusSuite) TestQueueWithPrefetch() {
+	tests := map[string]func(context.Context, *testing.T, *Queue){
+		"SendAndReceive": testQueueSendAndReceive,
+	}
+	suite.queueMessageTestWithQueueOptions(tests, QueueWithPrefetchCount(10))
+}
+
+func testQueueSendAndReceive(ctx context.Context, t *testing.T, q *Queue) {
+	messages := []string{"foo", "bar", "bazz", "buzz"}
+	for _, msg := range messages {
+		require.NoError(t, q.Send(ctx, NewMessageFromString(msg)))
+	}
+
+	count := 0
+	for idx := range messages {
+		err := q.ReceiveOne(ctx, HandlerFunc(func(ctx context.Context, msg *Message) error {
+			assert.Equal(t, messages[idx], string(msg.Data))
+			count++
+			return msg.Complete(ctx)
+		}))
+		assert.NoError(t, err)
+	}
+	assert.Len(t, messages, count)
+}
+
 func (suite *serviceBusSuite) TestIssue73QueueClient() {
 	tests := map[string]func(context.Context, *testing.T, *Queue){
 		"SimpleSend200_NoZeroCheck": func(ctx context.Context, t *testing.T, queue *Queue) {
diff --git a/subscription_test.go b/subscription_test.go
index 4e795a3..891c801 100644
--- a/subscription_test.go
+++ b/subscription_test.go
@@ -348,6 +348,32 @@ func testSubscriptionReceiveOneFromDeadLetter(ctx context.Context, t *testing.T,
 	assert.NoError(t, err)
 }
 
+func (suite *serviceBusSuite) TestSubscriptionWithPrefetch() {
+	tests := map[string]func(context.Context, *testing.T, *Topic, *Subscription){
+		"SendAndReceive": testSubscriptionSendAndReceive,
+	}
+	suite.subscriptionMessageTestWithOptions(tests, SubscriptionWithPrefetchCount(10))
+}
+
+func testSubscriptionSendAndReceive(ctx context.Context, t *testing.T, topic *Topic, s *Subscription) {
+	messages := []string{"foo", "bar", "bazz", "buzz"}
+	for _, msg := range messages {
+		require.NoError(t, topic.Send(ctx, NewMessageFromString(msg)))
+	}
+
+	count := 0
+	for idx := range messages {
+		err := s.ReceiveOne(ctx, HandlerFunc(func(ctx context.Context, msg *Message) error {
+			assert.Equal(t, messages[idx], string(msg.Data))
+			count++
+			return msg.Complete(ctx)
+		}))
+		assert.NoError(t, err)
+	}
+	assert.Len(t, messages, count)
+}
+
+
 func (suite *serviceBusSuite) TestSubscriptionSessionClient() {
 	tests := map[string]func(context.Context, *testing.T, *TopicSession, *SubscriptionSession){
 		"ReceiveOne": testSubscriptionSessionReceiveOne,