Fix address mappings for when more than two transports are being used (#526)

Author: abparticular

src/NServiceBus.MessagingBridge/AddressMap.cs

@@ -0,0 +1,61 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using NServiceBus;
+using NServiceBus.Raw;
+using NServiceBus.Transport;
+
+/// <summary>
+/// How to map an address depends upon which transport the message is being sent to, 
+/// so it is necessary to have an address mapping dictionary for each transport
+/// </summary>
+class AddressMap(IReadOnlyDictionary<string, IStartableRawEndpoint> dispatchers) : IAddressMap
+{
+    readonly Dictionary<string, TransportAddressMap> addressMap =
+        dispatchers.ToDictionary(dispatcher => dispatcher.Key, _ => new TransportAddressMap());
+
+    readonly HashSet<string> endpoints = [];
+
+    public void Add(BridgeTransport transport, BridgeEndpoint endpoint)
+    {
+        if (!endpoints.Add(endpoint.Name))
+        {
+            throw new InvalidOperationException($"{endpoint.Name} has already been added to the address map.");
+        }
+
+        var queueAddress = new QueueAddress(endpoint.Name);
+
+        foreach (var targetTransport in addressMap.Keys)
+        {
+            var targetAddress = DetermineTransportAddress(targetTransport, transport.Name, endpoint.QueueAddress, queueAddress);
+
+            foreach (var sourceTransport in addressMap.Keys)
+            {
+                var sourceAddress = DetermineTransportAddress(sourceTransport, transport.Name, endpoint.QueueAddress, queueAddress);
+
+                addressMap[targetTransport][sourceAddress] = targetAddress;
+            }
+        }
+    }
+
+    string DetermineTransportAddress(string transportToMap, string endpointTransport, string overriddenQueueAddress, QueueAddress queueAddress) => transportToMap == endpointTransport
+        ? overriddenQueueAddress ?? dispatchers[transportToMap].ToTransportAddress(queueAddress)
+        : dispatchers[transportToMap].ToTransportAddress(queueAddress);
+
+    public bool TryTranslate(string targetTransport, string address, out string bestMatch)
+    {
+        var transportAddressMappings = addressMap[targetTransport];
+
+        if (transportAddressMappings.TryGetValue(address, out bestMatch))
+        {
+            return true;
+        }
+
+        bestMatch = address.GetClosestMatch(transportAddressMappings.Keys);
+        return false;
+    }
+}
+
+class TransportAddressMap : Dictionary<string, string>
+{
+}

src/NServiceBus.MessagingBridge/EndpointRegistry.cs

@@ -15,59 +15,77 @@ public async Task<IEnumerable<ProxyRegistration>> Initialize(IReadOnlyCollection
     {
         // Assume that it is the number of endpoints that is more likely to scale up in size than the number of transports (typically only 2).
         // Therefore in cases where it might matter, it is more efficient to iterate over the transports multiple times.
-        transportConfigurationMappings = transportConfigurations.ToDictionary(t => t.Name, t => t);
 
         IList<ProxyRegistration> proxyRegistrations = [];
 
-        foreach (var targetTransport in transportConfigurations)
+        await CreateAndRegisterDispatchers(transportConfigurations, proxyRegistrations, cancellationToken)
+            .ConfigureAwait(false);
+
+        var allEndpoints = transportConfigurations.SelectMany(
+            transport => transport.Endpoints.Select(endpoint => (transport, endpoint)));
+
+        foreach (var (transport, endpoint) in allEndpoints)
         {
-            // Create the dispatcher for this transport
-            var dispatchEndpoint = await EndpointProxyFactory.CreateDispatcher(targetTransport, cancellationToken).ConfigureAwait(false);
+            AddressMap.Add(transport, endpoint);
+
+            var dispatcher = dispatchers[transport.Name];
+            var queueAddress = new QueueAddress(endpoint.Name);
+            var targetTransportAddress = dispatcher.ToTransportAddress(queueAddress);
+            endpointAddressMappings[endpoint.Name] = endpoint.QueueAddress ?? targetTransportAddress;
+            targetEndpointDispatchers[endpoint.Name] = new TargetEndpointDispatcher(transport.Name, dispatcher, targetTransportAddress);
+
+            await CreateAndRegisterProxies(transport, endpoint, transportConfigurations, proxyRegistrations, cancellationToken)
+                .ConfigureAwait(false);
+        }
+
+        return proxyRegistrations;
+    }
 
+    async Task CreateAndRegisterDispatchers(
+        IReadOnlyCollection<BridgeTransport> transportConfigurations,
+        IList<ProxyRegistration> proxyRegistrations,
+        CancellationToken cancellationToken)
+    {
+        foreach (var transport in transportConfigurations)
+        {
+            var dispatcher = await EndpointProxyFactory.CreateDispatcher(transport, cancellationToken).ConfigureAwait(false);
+            dispatchers.Add(transport.Name, dispatcher);
             proxyRegistrations.Add(new ProxyRegistration
             {
                 Endpoint = null,
-                TranportName = targetTransport.Name,
-                RawEndpoint = dispatchEndpoint
+                RawEndpoint = dispatcher
             });
+        }
+
+        AddressMap = new AddressMap(dispatchers);
+    }
 
-            // create required proxy endpoints on all transports
-            foreach (var endpointToSimulate in targetTransport.Endpoints)
+    async Task CreateAndRegisterProxies(
+        BridgeTransport targetTransport,
+        BridgeEndpoint targetEndpoint,
+        IReadOnlyCollection<BridgeTransport> transportConfigurations,
+        IList<ProxyRegistration> proxyRegistrations,
+        CancellationToken cancellationToken)
+    {
+        // Endpoint will need to be proxied on the other transports
+        foreach (var proxyTransport in transportConfigurations)
+        {
+            if (proxyTransport.Name == targetTransport.Name)
             {
-                // Endpoint will need to be proxied on the other transports
-                foreach (var proxyTransport in transportConfigurationMappings.Where(kvp => kvp.Key != targetTransport.Name).Select(kvp => kvp.Value))
-                {
-                    var startableEndpointProxy = await endpointProxyFactory.CreateProxy(endpointToSimulate, proxyTransport, cancellationToken)
-                        .ConfigureAwait(false);
-
-                    logger.LogInformation("Proxy for endpoint {endpoint} created on {transport}", endpointToSimulate.Name, proxyTransport.Name);
-
-                    var queueAddress = new QueueAddress(endpointToSimulate.Name);
-                    var targetTransportAddress = dispatchEndpoint.ToTransportAddress(queueAddress);
-                    var sourceTransportAddress = startableEndpointProxy.ToTransportAddress(queueAddress);
-
-                    endpointAddressMappings[endpointToSimulate.Name] = endpointToSimulate.QueueAddress ?? targetTransportAddress;
-                    targetEndpointAddressMappings[targetTransportAddress] = sourceTransportAddress;
-                    if (targetTransportAddress != sourceTransportAddress)
-                    {
-                        // Also add the reverse mapping so that any messages that were in-flight before a bridge configuration change
-                        // can still have their address translated correctly.  This also allows for the case where duplicate logical endpoints
-                        // are running as a competing consumer with the bridge in canary/parallel deployment scenarios.
-                        targetEndpointAddressMappings[sourceTransportAddress] = targetTransportAddress;
-                    }
-                    targetEndpointDispatchers[endpointToSimulate.Name] = new TargetEndpointDispatcher(targetTransport.Name, dispatchEndpoint, targetTransportAddress);
-
-                    proxyRegistrations.Add(new ProxyRegistration
-                    {
-                        Endpoint = endpointToSimulate,
-                        TranportName = proxyTransport.Name,
-                        RawEndpoint = startableEndpointProxy
-                    });
-                }
+                continue;
             }
-        }
 
-        return proxyRegistrations;
+            var startableEndpointProxy = await endpointProxyFactory.CreateProxy(targetEndpoint, proxyTransport, cancellationToken)
+                .ConfigureAwait(false);
+
+            logger.LogInformation("Proxy for endpoint {endpoint} created on {transport}", targetTransport.Name, proxyTransport.Name);
+
+            proxyRegistrations.Add(new ProxyRegistration
+            {
+                Endpoint = targetEndpoint,
+                RawEndpoint = startableEndpointProxy
+            });
+        }
     }
 
     public TargetEndpointDispatcher GetTargetEndpointDispatcher(string sourceEndpointName)
@@ -77,21 +95,12 @@ public TargetEndpointDispatcher GetTargetEndpointDispatcher(string sourceEndpoin
             return endpointDispatcher;
         }
 
-        var nearestMatch = GetClosestMatchForExceptionMessage(sourceEndpointName, targetEndpointDispatchers.Keys);
+        var nearestMatch = sourceEndpointName.GetClosestMatch(targetEndpointDispatchers.Keys);
 
         throw new Exception($"No target endpoint dispatcher could be found for endpoint: {sourceEndpointName}. Ensure names have correct casing as mappings are case-sensitive. Nearest configured match: {nearestMatch}");
     }
 
