Use Channel Instead of BufferBlock

build/Dependencies.props

@@ -9,6 +9,7 @@
     <SystemMemoryVersion>4.5.1</SystemMemoryVersion>
     <SystemReflectionEmitLightweightPackageVersion>4.3.0</SystemReflectionEmitLightweightPackageVersion>
     <SystemThreadingTasksDataflowPackageVersion>4.8.0</SystemThreadingTasksDataflowPackageVersion>
+    <SystemThreadingChannelsPackageVersion>4.7.1</SystemThreadingChannelsPackageVersion>
   </PropertyGroup>
 
   <!-- Other/Non-Core Product Dependencies -->

src/Microsoft.ML.Data/Microsoft.ML.Data.csproj

@@ -10,6 +10,7 @@
   <ItemGroup>
     <PackageReference Include="Newtonsoft.Json" Version="$(NewtonsoftJsonPackageVersion)" />
     <PackageReference Include="System.Collections.Immutable" Version="$(SystemCollectionsImmutableVersion)" />
+    <PackageReference Include="System.Threading.Channels" Version="$(SystemThreadingChannelsPackageVersion)" />
     <PackageReference Include="System.Threading.Tasks.Dataflow" Version="$(SystemThreadingTasksDataflowPackageVersion)" />
     <PackageReference Include="System.Reflection.Emit.Lightweight" Version="$(SystemReflectionEmitLightweightPackageVersion)" />
   </ItemGroup>

src/Microsoft.ML.Data/Transforms/RowShufflingTransformer.cs

@@ -5,8 +5,8 @@
 using System;
 using System.Collections.Generic;
 using System.Threading;
+using System.Threading.Channels;
 using System.Threading.Tasks;
-using System.Threading.Tasks.Dataflow;
 using Microsoft.ML;
 using Microsoft.ML.CommandLine;
 using Microsoft.ML.Data;
@@ -486,8 +486,8 @@ private static readonly FuncInstanceMethodInfo1<Cursor, int, Delegate> _createGe
             private int _liveCount;
             private bool _doneConsuming;
 
-            private readonly BufferBlock<int> _toProduce;
-            private readonly BufferBlock<int> _toConsume;
+            private readonly Channel<int> _toProduceChannel;
+            private readonly Channel<int> _toConsumeChannel;
             private readonly Task _producerTask;
             private Exception _producerTaskException;
 
@@ -541,13 +541,13 @@ public Cursor(IChannelProvider provider, int poolRows, DataViewRowCursor input,
                 _liveCount = 1;
 
                 // Set up the producer worker.
-                _toConsume = new BufferBlock<int>();
-                _toProduce = new BufferBlock<int>();
+                _toConsumeChannel = Channel.CreateUnbounded<int>(new UnboundedChannelOptions { SingleWriter = true });
+                _toProduceChannel = Channel.CreateUnbounded<int>(new UnboundedChannelOptions { SingleWriter = true });
                 // First request the pool - 1 + block size rows, to get us going.
-                PostAssert(_toProduce, _poolRows - 1 + _blockSize);
+                PostAssert(_toProduceChannel, _poolRows - 1 + _blockSize);
                 // Queue up the remaining capacity.
                 for (int i = 1; i < _bufferDepth; ++i)
-                    PostAssert(_toProduce, _blockSize);
+                    PostAssert(_toProduceChannel, _blockSize);
 
                 _producerTask = ProduceAsync();
             }
@@ -559,28 +559,28 @@ protected override void Dispose(bool disposing)
 
                 if (disposing)
                 {
-                    _toProduce.Complete();
+                    _toProduceChannel.Writer.Complete();
                     _producerTask.Wait();
 
                     // Complete the consumer after the producerTask has finished, since producerTask could
                     // have posted more items to _toConsume.
-                    _toConsume.Complete();
+                    _toConsumeChannel.Writer.Complete();
 
                     // Drain both BufferBlocks - this prevents what appears to be memory leaks when using the VS Debugger
                     // because if a BufferBlock still contains items, its underlying Tasks are not getting completed.
                     // See https://github.com/dotnet/corefx/issues/30582 for the VS Debugger issue.
                     // See also https://github.com/dotnet/machinelearning/issues/4399
-                    _toProduce.TryReceiveAll(out _);
-                    _toConsume.TryReceiveAll(out _);
+                    _toProduceChannel.Reader.ReadAsync();
+                    _toConsumeChannel.Reader.ReadAsync();
                 }
 
                 _disposed = true;
                 base.Dispose(disposing);
             }
 
-            public static void PostAssert<T>(ITargetBlock<T> target, T item)
+            public static void PostAssert<T>(Channel<T> target, T item)
             {
-                bool retval = target.Post(item);
+                bool retval = target.Writer.TryWrite(item);
                 Contracts.Assert(retval);
             }
 
