Merge branch 'profiling' into pipelines

StackExchange.Redis.StrongName/TypeForwarding.cs

@@ -36,8 +36,7 @@
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IConnectionMultiplexer))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IDatabase))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IDatabaseAsync))]
-[assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IProfiledCommand))]
-[assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IProfiler))]
+[assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.Profiling.IProfiledCommand))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IReconnectRetryPolicy))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IRedis))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.IRedisAsync))]
@@ -52,7 +51,8 @@
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.LuaScript))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.MigrateOptions))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.Order))]
-[assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.ProfiledCommandEnumerable))]
+[assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.Profiling.ProfiledCommandEnumerable))]
+[assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.Profiling.ProfilingSession))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.Proxy))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.RedisChannel))]
 [assembly: TypeForwardedTo(typeof(global::StackExchange.Redis.RedisCommandException))]

StackExchange.Redis.Tests/Cluster.cs

@@ -5,6 +5,7 @@
 using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
+using StackExchange.Redis.Profiling;
 using Xunit;
 using Xunit.Abstractions;
 
@@ -580,34 +581,31 @@ public void GetFromRightNodeBasedOnFlags(CommandFlags flags, bool isSlave)
 
         private static string Describe(EndPoint endpoint) => endpoint?.ToString() ?? "(unknown)";
 
-        private class TestProfiler : IProfiler
-        {
-            public object MyContext = new object();
-            public object GetContext() => MyContext;
-        }
-
         [Fact]
         public void SimpleProfiling()
         {
             using (var conn = Create())
             {
-                var profiler = new TestProfiler();
+                var profiler = new ProfilingSession();
                 var key = Me();
                 var db = conn.GetDatabase();
                 db.KeyDelete(key, CommandFlags.FireAndForget);
 
-                conn.RegisterProfiler(profiler);
-                conn.BeginProfiling(profiler.MyContext);
+                conn.RegisterProfiler(() => profiler);
                 db.StringSet(key, "world");
                 var val = db.StringGet(key);
                 Assert.Equal("world", val);
 
-                var msgs = conn.FinishProfiling(profiler.MyContext);
+                var msgs = profiler.GetCommands();
                 Log("Checking GET...");
                 Assert.Contains(msgs, m => m.Command == "GET");
                 Log("Checking SET...");
                 Assert.Contains(msgs, m => m.Command == "SET");
                 Assert.Equal(2, msgs.Count());
+
+                var arr = msgs.ToArray();
+                Assert.Equal("SET", arr[0].Command);
+                Assert.Equal("GET", arr[1].Command);
             }
         }
 

StackExchange.Redis/StackExchange/Redis/ConnectionMultiplexer.Profiling.cs

 using System;
