combine more core batch ops

3 files changed, 132 insertions, 69 deletions

diff --git a/src/ruby/lib/grpc/generic/active_call.rb b/src/ruby/lib/grpc/generic/active_call.rb
index dfc2644c46..f5c426ebfc 100644
--- a/src/ruby/lib/grpc/generic/active_call.rb
+++ b/src/ruby/lib/grpc/generic/active_call.rb
@@ -43,7 +43,8 @@ class Struct
         GRPC.logger.debug("Failing with status #{status}")
         # raise BadStatus, propagating the metadata if present.
         md = status.metadata
-        fail GRPC::BadStatus.new(status.code, status.details, md)
+        fail GRPC::BadStatus.new(status.code, status.details, md),
+             "status code: #{status.code}, details: #{status.details}"
       end
       status
     end
@@ -156,41 +157,25 @@ module GRPC
       Operation.new(self)
     end
 
-    # writes_done indicates that all writes are completed.
-    #
-    # It blocks until the remote endpoint acknowledges with at status unless
-    # assert_finished is set to false.  Any calls to #remote_send after this
-    # call will fail.
-    #
-    # @param assert_finished [true, false] when true(default), waits for
-    # FINISHED.
-    def writes_done(assert_finished = true)
-      ops = {
-        SEND_CLOSE_FROM_CLIENT => nil
-      }
-      ops[RECV_STATUS_ON_CLIENT] = nil if assert_finished
-      batch_result = @call.run_batch(ops)
-      return unless assert_finished
-      unless batch_result.status.nil?
-        @call.trailing_metadata = batch_result.status.metadata
-      end
-      @call.status = batch_result.status
-      op_is_done
-      batch_result.check_status
-    end
-
     # finished waits until a client call is completed.
     #
     # It blocks until the remote endpoint acknowledges by sending a status.
     def finished
       batch_result = @call.run_batch(RECV_STATUS_ON_CLIENT => nil)
-      unless batch_result.status.nil?
-        @call.trailing_metadata = batch_result.status.metadata
+      attach_status_results_and_complete_call(batch_result)
+    end
+
+    def attach_status_results_and_complete_call(recv_status_batch_result)
+      unless recv_status_batch_result.status.nil?
+        @call.trailing_metadata = recv_status_batch_result.status.metadata
       end
-      @call.status = batch_result.status
-      op_is_done
-      batch_result.check_status
+      @call.status = recv_status_batch_result.status
       @call.close
+      op_is_done
+
+      # The RECV_STATUS in run_batch always succeeds
+      # Check the status for a bad status or failed run batch
+      recv_status_batch_result.check_status
     end
 
     # remote_send sends a request to the remote endpoint.
@@ -226,6 +211,23 @@ module GRPC
       nil
     end
 
+    def server_unary_response(req, trailing_metadata: {},
+                              code: Core::StatusCodes::OK, details: 'OK')
+      ops = {}
+      @send_initial_md_mutex.synchronize do
+        ops[SEND_INITIAL_METADATA] = @metadata_to_send unless @metadata_sent
+        @metadata_sent = true
+      end
+
+      payload = @marshal.call(req)
+      ops[SEND_MESSAGE] = payload
+      ops[SEND_STATUS_FROM_SERVER] = Struct::Status.new(
+        code, details, trailing_metadata)
+      ops[RECV_CLOSE_ON_SERVER] = nil
+
+      @call.run_batch(ops)
+    end
+
     # remote_read reads a response from the remote endpoint.
     #
     # It blocks until the remote endpoint replies with a message or status.
@@ -240,9 +242,13 @@ module GRPC
         @call.metadata = batch_result.metadata
         @metadata_received = true
       end
-      unless batch_result.nil? || batch_result.message.nil?
-        res = @unmarshal.call(batch_result.message)
-        return res
+      get_message_from_batch_result(batch_result)
+    end
+
+    def get_message_from_batch_result(recv_message_batch_result)
+      unless recv_message_batch_result.nil? ||
+             recv_message_batch_result.message.nil?
+        return @unmarshal.call(recv_message_batch_result.message)
       end
       GRPC.logger.debug('found nil; the final response has been sent')
       nil
@@ -298,7 +304,6 @@ module GRPC
       return enum_for(:each_remote_read_then_finish) unless block_given?
       loop do
         resp = remote_read
-        break if resp.is_a? Struct::Status  # is an OK status
         if resp.nil?  # the last response was received, but not finished yet
           finished
           break
@@ -315,15 +320,25 @@ module GRPC
     # a list, multiple metadata for its key are sent
     # @return [Object] the response received from the server
     def request_response(req, metadata: {})
