iruby 0.2.9 → 0.6.1

data/lib/iruby/ostream.rb

@@ -25,22 +25,43 @@ module IRuby
     alias_method :next, :read
     alias_method :readline, :read
 
-    def write(s)
-      raise 'I/O operation on closed file' unless @session
-      @session.send(:publish, :stream, name: @name, text: s.to_s)
-      nil
+    def write(*obj)
+      str = build_string { |sio| sio.write(*obj) }
+      session_send(str)
     end
     alias_method :<<, :write
     alias_method :print, :write
 
-    def puts(*lines)
-      lines = [''] if lines.empty?
-      lines.each { |s| write("#{s}\n")}
-      nil
+    def printf(format, *obj)
+      str = build_string { |sio| sio.printf(format, *obj) }
+      session_send(str)
+    end
+
+    def puts(*obj)
+      str = build_string { |sio| sio.puts(*obj) }
+      session_send(str)
     end
 
     def writelines(lines)
       lines.each { |s| write(s) }
     end
+
+    # Called by irb
+    def set_encoding(extern, intern)
+      a = extern
+    end
+
+    private
+
+    def build_string
+      StringIO.open { |sio| yield(sio); sio.string }
+    end
+
+    def session_send(str)
+      raise 'I/O operation on closed file' unless @session
+
+      @session.send(:publish, :stream, name: @name, text: str)
+      nil
+    end
   end
 end

data/lib/iruby/session.rb

@@ -0,0 +1,116 @@
+require 'iruby/session_adapter'
+require 'iruby/session/mixin'
+
+require 'securerandom'
+
+module IRuby
+  class Session
+    include SessionSerialize
+
+    def initialize(config, adapter_name=nil)
+      @config = config
+      @adapter = create_session_adapter(config, adapter_name)
+
+      setup
+      setup_sockets
+      setup_heartbeat
+      setup_security
+    end
+
+    attr_reader :adapter, :config
+
+    def description
+      "#{@adapter.name} session adapter"
+    end
+
+    def setup
+    end
+
+    def setup_sockets
+      protocol, host = config.values_at('transport', 'ip')
+      shell_port = config['shell_port']
+      iopub_port = config['iopub_port']
+      stdin_port = config['stdin_port']
+
+      @shell_socket, @shell_port = @adapter.make_router_socket(protocol, host, shell_port)
+      @iopub_socket, @iopub_port = @adapter.make_pub_socket(protocol, host, iopub_port)
+      @stdin_socket, @stdin_port = @adapter.make_router_socket(protocol, host, stdin_port)
+
+      @sockets = {
+        publish: @iopub_socket,
+        reply:   @shell_socket,
+        stdin:   @stdin_socket
+      }
+    end
+
+    def setup_heartbeat
+      protocol, host = config.values_at('transport', 'ip')
+      hb_port = config['hb_port']
+      @hb_socket, @hb_port = @adapter.make_rep_socket(protocol, host, hb_port)
+      @heartbeat_thread = Thread.start do
+        begin
+          # NOTE: this loop is copied from CZTop's old session code
+          @adapter.heartbeat_loop(@hb_socket)
+        rescue Exception => e
+          IRuby.logger.fatal "Kernel heartbeat died: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+    end
+
+    def setup_security
+      @session_id = SecureRandom.uuid
+      unless config['key'].empty? || config['signature_scheme'].empty?
+        unless config['signature_scheme'] =~ /\Ahmac-/
+          raise "Unknown signature_scheme: #{config['signature_scheme']}"
+        end
+        digest_algorithm = config['signature_scheme'][/\Ahmac-(.*)\Z/, 1]
+        @hmac = OpenSSL::HMAC.new(config['key'], OpenSSL::Digest.new(digest_algorithm))
+      end
+    end
+
+    def send(socket_type, message_type, content)
+      sock = check_socket_type(socket_type)
+      idents = if socket_type == :reply && @last_recvd_msg
+                 @last_recvd_msg[:idents]
+               else
+                 message_type == :stream ? "stream.#{content[:name]}" : message_type
+               end
+      header = {
+        msg_type: message_type,
+        msg_id:   SecureRandom.uuid,
+        username: 'kernel',
+        session:  @session_id,
+        version:  '5.0'
+      }
+      @adapter.send(sock, serialize(idents, header, content))
+    end
+
+    def recv(socket_type)
+      sock = check_socket_type(socket_type)
+      data = @adapter.recv(sock)
+      @last_recvd_msg = unserialize(data)
+    end
+
+    def recv_input
+      sock = check_socket_type(:stdin)
+      data = @adapter.recv(sock)
+      unserialize(data)[:content]["value"]
+    end
+
+    private
+
+    def check_socket_type(socket_type)
+      case socket_type
+      when :publish, :reply, :stdin
+        @sockets[socket_type]
+      else
+        raise ArgumentError, "Invalid socket type #{socket_type}"
+      end
+    end
+
+    def create_session_adapter(config, adapter_name)
+      adapter_class = SessionAdapter.select_adapter_class(adapter_name)
+      adapter_class.new(config)
+    end
+  end
+end

