MES-ETL/MesETL.App/Services/DataRecordQueue.cs

using System.Collections.Concurrent;
using System.Diagnostics.CodeAnalysis;
using TaskExtensions = MesETL.App.Helpers.TaskExtensions;

namespace MesETL.App.Services;

/// <summary>
/// 数据队列
/// </summary>
public class DataRecordQueue : IDisposable
{
    private readonly BlockingCollection<DataRecord> _queue;

    private long _currentCharCount;
    private readonly long _maxCharCount = 2_147_483_648; // 4GiB

    public int Count => _queue.Count;
    public bool IsCompleted => _queue.IsCompleted;
    public bool IsAddingCompleted => _queue.IsAddingCompleted;

    public event Action? OnRecordWrite;
    public event Action? OnRecordRead;

    public DataRecordQueue() : this(500_000, 2_147_483_648) // 默认容量最大500K
    {
    }
    
    public DataRecordQueue(int boundedCapacity, long maxCharCount)
    {
        _queue = new BlockingCollection<DataRecord>(boundedCapacity); 
        _maxCharCount = maxCharCount;
    }
    
    public void CompleteAdding() => _queue.CompleteAdding();
    
    public bool TryDequeue([MaybeNullWhen(false)] out DataRecord record)
    {
        if (_queue.TryTake(out record))
        {
            Interlocked.Add(ref _currentCharCount, -record.FieldCharCount);
            OnRecordRead?.Invoke();
            return true;
        }

        return false;
    }


    public async Task EnqueueAsync(DataRecord record)
    {
        var charCount = record.FieldCharCount;
        if(_currentCharCount + charCount > _maxCharCount)
            await TaskExtensions.WaitUntil(() => _currentCharCount + charCount < _maxCharCount, 50);
        _queue.Add(record);
        Interlocked.Add(ref _currentCharCount, charCount);
        OnRecordWrite?.Invoke();
    }


    public void Dispose()
    {
        _queue.Dispose();
    }
}