Wrappers around ASP.NET Core Session

.README.md

Wrappers to ensure distributed session loading is done asynchronously as per advice in MS docs.

Note: It looks like currently it's kind of pointless to try and optimize LoadAsync to only occur where needed (as in my middleware below checking for at least the cookie, or as per suggestions in closed issues to have an attribute that could be sprinkled on controllers/actions) as the standard session middleware always runs CommitAsync (see here), which in turn checks ISession.IsAvailable, which calls Load and ends up doing a sync load... So as far as I can tell, as soon as you've got UseSession() in your pipeline, bar shortcircuited handling for things like static files, you're going to have this CommitAsync run, regardless of what your app code is doing (setting or getting values from Session). Still testing if there's something else that avoids the unnecessary load though.

Usage:

// Services setup
services.Configure<OurDistributedSessionSettings>(
    Config.GetSection("OurDistributedSessionSettings"));
services.AddTransient<DistributedSessionStore>();
services.AddTransient<ISessionStore, OurDistributedSessionStore>();
services.AddSession(options => {});

// Middleware setup
app.UseSession();
app.UseMiddleware<SessionLoaderMiddleware>();

OurDistributedSession.cs

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Session;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace MyCompany.AspNetCore.Session
{
    public class OurDistributedSession
        : ISession
    {
        private readonly ILogger<OurDistributedSession> _logger;
        private readonly OurDistributedSessionSettings _settings;
        private readonly DistributedSession _session;
        private bool _isLoaded = false;

        public OurDistributedSession(
            ILogger<OurDistributedSession> logger,
            IOptionsSnapshot<OurDistributedSessionSettings> settings,
            DistributedSession session)
        {
            _logger = logger;
            _settings = settings.Value;
            _session = session;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)] // hoping this will get us the callerMemberName properly
        private void CheckIfLoaded(
            string key = "unknown",
            [CallerMemberName] string? callerMemberName = null)
        {
            if (_isLoaded)
            {
                return;
            }

            switch (_settings.Mode)
            {
                case OurDistributedSessionMode.AllowSyncLoad:
                    return;
                case OurDistributedSessionMode.AllowSyncLoadButLogWarning:
                    _logger.LogWarning(
                        "Session will be loaded synchronously. Caller: {callerMemberName} Key: {key}",
                        callerMemberName ?? "unknown",
                        key);
                    return;
                case OurDistributedSessionMode.DontAllowSyncLoad:
                    callerMemberName ??= "unknown";
                    throw new Exception(
                        $"Session would've been loaded synchronously. Caller: {callerMemberName} Key: {key}");
            }
        }

        public bool IsAvailable
        {
            get
            {
                CheckIfLoaded("is available get");
                var isAvailable = _session.IsAvailable;
                _isLoaded = true;
                return isAvailable;
            }
        }

        public string Id
        {
            get
            {
                CheckIfLoaded("id get");
                var id = _session.Id;
                _isLoaded = true;
                return id;
            }
        }

        public IEnumerable<string> Keys
        {
            get
            {
                CheckIfLoaded("keys enumeration");
                var keys = _session.Keys;
                _isLoaded = true;
                return keys;
            }
        }

        public void Clear()
        {
            CheckIfLoaded("clear operation");
            _session.Clear();
            _isLoaded = true;
        }

        public async Task CommitAsync(CancellationToken cancellationToken = default)
        {
            CheckIfLoaded("commit operation");
            await _session.CommitAsync(cancellationToken);
            _isLoaded = true;
        }

        public async Task LoadAsync(CancellationToken cancellationToken = default)
        {
            await _session.LoadAsync(cancellationToken);
            _isLoaded = true;
        }
        
        public void Remove(string key)
        {
            CheckIfLoaded(key);
            _session.Remove(key);
            _isLoaded = true;
        }

        public void Set(string key, byte[] value)
        {
            if (value == null)
            {
                throw new ArgumentNullException(nameof(value));
            }

            CheckIfLoaded(key);
            _session.Set(key, value);
            _isLoaded = true;
        }

        public bool TryGetValue(string key, [NotNullWhen(true)] out byte[]? value)
        {
            CheckIfLoaded(key);
            var result = _session.TryGetValue(key, out value);
            _isLoaded = true;
            return result;
        }
    }
}

OurDistributedSessionMode.cs

namespace MyCompany.AspNetCore.Session
{
    public enum OurDistributedSessionMode
    {
        AllowSyncLoad,
        AllowSyncLoadButLogWarning,
        DontAllowSyncLoad
    }
}

OurDistributedSessionSettings.cs

namespace MyCompany.AspNetCore.Session
{
    public class OurDistributedSessionSettings
    {
        public OurDistributedSessionMode Mode { get; } = OurDistributedSessionMode.AllowSyncLoadButLogWarning;
    }
}

OurDistributedSessionStore.cs

using System;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Session;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace MyCompany.AspNetCore.Session
{
    public class OurDistributedSessionStore : ISessionStore
    {
        private readonly IServiceProvider _serviceProvider;
        private readonly DistributedSessionStore _sessionStore;

        public OurDistributedSessionStore(
            IServiceProvider serviceProvider,
            DistributedSessionStore sessionStore)
        {
            _serviceProvider = serviceProvider;
            _sessionStore = sessionStore;
        }

        public ISession Create(string sessionKey, TimeSpan idleTimeout, TimeSpan ioTimeout, Func<bool> tryEstablishSession, bool isNewSessionKey)
        {
            var defaultImpl = _sessionStore.Create(sessionKey, idleTimeout, ioTimeout, tryEstablishSession, isNewSessionKey)
                as DistributedSession;
            if (defaultImpl == null)
            {
                throw new Exception(
                    $"Default implementation of {nameof(ISession)} should've " +
                    $"been a {nameof(DistributedSession)}");
            }

            return new OurDistributedSession(
                _serviceProvider.GetRequiredService<ILogger<OurDistributedSession>>(),
                _serviceProvider.GetRequiredService<IOptionsSnapshot<OurDistributedSessionSettings>>(),
                defaultImpl);
        }
    }
}

SessionLoaderMiddleware.cs

using System;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Options;

namespace MyCompany.AspNetCore.Session
{
    public class SessionLoaderMiddleware
    {
        private readonly SessionOptions _sessionOptions;
        private readonly RequestDelegate _next;

        public SessionLoaderMiddleware(
            IOptions<SessionOptions> sessionOptions,
            RequestDelegate next)
        {
            _sessionOptions = sessionOptions.Value;
            _next = next ?? throw new ArgumentNullException(nameof(next));
        }

        public async Task InvokeAsync(HttpContext httpContext)
        {
            if (httpContext == null)
                throw new ArgumentNullException(nameof(httpContext));

            if (httpContext.Request.Cookies.ContainsKey(_sessionOptions.Cookie.Name!))
            {
                await httpContext.Session.LoadAsync();
            }

            await _next(httpContext);
        }
    }
}