Add ServiceBusTransportMetrics (Azure#27663)

* Add ServiceBusTransportMetrics

* PR fb

* Fix tests

* Fix test

Author: JoshLove-msft

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

@@ -87,11 +87,13 @@ public ServiceBusClient(string connectionString, Azure.Messaging.ServiceBus.Serv
         public virtual Azure.Messaging.ServiceBus.ServiceBusSessionProcessor CreateSessionProcessor(string queueName, Azure.Messaging.ServiceBus.ServiceBusSessionProcessorOptions options = null) { throw null; }
         public virtual Azure.Messaging.ServiceBus.ServiceBusSessionProcessor CreateSessionProcessor(string topicName, string subscriptionName, Azure.Messaging.ServiceBus.ServiceBusSessionProcessorOptions options = null) { throw null; }
         public virtual System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
+        public virtual Azure.Messaging.ServiceBus.ServiceBusTransportMetrics GetTransportMetrics() { throw null; }
     }
     public partial class ServiceBusClientOptions
     {
         public ServiceBusClientOptions() { }
         public bool EnableCrossEntityTransactions { get { throw null; } set { } }
+        public bool EnableTransportMetrics { get { throw null; } set { } }
         public Azure.Messaging.ServiceBus.ServiceBusRetryOptions RetryOptions { get { throw null; } set { } }
         public Azure.Messaging.ServiceBus.ServiceBusTransportType TransportType { get { throw null; } set { } }
         public System.Net.IWebProxy WebProxy { get { throw null; } set { } }
@@ -205,6 +207,7 @@ public static partial class ServiceBusModelFactory
         [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.ServiceBusTransportMetrics ServiceBusTransportMetrics(System.DateTimeOffset? lastHeartbeat = default(System.DateTimeOffset?), System.DateTimeOffset? lastConnectionOpen = default(System.DateTimeOffset?), System.DateTimeOffset? lastConnectionClose = default(System.DateTimeOffset?)) { 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)]
@@ -485,6 +488,13 @@ public ServiceBusSessionReceiverOptions() { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public override string ToString() { throw null; }
     }
+    public partial class ServiceBusTransportMetrics
+    {
+        protected internal ServiceBusTransportMetrics() { }
+        public System.DateTimeOffset? LastConnectionClose { get { throw null; } }
+        public System.DateTimeOffset? LastConnectionOpen { get { throw null; } }
+        public System.DateTimeOffset? LastHeartBeat { get { throw null; } }
+    }
     public enum ServiceBusTransportType
     {
         AmqpTcp = 0,

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

@@ -61,6 +61,8 @@ internal class AmqpClient : TransportClient
         ///
         private AmqpConnectionScope ConnectionScope { get; }
 
+        public override ServiceBusTransportMetrics TransportMetrics { get; }
+
         /// <summary>
         ///   Initializes a new instance of the <see cref="AmqpClient"/> class.
         /// </summary>
@@ -94,13 +96,18 @@ internal AmqpClient(
             }.Uri;
 
             Credential = credential;
+            if (options.EnableTransportMetrics)
+            {
+                TransportMetrics = new ServiceBusTransportMetrics();
+            }
             ConnectionScope = new AmqpConnectionScope(
                 ServiceEndpoint,
                 credential,
                 options.TransportType,
                 options.WebProxy,
                 options.EnableCrossEntityTransactions,
-                options.RetryOptions.TryTimeout);
+                options.RetryOptions.TryTimeout,
+                TransportMetrics);
         }
 
         /// <summary>

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

@@ -178,13 +178,15 @@ public override bool IsDisposed
         /// <param name="proxy">The proxy, if any, to use for communication.</param>
         /// <param name="useSingleSession">If true, all links will use a single session.</param>
         /// <param name="operationTimeout">The timeout for operations associated with the connection.</param>
+        /// <param name="metrics">The metrics instance to populate transport metrics. May be null.</param>
         public AmqpConnectionScope(
             Uri serviceEndpoint,
             ServiceBusTokenCredential credential,
             ServiceBusTransportType transport,
             IWebProxy proxy,
             bool useSingleSession,
-            TimeSpan operationTimeout)
+            TimeSpan operationTimeout,
+            ServiceBusTransportMetrics metrics)
         {
             Argument.AssertNotNull(serviceEndpoint, nameof(serviceEndpoint));
             Argument.AssertNotNull(credential, nameof(credential));
@@ -198,7 +200,7 @@ public AmqpConnectionScope(
             TokenProvider = new CbsTokenProvider(new ServiceBusTokenCredential(credential), AuthorizationTokenExpirationBuffer, OperationCancellationSource.Token);
             _useSingleSession = useSingleSession;
 #pragma warning disable CA2214 // Do not call overridable methods in constructors. This internal method is virtual for testing purposes.
-            Task<AmqpConnection> connectionFactory(TimeSpan timeout) => CreateAndOpenConnectionAsync(AmqpVersion, ServiceEndpoint, Transport, Proxy, Id, timeout);
+            Task<AmqpConnection> connectionFactory(TimeSpan timeout) => CreateAndOpenConnectionAsync(AmqpVersion, ServiceEndpoint, Transport, Proxy, Id, timeout, metrics);
 #pragma warning restore CA2214 // Do not call overridable methods in constructors
 
             ActiveConnection = new FaultTolerantAmqpObject<AmqpConnection>(
@@ -436,15 +438,16 @@ public override void Dispose()
         /// <param name="proxy">The proxy, if any, to use for communication.</param>
         /// <param name="scopeIdentifier">The unique identifier for the associated scope.</param>
         /// <param name="timeout">The timeout to consider when creating the connection.</param>
-        ///
+        /// <param name="metrics">The metrics instance to populate transport metrics. May be null.</param>
         /// <returns>An AMQP connection that may be used for communicating with the Service Bus service.</returns>
         protected virtual async Task<AmqpConnection> CreateAndOpenConnectionAsync(
             Version amqpVersion,
             Uri serviceEndpoint,
             ServiceBusTransportType transportType,
             IWebProxy proxy,
             string scopeIdentifier,
-            TimeSpan timeout)
+            TimeSpan timeout,
+            ServiceBusTransportMetrics metrics)
         {
             var hostName = serviceEndpoint.Host;
             AmqpSettings amqpSettings = CreateAmpqSettings(AmqpVersion);
@@ -463,6 +466,11 @@ protected virtual async Task<AmqpConnection> CreateAndOpenConnectionAsync(
             TransportBase transport = await initiator.ConnectTaskAsync(timeout).ConfigureAwait(false);
 
             var connection = new AmqpConnection(transport, amqpSettings, connectionSetings);
+            if (metrics != null)
+            {
+                connection.UsageMeter = new AmqpUsageMeter(metrics);
+            }
+
             await OpenAmqpObjectAsync(connection, timeout.CalculateRemaining(stopWatch.GetElapsedTime()), CancellationToken.None).ConfigureAwait(false);
 
             // Create the CBS link that will be used for authorization.  The act of creating the link will associate

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

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading;
+using Microsoft.Azure.Amqp;
+
+namespace Azure.Messaging.ServiceBus.Amqp
+{
+    internal class AmqpUsageMeter : IAmqpUsageMeter
+    {
+        private readonly ServiceBusTransportMetrics _metrics;
+
+        public AmqpUsageMeter(ServiceBusTransportMetrics metrics)
+        {
+            _metrics = metrics;
+        }
+        public void OnTransportWrite(int bufferSize, int writeSize, long queueSize, long latencyTicks)
+        {
+            // not implemented
+        }
+
+        public void OnTransportRead(int bufferSize, int readSize, int cacheHits, long latencyTicks)
+        {
+            // not implemented
+        }
+
+        public void OnRead(AmqpConnection connection, ulong frameCode, int numberOfBytes)
+        {
+            switch (frameCode)
+            {
+                case 0x00:
+                    _metrics.LastHeartBeat = DateTimeOffset.UtcNow;
+                    break;
+                case 0x10:
+                    _metrics.LastConnectionOpen = DateTimeOffset.UtcNow;
+                    break;
+                case 0x18:
+                    _metrics.LastConnectionClose = DateTimeOffset.UtcNow;
+                    break;
+            }
+        }
+
+        public void OnWrite(AmqpConnection connection, ulong frameCode, int numberOfBytes)
+        {
+            // not implemented
+        }
+    }
+}

sdk/servicebus/Azure.Messaging.ServiceBus/src/Client/ServiceBusClient.cs

@@ -87,17 +87,26 @@ public virtual async ValueTask DisposeAsync()
         }
 
         /// <summary>
-        /// Can be used for mocking.
+        /// Gets the metrics associated with this <see cref="ServiceBusClient"/> instance. The metrics returned represent a snapshot and will not be updated.
+        /// To get updated metrics, this method should be called again.
+        /// In order to use this property, <see cref="ServiceBusClientOptions.EnableTransportMetrics"/> must be set to <value>true</value>.
         /// </summary>
-        protected ServiceBusClient()
-        {
-        }
+        public virtual ServiceBusTransportMetrics GetTransportMetrics()
+            => Connection.InnerClient.TransportMetrics?.Clone() ??
+               throw new InvalidOperationException("Transport metrics are not enabled. To enable transport metrics, set the EnableTransportMetrics property on the ServiceBusClientOptions.");
 
         /// <summary>
         /// The connection that is used for the client.
         /// </summary>
         internal ServiceBusConnection Connection { get; }
 
+        /// <summary>
+        /// Can be used for mocking.
+        /// </summary>
+        protected ServiceBusClient()
+        {
+        }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="ServiceBusClient"/> class.
         /// </summary>

sdk/servicebus/Azure.Messaging.ServiceBus/src/Client/ServiceBusClientOptions.cs

@@ -61,6 +61,12 @@ public ServiceBusRetryOptions RetryOptions
         ///</value>
         public bool EnableCrossEntityTransactions { get; set; }
 
+        /// <summary>
+        /// Gets or sets whether or not to enable metrics for the associated <see cref="ServiceBusClient"/> instance.
+        /// If set to <value>true</value>, <see cref="ServiceBusClient.GetTransportMetrics"/> can be called.
+        /// </summary>
+        public bool EnableTransportMetrics { get; set; }
+
         /// <summary>
         ///   Determines whether the specified <see cref="System.Object" /> is equal to this instance.
         /// </summary>
@@ -103,6 +109,7 @@ internal ServiceBusClientOptions Clone() =>
                 WebProxy = WebProxy,
                 RetryOptions = RetryOptions.Clone(),
                 EnableCrossEntityTransactions = EnableCrossEntityTransactions,
+                EnableTransportMetrics = EnableTransportMetrics
             };
     }
 }

