[PM-25520] - Adding a transaction for mutliple repo calls

bitwarden_license/src/Scim/Users/PostUserCommand.cs

@@ -1,5 +1,6 @@
 #nullable enable
 
+using System.Data;
 using Bit.Core;
 using Bit.Core.AdminConsole.Enums;
 using Bit.Core.AdminConsole.OrganizationFeatures.OrganizationUsers.InviteUsers;
@@ -10,6 +11,7 @@
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
 using Bit.Core.Models.Data.Organizations.OrganizationUsers;
+using Bit.Core.Platform.Data;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
 using Bit.Scim.Context;
@@ -27,7 +29,8 @@ public class PostUserCommand(
     IScimContext scimContext,
     IFeatureService featureService,
     IInviteOrganizationUsersCommand inviteOrganizationUsersCommand,
-    TimeProvider timeProvider)
+    TimeProvider timeProvider,
+    ITransactionManager transactionManager)
     : IPostUserCommand
 {
     public async Task<OrganizationUserUserDetails?> PostUserAsync(Guid organizationId, ScimUserRequestModel model)
@@ -45,6 +48,8 @@ public class PostUserCommand(
         Guid organizationId,
         ScimProviderType scimProvider)
     {
+        await using var transactionScope = await transactionManager.BeginTransactionAsync(IsolationLevel.Serializable);
+
         var organization = await organizationRepository.GetByIdAsync(organizationId);
 
         if (organization is null)
@@ -81,6 +86,7 @@ public class PostUserCommand(
             ? await organizationUserRepository.GetDetailsByIdAsync(invitedOrganizationUserId.Value)
             : null;
 
+        await transactionScope.CommitAsync();
         return organizationUser;
     }
 

bitwarden_license/test/Scim.IntegrationTest/Controllers/v2/UsersControllerConcurrencyTests.cs

@@ -0,0 +1,93 @@
+using Bit.Core;
+using Bit.Core.Billing.Enums;
+using Bit.Core.Services;
+using Bit.Infrastructure.EntityFramework.Repositories;
+using Bit.IntegrationTestCommon;
+using Bit.Scim.IntegrationTest.Factories;
+using Bit.Scim.Models;
+using Bit.Scim.Utilities;
+using NSubstitute;
+using Xunit;
+
+namespace Bit.Scim.IntegrationTest.Controllers.v2;
+
+/// <summary>
+/// Verifies seat-count integrity when SCIM invite requests run concurrently.
+/// Requires a real SQL Server (vault_test) — SQLite serializes writes globally and
+/// cannot reproduce the read-modify-write race on Organization.Seats.
+/// </summary>
+public class UsersControllerConcurrencyTests
+{
+    [Fact]
+    public async Task Post_ConcurrentInvites_DoNotOvershootMaxAutoscaleSeats()
+    {
+        const short startingSeats = 3;
+        const int availableSeats = 2;
+        const int concurrentInvites = 6;
+
+        var factory = new ScimApplicationFactory
+        {
+            TestDatabase = new SqlServerTestDatabase()
+        };
+
+        factory.SubstituteService((IFeatureService f) => f.IsEnabled(FeatureFlagKeys.ScimInviteUserOptimization)
+            .Returns(true));
+
+        try
+        {
+            factory.ReinitializeDbForTests(factory.GetDatabaseContext());
+
+            using (var setupScope = factory.Services.CreateScope())
+            {
+                var setupContext = setupScope.ServiceProvider.GetRequiredService<DatabaseContext>();
+                var org = setupContext.Organizations.Single(o => o.Id == ScimApplicationFactory.TestOrganizationId1);
+                org.PlanType = PlanType.EnterpriseAnnually;
+                org.Plan = "Enterprise (Annually)";
+                org.Seats = startingSeats;
+                org.MaxAutoscaleSeats = startingSeats + availableSeats;
+                await setupContext.SaveChangesAsync();
+            }
+
+            var inputs = Enumerable.Range(0, concurrentInvites).Select(BuildInvite).ToArray();
+
+            var responses = await Task.WhenAll(
+                inputs.Select(input =>
+                    factory.UsersPostAsync(ScimApplicationFactory.TestOrganizationId1, input)));
+
+            var successfulInvites = responses.Count(r => r.Response.StatusCode == StatusCodes.Status201Created);
+
+            using var verifyScope = factory.Services.CreateScope();
+            var verifyContext = verifyScope.ServiceProvider.GetRequiredService<DatabaseContext>();
+            var finalOrg = verifyContext.Organizations
+                .Single(o => o.Id == ScimApplicationFactory.TestOrganizationId1);
+            var finalActiveUserCount = verifyContext.OrganizationUsers
+                .Count(ou => ou.OrganizationId == ScimApplicationFactory.TestOrganizationId1 && ou.Status >= 0);
+
+            Assert.All(responses, r => Assert.True(r.Response.StatusCode < 500,
+                $"Expected non-5xx status, got {r.Response.StatusCode}"));
+
+            Assert.Equal(startingSeats + successfulInvites, finalOrg.Seats);
+
+            Assert.Equal(startingSeats + successfulInvites, finalActiveUserCount);
+
+            Assert.True(finalOrg.Seats <= finalOrg.MaxAutoscaleSeats,
+                $"Seats {finalOrg.Seats} exceeded MaxAutoscaleSeats {finalOrg.MaxAutoscaleSeats}");
+        }
+        finally
+        {
+            await factory.DisposeAsync();
+        }
+    }
+
+    private static ScimUserRequestModel BuildInvite(int i) => new()
+    {
+        DisplayName = $"Concurrent User {i}",
+        Emails = new List<BaseScimUserModel.EmailModel>
+        {
+            new() { Primary = true, Type = "work", Value = $"concurrent-{i}@example.com" }
+        },
+        ExternalId = $"CONC-{i}",
+        Active = true,
+        Schemas = new List<string> { ScimConstants.Scim2SchemaUser }
+    };
+}

src/Core/AdminConsole/OrganizationFeatures/OrganizationUsers/InviteUsers/InviteOrganizationUsersCommand.cs

@@ -171,11 +171,13 @@ private async Task<CommandResult<InviteOrganizationUsersResponse>> InviteOrganiz
         {
             logger.LogError(ex, FailedToInviteUsersError.Code);
 
-            await organizationUserRepository.DeleteManyAsync(organizationUserToInviteEntities.Select(x => x.OrganizationUser.Id));
+            // this should already be done
+            //await organizationUserRepository.DeleteManyAsync(organizationUserToInviteEntities.Select(x => x.OrganizationUser.Id));
 
             // Do this first so that SmSeats never exceed PM seats (due to current billing requirements)
             await RevertSecretsManagerChangesAsync(validatedRequest, organization, validatedRequest.Value.InviteOrganization.SmSeats);
 
+            //this should already be done
             await RevertPasswordManagerChangesAsync(validatedRequest, organization);
 
             return new Failure<InviteOrganizationUsersResponse>(
@@ -206,7 +208,7 @@ private async Task RevertPasswordManagerChangesAsync(Valid<InviteOrganizationUse
         {
             organization.Seats = (short?)validatedResult.Value.PasswordManagerSubscriptionUpdate.Seats;
 
-            await organizationRepository.ReplaceAsync(organization);
+            //await organizationRepository.ReplaceAsync(organization);
             await applicationCacheService.UpsertOrganizationAbilityAsync(organization);
         }
     }

src/Core/Platform/Data/ITransactionManager.cs

@@ -0,0 +1,26 @@
+using System.Data;
+
+namespace Bit.Core.Platform.Data;
+
+/// <summary>
+/// Manages ambient database transactions that span multiple repository calls.
+/// Implementations are singleton-safe; transaction state is stored per async flow.
+/// </summary>
+public interface ITransactionManager
+{
+    /// <summary>
+    /// Begins a new ambient transaction. All repository operations on the current
+    /// async flow will use the same connection and transaction until disposed.
+    /// Supports nesting: inner calls increment a reference count; only the
+    /// outermost Dispose/Commit actually affects the database. The isolation level
+    /// on a nested call is ignored — the inner scope joins the outer transaction.
+    /// </summary>
+    Task<ITransactionScope> BeginTransactionAsync(
+        IsolationLevel isolationLevel = IsolationLevel.ReadCommitted,
+        CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Returns true if the current async flow has an active ambient transaction.
+    /// </summary>
+    bool HasActiveTransaction { get; }
+}

src/Core/Platform/Data/ITransactionScope.cs

@@ -0,0 +1,11 @@
+namespace Bit.Core.Platform.Data;
+
+/// <summary>
+/// Represents an ambient transaction scope. Commit must be called explicitly;
+/// disposing without committing triggers rollback.
+/// </summary>
+public interface ITransactionScope : IAsyncDisposable
+{
+    Task CommitAsync(CancellationToken cancellationToken = default);
+    Task RollbackAsync(CancellationToken cancellationToken = default);
+}

src/Core/Platform/Data/NestedTransactionScope.cs

@@ -0,0 +1,37 @@
+namespace Bit.Core.Platform.Data;
+
+public sealed class NestedTransactionScope : ITransactionScope
+{
+    private readonly TransactionHolder _holder;
+    private bool _disposed;
+
+    public NestedTransactionScope(TransactionHolder holder)
+    {
+        _holder = holder;
+    }
+
+    public Task CommitAsync(CancellationToken cancellationToken = default)
+    {
+        // Nested scope commit is a no-op; only the root scope commits.
+        return Task.CompletedTask;
+    }
+
+    public Task RollbackAsync(CancellationToken cancellationToken = default)
+    {
+        // Mark the transaction as doomed so the root scope cannot commit.
+        _holder.Doomed = true;
+        return Task.CompletedTask;
+    }
+
+    public ValueTask DisposeAsync()
+    {
+        if (_disposed)
+        {
+            return ValueTask.CompletedTask;
+        }
+
+        _disposed = true;
+        _holder.ReferenceCount--;
+        return ValueTask.CompletedTask;
+    }
+}

src/Core/Platform/Data/RootTransactionScope.cs

@@ -0,0 +1,42 @@
+namespace Bit.Core.Platform.Data;
+
+public sealed class RootTransactionScope : ITransactionScope
+{
+    private readonly TransactionHolder _holder;
+    private bool _disposed;
+
+    public RootTransactionScope(TransactionHolder holder)
+    {
+        _holder = holder;
+    }
+
+    public async Task CommitAsync(CancellationToken cancellationToken = default)
+    {
+        if (_holder.Doomed)
+        {
+            throw new InvalidOperationException(
+                "Cannot commit a transaction that has been marked for rollback by a nested scope.");
+        }
+
+        _holder.Committed = true;
+        await _holder.Transaction.CommitAsync(cancellationToken);
+    }
+
+    public async Task RollbackAsync(CancellationToken cancellationToken = default)
+    {
+        _holder.Doomed = true;
+        await _holder.Transaction.RollbackAsync(cancellationToken);
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (_disposed)
+        {
+            return;
+        }
+
+        _disposed = true;
+        TransactionState.Current = null;
+        await _holder.DisposeAsync();
+    }
+}

src/Core/Platform/Data/TransactionState.cs

@@ -0,0 +1,67 @@
+using System.Data.Common;
+
+namespace Bit.Core.Platform.Data;
+
+public static class TransactionState
+{
+    private static readonly AsyncLocal<TransactionHolder?> _current = new();
+
+    public static TransactionHolder? Current
+    {
+        get => _current.Value;
+        set => _current.Value = value;
+    }
+}
+
+public sealed class TransactionHolder : IAsyncDisposable
+{
+    public required DbConnection Connection { get; init; }
+    public required DbTransaction Transaction { get; init; }
+    public int ReferenceCount { get; set; } = 1;
+    public bool Committed { get; set; }
+    public bool Doomed { get; set; }
+
+    /// <summary>
+    /// True when this holder is responsible for disposing <see cref="Connection"/>.
+    /// EF reuses the DbContext's connection and must leave its lifetime to the scope;
+    /// Dapper opens its own connection and must dispose it here.
+    /// </summary>
+    public bool OwnsConnection { get; init; } = true;
+
+    /// <summary>
+    /// For EF: the DatabaseContext associated with this transaction.
+    /// </summary>
+    public object? DbContext { get; set; }
+
+    /// <summary>
+    /// For EF: the IServiceScope that owns the DatabaseContext.
+    /// </summary>
+    public IAsyncDisposable? Scope { get; set; }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (!Committed)
+        {
+            try
+            {
+                await Transaction.RollbackAsync();
+            }
+            catch
+            {
+                // Best-effort rollback; connection may already be broken
+            }
+        }
+
+        await Transaction.DisposeAsync();
+
+        if (OwnsConnection)
+        {
+            await Connection.DisposeAsync();
+        }
+
+        if (Scope is not null)
+        {
+            await Scope.DisposeAsync();
+        }
+    }
+}

src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationRepository.cs

@@ -250,15 +250,16 @@ public async Task<ICollection<Organization>> GetManyByIdsAsync(IEnumerable<Guid>
 
     public async Task<OrganizationSeatCounts> GetOccupiedSeatCountByOrganizationIdAsync(Guid organizationId)
     {
-        using (var connection = new SqlConnection(ConnectionString))
+        return await ExecuteWithConnectionAsync(async (connection, transaction) =>
         {
             var result = await connection.QueryAsync<OrganizationSeatCounts>(
                 "[dbo].[Organization_ReadOccupiedSeatCountByOrganizationId]",
                 new { OrganizationId = organizationId },
+                transaction: transaction,
                 commandType: CommandType.StoredProcedure);
 
             return result.SingleOrDefault() ?? new OrganizationSeatCounts();
-        }
+        });
     }
 
     public async Task<IEnumerable<Organization>> GetOrganizationsForSubscriptionSyncAsync()
@@ -285,11 +286,13 @@ await connection.ExecuteAsync("[dbo].[Organization_UpdateSubscriptionStatus]",
 
     public async Task IncrementSeatCountAsync(Guid organizationId, int increaseAmount, DateTime requestDate)
     {
-        await using var connection = new SqlConnection(ConnectionString);
-
-        await connection.ExecuteAsync("[dbo].[Organization_IncrementSeatCount]",
-            new { OrganizationId = organizationId, SeatsToAdd = increaseAmount, RequestDate = requestDate },
-            commandType: CommandType.StoredProcedure);
+        await ExecuteWithConnectionAsync(async (connection, transaction) =>
+        {
+            await connection.ExecuteAsync("[dbo].[Organization_IncrementSeatCount]",
+                new { OrganizationId = organizationId, SeatsToAdd = increaseAmount, RequestDate = requestDate },
+                transaction: transaction,
+                commandType: CommandType.StoredProcedure);
+        });
     }
 
     public async Task InitializeOrganizationAsync(Organization organization, Func<DbConnection, DbTransaction, Task> confirmOwnerAction)