Adding MessageState property to ServiceBusReceivedMessage (Azure#25446)

* Adding MessageState property to ServiceBusReceivedMessage

* Rename enum

* Hide previous method

* Augment a few tests

* Update API spec

* Apply suggestions from code review

Co-authored-by: JoshLove-msft <54595583+JoshLove-msft@users.noreply.github.com>

* Update ServiceBusModelFactory.cs

Co-authored-by: JoshLove-msft <54595583+JoshLove-msft@users.noreply.github.com>

Author: danielmarbach

sdk/servicebus/Azure.Messaging.ServiceBus/api/Azure.Messaging.ServiceBus.netstandard2.0.cs

@@ -185,6 +185,12 @@ internal ServiceBusMessageBatch() { }
         public void Dispose() { }
         public bool TryAddMessage(Azure.Messaging.ServiceBus.ServiceBusMessage message) { throw null; }
     }
+    public enum ServiceBusMessageState
+    {
+        Active = 0,
+        Deferred = 1,
+        Scheduled = 2,
+    }
     public static partial class ServiceBusModelFactory
     {
         public static Azure.Messaging.ServiceBus.Administration.QueueProperties QueueProperties(string name, System.TimeSpan lockDuration, long maxSizeInMegabytes, bool requiresDuplicateDetection, bool requiresSession, System.TimeSpan defaultMessageTimeToLive, System.TimeSpan autoDeleteOnIdle, bool deadLetteringOnMessageExpiration, System.TimeSpan duplicateDetectionHistoryTimeWindow, int maxDeliveryCount, bool enableBatchedOperations, Azure.Messaging.ServiceBus.Administration.EntityStatus status, string forwardTo, string forwardDeadLetteredMessagesTo, string userMetadata, bool enablePartitioning) { throw null; }
@@ -193,7 +199,9 @@ public static partial class ServiceBusModelFactory
         public static Azure.Messaging.ServiceBus.Administration.QueueRuntimeProperties QueueRuntimeProperties(string name, long activeMessageCount = (long)0, long scheduledMessageCount = (long)0, long deadLetterMessageCount = (long)0, long transferDeadLetterMessageCount = (long)0, long transferMessageCount = (long)0, long totalMessageCount = (long)0, long sizeInBytes = (long)0, System.DateTimeOffset createdAt = default(System.DateTimeOffset), System.DateTimeOffset updatedAt = default(System.DateTimeOffset), System.DateTimeOffset accessedAt = default(System.DateTimeOffset)) { throw null; }
         public static Azure.Messaging.ServiceBus.Administration.RuleProperties RuleProperties(string name, Azure.Messaging.ServiceBus.Administration.RuleFilter filter = null, Azure.Messaging.ServiceBus.Administration.RuleAction action = null) { throw null; }
         public static Azure.Messaging.ServiceBus.ServiceBusMessageBatch ServiceBusMessageBatch(long batchSizeBytes, System.Collections.Generic.IList<Azure.Messaging.ServiceBus.ServiceBusMessage> batchMessageStore, Azure.Messaging.ServiceBus.CreateMessageBatchOptions batchOptions = null, System.Func<Azure.Messaging.ServiceBus.ServiceBusMessage, bool> tryAddCallback = null) { throw null; }
-        public static Azure.Messaging.ServiceBus.ServiceBusReceivedMessage ServiceBusReceivedMessage(System.BinaryData body = null, string messageId = null, string partitionKey = null, string viaPartitionKey = null, string sessionId = null, string replyToSessionId = null, System.TimeSpan timeToLive = default(System.TimeSpan), string correlationId = null, string subject = null, string to = null, string contentType = null, string replyTo = null, System.DateTimeOffset scheduledEnqueueTime = default(System.DateTimeOffset), System.Collections.Generic.IDictionary<string, object> properties = null, System.Guid lockTokenGuid = default(System.Guid), int deliveryCount = 0, System.DateTimeOffset lockedUntil = default(System.DateTimeOffset), long sequenceNumber = (long)-1, string deadLetterSource = null, long enqueuedSequenceNumber = (long)0, System.DateTimeOffset enqueuedTime = default(System.DateTimeOffset)) { throw null; }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public static Azure.Messaging.ServiceBus.ServiceBusReceivedMessage ServiceBusReceivedMessage(System.BinaryData body, string messageId, string partitionKey, string viaPartitionKey, string sessionId, string replyToSessionId, System.TimeSpan timeToLive, string correlationId, string subject, string to, string contentType, string replyTo, System.DateTimeOffset scheduledEnqueueTime, System.Collections.Generic.IDictionary<string, object> properties, System.Guid lockTokenGuid, int deliveryCount, System.DateTimeOffset lockedUntil, long sequenceNumber, string deadLetterSource, long enqueuedSequenceNumber, System.DateTimeOffset enqueuedTime) { throw null; }
+        public static Azure.Messaging.ServiceBus.ServiceBusReceivedMessage ServiceBusReceivedMessage(System.BinaryData body = null, string messageId = null, string partitionKey = null, string viaPartitionKey = null, string sessionId = null, string replyToSessionId = null, System.TimeSpan timeToLive = default(System.TimeSpan), string correlationId = null, string subject = null, string to = null, string contentType = null, string replyTo = null, System.DateTimeOffset scheduledEnqueueTime = default(System.DateTimeOffset), System.Collections.Generic.IDictionary<string, object> properties = null, System.Guid lockTokenGuid = default(System.Guid), int deliveryCount = 0, System.DateTimeOffset lockedUntil = default(System.DateTimeOffset), long sequenceNumber = (long)-1, string deadLetterSource = null, long enqueuedSequenceNumber = (long)0, System.DateTimeOffset enqueuedTime = default(System.DateTimeOffset), Azure.Messaging.ServiceBus.ServiceBusMessageState serviceBusMessageState = Azure.Messaging.ServiceBus.ServiceBusMessageState.Active) { throw null; }
         public static Azure.Messaging.ServiceBus.Administration.SubscriptionProperties SubscriptionProperties(string topicName, string subscriptionName, System.TimeSpan lockDuration = default(System.TimeSpan), bool requiresSession = false, System.TimeSpan defaultMessageTimeToLive = default(System.TimeSpan), System.TimeSpan autoDeleteOnIdle = default(System.TimeSpan), bool deadLetteringOnMessageExpiration = false, int maxDeliveryCount = 0, bool enableBatchedOperations = false, Azure.Messaging.ServiceBus.Administration.EntityStatus status = default(Azure.Messaging.ServiceBus.Administration.EntityStatus), string forwardTo = null, string forwardDeadLetteredMessagesTo = null, string userMetadata = null) { throw null; }
         public static Azure.Messaging.ServiceBus.Administration.SubscriptionRuntimeProperties SubscriptionRuntimeProperties(string topicName, string subscriptionName, long activeMessageCount = (long)0, long deadLetterMessageCount = (long)0, long transferDeadLetterMessageCount = (long)0, long transferMessageCount = (long)0, long totalMessageCount = (long)0, System.DateTimeOffset createdAt = default(System.DateTimeOffset), System.DateTimeOffset updatedAt = default(System.DateTimeOffset), System.DateTimeOffset accessedAt = default(System.DateTimeOffset)) { throw null; }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
@@ -270,6 +278,7 @@ internal ServiceBusReceivedMessage() { }
         public System.DateTimeOffset ScheduledEnqueueTime { get { throw null; } }
         public long SequenceNumber { get { throw null; } }
         public string SessionId { get { throw null; } }
+        public Azure.Messaging.ServiceBus.ServiceBusMessageState State { get { throw null; } }
         public string Subject { get { throw null; } }
         public System.TimeSpan TimeToLive { get { throw null; } }
         public string To { get { throw null; } }

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpMessageConstants.cs

@@ -20,6 +20,7 @@ internal class AmqpMessageConstants
         internal const string PartitionIdName = "x-opt-partition-id";
         internal const string ViaPartitionKeyName = "x-opt-via-partition-key";
         internal const string DeadLetterSourceName = "x-opt-deadletter-source";
+        internal const string MessageStateName = "x-opt-message-state";
         internal const string TimeSpanName = AmqpConstants.Vendor + ":timespan";
         internal const string UriName = AmqpConstants.Vendor + ":uri";
         internal const string DateTimeOffsetName = AmqpConstants.Vendor + ":datetime-offset";

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpMessageConverter.cs

@@ -469,6 +469,13 @@ public static ServiceBusReceivedMessage AmqpMessageToSBMessage(AmqpMessage amqpM
                         case AmqpMessageConstants.DeadLetterSourceName:
                             annotatedMessage.MessageAnnotations[AmqpMessageConstants.DeadLetterSourceName] = pair.Value;
                             break;
+                        case AmqpMessageConstants.MessageStateName:
+                            int enumValue = (int)pair.Value;
+                            if (Enum.IsDefined(typeof(ServiceBusMessageState), enumValue))
+                            {
+                                annotatedMessage.MessageAnnotations[AmqpMessageConstants.MessageStateName] = (ServiceBusMessageState)enumValue;
+                            }
+                            break;
                         default:
                             if (TryGetNetObjectFromAmqpObject(pair.Value, MappingType.ApplicationProperty, out var netObject))
                             {

sdk/servicebus/Azure.Messaging.ServiceBus/src/Primitives/ServiceBusMessage.cs

@@ -115,7 +115,7 @@ public ServiceBusMessage(ServiceBusReceivedMessage receivedMessage)
             {
                 if (kvp.Key == AmqpMessageConstants.LockedUntilName || kvp.Key == AmqpMessageConstants.SequenceNumberName ||
                     kvp.Key == AmqpMessageConstants.DeadLetterSourceName || kvp.Key == AmqpMessageConstants.EnqueueSequenceNumberName ||
-                    kvp.Key == AmqpMessageConstants.EnqueuedTimeUtcName)
+                    kvp.Key == AmqpMessageConstants.EnqueuedTimeUtcName || kvp.Key == AmqpMessageConstants.MessageStateName)
                 {
                     continue;
                 }

sdk/servicebus/Azure.Messaging.ServiceBus/src/Primitives/ServiceBusMessageState.cs

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+namespace Azure.Messaging.ServiceBus
+{
+    /// <summary>Represents the message state of the <see cref="ServiceBusReceivedMessage"/></summary>
+    public enum ServiceBusMessageState
+    {
+        /// <summary>Specifies an active message state.</summary>
+        Active = 0,
+
+        /// <summary>Specifies a deferred message state.</summary>
+        Deferred = 1,
+
+        /// <summary>Specifies the scheduled message state.</summary>
+        Scheduled = 2
+    }
+}

sdk/servicebus/Azure.Messaging.ServiceBus/src/Primitives/ServiceBusModelFactory.cs

@@ -17,6 +17,36 @@ namespace Azure.Messaging.ServiceBus
     /// </summary>
     public static class ServiceBusModelFactory
     {
+        /// <summary>
+        /// Creates a new ServiceBusReceivedMessage instance for mocking.
+        /// </summary>
+        [EditorBrowsable(EditorBrowsableState.Never)]
+        public static ServiceBusReceivedMessage ServiceBusReceivedMessage(
+            BinaryData body,
+            string messageId,
+            string partitionKey,
+            string viaPartitionKey,
+            string sessionId,
+            string replyToSessionId,
+            TimeSpan timeToLive,
+            string correlationId,
+            string subject,
+            string to,
+            string contentType,
+            string replyTo,
+            DateTimeOffset scheduledEnqueueTime,
+            IDictionary<string, object> properties,
+            Guid lockTokenGuid,
+            int deliveryCount,
+            DateTimeOffset lockedUntil,
+            long sequenceNumber,
+            string deadLetterSource,
+            long enqueuedSequenceNumber,
+            DateTimeOffset enqueuedTime) =>
+            ServiceBusReceivedMessage(body, messageId, partitionKey, viaPartitionKey, sessionId, replyToSessionId,
+            timeToLive, correlationId, subject, to, contentType, replyTo, scheduledEnqueueTime, properties,
+            lockTokenGuid, deliveryCount, lockedUntil, sequenceNumber, deadLetterSource, enqueuedSequenceNumber, enqueuedTime, ServiceBusMessageState.Active);
+
         /// <summary>
         /// Creates a new ServiceBusReceivedMessage instance for mocking.
         /// </summary>
@@ -41,7 +71,8 @@ public static ServiceBusReceivedMessage ServiceBusReceivedMessage(
             long sequenceNumber = -1,
             string deadLetterSource = default,
             long enqueuedSequenceNumber = default,
-            DateTimeOffset enqueuedTime = default)
+            DateTimeOffset enqueuedTime = default,
+            ServiceBusMessageState serviceBusMessageState = default)
         {
             var amqpMessage = new AmqpAnnotatedMessage(new AmqpMessageBody(new ReadOnlyMemory<byte>[] { body }));
 
@@ -98,6 +129,7 @@ public static ServiceBusReceivedMessage ServiceBusReceivedMessage(
             amqpMessage.MessageAnnotations[AmqpMessageConstants.DeadLetterSourceName] = deadLetterSource;
             amqpMessage.MessageAnnotations[AmqpMessageConstants.EnqueueSequenceNumberName] = enqueuedSequenceNumber;
             amqpMessage.MessageAnnotations[AmqpMessageConstants.EnqueuedTimeUtcName] = enqueuedTime.UtcDateTime;
+            amqpMessage.MessageAnnotations[AmqpMessageConstants.MessageStateName] = serviceBusMessageState;
 
             return new ServiceBusReceivedMessage(amqpMessage)
             {

sdk/servicebus/Azure.Messaging.ServiceBus/src/Primitives/ServiceBusReceivedMessage.cs

@@ -408,6 +408,31 @@ public string DeadLetterErrorDescription
             }
         }
 
+        /// <summary>Gets the state of the message.</summary>
+        /// <value>The state of the message. </value>
+        /// <remarks>
+        ///    The state of the message can be Active, Deferred, or Scheduled. Deferred messages have Deferred state,
+        ///    scheduled messages have Scheduled state, all other messages have Active state.
+        /// </remarks>
+        public ServiceBusMessageState State
+        {
+            get
+            {
+                if (AmqpMessage.MessageAnnotations.TryGetValue(
+                    AmqpMessageConstants.MessageStateName,
+                    out object val))
+                {
+                    return (ServiceBusMessageState)val;
+                }
+
+                return default;
+            }
+            internal set
+            {
+                AmqpMessage.MessageAnnotations[AmqpMessageConstants.MessageStateName] = value;
+            }
+        }
+
         /// <summary>Returns a string that represents the current message.</summary>
         /// <returns>The string representation of the current message.</returns>
         public override string ToString()

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Message/MessageLiveTests.cs

@@ -152,6 +152,7 @@ public async Task CreateFromReceivedMessageCopiesProperties()
                 Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.EnqueueSequenceNumberName));
                 Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.EnqueuedTimeUtcName));
                 Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.DeadLetterSourceName));
