Implementing custom IValueTaskSource – async without allocations

This blog post was written as part of the preparations while writing the book about .NET, which will be announced in a few weeks. If you want to be informed about its publication and receive auxiliary materials, feel free to subscribe to my newsletter. Many thanks to Stephen Toub that helped in reviewing this code.

Async programming becomes more and more popular. While being very convenient in use, from performance perspective there are scenarios where regular Task-returning async methods have one serious drawback: they need to allocate a new Task to represent the operation (and its result). Such heap-allocated Task is unavoidable in the truly asynchronous path of execution because async continuations are not guaranteed to be executed on the same thread – thus the operation must persist on the heap, not on the stack.

However, there are cases where async operations may complete synchronously (because of really fast meeting some conditions). It would be nice to avoid heap-allocating Task in such case, created just to pass the result of an operation. Exactly for such purpose ValueTask type was introduced in .NET Core 2.0 (and the corresponding AsyncValueTaskMethodBuilder handling the underlying state machine). Initially, it was a struct made as a discriminated union, which could take one of two possible values:

  • ready to use the result (if the operation completed successfully synchronously)
  • a normal Task which may be awaited on (if the operation become truly asynchronous)

In other words, ValueTask helps in handling synchronous path of async method execution. Thanks to being a struct (which will be allocated on the stack or enregistered into CPU register) synchronous result of the operation can be returned without heap allocations. And only in case of an asynchronous path, a new Task will be heap-allocated eventually by underlying machinery:

But what if we were able to not heap-allocate Task even in case of an asynchronous path? This may be useful in very, very high-performance code, avoiding async-related allocations at all. As already said, something on the heap must represent our async operation because there is no thread affinity. But why not use heap-allocated, pooled objects for that, reused between successive async operations?

Indeed, since .NET Core 2.1, ValueTask can also wrap an object implementing IValueTaskSource interface. Such an object can be pooled and reused to minimize allocations. It represents our operation and underlying AsyncValueTaskMethodBuilder is aware of it, calling appropriate methods of IValueTaskSource interface. Here is how the previous example looks with the help of custom IValueTaskSource implementation described in this post:

Thanks to pooling, even on the asynchronous path there will be no allocations (at least most often, if our pool is used efficiently). We can await such method in a regular way and underlying machinery will take care of it.

Note. If you would like to hear more about Task, ValueTask and IValueTaskSource again, in similar but other words, please look at great Task, Async Await, ValueTask, IValueTaskSource and how to keep your sanity in modern .NET world post by Szymon Kulec.

Implementation details

Although the rationale behind IValueTaskSource seems to be clear, as well as its usage presented above, implementing it is not trivial. When implementing IValueTaskSource interface, we must implement three following methods:
* GetResult – called only once, when the async state machine needs to obtain the result of the operation
* GetStatus – called by the async state machine to check the status of the operation
* OnCompleted – called by the async state machine when wrapping ValueTask has been awaited. We should remember here the continuation to be called when the operation completes (but if it already has been completed, we should call the continuation immediately)

Seems easy, right? Read on to see if it really is! All source code described here is available in my PooledValueTaskSource repository on GitHub. There are quite many comments in the code but this post explains most relevant parts as well. Do not be also surprised with many diagnostic Console.Write in this code – it serves to illustrate the internal working of this class in prepared example program (also available in the repository).

In my custom implementation, I use object pooling based on ObjectPool class based on the internal class taken from Roslyn source code and a little refactored (with renaming mostly) – I’ve omitted it here for brevity as not so relevant. From our perspective here there are obvious Rent and Return methods, period.

In my implementation, I am also mostly based on code from AwaitableSocketAsyncEventArgs in System.Net.Sockets.Socket
and AsyncIOOperation in ASP.NET IIS Integration code. What I’ve tried to do is to provide Minimal Valuable Product that is correct and working (stripping as much as possible from the mentioned code).

My custom IValueTaskSource represents an operation that returns a string that is being read from the provided file. Obviously, one would probably like to introduce a more generic class with generic result type and action being provided as a lambda expression. However, to not clutter such example too much, I’ve decided to prepare it in such “hardcoded”, specific scenario. Feel free to contribute more generic versions!

Let’s start from fields that FileReadingPooledValueTaskSource contains:

The most important fields of FileReadingPooledValueTaskSource include:

  • Action< object > continuation – it represents a continuation to be executed when our operation ends
  • string result – it keeps the result of our operation (in case of success)
  • Exception exception – it keeps an Exception instance that happened during executing our operation (in case of failure)
  • short token – current token value given to a ValueTask and then verified against the value it passes back to us. This is not meant to be a completely reliable mechanism, doesn’t require additional synchronization, etc. It’s purely a best effort attempt to catch misuse, including awaiting for a value task twice and after it’s already being reused by someone else
  • object state – state internally used by asynchronous machinery
  • static readonly Action<object> CallbackCompleted – sentinel object used to indicate that the operation has completed prior to OnCompleted being called

