Files
UnrealEngine/Engine/Source/Programs/Shared/EpicGames.Core/ChannelWithRefCount.cs
2025-05-18 13:04:45 +08:00

64 lines
1.4 KiB
C#

// Copyright Epic Games, Inc. All Rights Reserved.
using System;
using System.Threading;
using System.Threading.Channels;
namespace EpicGames.Core
{
/// <summary>
/// Implements a <see cref="Channel{T}"/> which signals completion according to a reference count.
/// </summary>
public class ChannelWithRefCount<T> : Channel<T>
{
int _count = 1;
/// <summary>
/// Called when the channel is completed
/// </summary>
public event Action? OnComplete;
/// <summary>
/// Constructor
/// </summary>
internal ChannelWithRefCount(Channel<T> inner)
{
Reader = inner.Reader;
Writer = inner.Writer;
}
/// <summary>
/// Adds a new reference to the channel
/// </summary>
public ChannelWithRefCount<T> AddRef()
{
Interlocked.Increment(ref _count);
return this;
}
/// <summary>
/// Releases a reference to the channel
/// </summary>
public void Release()
{
if (Interlocked.Decrement(ref _count) == 0)
{
OnComplete?.Invoke();
Writer.Complete();
}
}
}
/// <summary>
/// Extension methods for <see cref="Channel{T}"/>
/// </summary>
public static class ChannelExtensions
{
/// <summary>
/// Creates a channel whose completion state is controlled by a reference count
/// </summary>
public static ChannelWithRefCount<T> WithRefCount<T>(this Channel<T> channel)
=> new ChannelWithRefCount<T>(channel);
}
}