Enumerable Timer with Stopwatch using an Extension Method

There have been numerous times when I’ve wanted to run a deferred query using either Linq-to-SQL (IQueryable<T>) or Linq-to-Objects (IEnumerable<T>) to capture the execution time for analysis.  Doing this for Linq-to-Objects is relatively easy using an extension method:

 public static IEnumerable<T> WithTimer<T>(this IEnumerable<T> source, Stopwatch watch)
 {
#if TRACE
  watch.Reset();
  watch.Start();
  foreach (var item in source)
   yield return item;
  watch.Stop();
#else
  return source;
#endif
 }

We simply pass in a StopWatch which is started and stopped pre- and post- enumeration respectively.  This allows us to get the total elapsed time outside of the enumeration as follows:

 Stopwatch watch = new Stopwatch();
 Enumerable.Range(0, 1000000).Where(n => n % 2 == 0).WithTimer(watch).ToList();
 Console.WriteLine(watch.ElapsedMilliseconds);

However, doing this for Linq-to-SQL is marginally more difficult as YIELD RETURN isn't supported for IQueryable<T> as it isn't an iterator interface type.  We therefore have to create our own Queryable wrapper (TimerQueryable<T>) for the IQueryable source and then wrap the enumerator which is returned by GetEnumerator into a TimerEnumerator<T> to control starting and stopping the timer.

    public class TimerQueryable<T> : IQueryable<T>
    {
        readonly IQueryable<T> _queryable;
        readonly IQueryProvider _provider;
        readonly Stopwatch _watch;

        public TimerQueryable(IQueryable<T> queryable, Stopwatch watch)
        {
            _watch = watch;
            _queryable = queryable;
            _provider = queryable.Provider;
        }

        public IEnumerator<T> GetEnumerator()
        {
            return new TimerEnumerator<T>(this._provider.CreateQuery<T>(Expression).GetEnumerator(), _watch);
        }

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return this.GetEnumerator();
        }

        public Type ElementType
        {
            get { return _queryable.ElementType; }
        }

        public System.Linq.Expressions.Expression Expression
        {
            get { return _queryable.Expression; }
        }

        public IQueryProvider Provider
        {
            get { return _queryable.Provider; }
        }
    }

    public class TimerEnumerator<T> : IEnumerator<T>
    {
        IEnumerator<T> _enumerator;
        Stopwatch _watch;

        public TimerEnumerator(IEnumerator<T> enumerator, Stopwatch watch)
        {
            _enumerator = enumerator;
            _watch = watch;
        }

        public T Current
        {
            get { return _enumerator.Current; }
        }

        public void Dispose()
        {
            _watch.Stop();
            _enumerator.Dispose();
        }

        object System.Collections.IEnumerator.Current
        {
            get { return this.Current; }
        }

        public bool MoveNext()
        {
            if (!_watch.IsRunning)
            {
                _watch.Reset();
                _watch.Start();
            }
            return this._enumerator.MoveNext();
        }

        public void Reset()
        {
            this._enumerator.Reset();
        }
    }

The StopWatch in the TimerEnumerator<T> class is started upon the enumeration in MoveNext and stopped upon the enumerator being disposed of in Dispose.  The timer isn't disposed of as we'd like to get the total elapsed time outside of the enumerator.  As IEnumerator<T> derives from IDisposable it would be fair to assume that the enumerator will be disposed correctly if the developer is adhering to the IDisposable pattern.

As there is only one instance of the StopWatch created and passed into multiple TimerEnumerator<T> instances it would be useful to mention that this isn't inherently thread-safe, but then again, do you know of an IQueryable or IEnumerable source which is?  Trying to execute two queries at the same time against an Entity Framework context in different threads would most certainly cause exceptions to be thrown.

Additionally, as we don't want this overhead in production code I have used the TRACE constant so that the timers will only work with the TRACE constant defined in the project - it doesn't make sense to run these timers when efficiency is paramount.

The full source code of the Enumerable Timer can be found on GitHub.