-      merge_metadata_to_send(metadata) && send_initial_metadata
-      remote_send(req)
-      writes_done(false)
-      response = remote_read
-      finished unless response.is_a? Struct::Status
-      response
-    rescue GRPC::Core::CallError => e
-      finished  # checks for Cancelled
-      raise e
+      ops = {
+        SEND_MESSAGE => @marshal.call(req),
+        SEND_CLOSE_FROM_CLIENT => nil,
+        RECV_INITIAL_METADATA => nil,
+        RECV_MESSAGE => nil,
+        RECV_STATUS_ON_CLIENT => nil
+      }
+      @send_initial_md_mutex.synchronize do
+        # Metadata might have already been sent if this is an operation view
+        unless @metadata_sent
+          ops[SEND_INITIAL_METADATA] = @metadata_to_send.merge!(metadata)
+        end
+        @metadata_sent = true
+      end
+      batch_result = @call.run_batch(ops)
+
+      @call.metadata = batch_result.metadata
+      attach_status_results_and_complete_call(batch_result)
+      get_message_from_batch_result(batch_result)
     end
 
     # client_streamer sends a stream of requests to a GRPC server, and
@@ -339,12 +354,20 @@ module GRPC
     # a list, multiple metadata for its key are sent
     # @return [Object] the response received from the server
     def client_streamer(requests, metadata: {})
-      merge_metadata_to_send(metadata) && send_initial_metadata
-      requests.each { |r| remote_send(r) }
-      writes_done(false)
-      response = remote_read
-      finished unless response.is_a? Struct::Status
-      response
+      # Metadata might have already been sent if this is an operation view
+      merge_metadata_and_send_if_not_already_sent(metadata)
+
+      requests.each { |r| @call.run_batch(SEND_MESSAGE => @marshal.call(r)) }
+      batch_result = @call.run_batch(
+        SEND_CLOSE_FROM_CLIENT => nil,
+        RECV_INITIAL_METADATA => nil,
+        RECV_MESSAGE => nil,
+        RECV_STATUS_ON_CLIENT => nil
+      )
+
+      @call.metadata = batch_result.metadata
+      attach_status_results_and_complete_call(batch_result)
+      get_message_from_batch_result(batch_result)
     rescue GRPC::Core::CallError => e
       finished  # checks for Cancelled
       raise e
@@ -365,9 +388,18 @@ module GRPC
     # a list, multiple metadata for its key are sent
     # @return [Enumerator|nil] a response Enumerator
     def server_streamer(req, metadata: {})
-      merge_metadata_to_send(metadata) && send_initial_metadata
-      remote_send(req)
-      writes_done(false)
+      ops = {
+        SEND_MESSAGE => @marshal.call(req),
+        SEND_CLOSE_FROM_CLIENT => nil
+      }
+      @send_initial_md_mutex.synchronize do
+        # Metadata might have already been sent if this is an operation view
+        unless @metadata_sent
+          ops[SEND_INITIAL_METADATA] = @metadata_to_send.merge!(metadata)
+        end
+        @metadata_sent = true
+      end
+      @call.run_batch(ops)
       replies = enum_for(:each_remote_read_then_finish)
       return replies unless block_given?
       replies.each { |r| yield r }
@@ -404,7 +436,8 @@ module GRPC
     # a list, multiple metadata for its key are sent
     # @return [Enumerator, nil] a response Enumerator
     def bidi_streamer(requests, metadata: {}, &blk)
-      merge_metadata_to_send(metadata) && send_initial_metadata
+      # Metadata might have already been sent if this is an operation view
+      merge_metadata_and_send_if_not_already_sent(metadata)
       bd = BidiCall.new(@call,
                         @marshal,
                         @unmarshal,
@@ -457,6 +490,15 @@ module GRPC
       end
     end
 
+    def merge_metadata_and_send_if_not_already_sent(new_metadata = {})
+      @send_initial_md_mutex.synchronize do
+        return if @metadata_sent
+        @metadata_to_send.merge!(new_metadata)
+        @call.run_batch(SEND_INITIAL_METADATA => @metadata_to_send)
+        @metadata_sent = true
+      end
+    end
+
     private
 
     # Starts the call if not already started
diff --git a/src/ruby/lib/grpc/generic/client_stub.rb b/src/ruby/lib/grpc/generic/client_stub.rb
index 0d7c1f7805..6934257cbc 100644
--- a/src/ruby/lib/grpc/generic/client_stub.rb
+++ b/src/ruby/lib/grpc/generic/client_stub.rb
@@ -168,6 +168,7 @@ module GRPC
 
       # return the operation view of the active_call; define #execute as a
       # new method for this instance that invokes #request_response.
+      c.merge_metadata_to_send(metadata)
       op = c.operation
       op.define_singleton_method(:execute) do
         c.request_response(req, metadata: metadata)
@@ -231,9 +232,10 @@ module GRPC
 
       # return the operation view of the active_call; define #execute as a
       # new method for this instance that invokes #client_streamer.
+      c.merge_metadata_to_send(metadata)
       op = c.operation
       op.define_singleton_method(:execute) do
-        c.client_streamer(requests, metadata: metadata)
+        c.client_streamer(requests)
       end
       op
     end
@@ -309,9 +311,10 @@ module GRPC
 
       # return the operation view of the active_call; define #execute
       # as a new method for this instance that invokes #server_streamer
+      c.merge_metadata_to_send(metadata)
       op = c.operation
       op.define_singleton_method(:execute) do
-        c.server_streamer(req, metadata: metadata, &blk)
+        c.server_streamer(req, &blk)
       end
       op
     end
@@ -417,15 +420,15 @@ module GRPC
                           deadline: deadline,
                           parent: parent,
                           credentials: credentials)