-    public bool TryTranslateToTargetAddress(string sourceAddress, out string bestMatch)
-    {
-        if (targetEndpointAddressMappings.TryGetValue(sourceAddress, out bestMatch))
-        {
-            return true;
-        }
-
-        bestMatch = GetClosestMatchForExceptionMessage(sourceAddress, targetEndpointAddressMappings.Keys);
-        return false;
-    }
+    public IAddressMap AddressMap { get; private set; }
 
     public string GetEndpointAddress(string endpointName)
     {
@@ -100,30 +109,18 @@ public string GetEndpointAddress(string endpointName)
             return address;
         }
 
-        var nearestMatch = GetClosestMatchForExceptionMessage(endpointName, endpointAddressMappings.Keys);
+        var nearestMatch = endpointName.GetClosestMatch(endpointAddressMappings.Keys) ?? "(No mappings registered)";
 
         throw new Exception($"No address mapping could be found for endpoint: {endpointName}. Ensure names have correct casing as mappings are case-sensitive. Nearest configured match: {nearestMatch}");
     }
 
-    static string GetClosestMatchForExceptionMessage(string sourceEndpointName, IEnumerable<string> items)
-    {
-        var calculator = new Levenshtein(sourceEndpointName.ToLower());
-        var nearestMatch = items
-            .OrderBy(x => calculator.DistanceFrom(x.ToLower()))
-            .FirstOrDefault();
-        return nearestMatch ?? "(No mappings registered)";
-    }
-
-    Dictionary<string, BridgeTransport> transportConfigurationMappings = [];
-
+    readonly Dictionary<string, IStartableRawEndpoint> dispatchers = [];
     readonly Dictionary<string, TargetEndpointDispatcher> targetEndpointDispatchers = [];