+                Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.MessageStateName));
                 Assert.IsFalse(toSend.ApplicationProperties.ContainsKey(AmqpMessageConstants.DeadLetterReasonHeader));
                 Assert.IsFalse(toSend.ApplicationProperties.ContainsKey(AmqpMessageConstants.DeadLetterErrorDescriptionHeader));
 
@@ -209,6 +210,7 @@ public async Task CreateFromReceivedMessageCopiesPropertiesTopic()
                 Assert.IsTrue(rawReceived.MessageAnnotations.ContainsKey(AmqpMessageConstants.SequenceNumberName));
                 Assert.IsTrue(rawReceived.MessageAnnotations.ContainsKey(AmqpMessageConstants.EnqueueSequenceNumberName));
                 Assert.IsTrue(rawReceived.MessageAnnotations.ContainsKey(AmqpMessageConstants.EnqueuedTimeUtcName));
+                Assert.IsTrue(rawReceived.MessageAnnotations.ContainsKey(AmqpMessageConstants.MessageStateName));
 
                 AssertMessagesEqual(msg, received);
                 var toSend = new ServiceBusMessage(received);
@@ -222,6 +224,7 @@ public async Task CreateFromReceivedMessageCopiesPropertiesTopic()
                 Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.EnqueueSequenceNumberName));
                 Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.EnqueuedTimeUtcName));
                 Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.DeadLetterSourceName));
