From b3f132efd1ac38694e20df270be681890bf225ca Mon Sep 17 00:00:00 2001
From: Dylan Thacker-Smith <Dylan.Smith@shopify.com>
Date: Tue, 17 Nov 2020 13:53:35 -0500
Subject: [PATCH] Fix total_render_time if a Profiler gets used for multiple
 renders

---
 lib/liquid/profiler.rb                    | 18 +++++++++++------
 lib/liquid/template.rb                    | 18 ++++++-----------
 test/integration/render_profiling_test.rb | 24 +++++++++++++++++++++++
 3 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/lib/liquid/profiler.rb b/lib/liquid/profiler.rb
index cde2d5b..7a7c376 100644
--- a/lib/liquid/profiler.rb
+++ b/lib/liquid/profiler.rb
@@ -90,14 +90,20 @@ module Liquid
     def initialize
       @root_timing  = Timing.new("", nil)
       @timing_stack = [@root_timing]
+      @render_start_at = nil
+      @total_render_time = 0.0
     end
 
-    def start
-      @render_start_at = monotonic_time
-    end
-
-    def stop
-      @total_render_time = monotonic_time - @render_start_at
+    def profile
+      return yield if @render_start_at
+      started_at = monotonic_time
+      begin
+        @started_at = started_at
+        yield
+      ensure
+        @started_at = nil
+        @total_render_time += monotonic_time - started_at
+      end
     end
 
     def each(&block)
diff --git a/lib/liquid/template.rb b/lib/liquid/template.rb
index 6818ba4..da28ca4 100644
--- a/lib/liquid/template.rb
+++ b/lib/liquid/template.rb
@@ -193,6 +193,10 @@ module Liquid
       # Retrying a render resets resource usage
       context.resource_limits.reset
 
+      if @profiling && context.profiler.nil?
+        @profiler = context.profiler = Liquid::Profiler.new
+      end
+
       begin
         # render the nodelist.
         with_profiling(context) do
@@ -221,18 +225,8 @@ module Liquid
     end
 
     def with_profiling(context)
-      if @profiling && context.profiler.nil?
-        @profiler = context.profiler = Profiler.new
-        @profiler.start
-
-        begin
-          yield
-        ensure
-          @profiler.stop
-        end
-      else
-        yield
-      end
+      return yield unless context.profiler
+      context.profiler.profile { yield }
     end
 
     def apply_options_to_context(context, options)
diff --git a/test/integration/render_profiling_test.rb b/test/integration/render_profiling_test.rb
index dc9f001..67ba213 100644
--- a/test/integration/render_profiling_test.rb
+++ b/test/integration/render_profiling_test.rb
@@ -88,6 +88,30 @@ class RenderProfilingTest < Minitest::Test
     assert(t.profiler.total_render_time >= 0, "Total render time was not calculated")
   end
 
+  class SleepTag < Liquid::Tag
+    def initialize(tag_name, markup, parse_context)
+      super
+      @duration = Float(markup)
+    end
+
+    def render_to_output_buffer(_context, _output)
+      sleep(@duration)
+    end
+  end
+
+  def test_profiling_multiple_renders
+    with_custom_tag('sleep', SleepTag) do
+      context = Liquid::Context.new
+      t = Liquid::Template.parse("{% sleep 0.001 %}", profile: true)
+      t.render!(context)
+      first_render_time = context.profiler.total_render_time
+      t.render!(context)
+
+      assert_operator(first_render_time, :>=, 0.001)
+      assert_operator(context.profiler.total_render_time, :>=, 0.001 + first_render_time)
+    end
+  end
+
   def test_profiling_uses_include_to_mark_children
     t = Template.parse("{{ 'a string' | upcase }}\n{% include 'a_template' %}", profile: true)
     t.render!