spark/src/DataStreamSuite/Bts.DataStreamWebApp.Server/DataStreamWebSocketServer.cs

namespace Bts.DataStreamWebApp.Server.Controllers;

using bts.websocket_core.json_rpc;

using Bts.DataStream;
using Bts.DataStreamWebApp.Server.Messages;

using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;

public class DataStreamWebSocketServer : IDisposable, IAsyncDisposable {
   readonly ILogger logger;
   readonly WebSocket webSocket;
   readonly Dictionary<string, DataStreamSubscription> subscriptions = new();
   readonly IDataStreamFactory streamFactory;

   public DataStreamWebSocketServer(WebSocket webSocket, IDataStreamFactory streamFactory, ILogger logger) {
      this.webSocket = webSocket;
      this.streamFactory = streamFactory;
      this.logger = logger;
   }

   public async Task<int> Run(CancellationToken cancellationToken) {
      logger.LogInformation("DataStreamWebSocketServer running");
      await foreach (var req in JsonRpcRequestAsyncEnumerable.Create(webSocket, logger, cancellationToken)) {
         await (req.Method switch {
            "subscribe" => onSubscribe(req, cancellationToken),
            "unsubscribe" => onUnsubscribe(req.Deserialize<UnsubscribeRequest>(), cancellationToken),
            _ => Task.CompletedTask,
         });
      }
      logger.LogInformation("DataStreamWebSocketServer stopped");
      return 0;
   }

   /// <summary>
   /// Disposes internal subscriptions. To additionally notify the client of the subscription status
   /// change(s), use <see cref="DisposeAsync"/>.
   /// </summary>
   public void Dispose() {
      GC.SuppressFinalize(this);
      disposeSubscriptions();
   }

   /// <summary>
   /// Disposes internal subscriptions (like <see cref="Dispose"/>), then sends 'subscription status change'
   /// notifications to the client (unlike <see cref="Dispose"/>).
   /// </summary>
   public async ValueTask DisposeAsync() {
      GC.SuppressFinalize(this);
      var finalSubs = subscriptions.Values.Select(s => (s.RequestId, s.Id)).ToList();
      disposeSubscriptions();
      // Send a 'subscription status change' msg to the client for each one.
      foreach (var (requestId, id) in finalSubs) {
         await sendSubscriptionStatus(requestId, id, "unsubscribe", SubscriptionStatus.Unsubscribed, CancellationToken.None);
      }
   }

   async Task onSubscribe(JsonRpcRequest<JsonNode> req, CancellationToken cancellationToken) {
      ArgumentNullException.ThrowIfNull(req.Id);
      ArgumentException.ThrowIfNullOrEmpty(req.Method);
      if (!SubscribeRequest.Parse(req, out var subReq, out var parseErr, out var errCode)) {
         logger.LogError("Invalid 'subscribe' request: {Reason}", parseErr);
         var rpcError = JsonRpcResponse.CreateError<SubscriptionError>(
            req.Id, new() { Code = (long)errCode, Message = "Failed to parse subscription request", Data = parseErr });
         await send(rpcError, cancellationToken);
         return;
      }
      var subId = subReq.SubscriptionId ?? throw new InvalidOperationException("SubscriptionId is null (parse false)");
      disposeSubscription(subId);
      if (!subReq.ParseQuery(streamFactory, out var vs, out parseErr, out errCode)) {
         logger.LogError("Invalid 'subscribe' request: {Reason}", parseErr);
         var rpcError = JsonRpcResponse.CreateError<SubscriptionError>(
            req.Id, new() { Code = (long)errCode, Message = "Failed to parse subscription request", Data = parseErr });
         await send(rpcError, cancellationToken);
         return;
      }
      if (subscriptions.TryGetValue(subId, out var oldSub)) {
         logger.LogInformation("Disposing existing subscription with id {SubscriptionId}", subId);
         oldSub.Dispose();
      }
      var sub = new DataStreamSubscription(subId, req.Id.Value, vs, webSocket, logger);
      subscriptions[subId] = sub;
      await sendSubscriptionStatus(req.Id.Value, subId, req.Method, SubscriptionStatus.Subscribed, cancellationToken);
      vs.ErrorChanged += OnStreamError;
      await sub.StartAsync(cancellationToken);

      async void OnStreamError(Exception exception) {
         logger.LogError("Error in ViewStream (disposing subscription): {Exception}", exception);
         subscriptions.Remove(sub.Id);
         sub.Dispose();
         var error = new SubscriptionError() {
            ErrorMessage = "Error in ViewStream (disposing subscription)"
         };
         if (exception is not null) {
            error.ErrorMessage += ". " + exception.Message;
            error.StackTrace = exception.StackTrace?.Split("\n")?.ToList();
         }
         var rpcError = JsonRpcResponse.CreateError<SubscriptionError>(
            req.Id, new() { Code = (long)ErrorCode.ViewStreamError, Message = "Error in ViewStream (disposing subscription)", Data = error });
         await send(rpcError, cancellationToken);
      }
   }

   async Task onUnsubscribe(JsonRpcRequest<UnsubscribeRequest> req, CancellationToken cancellationToken) {
      ArgumentNullException.ThrowIfNull(req.Id);
      ArgumentException.ThrowIfNullOrEmpty(req.Method);
      var subId = req.Params?.SubscriptionId ?? string.Empty;
      if (!disposeSubscription(subId)) {
         logger.LogInformation("Unable to dispose subscription. Subscription with id {Id} not found.", subId);
      }
      await sendSubscriptionStatus(req.Id.Value, subId, "unsubscribe", SubscriptionStatus.Unsubscribed, cancellationToken);
   }

   bool disposeSubscription(string subId) {
      if (!subscriptions.TryGetValue(subId, out var sub)) {
         return false;
      }
      logger.LogInformation("Disposing subscription with id {Id}", subId);
      sub.Dispose();
      subscriptions.Remove(subId);
      return true;
   }

   Task sendSubscriptionStatus(JsonRpcId reqId, string subId, string method, SubscriptionStatus status, CancellationToken cancellationToken) {
      logger.LogInformation("Sending {Method} ack. SubscriptionId = {SubscriptionId}", method, subId);
      var response = JsonRpcResponse.Create(reqId, new SubscriptionStatusMessage() {
         SubscriptionId = subId,
         Status = status,
      });
      return send(response, cancellationToken);
   }

   Task send<T>(T message, CancellationToken cancellationToken) {
      if (message is null) {
         return Task.CompletedTask;
      }
      var str = JsonSerializer.Serialize(message);
      var bytes = Encoding.UTF8.GetBytes(str);
      logger.LogTrace("Sending data. Length = {Length}", bytes.Length);
      return webSocket.SendAsync(bytes, WebSocketMessageType.Text, true, cancellationToken);
   }

   void disposeSubscriptions() {
      logger.LogInformation("Disposing DataStreamWebSocketServer subscriptions. Count: {Count}", subscriptions.Count);
      subscriptions.Values.DisposeAll();
      subscriptions.Clear();
   }
}