rpc.rb in riakpb-0.1.4

- old
+ new

@@ -17,14 +17,16 @@
       RECV_LIMIT=1638400
 
       attr_reader :req_message, :response, :resp_message_codes, :resp_message, :status
 
       # Establishes a Client ID with the Riak node, for the life of the RPC connection.
-      # @param [Client] the Riak::Client object in which this Rpc instance lives
-      def initialize(client)
+      # @param [Client] client the Riak::Client object in which this Rpc instance lives
+      # @param [Fixnum] limit the max size of an individual TCPSocket receive call.  Need to fix, later.
+      def initialize(client, limit=RECV_LIMIT)
         @status             = false
         @client             = client
+        @limit              = limit
         @client_id          = request(Util::MessageCode::GET_CLIENT_ID_REQUEST).client_id
         @set_client_id      = Riak::RpbSetClientIdReq.new(:client_id => @client_id)
 
         # Request / Response Data
         @resp_message_codes = -1
@@ -63,11 +65,11 @@
       def set_client_id(socket)
         @set_c_id_req     ||= assemble_request( Util::MessageCode::SET_CLIENT_ID_REQUEST,
                                                 @set_client_id.serialize_to_string)
 
         socket.send(@set_c_id_req, 0)
-        set_c_id_resp       = socket.recv(RECV_LIMIT)
+        set_c_id_resp       = socket.recv(@limit)
 
         resp_code, resp_msg = decode_message(set_c_id_resp)
 
         return(resp_code == Util::MessageCode::SET_CLIENT_ID_RESPONSE)
       end
@@ -82,13 +84,17 @@
         @req_message_code = mc
         @response         = RESPONSE_CLASS_FOR[mc].new unless RESPONSE_CLASS_FOR[mc].nil?
 
         with_socket do |socket|
           begin
-            @req_message  = assemble_request mc, (pb_msg.serialize_to_string rescue '')
+            begin
+              @req_message  = assemble_request mc, pb_msg.serialize_to_string
+            rescue NoMethodError
+              @req_message  = assemble_request mc
+            end
 
             socket.send(@req_message, 0)
-            self.response = socket.recv(RECV_LIMIT)
+            self.response = socket.recv(@limit)
 
           end while(false == (@response.done rescue true))
         end # with_socket
 
         return(@response)