Let’s now look at each of IValueTaskSource method implementation. GetResult is quite easy – it will be called only once by underlying state machine when we inform that our operation has completed (by GetStatus method explained soon). Thus, we need to reset the object state (to be reusable), return it to the pool and return the result (or throw an exception in case of failure):

GetStatus is called by the state machine to check the current status of our operation. In my case, I assume it is completed if result is no more null. Depends on the exception field, it is then succeeded or faulted:

The most complex is OnCompleted method implementation. It is being called by the underlying state machine if wrapped ValueTask is being awaited. Two scenarios may happen here that must be handled:

  • if an operation has not yet completed – we store the provided continuation to be executed once the operation is completed
  • if an operation has already completed – in such case our internal continuation should be already set to CallbackCompleted value. If it so, we simply invoke the continuation here

Please note how much code is dedicated to properly get the context of the continuation (with respect to provided ValueTaskSourceOnCompletedFlags):

This concludes implementing IValueTaskSource methods but we need to add two more crucial pieces into this puzzle: a method that starts our operation and method that is called when an operation completes asynchronously.

The first one, named by my as simple as RunAsync (called in our example at the beginning of the article) is responsible for executing the main work:

I’ve implemented here simulation of some operation that may both complete immediately (synchronously) and asynchronously. In case of asynchronous path, returned ValueTask pass the result so we avoid allocations again. In case of asynchronous case the key is to return ValueTask that wraps… ourselves. It may be then awaited on, while we also started asynchronous processing (simulated by thread pool work in our case).

When asynchronous operation finishes, NotifyAsyncWorkCompletion method will be called (remember – in the real-world scenario this would be some callback registered in asynchronous IO or other low-level API). The responsibility of this method is simple:

  • it stores result and/or exception
  • if the operation is not yet awaited (in such case this.continuation will be null) – it only sets continuation to CallbackCompleted. Continuation will be executed in OnCompleted method when ValueTask will be awaited
  • if the operation is already awaited (in such case this.continuation contains awaited continuation) – it executes continuation in the appropriate context (which again is quite a complex process)

 

Above-mentioned clearing and returning to the pool is implemented in ResetAndReleaseOperation (yes I know, SRP is dying here, refactor!). The only field we cannot clear is token, which is solely dedicated to detecting incorrect re-usage of those objects:

And… only such little code is necessary to avoid heap-allocating in case of async operations!

Remarks

  • I do not claim that my code in current form is ideal. Quite opposite, I still expect it to be by far ideal! It serves as an illustration and base for further development. Please, feel invited to comment and to contribute to making it better!
  • Current repository is oversimplified – due to the work on the book, I do not have time to reorganize it properly (especially include comprehensive unit tests). Again, feel free to contribute!

 

6 comments

  1. Pingback: dotnetomaniak.pl
  2. if (exception != null)
    {
    throw exception;
    }

    It’s probably a bad idea since it loses the stack trace. Consider using ExceptionDispatchInfo.Capture(e).Throw();

    Your implementation also makes ValueTask.Result property getter to effectively *mutate* the state by invalidating the ValueTask. People expect to be able to get Result value multiple times and not get exceptions.

    1. Thank you for great exception-related remark!

      Regarding multiple Result usage – You are right. But I believe such custom code will be heavily fine-tuned to specific purpose and such general-purpose assumptions will often not apply (if they complicate things). In other words, such class will be written for deep internal, per-project specific usage and heavily documented (to clearly describe what and how it handles). Due to pooling, I expect such single usage pattern will be quite common.

  3. Hi Konrad,

    I made some benchmarking of Task vs ValueTask vs ValueTaskSource in .NET Core 3.0 based on BenchmarkDotNet. Here are some results:

    Method Mean Error StdDev Rank Allocated
    ReadFileAsync_ValueTask_AsyncMethod 475.9 us 9.496 us 26.47 us 2 190520 B
    ReadFileAsync_ValueTask 450.0 us 14.933 us 44.03 us 1 190547 B
    ReadFileAsync_ValueTask_ValueTaskSource 1,001,244.0 us 315.576 us 295.19 us 4 424 B
    ReadFileAsync_TaskString_AsyncMethod 468.3 us 17.459 us 51.20 us 2 190586 B
    ReadFileAsync_TaskString 510.5 us 16.525 us 48.46 us 3 190672 B

    As you can see the memory allocation for ValueTaskSource is very low but at the same time, it is very slow. In the benchmark I read a file of size 40KB. And here is a code I used (based on you sample):

    private ValueTask ReadFileAsyncWithValueTaskAndValueTaskSource(string filename = “data”)
    {
    var cachedOp = _pool.Rent();
    return cachedOp.RunAsync(filename, _pool);
    }

    Can you comment on that? What might be the reason?

Leave a Reply

Your email address will not be published. Required fields are marked *