+using StackExchange.Redis.Profiling;
 
 namespace StackExchange.Redis
 {
     public partial class ConnectionMultiplexer
     {
-        private IProfiler profiler;
-
-        // internal for test purposes
-        internal ProfileContextTracker profiledCommands;
-
-        /// <summary>
-        /// <para>Sets an IProfiler instance for this ConnectionMultiplexer.</para>
-        /// <para>
-        /// An IProfiler instances is used to determine which context to associate an
-        /// IProfiledCommand with.  See BeginProfiling(object) and FinishProfiling(object)
-        /// for more details.
-        /// </para>
-        /// </summary>
-        /// <param name="profiler">The profiler to register.</param>
-        public void RegisterProfiler(IProfiler profiler)
-        {
-            if (this.profiler != null) throw new InvalidOperationException("IProfiler already registered for this ConnectionMultiplexer");
-
-            this.profiler = profiler ?? throw new ArgumentNullException(nameof(profiler));
-            profiledCommands = new ProfileContextTracker();
-        }
-
-        /// <summary>
-        /// <para>Begins profiling for the given context.</para>
-        /// <para>
-        /// If the same context object is returned by the registered IProfiler, the IProfiledCommands
-        /// will be associated with each other.
-        /// </para>
-        /// <para>Call FinishProfiling with the same context to get the assocated commands.</para>
-        /// <para>Note that forContext cannot be a WeakReference or a WeakReference&lt;T</para>&gt;
-        /// </summary>
-        /// <param name="forContext">The context to begin profiling.</param>
-        public void BeginProfiling(object forContext)
-        {
-            if (profiler == null) throw new InvalidOperationException("Cannot begin profiling if no IProfiler has been registered with RegisterProfiler");
-            if (forContext == null) throw new ArgumentNullException(nameof(forContext));
-            if (forContext is WeakReference) throw new ArgumentException("Context object cannot be a WeakReference", nameof(forContext));
-
-            if (!profiledCommands.TryCreate(forContext))
-            {
-                throw ExceptionFactory.BeganProfilingWithDuplicateContext(forContext);
-            }
-        }
+        Func<ProfilingSession> _profilingSessionProvider;
 
         /// <summary>
-        /// <para>Stops profiling for the given context, returns all IProfiledCommands associated.</para>
-        /// <para>By default this may do a sweep for dead profiling contexts, you can disable this by passing "allowCleanupSweep: false".</para>
+        /// Register a callback to provide an on-demand ambient session provider based on the
+        /// calling context; the implementing code is responsible for reliably resolving the same provider
+        /// based on ambient context, or returning null to not profile
         /// </summary>
-        /// <param name="forContext">The context to begin profiling.</param>
-        /// <param name="allowCleanupSweep">Whether to allow cleanup of old profiling sessions.</param>
-        public ProfiledCommandEnumerable FinishProfiling(object forContext, bool allowCleanupSweep = true)
-        {
-            if (profiler == null) throw new InvalidOperationException("Cannot begin profiling if no IProfiler has been registered with RegisterProfiler");
-            if (forContext == null) throw new ArgumentNullException(nameof(forContext));
-
-            if (!profiledCommands.TryRemove(forContext, out ProfiledCommandEnumerable ret))
-            {
-                throw ExceptionFactory.FinishedProfilingWithInvalidContext(forContext);
-            }
-
-            // conditional, because it could hurt and that may sometimes be unacceptable
-            if (allowCleanupSweep)
-            {
-                profiledCommands.TryCleanup();
-            }
-
-            return ret;
-        }
+        public void RegisterProfiler(Func<ProfilingSession> profilingSessionProvider) => _profilingSessionProvider = profilingSessionProvider;
     }
 }

StackExchange.Redis/StackExchange/Redis/ConnectionMultiplexer.cs

@@ -11,6 +11,7 @@
 using System.Reflection;
 using System.IO.Compression;
 using System.Runtime.CompilerServices;