data/lib/iruby/session/{rbczmq.rb → cztop.rb}

@@ -1,34 +1,32 @@
-require 'rbczmq'
+require 'cztop'
 
 module IRuby
   class Session
     include SessionSerialize
 
     def initialize(config)
-      c = ZMQ::Context.new
-
       connection = "#{config['transport']}://#{config['ip']}:%d"
-      reply_socket = c.socket(:ROUTER)
-      reply_socket.bind(connection % config['shell_port'])
-
-      pub_socket = c.socket(:PUB)
-      pub_socket.bind(connection % config['iopub_port'])
 
-      stdin_socket = c.socket(:ROUTER)
-      stdin_socket.bind(connection % config['stdin_port'])
+      reply_socket = CZTop::Socket::ROUTER.new(connection % config['shell_port'])
+      pub_socket = CZTop::Socket::PUB.new(connection % config['iopub_port'])
+      stdin_socket = CZTop::Socket::ROUTER.new(connection % config['stdin_port'])
 
       Thread.new do
         begin
-          hb_socket = c.socket(:REP)
-          hb_socket.bind(connection % config['hb_port'])
-          ZMQ.proxy(hb_socket, hb_socket)
+          hb_socket = CZTop::Socket::REP.new(connection % config['hb_port'])
+          loop do
+            message = hb_socket.receive
+            hb_socket << message
+          end
         rescue Exception => e
           IRuby.logger.fatal "Kernel heartbeat died: #{e.message}\n#{e.backtrace.join("\n")}"
         end
       end
 
-      @sockets = { 
-        publish: pub_socket, reply: reply_socket, stdin: stdin_socket
+      @sockets = {
+        publish: pub_socket,
+        reply: reply_socket,
+        stdin: stdin_socket,
       }
 
       @session = SecureRandom.uuid
@@ -38,6 +36,10 @@ module IRuby
       end
     end
 
+    def description
+      'old-stle session using cztop'
+    end
+
     # Build and send a message
     def send(socket, type, content)
       idents =
@@ -53,16 +55,16 @@ module IRuby
         session:  @session,
         version:  '5.0'
       }
-      @sockets[socket].send_message(ZMQ::Message(*serialize(idents, header, content)))
+      @sockets[socket] << serialize(idents, header, content)
     end
 
     # Receive a message and decode it
     def recv(socket)
-      @last_recvd_msg = unserialize(@sockets[socket].recv_message)
+      @last_recvd_msg = unserialize(@sockets[socket].receive)
     end
 
-    def recv_input 
-      unserialize(@sockets[:stdin].recv_message)[:content]["value"]
+    def recv_input
+      unserialize(@sockets[:stdin].receive)[:content]["value"]
     end
   end
 end

