fix(ServiceBus): Read max-message-batch-size vendor property for batch sizing

sdk/servicebus/Azure.Messaging.ServiceBus/CHANGELOG.md

@@ -16,6 +16,8 @@ Thank you to our developer community members who helped to make the Service Bus
 
 - Fixed a race condition in `AmqpSender` where concurrent calls to `CreateMessageBatchAsync` during initial AMQP link creation could observe an inconsistent `MaxBatchSize`, causing a spurious `ArgumentOutOfRangeException`. ([#56301](https://github.com/Azure/azure-sdk-for-net/issues/56301))
 
+- The sender now reads the `com.microsoft:max-message-batch-size` vendor property from the AMQP link to correctly limit batch size on Premium large-message entities, where `max-message-size` can be up to 100 MB but the batch limit is 1 MB. The 4,500 message count cap on batches has been removed as the service does not enforce a count limit. ([#44914](https://github.com/Azure/azure-sdk-for-net/issues/44914))
+
 ### Other Changes
 
 - Several areas of the AMQP transport integration have been cleaned up, modernized, and made more efficient.  _(A community contribution, courtesy of [danielmarbach](https://github.com/danielmarbach))_

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

@@ -38,6 +38,7 @@ internal class AmqpClientConstants
         public static readonly AmqpSymbol AttachEpoch = AmqpConstants.Vendor + ":epoch";
         public static readonly AmqpSymbol BatchFlushIntervalName = AmqpConstants.Vendor + ":batch-flush-interval";
         public static readonly AmqpSymbol EntityTypeName = AmqpConstants.Vendor + ":entity-type";
+        public static readonly AmqpSymbol MaxMessageBatchSizeName = AmqpConstants.Vendor + ":max-message-batch-size";
         public static readonly AmqpSymbol TransferDestinationAddress = AmqpConstants.Vendor + ":transfer-destination-address";
         public static readonly AmqpSymbol TimeoutName = AmqpConstants.Vendor + ":timeout";
         public static readonly AmqpSymbol TrackingIdName = AmqpConstants.Vendor + ":tracking-id";

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

@@ -93,12 +93,6 @@ internal class AmqpSender : TransportSender
         ///
         private (long MaxMessageSize, long MaxBatchSize) _linkLimits;
 
-        /// <summary>
-        ///   The maximum number of messages to allow in a single batch.
-        /// </summary>
-        ///
-        private int MaxMessageCount { get; set; }
-
         /// <summary>
         ///   Initializes a new instance of the <see cref="AmqpSender"/> class.
         /// </summary>
@@ -139,12 +133,6 @@ public AmqpSender(
             //
             //   Tracked by: https://github.com/Azure/azure-sdk-for-net/issues/44914
 
-            // NOTE:
-            //   This is a temporary work-around until Service Bus exposes a link property for
-            //   the maximum batch size.  The limit for batches differs from the limit for individual
-            //   messages.  Tracked by: https://github.com/Azure/azure-sdk-for-net/issues/44916
-            MaxMessageCount = 4500;
-
             _entityPath = entityPath;
             Identifier = identifier;
             _retryPolicy = retryPolicy;
@@ -230,7 +218,6 @@ internal async ValueTask<TransportMessageBatch> CreateMessageBatchInternalAsync(
             // default to the maximum size allowed by the link.
 
             options.MaxSizeInBytes ??= limits.MaxBatchSize;
-            options.MaxMessageCount ??= MaxMessageCount;
 
             Argument.AssertInRange(options.MaxSizeInBytes.Value, ServiceBusSender.MinimumBatchSizeLimit, limits.MaxBatchSize, nameof(options.MaxSizeInBytes));
             return new AmqpMessageBatch(_messageConverter, options);
@@ -667,15 +654,19 @@ protected virtual async Task<SendingAmqpLink> CreateLinkAndEnsureSenderStateAsyn
 
                 // Publish both limits as a single tuple assignment so that concurrent
                 // readers observe a consistent pair.
-                //
-                // Update with service metadata when available:
-                //  https://github.com/Azure/azure-sdk-for-net/issues/44914
-                //  https://github.com/Azure/azure-sdk-for-net/issues/44916
 
                 var maxMessageSize = (long)link.Settings.MaxMessageSize;
-                var currentLimits = _linkLimits;
 
-                _linkLimits = (maxMessageSize, Math.Min(maxMessageSize, currentLimits == default ? DefaultMaxBatchSize : currentLimits.MaxBatchSize));
+                // Read the batch size limit from the vendor property if available;
+                // fall back to the default for older service versions.
+                long maxBatchSize = DefaultMaxBatchSize;
+                if (link.Settings?.Properties.TryGetValue<long>(AmqpClientConstants.MaxMessageBatchSizeName, out var vendorBatchSize) == true
+                    && vendorBatchSize > 0)
+                {
+                    maxBatchSize = vendorBatchSize;
+                }
+
+                _linkLimits = (maxMessageSize, Math.Min(maxMessageSize, maxBatchSize));
 
                 ServiceBusEventSource.Log.CreateSendLinkComplete(Identifier);
                 link.Closed += OnSenderLinkClosed;

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Amqp/AmqpSenderTests.cs

@@ -124,15 +124,17 @@ private static CreateMessageBatchOptions GetEventBatchOptions(AmqpMessageBatch b
         ///   <see cref="AmqpSender.CreateLinkAndEnsureSenderStateAsync" />.
         /// </summary>
         ///
-        private static void SetLinkLimits(AmqpSender target, long value)
+        private static void SetLinkLimits(AmqpSender target, long value, long? maxBatchSize = null)
         {
             var defaultMaxBatchSize = (long)typeof(AmqpSender)
                 .GetField("DefaultMaxBatchSize", BindingFlags.Static | BindingFlags.NonPublic)
                 .GetValue(null);
 
+            var effectiveBatchSize = maxBatchSize ?? defaultMaxBatchSize;
+
             typeof(AmqpSender)
                 .GetField("_linkLimits", BindingFlags.Instance | BindingFlags.NonPublic)
-                .SetValue(target, (value, Math.Min(value, defaultMaxBatchSize)));
+                .SetValue(target, (value, Math.Min(value, effectiveBatchSize)));
         }
 
         /// <summary>
@@ -332,5 +334,146 @@ public async Task CreateBatchAsyncCapsMaxBatchSizeToLinkMaxMessageSize()
                 Is.EqualTo(standardTierMaxMessageSize),
                 "MaxBatchSize should be Math.Min(MaxMessageSize, DefaultMaxBatchSize).");
         }
+
+        /// <summary>
+        ///   Verifies that when the vendor property <c>com.microsoft:max-message-batch-size</c>
+        ///   reports a batch limit smaller than <c>MaxMessageSize</c>, the batch uses the
+        ///   vendor-reported limit.  This simulates a Premium tier large-message entity where
+        ///   <c>MaxMessageSize</c> is 100 MB but the batch limit is 1 MB.
+        /// </summary>
+        ///
+        [Test]
+        public async Task CreateBatchAsyncUsesVendorPropertyWhenSmallerThanMaxMessageSize()
+        {
+            var premiumLargeMaxMessageSize = 100L * 1024 * 1024; // 100 MB
+            var vendorBatchSize = 1_048_576L;                     // 1 MB
+            var retryPolicy = new BasicRetryPolicy(new ServiceBusRetryOptions { TryTimeout = TimeSpan.FromSeconds(17) });
+
+            var sender = new Mock<AmqpSender>("somePath", Mock.Of<AmqpConnectionScope>(), retryPolicy, "fake-id", new AmqpMessageConverter())
+            {
+                CallBase = true
+            };
+
+            sender
+                .Protected()
+                .Setup<Task<SendingAmqpLink>>("CreateLinkAndEnsureSenderStateAsync",
+                    ItExpr.IsAny<TimeSpan>(),
+                    ItExpr.IsAny<CancellationToken>())
+                .Callback(() => SetLinkLimits(sender.Object, premiumLargeMaxMessageSize, maxBatchSize: vendorBatchSize))
+                .Returns(Task.FromResult(new SendingAmqpLink(new AmqpLinkSettings())));
+
+            using TransportMessageBatch batch = await sender.Object.CreateMessageBatchAsync(
+                new CreateMessageBatchOptions(), default);
+
+            Assert.That(GetEventBatchOptions((AmqpMessageBatch)batch).MaxSizeInBytes,
+                Is.EqualTo(vendorBatchSize),
+                "MaxBatchSize should use the vendor property (1 MB), not MaxMessageSize (100 MB).");
+        }
+
+        /// <summary>
+        ///   Verifies that Standard tier batch sizing works when the vendor property reports
+        ///   256 KB — the same value as <c>MaxMessageSize</c>.
+        /// </summary>
+        ///
+        [Test]
+        public async Task CreateBatchAsyncUsesVendorPropertyForStandardTier()
+        {
+            var standardMaxMessageSize = 262_144L; // 256 KB
+            var vendorBatchSize = 262_144L;         // 256 KB (same as max-message-size on Standard)
+            var retryPolicy = new BasicRetryPolicy(new ServiceBusRetryOptions { TryTimeout = TimeSpan.FromSeconds(17) });
+
+            var sender = new Mock<AmqpSender>("somePath", Mock.Of<AmqpConnectionScope>(), retryPolicy, "fake-id", new AmqpMessageConverter())
+            {
+                CallBase = true
+            };
+
+            sender
+                .Protected()
+                .Setup<Task<SendingAmqpLink>>("CreateLinkAndEnsureSenderStateAsync",
+                    ItExpr.IsAny<TimeSpan>(),
+                    ItExpr.IsAny<CancellationToken>())
+                .Callback(() => SetLinkLimits(sender.Object, standardMaxMessageSize, maxBatchSize: vendorBatchSize))
+                .Returns(Task.FromResult(new SendingAmqpLink(new AmqpLinkSettings())));
+
+            using TransportMessageBatch batch = await sender.Object.CreateMessageBatchAsync(
+                new CreateMessageBatchOptions(), default);
+
+            Assert.That(GetEventBatchOptions((AmqpMessageBatch)batch).MaxSizeInBytes,
+                Is.EqualTo(vendorBatchSize),
+                "Standard tier batch size should be 256 KB.");
+        }
+
+        /// <summary>
+        ///   Verifies that when the vendor property is absent, the batch falls back to
+        ///   <c>DefaultMaxBatchSize</c> (1 MB) — not the link's <c>MaxMessageSize</c>
+        ///   when MaxMessageSize is larger.
+        /// </summary>
+        ///
+        [Test]
+        public async Task CreateBatchAsyncFallsBackToDefaultWhenVendorPropertyAbsent()
+        {
+            // Simulate a link that reports 100 MB max-message-size but no vendor property.
+            // The fallback DefaultMaxBatchSize caps batch at 1 MB.
+
+            var premiumLargeMaxMessageSize = 100L * 1024 * 1024;
+            var expectedBatchSize = 1_048_576L; // DefaultMaxBatchSize
+            var retryPolicy = new BasicRetryPolicy(new ServiceBusRetryOptions { TryTimeout = TimeSpan.FromSeconds(17) });
+
+            var sender = new Mock<AmqpSender>("somePath", Mock.Of<AmqpConnectionScope>(), retryPolicy, "fake-id", new AmqpMessageConverter())
+            {
+                CallBase = true
+            };
+
+            sender
+                .Protected()
+                .Setup<Task<SendingAmqpLink>>("CreateLinkAndEnsureSenderStateAsync",
+                    ItExpr.IsAny<TimeSpan>(),
+                    ItExpr.IsAny<CancellationToken>())
+                .Callback(() => SetLinkLimits(sender.Object, premiumLargeMaxMessageSize))
+                .Returns(Task.FromResult(new SendingAmqpLink(new AmqpLinkSettings())));
+
+            using TransportMessageBatch batch = await sender.Object.CreateMessageBatchAsync(
+                new CreateMessageBatchOptions(), default);
+
+            Assert.That(GetEventBatchOptions((AmqpMessageBatch)batch).MaxSizeInBytes,
+                Is.EqualTo(expectedBatchSize),
+                "Without vendor property, batch should fall back to DefaultMaxBatchSize (1 MB).");
+        }
+
+        /// <summary>
+        ///   Verifies that the message count is not capped — batches are limited only
+        ///   by byte size, not by message count.
+        /// </summary>
+        ///
+        [Test]
+        public async Task CreateBatchAsyncDoesNotEnforceMessageCountLimit()
+        {
+            var retryPolicy = new BasicRetryPolicy(new ServiceBusRetryOptions { TryTimeout = TimeSpan.FromSeconds(17) });
+
+            var sender = new Mock<AmqpSender>("somePath", Mock.Of<AmqpConnectionScope>(), retryPolicy, "fake-id", new AmqpMessageConverter())
+            {
+                CallBase = true
+            };
+
+            // Use a large batch size (10 MB) to ensure we can fit well over 4500
+            // tiny messages regardless of envelope overhead variations.
+            SetLinkLimits(sender.Object, 10L * 1024 * 1024, maxBatchSize: 10L * 1024 * 1024);
+
+            using TransportMessageBatch batch = await sender.Object.CreateMessageBatchAsync(
+                new CreateMessageBatchOptions(), default);
+
+            // Add 6000 tiny messages. Previously capped at 4500. With a 10 MB batch,
+            // all 6000 should fit easily (each ~200 bytes = ~1.2 MB total).
+            int added = 0;
+            for (int i = 0; i < 6000; i++)
+            {
+                if (!batch.TryAddMessage(new ServiceBusMessage(new byte[] { 0x01 })))
+                    break;
+                added++;
+            }
+
+            Assert.That(added, Is.EqualTo(6000),
+                "All 6000 messages should be accepted now that the count cap is removed.");
+        }
     }
 }