+                Assert.IsFalse(rawSend.MessageAnnotations.ContainsKey(AmqpMessageConstants.MessageStateName));
                 Assert.IsFalse(toSend.ApplicationProperties.ContainsKey(AmqpMessageConstants.DeadLetterReasonHeader));
                 Assert.IsFalse(toSend.ApplicationProperties.ContainsKey(AmqpMessageConstants.DeadLetterErrorDescriptionHeader));
 

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Receiver/ReceiverLiveTests.cs

@@ -490,6 +490,7 @@ public async Task DeferMessagesList()
                 {
                     Assert.AreEqual(messageList[i].MessageId, deferredMessages[i].MessageId);
                     Assert.AreEqual(messageList[i].Body.ToArray(), deferredMessages[i].Body.ToArray());
+                    Assert.AreEqual(ServiceBusMessageState.Deferred, deferredMessages[i]);
                 }
 
                 // verify that looking up a non-existent sequence number will throw
@@ -544,6 +545,7 @@ public async Task DeferMessagesArray()
                 {
                     Assert.AreEqual(messageList[i].MessageId, deferredMessages[i].MessageId);
                     Assert.AreEqual(messageList[i].Body.ToArray(), deferredMessages[i].Body.ToArray());
+                    Assert.AreEqual(ServiceBusMessageState.Deferred, deferredMessages[i]);
                 }
 
                 // verify that an empty array can be passed
