Handle thread safety for Dotenv.restore, Add Dotenv.modify

lib/dotenv.rb

@@ -3,22 +3,21 @@
 require "dotenv/missing_keys"
 require "dotenv/diff"
 
-# The top level Dotenv module. The entrypoint for the application logic.
+# Shim to load environment variables from `.env files into `ENV`.
 module Dotenv
-  class << self
-    attr_accessor :instrumenter
-  end
+  extend self
+
+  # An internal monitor to synchronize access to ENV in multi-threaded environments.
+  SEMAPHORE = Monitor.new
+  private_constant :SEMAPHORE
 
-  module_function
+  attr_accessor :instrumenter
 
   # Loads environment variables from one or more `.env` files. See `#parse` for more details.
-  def load(*filenames, **kwargs)
-    parse(*filenames, **kwargs) do |env|
+  def load(*filenames, overwrite: false, ignore: true)
+    parse(*filenames, overwrite: overwrite, ignore: ignore) do |env|
       instrument(:load, env: env) do |payload|
-        env_before = ENV.to_h
-        env.apply
-        payload[:diff] = Dotenv::Diff.new(env_before, ENV.to_h)
-        env
+        update(env, overwrite: overwrite)
       end
     end
   end
@@ -33,14 +32,12 @@ def overwrite(*filenames)
     load(*filenames, overwrite: true)
   end
   alias_method :overload, :overwrite
-  module_function :overload
 
   # same as `#overwrite`, but raises Errno::ENOENT if any files don't exist
   def overwrite!(*filenames)
     load(*filenames, overwrite: true, ignore: false)
   end
   alias_method :overload!, :overwrite!
-  module_function :overload!
 
   # Parses the given files, yielding for each file if a block is given.
   #
@@ -65,11 +62,60 @@ def parse(*filenames, overwrite: false, ignore: true, &block)
     end
   end
 
-  def instrument(name, payload = {}, &block)
-    if instrumenter
-      instrumenter.instrument("#{name}.dotenv", payload, &block)
-    else
-      block&.call payload
+  # Save the current `ENV` to be restored later
+  def save
+    instrument(:save) do |payload|
+      @diff = payload[:diff] = Dotenv::Diff.new
+    end
+  end
+
+  # Restore `ENV` to a given state
+  #
+  # @param env [Hash] Hash of keys and values to restore, defaults to the last saved state
+  # @param safe [Boolean] Is it safe to modify `ENV`? Defaults to `true` in the main thread, otherwise raises an error.
+  def restore(env = @diff&.a, safe: Thread.current == Thread.main)
+    diff = Dotenv::Diff.new(b: env)
+    return unless diff.any?
+
+    unless safe
+      raise ThreadError, <<~EOE.tr("\n", " ")
+        Dotenv.restore is not thread safe. Use `Dotenv.modify { }` to update ENV for the duration
+        of the block in a thread safe manner, or call `Dotenv.restore(safe: true)` to ignore
+        this error.
+      EOE
+    end
+    instrument(:restore, diff: diff) { ENV.replace(env) }
+  end
+
+  # Update `ENV` with the given hash of keys and values
+  #
+  # @param env [Hash] Hash of keys and values to set in `ENV`
+  # @param overwrite [Boolean] Overwrite existing `ENV` values
+  def update(env = {}, overwrite: false)
+    instrument(:update) do |payload|
+      diff = payload[:diff] = Dotenv::Diff.new do
+        ENV.update(env.transform_keys(&:to_s)) do |key, old_value, new_value|
+          # This block is called when a key exists. Return the new value if overwrite is true.
+          overwrite ? new_value : old_value
+        end
+      end
+      diff.env
+    end
+  end
+
+  # Modify `ENV` for the block and restore it to its previous state afterwards.
+  #
+  # Note that the block is synchronized to prevent concurrent modifications to `ENV`,
+  # so multiple threads will be executed serially.
+  #
+  # @param env [Hash] Hash of keys and values to set in `ENV`
+  def modify(env = {}, &block)
+    SEMAPHORE.synchronize do
+      diff = Dotenv::Diff.new
+      update(env, overwrite: true)
+      block.call
+    ensure
+      restore(diff.a, safe: true)
     end
   end
 