@@ -594,10 +594,10 @@ private async Task ProduceAsync()
                 try
                 {
                     int circularIndex = 0;
-                    while (await _toProduce.OutputAvailableAsync().ConfigureAwait(false))
+                    while (await _toProduceChannel.Reader.WaitToReadAsync().ConfigureAwait(false))
                     {
                         int requested;
-                        if (!_toProduce.TryReceive(out requested))
+                        if (!_toProduceChannel.Reader.TryRead(out requested))
                         {
                             // OutputAvailableAsync returned true, but TryReceive returned false -
                             // so loop back around and try again.
@@ -618,14 +618,14 @@ private async Task ProduceAsync()
                             if (circularIndex == _pipeIndices.Length)
                                 circularIndex = 0;
                         }
-                        PostAssert(_toConsume, numRows);
+                        PostAssert(_toConsumeChannel, numRows);
                         if (numRows < requested)
                         {
                             // We've reached the end of the cursor. Send the sentinel, then exit.
                             // This assumes that the receiver will receive things in Post order
                             // (so that the sentinel is received, after the last Post).
                             if (numRows > 0)
-                                PostAssert(_toConsume, 0);
+                                PostAssert(_toConsumeChannel, 0);
                             return;
                         }
                     }
@@ -634,7 +634,7 @@ private async Task ProduceAsync()
                 {
                     _producerTaskException = ex;
                     // Send the sentinel in this case as well, the field will be checked.
-                    PostAssert(_toConsume, 0);
+                    PostAssert(_toConsumeChannel, 0);
                 }
             }
 
@@ -652,14 +652,14 @@ protected override bool MoveNextCore()
                     // We should let the producer know it can give us more stuff.
                     // It is possible for int values to be sent beyond the
                     // end of the sentinel, but we suppose this is irrelevant.
-                    PostAssert(_toProduce, _deadCount);
+                    PostAssert(_toProduceChannel, _deadCount);
                     _deadCount = 0;
                 }
 
                 while (_liveCount < _poolRows && !_doneConsuming)
                 {
                     // We are under capacity. Try to get some more.
-                    int got = _toConsume.Receive();
+                    _toConsumeChannel.Reader.TryRead(out int got);
                     if (got == 0)
                     {
                         // We've reached the end sentinel. There's no reason

src/Microsoft.ML.Sweeper/AsyncSweeper.cs

@@ -5,8 +5,8 @@
 using System;
 using System.Collections.Generic;
 using System.Threading;
+using System.Threading.Channels;
 using System.Threading.Tasks;
-using System.Threading.Tasks.Dataflow;
 using Microsoft.ML;
 using Microsoft.ML.CommandLine;
 using Microsoft.ML.Internal.Utilities;
@@ -168,7 +168,7 @@ public sealed class Options
         private readonly object _lock;
         private readonly CancellationTokenSource _cts;
 
-        private readonly BufferBlock<ParameterSetWithId> _paramQueue;
+        private readonly Channel<ParameterSetWithId> _paramChannel;
         private readonly int _relaxation;
         private readonly ISweeper _baseSweeper;
         private readonly IHost _host;
@@ -208,7 +208,8 @@ public DeterministicSweeperAsync(IHostEnvironment env, Options options)
             _lock = new object();
             _results = new List<IRunResult>();
             _nullRuns = new HashSet<int>();
-            _paramQueue = new BufferBlock<ParameterSetWithId>();
+            _paramChannel = Channel.CreateUnbounded<ParameterSetWithId>(
+                new UnboundedChannelOptions { SingleWriter = true });
 
             PrepareNextBatch(null);
         }
@@ -220,12 +221,12 @@ private void PrepareNextBatch(IEnumerable<IRunResult> results)
             if (Utils.Size(paramSets) == 0)
             {
                 // Mark the queue as completed.
-                _paramQueue.Complete();
+                _paramChannel.Writer.Complete();
                 return;
             }
             // Assign an id to each ParameterSet and enque it.
             foreach (var paramSet in paramSets)
-                _paramQueue.Post(new ParameterSetWithId(_numGenerated++, paramSet));
+                _paramChannel.Writer.TryWrite(new ParameterSetWithId(_numGenerated++, paramSet));
             EnsureResultsSize();
         }
 
@@ -278,7 +279,7 @@ public async Task<ParameterSetWithId> ProposeAsync()
                 return null;
             try
             {
-                return await _paramQueue.ReceiveAsync(_cts.Token);
+                return await _paramChannel.Reader.ReadAsync(_cts.Token);
             }
             catch (InvalidOperationException)
             {

src/Microsoft.ML.Sweeper/Microsoft.ML.Sweeper.csproj

@@ -6,6 +6,10 @@
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
   </PropertyGroup>
 
+  <ItemGroup>
+    <PackageReference Include="System.Threading.Channels" Version="$(SystemThreadingChannelsPackageVersion)" />
+  </ItemGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\Microsoft.ML.Core\Microsoft.ML.Core.csproj" />
     <ProjectReference Include="..\Microsoft.ML.CpuMath\Microsoft.ML.CpuMath.csproj" />