Skip to content

Commit 7422e9d

Browse files
AnnaSasDevAnnaSasDev
committed
Fix: UnitOfWork
1 parent a4ddbc9 commit 7422e9d

File tree

8 files changed

+85
-47
lines changed

8 files changed

+85
-47
lines changed

‎src/CodeOfChaos.Types.UnitOfWork/ReadonlyUnitOfWork.cs‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,12 @@ public class ReadonlyUnitOfWork<TDbContext>(
1414
IServiceScope serviceScope
1515
) : UnitOfWork<TDbContext>(dbContextFactory, serviceScope), IReadonlyUnitOfWork
1616
where TDbContext : DbContext, IReadonlyCapableDbContext {
17-
18-
protected override AsyncLazy<TDbContext> LazyDb { get; } = new(async ct => {
19-
TDbContext dbContext = await dbContextFactory.CreateDbContextAsync(ct);
17+
18+
protected async override ValueTask<TDbContext> GetDbContextAsync(CancellationToken ct) {
19+
TDbContext dbContext = await base.GetDbContextAsync(ct);
2020
dbContext.SetAsReadonly();
2121
return dbContext;
22-
});
23-
22+
}
2423

2524
// -----------------------------------------------------------------------------------------------------------------
2625
// Methods

‎src/CodeOfChaos.Types.UnitOfWork/UnitOfWork.cs‎

Lines changed: 74 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -11,43 +11,71 @@ namespace CodeOfChaos.Types.UnitOfWork;
1111
// Code
1212
// ---------------------------------------------------------------------------------------------------------------------
1313
public class UnitOfWork<TDbContext>(IDbContextFactory<TDbContext> dbContextFactory, IServiceScope serviceScope) : IUnitOfWork where TDbContext : DbContext {
14-
protected virtual AsyncLazy<TDbContext> LazyDb { get; } = new(async ct => await dbContextFactory.CreateDbContextAsync(ct));
14+
private TDbContext? _dbContext;
1515
private IDbContextTransaction? _transaction;
16-
private ConcurrentDictionary<Type, IUnitOfWorkRepository> AttachedRepositories { get; } = [];
16+
private readonly ConcurrentDictionary<Type, IUnitOfWorkRepository> AttachedRepositories = [];
17+
private readonly SemaphoreSlim _initLock = new(1, 1);
18+
1719

1820
// -----------------------------------------------------------------------------------------------------------------
1921
// Methods
2022
// -----------------------------------------------------------------------------------------------------------------
21-
public virtual async ValueTask SaveChangesAsync(CancellationToken ct = default) {
22-
DbContext dbContext = await LazyDb.GetValueAsync(ct);
23-
await dbContext.SaveChangesAsync(ct);
24-
}
25-
26-
public virtual async ValueTask<bool> TryCommitTransactionAsync(CancellationToken ct = default) {
27-
if (_transaction == null) return false;
23+
protected virtual async ValueTask<TDbContext> GetDbContextAsync(CancellationToken ct) {
24+
if (_dbContext != null) return _dbContext;
2825

29-
await _transaction.CommitAsync(ct);
30-
_transaction.Dispose();
31-
_transaction = null;
26+
await _initLock.WaitAsync(ct);
27+
try {
28+
// Double-check pattern
29+
if (_dbContext != null) return _dbContext;
3230

33-
return true;
31+
_dbContext = await dbContextFactory.CreateDbContextAsync(ct);
32+
return _dbContext;
33+
}
34+
finally {
35+
_initLock.Release();
36+
}
3437
}
35-
38+
3639
public virtual async ValueTask<bool> TryCreateTransactionAsync(CancellationToken ct = default) {
3740
if (_transaction != null) return false;
3841

39-
TDbContext dbContext = await LazyDb.GetValueAsync(ct);
42+
TDbContext dbContext = await GetDbContextAsync(ct);
4043
if (dbContext.Database.CurrentTransaction != null) {
41-
// Something went wrong during saving before and the transaction wasn't set by the unit of work
4244
_transaction = dbContext.Database.CurrentTransaction;
4345
return true;
4446
}
4547

46-
_transaction = await dbContext.Database.BeginTransactionAsync(ct);
48+
await _initLock.WaitAsync(ct);
49+
try {
50+
_transaction = await dbContext.Database.BeginTransactionAsync(ct);
51+
return true;
52+
}
53+
finally {
54+
_initLock.Release();
55+
}
56+
}
57+
58+
public virtual async ValueTask SaveChangesAsync(CancellationToken ct = default) {
59+
DbContext dbContext = await GetDbContextAsync(ct);
60+
await dbContext.SaveChangesAsync(ct);
61+
}
62+
63+
public virtual async ValueTask<bool> TryCommitTransactionAsync(CancellationToken ct = default) {
64+
if (_transaction == null) return false;
4765

48-
return true;
66+
await _initLock.WaitAsync(ct);
67+
try {
68+
await _transaction.CommitAsync(ct);
69+
await _transaction.DisposeAsync();
70+
_transaction = null;
71+
return true;
72+
}
73+
finally {
74+
_initLock.Release();
75+
}
4976
}
5077

78+
5179
public virtual async ValueTask<bool> TryRollbackTransactionAsync(CancellationToken ct = default) {
5280
if (_transaction == null) return false;
5381

@@ -79,7 +107,7 @@ public virtual async ValueTask<bool> TryCreateSavepointAsync(Guid id, Cancellati
79107
public virtual async ValueTask<T> GetDbContextAsync<T>(CancellationToken ct = default) where T : DbContext {
80108
if (typeof(T) != typeof(TDbContext)) throw new NotSupportedException($"DbContext type '{typeof(T)}' is not supported by this UnitOfWork.");
81109

82-
TDbContext dbContext = await LazyDb.GetValueAsync(ct);
110+
TDbContext dbContext = await GetDbContextAsync(ct);
83111
return dbContext as T ?? throw new InvalidCastException($"Cannot cast DbContext of type '{dbContext.GetType()}' to '{typeof(T)}'");
84112
}
85113

@@ -88,7 +116,7 @@ public virtual async ValueTask<TRepo> GetRepositoryAsync<TRepo>(CancellationToke
88116

89117
// Cache miss so we create a new instance
90118
var repo = await CreateAndAttachRepositoryAsync<TRepo>(ct);
91-
119+
92120
AttachedRepositories.AddOrUpdate(typeof(TRepo), repo);
93121
return repo;
94122
}
@@ -102,22 +130,38 @@ private async ValueTask<TRepo> CreateAndAttachRepositoryAsync<TRepo>(Cancellatio
102130
}
103131

104132
public virtual async ValueTask DisposeAsync() {
105-
if (_transaction != null) await TryRollbackTransactionAsync();
133+
// ReSharper disable once ConditionIsAlwaysTrueOrFalseAccordingToNullableAPIContract
134+
if (_initLock == null) {
135+
GC.SuppressFinalize(this);
136+
return;
137+
}
106138

107-
if (!AttachedRepositories.IsEmpty) {
108-
foreach (IUnitOfWorkRepository repository in AttachedRepositories.Values) {
109-
if (repository is not UnitOfWorkRepository<TDbContext> castedRepo) continue;
139+
await _initLock.WaitAsync();
140+
try {
141+
if (_transaction != null) {
142+
await TryRollbackTransactionAsync();
143+
}
110144

111-
castedRepo.Detach();
145+
foreach (IUnitOfWorkRepository repository in AttachedRepositories.Values) {
146+
if (repository is UnitOfWorkRepository<TDbContext> castedRepo) {
147+
castedRepo.Detach();
148+
}
112149
}
113150

114151
AttachedRepositories.Clear();
115-
}
116-
117-
serviceScope.Dispose();
118152

119-
await LazyDb.DisposeAsync();
153+
if (_dbContext != null) {
154+
await _dbContext.DisposeAsync();
155+
_dbContext = null;
156+
}
120157

121-
GC.SuppressFinalize(this);
158+
serviceScope.Dispose();
159+
}
160+
finally {
161+
_initLock.Release();
162+
_initLock.Dispose();
163+
GC.SuppressFinalize(this);
164+
}
122165
}
166+
123167
}

‎tests/Benchmarks.CodeOfChaos.Types.TypedValueStore/Benchmarks.CodeOfChaos.Types.TypedValueStore.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
</PropertyGroup>
99

1010
<ItemGroup>
11-
<PackageReference Include="BenchmarkDotNet" Version="0.15.0" />
11+
<PackageReference Include="BenchmarkDotNet" Version="0.15.1" />
1212
<PackageReference Include="Microsoft.Extensions.Logging" Version="9.0.5" />
1313
</ItemGroup>
1414

‎tests/Tests.CodeOfChaos.Types.DataSeeder/Tests.CodeOfChaos.Types.DataSeeder.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
<PackageReference Include="JetBrains.Annotations" Version="2024.3.0"/>
1515
<PackageReference Include="Microsoft.Extensions.Logging" Version="9.0.5" />
1616
<PackageReference Include="Moq" Version="4.20.72"/>
17-
<PackageReference Include="TUnit" Version="0.24.0" />
17+
<PackageReference Include="TUnit" Version="0.25.21" />
1818
<PackageReference Include="Bogus" Version="35.6.3" />
1919
</ItemGroup>
2020

‎tests/Tests.CodeOfChaos.Types.TypedValueStore/Tests.CodeOfChaos.Types.TypedValueStore.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
<ItemGroup>
1515
<PackageReference Include="JetBrains.Annotations" Version="2024.3.0"/>
1616
<PackageReference Include="Moq" Version="4.20.72"/>
17-
<PackageReference Include="TUnit" Version="0.24.0" />
17+
<PackageReference Include="TUnit" Version="0.25.21" />
1818
<PackageReference Include="Bogus" Version="35.6.3" />
1919
</ItemGroup>
2020

‎tests/Tests.CodeOfChaos.Types.UnitOfWork/Tests.CodeOfChaos.Types.UnitOfWork.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
<PackageReference Include="Microsoft.EntityFrameworkCore" Version="9.0.5" />
1111
<PackageReference Include="Microsoft.EntityFrameworkCore.InMemory" Version="9.0.5" />
1212
<PackageReference Include="Moq" Version="4.20.72"/>
13-
<PackageReference Include="TUnit" Version="0.24.0" />
13+
<PackageReference Include="TUnit" Version="0.25.21" />
1414
</ItemGroup>
1515

1616
<ItemGroup>

‎tests/Tests.CodeOfChaos.Types.UnitOfWork/UnitOfWorkTests.cs‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ public void Setup() {
3939
_dbTransaction = new Mock<IDbContextTransaction>();
4040

4141
// Mock the DatabaseFacade for the DbContext and its transaction behavior
42-
Mock<DatabaseFacade> mockDatabaseFacade = new Mock<DatabaseFacade>(_dbContext.Object);
42+
var mockDatabaseFacade = new Mock<DatabaseFacade>(_dbContext.Object);
4343

4444
mockDatabaseFacade
4545
.Setup(db => db.BeginTransactionAsync(It.IsAny<CancellationToken>()))
@@ -69,11 +69,6 @@ public void Setup() {
6969

7070
}
7171

72-
[After(Test)]
73-
public async Task Cleanup() {
74-
await _unitOfWork.DisposeAsync();
75-
}
76-
7772
// -----------------------------------------------------------------------------------------------------------------
7873
// Tests
7974
// -----------------------------------------------------------------------------------------------------------------
@@ -116,7 +111,7 @@ public async Task TryCommitTransactionAsync_ShouldCommitTransaction() {
116111
// Assert
117112
await Assert.That(result).IsTrue();
118113
_dbTransaction.Verify(expression: transaction => transaction.CommitAsync(It.IsAny<CancellationToken>()), Times.Once);
119-
_dbTransaction.Verify(expression: transaction => transaction.Dispose(), Times.Once);
114+
_dbTransaction.Verify(expression: transaction => transaction.DisposeAsync(), Times.Once);
120115
}
121116

122117
[Test]

‎tests/Tests.CodeOfChaos.Types/Tests.CodeOfChaos.Types.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
<ItemGroup>
1414
<PackageReference Include="JetBrains.Annotations" Version="2024.3.0"/>
1515
<PackageReference Include="Moq" Version="4.20.72"/>
16-
<PackageReference Include="TUnit" Version="0.24.0" />
16+
<PackageReference Include="TUnit" Version="0.25.21" />
1717
<PackageReference Include="Bogus" Version="35.6.3" />
1818
</ItemGroup>
1919

0 commit comments

Comments
 (0)