66using AutoMapper ;
77using LeadCMS . Data ;
88using LeadCMS . DataAnnotations ;
9- using LeadCMS . Entities ;
109using Microsoft . EntityFrameworkCore ;
1110using Npgsql ;
1211
@@ -23,11 +22,12 @@ public class PostgresNotificationService : IHostedService, IDisposable
2322 private readonly IConfiguration configuration ;
2423 private readonly IMapper mapper ;
2524 private readonly HashSet < Type > supportedTypes ;
26-
25+
2726 private NpgsqlConnection ? notificationConnection ;
2827 private Task ? listeningTask ;
2928 private Timer ? pollingTimer ;
3029 private CancellationTokenSource ? cancellationTokenSource ;
30+ private CancellationTokenSource ? listeningCancellationTokenSource ;
3131 private bool isListening = false ;
3232 private int lastPolledChangeLogId = 0 ;
3333
@@ -49,10 +49,10 @@ public PostgresNotificationService(
4949 public Task StartAsync ( CancellationToken cancellationToken )
5050 {
5151 cancellationTokenSource = new CancellationTokenSource ( ) ;
52-
52+
5353 // Start monitoring client connections
5454 _ = Task . Run ( MonitorClientConnections , cancellationToken ) ;
55-
55+
5656 logger . LogInformation ( "PostgresNotificationService started" ) ;
5757 return Task . CompletedTask ;
5858 }
@@ -77,6 +77,7 @@ public void Dispose()
7777 finally
7878 {
7979 cancellationTokenSource ? . Dispose ( ) ;
80+ listeningCancellationTokenSource ? . Dispose ( ) ;
8081 notificationConnection ? . Dispose ( ) ;
8182 pollingTimer ? . Dispose ( ) ;
8283 }
@@ -92,7 +93,7 @@ private async Task MonitorClientConnections()
9293 try
9394 {
9495 var clientCount = clientManager . ConnectedClientCount ;
95-
96+
9697 if ( clientCount > 0 && ! isListening )
9798 {
9899 // First client connected - start listening
@@ -136,7 +137,7 @@ private async Task StartListening()
136137 // This determines what changes we need to poll for
137138 using var scope = serviceProvider . CreateScope ( ) ;
138139 var dbContext = scope . ServiceProvider . GetRequiredService < PgDbContext > ( ) ;
139-
140+
140141 var minClientLastId = clientManager . GetMinimumLastChangeLogId ( ) ;
141142 if ( minClientLastId . HasValue )
142143 {
@@ -168,9 +169,9 @@ private async Task StartListening()
168169 }
169170
170171 // Setup PostgreSQL NOTIFY listener
171- var connectionString = configuration . GetConnectionString ( "DefaultConnection" ) ??
172+ var connectionString = configuration . GetConnectionString ( "DefaultConnection" ) ??
172173 BuildConnectionString ( ) ;
173-
174+
174175 notificationConnection = new NpgsqlConnection ( connectionString ) ;
175176 await notificationConnection . OpenAsync ( cancellationTokenSource ! . Token ) ;
176177
@@ -188,8 +189,11 @@ private async Task StartListening()
188189
189190 logger . LogInformation ( "Successfully executed LISTEN entity_changes and draft_changes commands" ) ;
190191
191- // Start background listening task
192- listeningTask = Task . Run ( ListenForNotifications , cancellationTokenSource . Token ) ;
192+ // Create separate cancellation token for listening task
193+ listeningCancellationTokenSource = new CancellationTokenSource ( ) ;
194+
195+ // Start background listening task with separate cancellation token
196+ listeningTask = Task . Run ( ( ) => ListenForNotifications ( listeningCancellationTokenSource . Token ) , cancellationTokenSource . Token ) ;
193197
194198 // Start polling timer as Plan B (every 5 seconds)
195199 pollingTimer = new Timer ( async _ => await PollForChanges ( ) , null , TimeSpan . FromSeconds ( 5 ) , TimeSpan . FromSeconds ( 5 ) ) ;
@@ -222,18 +226,41 @@ private async Task StopListening()
222226 pollingTimer ? . Dispose ( ) ;
223227 pollingTimer = null ;
224228
225- // Stop listening task
229+ // Cancel the listening task (cancels WaitAsync)
230+ if ( listeningCancellationTokenSource != null && ! listeningCancellationTokenSource . IsCancellationRequested )
231+ {
232+ listeningCancellationTokenSource . Cancel ( ) ;
233+ }
234+
235+ // Wait for listening task to finish
236+ if ( listeningTask != null )
237+ {
238+ try
239+ {
240+ await listeningTask ;
241+ }
242+ catch
243+ {
244+ /* ignore */
245+ }
246+ finally
247+ {
248+ listeningTask = null ;
249+ }
250+ }
251+
252+ // Now it is safe to execute commands and close the connection
226253 if ( notificationConnection != null )
227254 {
228255 try
229256 {
230257 // Unlisten from both channels
231258 using var cmd1 = new NpgsqlCommand ( "UNLISTEN entity_changes" , notificationConnection ) ;
232259 await cmd1 . ExecuteNonQueryAsync ( ) ;
233-
260+
234261 using var cmd2 = new NpgsqlCommand ( "UNLISTEN draft_changes" , notificationConnection ) ;
235262 await cmd2 . ExecuteNonQueryAsync ( ) ;
236-
263+
237264 logger . LogInformation ( "Successfully executed UNLISTEN commands for both entity_changes and draft_changes" ) ;
238265 }
239266 catch ( Exception ex )
@@ -257,11 +284,9 @@ private async Task StopListening()
257284 }
258285 }
259286
260- if ( listeningTask != null )
261- {
262- await listeningTask ;
263- listeningTask = null ;
264- }
287+ // Dispose of the listening cancellation token source
288+ listeningCancellationTokenSource ? . Dispose ( ) ;
289+ listeningCancellationTokenSource = null ;
265290
266291 logger . LogInformation ( "Stopped PostgreSQL LISTEN and polling" ) ;
267292 }
@@ -274,22 +299,22 @@ private async Task StopListening()
274299 /// <summary>
275300 /// Background task that waits for PostgreSQL notifications.
276301 /// </summary>
277- private async Task ListenForNotifications ( )
302+ private async Task ListenForNotifications ( CancellationToken cancellationToken )
278303 {
279304 try
280305 {
281- while ( isListening && ! cancellationTokenSource ! . Token . IsCancellationRequested )
306+ while ( isListening && ! cancellationToken . IsCancellationRequested )
282307 {
283308 if ( notificationConnection != null && notificationConnection . State == System . Data . ConnectionState . Open )
284309 {
285- await notificationConnection . WaitAsync ( cancellationTokenSource . Token ) ;
310+ await notificationConnection . WaitAsync ( cancellationToken ) ;
286311 }
287312 else
288313 {
289314 // Connection is closed or null, wait a bit before retrying
290315 logger . LogWarning ( "PostgreSQL notification connection is not available, waiting before retry" ) ;
291- await Task . Delay ( 5000 , cancellationTokenSource . Token ) ;
292-
316+ await Task . Delay ( 5000 , cancellationToken ) ;
317+
293318 // Try to restart listening if connection is lost
294319 if ( isListening )
295320 {
@@ -315,7 +340,7 @@ private async Task ListenForNotifications()
315340 {
316341 try
317342 {
318- await Task . Delay ( 5000 , cancellationTokenSource . Token ) ;
343+ await Task . Delay ( 5000 , cancellationToken ) ;
319344 await StopListening ( ) ;
320345 await StartListening ( ) ;
321346 }
@@ -408,11 +433,11 @@ await clientManager.SendDraftNotificationAsync(
408433 draft . Data ) ;
409434
410435 var thisDraftAt = draft . UpdatedAt ?? draft . CreatedAt ;
411-
436+
412437 if ( maxSent == null || thisDraftAt > maxSent )
413438 {
414439 maxSent = thisDraftAt ;
415- }
440+ }
416441 }
417442
418443 if ( maxSent != null )
@@ -514,7 +539,7 @@ await clientManager.SendNotificationAsync(
514539 if ( maxSent == null || change . Id > maxSent )
515540 {
516541 maxSent = change . Id ;
517- }
542+ }
518543 }
519544 }
520545
@@ -542,7 +567,7 @@ private async Task<Dictionary<int, object>> GetEntityData(PgDbContext dbContext,
542567 // Find the entity type
543568 var assembly = typeof ( PgDbContext ) . Assembly ; // Use the correct assembly
544569 var type = assembly . GetTypes ( ) . FirstOrDefault ( t => t . Name == entityType ) ;
545-
570+
546571 if ( type == null )
547572 {
548573 logger . LogWarning ( "Entity type {EntityType} not found" , entityType ) ;
@@ -602,7 +627,7 @@ private async Task<Dictionary<int, object>> GetEntityData(PgDbContext dbContext,
602627 foreach ( var entity in entities )
603628 {
604629 var idValue = ( int ) entity . GetType ( ) . GetProperty ( "Id" ) ! . GetValue ( entity ) ! ;
605-
630+
606631 try
607632 {
608633 if ( detailsDtoType != null )
0 commit comments