lib/zmachine/connection.rb in zmachine-0.3.0 vs lib/zmachine/connection.rb in zmachine-0.3.2

@@ -26,10 +26,11 @@
       ZMachine.logger.debug("zmachine:connection:#{__method__}", connection: self) if ZMachine.debug
       klass = (address =~ %r{\w+://}) ? ZMQChannel : TCPChannel
       @channel = klass.new
       @channel.bind(address, port_or_type)
       @block = block
+      @block.call(self) if @block && @channel.is_a?(ZMQChannel)
       self
     end
 
     def connect(address, port_or_type, &block)
       ZMachine.logger.debug("zmachine:connection:#{__method__}", connection: self) if ZMachine.debug
@@ -163,10 +164,10 @@
 
     # selector registration
 
     def register(selector)
       ZMachine.logger.debug("zmachine:connection:#{__method__}", connection: self, fd: @channel.selectable_fd) if ZMachine.debug
-      @channel_key ||= @channel.selectable_fd.register(selector, current_events, self)
+      @channel_key = @channel.selectable_fd.register(selector, current_events, self)
     end
 
     def valid?
       @channel_key &&
       @channel_key.valid?