Refactor to use ThrottleDebounce package, fix throttle being leading/trailing instead of just trailing

Author: LucDeCaf

PowerSync/PowerSync.Common/Client/PowerSyncDatabase.cs

@@ -1,16 +1,16 @@
 namespace PowerSync.Common.Client;
 
-using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Text.RegularExpressions;
+using System.Threading.Channels;
 using System.Threading.Tasks;
 
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 
 using Newtonsoft.Json;
-
 using Nito.AsyncEx;
+using ThrottleDebounce;
 
 using PowerSync.Common.Client.Connection;
 using PowerSync.Common.Client.Sync.Bucket;
@@ -990,7 +990,7 @@ internal async Task<HashSet<string>> GetSourceTables(string sql, object?[]? para
     private async IAsyncEnumerable<WatchOnChangeEvent> OnRawTableChange(
         HashSet<string> watchedTables,
         IAsyncEnumerable<DBAdapterEvents.TablesUpdatedEvent> listener,
-        [EnumeratorCancellation] CancellationToken token,
+        [EnumeratorCancellation] CancellationToken signal,
         bool triggerImmediately = false,
         int throttleMs = DEFAULT_WATCH_THROTTLE_MS
     )
@@ -1014,102 +1014,66 @@ private async IAsyncEnumerable<WatchOnChangeEvent> OnRawTableChange(
             yield break;
         }
 
-        // Leading + trailing edge throttle
-        var enumerator = listener.GetAsyncEnumerator(token);
-        try
-        {
-            var accumulatedTables = new HashSet<string>();
-            var changedTables = new HashSet<string>();
-            long lastYieldTime = 0;
-
-            Task<bool> moveNextTask = enumerator.MoveNextAsync().AsTask();
-            Task? throttleTask = null;
-
-            while (true)
-            {
-                if (throttleTask != null)
-                    await Task.WhenAny(moveNextTask, throttleTask);
-                else
-                {
-                    try { await moveNextTask; }
-                    catch (OperationCanceledException) { break; }
-                }
+        // Throttled - publish via throttled call to an action that flushes accumulated changes into this channel
+        var channel = Channel.CreateUnbounded<WatchOnChangeEvent>();
+        var accumulatedTables = new HashSet<string>();
 
-                if (throttleTask != null && throttleTask.IsCompleted && !moveNextTask.IsCompleted)
+        _ = Task.Run(async () =>
+        {
+            using var throttledFlush = Throttler.Throttle(() =>
                 {
-                    // Throttle timer expired without a new event
-                    if (accumulatedTables.Count > 0)
+                    // Safe to lock directly on accumulatedTables because it's a local variable
+                    lock (accumulatedTables)
                     {
-                        lastYieldTime = Stopwatch.GetTimestamp();
-                        yield return new WatchOnChangeEvent { ChangedTables = [.. accumulatedTables] };
+                        if (accumulatedTables.Count == 0) return;
+                        channel.Writer.TryWrite(new WatchOnChangeEvent { ChangedTables = [.. accumulatedTables] });
                         accumulatedTables.Clear();
                     }
-                    throttleTask = null;
-                    continue;
-                }
-
-                // A new event arrived (possibly alongside throttle)
-                // Check if the event actually exists or if this is the end of the enumerator
-                bool hasNext;
-                try { hasNext = await moveNextTask; }
-                catch (OperationCanceledException) { break; }
-                if (!hasNext) break;
-
-                // Accumulate changed tables from the most recent OnTablesUpdated event
-                GetTablesFromNotification(enumerator.Current.TablesUpdated, changedTables);
-
-                // Filter only watched tables and add to accumulatedTables set
-                changedTables.IntersectWith(watchedTables);
-                accumulatedTables.UnionWith(changedTables);
+                },
+                TimeSpan.FromMilliseconds(throttleMs),
+                leading: false,
+                trailing: true
+            );
 
-                if (accumulatedTables.Count > 0)
+            try
+            {
+                var changedTables = new HashSet<string>();
+                await foreach (var e in listener)
                 {
-                    var now = Stopwatch.GetTimestamp();
+                    GetTablesFromNotification(e.TablesUpdated, changedTables);
+                    changedTables.IntersectWith(watchedTables);
+                    if (changedTables.Count == 0) continue;
 
-                    // There's a nice built-in method for this (Stopwatch.GetElapsedTime), but
-                    // it's not supported in .NET 6.0. :(
-                    var elapsedMs = (now - lastYieldTime) * 1000.0 / Stopwatch.Frequency;
-
-                    if (elapsedMs >= throttleMs)
+                    lock (accumulatedTables) { accumulatedTables.UnionWith(changedTables); }
+                    throttledFlush.Invoke();
+                }
+            }
+            catch (OperationCanceledException) { }
+            finally
+            {
+                // Flush any remaining events and close the channel
+                lock (accumulatedTables)
+                {
+                    if (accumulatedTables.Count > 0)
                     {
-                        // First event since throttle expiration
-                        // Fire immediately (leading edge) and reset throttle timer
-                        lastYieldTime = now;
-                        yield return new WatchOnChangeEvent { ChangedTables = [.. accumulatedTables] };
+                        channel.Writer.TryWrite(new WatchOnChangeEvent { ChangedTables = [.. accumulatedTables] });
                         accumulatedTables.Clear();
-                        throttleTask = null;
-                    }
-                    else
-                    {
-                        throttleTask ??= Task.Delay((int)(throttleMs - elapsedMs), token);
                     }
                 }
-
-                moveNextTask = enumerator.MoveNextAsync().AsTask();
+                channel.Writer.Complete();
             }
+        });
 
-            // Flush any remaining events
-            if (accumulatedTables.Count > 0)
-                yield return new WatchOnChangeEvent { ChangedTables = [.. accumulatedTables] };
-        }
-        finally
+        // Continuously pull values from channel and publish to the consumer
+        while (await channel.Reader.WaitToReadAsync(CancellationToken.None))
         {
-            await enumerator.DisposeAsync();
+            while (channel.Reader.TryRead(out var evt))
+            {
+                yield return evt;
+            }
         }
     }
 
