ParkingRobot/ClumsyPilot/TrajectoryPlanningVisualization/Runtime/LoopbackVisualizationServer.cs

using System;
using System.Collections.Generic;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace TrajectoryPlanningVisualization;

internal sealed class LoopbackVisualizationServer : IDisposable
{
    private readonly PlanningVisualizationOptionsSnapshot options;
    private readonly PlanningVisualizationStaticSnapshot staticSnapshot;
    private readonly string token;
    private readonly LatestVisualizationFrameStore frames = new LatestVisualizationFrameStore();
    private readonly BoundedCycleHistory history;
    private readonly CancellationTokenSource cancellation = new CancellationTokenSource();
    private readonly object clientsLock = new object();
    private readonly List<SseClientConnection> clients = new List<SseClientConnection>();
    private TcpListener listener;
    private Task acceptTask;
    private Task dispatcherTask;
    private int stopped;
    private long lastProcessedVersion;
    private int clientGeneration;

    public LoopbackVisualizationServer(PlanningVisualizationOptionsSnapshot options,
        PlanningVisualizationStaticSnapshot staticSnapshot, string token)
    {
        this.options = options ?? throw new ArgumentNullException(nameof(options));
        this.staticSnapshot = staticSnapshot ?? throw new ArgumentNullException(nameof(staticSnapshot));
        this.token = token ?? throw new ArgumentNullException(nameof(token));
        history = new BoundedCycleHistory(options.HistoryCycleLimit);
    }

    public int Port { get; private set; }
    public string FaultReason { get; private set; }
    public bool IsRunning => Volatile.Read(ref stopped) == 0 && listener != null;

    public void Start()
    {
        listener = new TcpListener(IPAddress.Loopback, options.Port);
        listener.Start();
        Port = ((IPEndPoint)listener.LocalEndpoint).Port;
        acceptTask = Task.Run((Func<Task>)AcceptLoop);
        dispatcherTask = Task.Run((Func<Task>)DispatchLoop);
    }

    public void Publish(PlanningVisualizationDynamicSnapshot snapshot)
    {
        if (snapshot == null)
        {
            throw new ArgumentNullException(nameof(snapshot));
        }

        frames.Publish(snapshot);
    }

    public void Stop()
    {
        if (Interlocked.Exchange(ref stopped, 1) != 0)
        {
            return;
        }

        SseClientConnection[] currentClients = SnapshotClients();
        const string terminal = "event: end\ndata: {\"reasonChinese\":\"会话已结束\"}\n\n";
        foreach (SseClientConnection client in currentClients)
        {
            client.OfferTerminal(terminal);
        }

        DateTime deadline = DateTime.UtcNow.AddMilliseconds(100);
        foreach (SseClientConnection client in currentClients)
        {
            TimeSpan remaining = deadline - DateTime.UtcNow;
            if (remaining <= TimeSpan.Zero)
            {
                break;
            }

            client.WaitForCompletion(remaining);
        }

        cancellation.Cancel();
        try
        {
            listener?.Stop();
        }
        catch (SocketException)
        {
        }

        foreach (SseClientConnection client in currentClients)
        {
            client.Dispose();
        }
    }

    public void Dispose()
    {
        Stop();
        cancellation.Dispose();
    }

    private async Task AcceptLoop()
    {
        try
        {
            while (!cancellation.IsCancellationRequested)
            {
                TcpClient client;
                try
                {
                    client = await listener.AcceptTcpClientAsync().ConfigureAwait(false);
                }
                catch (ObjectDisposedException)
                {
                    return;
                }
                catch (SocketException)
                {
                    if (cancellation.IsCancellationRequested)
                    {
                        return;
                    }

                    throw;
                }

                _ = Task.Run(() => HandleClient(client));
            }
        }
        catch (Exception exception)
        {
            Fault(exception);
        }
    }

