CryptoExchange.Net/CryptoExchange.Net/Sockets/Query.cs

using CryptoExchange.Net.Interfaces;
using CryptoExchange.Net.Objects;
using CryptoExchange.Net.Sockets.Default;
using CryptoExchange.Net.Sockets.Default.Routing;
using CryptoExchange.Net.Sockets.Interfaces;
using System;
using System.Diagnostics.CodeAnalysis;
using System.Threading;
using System.Threading.Tasks;

namespace CryptoExchange.Net.Sockets
{
    /// <summary>
    /// Query 
    /// </summary>
    public abstract class Query : IMessageProcessor
    {
        /// <summary>
        /// Unique identifier
        /// </summary>
        public int Id { get; } = ExchangeHelpers.NextId();

        /// <summary>
        /// Has this query been completed
        /// </summary>
        public bool Completed { get; set; }

        /// <summary>
        /// Whether this query completed successfully
        /// </summary>
        [MemberNotNullWhen(false, nameof(Error))]
        public abstract bool Success { get; }

        /// <summary>
        /// Error result for this query
        /// </summary>
        public abstract Error? Error { get; }

        /// <summary>
        /// The original data returned by the query, only available when `OutputOriginalData` is set to `true` in the client options
        /// </summary>
        public abstract string? OriginalData { get; }
        /// <summary>
        /// The request body content
        /// </summary>
        public string? RequestBody { get; set; }

        /// <summary>
        /// Timeout for the request
        /// </summary>
        public TimeSpan? RequestTimeout { get; set; }

        /// <summary>
        /// What should happen if the query times out
        /// </summary>
        public TimeoutBehavior TimeoutBehavior { get; set; } = TimeoutBehavior.Fail;

        /// <summary>
        /// The number of required responses. Can be more than 1 when for example subscribing multiple symbols streams in a single request,
        /// and each symbol receives it's own confirmation response
        /// </summary>
        public int RequiredResponses { get; set; } = 1;

        /// <summary>
        /// The current number of responses received on this query
        /// </summary>
        public int CurrentResponses { get; set; }

        /// <summary>
        /// Timestamp of when the request was send
        /// </summary>
        public DateTime RequestTimestamp { get; set; }
        
        /// <summary>
        /// Response
        /// </summary>
        public object? Response { get; set; }

        private MessageRouter _router;
        /// <summary>
        /// Router for this query
        /// </summary>
        public MessageRouter MessageRouter
        {
            get => _router;
            set
            {
                _router = value;
                _router.BuildQueryRouter();
                OnMessageRouterUpdated?.Invoke();
            }
        }

        /// <summary>
        /// The query request object
        /// </summary>
        public object Request { get; set; }

        /// <summary>
        /// If this is a private request
        /// </summary>
        public bool Authenticated { get; }

        /// <summary>
        /// Weight of the query
        /// </summary>
        public int Weight { get; }

        /// <summary>
        /// Whether the query should wait for a response or not
        /// </summary>
        public bool ExpectsResponse { get; set; } = true;

        /// <summary>
        /// Wait event for response
        /// </summary>
        protected AsyncResetEvent _event;

        /// <summary>
        /// Cancellation token
        /// </summary>
        protected CancellationTokenSource? _cts;

        /// <summary>
        /// On complete callback
        /// </summary>
        public Action<bool>? OnComplete { get; set; }

        /// <inheritdoc />
        public event Action? OnMessageRouterUpdated;

        /// <summary>
        /// ctor
        /// </summary>
#pragma warning disable CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider adding the 'required' modifier or declaring as nullable.
        public Query(object request, bool authenticated, int weight = 1)
#pragma warning restore CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider adding the 'required' modifier or declaring as nullable.
        {
            _event = new AsyncResetEvent(false, false);

            Authenticated = authenticated;
            Request = request;
            Weight = weight;
        }

