|
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// =+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
//
// OrderPreservingPipeliningSpoolingTask.cs
//
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Linq.Parallel;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
namespace System.Linq.Parallel
{
#if !FEATURE_WASM_MANAGED_THREADS
[System.Runtime.Versioning.UnsupportedOSPlatform("browser")]
#endif
internal sealed class OrderPreservingPipeliningSpoolingTask<TOutput, TKey> : SpoolingTaskBase
{
private readonly QueryTaskGroupState _taskGroupState; // State shared among tasks.
private readonly QueryOperatorEnumerator<TOutput, TKey> _partition; // The source partition.
private readonly bool[] _consumerWaiting; // Whether a consumer is waiting on a particular producer
private readonly bool[] _producerWaiting; // Whether a particular producer is waiting on the consumer
private readonly bool[] _producerDone; // Whether each producer is done
private readonly int _partitionIndex; // Index of the partition owned by this task.
private readonly Queue<Pair<TKey, TOutput>>[] _buffers; // The buffer for the results
private readonly object _bufferLock; // A lock for the buffer
/// <summary>
/// Whether the producer is allowed to buffer up elements before handing a chunk to the consumer.
/// If false, the producer will make each result available to the consumer immediately after it is
/// produced.
/// </summary>
private readonly bool _autoBuffered;
/// <summary>
/// The number of elements to accumulate on the producer before copying the elements to the
/// producer-consumer buffer. This constant is only used in the AutoBuffered mode.
///
/// Experimentally, 16 appears to be sufficient buffer size to compensate for the synchronization
/// cost.
/// </summary>
private const int PRODUCER_BUFFER_AUTO_SIZE = 16;
/// <summary>
/// Constructor
/// </summary>
internal OrderPreservingPipeliningSpoolingTask(
QueryOperatorEnumerator<TOutput, TKey> partition,
QueryTaskGroupState taskGroupState,
bool[] consumerWaiting,
bool[] producerWaiting,
bool[] producerDone,
int partitionIndex,
Queue<Pair<TKey, TOutput>>[] buffers,
object bufferLock,
bool autoBuffered)
: base(partitionIndex, taskGroupState)
{
Debug.Assert(partition != null);
Debug.Assert(taskGroupState != null);
Debug.Assert(consumerWaiting != null);
Debug.Assert(producerWaiting != null && producerWaiting.Length == consumerWaiting.Length);
Debug.Assert(producerDone != null && producerDone.Length == consumerWaiting.Length);
Debug.Assert(buffers != null && buffers.Length == consumerWaiting.Length);
Debug.Assert(partitionIndex >= 0 && partitionIndex < consumerWaiting.Length);
_partition = partition;
_taskGroupState = taskGroupState;
_producerDone = producerDone;
_consumerWaiting = consumerWaiting;
_producerWaiting = producerWaiting;
_partitionIndex = partitionIndex;
_buffers = buffers;
_bufferLock = bufferLock;
_autoBuffered = autoBuffered;
}
/// <summary>
/// This method is responsible for enumerating results and enqueuing them to
/// the output buffer as appropriate. Each base class implements its own.
/// </summary>
protected override void SpoolingWork()
{
TOutput element = default(TOutput)!;
TKey key = default(TKey)!;
int chunkSize = _autoBuffered ? PRODUCER_BUFFER_AUTO_SIZE : 1;
Pair<TKey, TOutput>[] chunk = new Pair<TKey, TOutput>[chunkSize];
var partition = _partition;
CancellationToken cancelToken = _taskGroupState.CancellationState.MergedCancellationToken;
int lastChunkSize;
do
{
lastChunkSize = 0;
while (lastChunkSize < chunkSize && partition.MoveNext(ref element!, ref key))
{
chunk[lastChunkSize] = new Pair<TKey, TOutput>(key, element);
lastChunkSize++;
}
if (lastChunkSize == 0) break;
lock (_bufferLock)
{
// Check if the query has been cancelled.
if (cancelToken.IsCancellationRequested)
{
break;
}
for (int i = 0; i < lastChunkSize; i++)
{
_buffers[_partitionIndex].Enqueue(chunk[i]);
}
if (_consumerWaiting[_partitionIndex])
{
Monitor.Pulse(_bufferLock);
_consumerWaiting[_partitionIndex] = false;
}
// If the producer buffer is too large, wait.
// Note: we already checked for cancellation after acquiring the lock on this producer.
// That guarantees that the consumer will eventually wake up the producer.
if (_buffers[_partitionIndex].Count >= OrderPreservingPipeliningMergeHelper<TOutput, TKey>.MAX_BUFFER_SIZE)
{
_producerWaiting[_partitionIndex] = true;
Monitor.Wait(_bufferLock);
}
}
} while (lastChunkSize == chunkSize);
}
/// <summary>
/// Creates and begins execution of a new set of spooling tasks.
/// </summary>
public static void Spool(
QueryTaskGroupState groupState, PartitionedStream<TOutput, TKey> partitions,
bool[] consumerWaiting, bool[] producerWaiting, bool[] producerDone,
Queue<Pair<TKey, TOutput>>[] buffers, object[] bufferLocks,
TaskScheduler taskScheduler, bool autoBuffered)
{
Debug.Assert(groupState != null);
Debug.Assert(partitions != null);
Debug.Assert(producerDone != null && producerDone.Length == partitions.PartitionCount);
Debug.Assert(buffers != null && buffers.Length == partitions.PartitionCount);
Debug.Assert(bufferLocks != null);
int degreeOfParallelism = partitions.PartitionCount;
// Initialize the buffers and buffer locks.
for (int i = 0; i < degreeOfParallelism; i++)
{
buffers[i] = new Queue<Pair<TKey, TOutput>>(OrderPreservingPipeliningMergeHelper<TOutput, TKey>.INITIAL_BUFFER_SIZE);
bufferLocks[i] = new object();
}
// Ensure all tasks in this query are parented under a common root. Because this
// is a pipelined query, we detach it from the parent (to avoid blocking the calling
// thread), and run the query on a separate thread.
Task rootTask = new Task(
() =>
{
for (int i = 0; i < degreeOfParallelism; i++)
{
QueryTask asyncTask = new OrderPreservingPipeliningSpoolingTask<TOutput, TKey>(
partitions[i], groupState, consumerWaiting, producerWaiting,
producerDone, i, buffers, bufferLocks[i], autoBuffered);
asyncTask.RunAsynchronously(taskScheduler);
}
});
// Begin the query on the calling thread.
groupState.QueryBegin(rootTask);
// And schedule it for execution. This is done after beginning to ensure no thread tries to
// end the query before its root task has been recorded properly.
rootTask.Start(taskScheduler);
// We don't call QueryEnd here; when we return, the query is still executing, and the
// last enumerator to be disposed of will call QueryEnd for us.
}
/// <summary>
/// Dispose the underlying enumerator and wake up the consumer if necessary.
/// </summary>
protected override void SpoolingFinally()
{
// Let the consumer know that this producer is done.
lock (_bufferLock)
{
_producerDone[_partitionIndex] = true;
if (_consumerWaiting[_partitionIndex])
{
Monitor.Pulse(_bufferLock);
_consumerWaiting[_partitionIndex] = false;
}
}
// Call the base implementation.
base.SpoolingFinally();
// Dispose of the source enumerator *after* signaling that the task is done.
// We call Dispose() last to ensure that if it throws an exception, we will not cause a deadlock.
_partition.Dispose();
}
}
}
|