    private async Task HandleClient(TcpClient client)
    {
        SseClientConnection sseClient = null;
        try
        {
            NetworkStream stream = client.GetStream();
            if (!LoopbackHttpRequestReader.TryRead(stream, out LoopbackHttpRequest request))
            {
                WriteResponse(stream, 400, "Bad Request", "text/plain; charset=utf-8", "bad request");
                return;
            }

            if (request.Method != "GET")
            {
                WriteResponse(stream, 405, "Method Not Allowed", "text/plain; charset=utf-8", "method not allowed");
                return;
            }

            if (!ConstantTimeEquals(token, request.Token))
            {
                WriteResponse(stream, 403, "Forbidden", "text/plain; charset=utf-8", "forbidden");
                return;
            }

            if (request.Path == "/api/bootstrap")
            {
                WriteResponse(stream, 200, "OK", "application/json; charset=utf-8", VisualizationJson.Serialize(staticSnapshot));
                return;
            }

            if (request.Path != "/api/events")
            {
                WriteResponse(stream, 404, "Not Found", "text/plain; charset=utf-8", "not found");
                return;
            }

            if (!TryAddClient(client, out sseClient))
            {
                WriteResponse(stream, 503, "Service Unavailable", "text/plain; charset=utf-8", "too many clients");
                return;
            }

            WriteSseHeaders(stream);
            await sseClient.Completion.ConfigureAwait(false);
        }
        catch (SocketException)
        {
        }
        catch (System.IO.IOException)
        {
        }
        catch (ObjectDisposedException)
        {
        }
        finally
        {
            if (sseClient != null)
            {
                RemoveClient(sseClient);
                sseClient.Dispose();
            }
            else
            {
                client.Dispose();
            }
        }
    }

    private async Task DispatchLoop()
    {
        try
        {
            int deliveredGeneration = -1;
            while (!cancellation.IsCancellationRequested)
            {
                await Task.Delay(TimeSpan.FromSeconds(1d / options.RefreshRateHz), cancellation.Token).ConfigureAwait(false);
                if (frames.TryReadAfter(lastProcessedVersion, out VisualizationFrame latest))
                {
                    lastProcessedVersion = latest.Version;
                    history.Add(latest.Snapshot.CycleSummary);
                }

                SseClientConnection[] currentClients = SnapshotClients();
                int currentGeneration = Volatile.Read(ref clientGeneration);
                if (currentClients.Length == 0 || (latest == null && deliveredGeneration == currentGeneration))
                {
                    continue;
                }

                if (latest == null && !frames.TryReadAfter(-1, out latest))
                {
                    continue;
                }

                string payload = "event: frame\ndata: " + VisualizationJson.Serialize(new
                {
                    snapshot = latest.Snapshot,
                    history = history.Snapshot()
                }) + "\n\n";
                foreach (SseClientConnection client in currentClients)
                {
                    client.Offer(payload);
                }

                deliveredGeneration = currentGeneration;
            }
        }
        catch (OperationCanceledException)
        {
        }
        catch (Exception exception)
        {
            Fault(exception);
        }
    }

    private bool TryAddClient(TcpClient tcpClient, out SseClientConnection client)
    {
        lock (clientsLock)
        {
            if (clients.Count >= options.MaximumClients || cancellation.IsCancellationRequested)
            {
                client = null;
                return false;
            }

            client = new SseClientConnection(tcpClient, cancellation.Token);
            clients.Add(client);
            Interlocked.Increment(ref clientGeneration);
            return true;
        }
    }

    private void RemoveClient(SseClientConnection client)
    {
        lock (clientsLock)
        {
            clients.Remove(client);
            Interlocked.Increment(ref clientGeneration);
        }
    }

    private SseClientConnection[] SnapshotClients()
    {
        lock (clientsLock)
        {
            return clients.ToArray();
        }
    }

    private void Fault(Exception exception)
    {
        FaultReason = exception.Message;
        Stop();
    }

    private static bool ConstantTimeEquals(string expected, string actual)
    {
        if (actual == null || expected.Length != actual.Length)
        {
            return false;
        }

        int difference = 0;
        for (int index = 0; index < expected.Length; index++)
        {
            difference |= expected[index] ^ actual[index];
        }

        return difference == 0;
    }

    private static void WriteResponse(NetworkStream stream, int status, string reason, string contentType, string body)
    {
        byte[] bodyBytes = Encoding.UTF8.GetBytes(body);
        string headers = "HTTP/1.1 " + status + " " + reason + "\r\nContent-Type: " + contentType +
            "\r\nContent-Length: " + bodyBytes.Length + "\r\nConnection: close\r\n\r\n";
        byte[] headerBytes = Encoding.ASCII.GetBytes(headers);
        stream.WriteTimeout = 1000;
        stream.Write(headerBytes, 0, headerBytes.Length);
        stream.Write(bodyBytes, 0, bodyBytes.Length);
        stream.Flush();
    }

    private static void WriteSseHeaders(NetworkStream stream)
    {
        const string headers = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream; charset=utf-8\r\nCache-Control: no-cache\r\nConnection: keep-alive\r\n\r\n";
        byte[] bytes = Encoding.ASCII.GetBytes(headers);
        stream.WriteTimeout = 1000;
        stream.Write(bytes, 0, bytes.Length);
        stream.Flush();
    }
}