Remove finished dispatcher handlers, order dispatch handlers

Author: jeffschoner-stripe

lib/temporal/workflow/context.rb

@@ -254,7 +254,7 @@ def wait_for_any(*futures)
         fiber = Fiber.current
         blocked = true
 
-        futures.each do |future|
+        handlers = futures.map do |future|
           dispatcher.register_handler(future.target, Dispatcher::WILDCARD) do
             # Because any of the futures can resume the fiber, ignore any callbacks
             # from other futures after unblocking has occurred
@@ -266,6 +266,7 @@ def wait_for_any(*futures)
         end
 
         Fiber.yield
+        handlers.each(&:unregister)
 
         return
       end
@@ -279,7 +280,7 @@ def wait_until(&unblock_condition)
         fiber = Fiber.current
         blocked = true
 
-        dispatcher.register_handler(Dispatcher::TARGET_WILDCARD, Dispatcher::WILDCARD) do
+        handler = dispatcher.register_handler(Dispatcher::TARGET_WILDCARD, Dispatcher::WILDCARD) do
           # Because this block can run for any dispatch, ensure the fiber is only
           # resumed one time by checking if it's already been unblocked.
           if blocked && unblock_condition.call
@@ -289,6 +290,7 @@ def wait_until(&unblock_condition)
         end
 
         Fiber.yield
+        handler.unregister
 
         return
       end
@@ -316,6 +318,8 @@ def on_signal(signal_name = nil, &block)
             call_in_fiber(block, signal, input)
           end
         end
+
+        return
       end
 
       def on_query(query, &block)

lib/temporal/workflow/dispatcher.rb

@@ -13,18 +13,36 @@ class Workflow
     # the event_name. The order of this dispatch is not guaranteed.
     #
     class Dispatcher
+      class DispatchHandler
+        def initialize(handlers_for_target, id)
+          @handlers_for_target = handlers_for_target
+          @id = id
+        end
+
+        # Unregister the handler from the dispatcher
+        def unregister
+          handlers_for_target.delete(id)
+        end
+
+        private
+
+        attr_reader :handlers_for_target, :id
+      end
+
       WILDCARD = '*'.freeze
       TARGET_WILDCARD = '*'.freeze
 
       EventStruct = Struct.new(:event_name, :handler)
 
       def initialize
-        @handlers = Hash.new { |hash, key| hash[key] = [] }
+        @handlers = Hash.new { |hash, key| hash[key] = {} }
+        @next_id = 0
       end
 
       def register_handler(target, event_name, &handler)
-        handlers[target] << EventStruct.new(event_name, handler)
-        self
+        @next_id += 1
+        handlers[target][@next_id] = EventStruct.new(event_name, handler)
+        DispatchHandler.new(handlers[target], @next_id)
       end
 
       def dispatch(target, event_name, args = nil)
@@ -39,9 +57,10 @@ def dispatch(target, event_name, args = nil)
 
       def handlers_for(target, event_name)
         handlers[target]
-          .concat(handlers[TARGET_WILDCARD])
-          .select { |event_struct| match?(event_struct, event_name) }
-          .map(&:handler)
+          .merge(handlers[TARGET_WILDCARD]) { raise 'Cannot resolve duplicate dispatcher handler IDs' }
+          .select { |_, event_struct| match?(event_struct, event_name) }
+          .sort
+          .map { |_, event_struct| event_struct.handler }
       end
 
       def match?(event_struct, event_name)

spec/unit/lib/temporal/workflow/dispatcher_spec.rb

@@ -8,7 +8,10 @@
   describe '#register_handler' do
     let(:block) { -> { 'handler body' } }
     let(:event_name) { 'signaled' }
-    let(:dispatcher) { subject.register_handler(target, event_name, &block) }
+    let(:dispatcher) do
+      subject.register_handler(target, event_name, &block)
+      subject
+    end
     let(:handlers) { dispatcher.send(:handlers) }
 
     context 'with default handler_name' do
@@ -19,17 +22,17 @@
       end
 
       it 'stores the target and handler once' do
-        expect(handlers[target]).to be_kind_of(Array)
+        expect(handlers[target]).to be_kind_of(Hash)
         expect(handlers[target].count).to eq 1
       end
 
       it 'associates the event name with the target' do
-        event = handlers[target].first
+        event = handlers[target][1]
         expect(event.event_name).to eq(event_name)
       end
 
       it 'associates the handler with the target' do
-        event = handlers[target].first
+        event = handlers[target][1]
         expect(event.handler).to eq(block)
       end
     end
@@ -43,20 +46,42 @@
       end
 
       it 'stores the target and handler once' do
-        expect(handlers[target]).to be_kind_of(Array)
+        expect(handlers[target]).to be_kind_of(Hash)
         expect(handlers[target].count).to eq 1
       end
 
       it 'associates the event name and handler name with the target' do
-        event = handlers[target].first
+        event = handlers[target][1]
         expect(event.event_name).to eq(event_name)
       end
 
       it 'associates the handler with the target' do
-        event = handlers[target].first
+        event = handlers[target][1]
         expect(event.handler).to eq(block)
       end
     end
+
+    it 'removes a given handler against the target' do
+      block1 = -> { 'handler body' }
+      block2 = -> { 'other handler body' }
+      block3 = -> { 'yet another handler body' }
+
+      handle1 = subject.register_handler(target, 'signaled', &block1)
+      subject.register_handler(target, 'signaled', &block2)
+      subject.register_handler(other_target, 'signaled', &block3)
+
+      expect(subject.send(:handlers)[target][1].event_name).to eq('signaled')
+      expect(subject.send(:handlers)[target][1].handler).to be(block1)
+
+      expect(subject.send(:handlers)[target][2].event_name).to eq('signaled')
+      expect(subject.send(:handlers)[target][2].handler).to be(block2)
+
+      expect(subject.send(:handlers)[other_target][3].event_name).to eq('signaled')
+      expect(subject.send(:handlers)[other_target][3].handler).to be(block3)
+
+      handle1.unregister
+      expect(subject.send(:handlers)[target][1]).to be(nil)
+    end
   end
 
   describe '#dispatch' do
@@ -114,10 +139,13 @@
 
     context 'with TARGET_WILDCARD target handler' do
       let(:handler_6) { -> { 'sixth block' } }
+      let(:handler_7) { -> { 'seventh block' } }
       before do
         allow(handler_6).to receive(:call)
+        allow(handler_7).to receive(:call)
 
         subject.register_handler(described_class::TARGET_WILDCARD, described_class::WILDCARD, &handler_6)
+        subject.register_handler(target, 'completed', &handler_7)
       end
 
       it 'calls the handler' do
@@ -127,6 +155,7 @@
         expect(handler_1).to have_received(:call).ordered
         expect(handler_4).to have_received(:call).ordered
         expect(handler_6).to have_received(:call).ordered
+        expect(handler_7).to have_received(:call).ordered
       end
 
       it 'TARGET_WILDCARD can be compared to an EventTarget object' do
@@ -157,4 +186,4 @@
       end
     end
   end
-end
+end