rb/lib/thrift/processor: Refactor the processor to handle properly unary requests

Author: AlexisMontagne

compiler/cpp/src/generate/t_rb_generator.cc

@@ -826,9 +826,9 @@ void t_rb_generator::generate_service(t_service* tservice) {
   f_service_.indent() << "NAMESPACE = '" << tservice->get_program()->get_namespace("*") << "'.freeze" << endl << endl;
 
   // Generate the three main parts of the service (well, two for now in PHP)
+  generate_service_helpers(tservice);
   generate_service_client(tservice);
   generate_service_server(tservice);
-  generate_service_helpers(tservice);
 
   f_service_.indent() << "::Thrift.register_service_type(self)"<< endl << endl;
 
@@ -1019,6 +1019,16 @@ void t_rb_generator::generate_service_server(t_service* tservice) {
   f_service_.indent_down();
   f_service_.indent() << "end" << endl << endl;
 
+  f_service_.indent() << "METHODS = {" << endl;
+  f_service_.indent_up();
+
+  for (f_iter = functions.begin(); f_iter != functions.end(); ++f_iter) {
+  f_service_.indent() << "'" << (*f_iter)->get_name() << "' => { args_klass: " << capitalize((*f_iter)->get_name()) << "_args, oneway: " << ((*f_iter)->is_oneway() ? "true" : "false") << "}," << endl;
+  }
+
+  f_service_.indent_down();
+  f_service_.indent() << "}.freeze" << endl << endl;
+
   // Generate the process subfunctions
   for (f_iter = functions.begin(); f_iter != functions.end(); ++f_iter) {
     generate_process_function(tservice, *f_iter);
@@ -1035,34 +1045,20 @@ void t_rb_generator::generate_service_server(t_service* tservice) {
  */
 void t_rb_generator::generate_process_function(t_service* tservice, t_function* tfunction) {
   (void)tservice;
-  // Open function
-  f_service_.indent() << "def process_" << tfunction->get_name() << "(seqid, iprot, oprot)" << endl;
-
-  f_service_.indent_up();
-
-  string argsname = capitalize(tfunction->get_name()) + "_args";
   string resultname = capitalize(tfunction->get_name()) + "_result";
 
-  f_service_.indent() << "args = read_args(iprot, " << argsname << ")" << endl;
-
-  f_service_.indent();
-
-  if (!tfunction->is_oneway()) {
-    f_service_ << "result = ";
-  }
-  f_service_ << "@middleware.handle_" << (tfunction->is_oneway() ? "unary" : "binary") << "('" << tfunction->get_name() << "', args) do |args|" << endl;
-  f_service_.indent_up();
-
   t_struct* xs = tfunction->get_xceptions();
   const std::vector<t_field*>& xceptions = xs->get_members();
   vector<t_field*>::const_iterator x_iter;
 
-  // Declare result for non oneway function
+  // Open function
+  f_service_.indent() << "def execute_" << tfunction->get_name() << "(args)" << endl;
+  f_service_.indent_up();
+
   if (!tfunction->is_oneway()) {
-    f_service_.indent() << "result = " << resultname << ".new()" << endl;
+    f_service_.indent() << "result = " << resultname << ".new()" << endl << endl;
   }
 
-  // Try block for a function with exceptions
   if (xceptions.size() > 0) {
     f_service_.indent() << "begin" << endl;
     f_service_.indent_up();
@@ -1104,24 +1100,39 @@ void t_rb_generator::generate_process_function(t_service* tservice, t_function*
     f_service_.indent() << "end" << endl;
   }
 
-  // Shortcut out here for oneway functions
+  f_service_.indent() << endl;
+
   if (tfunction->is_oneway()) {
     f_service_.indent() << "nil" << endl;
-    f_service_.indent_down();
+  } else {
+    f_service_.indent() << "result" << endl;
+  }
 
-    f_service_.indent() << "end" << endl;
-    f_service_.indent_down();
+  f_service_.indent_down();
+  f_service_.indent() << "end" << endl << endl;
 
-    f_service_.indent() << "end" << endl << endl;
-    return;
-  }
+  f_service_.indent() << "def process_" << tfunction->get_name() << "(seqid, iprot, oprot)" << endl;
+
+  f_service_.indent_up();
+
+  string argsname = capitalize(tfunction->get_name()) + "_args";
 
+  f_service_.indent() << "args = read_args(iprot, " << argsname << ")" << endl;
 
-  f_service_.indent() << "result" << endl;
+  f_service_.indent();
+
+  if (!tfunction->is_oneway()) {
+    f_service_ << "result = ";
+  }
+  f_service_ << "@middleware.handle_" << (tfunction->is_oneway() ? "unary" : "binary") << "('" << tfunction->get_name() << "', args) do |args|" << endl;
+  f_service_.indent_up();
+  f_service_.indent() << "execute_" << tfunction->get_name() << "(args)" << endl;
   f_service_.indent_down();
   f_service_.indent() << "end" << endl << endl;
-  f_service_.indent() << "write_result(result, oprot, '" << tfunction->get_name() << "', seqid)"
-                      << endl;
+  if (!tfunction->is_oneway()) {
+    f_service_.indent() << "write_result(result, oprot, '" << tfunction->get_name() << "', seqid)"
+                        << endl;
+  }
 
   // Close function
   f_service_.indent_down();

lib/rb/lib/thrift/processor.rb

@@ -19,13 +19,109 @@
 
 module Thrift
   module Processor
+    class BaseProcessorFunction
+      def initialize(fname, middleware, args_klass, method)
+        @fname = fname
+        @middleware = middleware
+        @args_klass = args_klass
+        @method = method
+      end
+
+      protected
+
+      def read_args(iprot)
+        args = @args_klass.new
+
+        args.read(iprot)
+        iprot.read_message_end
+
+        args
+      end
+    end
+
+    class BinaryProcessorFunction < BaseProcessorFunction
+      def process(seqid, iprot, oprot)
+        execute(seqid, iprot, oprot)
+
+        true
+      end
+
+      private
+
+      def execute(seqid, iprot, oprot)
+        args = read_args(iprot)
+
+        result = @middleware.handle_binary(@fname, args) do |args|
+          @method.call(args)
+        end
+
+        write_result(result, oprot, seqid)
+      rescue => e
+        write_exception(e, oprot, seqid)
+      end
+
+      def write_exception(exception, oprot, seqid)
+        oprot.write_message_begin(@fname, MessageTypes::EXCEPTION, seqid)
+
+        unless exception.is_a? ApplicationException
+          exception = ApplicationException.new(
+            ApplicationException::INTERNAL_ERROR,
+            "Internal error processing #{@fname}: #{exception.class}: #{exception}"
+          )
+        end
+
+        exception.write(oprot)
+        oprot.write_message_end
+        oprot.trans.flush
+      end
+
+      def write_result(result, oprot, seqid)
+        oprot.write_message_begin(@fname, MessageTypes::REPLY, seqid)
+        result.write(oprot)
+        oprot.write_message_end
+        oprot.trans.flush
+      end
+    end
+
+    class UnaryProcessorFunction < BaseProcessorFunction
+      def process(_seqid, iprot, _oprot)
+        args = read_args(iprot)
+
+        @middleware.handle_unary(@fname, args) do |args|
+          @method.call(args)
+        end
+
+        true
+      end
+    end
+
     def initialize(handler, middlewares = [])
       @handler = handler
       @middleware = Middleware.wrap(middlewares)
+
+      @functions = if self.class.const_defined?(:METHODS)
+        self.class::METHODS.reduce({}) do |acc, (key, args)|
+          klass = args[:oneway] ? UnaryProcessorFunction : BinaryProcessorFunction
+
+          acc.merge key => klass.new(
+            key,
+            @middleware,
+            args[:args_klass],
+            method("execute_#{key}")
+          )
+        end
+      end || {}
     end
 
     def process(iprot, oprot)
       name, _type, seqid = iprot.read_message_begin
+
+      func = @functions[name]
+
+      return func.process(seqid, iprot, oprot) if func
+
+      # TODO: once all the stubs will be generated w thrift >=2.5 the next lines
+      # can be deleted
       if respond_to?("process_#{name}")
         begin
           send("process_#{name}", seqid, iprot, oprot)