How can I return XML from BeforeSendRequest and AfterReceiveReply to the calling method in a thread-safe way?

Question

We have a console application using the Azure WebJob SDK. The WebJob relies on a WCF service using SOAP, which it accesses through a DLL we wrote that wraps the auto-generated WCF types in something a bit more friendly.

For logging purposes, we want to save the request and response XML bodies for requests that we make. These XML bodies would be saved in our database. But, because the WCF code lives in a low-level DLL, it has no concept of our database and can't save to it.

The DLL uses Microsoft's DI extensions to register types, and the WebJob calls into it like this:

class WebJobClass
{
    IWCFWrapperClient _wcfWrapperClient;

    public WebJobClass(IWCFWrapperClient wcfWrapperClient)
    {
        _wcfWrapperClient = wcfWrapperClient;
    }

    public async Task DoThing()
    {
        var callResult = await _wcfWrapperClient.CallWCFService();
    }
}

IWCFWrapperClient looks like this:

class WCFWrapperClient : IWCFWrapperClient
{
    IWCF _wcf;    // auto-generated by VS, stored in Reference.cs

    public async Task<object> CallWCFService()
    {
        return await _wcf.Call();    // another auto-generated method
    }
}

I've implemented an IClientMessageInspector, and it works fine to get me the XML request/response, but I don't have a way to pass it back up to WCFWrapperClient.CallWCFService so that it can be returned to WebJobClass.DoThing(), who could then save it to the database.

The problem is multithreading. WebJobs, IIRC, will run multiple requests in parallel, calling into the DLL from multiple threads. This means we can't, say, share a static property LastRequestXmlBody since multiple threads could overwrite it. We also can't, say, give each call a Guid or something since there's no way to pass anything from IWCFWrapperClient.CallWCFService into the auto-generated IWCF.Call except what was auto-generated.

So, how can I return XML to WebJobClass.DoThing in a thread-safe way?

Answer 1

I was able to find a solution that uses ConcurrentDictionary<TKey, TValue>, but it's a bit ugly.

First, I amended the auto-generated classes in Reference.cs with a new property Guid InternalCorrelationId. Since the auto-generated classes are partial, this can be done in separate files that aren't changed when the client is regenerated.

public partial class AutoGeneratedWCFType
{
    private Guid InternalCorrelationIdField;

    [System.Runtime.Serialization.DataMember()]
    public Guid InternalCorrelationId
    {
        get { return InternalCorrelationIdField; }
        set { InternalCorrelationIdField = value; }
    }
}

Next, I made all my request DTO types derive from a type named RequestBase, and all my response DTO types derive from a typed named ResponseBase, so I could handle them generically:

public abstract class RequestBase
{
    public Guid InternalCorrelationId { get; set; }
}

public abstract class ResponseBase
{
    public string RequestXml { get; set; }
    public string ResponseXml { get; set; }
}

I then added a type RequestCorrelator that simply holds on to a ConcurrentDictionary<Guid, XmlRequestResponse>:

public sealed class RequestCorrelator : IRequestCorrelator
{
    public ConcurrentDictionary<Guid, XmlRequestResponse> PendingCalls { get; }

    public RequestCorrelator() => PendingCalls = new ConcurrentDictionary<Guid, XmlRequestResponse>();
}

public sealed class XmlRequestResponse
{
    public string RequestXml { get; set; }
    public string ResponseXml { get; set; }
}

RequestCorrelator is its own type for DI purposes - you may just be able to use a ConcurrentDictionary<TKey, TValue> directly.

Finally, we have the code that actually grabs the XML, a type implementing IClientMessageInspector:

public sealed class ClientMessageProvider : IClientMessageInspector
{
    private readonly IRequestCorrelator _requestCorrelator;

    public ClientMessageProvider(IRequestCorrelator requestCorrelator) =>
        _requestCorrelator = requestCorrelator;

    public object BeforeSendRequest(ref Message request, IClientChannel channel)
    {
        var requestXml = request.ToString();
        var internalCorrelationId = GetInternalCorrelationId(requestXml);

        if (internalCorrelationId != null)
        {
            if (_requestCorrelator.PendingCalls.TryGetValue(internalCorrelationId.Value,
                out var requestResponse))
            {
                requestResponse.RequestXml = requestXml;
            }

            request = RemoveInternalCorrelationId(request);
        }
        