-
       return c.bidi_streamer(requests, metadata: metadata,
                              &blk) unless return_op
 
       # return the operation view of the active_call; define #execute
       # as a new method for this instance that invokes #bidi_streamer
+      c.merge_metadata_to_send(metadata)
       op = c.operation
       op.define_singleton_method(:execute) do
-        c.bidi_streamer(requests, metadata: metadata, &blk)
+        c.bidi_streamer(requests, &blk)
       end
       op
     end
@@ -445,7 +448,6 @@ module GRPC
                         deadline: nil,
                         parent: nil,
                         credentials: nil)
-
       deadline = from_relative_time(@timeout) if deadline.nil?
       # Provide each new client call with its own completion queue
       call = @ch.create_call(parent, # parent call
diff --git a/src/ruby/lib/grpc/generic/rpc_desc.rb b/src/ruby/lib/grpc/generic/rpc_desc.rb
index 584fe78169..cd17aed8e7 100644
--- a/src/ruby/lib/grpc/generic/rpc_desc.rb
+++ b/src/ruby/lib/grpc/generic/rpc_desc.rb
@@ -62,25 +62,44 @@ module GRPC
       proc { |o| unmarshal_class.method(unmarshal_method).call(o) }
     end
 
+    def handle_request_response(active_call, mth)
+      req = active_call.remote_read
+      resp = mth.call(req, active_call.single_req_view)
+      active_call.server_unary_response(
+        resp, trailing_metadata: active_call.output_metadata)
+    end
+
+    def handle_client_streamer(active_call,  mth)
+      resp = mth.call(active_call.multi_req_view)
+      active_call.server_unary_response(
+        resp, trailing_metadata: active_call.output_metadata)
+    end
+
+    def handle_server_streamer(active_call, mth)
+      req = active_call.remote_read
+      replys = mth.call(req, active_call.single_req_view)
+      replys.each { |r| active_call.remote_send(r) }
+      send_status(active_call, OK, 'OK', active_call.output_metadata)
+    end
+
+    def handle_bidi_streamer(active_call, mth)
+      active_call.run_server_bidi(mth)
+      send_status(active_call, OK, 'OK', active_call.output_metadata)
+    end
+
     def run_server_method(active_call, mth)
       # While a server method is running, it might be cancelled, its deadline
       # might be reached, the handler could throw an unknown error, or a
       # well-behaved handler could throw a StatusError.
       if request_response?
-        req = active_call.remote_read
-        resp = mth.call(req, active_call.single_req_view)
-        active_call.remote_send(resp)
+        handle_request_response(active_call, mth)
       elsif client_streamer?
-        resp = mth.call(active_call.multi_req_view)
-        active_call.remote_send(resp)
+        handle_client_streamer(active_call, mth)
       elsif server_streamer?
-        req = active_call.remote_read
-        replys = mth.call(req, active_call.single_req_view)
-        replys.each { |r| active_call.remote_send(r) }
+        handle_server_streamer(active_call, mth)
       else  # is a bidi_stream
-        active_call.run_server_bidi(mth)
+        handle_bidi_streamer(active_call, mth)
       end
-      send_status(active_call, OK, 'OK', active_call.output_metadata)
     rescue BadStatus => e
       # this is raised by handlers that want GRPC to send an application error
       # code and detail message and some additional app-specific metadata.
@@ -91,7 +110,7 @@ module GRPC
       # Log it, but don't notify the other endpoint..
       GRPC.logger.warn("failed call: #{active_call}\n#{e}")
     rescue Core::OutOfTime
-      # This is raised when active_call#method.call exceeeds the deadline
+      # This is raised when active_call#method.call exceeds the deadline
       # event.  Send a status of deadline exceeded
       GRPC.logger.warn("late call: #{active_call}")
       send_status(active_call, DEADLINE_EXCEEDED, 'late')
@@ -100,7 +119,7 @@ module GRPC
       # Send back a UNKNOWN status to the client
       GRPC.logger.warn("failed handler: #{active_call}; sending status:UNKNOWN")
       GRPC.logger.warn(e)
-      send_status(active_call, UNKNOWN, 'no reason given')
+      send_status(active_call, UNKNOWN, 'unkown error handling call on server')
     end
 
     def assert_arity_matches(mth)