rpc.rb in riakpb-0.3.0

- old
+ new

@@ -14,11 +14,11 @@
       include Riakpb::Util::Encode
       include Riakpb::Util::Decode
 
       RECV_LIMIT=1073741824
 
-      attr_reader :req_message, :response, :resp_message_codes, :resp_message, :status
+      attr_reader :req_message, :response, :resp_message_codes, :resp_message, :status, :socket
 
       # Establishes a Client ID with the Riakpb node, for the life of the RPC connection.
       # @param [Client] client the Riakpb::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)
@@ -48,29 +48,30 @@
       end
 
       # Opens a TCPSocket connection with the riak host/node
       # @yield [TCPSocket] hands off the socket connection
       # @return [TCPSocket] data that was exchanged with the host/node
-      def with_socket(&block)
-        socket              = TCPSocket.open(@client.host, @client.port)
-        set_client_id(socket) if @set_client_id
+      def with_socket
+        set_client_id unless @set_client_id.nil?
 
-        out                 = yield(socket)
-        socket.close
+        yield(socket)
+      end
 
-        return(out)
+      # @return [TCPSocket] The TCPSocket of the remote riak node
+      def socket
+        @socket ||= TCPSocket.open(@client.host, @client.port)
       end
 
       # Sets the Client ID for the TCPSocket session
       # @param [TCPSocket] socket connection for which the Client ID will be set
       # @return [True/False] whether or not the set client id request succeeded
-      def set_client_id(socket)
+      def set_client_id
         @set_c_id_req     ||= assemble_request( Util::MessageCode::SET_CLIENT_ID_REQUEST,
                                                 @set_client_id.serialize_to_string)
 
         socket.write(@set_c_id_req)
-        set_c_id_resp       = socket.sysread(@limit)
+        set_c_id_resp = socket.sysread(@limit)
 
         resp_code, resp_msg = decode_message(set_c_id_resp)
 
         return(resp_code == Util::MessageCode::SET_CLIENT_ID_RESPONSE)
       end
@@ -109,11 +110,9 @@
       # @return [Protobuf::Message] @response the processed response (if any) from the Riakpb node
       def parse_response(value)
         @resp_message << value
 
         value = @buffer + value
-
-#        return {:done => false} if message_remaining?(@resp_message)
 
         response_chunk, @resp_message_codes, @buffer = decode_message(value)
 
         @resp_message_codes.each do |resp_mc|
           if resp_mc.equal?(ERROR_RESPONSE)