Propagating MongoDB session using AsyncLocal in C#

There are some scenarios in which would not be possible to use or would require extensive code change to propagate the MongoDB’s session among different parts of the system using a simple approach (an example can be found on the article Getting started with .NET Core API, MongoDB, and Transactions). However, for these edge cases scenarios, the AsyncLocal class in .Net provides a viable solution. By persisting instances of MongoDB’s session across threads using AsyncLocal, these instances become readily available for performing database operations in various parts of the codebase.

ShowMeTheCode

In order to demonstrate the usage of AsyncLocal for propagating MongoDB transactions, I have developed a straightforward console application. This application performs two main tasks: inserting a Bank Transaction and a Push Notification to a MongoDB database. Let's take a look at the code:

using MongoDB.Driver;

public class DatabaseContext
{
    public DatabaseContext()
    {
        MongoClient = new MongoClient("mongodb://localhost:27017/Test");
    }

    public IMongoClient MongoClient { get; set; }

    public IMongoCollection<BankTransaction> BankTransactionCollection =>
        MongoClient.GetDatabase("Test").GetCollection<BankTransaction>("BankTransaction");

    public IMongoCollection<PushNotification> PushNotificationCollection =>
        MongoClient.GetDatabase("Test").GetCollection<PushNotification>("PushNotification");
}

The DatabaseContext contains the IMongoClient interface, serving as the connection to the MongoDB database. Within this context, we have access to two collections: BankTransaction and PushNotification. These collections are utilized to store information for our application.

using MongoDB.Driver;

public static class SessionContainer
{
    private static AsyncLocal<IClientSessionHandle> _ambientSession = new();

    public static IClientSessionHandle AmbientSession => _ambientSession?.Value;

    public static void SetSession(IClientSessionHandle session)
    {
        _ambientSession.Value = session;
    }
}

The SessionContainer class is encapsulating the AsyncLocal which will be used to store instances of the IClientSessionHandle (MongoDB’s session).

public class CreateBankTransaction
{
    public CreateBankTransaction(DatabaseContext context)
    {
       _context = context; 
    }

    private readonly DatabaseContext _context;

    public void Handle(BankTransaction bankTransaction)
    {
        Console.WriteLine($"Insert bank transaction for operation {bankTransaction.Id} with MongoDB session {SessionContainer.AmbientSession.ServerSession.Id}");
        _context.BankTransactionCollection.InsertOne(SessionContainer.AmbientSession, bankTransaction);
    }
}

public class CreatePushNotification
{
    public CreatePushNotification(DatabaseContext context)
    {
        _context = context;
    }

    private readonly DatabaseContext _context;

    public void Handle(int operation, PushNotification pushNotification)
    {
        Console.WriteLine($"Insert push notification for operation {operation} with MongoDB session {SessionContainer.AmbientSession.ServerSession.Id}");
        _context.PushNotificationCollection.InsertOne(SessionContainer.AmbientSession, pushNotification);
    }
}

The classes CreateBankTransaction and CreatePushNotification are simulating a codebase in which database’s operations have been done in many places. It is important to notice that both methods are getting the MongoDB session from SessionContainer class.

var databaseContext = new DatabaseContext();
var createBankTransaction = new CreateBankTransaction(databaseContext);
var createPushNotification = new CreatePushNotification(databaseContext);

var operations = new List<BankTransaction>
{
    new BankTransaction(4 , "1000", "2000", 304),
    new BankTransaction(5, "3000", "4000", 500),
    new BankTransaction(6, "4000", "2000", 200)
};


Parallel.ForEach(operations, operation =>
{
    var session = databaseContext.MongoClient.StartSession();
    session.StartTransaction();
    SessionContainer.SetSession(session);
    Console.WriteLine($"Start operation {operation.Id} with mongo session {session.ServerSession.Id}");
    createBankTransaction.Handle(operation);
    createPushNotification.Handle(operation.Id, new PushNotification($"Transfer from account {operation.From} to {operation.To} with amount {operation.Amount}"));
    Console.WriteLine($"Commit operation {operation.Id} with mongo session {session.ServerSession.Id}");
    SessionContainer.AmbientSession.CommitTransaction();
});

Console.ReadKey();

public record PushNotification(string Message);
public record BankTransaction(int Id, string From, string To, decimal Amount);

In this example we are simulating a scenario where the classes are injected as singleton, as evident from lines 1, 2, and 3. Otherwise, if the classes were injected as scoped, the IClientSessionHandle instance could have been stored within the DatabaseContext, instead of in the AsyncLocal, because each operation would have a different instance of the DatabaseContext.

The code triggers three operations in parallel, for each operation it starts a new MongoDB session, creates a transaction for the session and saves the session in the AsyncLocal (through SessionContainer class), making it accessible within the current thread. Subsequently, the application inserts BankTransacation and PushNotification into our database. Finally, it commits the created transaction. Below are the logs generated when the application runs:

Upon reviewing the logs, it becomes apparent that a distinct instance of MongoDB's session was created for each operation. Both classes, CreateBankTransaction and CreatePushNotification, effectively utilized the available IClientSessionHandle instance within the thread to perform their respective database operations. As a result, no conflicts arose, since each operation operated within its isolated thread.

Conclusion

Although for most projects I would not recommend this approach, it can prove useful in certain existing codebases where implementing MongoDB transactions might be challenging. For example, if implementing MongoDB transaction would necessitate modifying methods signatures across multiple classes to pass an instance of IClientSessionHandle as an input argument.

I hope you found this article useful! If you have any doubts or questions, feel free to comment below.