PageRenderTime 7ms CodeModel.GetById 12ms app.highlight 52ms RepoModel.GetById 1ms app.codeStats 1ms

/ruby/ffi-rxs/socket.rb

https://bitbucket.org/hengestone/whip
Ruby | 716 lines | 261 code | 85 blank | 370 comment | 38 complexity | 8ca4e11bebc7aba32368608665f1f83d MD5 | raw file
  1# encoding: utf-8
  2
  3module XS
  4
  5  module CommonSocketBehavior
  6
  7    attr_reader :socket, :name
  8
  9    # Allocates a socket of type +type+ for sending and receiving data.
 10    #
 11    # By default, this class uses XS::Message for manual
 12    # memory management. For automatic garbage collection of received messages,
 13    # it is possible to override the :receiver_class to use XS::ManagedMessage.
 14    #
 15    # @example Socket creation
 16    #   sock = Socket.create(Context.create, XS::REQ, :receiver_class => XS::ManagedMessage)
 17    #
 18    # Advanced users may want to replace the receiver class with their
 19    # own custom class. The custom class must conform to the same public API
 20    # as XS::Message.
 21    #
 22    # 
 23    # @example
 24    #   if (socket = Socket.new(context.pointer, XS::REQ))
 25    #     ...
 26    #   else
 27    #     STDERR.puts "Socket creation failed"
 28    #   end
 29    #
 30    # @param pointer
 31    # @param [Constant] type
 32    #   One of @XS::REQ@, @XS::REP@, @XS::PUB@, @XS::SUB@, @XS::PAIR@,
 33    #          @XS::PULL@, @XS::PUSH@, @XS::XREQ@, @XS::REP@,
 34    #          @XS::DEALER@ or @XS::ROUTER@
 35    # @param [Hash] options
 36    #
 37    # @return [Socket] when successful
 38    # @return nil when unsuccessful
 39    def self.create context_ptr, type, opts = {:receiver_class => XS::Message}
 40      new(context_ptr, type, opts) rescue nil
 41    end
 42
 43    # Allocates a socket of type +type+ for sending and receiving data.
 44    #
 45    # To avoid rescuing exceptions, use the factory method #create for
 46    # all socket creation.
 47    #
 48    # By default, this class uses XS::Message for manual
 49    # memory management. For automatic garbage collection of received messages,
 50    # it is possible to override the :receiver_class to use XS::ManagedMessage.
 51    #
 52    # @example Socket creation
 53    #   sock = Socket.new(Context.new, XS::REQ, :receiver_class => XS::ManagedMessage)
 54    #
 55    # Advanced users may want to replace the receiver class with their
 56    # own custom class. The custom class must conform to the same public API
 57    # as XS::Message.
 58    #
 59    # Creation of a new Socket object can raise an exception. This occurs when the
 60    # +context_ptr+ is null or when the allocation of the Crossroads socket within the
 61    # context fails.
 62    #
 63    # @example
 64    #   begin
 65    #     socket = Socket.new(context.pointer, XS::REQ)
 66    #   rescue ContextError => e
 67    #     # error handling
 68    #   end
 69    #
 70    # @param pointer
 71    # @param [Constant] type
 72    #   One of @XS::REQ@, @XS::REP@, @XS::PUB@, @XS::SUB@, @XS::PAIR@,
 73    #          @XS::PULL@, @XS::PUSH@, @XS::XREQ@, @XS::REP@,
 74    #          @XS::DEALER@ or @XS::ROUTER@
 75    # @param [Hash] options
 76    #
 77    # @return [Socket] when successful
 78    # @return nil when unsuccessful
 79    def initialize context_ptr, type, opts = {:receiver_class => XS::Message}
 80      # users may override the classes used for receiving; class must conform to the
 81      # same public API as XS::Message
 82      @receiver_klass = opts[:receiver_class]
 83
 84      context_ptr = context_ptr.pointer if context_ptr.kind_of?(XS::Context)
 85
 86      unless context_ptr.null?
 87        @socket = LibXS.xs_socket context_ptr, type
 88        if @socket && !@socket.null?
 89          @name = SocketTypeNameMap[type]
 90        else
 91          raise ContextError.new 'xs_socket', 0, ETERM, "Socket pointer was null"
 92        end
 93      else
 94        raise ContextError.new 'xs_socket', 0, ETERM, "Context pointer was null"
 95      end
 96
 97      @longlong_cache = @int_cache = nil
 98      @more_parts_array = []
 99      @option_lookup = []