@@ -79,15 +125,14 @@ def require_keys(*keys)
     raise MissingKeys, missing_keys
   end
 
-  # Save a snapshot of the current `ENV` to be restored later
-  def save
-    @snapshot = ENV.to_h.freeze
-    instrument(:save, snapshot: @snapshot)
-  end
+  private
 
-  # Restore the previous snapshot of `ENV`
-  def restore
-    instrument(:restore, diff: Dotenv::Diff.new(ENV.to_h, @snapshot)) { ENV.replace(@snapshot) }
+  def instrument(name, payload = {}, &block)
+    if instrumenter
+      instrumenter.instrument("#{name}.dotenv", payload, &block)
+    else
+      block&.call payload
+    end
   end
 end
 

lib/dotenv/autorestore.rb

@@ -17,7 +17,13 @@
       setup { Dotenv.save }
 
       # Restore ENV after each test
-      teardown { Dotenv.restore }
+      teardown do
+        Dotenv.restore
+      rescue ThreadError => e
+        # Restore will fail if running tests in parallel.
+        warn e.message
+        warn "Set `config.dotenv.autorestore = false` in `config/initializers/test.rb`" if defined?(Dotenv::Rails)
+      end
     end
   end
 end

lib/dotenv/diff.rb

@@ -1,10 +1,28 @@
 module Dotenv
-  # Compare two hashes and return the differences
+  # A diff between multiple states of ENV.
   class Diff
-    attr_reader :a, :b
+    attr_reader :a
 
-    def initialize(a, b)
+    # Create a new diff. The initial state defaults to a clone of current ENV. If given a block,
+    # the state of ENV will be preserved as the final state for comparison. Otherwise, the current
+    # ENV will be the current state.
+    def initialize(a: current_env, b: nil, &block)
       @a, @b = a, b
+      if block
+        begin
+          block.call self
+        ensure
+          @b = current_env
+        end
+      end
+    end
+
+    def b
+      @b || current_env
+    end
+
+    def current_env
+      ENV.to_h.freeze
     end
 
     # Return a Hash of keys added with their new values
@@ -23,5 +41,14 @@ def changed
         [k, [a[k], v]]
       end.to_h
     end
+
+    # Returns a Hash of all added, changed, and removed keys and their new values
+    def env
+      @env ||= b.slice(*(added.keys + changed.keys)).merge(removed.transform_values { |v| nil })
+    end
+
+    def any?
+      [added, removed, changed].any?(&:any?)
+    end
   end
 end

lib/dotenv/environment.rb

@@ -12,21 +12,11 @@ def initialize(filename, overwrite: false)
     end
 
     def load
-      update Parser.call(read, overwrite: @overwrite)
+      update Parser.call(read, overwrite: overwrite)
     end
 
     def read
       File.open(@filename, "rb:bom|utf-8", &:read)
     end
-
-    def apply
-      each do |k, v|
-        if @overwrite
-          ENV[k] = v
-        else
-          ENV[k] ||= v
-        end
-      end
-    end
   end
 end

lib/dotenv/log_subscriber.rb

@@ -9,13 +9,15 @@ def logger
     end
 
     def load(event)
-      diff = event.payload[:diff]
       env = event.payload[:env]
 
-      # Only show the keys that were added or changed
-      changed = env.slice(*(diff.added.keys + diff.changed.keys)).keys.map { |key| color_var(key) }
+      info "Loaded #{color_filename(env.filename)}"
+    end
 
-      info "Set #{changed.to_sentence} from #{color_filename(env.filename)}" if changed.any?
+    def update(event)
+      diff = event.payload[:diff]
+      changed = diff.env.keys.map { |key| color_var(key) }
+      debug "Set #{changed.to_sentence}" if diff.any?
     end
 
     def save(event)

spec/dotenv/diff_spec.rb

@@ -3,7 +3,7 @@
 describe Dotenv::Diff do
   let(:before) { {} }
   let(:after) { {} }
-  subject { Dotenv::Diff.new(before, after) }
+  subject { Dotenv::Diff.new(a: before, b: after) }
 
   context "no changes" do
     let(:before) { {"A" => 1} }
@@ -12,6 +12,8 @@
     it { expect(subject.added).to eq({}) }
     it { expect(subject.removed).to eq({}) }
     it { expect(subject.changed).to eq({}) }
