Refactor AsyncMemoize, introduce AsyncLazy

[majocha]

Maintainability is the goal here. The code is smaller now for sure but that not necessarily implies more understandable, so let's see how it goes.

The state machine of running / tracking requests / canceling jobs is now moved to a separate AsyncLazy type.

This is not just a refactor because cancellation behavior changed:
A running job that is awaited by a client will never get canceled "from the inside" - from the requestor's pov it'd be just an exception thrown from the job, so this is not allowed.

Completion behavior is configurable now. If all clients cancel their requests, AsyncLazy can either cancel the job or follow to completion and store result anyway. This may or may not be useful.

The LRU cache now just holds AsyncLazy values. State transitions happen inside the AsyncLazy so there is no corresponding cache.Set like before.

Gist of the file for easier reading is here: https://gist.github.com/majocha/6c075dcdabf45c3fd92b03e2cf3a07eb

Todo:

Investigate passing ExceptionDispatchInfo from the job instead of reraising exns.

[github-actions]

⚠️ Release notes required, but author opted out

Warning

Author opted out of release notes, check is disabled for this pull request.
cc @dotnet/fsharp-team-msft

[auduchinok]

@majocha Does it interact with ability to cancel requests from inside accessing IL data in any way? It was introduced in #16137, #16179 and is crucial for us in Rider.
I.e. it's expected that we throw OCE when analysis in cancelled and the C#/VB metadata is considered outdated.

[majocha]

@majocha Does it interact with ability to cancel requests from inside accessing IL data in any way? It was introduced in #16137, #16179 and is crucial for us in Rider. I.e. it's expected that we throw OCE when analysis in cancelled and the C#/VB metadata is considered outdated.

Now, that rings a bell. I'll make sure it still works.
Btw does Rider use TransparentCompiler?

[auduchinok]

Btw does Rider use TransparentCompiler?

Not yet, but we plan to start using it when it's considered stable/recommended.

[majocha]

@auduchinok, this change here should affect only TransparentCompiler, nothing else. OCEs originating from compiler tasks cached by AsyncMemoize should propagate like any other exception, without special treatment. I think this should be ok in this regard.

[auduchinok]

OCEs originating from compiler tasks cached by AsyncMemoize should propagate like any other exception, without special treatment

But we never want to cache OCEs, do we? Or is it less relevant with the transparent compiler, because the change in files will invalidate things in a natural way?

[majocha]

OCEs originating from compiler tasks cached by AsyncMemoize should propagate like any other exception, without special treatment

But we never want to cache OCEs, do we? Or is it less relevant with the transparent compiler, because the change in files will invalidate things in a natural way?

I'm not super familiar with this part, but IIUC AsyncMemoize is keyed by partial project snapshots. Ideally a snapshot should encode all of the deps, such as to make AsyncMemoize.Get a pure function. In practice many things can happen, so we should probably evict exceptional results pretty quickly from the cache. Again, I still don't know much about it, I don't want to spout nonsense :)

[majocha]

Cancellable does not throw OCE on the async boundary, it just calls cancel continuation, very neat:

fsharp/src/Compiler/Utilities/Cancellable.fs

Lines 72 to 82 in fb69e58

	let toAsync c =
	async {
	let! ct = Async.CancellationToken
	let res = run ct c
	return!
	Async.FromContinuations(fun (cont, _econt, ccont) ->
	match res with
	| ValueOrCancelled.Value v -> cont v
	| ValueOrCancelled.Cancelled ce -> ccont ce)
	}

[0101]

Nice!

Can we run some benchmarks of before and after to see if anything changed? (hopefully improved)

OCEs originating from compiler tasks cached by AsyncMemoize should propagate like any other exception, without special treatment

But we never want to cache OCEs, do we? Or is it less relevant with the transparent compiler, because the change in files will invalidate things in a natural way?

I haven't studied this yet, but intuitively I would say OCEs might deserve special treatment. User might open a file, checking starts, they decide to close it so it gets cancelled. Then without changing anything that affects it they open it again. Same snapshot should be used and we need to restart the computation. We should probably still restart even in case of a failed job, but there theoretically we could decide not to, so we should be able to tell the difference between failed and cancelled.

[0101]

How does it look with stack traces here? Are we able to preserve them over the request? Would be nice if we could, though I didn't figure it out before so if not it's not a downgrade...

[majocha]

How does it look with stack traces here? Are we able to preserve them over the request? Would be nice if we could, though I didn't figure it out before so if not it's not a downgrade...

It is possible to store exception in ExceptionDispatchInfo and rethrow it from there instead of using naked exceptions. I tried it, but then I realized it is what async computations already do currently:

fsharp/src/FSharp.Core/async.fs

Line 44 in 3dff60a

let associationTable = ConditionalWeakTable<exn, ExceptionDispatchInfo>()

if you catch an exception in async CE and then rethrow it in another async, it should restore the stacktrace from ExceptionDispatchInfo cached in that ConditionalWeakTable.

I think we need a test case just for this. It would allow for some experiments.

[majocha]

Nice!

Can we run some benchmarks of before and after to see if anything changed? (hopefully improved)

OCEs originating from compiler tasks cached by AsyncMemoize should propagate like any other exception, without special treatment

But we never want to cache OCEs, do we? Or is it less relevant with the transparent compiler, because the change in files will invalidate things in a natural way?

I haven't studied this yet, but intuitively I would say OCEs might deserve special treatment. User might open a file, checking starts, they decide to close it so it gets cancelled. Then without changing anything that affects it they open it again. Same snapshot should be used and we need to restart the computation. We should probably still restart even in case of a failed job, but there theoretically we could decide not to, so we should be able to tell the difference between failed and cancelled.