        /// <summary>
        /// Signal that the request has been send and the timeout timer should start
        /// </summary>
        public void IsSend(TimeSpan timeout)
        {
            RequestTimestamp = DateTime.UtcNow;
            if (ExpectsResponse)
            {
                // Start timeout countdown
                _cts = new CancellationTokenSource(timeout);
                _cts.Token.Register(Timeout, false);
            }
            else
            {
                Complete();
            }
        }

        /// <summary>
        /// Wait until timeout or the request is completed
        /// </summary>
        /// <param name="timeout"></param>
        /// <param name="ct">Cancellation token</param>
        /// <returns></returns>
        public async Task WaitAsync(TimeSpan timeout, CancellationToken ct) => await _event.WaitAsync(timeout, ct).ConfigureAwait(false);

        /// <summary>
        /// Mark request as timeout
        /// </summary>
        public abstract void Timeout();

        /// <summary>
        /// Mark request as successfully completed
        /// </summary>
        public abstract void Complete();

        /// <summary>
        /// Mark request as failed
        /// </summary>
        /// <param name="error"></param>
        public abstract void Fail(Error error);

        /// <summary>
        /// Handle a response message
        /// </summary>
        public abstract bool Handle(string typeIdentifier, string? topicFilter, SocketConnection connection, DateTime receiveTime, string? originalData, object message);

    }

    /// <summary>
    /// Query
    /// </summary>
    /// <typeparam name="THandlerResponse">The type to be returned to the caller</typeparam>
    public abstract class Query<THandlerResponse> : Query
    {
        /// <summary>
        /// Result
        /// </summary>
        public CallResult<THandlerResponse>? Result { get; set; }

        /// <inheritdoc />
        [MemberNotNullWhen(false, nameof(Error))]
        [MemberNotNullWhen(true, nameof(Result))]
        public override bool Success => Result?.Success == true;
        /// <inheritdoc />
        public override Error? Error => Result?.Error;
        /// <inheritdoc />
        public override string? OriginalData => Result?.OriginalData;

        /// <summary>
        /// ctor
        /// </summary>
        /// <param name="request"></param>
        /// <param name="authenticated"></param>
        /// <param name="weight"></param>
        protected Query(
            object request,
            bool authenticated,
            int weight = 1)
            : base(request, authenticated, weight)
        {
        }

        /// <inheritdoc />
        public override bool Handle(string typeIdentifier, string? topicFilter, SocketConnection connection, DateTime receiveTime, string? originalData, object message)
        {
            if (Completed)
                return false;

            CurrentResponses++;
            if (CurrentResponses == RequiredResponses)
                Response = message;

            var handled = false;
            if (Result?.Success != false)
            {
                // If an error result is already set don't override that
                MessageRouter.Handle(typeIdentifier, topicFilter, connection, receiveTime, originalData, message, out var result);
                Result = (CallResult<THandlerResponse>?)result;
                handled = Result != null;
                if (!handled)
                    // Null from Handle means it wasn't actually for this query
                    CurrentResponses -= 1;
            }

            if (CurrentResponses == RequiredResponses)
            {
                Completed = true;
                _event.Set();
                OnComplete?.Invoke(Success);
            }

            return handled;
        }

        /// <inheritdoc />
        public override void Timeout()
        {
            if (Completed)
                return;
                        
            if (TimeoutBehavior == TimeoutBehavior.Fail)
                Result = CallResult<THandlerResponse>.Fail(new TimeoutError());
            else
                Result = CallResult<THandlerResponse>.Ok(default!);

            Completed = true;
            _event.Set();
            OnComplete?.Invoke(Success);
        }

        /// <inheritdoc />
        public override void Complete()
        {
            if (Completed)
                return;

            Result = CallResult<THandlerResponse>.Ok(default!);
            Completed = true;
            _event.Set();
            OnComplete?.Invoke(Success);
        }

        /// <inheritdoc />
        public override void Fail(Error error)
        {
            if (Completed)
                return;

            Result = CallResult<THandlerResponse>.Fail(error);
            Completed = true;

            _event.Set();
            OnComplete?.Invoke(Success);
        }
    }
}