-    readonly Dictionary<string, string> targetEndpointAddressMappings = [];
     readonly Dictionary<string, string> endpointAddressMappings = [];
 
     public class ProxyRegistration
     {
         public BridgeEndpoint Endpoint;
-        public string TranportName;
         public IStartableRawEndpoint RawEndpoint;
     }
 }

src/NServiceBus.MessagingBridge/IAddressMap.cs

@@ -0,0 +1,8 @@
+using NServiceBus;
+
+interface IAddressMap
+{
+    void Add(BridgeTransport transport, BridgeEndpoint endpoint);
+
+    bool TryTranslate(string targetTransport, string address, out string bestMatch);
+}

src/NServiceBus.MessagingBridge/IEndpointRegistry.cs

@@ -2,7 +2,7 @@
 {
     TargetEndpointDispatcher GetTargetEndpointDispatcher(string sourceEndpointName);
 
-    bool TryTranslateToTargetAddress(string sourceAddress, out string bestMatch);
+    IAddressMap AddressMap { get; }
 
     string GetEndpointAddress(string endpointName);
 }

src/NServiceBus.MessagingBridge/MessageShovel.cs

@@ -6,18 +6,11 @@
 using NServiceBus.Faults;
 using NServiceBus.Transport;
 
-sealed class MessageShovel : IMessageShovel
+sealed class MessageShovel(
+    ILogger<MessageShovel> logger,
+    IEndpointRegistry targetEndpointRegistry,
+    FinalizedBridgeConfiguration finalizedBridgeConfiguration) : IMessageShovel
 {
-    public MessageShovel(
-        ILogger<MessageShovel> logger,
-        IEndpointRegistry targetEndpointRegistry,
-        FinalizedBridgeConfiguration finalizedBridgeConfiguration)
-    {
-        this.logger = logger;
-        this.targetEndpointRegistry = targetEndpointRegistry;
-        translateReplyToAddressForFailedMessages = finalizedBridgeConfiguration.TranslateReplyToAddressForFailedMessages;
-    }
-
     public async Task TransferMessage(TransferContext transferContext, CancellationToken cancellationToken = default)
     {
         TargetEndpointDispatcher targetEndpointDispatcher = null;
@@ -26,13 +19,14 @@ public async Task TransferMessage(TransferContext transferContext, CancellationT
             targetEndpointDispatcher = targetEndpointRegistry.GetTargetEndpointDispatcher(transferContext.SourceEndpointName);
 
             var messageContext = transferContext.MessageToTransfer;
+            var targetTransport = targetEndpointDispatcher.TransportName;
 
             var messageToSend = new OutgoingMessage(messageContext.NativeMessageId, messageContext.Headers, messageContext.Body);
             messageToSend.Headers.Remove(BridgeHeaders.FailedQ);
 
-            var length = transferContext.SourceTransport.Length + targetEndpointDispatcher.TransportName.Length + 2 /* ->*/;
+            var length = transferContext.SourceTransport.Length + targetTransport.Length + 2 /* ->*/;
             var transferDetails = string.Create(length,
-                (Source: transferContext.SourceTransport, Target: targetEndpointDispatcher.TransportName),
+                (Source: transferContext.SourceTransport, Target: targetTransport),
                 static (chars, context) =>
                 {
                     var position = 0;
@@ -48,12 +42,12 @@ public async Task TransferMessage(TransferContext transferContext, CancellationT
                 //This is a failed message forwarded to ServiceControl. We need to transform the FailedQ header so that ServiceControl returns the message
                 //to the correct queue/transport on the other side
 
-                TransformAddressHeader(messageToSend, targetEndpointRegistry, FaultsHeaderKeys.FailedQ);
+                TransformAddressHeader(messageToSend, targetTransport, FaultsHeaderKeys.FailedQ);
 
                 if (translateReplyToAddressForFailedMessages)
                 {
                     //Try to translate the ReplyToAddress, this is needed when an endpoint is migrated to the ServiceControl side before this message is retried
-                    TransformAddressHeader(messageToSend, targetEndpointRegistry, Headers.ReplyToAddress, throwOnError: false);
+                    TransformAddressHeader(messageToSend, targetTransport, Headers.ReplyToAddress, throwOnError: false);
                 }
             }
             else if (IsAuditMessage(messageToSend))
@@ -64,27 +58,27 @@ public async Task TransferMessage(TransferContext transferContext, CancellationT
             else if (IsRetryMessage(messageToSend))
             {
                 //Transform the retry ack queue address
-                TransformAddressHeader(messageToSend, targetEndpointRegistry, "ServiceControl.Retry.AcknowledgementQueue");
+                TransformAddressHeader(messageToSend, targetTransport, "ServiceControl.Retry.AcknowledgementQueue");
 
                 if (translateReplyToAddressForFailedMessages)
                 {
                     //This is a message retried from ServiceControl. We try to translate its ReplyToAddress.
-                    TransformAddressHeader(messageToSend, targetEndpointRegistry, Headers.ReplyToAddress, throwOnError: false);
+                    TransformAddressHeader(messageToSend, targetTransport, Headers.ReplyToAddress, throwOnError: false);
                 }
             }
             else if (IsRetryEditedMessage(messageToSend))
             {
                 // This is a retry message that has been edited going from one side of the bridge to another
                 // The ReplyToAddress is transformed to allow for replies to be delivered
                 messageToSend.Headers[BridgeHeaders.Transfer] = transferDetails;
-                TransformAddressHeader(messageToSend, targetEndpointRegistry, Headers.ReplyToAddress, !translateReplyToAddressForFailedMessages);
+                TransformAddressHeader(messageToSend, targetTransport, Headers.ReplyToAddress, !translateReplyToAddressForFailedMessages);
             }
             else
             {
                 // This is a regular message sent between the endpoints on different sides of the bridge.
                 // The ReplyToAddress is transformed to allow for replies to be delivered
                 messageToSend.Headers[BridgeHeaders.Transfer] = transferDetails;
-                TransformRegularMessageReplyToAddress(transferContext, messageToSend, targetEndpointRegistry);
+                TransformRegularMessageReplyToAddress(transferContext, messageToSend, targetTransport);
             }
 
             await targetEndpointDispatcher.Dispatch(
@@ -118,7 +112,7 @@ await targetEndpointDispatcher.Dispatch(
     void TransformRegularMessageReplyToAddress(
         TransferContext transferContext,
         OutgoingMessage messageToSend,
-        IEndpointRegistry targetEndpointRegistry)
+        string targetTransport)
     {
         if (!messageToSend.Headers.TryGetValue(Headers.ReplyToAddress, out var headerValue))
         {
@@ -133,36 +127,34 @@ void TransformRegularMessageReplyToAddress(
         }
         else
         {
-            TransformAddressHeader(messageToSend, targetEndpointRegistry, Headers.ReplyToAddress);
+            TransformAddressHeader(messageToSend, targetTransport, Headers.ReplyToAddress);
         }
     }
 
     void TransformAddressHeader(
         OutgoingMessage messageToSend,
-        IEndpointRegistry endpointRegistry,
+        string targetTransport,
         string addressHeaderKey,
         bool throwOnError = true)
     {
-        if (!messageToSend.Headers.TryGetValue(addressHeaderKey, out var sourceAddress))
+        if (!messageToSend.Headers.TryGetValue(addressHeaderKey, out var address))
         {
             return;
         }
 
-        if (endpointRegistry.TryTranslateToTargetAddress(sourceAddress, out string bestMatch))
+        if (targetEndpointRegistry.AddressMap.TryTranslate(targetTransport, address, out string bestMatch))
         {
             messageToSend.Headers[addressHeaderKey] = bestMatch;
         }
         else if (throwOnError == false)
         {
-            logger.LogWarning($"Could not translate {sourceAddress} address. Consider using `.HasEndpoint()` method to add missing endpoint declaration.");
+            logger.LogWarning($"Could not translate {address} address. Consider using `.HasEndpoint()` method to add missing endpoint declaration.");
         }
         else
         {
-            throw new Exception($"No target address mapping could be found for source address: {sourceAddress}. Ensure names have correct casing as mappings are case-sensitive. Nearest configured match: {bestMatch}");
+            throw new Exception($"No target address mapping could be found for source address: {address}. Ensure names have correct casing as mappings are case-sensitive. Nearest configured match: {bestMatch ?? "(No mappings registered)"}");
         }
     }
 
-    readonly ILogger<MessageShovel> logger;
-    readonly IEndpointRegistry targetEndpointRegistry;
-    readonly bool translateReplyToAddressForFailedMessages;
+    readonly bool translateReplyToAddressForFailedMessages = finalizedBridgeConfiguration.TranslateReplyToAddressForFailedMessages;
 }