Profile commands sent in a batch (#2175)

src/StackExchange.Redis/ConnectionMultiplexer.Profiling.cs

@@ -14,4 +14,18 @@ public partial class ConnectionMultiplexer
     /// </summary>
     /// <param name="profilingSessionProvider">The session provider to register.</param>
     public void RegisterProfiler(Func<ProfilingSession?> profilingSessionProvider) => _profilingSessionProvider = profilingSessionProvider;
+
+    /// <summary>
+    /// Gives the message somewhere to record its timings, if a profiler is registered and has a session for the
+    /// calling context. Every path that hands a message to a server calls this once the server is known: single
+    /// commands via <c>PrepareToPushMessageToBridge</c>, and batches, which go to the bridge directly (#2175).
+    /// </summary>
+    internal void AttachProfiling(Message message, ServerEndPoint server)
+    {
+        var profilingSession = _profilingSessionProvider?.Invoke();
+        if (profilingSession != null)
+        {
+            message.SetProfileStorage(ProfiledCommand.NewWithContext(profilingSession, server));
+        }
+    }
 }

src/StackExchange.Redis/ConnectionMultiplexer.cs

@@ -13,7 +13,6 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
-using StackExchange.Redis.Profiling;
 
 namespace StackExchange.Redis
 {
@@ -2629,11 +2628,7 @@ private bool PrepareToPushMessageToBridge<T>(Message message, ResultProcessor<T>
 
             if (server != null)
             {
-                var profilingSession = _profilingSessionProvider?.Invoke();
-                if (profilingSession != null)
-                {
-                    message.SetProfileStorage(ProfiledCommand.NewWithContext(profilingSession, server));
-                }
+                AttachProfiling(message, server);
 
                 if (message.Db >= 0)
                 {

src/StackExchange.Redis/RedisBatch.cs

@@ -75,6 +75,12 @@ public void Execute()
 
             foreach (var pair in byBridge)
             {
+                // attached here, not in the loop above: that loop can still abandon the whole batch for want of a
+                // server, and commands that never reach a bridge should not show up in a profiling session
+                foreach (var message in pair.Value)
+                {
+                    multiplexer.AttachProfiling(message, pair.Key.ServerEndPoint);
+                }
                 if (!pair.Key.TryEnqueue(pair.Value, pair.Key.ServerEndPoint.IsReplica))
                 {
                     FailNoServer(multiplexer, pair.Value);

tests/StackExchange.Redis.Tests/ProfilingTests.cs

@@ -95,6 +95,35 @@ private static void AssertProfiledCommandValues(IProfiledCommand command, IConne
         Assert.True(command.RetransmissionReason == null, nameof(command.RetransmissionReason));
     }
 
+    [Fact]
+    public async Task Batch()
+    {
+        await using var conn = Create();
+
+        var key = Me();
+        var session = new ProfilingSession();
+        conn.RegisterProfiler(() => session);
+
+        var dbId = TestConfig.GetDedicatedDB(conn);
+        var db = conn.GetDatabase(dbId);
+
+        // a batch is handed to the bridge directly, so it used to skip profiling altogether (#2175)
+        var batch = db.CreateBatch();
+        var set = batch.StringSetAsync(key, "world");
+        var get = batch.StringGetAsync(key);
+        var echo = batch.ExecuteAsync("ECHO", "fii");
+        batch.Execute();
+        await Task.WhenAll(set, get, echo);
+        Assert.Equal("world", await get);
+
+        var cmds = session.FinishProfiling();
+        Assert.Equal("ECHO,GET,SET", string.Join(",", cmds.Select(x => x.Command).OrderBy(x => x)));
+
+        AssertProfiledCommandValues(cmds.Single(cmd => cmd.Command == "SET"), conn, dbId);
+        AssertProfiledCommandValues(cmds.Single(cmd => cmd.Command == "GET"), conn, dbId);
+        AssertProfiledCommandValues(cmds.Single(cmd => cmd.Command == "ECHO"), conn, -1);
+    }
+
     [Fact]
     public async Task ManyThreads()
     {