Task manager object for lazy initialization

Question

The Problem

I am developing a Visual Studio extension for importing code from a server.

Because of the way Visual Studio works, multiple threads could try to perform the same operation at the same time.

Thus, I wanted a class that would ensure that a) each Task would only be performed when needed and b) each Task would only be performed once.

The only exception to b) is if the requesters cancel and then later another request is made in which case if the operation was not finished, it should be retried.

Intended functionality:

1. A TaskManager object is created.
2. At some point, TaskManager.Value is called on the object and the Task is set to run.
3. If TaskManger.Value is called multiple times (possibly on other threads), the other executions are set waiting on the Task that is already in progress.
4. If the CancellationToken indicates a cancellation request, only if all the requestors agree should the Task be cancelled.
5. If the TaskManager object is cancelled before completion and later is requested again, the Task should be reinitialized.

The Code

public class TaskManager<T>
{
    readonly Func<CancellationToken, Task<T>> _initializer;

    public TaskManager(Func<CancellationToken, Task<T>> initializer) => _initializer = initializer;

    Task<T>? _task;
    CancellationTokenSource _cancellationToken = new CancellationTokenSource();

    readonly object _initializationLock = new();
    readonly ConcurrentBag<CancellationToken> _cancellationTokens = new();

    public Task<T> Value (CancellationToken cancellationToken)
    {
        if (cancellationToken != default)
        {
            if (cancellationToken.IsCancellationRequested) return Task.FromResult<T>(default);

            _cancellationTokens.Add(cancellationToken);
            cancellationToken.Register(() =>
            {
                if (_task is not { IsCompleted: true } && _cancellationTokens.All(c => c.IsCancellationRequested))
                    _cancellationToken.Cancel();
            });
        }

        if (_task != null && !_cancellationToken.IsCancellationRequested) return _task;
        lock (_initializationLock)
        {
            if (_task != null) return _task;

            _cancellationToken = new CancellationTokenSource();

            _task = _initializer(_cancellationToken.Token);
        }
        return _task;
    }
}

Example Usage

public static class CodepediaApi
{
    private static readonly RestClient client = new RestClient("https://localhost:5001"); // Test Server
    private static readonly Dictionary<string, TaskManager<SearchResult[]>> responses = new Dictionary<string, TaskManager<SearchResult[]>>();

    public static Task<SearchResult[]> Search(string query, CancellationToken cancellationToken) =>
        responses.TryGetOrAdd(query,
            () => new TaskManager<SearchResult[]>(async c =>
                (await client.GetAsync<SearchResult[]>(new RestRequest("/api/search").AddQueryParameter("q", query), c)) ?? new SearchResult[0]
            )
        ).Value(cancellationToken);
}

where TryGetOrAdd is defined as follows:

public static class Extensions
{
    public static TValue TryGetOrAdd<TKey, TValue>(this IDictionary<TKey, TValue> d, TKey key, Func<TValue> toAdd)
    {
        if (d.TryGetValue(key, out TValue value))
            return value;
        TValue adding = toAdd();
        lock (d)
        {
            if (d.TryGetValue(key, out value))
                return value;
            d.Add(key, adding);
        }
        return adding;
    }
}

Another Example Usage

Here is one of my classes:

public class CodepediaUtil
{
    public readonly Workspace Workspace;
    public readonly SnapshotSpan Range;

    public readonly TaskManager<(Document document, SemanticModel semanticModel, SyntaxNode outerNode, UnresolvedMethodCallInfo unresolvedInvocation)?> CodeAnalysisInfo;
    public readonly TaskManager<SearchResult[]> SearchResults;
    public readonly TaskManager<Dictionary<SearchResult, WikiEntry>> SearchResultsInfo;

    public async Task<UnresolvedMethodCallInfo?> UnresolvedInvocation (CancellationToken cancellationToken)
    {
        var unresolvedInvocationInfo = await CodeAnalysisInfo.Value(cancellationToken);
        if (unresolvedInvocationInfo is not (_, _, _, UnresolvedMethodCallInfo unresolvedInvocation)) return null;

        return unresolvedInvocation;
    }

