mcdiscordbot/MinecraftDiscordBot/ChunkWaiter.cs

using MinecraftDiscordBot.Models;
using MinecraftDiscordBot.Services;

namespace MinecraftDiscordBot;

public class ChunkWaiter<T> : IChunkWaiter {
    public int ID { get; }
    private readonly CancellationToken _ct;
    public ChunkWaiter(int id, Func<string, T> resultParser, CancellationToken ct) {
        ID = id;
        this.resultParser = resultParser;
        _ct = ct;
    }
    private readonly TaskCompletionSource<T> tcs = new();
    private readonly Func<string, T> resultParser;
    public Task<T> Task => tcs.Task.WaitAsync(_ct);
    public bool Finished { get; private set; } = false;
    public bool IsCancellationRequested => _ct.IsCancellationRequested;
    private string?[]? _chunks = null;
    private int _receivedChunks = 0;
    private ResultState? _state = null;
    private readonly object _syncRoot = new();
    public void AddChunk(int chunkId, int totalChunks, string value) {
        lock (_syncRoot) {
            if (_chunks is null) _chunks = new string[totalChunks];
            else if (_chunks.Length != totalChunks) {
                Program.LogErrorAsync(Program.WebSocketSource, new InvalidOperationException("Different numbers of chunks in same message ID!"));
                return;
            }
            ref string? chunk = ref _chunks[chunkId - 1]; // Lua 1-indexed
            if (chunk is not null) {
                Program.LogErrorAsync(Program.WebSocketSource, new InvalidOperationException($"Chunk with ID {chunkId} was already received!"));
                return;
            }
            chunk = value;
        }
        if (++_receivedChunks == totalChunks) FinalizeResult(_chunks);
    }
    private void FinalizeResult(string?[] _chunks) {
        var resultString = string.Concat(_chunks);
        switch (_state) {
        case ResultState.Successful: tcs.SetResult(resultParser(resultString)); break;
        case ResultState.Unsuccessful: tcs.SetException(new ReplyException(resultString)); break;
        case ResultState.Fatal: tcs.SetException(new InvalidProgramException($"Client script failed: {resultString}")); break;
        default: throw new InvalidProgramException($"Program cannot handle result state '{_state}'!");
        }
        Finished = true;
    }
    public void SetResultState(ResultState state) => _state = _state is ResultState oldState && state != oldState
    ? throw new InvalidOperationException("Cannot set two different result states for same message!")
    : state;
}