data/lib/iruby/session/ffi_rzmq.rb

@@ -67,7 +67,7 @@ module IRuby
       while msg.empty? || socket.more_parts?
         begin
           frame = ''
-          rc = socket.recv_string(frame, 1)
+          rc = socket.recv_string(frame)
           ZMQ::Util.error_check('zmq_msg_send', rc)
           msg << frame
         rescue

data/lib/iruby/session_adapter.rb

@@ -0,0 +1,72 @@
+module IRuby
+  class SessionAdapterNotFound < RuntimeError; end
+
+  module SessionAdapter
+    class BaseAdapter
+      def self.available?
+        load_requirements
+        true
+      rescue LoadError
+        false
+      end
+
+      def self.load_requirements
+        # Do nothing
+      end
+
+      def initialize(config)
+        @config = config
+      end
+
+      def name
+        self.class.name[/::(\w+)Adapter\Z/, 1].downcase
+      end
+
+      def make_router_socket(protocol, host, port)
+        socket, port = make_socket(:ROUTER, protocol, host, port)
+        [socket, port]
+      end
+
+      def make_pub_socket(protocol, host, port)
+        socket, port = make_socket(:PUB, protocol, host, port)
+        [socket, port]
+      end
+
+      def make_rep_socket(protocol, host, port)
+        socket, port = make_socket(:REP, protocol, host, port)
+        [socket, port]
+      end
+    end
+
+    require_relative 'session_adapter/ffirzmq_adapter'
+    require_relative 'session_adapter/cztop_adapter'
+    require_relative 'session_adapter/pyzmq_adapter'
+    require_relative 'session_adapter/test_adapter'
+
+    def self.select_adapter_class(name=nil)
+      classes = {
+        'ffi-rzmq' => SessionAdapter::FfirzmqAdapter,
+        'cztop' => SessionAdapter::CztopAdapter,
+        # 'pyzmq' => SessionAdapter::PyzmqAdapter
+        'test' => SessionAdapter::TestAdapter,
+      }
+      if (name ||= ENV.fetch('IRUBY_SESSION_ADAPTER', nil))
+        cls = classes[name]
+        unless cls.available?
+          if ENV['IRUBY_SESSION_ADAPTER']
+            raise SessionAdapterNotFound,
+                  "Session adapter `#{name}` from IRUBY_SESSION_ADAPTER is unavailable"
+          else
+            raise SessionAdapterNotFound,
+                  "Session adapter `#{name}` is unavailable"
+          end
+        end
+        return cls
+      end
+      classes.each_value do |cls|
+        return cls if cls.available?
+      end
+      raise SessionAdapterNotFound, "No session adapter is available"
+    end
+  end
+end

data/lib/iruby/session_adapter/cztop_adapter.rb

@@ -0,0 +1,45 @@
+module IRuby
+  module SessionAdapter
+    class CztopAdapter < BaseAdapter
+      def self.load_requirements
+        require 'cztop'
+      end
+
+      def send(sock, data)
+        sock << data
+      end
+
+      def recv(sock)
+        sock.receive
+      end
+
+      def heartbeat_loop(sock)
+        loop do
+          message = sock.receive
+          sock << message
+        end
+      end
+
+      private
+
+      def socket_type_class(type_symbol)
+        case type_symbol
+        when :ROUTER, :PUB, :REP
+          CZTop::Socket.const_get(type_symbol)
+        else
+          if CZTop::Socket.const_defined?(type_symbol)
+            raise ArgumentError, "Unsupported ZMQ socket type: #{type_symbol}"
+          else
+            raise ArgumentError, "Invalid ZMQ socket type: #{type_symbol}"
+          end
+        end
+      end
+
+      def make_socket(type_symbol, protocol, host, port)
+        uri = "#{protocol}://#{host}:#{port}"
+        socket_class = socket_type_class(type_symbol)
+        socket_class.new(uri)
+      end
+    end
+  end
+end

