From 60214b957c8c4c7d7242744304f734bc75eb31e1 Mon Sep 17 00:00:00 2001
From: Dylan Thacker-Smith <dylan.smith@shopify.com>
Date: Wed, 9 Dec 2020 10:04:20 -0500
Subject: [PATCH] Store the profiler in the context instead of a thread-local
 variable (#1364)

---
 lib/liquid/profiler.rb       | 28 +++++++++-------------------
 lib/liquid/profiler/hooks.rb | 17 ++++++++++++++++-
 lib/liquid/template.rb       | 14 ++++++++------
 3 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/lib/liquid/profiler.rb b/lib/liquid/profiler.rb
index 762aa8d..cde2d5b 100644
--- a/lib/liquid/profiler.rb
+++ b/lib/liquid/profiler.rb
@@ -85,21 +85,6 @@ module Liquid
       end
     end
 
-    def self.profile_node_render(node, template_name)
-      if Profiler.current_profile && node.respond_to?(:render)
-        Profiler.current_profile.start_node(node, template_name)
-        output = yield
-        Profiler.current_profile.end_node
-        output
-      else
-        yield
-      end
-    end
-
-    def self.current_profile
-      Thread.current[:liquid_profiler]
-    end
-
     attr_reader :total_render_time
 
     def initialize
@@ -108,12 +93,10 @@ module Liquid
     end
 
     def start
-      Thread.current[:liquid_profiler] = self
       @render_start_at = monotonic_time
     end
 
     def stop
-      Thread.current[:liquid_profiler] = nil
       @total_render_time = monotonic_time - @render_start_at
     end
 
@@ -129,6 +112,15 @@ module Liquid
       @root_timing.children.length
     end
 
+    def profile_node(node, template_name)
+      start_node(node, template_name)
+      yield
+    ensure
+      end_node
+    end
+
+    private
+
     def start_node(node, template_name)
       @timing_stack.push(Timing.start(node, template_name))
     end
@@ -140,8 +132,6 @@ module Liquid
       @timing_stack.last.children << timing
     end
 
-    private
-
     def monotonic_time
       Process.clock_gettime(Process::CLOCK_MONOTONIC)
     end
diff --git a/lib/liquid/profiler/hooks.rb b/lib/liquid/profiler/hooks.rb
index b6cc36e..b07c580 100644
--- a/lib/liquid/profiler/hooks.rb
+++ b/lib/liquid/profiler/hooks.rb
@@ -3,10 +3,25 @@
 module Liquid
   module BlockBodyProfilingHook
     def render_node(context, output, node)
-      Profiler.profile_node_render(node, context.template_name) do
+      if (profiler = context.profiler)
+        profiler.profile_node(node, context.template_name) do
+          super
+        end
+      else
         super
       end
     end
   end
   BlockBody.prepend(BlockBodyProfilingHook)
+
+  module ContextProfilingHook
+    attr_accessor :profiler
+
+    def new_isolated_subcontext
+      new_context = super
+      new_context.profiler = profiler
+      new_context
+    end
+  end
+  Context.prepend(ContextProfilingHook)
 end
diff --git a/lib/liquid/template.rb b/lib/liquid/template.rb
index fe47623..6818ba4 100644
--- a/lib/liquid/template.rb
+++ b/lib/liquid/template.rb
@@ -106,9 +106,13 @@ module Liquid
     # Parse source code.
     # Returns self for easy chaining
     def parse(source, options = {})
+      if (profiling = options[:profile])
+        raise "Profiler not loaded, require 'liquid/profiler' first" unless defined?(Liquid::Profiler)
+      end
+
       @options      = options
-      @profiling    = options[:profile]
-      @line_numbers = options[:line_numbers] || @profiling
+      @profiling    = profiling
+      @line_numbers = options[:line_numbers] || profiling
       parse_context = options.is_a?(ParseContext) ? options : ParseContext.new(options)
       @root         = Document.parse(tokenize(source), parse_context)
       @warnings     = parse_context.warnings
@@ -217,10 +221,8 @@ module Liquid
     end
 
     def with_profiling(context)
-      if @profiling && !context.partial
-        raise "Profiler not loaded, require 'liquid/profiler' first" unless defined?(Liquid::Profiler)
-
-        @profiler = Profiler.new
+      if @profiling && context.profiler.nil?
+        @profiler = context.profiler = Profiler.new
         @profiler.start
 
         begin