Out of memory exception when using cancellation token

[lee-baker]

Hi,

We have recently added cancellation tokens to our restsharp requests and have been leaking memory since. The stackoverflow question below suggests that the token.Register() call returns a CancellationTokenRegistration that should be disposed when the task is completed.

https://stackoverflow.com/questions/14627226/why-is-my-async-await-with-cancellationtokensource-leaking-memory

Is there any reason this process is not followed in the RestClient.Async class ExecuteTaskAsync method?

[LukeBaulch]

To solve this for my project I have created a new class which inherits the RestClient and overrides the ExecuteTaskAsync methods. These overrides are duplicated from the RestClient source except the following line/s have been replaced.

    // original RestSharp code
    token.Register(() =>
        {
            async.Abort();
            taskCompletionSource.TrySetCanceled();
        });

    // updated code to dispose the registration after task complete
    var registration = token.Register(() =>
        {
            async.Abort();
            taskCompletionSource.TrySetCanceled();
        });
    taskCompletionSource.Task.ContinueWith(t => registration.Dispose());

It would be good if I didn't have to do this, so would be nice if a RestSharp contributor could incorporate this (or equivalent) change.

[hallem]

I'll get this fixed in the next release.

[hallem]

@lee-baker @LukeBaulch i just committed a fix that should address this. can one of you test it out please to make sure things are good?

[lee-baker]

Hi Michael, Thanks for having a go at this. The code that Luke has above shows the CancellationTokenRegistration being disposed, it looks like you are just disposing the Task, which in my understanding would not clean up the token registration. If I can find the time, I will write a little test app and see if your solution cleans it up but for now I am going to keep using Luke's solution in prod.
Thanks

[hallem]

@lee-baker lukes code didn't make sense. it was giving an ambiguous error when i put that line in. can you provide a larger snippit?

[lee-baker]

using System;
using System.Threading;
using System.Threading.Tasks;
using RestSharp;

namespace PPS.PortfolioService.Client
{
    // NOTE: ********************************************************************************************************
    // NOTE: This code was copied directly from the RestSharp source code but modified only slightly
    // NOTE: ********************************************************************************************************

    internal class CancellationTokenFixedRestClient : RestClient
    {
        public CancellationTokenFixedRestClient(string baseUrl)
            : base(baseUrl)
        {
        }

        /// <summary>
        /// Executes the request asynchronously, authenticating if needed
        /// </summary>
        /// <typeparam name="T">Target deserialization type</typeparam>
        /// <param name="request">Request to be executed</param>
        /// <param name="token">The cancellation token</param>
        public override Task<IRestResponse<T>> ExecuteTaskAsync<T>(IRestRequest request, CancellationToken token)
        {
            if (request == null)
            {
                throw new ArgumentNullException("request");
            }

            var taskCompletionSource = new TaskCompletionSource<IRestResponse<T>>();

            try
            {
                var async = ExecuteAsync<T>(request, (response, _) =>
                                                     {
                                                         if (token.IsCancellationRequested)
                                                         {
                                                             taskCompletionSource.TrySetCanceled();
                                                         }
                                                         //Don't run TrySetException, since we should set Error properties and swallow exceptions to be consistent with sync methods
                                                         else
                                                         {
                                                             taskCompletionSource.TrySetResult(response);
                                                         }
                                                     });
                // MODIFICATION: This code was altered from the RestSharp source
                var registration = token.Register(() =>
                                                  {
                                                      async.Abort();
                                                      taskCompletionSource.TrySetCanceled();
                                                  });
                // ReSharper disable once MethodSupportsCancellation
                taskCompletionSource.Task.ContinueWith(t => registration.Dispose());
                // MODIFICATION: End
            }
            catch (Exception ex)
            {
                taskCompletionSource.TrySetException(ex);
            }
            return taskCompletionSource.Task;
        }

        /// <summary>
        /// Executes the request asynchronously, authenticating if needed
        /// </summary>
        /// <param name="request">Request to be executed</param>
        /// <param name="token">The cancellation token</param>
        public override Task<IRestResponse> ExecuteTaskAsync(IRestRequest request, CancellationToken token)
        {
            if (request == null)
            {
                throw new ArgumentNullException("request");
            }

            var taskCompletionSource = new TaskCompletionSource<IRestResponse>();

            try
            {
                var async = this.ExecuteAsync(request, (response, _) =>
                                                       {
                                                           if (token.IsCancellationRequested)
                                                           {
                                                               taskCompletionSource.TrySetCanceled();
                                                           }
                                                           //Don't run TrySetException, since we should set Error properties and swallow exceptions to be consistent with sync methods
                                                           else
                                                           {
                                                               taskCompletionSource.TrySetResult(response);
                                                           }
                                                       });
                // MODIFICATION: This code was altered from the RestSharp source
                var registration = token.Register(() =>
                                                  {
                                                      async.Abort();
                                                      taskCompletionSource.TrySetCanceled();
                                                  });
                // ReSharper disable once MethodSupportsCancellation
                taskCompletionSource.Task.ContinueWith(t => registration.Dispose());
                // MODIFICATION: End
            }
            catch (Exception ex)
            {
                taskCompletionSource.TrySetException(ex);
            }

            return taskCompletionSource.Task;
        }
    }
}

[lee-baker]

That is the exact file I am using in production

[hallem]

thanks, i missed the variable assignment for registration... i was a little cross-eyed

[hallem]

just pushed the changes... let me know if it works

[lee-baker]

Thanks Michael, looks like you are making lots of fixes 👍