rudder-unity-sdk/Packages/rudder.sdk/Runtime/Sources/Realtime/NativeWebSocketAdapter.cs

using System;
using System.Collections.Concurrent;
using System.IO;
using System.Net.WebSockets;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;

namespace RudderSdk.Unity
{
    /// <summary>
    /// Native WebSocket adapter using System.Net.WebSockets.ClientWebSocket.
    /// Dispatches received events to the Unity main thread via a MonoBehaviour Update loop.
    /// </summary>
    public sealed class NativeWebSocketAdapter : IWebSocketAdapter, IDisposable
    {
        private ClientWebSocket _ws;
        private CancellationTokenSource _cts;
        private readonly WebSocketDispatcher _dispatcher;
        private readonly ConcurrentQueue<Action> _eventQueue = new();
        private bool _disposed;

        /// <summary>
        /// Fired when the WebSocket connection is closed.
        /// </summary>
        public event Action Closed;

        /// <summary>
        /// Fired when binary data is received from the server.
        /// </summary>
        public event Action<ArraySegment<byte>> Received;

        /// <summary>
        /// Fired when an error occurs on the receive loop.
        /// </summary>
        public event Action<Exception> ReceivedError;

        /// <summary>
        /// True when the WebSocket is in the Open state.
        /// </summary>
        public bool IsConnected => _ws?.State == WebSocketState.Open;

        /// <summary>
        /// True when the WebSocket is in the Connecting state.
        /// </summary>
        public bool IsConnecting => _ws?.State == WebSocketState.Connecting;

        public NativeWebSocketAdapter()
        {
            var go = new GameObject("WebSocketDispatcher");
            UnityEngine.Object.DontDestroyOnLoad(go);
            _dispatcher = go.AddComponent<WebSocketDispatcher>();
            _dispatcher.Init(this);
        }

        /// <summary>
        /// Connects to the specified WebSocket URI.
        /// </summary>
        /// <param name="uri">The WebSocket server URI (ws:// or wss://).</param>
        /// <param name="timeoutSeconds">Connection timeout in seconds.</param>
        /// <param name="cancellationToken">Optional cancellation token.</param>
        /// <exception cref="ArgumentNullException">Thrown when uri is null.</exception>
        /// <exception cref="ArgumentOutOfRangeException">Thrown when timeoutSeconds is not positive.</exception>
        /// <exception cref="TimeoutException">Thrown when the connection attempt times out.</exception>
        public async Task ConnectAsync(Uri uri, int timeoutSeconds, CancellationToken cancellationToken = default)
        {
            if (uri == null)
                throw new ArgumentNullException(nameof(uri));
            if (timeoutSeconds <= 0)
                throw new ArgumentOutOfRangeException(nameof(timeoutSeconds), "Timeout must be greater than zero.");

            _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            _ws = new ClientWebSocket();
            _ws.Options.KeepAliveInterval = TimeSpan.FromSeconds(30);

            using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(timeoutSeconds));
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, timeoutCts.Token);

            try
            {
                await _ws.ConnectAsync(uri, linked.Token);
                _ = ReceiveLoop();
            }
            catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
            {
                throw new TimeoutException($"WebSocket connection timed out after {timeoutSeconds} seconds.");
            }
        }

        /// <summary>
        /// Gracefully closes the WebSocket connection.
        /// </summary>
        public async Task CloseAsync()
        {
            if (_ws?.State == WebSocketState.Open)
            {
                try
                {
                    await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "Client closing", CancellationToken.None);
                }
                catch
                {
                    // Best-effort close
                }
            }

            _cts?.Cancel();
            EnqueueEvent(() => Closed?.Invoke());
        }

        /// <summary>
        /// Sends binary data over the WebSocket. WebSocket transport is always reliable,
        /// so no reliability parameter is needed.
        /// </summary>
        public async Task SendAsync(ArraySegment<byte> data, CancellationToken cancellationToken = default)
        {
            if (_ws?.State != WebSocketState.Open)
                return;

            await _ws.SendAsync(data, WebSocketMessageType.Binary, true, cancellationToken);
        }

        /// <summary>
        /// Disposes the adapter, cancelling any in-flight operations and cleaning up the dispatcher GameObject.
        /// </summary>
        public void Dispose()
        {
            if (_disposed) return;
            _disposed = true;

            _cts?.Cancel();
            _cts?.Dispose();
            _ws?.Dispose();

            if (_dispatcher != null && _dispatcher.gameObject != null)
            {
                UnityEngine.Object.Destroy(_dispatcher.gameObject);
            }
        }

        private async Task ReceiveLoop()
        {
            var buffer = new byte[65536];
            try
            {
                while (!_cts.IsCancellationRequested && _ws.State == WebSocketState.Open)
                {
                    using (var ms = new MemoryStream())
                    {
                        WebSocketReceiveResult result;
                        do
                        {
                            result = await _ws.ReceiveAsync(new ArraySegment<byte>(buffer), _cts.Token);
                            ms.Write(buffer, 0, result.Count);
                        } while (!result.EndOfMessage);

                        if (result.MessageType == WebSocketMessageType.Close)
                        {
                            try
                            {
                                await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
                            }
                            catch
                            {
                                // Best-effort close
                            }

                            EnqueueEvent(() => Closed?.Invoke());
                            return;
                        }

                        var data = ms.ToArray();
                        EnqueueEvent(() => Received?.Invoke(new ArraySegment<byte>(data)));
                    }
                }
            }
            catch (OperationCanceledException)
            {
                // Normal cancellation, no action needed
            }
            catch (Exception ex) when (!_disposed)
            {
                EnqueueEvent(() => ReceivedError?.Invoke(ex));
                EnqueueEvent(() => Closed?.Invoke());
            }
        }

        internal void EnqueueEvent(Action action)
        {
            _eventQueue.Enqueue(action);
        }

        internal void ProcessEvents()
        {
            while (_eventQueue.TryDequeue(out var action))
            {
                try
                {
                    action?.Invoke();
                }
                catch (Exception ex)
                {
                    Debug.LogException(ex);
                }
            }
        }
    }

    /// <summary>
    /// MonoBehaviour used to dispatch WebSocket events on the Unity main thread via Update.
    /// Created and managed by NativeWebSocketAdapter.
    /// </summary>
    internal class WebSocketDispatcher : MonoBehaviour
    {
        private NativeWebSocketAdapter _adapter;

        public void Init(NativeWebSocketAdapter adapter)
        {
            _adapter = adapter;
        }

        private void Update()
        {
            _adapter?.ProcessEvents();
        }

        private void OnDestroy()
        {
            _adapter?.Dispose();
            _adapter = null;
        }
    }
}