+using StackExchange.Redis.Profiling;
 
 namespace StackExchange.Redis
 {
@@ -1868,10 +1869,10 @@ private WriteResult TryPushMessageToBridge<T>(Message message, ResultProcessor<T
 
             if (server != null)
             {
-                var profCtx = profiler?.GetContext();
-                if (profCtx != null && profiledCommands.TryGetValue(profCtx, out ConcurrentProfileStorageCollection inFlightForCtx))
+                var profilingSession = _profilingSessionProvider?.Invoke();
+                if (profilingSession != null)
                 {
-                    message.SetProfileStorage(ProfileStorage.NewWithContext(inFlightForCtx, server));
+                    message.SetProfileStorage(ProfiledCommand.NewWithContext(profilingSession, server));
                 }
 
                 if (message.Db >= 0)
@@ -1952,6 +1953,7 @@ public bool IsConnecting
         public void Close(bool allowCommandsToComplete = true)
         {
             isDisposed = true;
+            _profilingSessionProvider = null;
             using (var tmp = pulse)
             {
                 pulse = null;

StackExchange.Redis/StackExchange/Redis/Interfaces/IConnectionMultiplexer.cs

@@ -2,6 +2,7 @@
 using System.IO;
 using System.Net;
 using System.Threading.Tasks;
+using StackExchange.Redis.Profiling;
 
 namespace StackExchange.Redis
 {
@@ -69,36 +70,11 @@ public interface IConnectionMultiplexer
         int StormLogThreshold { get; set; }
 
         /// <summary>
-        /// Sets an IProfiler instance for this ConnectionMultiplexer.
-        /// 
-        /// An IProfiler instances is used to determine which context to associate an
-        /// IProfiledCommand with.  See BeginProfiling(object) and FinishProfiling(object)
-        /// for more details.
+        /// Register a callback to provide an on-demand ambient session provider based on the
+        /// calling context; the implementing code is responsible for reliably resolving the same provider
+        /// based on ambient context, or returning null to not profile
         /// </summary>
-        /// <param name="profiler">The profiler to register.</param>
-        void RegisterProfiler(IProfiler profiler);
-
-        /// <summary>
-        /// Begins profiling for the given context.
-        /// 
-        /// If the same context object is returned by the registered IProfiler, the IProfiledCommands
-        /// will be associated with each other.
-        /// 
-        /// Call FinishProfiling with the same context to get the assocated commands.
-        /// 
-        /// Note that forContext cannot be a WeakReference or a WeakReference&lt;T&gt;
-        /// </summary>
-        /// <param name="forContext">The context to begin profiling for.</param>
-        void BeginProfiling(object forContext);
-
-        /// <summary>
-        /// Stops profiling for the given context, returns all IProfiledCommands associated.
-        /// 
-        /// By default this may do a sweep for dead profiling contexts, you can disable this by passing "allowCleanupSweep: false".
-        /// </summary>
-        /// <param name="forContext">The context to finish profiling for.</param>
-        /// <param name="allowCleanupSweep">Whether to allow a cleanup sweep of dead profiling contexts.</param>
-        ProfiledCommandEnumerable FinishProfiling(object forContext, bool allowCleanupSweep = true);
+        void RegisterProfiler(Func<ProfilingSession> profilingSessionProvider);
 
         /// <summary>
         /// Get summary statistics associates with this server

StackExchange.Redis/StackExchange/Redis/Message.cs

@@ -5,6 +5,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using StackExchange.Redis.Profiling;
 
 namespace StackExchange.Redis
 {
@@ -81,7 +82,7 @@ internal abstract class Message : ICompletable
         private ResultProcessor resultProcessor;
 
         // All for profiling purposes
-        private ProfileStorage performance;
+        private ProfiledCommand performance;
         internal DateTime createdDateTime;
         internal long createdTimestamp;
 
@@ -135,7 +136,7 @@ internal void SetMasterOnly()
             }
         }
 
-        internal void SetProfileStorage(ProfileStorage storage)
+        internal void SetProfileStorage(ProfiledCommand storage)
         {
             performance = storage;
             performance.SetMessage(this);
@@ -152,7 +153,7 @@ internal void PrepareToResend(ServerEndPoint resendTo, bool isMoved)
 
             createdDateTime = DateTime.UtcNow;
             createdTimestamp = System.Diagnostics.Stopwatch.GetTimestamp();
-            performance = ProfileStorage.NewAttachedToSameContext(oldPerformance, resendTo, isMoved);
+            performance = ProfiledCommand.NewAttachedToSameContext(oldPerformance, resendTo, isMoved);
             performance.SetMessage(this);
             Status = CommandStatus.WaitingToBeSent;
         }
@@ -437,33 +438,32 @@ public override string ToString()
             return $"[{Db}]:{CommandAndKey} ({resultProcessor?.GetType().Name ?? "(n/a)"})";
         }
 
-        public void SetResponseReceived()
-        {
-            performance?.SetResponseReceived();
-        }
+        public void SetResponseReceived() => performance?.SetResponseReceived();
 
         public bool TryComplete(bool isAsync)
         {
             //Ensure we can never call TryComplete on the same resultBox from two threads by grabbing it now
             var currBox = Interlocked.Exchange(ref resultBox, null);
+
+            if (!isAsync)
+            {   // set the performance completion the first chance we get (sync comes first)
+                performance?.SetCompleted();
+            }
             if (currBox != null)
             {
                 var ret = currBox.TryComplete(isAsync);
 
                 //in async mode TryComplete will have unwrapped and recycled resultBox
-                if (!(ret && isAsync))
+                if (!(ret || isAsync))
                 {
                     //put result box back if it was not already recycled
                     Interlocked.Exchange(ref resultBox, currBox);
                 }
-
-                performance?.SetCompleted();
                 return ret;
             }
             else
             {
                 ConnectionMultiplexer.TraceWithoutContext("No result-box to complete for " + Command, "Message");
-                performance?.SetCompleted();
                 return true;
             }
         }
@@ -612,10 +612,7 @@ internal void SetException(Exception exception)
             resultBox?.SetException(exception);
         }
 
-        internal void SetEnqueued()
-        {
-            performance?.SetEnqueued();
-        }
+        internal void SetEnqueued()=> performance?.SetEnqueued();
 
 
 

StackExchange.Redis/StackExchange/Redis/ProfileContextTracker.cs

-using System;
-using System.Collections.Concurrent;
-using System.Collections.Generic;
-using System.Threading;
-
-namespace StackExchange.Redis
-{
-    /// <summary>
-    /// Big ol' wrapper around most of the profiling storage logic, 'cause it got too big to just live in ConnectionMultiplexer.
-    /// </summary>
-    internal sealed class ProfileContextTracker
-    {
-        /// <summary>
-        /// <para>Necessary, because WeakReference can't be readily comparable (since the reference is... weak).</para>
-        /// <para>This lets us detect leaks* with some reasonable confidence, and cleanup periodically.</para>
-        /// <para>
-        /// Some calisthenics are done to avoid allocating WeakReferences for no reason, as often
-        /// we're just looking up ProfileStorage.
-        /// </para>
-        /// <para>* Somebody starts profiling, but for whatever reason never *stops* with a context object</para>
-        /// </summary>
-        private readonly struct ProfileContextCell : IEquatable<ProfileContextCell>
-        {
-            // This is a union of (object|WeakReference); if it's a WeakReference
-            //   then we're actually interested in it's Target, otherwise
-            //   we're concerned about the actual value of Reference
-            private readonly object Reference;
-
-            // It is absolutely crucial that this value **never change** once instantiated
-            private readonly int HashCode;
-
-            public bool IsContextLeaked => !TryGetTarget(out _);
-
-            private ProfileContextCell(object forObj, bool isEphemeral)
-            {
-                HashCode = forObj.GetHashCode();
-
-                if (isEphemeral)
-                {
-                    Reference = forObj;
-                }
-                else
-                {
-                    Reference = new WeakReference(forObj, trackResurrection: true); // ughhh, have to handle finalizers
-                }
-            }
-
-            /// <summary>
-            /// <para>Suitable for use as a key into something.</para>
-            /// <para>
-            /// This instance **WILL NOT** keep forObj alive, so it can
-            /// be copied out of the calling method's scope.
-            /// </para>
-            /// </summary>
-            /// <param name="forObj">The object to get a context for.</param>
-            public static ProfileContextCell ToStoreUnder(object forObj) => new ProfileContextCell(forObj, isEphemeral: false);
-
-            /// <summary>
-            /// <para>Only suitable for looking up.</para>
-            /// <para>
-            /// This instance **ABSOLUTELY WILL** keep forObj alive, so this
-            /// had better not be copied into anything outside the scope of the
-            /// calling method.
-            /// </para>
-            /// </summary>
-            /// <param name="forObj">The object to lookup a context by.</param>
-            public static ProfileContextCell ToLookupBy(object forObj) => new ProfileContextCell(forObj, isEphemeral: true);
-
-            private bool TryGetTarget(out object target)
-            {
-                var asWeakRef = Reference as WeakReference;
-
-                if (asWeakRef == null)
-                {
-                    target = Reference;
-                    return true;
-                }
-
-                // Do not use IsAlive here, it's race city
-                target = asWeakRef.Target;
-                return target != null;
-            }
-
-            public override bool Equals(object obj)
-            {
-                if (!(obj is ProfileContextCell)) return false;
-
-                return Equals((ProfileContextCell)obj);
-            }
-
-            public override int GetHashCode() => HashCode;
-
-            public bool Equals(ProfileContextCell other)
-            {
-                if (other.TryGetTarget(out object otherObj) != TryGetTarget(out object thisObj)) return false;
-
-                // dead references are equal
-                if (thisObj == null) return true;
-
-                return thisObj.Equals(otherObj);
-            }
-        }
-
-        // provided so default behavior doesn't do any boxing, for sure
-        private sealed class ProfileContextCellComparer : IEqualityComparer<ProfileContextCell>
-        {
-            public static readonly ProfileContextCellComparer Singleton = new ProfileContextCellComparer();
-
-            private ProfileContextCellComparer() { }
-
-            public bool Equals(ProfileContextCell x, ProfileContextCell y)
-            {
-                return x.Equals(y);
-            }
-
-            public int GetHashCode(ProfileContextCell obj)
-            {
-                return obj.GetHashCode();
-            }
-        }
-
-        private long lastCleanupSweep;
-        private readonly ConcurrentDictionary<ProfileContextCell, ConcurrentProfileStorageCollection> profiledCommands;
-
-        public int ContextCount => profiledCommands.Count;
-
-        public ProfileContextTracker()
-        {
-            profiledCommands = new ConcurrentDictionary<ProfileContextCell, ConcurrentProfileStorageCollection>(ProfileContextCellComparer.Singleton);
-            lastCleanupSweep = DateTime.UtcNow.Ticks;
-        }
-
-        /// <summary>
-        /// <para>Registers the passed context with a collection that can be retried with subsequent calls to TryGetValue.</para>
-        /// <para>Returns false if the passed context object is already registered.</para>
-        /// </summary>
-        /// <param name="ctx">The context to use.</param>
-        public bool TryCreate(object ctx)
-        {
-            var cell = ProfileContextCell.ToStoreUnder(ctx);
-
-            // we can't pass this as a delegate, because TryAdd may invoke the factory multiple times,
-            //   which would lead to over allocation.
-            var storage = ConcurrentProfileStorageCollection.GetOrCreate();
-            return profiledCommands.TryAdd(cell, storage);
-        }
-
-        /// <summary>
-        /// <para>
-        /// Returns true and sets val to the tracking collection associated with the given context if the context
-        /// was registered with TryCreate.
-        /// </para>
-        /// <para>Otherwise returns false and sets val to null.</para>
-        /// </summary>
-        /// <param name="ctx">The context to get a value for.</param>
-        /// <param name="val">The collection (if present) for <paramref name="ctx"/>.</param>
-        public bool TryGetValue(object ctx, out ConcurrentProfileStorageCollection val)
-        {
-            var cell = ProfileContextCell.ToLookupBy(ctx);
-            return profiledCommands.TryGetValue(cell, out val);
-        }
-
-        /// <summary>
-        /// <para>
-        /// Removes a context, setting all commands to a (non-thread safe) enumerable of
-        /// all the commands attached to that context.
-        /// </para>
-        /// <para>If the context was never registered, will return false and set commands to null.</para>
-        /// <para>
-        /// Subsequent calls to TryRemove with the same context will return false unless it is
-        /// re-registered with TryCreate.
-        /// </para>
-        /// </summary>
-        /// <param name="ctx">The context to remove for.</param>
-        /// <param name="commands">The commands to remove.</param>
-        public bool TryRemove(object ctx, out ProfiledCommandEnumerable commands)
-        {
-            var cell = ProfileContextCell.ToLookupBy(ctx);
-            if (!profiledCommands.TryRemove(cell, out ConcurrentProfileStorageCollection storage))
-            {
-                commands = default(ProfiledCommandEnumerable);
-                return false;
-            }
-
-            commands = storage.EnumerateAndReturnForReuse();
-            return true;
-        }
-
-        /// <summary>
-        /// If enough time has passed (1 minute) since the last call, this does walk of all contexts
-        /// and removes those that the GC has collected.
-        /// </summary>
-        public bool TryCleanup()
-        {
-            const long SweepEveryTicks = 600000000; // once a minute, tops
-
-            var now = DateTime.UtcNow.Ticks;    // resolution on this isn't great, but it's good enough
-            var last = lastCleanupSweep;
-            var since = now - last;
-            if (since < SweepEveryTicks) return false;
-
-            // this is just to keep other threads from wasting time, in theory
-            //  it'd be perfectly safe for this to run concurrently
-            var saw = Interlocked.CompareExchange(ref lastCleanupSweep, now, last);
-            if (saw != last) return false;
-
-            if (profiledCommands.Count == 0) return false;
-
-            using (var e = profiledCommands.GetEnumerator())
-            {
-                while (e.MoveNext())
-                {
-                    var pair = e.Current;
-                    if (pair.Key.IsContextLeaked && profiledCommands.TryRemove(pair.Key, out ConcurrentProfileStorageCollection abandoned))
-                    {
-                        // shove it back in the pool, but don't bother enumerating
-                        abandoned.ReturnForReuse();
-                    }
-                }
-            }
-
-            return true;
-        }
-    }
-}

StackExchange.Redis/StackExchange/Redis/Interfaces/IProfiler.cs → StackExchange.Redis/StackExchange/Redis/Profiling/IProfiledCommand.cs

 using System;
 using System.Net;
 
-namespace StackExchange.Redis
+namespace StackExchange.Redis.Profiling
 {
     /// <summary>
     /// <para>A profiled command against a redis instance.</para>
@@ -89,21 +89,4 @@ public interface IProfiledCommand
         /// </summary>
         RetransmissionReasonType? RetransmissionReason { get; }
     }
-
-    /// <summary>
-    /// Interface for profiling individual commands against an Redis ConnectionMulitplexer.
-    /// </summary>
-    public interface IProfiler
-    {
-        /// <summary>
-        /// Called to provide a context object.
-        /// 
-        /// This method is called before the method which triggers work against redis (such as StringSet(Async)) returns,
-        /// and will always be called on the same thread as that method.
-        /// 
-        /// Note that GetContext() may be called even if ConnectionMultiplexer.BeginProfiling() has not been called.
-        /// You may return `null` to prevent any tracking of commands.
-        /// </summary>
-        object GetContext();
-    }
 }

StackExchange.Redis/StackExchange/Redis/ProfileStorage.cs → StackExchange.Redis/StackExchange/Redis/Profiling/ProfiledCommand.cs

 using System;
 using System.Diagnostics;
 using System.Net;
+using System.Runtime.CompilerServices;
 using System.Threading;
 
-namespace StackExchange.Redis
+namespace StackExchange.Redis.Profiling
 {
-    internal class ProfileStorage : IProfiledCommand
+    internal sealed class ProfiledCommand : IProfiledCommand
     {
         #region IProfiledCommand Impl
         public EndPoint EndPoint => Server.EndPoint;
 
         public int Db => Message.Db;
 
-        public string Command => Message.Command.ToString();
+        public string Command => Message is RedisDatabase.ExecuteMessage em ? em.Command : Message.Command.ToString();
 
         public CommandFlags Flags => Message.Flags;
 
@@ -34,11 +35,11 @@ internal class ProfileStorage : IProfiledCommand
 
         #endregion
 
-        public ProfileStorage NextElement { get; set; }
+        public ProfiledCommand NextElement { get; set; }
 
         private Message Message;
         private readonly ServerEndPoint Server;
-        private readonly ProfileStorage OriginalProfiling;
+        private readonly ProfiledCommand OriginalProfiling;
 
         private DateTime MessageCreatedDateTime;
         private long MessageCreatedTimeStamp;
@@ -47,9 +48,9 @@ internal class ProfileStorage : IProfiledCommand
         private long ResponseReceivedTimeStamp;
         private long CompletedTimeStamp;
 
-        private readonly ConcurrentProfileStorageCollection PushToWhenFinished;
+        private readonly ProfilingSession PushToWhenFinished;
 
-        private ProfileStorage(ConcurrentProfileStorageCollection pushTo, ServerEndPoint server, ProfileStorage resentFor, RetransmissionReasonType? reason)
+        private ProfiledCommand(ProfilingSession pushTo, ServerEndPoint server, ProfiledCommand resentFor, RetransmissionReasonType? reason)
         {
             PushToWhenFinished = pushTo;
             OriginalProfiling = resentFor;
@@ -57,14 +58,14 @@ private ProfileStorage(ConcurrentProfileStorageCollection pushTo, ServerEndPoint
             RetransmissionReason = reason;
         }
 
-        public static ProfileStorage NewWithContext(ConcurrentProfileStorageCollection pushTo, ServerEndPoint server)
+        public static ProfiledCommand NewWithContext(ProfilingSession pushTo, ServerEndPoint server)
         {
-            return new ProfileStorage(pushTo, server, null, null);
+            return new ProfiledCommand(pushTo, server, null, null);
         }
 
-        public static ProfileStorage NewAttachedToSameContext(ProfileStorage resentFor, ServerEndPoint server, bool isMoved)
+        public static ProfiledCommand NewAttachedToSameContext(ProfiledCommand resentFor, ServerEndPoint server, bool isMoved)
         {
-            return new ProfileStorage(resentFor.PushToWhenFinished, server, resentFor, isMoved ? RetransmissionReasonType.Moved : RetransmissionReasonType.Ask);
+            return new ProfiledCommand(resentFor.PushToWhenFinished, server, resentFor, isMoved ? RetransmissionReasonType.Moved : RetransmissionReasonType.Ask);
         }
 
         public void SetMessage(Message msg)
@@ -77,27 +78,17 @@ public void SetMessage(Message msg)
             MessageCreatedTimeStamp = msg.createdTimestamp;
         }
 
-        public void SetEnqueued()
-        {
-            // This method should never be called twice
-            if (EnqueuedTimeStamp > 0) throw new InvalidOperationException($"{nameof(SetEnqueued)} called more than once");
-
-            EnqueuedTimeStamp = Stopwatch.GetTimestamp();
-        }
+        public void SetEnqueued() => SetTimestamp(ref EnqueuedTimeStamp);
 
-        public void SetRequestSent()
-        {
-            // This method should never be called twice
-            if (RequestSentTimeStamp > 0) throw new InvalidOperationException($"{nameof(SetRequestSent)} called more than once");
+        public void SetRequestSent() => SetTimestamp(ref RequestSentTimeStamp);
 
-            RequestSentTimeStamp = Stopwatch.GetTimestamp();
-        }
+        public void SetResponseReceived() => SetTimestamp(ref ResponseReceivedTimeStamp);
 
-        public void SetResponseReceived()
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static void SetTimestamp(ref long field)
         {
-            if (ResponseReceivedTimeStamp > 0) throw new InvalidOperationException($"{nameof(SetResponseReceived)} called more than once");
-
-            ResponseReceivedTimeStamp = Stopwatch.GetTimestamp();
+            var now = Stopwatch.GetTimestamp();
+            Interlocked.CompareExchange(ref field, now, 0);
         }
 
         public void SetCompleted()
@@ -108,11 +99,13 @@ public void SetCompleted()
             var now = Stopwatch.GetTimestamp();
             var oldVal = Interlocked.CompareExchange(ref CompletedTimeStamp, now, 0);
 
-            // second call
-            if (oldVal != 0) return;
-
             // only push on the first call, no dupes!
-            PushToWhenFinished.Add(this);
+            if (oldVal == 0)
+            {
+                // fake a response if we completed prematurely (timeout, broken connection, etc)
+                Interlocked.CompareExchange(ref ResponseReceivedTimeStamp, now, 0);
+                PushToWhenFinished?.Add(this);
+            }
         }
 
         public override string ToString()

StackExchange.Redis/StackExchange/Redis/ConcurrentProfileStorageCollection.cs → StackExchange.Redis/StackExchange/Redis/Profiling/ProfiledCommandEnumerable.cs

 using System.Collections.Generic;
-using System.Threading;
 
-namespace StackExchange.Redis
+namespace StackExchange.Redis.Profiling
 {
     /// <summary>
     /// <para>A collection of IProfiledCommands.</para>
@@ -24,13 +23,13 @@ namespace StackExchange.Redis
         /// </summary>
         public struct Enumerator : IEnumerator<IProfiledCommand>
         {
-            private ProfileStorage Head;
-            private ProfileStorage CurrentBacker;
+            private ProfiledCommand Head;
+            private ProfiledCommand CurrentBacker;
 
             private bool IsEmpty => Head == null;
             private bool IsUnstartedOrFinished => CurrentBacker == null;
 
-            internal Enumerator(ProfileStorage head)
+            internal Enumerator(ProfiledCommand head)
             {
                 Head = head;
                 CurrentBacker = null;
@@ -81,9 +80,9 @@ public void Dispose()
             }
         }
 
-        private readonly ProfileStorage Head;
+        private readonly ProfiledCommand Head;
 
-        internal ProfiledCommandEnumerable(ProfileStorage head)
+        internal ProfiledCommandEnumerable(ProfiledCommand head)
         {
             Head = head;
         }
@@ -101,117 +100,4 @@ internal ProfiledCommandEnumerable(ProfileStorage head)
 
         System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() => GetEnumerator();
     }
-
-    /// <summary>
-    /// <para>
-    /// A thread-safe collection tailored to the "always append, with high contention, then enumerate once with no contention"
-    /// behavior of our profiling.
-    /// </para>
-    /// <para>Performs better than ConcurrentBag, which is important since profiling code shouldn't impact timings.</para>
-    /// </summary>
-    internal sealed class ConcurrentProfileStorageCollection
-    {
-        // internal for test purposes
-        internal static int AllocationCount = 0;
-
-        // It is, by definition, impossible for an element to be in 2 intrusive collections
-        //   and we force Enumeration to release any reference to the collection object
-        //   so we can **always** pool these.
-        private const int PoolSize = 64;
-        private static readonly ConcurrentProfileStorageCollection[] Pool = new ConcurrentProfileStorageCollection[PoolSize];
-
-        private volatile ProfileStorage Head;
-
-        private ConcurrentProfileStorageCollection() { }
-
-        // for testing purposes only
-        internal static int CountInPool()
-        {
-            var ret = 0;
-
-            for (var i = 0; i < PoolSize; i++)
-            {
-                var inPool = Pool[i];
-                if (inPool != null) ret++;
-            }
-
-            return ret;
-        }
-
-        /// <summary>
-        /// <para>This method is thread-safe.</para>
-        /// <para>Adds an element to the bag.</para>
-        /// <para>Order is not preserved.</para>
-        /// <para>The element can only be a member of *one* bag.</para>
-        /// </summary>
-        /// <param name="command">The command to add.</param>
-        public void Add(ProfileStorage command)
-        {
-            while (true)
-            {
-                var cur = Head;
-                command.NextElement = cur;
-
-                // Interlocked references to volatile fields are perfectly cromulent
-#pragma warning disable 420
-                var got = Interlocked.CompareExchange(ref Head, command, cur);
-#pragma warning restore 420
-
-                if (object.ReferenceEquals(got, cur)) break;
-            }
-        }
-
-        /// <summary>
-        /// <para>
-        /// This method returns an enumerable view of the bag, and returns it to 
-        /// an internal pool for reuse by GetOrCreate().
-        /// </para>
-        /// <para>It is not thread safe.</para>
-        /// <para>It should only be called once the bag is finished being mutated.</para>
-        /// </summary>
-        public ProfiledCommandEnumerable EnumerateAndReturnForReuse()
-        {
-            var ret = new ProfiledCommandEnumerable(Head);
-
-            ReturnForReuse();
-
-            return ret;
-        }
-
-        /// <summary>
-        /// This returns the ConcurrentProfileStorageCollection to an internal pool for reuse by GetOrCreate().
-        /// </summary>
-        public void ReturnForReuse()
-        {
-            // no need for interlocking, this isn't a thread safe method
-            Head = null;
-
-            for (var i = 0; i < PoolSize; i++)
-            {
-                if (Interlocked.CompareExchange(ref Pool[i], this, null) == null) break;
-            }
-        }
-
-        /// <summary>
-        /// <para>Returns a ConcurrentProfileStorageCollection to use.</para>
-        /// <para>
-        /// It *may* have allocated a new one, or it may return one that has previously been released.
-        /// To return the collection, call EnumerateAndReturnForReuse()
-        /// </para>
-        /// </summary>
-        public static ConcurrentProfileStorageCollection GetOrCreate()
-        {
-            ConcurrentProfileStorageCollection found;
-            for (int i = 0; i < PoolSize; i++)
-            {
-                if ((found = Interlocked.Exchange(ref Pool[i], null)) != null)
-                {
-                    return found;
-                }
-            }
-
-            Interlocked.Increment(ref AllocationCount);
-            return new ConcurrentProfileStorageCollection();
-        }
-    }
 }

StackExchange.Redis/StackExchange/Redis/Profiling/ProfilingSession.cs

+using System.Threading;
+
+namespace StackExchange.Redis.Profiling
+{
+    /// <summary>
+    /// Lightweight profiling session that can be optionally registered (via ConnectionMultiplexer.RegisterProfiler) to track messages
+    /// </summary>
+    public sealed class ProfilingSession
+    {
+        /// <summary>
+        /// Caller-defined state object
+        /// </summary>
+        public object UserToken { get; }
+        /// <summary>
+        /// Create a new profiling session, optionally including a caller-defined state object
+        /// </summary>
+        public ProfilingSession(object userToken = null) => UserToken = userToken;
+
+        object _untypedHead;
+
+        internal void Add(ProfiledCommand command)
+        {
+            if (command == null) return;
+
+            object cur = Thread.VolatileRead(ref _untypedHead); ;
+            while (true)
+            {
+                command.NextElement = (ProfiledCommand)cur;
+                var got = Interlocked.CompareExchange(ref _untypedHead, command, cur);
+                if (ReferenceEquals(got, cur)) break; // successful update
+                cur = got; // retry; no need to re-fetch the field, we just did that
+            }
+        }
+
+        /// <summary>
+        /// Yield the commands that were captured as part of this session, resetting the session
+        /// </summary>
+        public ProfiledCommandEnumerable GetCommands()
+        {
+            var head = (ProfiledCommand)Interlocked.Exchange(ref _untypedHead, null);
+
+            // reverse the list so everything is ordered the way the consumer expected them
+            ProfiledCommand previous = null, current = head, next;
+            while(current != null)
+            {
+                next = current.NextElement;
+                current.NextElement = previous;
+                previous = current;
+                current = next;
+            }
+
+            return new ProfiledCommandEnumerable(previous);
+        }
+    }
+}

StackExchange.Redis/StackExchange/Redis/RedisDatabase.cs

@@ -3212,6 +3212,8 @@ internal sealed class ExecuteMessage : Message
         {
             private readonly string _command;
             private readonly ICollection<object> args;
+
+            public new string Command => _command;
             public ExecuteMessage(int db, CommandFlags flags, string command, ICollection<object> args) : base(db, flags, RedisCommand.UNKNOWN)
             {
                 _command = command;