New API consideration: JoinableTaskContext.OnMainThreadBlocked

Is your feature request related to a problem? Please describe.

When solution loading moves towards async, it exposes new issues inside VS product. Some original background tasks would be blocked to switch to UI thread now have a chance to run in the middle of solution loading time and failed in some odd points due to inconsistent solution state (for example, GetProjectOfGuid/Project.UniqueName could fail and leads those tasks to fail.) We want to introduce simple contract/API, so those tasks can await until the solution/project loading time is over explicitly, instead of depending on SwitchToMainThread would be blocked for background tasks during that period of time.

This, however, carries a new risk, when it is being adopted in the product. The original SwitchToMainThread would be unblocked when the background task is blocking the UI thread, while the additional waiting point would not. As the other task now has chance to run in the middle of loading, this could happen. We think an API to allow the new solution await API could abort the waiting point could help, and the API could be useful in other scenarios. Actually, CPS has already implemented something closer to promote UI thread blocking tasks in some places.

Describe the solution you'd like

IDisposable JoinableTaskContext.OnMainThreadBlocked<T>(Action<T> callback, T callbackState)

The usage pattern:

using (JTFContext.OnMainThreadBlocked<TaskCompletionSource<bool>>(
  t => t.TrySetResult(false),
  currentWaitingTaskCompletionSource))
  {
    await currentWaitingTaskCompletionSource.Task;
  }

another extension method, which throws OperationCancelledException when the task is blocking the UI thread -- I am not sure this type of exception is desirable, and want to hear more feedback.

Task WaitUnlessBlockingMainThreadAsync(this Task slowTask, JoinableTaskContext context, CancellationToken cancellationToken = default);

The usage pattern:

await solutionLoadedCompletionSource.WaitUnlessBlockingMainThreadAsync(JTF.Context, cancellationToken)

Describe alternatives you've considered

This doesn't have to be in the JTF library, but a built-in one might reduce the chance to have duplicated/incorrect implementation in different places.

Additional context

The implementation is simple. Inside the function, it creates a background task and SwitchToMainThread with a special JTF factory, which does not PostToUnderlyingSynchronizationContext, so the inner task only runs, when the current task is blocking the main thread. It needs to chain clean-up logic correctly to prevent memory leaks.

The primary reason is that we have various of code to use SwitchToMainThread as a way to wait project/solution to be loaded, which is now broken after solution load can be async. Waiting on solution loading completion could lead to deadlock, if the current task is being waited (and blocking UI thread earlier).

One example is that the logic to delay design time builds after the solution is loaded, we do want it to resume if another code is blocking the UI thread to wait build result earlier. WaitUnlessBlockingMainThreadAsync can also be used to prevent UI thread to block on long/slow task, which is not expected to block the UI thread. Or we can use OnMainThreadBlocked to write ETW logging in a slow task, or raise priority of certain work being throttled in a queue.

Start with the proposed JoinableTaskContext.OnMainThreadBlocked API and the WaitUnlessBlockingMainThreadAsync usage examples. Review the deadlock, cancellation, callback cleanup, and memory-leak concerns raised in the issue, along with the requested feedback on exception behavior and alternatives. Done requires an agreed API design and implementation scope.