Stop poller cooperatively

lib/flipper/poller.rb

@@ -18,13 +18,14 @@ def self.get(key, options = {})
     end
 
     def self.reset
-      instances.each do |_, instance|
+      instances.each do |key, instance|
         instance.stop
-        instance.thread&.join(1)
-      end.clear
+        instances.delete(key) unless instance.thread&.alive?
+      end
     end
 
     MINIMUM_POLL_INTERVAL = 10
+    STOP_JOIN_TIMEOUT = 1
 
     def initialize(options = {})
       @thread = nil
@@ -35,6 +36,9 @@ def initialize(options = {})
       @last_synced_at = Concurrent::AtomicFixnum.new(0)
       @adapter = Adapters::Memory.new(nil, threadsafe: true)
       @shutdown_requested = Concurrent::AtomicBoolean.new(false)
+      @stop_requested = Concurrent::AtomicBoolean.new(false)
+      @stop_mutex = Mutex.new
+      @stop_condition = ConditionVariable.new
 
       self.interval = options.fetch(:interval, 10)
       @initial_interval = @interval
@@ -56,21 +60,44 @@ def stop
       @instrumenter.instrument("poller.#{InstrumentationNamespace}", {
         operation: :stop,
       })
-      @thread&.kill
+
+      @stop_mutex.synchronize do
+        @stop_requested.make_true
+        @stop_condition.broadcast
+      end
+
+      thread_to_stop = @thread
+      unless thread_to_stop
+        @stop_requested.make_false
+        return
+      end
+      return if thread_to_stop.equal?(Thread.current)
+
+      thread_to_stop.join(STOP_JOIN_TIMEOUT)
+      unless thread_to_stop.alive?
+        @thread = nil if @thread.equal?(thread_to_stop)
+        @stop_requested.make_false unless @shutdown_requested.true?
+      end
     end
 
     def run
       loop do
-        sleep jitter
+        break if stop_requested?
+
+        break if wait_for_stop(jitter)
 
         begin
           sync
         rescue
           # you can instrument these using poller.flipper
         end
+        break if stop_requested?
 
-        sleep interval
+        break if wait_for_stop(interval)
       end
+    ensure
+      @thread = nil if @thread.equal?(Thread.current)
+      @stop_requested.make_false unless @shutdown_requested.true?
     end
 
     def sync
@@ -130,9 +157,25 @@ def thread_alive?
       @thread && @thread.alive?
     end
 
+    def stop_requested?
+      @stop_requested.true? || @shutdown_requested.true?
+    end
+
+    def wait_for_stop(timeout)
+      return true if stop_requested?
+
+      @stop_mutex.synchronize do
+        return true if stop_requested?
+
+        @stop_condition.wait(@stop_mutex, timeout)
+        stop_requested?
+      end
+    end
+
     def reset
       @pid = Process.pid
       @shutdown_requested.make_false
+      @stop_requested.make_false
       mutex.unlock if mutex.locked?
     end
 

spec/flipper/poller_spec.rb

@@ -10,6 +10,7 @@
     described_class.new(
       remote_adapter: remote_adapter,
       start_automatically: false,
+      shutdown_automatically: false,
       interval: 3600 # 1 hour
     )
   end
@@ -20,6 +21,7 @@
 
     allow(subject).to receive(:loop).and_yield # Make loop just call once
     allow(subject).to receive(:sleep)          # Disable sleep
+    allow(subject).to receive(:wait_for_stop).and_return(false) # Disable condition waits
     allow(Thread).to receive(:new).and_yield   # Disable separate thread
   end
 
@@ -30,6 +32,27 @@
     end
   end
 
+  describe ".reset" do
+    it "keeps instances tracked when their worker thread is still running" do
+      poller = described_class.get("stuck", {
+        remote_adapter: remote_adapter,
+        start_automatically: false,
+        shutdown_automatically: false,
+      })
+      thread = instance_double(Thread, alive?: true)
+      poller.instance_variable_set(:@thread, thread)
+
+      expect(thread).to receive(:join).with(Flipper::Poller::STOP_JOIN_TIMEOUT)
+
+      described_class.reset
+
+      expect(described_class.get("stuck")).to be(poller)
+    ensure
+      poller&.instance_variable_set(:@thread, nil)
+      described_class.reset
+    end
+  end
+
   describe "#sync" do
     it "syncs remote adapter to local adapter" do
       stub_request(:get, "#{url}/features?exclude_gate_names=true")
@@ -348,6 +371,77 @@
     end
   end
 
+  describe "#stop" do
+    it "requests stop, joins the worker thread, and clears the thread reference" do
+      thread = instance_double(Thread)
+      subject.instance_variable_set(:@thread, thread)
+
+      expect(thread).not_to receive(:kill)
+      expect(thread).to receive(:alive?).and_return(false)
+      expect(thread).to receive(:join).with(Flipper::Poller::STOP_JOIN_TIMEOUT)
+
+      subject.stop
+
+      expect(subject.thread).to be(nil)
+      expect(subject.instance_variable_get(:@shutdown_requested)).to be_false
+      expect(subject.instance_variable_get(:@stop_requested)).to be_false
+    end
+
+    it "allows the poller to start again after a normal stop" do
+      thread = instance_double(Thread)
+      subject.instance_variable_set(:@thread, thread)
+
+      allow(thread).to receive(:alive?).and_return(false)
+      allow(thread).to receive(:join).with(Flipper::Poller::STOP_JOIN_TIMEOUT)
+
+      subject.stop
+
+      expect(Thread).to receive(:new).and_yield
+      expect(subject).to receive(:loop).and_yield
+      expect(subject).to receive(:sync)
+      subject.start
+    end
+
+    it "does not wait indefinitely for a worker that has not stopped yet" do
+      thread = instance_double(Thread, alive?: true)
+      subject.instance_variable_set(:@thread, thread)
+
+      expect(thread).to receive(:join).with(Flipper::Poller::STOP_JOIN_TIMEOUT)
+
+      subject.stop
+
+      expect(subject.thread).to be(thread)
+      expect(subject.instance_variable_get(:@stop_requested)).to be_true
+    end
+
+    it "does not kill itself when shutdown is requested from the poller thread" do
+      stub_request(:get, "#{url}/features?exclude_gate_names=true")
+        .to_return(
+          status: 200,
+          body: JSON.generate(features: []),
+          headers: { "poll-shutdown" => "true" }
+        )
+
+      subject.instance_variable_set(:@thread, Thread.current)
+
+      expect(Thread.current).not_to receive(:kill)
+
+      subject.run
+
+      expect(subject.thread).to be(nil)
+    end
+
+    it "does not wait for the interval if stop was requested before the wait begins" do
+      allow(subject).to receive(:wait_for_stop).and_call_original
+      subject.instance_variable_get(:@stop_requested).make_true
+
+      started_at = Concurrent.monotonic_time
+      expect(subject.send(:wait_for_stop, 3600)).to be(true)
+
+      expect(Concurrent.monotonic_time - started_at).to be < 0.1
+    end
+  end
+
   describe "#start" do
     it "starts the poller thread" do
       expect(Thread).to receive(:new).and_yield