119 lines
3.5 KiB
C#
119 lines
3.5 KiB
C#
using System.Diagnostics;
|
|
|
|
namespace LANCommander.Server.Endpoints;
|
|
|
|
/// <summary>
|
|
/// Wraps a source stream and limits how fast it can be read by sleeping when a per-second byte budget is exceeded.
|
|
/// Seek and length members are delegated to the inner stream so HTTP range/resume and Content-Length keep working.
|
|
/// A <see cref="_bytesPerSecond"/> of 0 (or less) disables throttling.
|
|
/// </summary>
|
|
internal sealed class ThrottledStream : Stream
|
|
{
|
|
private readonly Stream _inner;
|
|
private readonly long _bytesPerSecond;
|
|
|
|
private long _windowStartTimestamp;
|
|
private long _bytesInWindow;
|
|
|
|
public ThrottledStream(Stream inner, long bytesPerSecond)
|
|
{
|
|
_inner = inner ?? throw new ArgumentNullException(nameof(inner));
|
|
_bytesPerSecond = bytesPerSecond;
|
|
_windowStartTimestamp = Stopwatch.GetTimestamp();
|
|
}
|
|
|
|
public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
|
|
{
|
|
if (_bytesPerSecond <= 0)
|
|
return await _inner.ReadAsync(buffer, cancellationToken);
|
|
|
|
await ThrottleAsync(cancellationToken);
|
|
|
|
var read = await _inner.ReadAsync(buffer, cancellationToken);
|
|
|
|
Accumulate(read);
|
|
|
|
return read;
|
|
}
|
|
|
|
public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
|
|
{
|
|
return await ReadAsync(buffer.AsMemory(offset, count), cancellationToken);
|
|
}
|
|
|
|
public override int Read(byte[] buffer, int offset, int count)
|
|
{
|
|
if (_bytesPerSecond <= 0)
|
|
return _inner.Read(buffer, offset, count);
|
|
|
|
ThrottleAsync(CancellationToken.None).AsTask().GetAwaiter().GetResult();
|
|
|
|
var read = _inner.Read(buffer, offset, count);
|
|
|
|
Accumulate(read);
|
|
|
|
return read;
|
|
}
|
|
|
|
private async ValueTask ThrottleAsync(CancellationToken cancellationToken)
|
|
{
|
|
if (_bytesInWindow < _bytesPerSecond)
|
|
return;
|
|
|
|
var elapsed = Stopwatch.GetElapsedTime(_windowStartTimestamp);
|
|
var remaining = TimeSpan.FromSeconds(1) - elapsed;
|
|
|
|
if (remaining > TimeSpan.Zero)
|
|
await Task.Delay(remaining, cancellationToken);
|
|
|
|
_windowStartTimestamp = Stopwatch.GetTimestamp();
|
|
_bytesInWindow = 0;
|
|
}
|
|
|
|
private void Accumulate(int read)
|
|
{
|
|
if (read <= 0)
|
|
return;
|
|
|
|
if (Stopwatch.GetElapsedTime(_windowStartTimestamp) >= TimeSpan.FromSeconds(1))
|
|
{
|
|
_windowStartTimestamp = Stopwatch.GetTimestamp();
|
|
_bytesInWindow = 0;
|
|
}
|
|
|
|
_bytesInWindow += read;
|
|
}
|
|
|
|
public override bool CanRead => _inner.CanRead;
|
|
public override bool CanSeek => _inner.CanSeek;
|
|
public override bool CanWrite => false;
|
|
public override long Length => _inner.Length;
|
|
|
|
public override long Position
|
|
{
|
|
get => _inner.Position;
|
|
set => _inner.Position = value;
|
|
}
|
|
|
|
public override long Seek(long offset, SeekOrigin origin) => _inner.Seek(offset, origin);
|
|
|
|
public override void Flush() => _inner.Flush();
|
|
|
|
public override void SetLength(long value) => throw new NotSupportedException();
|
|
|
|
public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
|
|
|
|
protected override void Dispose(bool disposing)
|
|
{
|
|
if (disposing)
|
|
_inner.Dispose();
|
|
|
|
base.Dispose(disposing);
|
|
}
|
|
|
|
public override async ValueTask DisposeAsync()
|
|
{
|
|
await _inner.DisposeAsync();
|
|
await base.DisposeAsync();
|
|
}
|
|
}
|