vfxreview/RenderWorker/VFXReviewWorker.Tests/MultipartUploadTests.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Text.Json;
using Microsoft.Extensions.Logging.Abstractions;
using VFXReviewWorker;
using Xunit;

namespace VFXReviewWorker.Tests;

/// <summary>
/// Large artifacts go up as parallel multipart parts: from the studio a single
/// connection to object storage is capped near 0.1 MB/s, so a ProRes MOV could
/// not finish inside the 30 min request timeout. These drive the real upload
/// code against a local fake of the server + storage.
/// </summary>
public class MultipartUploadTests : IDisposable
{
    private const int PartSize = 8 * 1024 * 1024;

    private readonly HttpListener _listener = new();
    private readonly string _base;
    private readonly ConcurrentDictionary<int, byte[]> _parts = new();
    private readonly ConcurrentDictionary<int, int> _failuresLeft = new();
    private readonly ConcurrentBag<string> _completions = new();
    private int _inFlight, _maxInFlight;
    private bool _multipartSupported = true;
    private byte[]? _singlePut;

    public MultipartUploadTests()
    {
        var port = FreePort();
        _base = $"http://localhost:{port}/";
        _listener.Prefixes.Add(_base);
        _listener.Start();
        _ = Task.Run(Serve);
    }

    private static int FreePort()
    {
        var l = new TcpListener(IPAddress.Loopback, 0);
        l.Start();
        var port = ((IPEndPoint)l.LocalEndpoint).Port;
        l.Stop();
        return port;
    }

    private async Task Serve()
    {
        while (_listener.IsListening)
        {
            HttpListenerContext ctx;
            try { ctx = await _listener.GetContextAsync(); } catch { return; }
            _ = Task.Run(() => Handle(ctx));
        }
    }

    private async Task Handle(HttpListenerContext ctx)
    {
        var path = ctx.Request.Url!.AbsolutePath;
        using var ms = new MemoryStream();
        await ctx.Request.InputStream.CopyToAsync(ms);
        var body = ms.ToArray();
        var res = ctx.Response;

        if (path.EndsWith("/artifact-presign"))
        {
            var req = JsonDocument.Parse(body).RootElement;
            var size = req.GetProperty("size").GetInt64();
            object payload = _multipartSupported && req.GetProperty("multipart").GetBoolean()
                ? new
                {
                    key = "highres/x.mov",
                    uploadId = "up1",
                    partSize = PartSize,
                    partUrls = Enumerable.Range(1, (int)((size + PartSize - 1) / PartSize))
                        .Select(n => $"{_base}part/{n}").ToList(),
                }
                : new { key = "highres/x.mov", presignedUrl = $"{_base}single" };
            await Json(res, payload);
        }
        else if (path.StartsWith("/part/"))
        {
            var n = int.Parse(path["/part/".Length..]);
            var now = Interlocked.Increment(ref _inFlight);
            InterlockedMax(ref _maxInFlight, now);
            await Task.Delay(50);
            Interlocked.Decrement(ref _inFlight);
            if (_failuresLeft.TryGetValue(n, out var left) && left > 0)
            {
                _failuresLeft[n] = left - 1;
                res.StatusCode = 500;
                res.Close();
                return;
            }
            _parts[n] = body;
            res.Headers["ETag"] = $"\"etag-{n}\"";
            res.Close();
        }
        else if (path == "/single")
        {
            _singlePut = body;
            res.Close();
        }
        else if (path.EndsWith("/artifact-complete"))
        {
            _completions.Add(System.Text.Encoding.UTF8.GetString(body));
            await Json(res, new { ok = true });
        }
        else
        {
            res.StatusCode = 404;
            res.Close();
        }
    }

    private static void InterlockedMax(ref int target, int value)
    {
        int seen;
        while ((seen = Volatile.Read(ref target)) < value && Interlocked.CompareExchange(ref target, value, seen) != seen) { }
    }

    private static async Task Json(HttpListenerResponse res, object payload)
    {
        var bytes = JsonSerializer.SerializeToUtf8Bytes(payload, WorkerOptions.JsonOpts);
        res.ContentType = "application/json";
        await res.OutputStream.WriteAsync(bytes);
        res.Close();
    }

    private ApiClient Client(int concurrency = 4) =>
        new(new WorkerOptions { ServerUrl = _base, ApiKey = "k", UploadConcurrency = concurrency },
            NullLogger<ApiClient>.Instance);

    private static string TempFile(int bytes)
    {
        var path = Path.GetTempFileName();
        var data = new byte[bytes];
        new Random(42).NextBytes(data);
        File.WriteAllBytes(path, data);
        return path;
    }

    [Fact]
    public async Task Uploads_parts_in_parallel_and_stitches_them_in_order()
    {
        var file = TempFile(PartSize * 5 + 12345);
        long lastProgress = 0;
        var key = await Client(4).UploadArtifactAsync("job1", "m1", "highres", file, "video/quicktime",
            CancellationToken.None, sent => Interlocked.Exchange(ref lastProgress, sent));

        Assert.Equal("highres/x.mov", key);
        var stitched = Enumerable.Range(1, 6).SelectMany(n => _parts[n]).ToArray();
        Assert.Equal(File.ReadAllBytes(file), stitched);
        Assert.True(_maxInFlight > 1, "parts should upload concurrently");
        Assert.True(_maxInFlight <= 4, "concurrency is capped by UploadConcurrency");
        Assert.Equal(new FileInfo(file).Length, Interlocked.Read(ref lastProgress));

        var complete = JsonDocument.Parse(Assert.Single(_completions)).RootElement;
        Assert.False(complete.TryGetProperty("abort", out _));
        Assert.Equal(6, complete.GetProperty("parts").GetArrayLength());
        Assert.Equal("\"etag-1\"", complete.GetProperty("parts")[0].GetProperty("etag").GetString());
        File.Delete(file);
    }

    [Fact]
    public async Task A_failed_part_is_retried_on_its_own()
    {
        var file = TempFile(PartSize * 3);
        _failuresLeft[2] = 2;
        var key = await Client().UploadArtifactAsync("job1", "m1", "highres", file, "video/quicktime", CancellationToken.None);

        Assert.NotNull(key);
        Assert.Equal(File.ReadAllBytes(file), Enumerable.Range(1, 3).SelectMany(n => _parts[n]).ToArray());
        File.Delete(file);
    }

    [Fact]
    public async Task A_part_that_keeps_failing_aborts_the_upload()
    {
        var file = TempFile(PartSize * 3);
        _failuresLeft[3] = 99;
        var key = await Client().UploadArtifactAsync("job1", "m1", "highres", file, "video/quicktime", CancellationToken.None);

        Assert.Null(key);
        var abort = JsonDocument.Parse(Assert.Single(_completions)).RootElement;
        Assert.True(abort.GetProperty("abort").GetBoolean());
        File.Delete(file);
    }

    [Fact]
    public async Task Falls_back_to_a_single_put_on_a_server_without_multipart()
    {
        _multipartSupported = false;
        var file = TempFile(PartSize * 3);
        var key = await Client().UploadArtifactAsync("job1", "m1", "highres", file, "video/quicktime", CancellationToken.None);

        Assert.Equal("highres/x.mov", key);
        Assert.Equal(File.ReadAllBytes(file), _singlePut);
        Assert.Empty(_completions);
        File.Delete(file);
    }

    public void Dispose() => _listener.Close();
}