fix(client): propagate CancellationToken through safe-mode setup (#123)

CHANGELOG.md

@@ -81,6 +81,13 @@ Behaviour changes that callers may notice. None of them changes a public signatu
 - `SSEClient` dispose race: the response handed off under the lock is now used via the
   local, so `Dispose()` during connect can no longer cause an NRE or a spurious
   failure report.
+- Safe mode's one-time setup now honours the `CancellationToken` passed to `ExecuteAsync`,
+  `ExecuteWithResultAsync` and the pager (#123). Before, a cancelled first query still
+  waited for the setup lock and the schema query. A cancelled schema query is not cached:
+  the next query runs it again. New overloads `GetRequestAsync(CancellationToken)`,
+  `GetRequestAsJsonAsync(CancellationToken)` and
+  `GraphClient.GetSchemaForSafeModeAsync(CancellationToken)` take the token; the existing
+  parameterless methods are unchanged.
 
 ### Performance
 

src/Linq2GraphQL.Client/GraphBaseExecute.cs

@@ -32,20 +32,22 @@ public GraphBaseExecute(GraphClient client, OperationType operationType, QueryNo
         queryExecutor = new QueryExecutor<T>(client);
     }
 
-    private async Task InitQueryAsync()
+    private async Task InitQueryAsync(CancellationToken cancellationToken)
     {
         // Fast path: initialization is one-shot; skip the semaphore once done.
         if (intialized)
         {
             return;
         }
 
-        await _lock.WaitAsync();
+        // Outside the try: a cancelled wait never acquired the lock, so it must not release it.
+        await _lock.WaitAsync(cancellationToken);
         try
         {
             if (!intialized)
             {
-                schema = await client.GetSchemaForSafeModeAsync();
+                // A cancelled schema fetch throws before intialized is set, so the next caller retries.
+                schema = await client.GetSchemaForSafeModeAsync(cancellationToken);
                 QueryNode.SetAllUniqueVariableNames();
                 QueryNode.AddPrimitiveChildren(true, schema);
                 intialized = true;
@@ -76,11 +78,18 @@ private string GetQueryVariablesString(IReadOnlyList<ArgumentValue> args)
         return $"({text})";
     }
 
-    public async Task<string> GetRequestAsJsonAsync()
+    public Task<string> GetRequestAsJsonAsync()
     {
-
+        return GetRequestAsJsonAsync(CancellationToken.None);
+    }
 
-        var request = await GetRequestAsync();
+    /// <summary>
+    /// Like <see cref="GetRequestAsJsonAsync()"/>, but the one-time safe-mode setup (waiting for
+    /// initialization and the schema query) stops when <paramref name="cancellationToken"/> is cancelled.
+    /// </summary>
+    public async Task<string> GetRequestAsJsonAsync(CancellationToken cancellationToken)
+    {
+        var request = await GetRequestAsync(cancellationToken);
 
 
         var result = request.Query;
@@ -124,9 +133,20 @@ public GraphQLRequest GetRequest()
         };
     }
 
-    public async Task<GraphQLRequest> GetRequestAsync()
+    public Task<GraphQLRequest> GetRequestAsync()
+    {
+        return GetRequestAsync(CancellationToken.None);
+    }
+
+    /// <summary>
+    /// Builds the request this query sends. The first call does one-time setup: with
+    /// <see cref="GraphClientOptions.UseSafeMode"/> on, that includes the schema query (unless the schema is
+    /// already cached). Cancelling <paramref name="cancellationToken"/> stops waiting for that setup and cancels
+    /// this caller's schema query; nothing is cached, so a later call starts the setup again.
+    /// </summary>
+    public async Task<GraphQLRequest> GetRequestAsync(CancellationToken cancellationToken)
     {
-        await InitQueryAsync();
+        await InitQueryAsync(cancellationToken);
 
         // Compute active arguments once for both the query string and the variables.
         var args = QueryNode.GetAllActiveArguments();

src/Linq2GraphQL.Client/GraphClient.cs

@@ -135,7 +135,24 @@ private string GetSubscriptionUrl()
         };
     }
 
-    public async Task<GraphQLSchema> GetSchemaForSafeModeAsync()
+    public Task<GraphQLSchema> GetSchemaForSafeModeAsync()
+    {
+        return GetSchemaForSafeModeAsync(CancellationToken.None);
+    }
+
+    /// <summary>
+    /// The schema safe mode validates against, fetched once and cached per base address, deprecated-fields
+    /// flag and schema cache key. Returns null when <see cref="GraphClientOptions.UseSafeMode"/> is off.
+    /// </summary>
+    /// <remarks>
+    /// <paramref name="cancellationToken"/> cancels only this caller's schema query. That is safe with
+    /// <c>IMemoryCache.GetOrCreateAsync</c>: it does not share one in-flight fetch between callers (each
+    /// cache miss runs its own factory), and a factory that throws, including on cancellation, stores
+    /// nothing, so a cancelled fetch is never cached. If the fetch is ever made shared (single-flight),
+    /// it must not run on one caller's token; each caller would then stop waiting with
+    /// <c>Task.WaitAsync(cancellationToken)</c> instead.
+    /// </remarks>
+    public async Task<GraphQLSchema> GetSchemaForSafeModeAsync(CancellationToken cancellationToken)
     {
         if (!options.Value.UseSafeMode)
         {
@@ -159,7 +176,7 @@ public async Task<GraphQLSchema> GetSchemaForSafeModeAsync()
             }
 
             var graphRequest = new GraphQLRequest { Query = query };
-            return await executor.ExecuteRequestAsync("__schema", graphRequest);
+            return await executor.ExecuteRequestAsync("__schema", graphRequest, cancellationToken);
         });
     }
 }