+    it { expect(subject.any?).to eq(false) }
+    it { expect(subject.env).to eq({}) }
   end
 
   context "key added" do
@@ -20,6 +22,8 @@
     it { expect(subject.added).to eq("A" => 1) }
     it { expect(subject.removed).to eq({}) }
     it { expect(subject.changed).to eq({}) }
+    it { expect(subject.any?).to eq(true) }
+    it { expect(subject.env).to eq("A" => 1) }
   end
 
   context "key removed" do
@@ -28,6 +32,8 @@
     it { expect(subject.added).to eq({}) }
     it { expect(subject.removed).to eq("A" => 1) }
     it { expect(subject.changed).to eq({}) }
+    it { expect(subject.any?).to eq(true) }
+    it { expect(subject.env).to eq("A" => nil) }
   end
 
   context "key changed" do
@@ -37,5 +43,7 @@
     it { expect(subject.added).to eq({}) }
     it { expect(subject.removed).to eq({}) }
     it { expect(subject.changed).to eq("A" => [1, 2]) }
+    it { expect(subject.any?).to eq(true) }
+    it { expect(subject.env).to eq("A" => 2) }
   end
 end

spec/dotenv/environment_spec.rb

@@ -16,34 +16,6 @@
     end
   end
 
-  describe "apply" do
-    it "sets variables in ENV" do
-      subject.apply
-      expect(ENV["OPTION_A"]).to eq("1")
-    end
-
-    it "does not overwrite defined variables" do
-      ENV["OPTION_A"] = "predefined"
-      subject.apply
-      expect(ENV["OPTION_A"]).to eq("predefined")
-    end
-
-    context "with overwrite: true" do
-      subject { env("OPTION_A=1\nOPTION_B=2", overwrite: true) }
-
-      it "sets variables in the ENV" do
-        subject.apply
-        expect(ENV["OPTION_A"]).to eq("1")
-      end
-
-      it "overwrites defined variables" do
-        ENV["OPTION_A"] = "predefined"
-        subject.apply
-        expect(ENV["OPTION_A"]).to eq("1")
-      end
-    end
-  end
-
   require "tempfile"
   def env(text, ...)
     file = Tempfile.new("dotenv")

spec/dotenv/log_subscriber_spec.rb

@@ -11,31 +11,35 @@
     Dotenv::Rails.logger = Logger.new(logs)
   end
 
-  context "set" do
-    it "logs when a new instance variable is set" do
+  describe "load" do
+    it "logs when a file is loaded" do
       Dotenv.load(fixture_path("plain.env"))
-      expect(logs.string).to match(/Set.*PLAIN.*from.*plain.env/)
+      expect(logs.string).to match(/Loaded.*plain.env/)
+      expect(logs.string).to match(/Set.*PLAIN/)
+    end
+  end
+
+  context "update" do
+    it "logs when a new instance variable is set" do
+      Dotenv.update({"PLAIN" => "true"})
+      expect(logs.string).to match(/Set.*PLAIN/)
     end
 
     it "logs when an instance variable is overwritten" do
       ENV["PLAIN"] = "nope"
-      Dotenv.load(fixture_path("plain.env"), overwrite: true)
-      expect(logs.string).to match(/Set.*PLAIN.*from.*plain.env/)
+      Dotenv.update({"PLAIN" => "true"}, overwrite: true)
+      expect(logs.string).to match(/Set.*PLAIN/)
     end
 
     it "does not log when an instance variable is not overwritten" do
-      # load everything once and clear the logs
-      Dotenv.load(fixture_path("plain.env"))
-      logs.truncate(0)
-
-      # load again
-      Dotenv.load(fixture_path("plain.env"))
-      expect(logs.string).not_to match(/Set.*plain.env/i)
+      ENV["FOO"] = "existing"
+      Dotenv.update({"FOO" => "new"})
+      expect(logs.string).not_to match(/FOO/)
     end
 
     it "does not log when an instance variable is unchanged" do
       ENV["PLAIN"] = "true"
-      Dotenv.load(fixture_path("plain.env"), overwrite: true)
+      Dotenv.update({"PLAIN" => "true"}, overwrite: true)
       expect(logs.string).not_to match(/PLAIN/)
     end
   end