data/lib/iruby/session_adapter/ffirzmq_adapter.rb

@@ -0,0 +1,55 @@
+module IRuby
+  module SessionAdapter
+    class FfirzmqAdapter < BaseAdapter
+      def self.load_requirements
+        require 'ffi-rzmq'
+      end
+
+      def send(sock, data)
+        data.each_with_index do |part, i|
+          sock.send_string(part, i == data.size - 1 ? 0 : ZMQ::SNDMORE)
+        end
+      end
+
+      def recv(sock)
+        msg = []
+        while msg.empty? || sock.more_parts?
+          begin
+            frame = ''
+            rc = sock.recv_string(frame)
+            ZMQ::Util.error_check('zmq_msg_recv', rc)
+            msg << frame
+          rescue
+          end
+        end
+        msg
+      end
+
+      def heartbeat_loop(sock)
+        @heartbeat_device = ZMQ::Device.new(sock, sock)
+      end
+
+      private
+
+      def make_socket(type, protocol, host, port)
+        case type
+        when :ROUTER, :PUB, :REP
+          type = ZMQ.const_get(type)
+        else
+          if ZMQ.const_defined?(type)
+            raise ArgumentError, "Unsupported ZMQ socket type: #{type_symbol}"
+          else
+            raise ArgumentError, "Invalid ZMQ socket type: #{type_symbol}"
+          end
+        end
+        zmq_context.socket(type).tap do |sock|
+          sock.bind("#{protocol}://#{host}:#{port}")
+        end
+      end
+
+      def zmq_context
+        @zmq_context ||= ZMQ::Context.new
+      end
+    end
+  end
+end

data/lib/iruby/session_adapter/pyzmq_adapter.rb

@@ -0,0 +1,77 @@
+module IRuby
+  module SessionAdapter
+    class PyzmqAdapter < BaseAdapter
+
+      class << self
+        def load_requirements
+          require 'pycall'
+          import_pyzmq
+        end
+
+        def import_pyzmq
+          @zmq = PyCall.import_module('zmq')
+        rescue PyCall::PyError => error
+          raise LoadError, error.message
+        end
+
+        attr_reader :zmq
+      end
+
+      def make_router_socket(protocol, host, port)
+        make_socket(:ROUTER, protocol, host, port)
+      end
+
+      def make_pub_socket(protocol, host, port)
+        make_socket(:PUB, protocol, host, port)
+      end
+
+      def heartbeat_loop(sock)
+        PyCall.sys.path.append(File.expand_path('../pyzmq', __FILE__))
+        heartbeat = PyCall.import_module('iruby.heartbeat')
+        @heartbeat_thread = heartbeat.Heartbeat.new(sock)
+        @heartbeat_thread.start
+      end
+
+      private
+
+      def socket_type(type_symbol)
+        case type_symbol
+        when :ROUTER, :PUB, :REP
+          zmq[type_symbol]
+        else
+          raise ArgumentError, "Unknown ZMQ socket type: #{type_symbol}"
+        end
+      end
+
+      def make_socket(type_symbol, protocol, host, port)
+        type = socket_type(type_symbol)
+        sock = zmq_context.socket(type)
+        bind_socket(sock, protocol, host, port)
+        sock
+      end
+
+      def bind_socket(sock, protocol, host, port)
+        iface = "#{protocol}://#{host}"
+        case protocol
+        when 'tcp'
+          if port <= 0
+            port = sock.bind_to_random_port(iface)
+          else
+            sock.bind("#{iface}:#{port}")
+          end
+        else
+          raise ArgumentError, "Unsupported protocol: #{protocol}"
+        end
+        [sock, port]
+      end
+
+      def zmq_context
+        zmq.Context.instance
+      end
+
+      def zmq
+        self.class.zmq
+      end
+    end
+  end
+end