-    private static void AccumulateMatchingTables(
-        DBAdapterEvents.TablesUpdatedEvent e,
-        HashSet<string> watchedTables,
-        HashSet<string> accumulated
-    )
-    {
-        var tables = new HashSet<string>();
-        GetTablesFromNotification(e.TablesUpdated, tables);
-        tables.IntersectWith(watchedTables);
-        accumulated.UnionWith(tables);
-    }
-
     private static void GetTablesFromNotification(INotification updateNotification, HashSet<string> changedTables)
     {
         changedTables.Clear();

PowerSync/PowerSync.Common/PowerSync.Common.csproj

@@ -35,6 +35,7 @@
     <PackageReference Include="Newtonsoft.Json" Version="13.0.3" />
     <PackageReference Include="Nito.AsyncEx" Version="5.1.2" />
     <PackageReference Include="System.Threading.Channels" Version="8.0.0" />
+    <PackageReference Include="ThrottleDebounce" Version="2.0.1" />
   </ItemGroup>
 
   <ItemGroup>

Tests/PowerSync/PowerSync.Common.Tests/Client/PowerSyncDatabaseTests.cs

@@ -867,13 +867,13 @@ public async Task Watch_CancelsOnTokenCancellation()
     public async Task OnChange_ThrottlesBatchesRapidChanges()
     {
         int eventCount = 0;
+        using var sem = new SemaphoreSlim(0);
         var tcs = new TaskCompletionSource<bool>();
 
         var listener = db.OnChange(new SQLWatchOptions
         {
             Tables = ["assets"],
             Signal = testCts.Token,
-            ThrottleMs = 200,
         });
 
         _ = Task.Run(async () =>
@@ -883,6 +883,7 @@ public async Task OnChange_ThrottlesBatchesRapidChanges()
                 await foreach (var _ in listener)
                 {
                     Interlocked.Increment(ref eventCount);
+                    sem.Release();
                 }
                 tcs.TrySetResult(true);
             }
@@ -897,23 +898,25 @@ public async Task OnChange_ThrottlesBatchesRapidChanges()
             await TestUtils.InsertRandomAsset(db);
         }
 
+        // Wait for the throttled event to arrive
+        Assert.True(await sem.WaitAsync(200));
         testCts.Cancel();
         Assert.True(await tcs.Task);
-
-        Assert.True(eventCount < 5, $"Expected fewer than 5 events but got {eventCount}");
+        Assert.Equal(1, eventCount);
     }
 
     [Fact(Timeout = 5000)]
     public async Task Watch_ThrottlesBatchesRapidChanges()
     {
         int eventCount = 0;
         long lastCount = 0;
+        using var sem = new SemaphoreSlim(0);
         var tcs = new TaskCompletionSource<bool>();
 
         var listener = db.Watch<CountResult>(
             "SELECT COUNT(*) AS count FROM assets",
             null,
-            new() { Signal = testCts.Token, ThrottleMs = 200 });
+            new() { Signal = testCts.Token });
 
         _ = Task.Run(async () =>
         {
@@ -923,6 +926,7 @@ public async Task Watch_ThrottlesBatchesRapidChanges()
                 {
                     lastCount = rows.First().count;
                     Interlocked.Increment(ref eventCount);
+                    sem.Release();
                 }
                 tcs.TrySetResult(true);
             }
@@ -937,6 +941,7 @@ public async Task Watch_ThrottlesBatchesRapidChanges()
             await TestUtils.InsertRandomAsset(db);
         }
 
