InetLab

Code: Select all

    public static class SmppClientExtension
    {
        private static ConcurrentDictionary<SmppClient, BoundWaiter> _boundWaiters = new ConcurrentDictionary<SmppClient, BoundWaiter>();

        class BoundWaiter
        {
            public TaskCompletionSource<bool> TaskSource;
            public CancellationTokenRegistration CancelRegistration;
        }

        public static Task WaitForBound(this SmppClient client)
        {
            return WaitForBound(client, CancellationToken.None);
        }
        public static Task WaitForBound(this SmppClient client, CancellationToken cancellationToken)
        {
            if (client.Status == ConnectionStatus.Bound) return Task.FromResult(true);

           
            if (client.ConnectionRecovery)
            {
                lock (_boundWaiters)
                {
                    return _boundWaiters.GetOrAdd(client, c =>
                    {
                        BoundWaiter waiter = new BoundWaiter();
                        waiter.TaskSource =
                            new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

                        waiter.CancelRegistration = cancellationToken.Register(c1 =>
                        {
                            var cancelClient = (SmppClient) c1;

                            BoundWaiter w;
                            if (_boundWaiters.TryRemove(cancelClient, out w))
                            {
                                cancelClient.evRecoverySucceeded -= OnRecoverySucceeded;
                                w.TaskSource.TrySetCanceled(cancellationToken);
                            }


                        }, c);



                        c.evRecoverySucceeded += OnRecoverySucceeded;

                        return waiter;

                    }).TaskSource.Task;

                }
            }

            return Task.FromCanceled(cancellationToken);

        }

        static void OnRecoverySucceeded(object sender, BindResp data)
        {
            SmppClient recoveryClient = (SmppClient) sender;

            BoundWaiter w;
            if (_boundWaiters.TryRemove(recoveryClient, out w))
            {
                recoveryClient.evRecoverySucceeded -= OnRecoverySucceeded;
                w.CancelRegistration.Dispose();

                w.TaskSource.TrySetResult(true);
            }
        }

    }