sdk/servicebus/Azure.Messaging.ServiceBus/src/Client/ServiceBusTransportMetrics.cs

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+
+namespace Azure.Messaging.ServiceBus
+{
+    /// <summary>
+    /// A set of metrics that can be used to monitor communication between the client and service.
+    /// </summary>
+    public class ServiceBusTransportMetrics
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ServiceBusTransportMetrics"/> class for mocking.
+        /// </summary>
+        protected internal ServiceBusTransportMetrics()
+        {
+        }
+
+        /// <summary>
+        /// Gets the last time that a heartbeat was received from the Service Bus service. These heartbeats are sent from the
+        /// service approximately every 30 seconds.
+        /// </summary>
+        public DateTimeOffset? LastHeartBeat { get; internal set; }
+
+        /// <summary>
+        /// Gets the last time that a connection was opened for the associated <see cref="ServiceBusClient"/> instance.
+        /// </summary>
+        public DateTimeOffset? LastConnectionOpen { get; internal set; }
+
+        /// <summary>
+        /// Gets the last time that a connection was closed for the associated <see cref="ServiceBusClient"/> instance. If the <see cref="ServiceBusClient"/>
+        /// was disposed, then this time will not be updated again. It may be updated multiple times if the close is initiated by the service.
+        /// </summary>
+        public DateTimeOffset? LastConnectionClose { get; internal set; }
+
+        internal ServiceBusTransportMetrics Clone() =>
+            new()
+            {
+                LastHeartBeat = LastHeartBeat,
+                LastConnectionOpen = LastConnectionOpen,
+                LastConnectionClose = LastConnectionClose
+            };
+    }
+}

sdk/servicebus/Azure.Messaging.ServiceBus/src/Core/TransportClient.cs

@@ -35,6 +35,11 @@ internal abstract class TransportClient : IAsyncDisposable
         ///
         public virtual Uri ServiceEndpoint { get; }
 
+        /// <summary>
+        /// The metrics related to the client.
+        /// </summary>
+        public virtual ServiceBusTransportMetrics TransportMetrics { get; }
+
         /// <summary>
         ///   Creates a sender strongly aligned with the active protocol and transport,
         ///   responsible for sending <see cref="ServiceBusMessage" /> to the entity.