+        Assert.True(await sem.WaitAsync(200));
         testCts.Cancel();
         Assert.True(await tcs.Task);
 
@@ -976,7 +981,7 @@ public async Task OnChange_NoThrottleWhenZero()
     }
 
     [Fact(Timeout = 5000)]
-    public async Task OnChange_FirstChangeIsNotDelayed()
+    public async Task OnChange_FirstChangeIsDelayedByThrottle()
     {
         using var sem = new SemaphoreSlim(0);
         var sw = Stopwatch.StartNew();
@@ -1000,26 +1005,30 @@ public async Task OnChange_FirstChangeIsNotDelayed()
         await db.Execute("INSERT INTO assets(id, description) VALUES(?, ?)", [Guid.NewGuid().ToString(), "test"]);
 
         Assert.True(await sem.WaitAsync(2000));
-        Assert.True(sw.ElapsedMilliseconds < 200, $"First event took {sw.ElapsedMilliseconds}ms, expected <200ms");
+        Assert.True(sw.ElapsedMilliseconds >= 400, $"First event took {sw.ElapsedMilliseconds}ms, expected >=400ms (trailing-only throttle)");
     }
 
     [Fact(Timeout = 5000)]
     public async Task OnChange_ThrottleCancelledCleanly()
     {
+        int eventCount = 0;
         var tcs = new TaskCompletionSource<bool>();
 
         var listener = db.OnChange(new SQLWatchOptions
         {
             Tables = ["assets"],
             Signal = testCts.Token,
-            ThrottleMs = 500,
+            ThrottleMs = 200,
         });
 
         _ = Task.Run(async () =>
         {
             try
             {
-                await foreach (var _ in listener) { }
+                await foreach (var _ in listener)
+                {
+                    Interlocked.Increment(ref eventCount);
+                }
                 tcs.TrySetResult(true);
             }
             catch (Exception ex)
@@ -1028,19 +1037,22 @@ public async Task OnChange_ThrottleCancelledCleanly()
             }
         });
 
-        // Insert to trigger the throttle delay, then cancel before the window expires
+        // Insert to trigger the throttle, then cancel before the window expires
         await db.Execute("INSERT INTO assets(id, description) VALUES(?, ?)", [Guid.NewGuid().ToString(), "test"]);
-        await Task.Delay(50);
+        await Task.Delay(100);
         testCts.Cancel();
 
         Assert.True(await tcs.Task);
+        // The flush-on-cancel should still deliver the accumulated event
+        Assert.Equal(1, eventCount);
     }
 
     [Fact]
     public async Task Watch_SingleEventForBatchedQuery()
     {
         int eventCount = 0;
         long lastCount = 0;
+        using var sem = new SemaphoreSlim(0);
         var tcs = new TaskCompletionSource<bool>();
 
         var listener = db.Watch<CountResult>(
@@ -1054,14 +1066,15 @@ public async Task Watch_SingleEventForBatchedQuery()
             {
                 lastCount = rows[0].count;
                 Interlocked.Increment(ref eventCount);
+                sem.Release();
             }
             tcs.TrySetResult(true);
         });
 
-        // Long batched query
-        const int QUERY_COUNT = 10000;
+        const int QUERY_COUNT = 1000;
         await TestUtils.InsertRandomAssets(db, QUERY_COUNT);
 
+        Assert.True(await sem.WaitAsync(200));
         testCts.Cancel();
         Assert.True(await tcs.Task);
         Assert.Equal(1, eventCount);
@@ -1072,6 +1085,7 @@ public async Task Watch_SingleEventForBatchedQuery()
     public async Task OnChange_SingleEventsForBatchedQuery()
     {
         int eventCount = 0;
+        using var sem = new SemaphoreSlim(0);
         var tcs = new TaskCompletionSource<bool>();
 
         var listener = db.OnChange(new()
@@ -1085,14 +1099,15 @@ public async Task OnChange_SingleEventsForBatchedQuery()
             await foreach (var _ in listener)
             {
                 Interlocked.Increment(ref eventCount);
+                sem.Release();
             }
             tcs.TrySetResult(true);
         });
 
-        // Long batched query
-        const int QUERY_COUNT = 10000;
+        const int QUERY_COUNT = 1000;
         await TestUtils.InsertRandomAssets(db, QUERY_COUNT);
 
+        Assert.True(await sem.WaitAsync(200));
         testCts.Cancel();
         Assert.True(await tcs.Task);