@@ -600,6 +602,7 @@ IEnumerable<long> GetEnumerable()
                 {
                     Assert.AreEqual(messageList[i].MessageId, deferredMessages[i].MessageId);
                     Assert.AreEqual(messageList[i].Body.ToArray(), deferredMessages[i].Body.ToArray());
+                    Assert.AreEqual(ServiceBusMessageState.Deferred, deferredMessages[i]);
                 }
 
                 // verify that an empty enumerable can be passed
@@ -626,10 +629,12 @@ public async Task CanPeekADeferredMessage()
                 var peekedMsg = await receiver.PeekMessageAsync();
                 Assert.AreEqual(receivedMsg.MessageId, peekedMsg.MessageId);
                 Assert.AreEqual(receivedMsg.SequenceNumber, peekedMsg.SequenceNumber);
+                Assert.AreEqual(ServiceBusMessageState.Deferred, receivedMsg.State);
 
                 var deferredMsg = await receiver.ReceiveDeferredMessageAsync(peekedMsg.SequenceNumber);
                 Assert.AreEqual(peekedMsg.MessageId, deferredMsg.MessageId);
                 Assert.AreEqual(peekedMsg.SequenceNumber, deferredMsg.SequenceNumber);
+                Assert.AreEqual(peekedMsg.State, receivedMsg.State);
             }
         }
 

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Sender/SenderLiveTests.cs

