jellyfin-ha-src/tests/Jellyfin.Database.Tests.PostgreSQL/PostgreSqlConcurrencyTests.cs

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Jellyfin.Database.Implementations;
using Jellyfin.Database.Implementations.DbConfiguration;
using Jellyfin.Database.Implementations.Entities;
using Jellyfin.Database.Implementations.Locking;
using Jellyfin.Database.Providers.PostgreSQL;
using Jellyfin.Server.Tests.Migrations;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using Npgsql;
using Xunit;

namespace Jellyfin.Database.Tests.PostgreSQL;

/// <summary>
/// Integration tests that verify concurrent access patterns against a real PostgreSQL server.
/// </summary>
[Xunit.Trait("Category", "RequiresDocker")]
public sealed class PostgreSqlConcurrencyTests : IAsyncLifetime
{
    private static int _databaseSequence;

    private PostgreSqlTestServer? _server;
    private NpgsqlDataSource? _dataSource;
    private PostgreSqlDatabaseProvider? _provider;

    /// <summary>
    /// Attaches to the test server, hands this test a database of its own and applies migrations to it.
    /// </summary>
    /// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
    public async ValueTask InitializeAsync()
    {
        _server = await PostgreSqlTestServer.StartAsync().ConfigureAwait(false);

        var databaseName = FormattableString.Invariant($"pg_concurrency_{Interlocked.Increment(ref _databaseSequence)}");
        var connectionString = await _server.CreateDatabaseAsync(databaseName, TestContext.Current.CancellationToken).ConfigureAwait(false);
        _dataSource = new NpgsqlDataSourceBuilder(connectionString).Build();
        _provider = new PostgreSqlDatabaseProvider(_dataSource);

        var context = CreateContext();
        await using (context.ConfigureAwait(false))
        {
            await context.Database.MigrateAsync(TestContext.Current.CancellationToken).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Releases the data source and the test server.
    /// </summary>
    /// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
    public async ValueTask DisposeAsync()
    {
        // InitializeAsync can fail before the data source exists; its error must not be masked by an NRE here.
        if (_dataSource is not null)
        {
            await _dataSource.DisposeAsync().ConfigureAwait(false);
        }

        if (_server is not null)
        {
            await _server.DisposeAsync().ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Verifies that concurrent inserts on <see cref="ActivityLog"/> from four parallel tasks succeed without deadlock.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task ConcurrentInserts_ActivityLogs_SucceedWithoutDeadlock()
    {
        const int parallelTasks = 4;
        const int insertsPerTask = 10;

        var tasks = new List<Task>(parallelTasks);
        for (var i = 0; i < parallelTasks; i++)
        {
            var taskIndex = i;
            tasks.Add(Task.Run(
                () => InsertBatchAsync(taskIndex, insertsPerTask),
                TestContext.Current.CancellationToken));
        }

        await Task.WhenAll(tasks);

        // Verify all rows were inserted
        var verifyCtx = CreateContext();
        await using (verifyCtx)
        {
            var count = await verifyCtx.ActivityLogs
                .CountAsync(l => l.Type == "ConcurrencyTest", TestContext.Current.CancellationToken);
            Assert.Equal(parallelTasks * insertsPerTask, count);
        }
    }

    private async Task InsertBatchAsync(int taskIndex, int insertsPerTask)
    {
        var ctx = CreateContext();
        await using (ctx.ConfigureAwait(false))
        {
            for (var j = 0; j < insertsPerTask; j++)
            {
                ctx.ActivityLogs.Add(new ActivityLog(
                    $"Task {taskIndex} Insert {j}",
                    "ConcurrencyTest",
                    Guid.Empty));
            }

            await ctx.SaveChangesAsync().ConfigureAwait(false);
        }
    }

    private JellyfinDbContext CreateContext()
    {
        var optionsBuilder = new DbContextOptionsBuilder<JellyfinDbContext>();
        _provider!.Initialise(optionsBuilder, new DatabaseConfigurationOptions { DatabaseType = "PostgreSQL" });
        return new JellyfinDbContext(
            optionsBuilder.Options,
            NullLogger<JellyfinDbContext>.Instance,
            _provider,
            new NoLockBehavior(NullLogger<NoLockBehavior>.Instance));
    }
}