Handle request cancellation in the integrated HTTP server

app/Server/Database/Export/ViewerJsonExport.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Data;
 using DHT.Server.Data.Filters;
@@ -13,14 +14,14 @@ namespace DHT.Server.Database.Export;
 static class ViewerJsonExport {
 	private static readonly Log Log = Log.ForType(typeof(ViewerJsonExport));
 
-	public static async Task Generate(Stream stream, IDatabaseFile db, MessageFilter? filter = null) {
+	public static async Task Generate(Stream stream, IDatabaseFile db, MessageFilter? filter = null, CancellationToken cancellationToken = default) {
 		var perf = Log.Start();
 
 		var includedUserIds = new HashSet<ulong>();
 		var includedChannelIds = new HashSet<ulong>();
 		var includedServerIds = new HashSet<ulong>();
 
-		var includedMessages = await db.Messages.Get(filter).ToListAsync();
+		var includedMessages = await db.Messages.Get(filter, cancellationToken).ToListAsync(cancellationToken);
 		var includedChannels = new List<Channel>();
 
 		foreach (var message in includedMessages) {
@@ -28,7 +29,7 @@ static class ViewerJsonExport {
 			includedChannelIds.Add(message.Channel);
 		}
 
-		await foreach (var channel in db.Channels.Get()) {
+		await foreach (var channel in db.Channels.Get(cancellationToken)) {
 			if (includedChannelIds.Contains(channel.Id)) {
 				includedChannels.Add(channel);
 				includedServerIds.Add(channel.Server);
@@ -53,7 +54,7 @@ static class ViewerJsonExport {
 
 		perf.Step("Generate value object");
 
-		await JsonSerializer.SerializeAsync(stream, value, ViewerJsonContext.Default.ViewerJson);
+		await JsonSerializer.SerializeAsync(stream, value, ViewerJsonContext.Default.ViewerJson, cancellationToken);
 
 		perf.Step("Serialize to JSON");
 		perf.End();

app/Server/Database/Repositories/IChannelRepository.cs

@@ -15,7 +15,7 @@ public interface IChannelRepository {
 	
 	Task<long> Count(CancellationToken cancellationToken = default);
 	
-	IAsyncEnumerable<Channel> Get();
+	IAsyncEnumerable<Channel> Get(CancellationToken cancellationToken = default);
 
 	internal sealed class Dummy : IChannelRepository {
 		public IObservable<long> TotalCount { get; } = Observable.Return(0L);
@@ -28,7 +28,7 @@ public interface IChannelRepository {
 			return Task.FromResult(0L);
 		}
 
-		public IAsyncEnumerable<Channel> Get() {
+		public IAsyncEnumerable<Channel> Get(CancellationToken cancellationToken) {
 			return AsyncEnumerable.Empty<Channel>();
 		}
 	}

app/Server/Database/Repositories/IDownloadRepository.cs

@@ -24,7 +24,7 @@ public interface IDownloadRepository {
 
 	Task<bool> GetDownloadData(string normalizedUrl, Func<Stream, Task> dataProcessor);
 	
-	Task<bool> GetSuccessfulDownloadWithData(string normalizedUrl, Func<Data.Download, Stream, Task> dataProcessor);
+	Task<bool> GetSuccessfulDownloadWithData(string normalizedUrl, Func<Data.Download, Stream, CancellationToken, Task> dataProcessor, CancellationToken cancellationToken = default);
 
 	IAsyncEnumerable<DownloadItem> PullPendingDownloadItems(int count, DownloadItemFilter filter, CancellationToken cancellationToken = default);
 	
@@ -55,7 +55,7 @@ public interface IDownloadRepository {
 			return Task.FromResult(false);
 		}
 
-		public Task<bool> GetSuccessfulDownloadWithData(string normalizedUrl, Func<Data.Download, Stream, Task> dataProcessor) {
+		public Task<bool> GetSuccessfulDownloadWithData(string normalizedUrl, Func<Data.Download, Stream, CancellationToken, Task> dataProcessor, CancellationToken cancellationToken) {
 			return Task.FromResult(false);
 		}
 

app/Server/Database/Repositories/IMessageRepository.cs

@@ -16,7 +16,7 @@ public interface IMessageRepository {
 	
 	Task<long> Count(MessageFilter? filter = null, CancellationToken cancellationToken = default);
 	
-	IAsyncEnumerable<Message> Get(MessageFilter? filter = null);
+	IAsyncEnumerable<Message> Get(MessageFilter? filter = null, CancellationToken cancellationToken = default);
 	
 	IAsyncEnumerable<ulong> GetIds(MessageFilter? filter = null);
 	
@@ -33,7 +33,7 @@ public interface IMessageRepository {
 			return Task.FromResult(0L);
 		}
 
-		public IAsyncEnumerable<Message> Get(MessageFilter? filter) {
+		public IAsyncEnumerable<Message> Get(MessageFilter? filter, CancellationToken cancellationToken) {
 			return AsyncEnumerable.Empty<Message>();
 		}
 

app/Server/Database/Sqlite/Repositories/SqliteChannelRepository.cs

@@ -1,4 +1,5 @@
 using System.Collections.Generic;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Data;
@@ -54,13 +55,13 @@ sealed class SqliteChannelRepository : BaseSqliteRepository, IChannelRepository
 		return await conn.ExecuteReaderAsync("SELECT COUNT(*) FROM channels", static reader => reader?.GetInt64(0) ?? 0L, cancellationToken);
 	}
 
-	public async IAsyncEnumerable<Channel> Get() {
+	public async IAsyncEnumerable<Channel> Get([EnumeratorCancellation] CancellationToken cancellationToken) {
 		await using var conn = await pool.Take();
 
 		await using var cmd = conn.Command("SELECT id, server, name, parent_id, position, topic, nsfw FROM channels");
-		await using var reader = await cmd.ExecuteReaderAsync();
+		await using var reader = await cmd.ExecuteReaderAsync(cancellationToken);
 
-		while (await reader.ReadAsync()) {
+		while (await reader.ReadAsync(cancellationToken)) {
 			yield return new Channel {
 				Id = reader.GetUint64(0),
 				Server = reader.GetUint64(1),

app/Server/Database/Sqlite/Repositories/SqliteDownloadRepository.cs

@@ -210,7 +210,7 @@ sealed class SqliteDownloadRepository(SqliteConnectionPool pool) : BaseSqliteRep
 		return true;
 	}
 
-	public async Task<bool> GetSuccessfulDownloadWithData(string normalizedUrl, Func<Data.Download, Stream, Task> dataProcessor) {
+	public async Task<bool> GetSuccessfulDownloadWithData(string normalizedUrl, Func<Data.Download, Stream, CancellationToken, Task> dataProcessor, CancellationToken cancellationToken) {
 		await using var conn = await pool.Take();
 
 		await using var cmd = conn.Command(
@@ -228,8 +228,8 @@ sealed class SqliteDownloadRepository(SqliteConnectionPool pool) : BaseSqliteRep
 		string? type;
 		long rowid;
 		
-		await using (var reader = await cmd.ExecuteReaderAsync()) {
+		await using (var reader = await cmd.ExecuteReaderAsync(cancellationToken)) {
-			if (!await reader.ReadAsync()) {
+			if (!await reader.ReadAsync(cancellationToken)) {
 				return false;
 			}
 
@@ -239,7 +239,7 @@ sealed class SqliteDownloadRepository(SqliteConnectionPool pool) : BaseSqliteRep
 		}
 		
 		await using (var blob = new SqliteBlob(conn.InnerConnection, "download_blobs", "blob", rowid, readOnly: true)) {
-			await dataProcessor(new Data.Download(normalizedUrl, downloadUrl, DownloadStatus.Success, type, (ulong) blob.Length), blob);
+			await dataProcessor(new Data.Download(normalizedUrl, downloadUrl, DownloadStatus.Success, type, (ulong) blob.Length), blob, cancellationToken);
 		}
 
 		return true;

app/Server/Database/Sqlite/Repositories/SqliteMessageRepository.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Data;
@@ -213,7 +214,7 @@ sealed class SqliteMessageRepository : BaseSqliteRepository, IMessageRepository
 		}
 	}
 
-	public async IAsyncEnumerable<Message> Get(MessageFilter? filter) {
+	public async IAsyncEnumerable<Message> Get(MessageFilter? filter, [EnumeratorCancellation] CancellationToken cancellationToken) {
 		await using var conn = await pool.Take();
 
 		const string AttachmentSql =
@@ -269,9 +270,9 @@ sealed class SqliteMessageRepository : BaseSqliteRepository, IMessageRepository
 			 """
 		);
 
-		await using var reader = await messageCmd.ExecuteReaderAsync();
+		await using var reader = await messageCmd.ExecuteReaderAsync(cancellationToken);
 
-		while (await reader.ReadAsync()) {
+		while (await reader.ReadAsync(cancellationToken)) {
 			ulong messageId = reader.GetUint64(0);
 
 			yield return new Message {

app/Server/Endpoints/BaseEndpoint.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Net;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Database;
 using DHT.Utils.Http;
@@ -19,7 +20,9 @@ abstract class BaseEndpoint(IDatabaseFile db) {
 
 		try {
 			response.StatusCode = (int) HttpStatusCode.OK;
-			await Respond(ctx.Request, response);
+			await Respond(ctx.Request, response, ctx.RequestAborted);
+		} catch (OperationCanceledException) {
+			throw;
 		} catch (HttpException e) {
 			Log.Error(e);
 			response.StatusCode = (int) e.StatusCode;
@@ -35,7 +38,7 @@ abstract class BaseEndpoint(IDatabaseFile db) {
 		}
 	}
 
-	protected abstract Task Respond(HttpRequest request, HttpResponse response);
+	protected abstract Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken);
 
 	protected static async Task<JsonElement> ReadJson(HttpRequest request) {
 		try {

app/Server/Endpoints/GetDownloadedFileEndpoint.cs

@@ -1,4 +1,7 @@
+using System;
+using System.IO;
 using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Database;
 using DHT.Server.Download;
@@ -8,12 +11,16 @@ using Microsoft.AspNetCore.Http;
 namespace DHT.Server.Endpoints;
 
 sealed class GetDownloadedFileEndpoint(IDatabaseFile db) : BaseEndpoint(db) {
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override async Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		string url = WebUtility.UrlDecode((string) request.RouteValues["url"]!);
 		string normalizedUrl = DiscordCdn.NormalizeUrl(url);
 		
-		if (!await Db.Downloads.GetSuccessfulDownloadWithData(normalizedUrl, (download, stream) => response.WriteStreamAsync(download.Type, download.Size, stream))) {
+		if (!await Db.Downloads.GetSuccessfulDownloadWithData(normalizedUrl, WriteDataTo(response), cancellationToken)) {
 			response.Redirect(url, permanent: false);
 		}
 	}
+	
+	private static Func<Data.Download, Stream, CancellationToken, Task> WriteDataTo(HttpResponse response) {
+		return (download, stream, cancellationToken) => response.WriteStreamAsync(download.Type, download.Size, stream, cancellationToken);
+	}
 }

app/Server/Endpoints/GetTrackingScriptEndpoint.cs

@@ -1,4 +1,5 @@
 using System.Net.Mime;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Web;
 using DHT.Server.Database;
@@ -10,7 +11,7 @@ using Microsoft.AspNetCore.Http;
 namespace DHT.Server.Endpoints;
 
 sealed class GetTrackingScriptEndpoint(IDatabaseFile db, ServerParameters parameters, ResourceLoader resources) : BaseEndpoint(db) {
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override async Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		string bootstrap = await resources.ReadTextAsync("Tracker/bootstrap.js");
 		string script = bootstrap.Replace("= 0; /*[PORT]*/", "= " + parameters.Port + ";")
 		                         .Replace("/*[TOKEN]*/", HttpUtility.JavaScriptStringEncode(parameters.Token))
@@ -20,6 +21,6 @@ sealed class GetTrackingScriptEndpoint(IDatabaseFile db, ServerParameters parame
 		                         .Replace("/*[DEBUGGER]*/", request.Query.ContainsKey("debug") ? "debugger;" : "");
 		
 		response.Headers.Append("X-DHT", "1");
-		await response.WriteTextAsync(MediaTypeNames.Text.JavaScript, script);
+		await response.WriteTextAsync(MediaTypeNames.Text.JavaScript, script, cancellationToken);
 	}
 }

app/Server/Endpoints/GetViewerDataEndpoint.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Net;
 using System.Net.Mime;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Database;
 using DHT.Server.Database.Export;
@@ -11,7 +12,7 @@ using Microsoft.AspNetCore.Http;
 namespace DHT.Server.Endpoints; 
 
 sealed class GetViewerDataEndpoint(IDatabaseFile db, ViewerSessions viewerSessions) : BaseEndpoint(db) {
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		if (!request.Query.TryGetValue("session", out var sessionIdValue) || sessionIdValue.Count != 1 || !Guid.TryParse(sessionIdValue[0], out Guid sessionId)) {
 			throw new HttpException(HttpStatusCode.BadRequest, "Invalid session ID.");
 		}
@@ -19,6 +20,6 @@ sealed class GetViewerDataEndpoint(IDatabaseFile db, ViewerSessions viewerSessio
 		response.ContentType = MediaTypeNames.Application.Json;
 		
 		var session = viewerSessions.Get(sessionId);
-		await ViewerJsonExport.Generate(response.Body, Db, session.MessageFilter);
+		return ViewerJsonExport.Generate(response.Body, Db, session.MessageFilter, cancellationToken);
 	}
 }

app/Server/Endpoints/TrackChannelEndpoint.cs

@@ -1,5 +1,6 @@
 using System.Net;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Data;
 using DHT.Server.Database;
@@ -9,7 +10,7 @@ using Microsoft.AspNetCore.Http;
 namespace DHT.Server.Endpoints;
 
 sealed class TrackChannelEndpoint(IDatabaseFile db) : BaseEndpoint(db) {
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override async Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		var root = await ReadJson(request);
 		var server = ReadServer(root.RequireObject("server"), "server");
 		var channel = ReadChannel(root.RequireObject("channel"), "channel", server.Id);

app/Server/Endpoints/TrackMessagesEndpoint.cs

@@ -4,6 +4,7 @@ using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Net;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Data;
 using DHT.Server.Data.Filters;
@@ -19,7 +20,7 @@ sealed class TrackMessagesEndpoint(IDatabaseFile db) : BaseEndpoint(db) {
 	private const string HasNewMessages = "1";
 	private const string NoNewMessages = "0";
 
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override async Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		var root = await ReadJson(request);
 
 		if (root.ValueKind != JsonValueKind.Array) {
@@ -37,11 +38,11 @@ sealed class TrackMessagesEndpoint(IDatabaseFile db) : BaseEndpoint(db) {
 		}
 
 		var addedMessageFilter = new MessageFilter { MessageIds = addedMessageIds };
-		bool anyNewMessages = await Db.Messages.Count(addedMessageFilter) < addedMessageIds.Count;
+		bool anyNewMessages = await Db.Messages.Count(addedMessageFilter, CancellationToken.None) < addedMessageIds.Count;
 
 		await Db.Messages.Add(messages);
 
-		await response.WriteTextAsync(anyNewMessages ? HasNewMessages : NoNewMessages);
+		await response.WriteTextAsync(anyNewMessages ? HasNewMessages : NoNewMessages, cancellationToken);
 	}
 
 	private static Message ReadMessage(JsonElement json, string path) => new () {

app/Server/Endpoints/TrackUsersEndpoint.cs

@@ -1,5 +1,6 @@
 using System.Net;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using DHT.Server.Data;
 using DHT.Server.Database;
@@ -9,7 +10,7 @@ using Microsoft.AspNetCore.Http;
 namespace DHT.Server.Endpoints;
 
 sealed class TrackUsersEndpoint(IDatabaseFile db) : BaseEndpoint(db) {
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override async Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		var root = await ReadJson(request);
 
 		if (root.ValueKind != JsonValueKind.Array) {

app/Server/Endpoints/ViewerEndpoint.cs

@@ -16,13 +16,13 @@ sealed class ViewerEndpoint(IDatabaseFile db, ResourceLoader resources) : BaseEn
 	private readonly Dictionary<string, byte[]?> cache = new ();
 	private readonly SemaphoreSlim cacheSemaphore = new (1);
 
-	protected override async Task Respond(HttpRequest request, HttpResponse response) {
+	protected override async Task Respond(HttpRequest request, HttpResponse response, CancellationToken cancellationToken) {
 		string path = (string?) request.RouteValues["path"] ?? "index.html";
 		string resourcePath = "Viewer/" + path;
 		
 		byte[]? resourceBytes;
 
-		await cacheSemaphore.WaitAsync();
+		await cacheSemaphore.WaitAsync(cancellationToken);
 		try {
 			if (!cache.TryGetValue(resourcePath, out resourceBytes)) {
 				cache[resourcePath] = resourceBytes = await resources.ReadBytesAsyncIfExists(resourcePath);
@@ -36,7 +36,7 @@ sealed class ViewerEndpoint(IDatabaseFile db, ResourceLoader resources) : BaseEn
 		}
 		else {
 			var contentType = ContentTypeProvider.TryGetContentType(path, out string? type) ? type : null;
-			await response.WriteFileAsync(contentType, resourceBytes);
+			await response.WriteFileAsync(contentType, resourceBytes, cancellationToken);
 		}
 	}
 }

app/Server/Service/Middlewares/ServerLoggingMiddleware.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Diagnostics;
 using System.Threading.Tasks;
 using DHT.Utils.Logging;
@@ -6,24 +7,34 @@ using Microsoft.AspNetCore.Http.Extensions;
 
 namespace DHT.Server.Service.Middlewares;
 
-sealed class ServerLoggingMiddleware {
+sealed class ServerLoggingMiddleware(RequestDelegate next) {
 	private static readonly Log Log = Log.ForType<ServerLoggingMiddleware>();
 
-	private readonly RequestDelegate next;
-	
-	public ServerLoggingMiddleware(RequestDelegate next) {
-		this.next = next;
-	}
-
 	public async Task InvokeAsync(HttpContext context) {
 		var stopwatch = Stopwatch.StartNew();
+		try {
 			await next(context);
+		} catch (OperationCanceledException) {
+			OnFinished(stopwatch, context);
+			throw;
+		}
+
+		OnFinished(stopwatch, context);
+	}
+
+	private static void OnFinished(Stopwatch stopwatch, HttpContext context) {
 		stopwatch.Stop();
 
 		var request = context.Request;
 		var requestLength = request.ContentLength ?? 0L;
-		var responseStatus = context.Response.StatusCode;
 		var elapsedMs = stopwatch.ElapsedMilliseconds;
+
+		if (context.RequestAborted.IsCancellationRequested) {
+			Log.Debug("Request to " + request.GetEncodedPathAndQuery() + " (" + requestLength + " B) was cancelled after " + elapsedMs + " ms");
+		}
+		else {
+			var responseStatus = context.Response.StatusCode;
 			Log.Debug("Request to " + request.GetEncodedPathAndQuery() + " (" + requestLength + " B) returned " + responseStatus + ", took " + elapsedMs + " ms");
 		}
 	}
+}

app/Utils/Http/HttpExtensions.cs

@@ -1,33 +1,34 @@
 using System.IO;
 using System.Net.Mime;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 
 namespace DHT.Utils.Http; 
 
 public static class HttpExtensions {
-	public static Task WriteTextAsync(this HttpResponse response, string text) {
+	public static Task WriteTextAsync(this HttpResponse response, string text, CancellationToken cancellationToken) {
-		return WriteTextAsync(response, MediaTypeNames.Text.Plain, text);
+		return WriteTextAsync(response, MediaTypeNames.Text.Plain, text, cancellationToken);
 	}
 	
-	public static async Task WriteTextAsync(this HttpResponse response, string contentType, string text) {
+	public static async Task WriteTextAsync(this HttpResponse response, string contentType, string text, CancellationToken cancellationToken) {
 		response.ContentType = contentType;
-		await response.StartAsync();
+		await response.StartAsync(cancellationToken);
-		await response.WriteAsync(text, Encoding.UTF8);
+		await response.WriteAsync(text, Encoding.UTF8, cancellationToken);
 	}
 	
-	public static async Task WriteFileAsync(this HttpResponse response, string? contentType, byte[] bytes) {
+	public static async Task WriteFileAsync(this HttpResponse response, string? contentType, byte[] bytes, CancellationToken cancellationToken) {
 		response.ContentType = contentType ?? string.Empty;
 		response.ContentLength = bytes.Length;
-		await response.StartAsync();
+		await response.StartAsync(cancellationToken);
-		await response.Body.WriteAsync(bytes);
+		await response.Body.WriteAsync(bytes, cancellationToken);
 	}
 	
-	public static async Task WriteStreamAsync(this HttpResponse response, string? contentType, ulong? contentLength, Stream source) {
+	public static async Task WriteStreamAsync(this HttpResponse response, string? contentType, ulong? contentLength, Stream source, CancellationToken cancellationToken) {
 		response.ContentType = contentType ?? string.Empty;
 		response.ContentLength = (long?) contentLength;
-		await response.StartAsync();
+		await response.StartAsync(cancellationToken);
-		await source.CopyToAsync(response.Body);
+		await source.CopyToAsync(response.Body, cancellationToken);
 	}
 }