@@ -190,6 +190,7 @@ public async Task Schedule()
                 await using var receiver = client.CreateReceiver(scope.QueueName);
                 ServiceBusReceivedMessage msg = await receiver.PeekMessageAsync(seq);
                 Assert.AreEqual(0, Convert.ToInt32(new TimeSpan(scheduleTime.Ticks - msg.ScheduledEnqueueTime.Ticks).TotalSeconds));
+                Assert.AreEqual(ServiceBusMessageState.Scheduled, msg.State);
 
                 await sender.CancelScheduledMessageAsync(seq);
                 msg = await receiver.PeekMessageAsync(seq);
@@ -211,6 +212,7 @@ public async Task ScheduleMultipleArray()
                 {
                     ServiceBusReceivedMessage msg = await receiver.PeekMessageAsync(seq);
                     Assert.AreEqual(0, Convert.ToInt32(new TimeSpan(scheduleTime.Ticks - msg.ScheduledEnqueueTime.Ticks).TotalSeconds));
+                    Assert.AreEqual(ServiceBusMessageState.Scheduled, msg.State);
                 }
                 await sender.CancelScheduledMessagesAsync(sequenceNumbers: sequenceNums);
 
@@ -244,6 +246,7 @@ public async Task ScheduleMultipleList()
                 {
                     ServiceBusReceivedMessage msg = await receiver.PeekMessageAsync(seq);
                     Assert.AreEqual(0, Convert.ToInt32(new TimeSpan(scheduleTime.Ticks - msg.ScheduledEnqueueTime.Ticks).TotalSeconds));
+                    Assert.AreEqual(ServiceBusMessageState.Scheduled, msg.State);
                 }
                 await sender.CancelScheduledMessagesAsync(sequenceNumbers: new List<long>(sequenceNums));
 
@@ -277,6 +280,7 @@ public async Task ScheduleMultipleEnumerable()
                 {
                     ServiceBusReceivedMessage msg = await receiver.PeekMessageAsync(seq);
                     Assert.AreEqual(0, Convert.ToInt32(new TimeSpan(scheduleTime.Ticks - msg.ScheduledEnqueueTime.Ticks).TotalSeconds));
+                    Assert.AreEqual(ServiceBusMessageState.Scheduled, msg.State);
                 }
 
                 // use an enumerable