Discovering AsyncLocal

My team was the first LoopNet team to deploy a project to AWS. As such we, had quite the learning curve working with Lambdas, DynamoDB, and Step Functions. There was also a new micro-service framework that was developed by another team at CoStar that also added to our learning curve, but since it had already been hosted in AWS Lambda before, it ultimately cut down on the challenge of deploying our APIs.

One of the features of this new framework was a Correlation ID, which could be sent to an API to tag all log statements for a given request. The hosts in the new framework were also configured to propagate the Correlation ID to other hosts, so if one API called another, then the Correlation ID would help stitch together all of the logging. Unfortunately, we had a Step Function in between two of our APIs that orchestrated some important business logic.

That step function needed the Correlation ID, and needed to pass on the Correlation ID to the APIs that it used. The developer who worked on this component leveraged a tool to generate client logic for each API that the Step Function consumed. The generated client logic also provided hooks so we could inject data into our requests. He ended up writing something like the following:

public async Task<Request> FnStep(Request request, ILambdaContext ctx)
{
   CorrelationHook h = new CorrelationHook(request.CorrelationId);
   IApiClient client = new ApiClient(new List<IRequestHooks> { h });

   var apiReq = new ApiRequest { /* properties for ApiRequest */ };
   var apiResp = await client.MakeApiCallAsync(apiReq).ConfigureAwait(false);

   // Use the API response.

   return request;
}

The code above solved our main challenge: we have passed in the Correlation ID by having it included in the request object for the Step Function, and we have created CorrelationHook to inject the Correlation ID on outbound HTTP calls. Unfortunately, this approach has some drawbacks due to a lack of Dependency Inversion. Since we’re not relying on an Inversion of Control container to allocate classes that should ostensibly be singletons (ApiClient), we’re incurring extra costs for the allocation of these classes on every single request made to that particular Lambda. I have had instances where I was using a library that I wrapped in an Adapter, and allocated per request. When the underlying library changed, we incurred a larger compute cost for that single Lambda by an order of magnitude. I fixed that cost problem by making sure that the Adapter for the library was allocated as a Singleton. Ideally singletons should be allocated once during a Lambda’s life cycle.

Another drawback to this approach is that we can’t write a unit test for this method. Since the method allocates an instance of IApiClient we have no means by which we can allocate and inject a mock of the API to isolate FnStep as a testable unit. If this was written with an Inversion of Control container, then we would have the ability to use constructor injection to supply FnStep with a Mock<IApiClient>, which then allows us to use SetUp and Verify to ensure that our code is doing what it ought to.

Finally, this approach also just creates a lot of clutter and boiler plate for the code. Let’s say that we have a class that encapsulates 5 function steps that each rely on 3 different API clients. This approach would necessitate that we have to allocate the API clients 15 times. That’s a lot of repetition, and a lot of extra noise in the code. If the class encapsulating these steps either had constructor injection, or leveraged a service provider during construction, then we could have the API clients as private fields. We would only have to reference the API clients to use them.

That brings us to the interesting challenge that we faced cleaning up this code: how can we have the Correlation ID from the request object injected into the ApiClient via the CorrelationHook if we’re not allocating these classes?

While working on another project, I encountered a class HttpContextAccessor and it seemed like it allowed us to access the HttpContext from Singletons. When I returned to the Step Functions a while later, I thought to take a look at the implementation of the HttpContextAccessor and found out that it was using a generic type called AsyncLocal<T>. With this type, you can assign a value that is local to the asynchronous context. This was the missing piece of the puzzle for passing the Correlation ID to the APIs while leveraging Dependency Inversion.

I created a class that allowed access to the Correlation ID:

public class CorrelationIdAccessor : ICorrelationIdAccessor
{
   private static AsyncLocal<string> correlationId = new AsyncLocal<string>();

   public string CorrelationId
   {
      get
      {
         return correlationId;
      }

      set
      {
         correlationId = value;
      }
   }
}

Then we could refactor the CorrelationHook to use the new ICorrelationIdAccessor:

public class CorrelationHook : IRequestHook
{
   private ICorrelationIdAccessor correlationIdAccessor;

   public CorrelationHook(ICorrelationIdAccessor correlationIdAccessor)
   {
      this.correlationIdAccessor = correlationIdAccessor;
   }

   public void UpdateRequest(HttpRequest request)
   {
      var correlationId = correaltionIdAccessor.CorrelationId;

      if (!string.IsNullOrEmpty(correlationId))
      {
         request.Headers.Add(Constants.CorrelationHeaderKey, correlationId);
      }
   }
}

Now that the CorrelationHook has been refactored to use Dependency Injection, we can leverage the Inversion of Control container, and refactor the example FnStep:

public async Task<Request> FnStep(Request request, ILambdaContext ctx)
{
   correlationIdAccessor.CorrelationId = request.CorrelationId;

   var apiReq = new ApiRequest { /* properties for ApiRequest */ };
   var apiResp = await client.MakeApiCallAsync(apiReq).ConfigureAwait(false);

   // Use the API response.

   return request;
}

Now we can have the Correlation ID propagate to all the APIs we rely on, and we can have all of the code dependencies as Singletons. This is much cleaner and easier to test.