        return internalCorrelationId;
    }

    public void AfterReceiveReply(ref Message reply, object correlationState)
    {
        // WCF can internally correlate a request between BeforeSendRequest and
        // AfterReceiveReply. We reuse the same correlation ID we added to the
        // XML as our correlation state.
    
        var responseXml = reply.ToString();
        var internalCorrelationId = (correlationState is Guid guid)
            ? guid
            : default;

        if (_requestCorrelator.PendingCalls.TryGetValue(internalCorrelationId,
            out var requestResponse))
        {
            requestResponse.ResponseXml = responseXml;
        }
    }

    private static Guid? GetInternalCorrelationId(string requestXml)
    {
        var document = XDocument.Parse(requestXml);
        var internalCorrelationIdElement = /* You'll have to write this yourself;
        every WCF XML request is different. */

        return internalCorrelationIdElement != null
            ? Guid.Parse(internalCorrelationIdElement.Value)
            : null;
    }
    
    private static Message RemoveInternalCorrelationId(Message oldMessage)
    {
        // https://stackoverflow.com/a/35639900/2709212
        var buffer = oldMessage.CreateBufferedCopy(2 * 1024 * 1024);
        var tempMessage = buffer.CreateMessage();
        var dictionaryReader = tempMessage.GetReaderAtBodyContents();
        var document = new XmlDocument();
        
        document.Load(dictionaryReader);
        dictionaryReader.Close();

        var internalCorrelationIdNode = /* You'll also have to write this yourself. */
        
        var parent = internalCorrelationIdNode.ParentNode;
        parent.RemoveChild(internalCorrelationIdNode);

        var memoryStream = new MemoryStream();
        var xmlWriter = XmlWriter.Create(memoryStream);
        document.Save(xmlWriter);
        xmlWriter.Flush();
        xmlWriter.Close();

        memoryStream.Position = 0;
        var xmlReader = XmlReader.Create(memoryStream);

        var newMessage = Message.CreateMessage(oldMessage.Version, null, xmlReader);
        newMessage.Headers.CopyHeadersFrom(oldMessage);
        newMessage.Properties.CopyProperties(oldMessage.Properties);

        return newMessage;
    }
}

In short, this type:

1. Finds the correlation ID in the XML request.
2. Finds the XmlRequestResponse with the same correlation ID and adds the request to it.
3. Removes the correlation ID element so that the service doesn't get elements they didn't expect.
4. After receiving a reply, uses correlationState to find the XmlRequestResponse and write the response XML to it.

Now all we have to do is change IWCFWrapperClient:

private async Task<TDtoResult> ExecuteCallWithLogging<TDtoRequest,
    TWcfRequest,
    TWcfResponse,
    TDtoResult>(TDtoRequest request,
    Func<TDtoRequest, TWcfRequest> dtoToWcfConverter,
    Func<TWcfRequest, Task<TWcfResponse>> wcfCall,
    Func<TWcfResponse, TDtoResult> wcfToDtoConverter)
    where TDtoRequest : CorrelationBase
    where TDtoResult : WcfBase
{
    request.InternalCorrelationId = Guid.NewGuid();
    var xmlRequestResponse = new XmlRequestResponse();
    _requestCorrelator.PendingCalls.GetOrAdd(request.InternalCorrelationId,
        xmlRequestResponse);

    var response = await contractingCall(dtoToWcfConverter(request));
    _requestCorrelator.PendingCalls.TryRemove(request.InternalCorrelationId, out _);
    return wcfToDtoConverter(response).WithRequestResponse(xmlRequestResponse);
}

public async Task<DoThingResponseDto> DoThing(DoThingRequestDto request) =>
    await ExecuteCallWithLogging(request,
        r => r.ToWcfModel(),
        async d => await _wcf.Call(d),
        d => d.ToDtoModel());

WithRequestResponse is implemented as follows:

public static T WithRequestResponse<T>(this T item, XmlRequestResponse requestResponse)
    where T : ResponseBase
{
    item.RequestXml = requestResponse?.RequestXml;
    item.ResponseXml = requestResponse?.ResponseXml;

    return item;
}

And there we go. WCF calls that return their XML in the response object rather than just something you can print to console or log to a file.