    public CodepediaUtil (Workspace workspace, SnapshotSpan range)
    {
        Workspace = workspace;
        Range = range;

        CodeAnalysisInfo = new TaskManager<(Document document, SemanticModel semanticModel, SyntaxNode outerNode, UnresolvedMethodCallInfo unresolvedInvocation)?>
        (
            async (CancellationToken cancellationToken) =>
            {
                Document? document = Range.Snapshot.TextBuffer.GetRelatedDocuments().FirstOrDefault();
                if (document == null) return null;

                SemanticModel? semanticModel = await document.GetSemanticModelAsync(cancellationToken);
                if (semanticModel == null) return null;

                SyntaxNode? outerNode = await CodeAnalyzer.GetOuterNodeAsync(semanticModel, Range.Start, cancellationToken);
                if (outerNode == null) return null;

                UnresolvedMethodCallInfo? unresolvedInvocation = outerNode.FindUnresolvedInvocation(semanticModel);
                if (unresolvedInvocation == null) return null;

                return (document, semanticModel, outerNode, unresolvedInvocation);
            }
        );
        SearchResults = new TaskManager<SearchResult[]>
        (
            async (CancellationToken cancellationToken) =>
            {
                UnresolvedMethodCallInfo? unresolvedInvocation = await UnresolvedInvocation(cancellationToken);
                if (unresolvedInvocation == null) return new SearchResult[0];

                return await CodepediaApi.Search(unresolvedInvocation.Name, cancellationToken);
            }
        );
        SearchResultsInfo = new TaskManager<Dictionary<SearchResult, WikiEntry>>
        (
            async (CancellationToken cancellationToken) =>
            {
                // UnresolvedMethodCallInfo? unresolvedInvocation = await UnresolvedInvocation(cancellationToken);
                // if (unresolvedInvocation == null) return new Dictionary<SearchResult, WikiEntry>();
                // string invocationName = unresolvedInvocation.Name;

                return await CodepediaApi.InterpretSearchResults(await SearchResults.Value(cancellationToken), cancellationToken);
            }
        );
    }
}

This coding pattern should allow other classes to request exactly what they need on a per-need basis.

What I want reviewed

• Is my code useful? (Is there a better way to do this?)
• Could the performance/elegance of my code improved? (Any comments to make it better would be appreciated.)
• Did I make a mistake? (Is my code thread-safe? Are there any race conditions or ways it could get stuck?)

Answer 1

Here are my observations:

TryGetOrAdd of Extensions

• Naming a parameter to d is not really a good choice
• I guess you have performed the TryGetValue twice due to the double-checked lock
  • I would like to warn you with Jon Skeet's words

Double-checking locking now works in Java as well as C# (the Java memory model changed and this is one of the effects). However, you have to get it exactly right. If you mess things up even slightly, you may well end up losing the thread safety.

• Locking on the d is not a good idea since it is provided by the caller (so, it can be considered as public)
  • Related SO topic
  • Pluralsight Best practices
• Instead of having three return statements you can achieve the same with two

private static readonly object syncObject = new();
public static TValue TryGetOrAdd<TKey, TValue>(this IDictionary<TKey, TValue> dict, TKey key, Func<TValue> toAdd)
{
    if (dict.TryGetValue(key, out TValue value))
        return value;
            
    lock (syncObject)
    {
        if (!dict.TryGetValue(key, out value))
        {
            value = toAdd();
            dict.Add(key, value);
        }
        return value;
    }
}

TaskManager

• Is there any particular reason why did you not use target-typed new expression for _cancellationToken?
  • By the way _cancellationToken is a really bad name for this member since it is a CancellationTokenSource not a CancellationToken
• It is not absolutely clear for me why did you pass the initializer via the ctor and the cancellationToken via the Value method
  • Can't you pass the the two either via ctor or via the method?
• In case of C# a method usually starts with verb, so GetValue or RetrieveValue might be better
• If the IsCancellationRequested is true then would it make more sense to call Task.FromCanceled<T>(cancellationToken) instead?
• Here yet again you can simplify your lock block to have a single return statement

lock (_initializationLock)
{
    if (_task is null)
    {
        _cancellationTokenSource = new();
        _task = _initializer(_cancellationTokenSource.Token);
    }
    return _task;
}

CodepediaApi's Search

• Here I've only two tiny suggestions
  • c is not the best name for CancellationToken
  • Instead of allocating empty array (new SearchResult[0] rather prefer Array.Empty<SearchResult>()

---

UPDATE #1: Reflect to the first comment

The method of creating a static variable called syncObject would mean that only one Dictionary could be accessed at a time by the TryGetOrAdd function. In my code's case that work fine as the only Dictionary is the search dictionary. The issue with using another object as a sync object is that since dict is not locked, the first call to dict.TryGetValue(key, out TValue value) could be accessing the Dictionary while it's an invalid state. So, the only solution would be to surround the whole function with lock (syncObject) which would decrease performance

The core problem here is ownership. Since the dictionary is not owned by the method or the class of the method anyone can access it (and potentially modify it) since the lock scope is the method, not the object. Let me demonstrate with a simple example:

public static void Main()
{
    var dict = new Dictionary<string, int>();
    var lockAcquired = new TaskCompletionSource<object>();
    _ = Task.Run(async () => { await lockAcquired.Task; dict.Add("1", 1);});
    Add(dict, lockAcquired);
}

public static void Add(Dictionary<string, int> dict, TaskCompletionSource<object> signal)
{
    lock(dict)
    {
        signal.SetResult(null);
        Thread.Sleep(100);
        dict.Add("1", 1);
    }
}

This code will throw an exception inside the Add method:

System.ArgumentException: An item with the same key has already been added. Key: 1

So, as you can see neither mine nor your solution is fine.