src/Linq2GraphQL.Client/GraphQueryExecute.cs

@@ -24,13 +24,15 @@ public GraphQueryExecute(GraphClient client, OperationType operationType, QueryN
     public async Task<T> ExecuteBaseAsync(CancellationToken cancellationToken = default)
     {
         BaseResult =
-            await queryExecutor.ExecuteRequestAsync(QueryNode.Name, await GetRequestAsync(), cancellationToken);
+            await queryExecutor.ExecuteRequestAsync(QueryNode.Name, await GetRequestAsync(cancellationToken),
+                cancellationToken);
         return BaseResult;
     }
 
     public async Task<GraphResult<T>> ExecuteBaseWithResultAsync(CancellationToken cancellationToken = default)
     {
-        var result = await queryExecutor.ExecuteRawAsync(QueryNode.Name, await GetRequestAsync(), cancellationToken);
+        var result = await queryExecutor.ExecuteRawAsync(QueryNode.Name, await GetRequestAsync(cancellationToken),
+            cancellationToken);
         BaseResult = result.Data;
         return result;
     }

test/Linq2GraphQL.Tests/GraphBaseExecuteTests.cs

@@ -49,6 +49,205 @@ protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage reques
         }
     }
 
+    private const string EmptySchemaJson = """{"data":{"__schema":{"types":[]}}}""";
+
+    [Fact]
+    public async Task GetRequestAsync_CancelledDuringIntrospection_ThrowsOperationCanceled()
+    {
+        var handler = new ScriptedHandler(HangUntilCancelled);
+        var exec = new TestExecute(CreateSafeModeClient(handler), e => e);
+        using var cts = new CancellationTokenSource();
+
+        var pending = exec.GetRequestAsync(cts.Token);
+        await handler.FirstRequestStarted;
+        cts.Cancel();
+
+        await AssertCancelledAsync(pending);
+    }
+
+    [Fact]
+    public async Task GetRequestAsync_PreCancelledToken_SendsNoRequest()
+    {
+        var handler = new ScriptedHandler(Respond(EmptySchemaJson));
+        var exec = new TestExecute(CreateSafeModeClient(handler), e => e);
+
+        await AssertCancelledAsync(exec.GetRequestAsync(new CancellationToken(true)));
+
+        Assert.Equal(0, handler.RequestCount);
+    }
+
+    [Fact]
+    public async Task GetRequestAsync_AfterCancelledIntrospection_RetriesAndSucceeds()
+    {
+        // A cancelled schema fetch must not be cached, and must not leave the query half-initialized.
+        var handler = new ScriptedHandler(HangUntilCancelled, Respond(EmptySchemaJson));
+        var exec = new TestExecute(CreateSafeModeClient(handler), e => e);
+        using var cts = new CancellationTokenSource();
+
+        var cancelled = exec.GetRequestAsync(cts.Token);
+        await handler.FirstRequestStarted;
+        cts.Cancel();
+        await AssertCancelledAsync(cancelled);
+
+        var request = await exec.GetRequestAsync(CancellationToken.None);
+
+        Assert.NotNull(request.Query);
+        Assert.Equal(2, handler.RequestCount);
+        Assert.NotNull(exec.Schema);
+    }
+
+    [Fact]
+    public async Task GetRequestAsync_CancelWhileWaitingForLock_DoesNotDisturbLockHolder()
+    {
+        // The first caller holds the init lock while its introspection request is blocked on a gate it
+        // does not cancel. A second caller waiting for the lock must be able to give up on its own.
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var handler = new ScriptedHandler(async _ =>
+        {
+            await gate.Task;
+            return Ok(EmptySchemaJson);
+        });
+        var exec = new TestExecute(CreateSafeModeClient(handler), e => e);
+
+        var holder = exec.GetRequestAsync(CancellationToken.None);
+        await handler.FirstRequestStarted;
+
+        using var cts = new CancellationTokenSource();
+        var waiter = exec.GetRequestAsync(cts.Token);
+        cts.Cancel();
+
+        await AssertCancelledAsync(waiter);
+        Assert.False(holder.IsCompleted);
+
+        gate.SetResult();
+        var request = await holder;
+
+        Assert.NotNull(request.Query);
+        Assert.Equal(1, handler.RequestCount);
+    }
+
+    [Fact]
+    public async Task GetRequestAsync_WithoutToken_StillInitializes()
+    {
+        // The parameterless overload is kept for binary compatibility and must behave as before.
+        var handler = new ScriptedHandler(Respond(EmptySchemaJson));
+        var exec = new TestExecute(CreateSafeModeClient(handler), e => e);
+
+        var request = await exec.GetRequestAsync();
+
+        Assert.NotNull(request.Query);
+        Assert.Equal(1, handler.RequestCount);
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_CancelledDuringSafeModeSetup_ThrowsOperationCanceled()
+    {
+        var handler = new ScriptedHandler(HangUntilCancelled);
+        var exec = new GraphQueryExecute<TestType, TestType>(CreateSafeModeClient(handler), OperationType.Query,
+            TestExecute.CreateNode(), e => e);
+        using var cts = new CancellationTokenSource();
+
+        var pending = exec.ExecuteAsync(cts.Token);
+        await handler.FirstRequestStarted;
+        cts.Cancel();
+
+        await AssertCancelledAsync(pending);
+        Assert.Equal(1, handler.RequestCount);
+    }
+
+    [Fact]
+    public async Task ExecuteWithResultAsync_CancelledDuringSafeModeSetup_ThrowsOperationCanceled()
+    {
+        var handler = new ScriptedHandler(HangUntilCancelled);
+        var exec = new GraphQueryExecute<TestType, TestType>(CreateSafeModeClient(handler), OperationType.Query,
+            TestExecute.CreateNode(), e => e);
+        using var cts = new CancellationTokenSource();
+
+        var pending = exec.ExecuteWithResultAsync(cts.Token);
+        await handler.FirstRequestStarted;
+        cts.Cancel();
+
+        await AssertCancelledAsync(pending);
+        Assert.Equal(1, handler.RequestCount);
+    }
+
+    [Fact]
+    public async Task GetSchemaForSafeModeAsync_Cancelled_IsNotCached()
+    {
+        var handler = new ScriptedHandler(HangUntilCancelled, Respond(EmptySchemaJson));
+        var client = CreateSafeModeClient(handler);
+        using var cts = new CancellationTokenSource();
+
+        var cancelled = client.GetSchemaForSafeModeAsync(cts.Token);
+        await handler.FirstRequestStarted;
+        cts.Cancel();
+        await AssertCancelledAsync(cancelled);
+
+        Assert.NotNull(await client.GetSchemaForSafeModeAsync());
+        Assert.NotNull(await client.GetSchemaForSafeModeAsync());
+        Assert.Equal(2, handler.RequestCount);
+    }
+
+    /// <summary>
+    /// The task must end as cancelled by the caller's token, promptly. Bounded so a dropped token fails the
+    /// test fast — <see cref="Task.WaitAsync(TimeSpan)"/> throws <see cref="TimeoutException"/>, which is no
+    /// <see cref="OperationCanceledException"/> — and the client's HTTP timeout is infinite (see
+    /// <see cref="CreateSafeModeClient"/>) so a timeout can never pass as a cancellation.
+    /// </summary>
+    private static async Task AssertCancelledAsync(Task task)
+    {
+        var ex = await Record.ExceptionAsync(() => task.WaitAsync(TimeSpan.FromSeconds(10)));
+
+        Assert.IsAssignableFrom<OperationCanceledException>(ex);
+    }
+
+    private static GraphClient CreateSafeModeClient(HttpMessageHandler handler)
+    {
+        return new GraphClient(new HttpClient(handler)
+            {
+                BaseAddress = new Uri("https://localhost/graphql/"),
+                Timeout = Timeout.InfiniteTimeSpan
+            },
+            Microsoft.Extensions.Options.Options.Create(new GraphClientOptions { UseSafeMode = true }),
+            new ServiceCollection().AddMemoryCache().BuildServiceProvider());
+    }
+
+    private static HttpResponseMessage Ok(string json) => new(System.Net.HttpStatusCode.OK)
+    {
+        Content = new StringContent(json, System.Text.Encoding.UTF8, "application/json")
+    };
+
+    private static Func<CancellationToken, Task<HttpResponseMessage>> Respond(string json) =>
+        _ => Task.FromResult(Ok(json));
+
+    private static async Task<HttpResponseMessage> HangUntilCancelled(CancellationToken cancellationToken)
+    {
+        await Task.Delay(Timeout.Infinite, cancellationToken);
+        throw new InvalidOperationException("unreachable");
+    }
+
+    /// <summary>Plays one response script per request, in order; the last one repeats.</summary>
+    private sealed class ScriptedHandler(params Func<CancellationToken, Task<HttpResponseMessage>>[] script)
+        : HttpMessageHandler
+    {
+        private readonly TaskCompletionSource firstRequestStarted =
+            new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        private int requestCount;
+
+        public int RequestCount => Volatile.Read(ref requestCount);
+
+        public Task FirstRequestStarted => firstRequestStarted.Task.WaitAsync(TimeSpan.FromSeconds(10));
+
+        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request,
+            CancellationToken cancellationToken)
+        {
+            var index = Interlocked.Increment(ref requestCount) - 1;
+            firstRequestStarted.TrySetResult();
+            return script[Math.Min(index, script.Length - 1)](cancellationToken);
+        }
+    }
+
     [Fact]
     public void NonIdentitySelector_IsCompiled()
     {
@@ -99,7 +298,7 @@ private sealed class TestExecute(GraphClient client, Expression<Func<TestType, T
     {
         public bool HasCompiledMapper => mapper != null;
 
-        private static QueryNode CreateNode()
+        internal static QueryNode CreateNode()
         {
             var member = typeof(TestFields).GetProperty(nameof(TestFields.Hello));
             return new QueryNode(member, topLevel: true);