This is kind of what this PR already do:

Failure is a failure - currently it's just cached as (Result.Error(exn), logger) value.

Cancellation happens when the client cancels the request. We detach from the job and it is configurable whether is runs to completion unawaited or is itself canceled.

What is still unclear to me is the situation with Cancelable.CheckAndThrow(). Does the thrown OCE ever reach the surface, so to say?
Is all of compiler code that potentially do CheckAndThrow() wrapped in Cancellable before being called from TransparentCompiler? If it is, this solves things, because Cancellable.toAsync catches OCE and just runs async cancel continuation, so it will never surface as exception at the AsyncMemoize layer here.

TBH I already built a vsix out of this and I'm using it for some time. Never saw any exceptions so I hope this is a non-issue. On the plus side, I had codefix analyzers throwing each time I switched branches in VS and with this PR this never happens, although it may be unrelated.

[0101]

if you catch an exception in async CE and then rethrow it in another async, it should restore the stacktrace from ExceptionDispatchInfo cached in that ConditionalWeakTable

Previously, the stack trace was always cut by a jump to the thread pool. You could always only see up to the point where a worker thread picked it up from there.

What is still unclear to me is the situation with Cancelable.CheckAndThrow(). Does the thrown OCE ever reach the surface, so to say? Is all of compiler code that potentially do CheckAndThrow() wrapped in Cancellable before being called from TransparentCompiler? If it is, this solves things, because Cancellable.toAsync catches OCE and just runs async cancel continuation, so it will never surface as exception at the AsyncMemoize layer here.

TBH I already built a vsix out of this and I'm using it for some time. Never saw any exceptions so I hope this is a non-issue. On the plus side, I had codefix analyzers throwing each time I switched branches in VS and with this PR this never happens, although it may be unrelated.

That sounds promising! Maybe regular operation is already covered and any other OCEs are actual errors, i.e. they shouldn't have been thrown. However, there is always some risk that some future changes will introduce an issue by throwing somewhere we don't expect it and by then no one will remember this 😂

[0101]

This looks much better than before!

[majocha]

Something to investigate. Script load closure tests don't like this version for some reason:

    member _.TryGet(key: 'TKey, predicate: 'TVersion -> bool) : 'TValue option =
        lock cache <| fun () ->
            cache.GetAll(key)
            |> Seq.tryPick (fun (version, job) ->
                match predicate version, job.TryResult with
                | true, Some(Ok result, _) -> Some result
                | _ -> None)

I had to take a snapshot with toList and do tryPick outside of the lock instead.

[auduchinok]

OCEs originating from compiler tasks cached by AsyncMemoize should propagate like any other exception, without special treatment

But we never want to cache OCEs, do we? Or is it less relevant with the transparent compiler, because the change in files will invalidate things in a natural way?

I haven't studied this yet, but intuitively I would say OCEs might deserve special treatment. User might open a file, checking starts, they decide to close it so it gets cancelled. Then without changing anything that affects it they open it again. Same snapshot should be used and we need to restart the computation.

That's how I thought it should work, i.e. we should not cache the cancellations and recalculate it in future if needed, even on the same snapshots.

I hope we get more Cancelable.CheckAndThrow() here and there in the compiler in future, so computations could stop earlier without using the CPU or blocking other features (as it can be in our case in Rider).

Is all of compiler code that potentially do CheckAndThrow() wrapped in Cancellable before being called from TransparentCompiler? If it is, this solves things, because Cancellable.toAsync catches OCE and just runs async cancel continuation, so it will never surface as exception at the AsyncMemoize layer here.

@majocha Is there any chance you could add tests covering this behavior if there aren't any already?
We can also add an assert inside CheckAndThrow which would check that it's run inside Cancellable only.

[majocha]

@majocha Is there any chance you could add tests covering this behavior if there aren't any already? We can also add an assert inside CheckAndThrow which would check that it's run inside Cancellable only.

Yes, this would be great!

Currently, because of this code:

fsharp/src/Compiler/Utilities/Cancellable.fs

Lines 9 to 11 in a95a9ea

	static let token = AsyncLocal<CancellationToken>()
	static member Token = token.Value

If there is no ambient Cancellable, Token will have default value which probably is CancellationToken.None (?). This is not optimal, because we cannot know if it was really initialized.
This probably should be something like AsyncLocal<CancellationToken voption>() instead.

[majocha]

Added back

    | Completed of result: 't
    | Faulted of exn

to the state machine. Also basic optimizations: computation and task will deallocate on completion and only the result is stored.

[psfinaki]

@majocha - great job here (as usual), also thanks for a solid PR description and GH gist, it helps reviewing things a lot.

The more tests the better here, but this can be done as a follow up - as you wish.
Also, I think AsyncLazy could be isolated to a separate file, as another step to less chaos in the compiler :)

[0101]

Also, I think AsyncLazy could be isolated to a separate file, as another step to less chaos in the compiler :)

I don't think that's necessary at the moment. It can be considered an implementation detail of AsyncMemoize. It would make sense if we find we need it also somewhere else.

[psfinaki]

Also, I think AsyncLazy could be isolated to a separate file, as another step to less chaos in the compiler :)

I don't think that's necessary at the moment. It can be considered an implementation detail of AsyncMemoize. It would make sense if we find we need it also somewhere else.

Yeah, fair, can be done later.

[psfinaki]

Thanks for this!