fix: prevent context loss by always setting thread and fiber locals (#497)

* Fix lost thread and fiber locals

Thread and fiber locals are being cleared on completion of every future.
If the promise ran on the same thread as the request, or the same thread
as the last promise, the locals are then not available. The execution
context hasn't changed in those circumstances so the code that sets the
locals at the start of a promise doesn't run.

1. Thread 1 runs
2. Sets context
3. Loads scope
4. Clears context
5. Thread 1 runs next scope
6. Context hasn’t changed so doesn’t set it again
7. Loads scope
8. Fails because the code relying on the context causes an exception
   because the locals aren't available.

Instead always set the thread and fiber locals and only clear those that
we know are new. This should avoid the lost locals and also avoid
memory leaks between promises.

See: #496

* Fix account for Fiber.current.storage not being available

* Fix account for Fiber.current.storage not being available 2/

Author: MattFenelon

lib/graphiti/scope.rb

@@ -140,7 +140,7 @@ def future_with_context(*args)
       end
       fiber_storage =
         if Fiber.current.respond_to?(:storage)
-          Fiber.current&.storage&.keys&.each_with_object({}) do |key, memo|
+          Fiber.current.storage&.keys&.each_with_object({}) do |key, memo|
             memo[key] = Fiber[key]
           end
         end
@@ -149,27 +149,46 @@ def future_with_context(*args)
         self.class.global_thread_pool_executor, Thread.current.object_id, thread_storage, fiber_storage, *args
       ) do |thread_id, thread_storage, fiber_storage, *args|
         wrap_in_rails_executor do
-          execution_context_changed = thread_id != Thread.current.object_id
-          if execution_context_changed
-            thread_storage&.keys&.each_with_object(Thread.current) do |key, thread_current|
-              thread_current[key] = thread_storage[key]
-            end
-            fiber_storage&.keys&.each_with_object(Fiber) do |key, fiber_current|
-              fiber_current[key] = fiber_storage[key]
+          with_thread_locals(thread_storage) do
+            with_fiber_locals(fiber_storage) do
+              Graphiti.broadcast(:global_thread_pool_task_run, self.class.global_thread_pool_stats) do
+                yield(*args)
+              end
             end
           end
+        end
+      end
+    end
 
-          result = Graphiti.broadcast(:global_thread_pool_task_run, self.class.global_thread_pool_stats) do
-            yield(*args)
-          end
+    def with_thread_locals(thread_locals)
+      new_thread_locals = []
+      thread_locals.each do |key, value|
+        if !Thread.current[key]
+          new_thread_locals << key
+        end
+        Thread.current[key] = value
+      end
+      yield
+    ensure
+      new_thread_locals.each do |key|
+        Thread.current[key] = nil
+      end
+    end
 
-          if execution_context_changed
-            thread_storage&.keys&.each { |key| Thread.current[key] = nil }
-            fiber_storage&.keys&.each { |key| Fiber[key] = nil }
-          end
+    def with_fiber_locals(fiber_locals)
+      return yield unless fiber_locals
 
-          result
+      new_fiber_locals = []
+      fiber_locals.each do |key, value|
+        if !Fiber[key]
+          new_fiber_locals << key
         end
+        Fiber[key] = value
+      end
+      yield
+    ensure
+      new_fiber_locals&.each do |key|
+        Fiber[key] = nil
       end
     end