100      populate_option_lookup
101
102      define_finalizer
103    end
104
105    # Set the queue options on this socket
106    #
107    # @param [Constant] name numeric values
108    #   One of @XS::AFFINITY@, @XS::RATE@, @XS::RECOVERY_IVL@,
109    #          @XS::LINGER@, @XS::RECONNECT_IVL@, @XS::BACKLOG@,
110    #          @XS::RECONNECT_IVL_MAX@, @XS::MAXMSGSIZE@, @XS::SNDHWM@,
111    #          @XS::RCVHWM@, @XS::MULTICAST_HOPS@, @XS::RCVTIMEO@,
112    #          @XS::SNDTIMEO@, @XS::IPV4ONLY@, @XS::KEEPALIVE@,
113    #          @XS::SUBSCRIBE@, @XS::UNSUBSCRIBE@, @XS::IDENTITY@,
114    #          @XS::SNDBUF@, @XS::RCVBUF@
115    # @param [Constant] name string values
116    #   One of @XS::IDENTITY@, @XS::SUBSCRIBE@ or @XS::UNSUBSCRIBE@
117    # @param value
118    #
119    # @return 0 when the operation completed successfully
120    # @return -1 when this operation fails
121    #
122    # With a -1 return code, the user must check XS.errno to determine the
123    # cause.
124    #
125    # @example
126    #   rc = socket.setsockopt(XS::LINGER, 1_000)
127    #   XS::Util.resultcode_ok?(rc) ? puts("succeeded") : puts("failed")
128    def setsockopt name, value, length = nil
129      if 1 == @option_lookup[name]
130        length = 8
131        pointer = LibC.malloc length
132        pointer.write_long_long value
133
134      elsif 0 == @option_lookup[name]
135        length = 4
136        pointer = LibC.malloc length
137        pointer.write_int value
138
139      elsif 2 == @option_lookup[name]
140        length ||= value.size
141
142        # note: not checking errno for failed memory allocations :(
143        pointer = LibC.malloc length
144        pointer.write_string value
145      end
146
147      rc = LibXS.xs_setsockopt @socket, name, pointer, length
148      LibC.free(pointer) unless pointer.nil? || pointer.null?
149      rc
150    end
151
152    # Convenience method for checking on additional message parts.
153    #
154    # Equivalent to calling Socket#getsockopt with XS::RCVMORE.
155    #
156    # Warning: if the call to #getsockopt fails, this method will return
157    # false and swallow the error.
158    #
159    # @example
160    #   message_parts = []
161    #   message = Message.new
162    #   rc = socket.recvmsg(message)
163    #   if XS::Util.resultcode_ok?(rc)
164    #     message_parts << message
165    #     while more_parts?
166    #       message = Message.new
167    #       rc = socket.recvmsg(message)
168    #       message_parts.push(message) if resultcode_ok?(rc)
169    #     end
170    #   end
171    #
172    # @return true if more message parts
173    # @return false if not
174    def more_parts?
175      rc = getsockopt XS::RCVMORE, @more_parts_array
176
177      Util.resultcode_ok?(rc) ? @more_parts_array.at(0) : false
178    end
179
180    # Binds the socket to an +address+.
181    #
182    # @example
183    #   socket.bind("tcp://127.0.0.1:5555")
184    #
185    # @param address
186    #
187    # @return 0 or greater if successful
188    # @return < 0 if unsuccessful
189    def bind address
190      LibXS.xs_bind @socket, address
191    end
192
193    # Connects the socket to an +address+.
194    #
195    # @example
196    #   rc = socket.connect("tcp://127.0.0.1:5555")
197    #
198    # @param address
199    #
200    # @return 0 or greater if successful
201    # @return < 0 if unsuccessful
202    def connect address
203      LibXS.xs_connect @socket, address
204    end
205
206    # Closes the socket. Any unprocessed messages in queue are sent or dropped
207    # depending upon the value of the socket option XS::LINGER.
208    #
209    # @example
210    #   rc = socket.close
211    #   puts("Given socket was invalid!") unless 0 == rc
212    #
213    # @return 0 upon success *or* when the socket has already been closed
214    # @return -1 when the operation fails. Check XS.errno for the error code
215    def close
216      if @socket
217        remove_finalizer
218        rc = LibXS.xs_close @socket
219        @socket = nil
220        release_cache
221        rc
222      else
223        0
224      end
225    end
226
227    # Queues the message for transmission. Message is assumed to conform to the
228    # same public API as #Message.
229    #
230    # @param message
231    # @param flag
232    #   One of @0 (default) - blocking operation@, @XS::NonBlocking - non-blocking operation@,
233    #          @XS::SNDMORE - this message is part of a multi-part message@
234    #
235    # @return 0 when the message was successfully enqueued
236    # @return -1 under two conditions
237    #   1. The message could not be enqueued
238    #   2. When +flag+ is set with XS::NonBlocking and the socket returned EAGAIN.
239    #
240    # With a -1 return code, the user must check XS.errno to determine the
241    # cause.
242    def sendmsg message, flag = 0
243      __sendmsg__(@socket, message.address, flag)
244    end
245
246    # Helper method to make a new #Message instance out of the +string+ passed
247    # in for transmission.
248    #
249    # @param message
250    # @param flag
251    #   One of @0 (default)@, @XS::NonBlocking@ and @XS::SNDMORE@
252    #
253    # @return 0 when the message was successfully enqueued
254    # @return -1 under two conditions
255    #   1. The message could not be enqueued
256    #   2. When +flag+ is set with XS::NonBlocking and the socket returned EAGAIN.
257    #
258    # With a -1 return code, the user must check XS.errno to determine the
259    # cause.
260    def send_string string, flag = 0
261      message = Message.new string
262      send_and_close message, flag
263    end
264
265    # Send a sequence of strings as a multipart message out of the +parts+
266    # passed in for transmission. Every element of +parts+ should be
267    # a String.
268    #
269    # @param [Array] parts
270    # @param flag
271    #   One of @0 (default)@ and @XS::NonBlocking@
272    #
273    # @return 0 when the messages were successfully enqueued
274    # @return -1 under two conditions
275    #   1. A message could not be enqueued
276    #   2. When +flag+ is set with XS::NonBlocking and the socket returned EAGAIN.
277    #
278    # With a -1 return code, the user must check XS.errno to determine the
279    # cause.
280    def send_strings parts, flag = 0
281      return -1 if !parts || parts.empty?
282      flag = NonBlocking if dontwait?(flag)
283
284      parts[0..-2].each do |part|
285        rc = send_string part, (flag | XS::SNDMORE)
286        return rc unless Util.resultcode_ok?(rc)
287      end
288
289      send_string parts[-1], flag
290    end
291
292    # Send a sequence of messages as a multipart message out of the +parts+
293    # passed in for transmission. Every element of +parts+ should be
294    # a Message (or subclass).
295    #
296    # @param [Array] parts
297    # @param flag
298    #   One of @0 (default)@ and @XS::NonBlocking@
299    #
300    # @return 0 when the messages were successfully enqueued
301    # @return -1 under two conditions
302    #   1. A message could not be enqueued
303    #   2. When +flag+ is set with XS::NonBlocking and the socket returned EAGAIN
304    #
305    # With a -1 return code, the user must check XS.errno to determine the
306    # cause.
307    def sendmsgs parts, flag = 0
308      return -1 if !parts || parts.empty?
309      flag = NonBlocking if dontwait?(flag)
310
311      parts[0..-2].each do |part|
312        rc = sendmsg part, (flag | XS::SNDMORE)
313        return rc unless Util.resultcode_ok?(rc)
314      end
315
316      sendmsg parts[-1], flag
317    end
318
319    # Sends a message. This will automatically close the +message+ for both successful
320    # and failed sends.
321    #
322    # @param message
323    # @param flag
324    #   One of @0 (default)@ and @XS::NonBlocking
325    #
326    # @return 0 when the message was successfully enqueued
327    # @return -1 under two conditions
328    #   1. The message could not be enqueued
329    #   2. When +flag+ is set with XS::NonBlocking and the socket returned EAGAIN.
330    #
331    # With a -1 return code, the user must check XS.errno to determine the
332    # cause.
333    def send_and_close message, flag = 0
334      rc = sendmsg message, flag
335      message.close
336      rc
337    end
338    
339    # Unbinds/disconnects the socket from an endpoint.
340    #
341    # @example
342    #   socket.shutdown(endpoint_id)
343    #
344    # @param endpoint_id (endpoint id returned from socket.bind/socket.connect)
345    #
346    # @return 0 if successful
347    # @return -1 if unsuccessful
348    def shutdown endpoint_id
349      LibXS.xs_shutdown @socket, endpoint_id
350    end
351
352    # Dequeues a message from the underlying queue. By default, this is a blocking operation.
353    #
354    # @param message
355    # @param flag
356    #   One of @0 (default) - blocking operation@ and @XS::NonBlocking - non-blocking operation@
357    #
358    # @return 0 when the message was successfully dequeued
359    # @return -1 under two conditions
360    #   1. The message could not be dequeued
361    #   2. When +flags+ is set with XS::NonBlocking and the socket returned EAGAIN
362    #
363    # With a -1 return code, the user must check XS.errno to determine the
364    # cause.
365    #
366    # The application code is responsible for handling the +message+ object lifecycle
367    # when #recv returns an error code.
368    def recvmsg message, flag = 0
369      __recvmsg__(@socket, message.address, flag)
370    end
371
372    # Helper method to make a new #Message instance and convert its payload
373    # to a string.
374    #
375    # @param string
376    # @param flag
377    #   One of @0 (default)@ and @XS::NonBlocking@
378    #
379    # @return 0 when the message was successfully dequeued
380    # @return -1 under two conditions
381    #   1. The message could not be dequeued
382    #   2. When +flag+ is set with XS::NonBlocking and the socket returned EAGAIN
383    #
384    # With a -1 return code, the user must check XS.errno to determine the
385    # cause.
386    #
387    # The application code is responsible for handling the +message+ object lifecycle
388    # when #recv returns an error code.
389    def recv_string string, flag = 0
390      message = @receiver_klass.new
391      rc = recvmsg message, flag
392      string.replace(message.copy_out_string) if Util.resultcode_ok?(rc)
393      message.close
394      rc
395    end
396
397    # Receive a multipart message as a list of strings.
398    #
399    # @param [Array] list
400    # @param flag
401    #   One of @0 (default)@ and @XS::NonBlocking@. Any other flag will be
402    #   removed.
403    #
404    # @return 0 if successful
405    # @return -1 if unsuccessful
406    def recv_strings list, flag = 0
407      array = []
408      rc = recvmsgs array, flag
409
410      if Util.resultcode_ok?(rc)
411        array.each do |message|
412          list << message.copy_out_string
413          message.close
414        end
415      end
416
417      rc
418    end
419
420    # Receive a multipart message as an array of objects
421    # (by default these are instances of Message).
422    #
423    # @param [Array] list
424    # @param flag
425    #   One of @0 (default)@ and @XS::NonBlocking@. Any other flag will be
426    #   removed.
427    #
428    # @return 0 if successful
429    # @return -1 if unsuccessful
430    def recvmsgs list, flag = 0
431      flag = NonBlocking if dontwait?(flag)
432
433      message = @receiver_klass.new
434      rc = recvmsg message, flag
435
436      if Util.resultcode_ok?(rc)
437        list << message
438
439        # check rc *first*; necessary because the call to #more_parts? can reset
440        # the xs_errno to a weird value, so the xs_errno that was set on the
441        # call to #recv gets lost
442        while Util.resultcode_ok?(rc) && more_parts?
443          message = @receiver_klass.new
444          rc = recvmsg message, flag
445
446          if Util.resultcode_ok?(rc)
447            list << message
448          else
449            message.close
450            list.each { |msg| msg.close }
451            list.clear
452          end
453        end
454      else
455        message.close
456      end
457
458      rc
459    end
460
461    # Should only be used for XREQ, XREP, DEALER and ROUTER type sockets. Takes
462    # a +list+ for receiving the message body parts and a +routing_envelope+
463    # for receiving the message parts comprising the 0mq routing information.
464    #
465    # @param [Array] list
466    # @param routing_envelope
467    # @param flag
468    #   One of @0 (default)@ and @XS::NonBlocking@
469    #
470    # @return 0 if successful
471    # @return -1 if unsuccessful
472    def recv_multipart list, routing_envelope, flag = 0
473      parts = []
474      rc = recvmsgs parts, flag
475
476      if Util.resultcode_ok?(rc)
477        routing = true
478        parts.each do |part|
479          if routing
480            routing_envelope << part
481            routing = part.size > 0
482          else
483            list << part
484          end
485        end
486      end
487
488      rc
489    end
490
491
492    private
493
494    # Gets socket option
495    #
496    # @param name
497    # @param array
498    #
499    # @return option number
500    def __getsockopt__ name, array
501      # a small optimization so we only have to determine the option
502      # type a single time; gives approx 5% speedup to do it this way.
503      option_type = @option_lookup[name]
504
505      value, length = sockopt_buffers option_type
506
507      rc = LibXS.xs_getsockopt @socket, name, value, length
508
509      if Util.resultcode_ok?(rc)
510        array[0] = if 1 == option_type
511          value.read_long_long
512        elsif 0 == option_type
513          value.read_int
514        elsif 2 == option_type
515          value.read_string(length.read_int)
516        end
517      end
518
519      rc
520    end
521
522    # Calls to xs_getsockopt require us to pass in some pointers. We can cache and save those buffers
523    # for subsequent calls. This is a big perf win for calling RCVMORE which happens quite often.
524    # Cannot save the buffer for the IDENTITY.
525    #
526    # @param option_type
527    # @return cached number or string
528    def sockopt_buffers option_type
529      if 1 == option_type
530        # int64_t or uint64_t
531        unless @longlong_cache
532          length = FFI::MemoryPointer.new :size_t
533          length.write_int 8
534          @longlong_cache = [FFI::MemoryPointer.new(:int64), length]
535        end
536
537        @longlong_cache
538
539      elsif 0 == option_type
540        # int, Crossroads assumes int is 4-bytes
541        unless @int_cache
542          length = FFI::MemoryPointer.new :size_t
543          length.write_int 4
544          @int_cache = [FFI::MemoryPointer.new(:int32), length]
545        end
546
547        @int_cache
548
549      elsif 2 == option_type
550        length = FFI::MemoryPointer.new :size_t
551        # could be a string of up to 255 bytes
552        length.write_int 255
553        [FFI::MemoryPointer.new(255), length]
554
555      else
556        # uh oh, someone passed in an unknown option; use a slop buffer
557        unless @int_cache
558          length = FFI::MemoryPointer.new :size_t
559          length.write_int 4
560          @int_cache = [FFI::MemoryPointer.new(:int32), length]
561        end
562
563        @int_cache
564      end
565    end
566
567    # Populate socket option lookup array
568    def populate_option_lookup
569      # integer options
570      [EVENTS, LINGER, RECONNECT_IVL, FD, TYPE, BACKLOG, KEEPALIVE, IPV4ONLY, SURVEY_TIMEOUT].each { |option| @option_lookup[option] = 0 }
571
572      # long long options
573      [RCVMORE, AFFINITY].each { |option| @option_lookup[option] = 1 }
574
575      # string options
576      [SUBSCRIBE, UNSUBSCRIBE].each { |option| @option_lookup[option] = 2 }
577    end
578
579    # Initialize caches
580    def release_cache
581      @longlong_cache = nil
582      @int_cache = nil
583    end
584
585    # Convenience method to decide whether flag is DONTWAIT
586    #
587    # @param flag
588    #
589    # @return true if is DONTWAIT
590    # @return false if not
591    def dontwait?(flag)
592      (NonBlocking & flag) == NonBlocking
593    end
594    alias :noblock? :dontwait?
595  end # module CommonSocketBehavior
596
597
598  module IdentitySupport
599
600    # Convenience method for getting the value of the socket IDENTITY.
601    #
602    # @return identity
603    def identity
604      array = []
605      getsockopt IDENTITY, array
606      array.at(0)
607    end
608
609    # Convenience method for setting the value of the socket IDENTITY.
610    #
611    # @param value
612    def identity=(value)
613      setsockopt IDENTITY, value.to_s
614    end
615
616
617    private
618
619    # Populate option lookup array
620    def populate_option_lookup
621      super()
622
623      # string options
624      [IDENTITY].each { |option| @option_lookup[option] = 2 }
625    end
626
627  end # module IdentitySupport
628
629  class Socket
630    include CommonSocketBehavior
631    include IdentitySupport
632
633    # Get the options set on this socket.
634    #
635    # @param name
636    #   One of @XS::RCVMORE@, @XS::SNDHWM@, @XS::AFFINITY@, @XS::IDENTITY@,
637    #          @XS::RATE@, @XS::RECOVERY_IVL@, @XS::SNDBUF@,
638    #          @XS::RCVBUF@, @XS::FD@, @XS::EVENTS@, @XS::LINGER@,
639    #          @XS::RECONNECT_IVL@, @XS::BACKLOG@, XS::RECONNECT_IVL_MAX@,
640    #          @XS::RCVTIMEO@, @XS::SNDTIMEO@, @XS::IPV4ONLY@, @XS::TYPE@,
641    #          @XS::RCVHWM@, @XS::MAXMSGSIZE@, @XS::MULTICAST_HOPS@,
642    #          @XS::KEEPALIVE@
643    # @param array should be an empty array; a result of the proper type
644    #   (numeric, string, boolean) will be inserted into
645    #   the first position.
646    #
647    # @return 0 when the operation completed successfully
648    # @return -1 when this operation failed
649    #
650    # With a -1 return code, the user must check XS.errno to determine the
651    # cause.
652    #
653    # @example Retrieve send high water mark
654    #   array = []
655    #   rc = socket.getsockopt(XS::SNDHWM, array)
656    #   sndhwm = array.first if XS::Util.resultcode_ok?(rc)
657    def getsockopt name, array
658      rc = __getsockopt__ name, array
659
660      if Util.resultcode_ok?(rc) && (RCVMORE == name)
661        # convert to boolean
662        array[0] = 1 == array[0]
663      end
664
665      rc
666    end
667
668
669    private
670
671    # Queue message to send
672    #
673    # @param socket
674    # @param address
675    # @param flag
676    def __sendmsg__(socket, address, flag)
677      LibXS.xs_sendmsg(socket, address, flag)
678    end
679
680    # Receive message
681    #
682    # @param socket
683    # @param address
684    # @param flag
685    def __recvmsg__(socket, address, flag)
686      LibXS.xs_recvmsg(socket, address, flag)
687    end
688
689    # Populate socket option lookup array
690    def populate_option_lookup
691      super()
692
693      # integer options
694      [RECONNECT_IVL_MAX, RCVHWM, SNDHWM, RATE, RECOVERY_IVL, SNDBUF, RCVBUF].each { |option| @option_lookup[option] = 0 }
695    end
696
697    # these finalizer-related methods cannot live in the CommonSocketBehavior
698    # module; they *must* be in the class definition directly
699    #
700    # Deletes native resources after object has been destroyed
701    def define_finalizer
702      ObjectSpace.define_finalizer(self, self.class.close(@socket))
703    end
704
705    # Removes all finalizers for object
706    def remove_finalizer
707      ObjectSpace.undefine_finalizer self
708    end
709
710    # Closes the socket
711    def self.close socket
712      Proc.new { LibXS.xs_close socket }
713    end
714  end
715
716end # module XS