Fix perf regression in ShuffleRows (#5417)

RowShufflingTransformer is using ChannelReader incorrectly. It needs to block waiting for items to read and was Thread.Sleeping in order to wait, but not spin the current core. This caused a major perf regression.

The fix is to block synchronously correctly - by calling AsTask() on the ValueTask that is returned from the ChannelReader and block on the Task.

Fix #5416

Author: eerhardt

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

@@ -634,26 +634,25 @@ protected override bool MoveNextCore()
                 while (_liveCount < _poolRows && !_doneConsuming)
                 {
                     // We are under capacity. Try to get some more.
-                    var hasReadItem = _toConsumeChannel.Reader.TryRead(out int got);
-                    if (hasReadItem)
+                    ValueTask<int> readTask = _toConsumeChannel.Reader.ReadAsync();
+
+                    // Note you can't wait synchronously on a ValueTask. So if it
+                    // hasn't been completed yet, need to call AsTask() to get a Task
+                    // which can be waited on synchronously.
+                    int got = readTask.IsCompletedSuccessfully ?
+                        readTask.Result :
+                        readTask.AsTask().GetAwaiter().GetResult();
+                    if (got == 0)
                     {
-                        if (got == 0)
-                        {
-                            // We've reached the end of the Channel. There's no reason
-                            // to attempt further communication with the producer.
-                            // Check whether something horrible happened.
-                            if (_producerTaskException != null)
-                                throw Ch.Except(_producerTaskException, "Shuffle input cursor reader failed with an exception");
-                            _doneConsuming = true;
-                            break;
-                        }
-                        _liveCount += got;
-                    }
-                    else
-                    {
-                        // Sleeping for one millisecond to stop the thread from spinning while waiting for the producer.
-                        Thread.Sleep(1);
+                        // We've reached the end of the Channel. There's no reason
+                        // to attempt further communication with the producer.
+                        // Check whether something horrible happened.
+                        if (_producerTaskException != null)
+                            throw Ch.Except(_producerTaskException, "Shuffle input cursor reader failed with an exception");
+                        _doneConsuming = true;
+                        break;
                     }
+                    _liveCount += got;
                 }
                 if (_liveCount == 0)
                     return false;

test/Microsoft.ML.Benchmarks/ShuffleRowsBench.cs

@@ -0,0 +1,53 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using BenchmarkDotNet.Attributes;
+using Microsoft.ML.Benchmarks.Harness;
+using Microsoft.ML.Data;
+
+namespace Microsoft.ML.Benchmarks
+{
+    [CIBenchmark]
+    public class ShuffleRowsBench : BenchmarkBase
+    {
+        private TrainRow[] _rows;
+        private MLContext _context;
+        
+        [GlobalSetup]
+        public void Setup()
+        {
+            _rows = new TrainRow[10_000];
+            for (var i = 0; i < _rows.Length; i++)
+            {
+                _rows[i] = new TrainRow() { Sample = i.ToString(), Week = i, Label = i / 2 };
+            }
+
+            _context = new MLContext();
+        }
+
+        [Benchmark]
+        public void ShuffleRows()
+        {
+            IDataView data = _context.Data.LoadFromEnumerable(_rows);
+
+            IDataView shuffledData = _context.Data.ShuffleRows(data, seed: 0);
+
+            foreach (string sample in shuffledData.GetColumn<string>("Sample"))
+            {
+            }
+        }
+
+        private class TrainRow
+        {
+            [ColumnName("Sample")]
+            public string Sample;
+
+            [ColumnName("Week")]
+            public float Week;
+
+            [ColumnName("Label")]
+            public float Label;
+        }
+    }
+}