Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize AsyncLazy<T> #511

Merged
merged 3 commits into from
Dec 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion Package/Core/Promises/Internal/PromiseFieldsInternal.cs
Original file line number Diff line number Diff line change
Expand Up @@ -192,12 +192,16 @@ partial class PromiseMultiAwait<TResult> : PromiseRef<TResult>
private int _retainCounter;
}

partial class PromiseRetainer<TResult> : PromiseRef<TResult>
partial class RetainedPromiseBase<TResult> : PromiseRef<TResult>
{
private TempCollectionBuilder<HandleablePromiseBase> _nextBranches;
private int _retainCounter;
}

partial class PromiseRetainer<TResult> : RetainedPromiseBase<TResult>
{
}

partial class PromiseWaitPromise<TResult> : PromiseSingleAwait<TResult>
{
}
Expand Down
121 changes: 80 additions & 41 deletions Package/Core/Promises/Internal/PromiseRetainerInternal.cs
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@
#undef PROMISE_DEBUG
#endif

#pragma warning disable IDE0016 // Use 'throw' expression

using Proto.Promises.Collections;
using System;
using System.Diagnostics;
Expand All @@ -20,11 +18,9 @@ internal abstract partial class PromiseRefBase
#if !PROTO_PROMISE_DEVELOPER_MODE
[DebuggerNonUserCode, StackTraceHidden]
#endif
internal sealed partial class PromiseRetainer<TResult> : PromiseRef<TResult>
internal abstract partial class RetainedPromiseBase<TResult> : PromiseRef<TResult>
{
private PromiseRetainer() { }

~PromiseRetainer()
~RetainedPromiseBase()
{
if (!WasAwaitedOrForgotten)
{
Expand All @@ -40,37 +36,62 @@ private PromiseRetainer() { }
}

[MethodImpl(InlineOption)]
new private void Reset()
new protected void Reset()
{
_retainCounter = 2; // 1 for dispose, 1 for completion.
base.Reset();
}

[MethodImpl(InlineOption)]
private static PromiseRetainer<TResult> GetOrCreate()
protected void ResetForInternalUse()
{
var obj = ObjectPool.TryTakeOrInvalid<PromiseRetainer<TResult>>();
return obj == InvalidAwaitSentinel.s_instance
? new PromiseRetainer<TResult>()
: obj.UnsafeAs<PromiseRetainer<TResult>>();
_retainCounter = 1; // 1 for completion/dispose.
_nextBranches = new TempCollectionBuilder<HandleablePromiseBase>(0);
base.Reset();
}

[MethodImpl(InlineOption)]
internal static PromiseRetainer<TResult> GetOrCreateAndHookup(PromiseRefBase previous, short id)
protected void Hookup(PromiseRefBase previous, short id)
=> previous.HookupNewPromise(id, this);

[MethodImpl(InlineOption)]
protected void ResetAndHookup(PromiseRefBase previous, short id)
{
var promise = GetOrCreate();
promise.Reset();
previous.HookupNewPromise(id, promise);
Reset();
Hookup(previous, id);
// We create the temp collection after we hook up in case the operation is invalid.
promise._nextBranches = new TempCollectionBuilder<HandleablePromiseBase>(0);
return promise;
_nextBranches = new TempCollectionBuilder<HandleablePromiseBase>(0);
}

[MethodImpl(InlineOption)]
private void Retain()
=> InterlockedAddWithUnsignedOverflowCheck(ref _retainCounter, 1);

internal override void MaybeDispose()
internal void Dispose(short promiseId)
{
lock (this)
{
if (promiseId != Id | WasAwaitedOrForgotten)
{
throw new ObjectDisposedException(nameof(Promise.Retainer), "The promise retainer was already disposed.");
}
ThrowIfInPool(this);

WasAwaitedOrForgotten = true;
}
MaybeDispose();
}

// Same as Dispose, but skips validation.
[MethodImpl(InlineOption)]
protected void DisposeUnsafe()
{
ThrowIfInPool(this);
WasAwaitedOrForgotten = true;
MaybeDispose();
}

internal sealed override void MaybeDispose()
{
if (InterlockedAddWithUnsignedOverflowCheck(ref _retainCounter, -1) != 0)
{
Expand All @@ -87,23 +108,10 @@ internal override void MaybeDispose()
// We handle this directly here because we don't add the PromiseForgetSentinel to this type when it is disposed.
MaybeReportUnhandledRejection(State);
Dispose();
ObjectPool.MaybeRepool(this);
MaybeRepool();
}

internal void Dispose(short promiseId)
{
lock (this)
{
if (promiseId != Id | WasAwaitedOrForgotten)
{
throw new ObjectDisposedException("The promise retainer was already disposed.", GetFormattedStacktrace(2));
}
ThrowIfInPool(this);

WasAwaitedOrForgotten = true;
}
MaybeDispose();
}
protected abstract void MaybeRepool();

internal override bool GetIsCompleted(short promiseId)
{
Expand Down Expand Up @@ -172,6 +180,12 @@ internal override void Handle(PromiseRefBase handler, Promise.State state)
SetCompletionState(state);
handler.MaybeDispose();

HandleBranches(state);
MaybeDispose();
}

protected void HandleBranches(Promise.State state)
{
TempCollectionBuilder<HandleablePromiseBase> branches;
lock (this)
{
Expand All @@ -181,7 +195,6 @@ internal override void Handle(PromiseRefBase handler, Promise.State state)
{
branches[i].Handle(this, state);
}
MaybeDispose();
}

internal Promise<TResult> WaitAsync(short promiseId)
Expand All @@ -190,7 +203,7 @@ internal Promise<TResult> WaitAsync(short promiseId)
{
if (!GetIsValid(promiseId))
{
throw new ObjectDisposedException("The promise retainer was already disposed.", GetFormattedStacktrace(2));
throw new ObjectDisposedException(nameof(Promise.Retainer), "The promise retainer was already disposed.");
}
ThrowIfInPool(this);

Expand All @@ -199,15 +212,13 @@ internal Promise<TResult> WaitAsync(short promiseId)
return GetWaitAsync(promiseId);
}

// Same as WaitAsync, but skips validation.
[MethodImpl(InlineOption)]
internal Promise<TResult> WaitAsyncSkipValidation()
internal Promise<TResult> WaitAsyncUnsafe()
{
#if PROMISE_DEBUG
return WaitAsync(Id);
#else
ThrowIfInPool(this);
Retain();
return GetWaitAsync(Id);
#endif
}

[MethodImpl(InlineOption)]
Expand All @@ -224,6 +235,34 @@ private Promise<TResult> GetWaitAsync(short promiseId)
#endif
}
}

#if !PROTO_PROMISE_DEVELOPER_MODE
[DebuggerNonUserCode, StackTraceHidden]
#endif
internal sealed partial class PromiseRetainer<TResult> : RetainedPromiseBase<TResult>
{
private PromiseRetainer() { }

[MethodImpl(InlineOption)]
private static PromiseRetainer<TResult> GetOrCreate()
{
var obj = ObjectPool.TryTakeOrInvalid<PromiseRetainer<TResult>>();
return obj == InvalidAwaitSentinel.s_instance
? new PromiseRetainer<TResult>()
: obj.UnsafeAs<PromiseRetainer<TResult>>();
}

[MethodImpl(InlineOption)]
internal static PromiseRetainer<TResult> GetOrCreateAndHookup(PromiseRefBase previous, short id)
{
var promise = GetOrCreate();
promise.ResetAndHookup(previous, id);
return promise;
}

protected override void MaybeRepool()
=> ObjectPool.MaybeRepool(this);
}
} // class PromiseRefBase
} // class Internal
}
Loading
Loading