130 lines
3.4 KiB
C#
130 lines
3.4 KiB
C#
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
using System.Collections.Concurrent;
|
|
using EpicGames.Core;
|
|
using Microsoft.Extensions.Logging;
|
|
using MongoDB.Driver;
|
|
|
|
namespace HordeServer.Utilities
|
|
{
|
|
/// <summary>
|
|
/// Utility class for buffering document writes to a Mongo collection
|
|
/// </summary>
|
|
public sealed class MongoBufferedWriter<TDocument> : IAsyncDisposable
|
|
where TDocument : class
|
|
{
|
|
readonly IMongoCollection<TDocument> _collection;
|
|
readonly int _flushCount;
|
|
readonly TimeSpan _flushTime;
|
|
readonly ConcurrentQueue<TDocument> _queue = new ConcurrentQueue<TDocument>();
|
|
readonly BackgroundTask _backgroundTask;
|
|
readonly AsyncEvent _newDataEvent = new AsyncEvent();
|
|
readonly AsyncEvent _flushEvent = new AsyncEvent();
|
|
readonly ILogger _logger;
|
|
|
|
/// <summary>
|
|
/// Constructor
|
|
/// </summary>
|
|
public MongoBufferedWriter(IMongoCollection<TDocument> collection, ILogger logger)
|
|
: this(collection, 50, TimeSpan.FromSeconds(5.0), logger)
|
|
{
|
|
}
|
|
|
|
/// <summary>
|
|
/// Constructor
|
|
/// </summary>
|
|
public MongoBufferedWriter(IMongoCollection<TDocument> collection, int flushCount, TimeSpan flushTime, ILogger logger)
|
|
{
|
|
_collection = collection;
|
|
_flushCount = flushCount;
|
|
_flushTime = flushTime;
|
|
_backgroundTask = new BackgroundTask(BackgroundFlushAsync);
|
|
_logger = logger;
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public async ValueTask DisposeAsync()
|
|
{
|
|
await _backgroundTask.DisposeAsync();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Start the background task to periodically flush data to the DB
|
|
/// </summary>
|
|
public ValueTask StartAsync()
|
|
{
|
|
_backgroundTask.Start();
|
|
return new ValueTask();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Stops the background task
|
|
/// </summary>
|
|
public async ValueTask StopAsync(CancellationToken cancellationToken)
|
|
{
|
|
await _backgroundTask.StopAsync(cancellationToken);
|
|
await FlushAsync(cancellationToken);
|
|
}
|
|
|
|
// Flushes the sink in the background
|
|
async Task BackgroundFlushAsync(CancellationToken cancellationToken)
|
|
{
|
|
Task newDataTask = _newDataEvent.Task;
|
|
Task flushTask = _flushEvent.Task;
|
|
|
|
while (!cancellationToken.IsCancellationRequested)
|
|
{
|
|
try
|
|
{
|
|
await newDataTask.WaitAsync(cancellationToken);
|
|
await Task.WhenAny(flushTask, Task.Delay(_flushTime, cancellationToken));
|
|
|
|
newDataTask = _newDataEvent.Task;
|
|
flushTask = _flushEvent.Task;
|
|
|
|
await FlushAsync(cancellationToken);
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
break;
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
_logger.LogError(ex, "Exception in buffered writer: {Message}", ex.Message);
|
|
await Task.Delay(TimeSpan.FromSeconds(30.0), cancellationToken);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public async ValueTask FlushAsync(CancellationToken cancellationToken)
|
|
{
|
|
// Copy all the event documents from the queue
|
|
List<TDocument> documents = new List<TDocument>(_queue.Count);
|
|
while (_queue.TryDequeue(out TDocument? document))
|
|
{
|
|
documents.Add(document);
|
|
}
|
|
|
|
// Insert them into the database
|
|
if (documents.Count > 0)
|
|
{
|
|
_logger.LogDebug("Writing {NumEvents} new telemetry events to {CollectionName}.", documents.Count, _collection.CollectionNamespace.CollectionName);
|
|
await _collection.InsertManyAsync(documents, cancellationToken: cancellationToken);
|
|
}
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public void Write(TDocument document)
|
|
{
|
|
_queue.Enqueue(document);
|
|
_newDataEvent.Set();
|
|
|
|
if (_queue.Count > _flushCount)
|
|
{
|
|
_flushEvent.Set();
|
|
}
|
|
}
|
|
}
|
|
}
|