Skip to content

Instantly share code, notes, and snippets.

@ramonsmits
Created March 16, 2018 12:15
Show Gist options
  • Select an option

  • Save ramonsmits/bef98d3a358826fedfa8140b5e06a00e to your computer and use it in GitHub Desktop.

Select an option

Save ramonsmits/bef98d3a358826fedfa8140b5e06a00e to your computer and use it in GitHub Desktop.
Code sample provided by RavenHQ for batching up writes
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Raven.Client;
/// <summary>
/// This class allows performing high throughput document store operations, using a batching mechanism
/// That mechanism collects up to 128 documents into a single batch and stores them in a single operation, reducing
/// network io compared to an atomic operation. The mechanism is also capable of treating single store operation failure
/// by retrying a batch, executing it item by item
/// </summary>
internal class RavenBatchOptimizer : IDisposable
{
private readonly BlockingCollection<Work> _pending = new BlockingCollection<Work>();
private readonly IDocumentStore _store;
private readonly CancellationToken _ct;
private readonly ConcurrentQueue<ManualResetEventSlim> _mresQueue = new ConcurrentQueue<ManualResetEventSlim>();
private readonly List<Task> _processingTasks;
private readonly CancellationTokenSource _cts;
public RavenBatchOptimizer(IDocumentStore store, CancellationToken ct, int parallelism = 2)
{
_store = store;
_cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
_ct = _cts.Token;
_processingTasks = new List<Task>(parallelism);
for (var i = 0; i < parallelism; i++)
{
_processingTasks.Add(Task.Factory.StartNew(DoWork, TaskCreationOptions.LongRunning));
}
}
/// <summary>
/// Enqueue the entity to be stored and waits until it will be proccessed
/// </summary>
/// <param name="entity"></param>
public void Write(object entity)
{
_cts.Token.ThrowIfCancellationRequested();
if (_mresQueue.TryDequeue(out var mres) == false)
{
mres = new ManualResetEventSlim();
}
var work = new Work
{
Entity = entity,
Event = mres
};
work.Event.Reset();
_pending.Add(work, CancellationToken.None);
// wait for entity to be stored in the next batch
work.Event.Wait(CancellationToken.None);
_mresQueue.Enqueue(mres);
if (work.Exception != null)
throw new InvalidOperationException("Failed to write", work.Exception);
}
public void Dispose()
{
if (_cts.IsCancellationRequested == false)
{
var proccesses = _processingTasks.ToArray();
_processingTasks.Clear();
_cts.Cancel();
_pending.CompleteAdding();
try
{
Task.WaitAll(proccesses);
}
catch (OperationCanceledException)
{
// do nothing
}
}
}
/// <summary>
/// Long running method that performs the batch grouping and execution operations
/// </summary>
/// <returns></returns>
async Task DoWork()
{
var list = new List<Work>();
while (_cts.IsCancellationRequested == false || _pending.IsCompleted == false)
{
list.Clear();
var work = _pending.Take(_ct);
list.Add(work);
while (list.Count < 128 && _pending.TryTake(out var item)) list.Add(item);
using (var session = _store.OpenAsyncSession())
{
var failed = false;
try
{
foreach (var p in list) session.Advanced.Store(p.Entity);
await session.SaveChangesAsync(CancellationToken.None).ConfigureAwait(false);
}
catch (Exception)
{
failed = true;
}
// if we've failed to store the batch (which was entirely rolled back), we try executing each of the
// items seperately, in case we've failed because of a single item
if (failed) await HandleSaveFailure(list).ConfigureAwait(false);
}
foreach (var p in list) p.Event.Set();
}
}
/// <summary>
/// Store entities list one by one, setting the errorous message state to the particular message
/// </summary>
/// <param name="list"></param>
/// <returns></returns>
async Task HandleSaveFailure(IEnumerable<Work> list)
{
foreach (var p in list)
using (var recoverySession = _store.OpenAsyncSession())
{
await recoverySession.StoreAsync(p.Entity, CancellationToken.None).ConfigureAwait(false);
try
{
await recoverySession.SaveChangesAsync(CancellationToken.None).ConfigureAwait(false);
}
catch (Exception e)
{
p.Exception = e;
}
}
}
class Work
{
public object Entity;
public ManualResetEventSlim Event;
public Exception Exception;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment