upfluence-thrift 2.6.1 → 2.7.4

data/lib/thrift/middleware.rb

@@ -8,6 +8,18 @@ module Thrift
       def handle_unary(_mth, args = {},  &block)
         block.call(args)
       end
+
+      def handle_bidi_stream(_mth, args, istream, ostream, &block)
+        block.call(args, istream, ostream)
+      end
+
+      def handle_inbound_stream(_mth, args, istream, &block)
+        block.call(args, istream)
+      end
+
+      def handle_outbound_stream(_mth, args, ostream, &block)
+        block.call(args, ostream)
+      end
     end
 
     class MultiMiddleware
@@ -30,6 +42,36 @@ module Thrift
           end
         end.call(args)
       end
+
+      def handle_bidi_stream(_mth, args, istream, ostream, &block)
+        @middlewares.reverse.reduce(block) do |acc, m|
+          Proc.new do |args, istream, ostream|
+            m.handle_bidi_stream(mth, args, istream, ostream) do |args, istream, ostream|
+              acc.call(args, istream, ostream)
+            end
+          end
+        end.call(args, istream, ostream)
+      end
+
+      def handle_inbound_stream(_mth, args, istream, &block)
+        @middlewares.reverse.reduce(block) do |acc, m|
+          Proc.new do |args, istream|
+            m.handle_bidi_stream(mth, args, istream) do |args, istream|
+              acc.call(args, istream)
+            end
+          end
+        end.call(args, istream)
+      end
+
+      def handle_outbound_stream(_mth, args, ostream, &block)
+        @middlewares.reverse.reduce(block) do |acc, m|
+          Proc.new do |args, ostream|
+            m.handle_bidi_stream(mth, args, ostream) do |args, ostream|
+              acc.call(args, ostream)
+            end
+          end
+        end.call(args, ostream)
+      end
     end
 
     NOP_MIDDLEWARE = NopMiddleware.new

data/lib/thrift/outbound_stream.rb

@@ -0,0 +1,52 @@
+module Thrift
+  class TOutboundStream < TBaseStream
+    def initialize(iprot, oprot, name, seqid, klass, message_type, unlock = nil)
+      @message_type = message_type
+      @klass = klass
+      @unlock = unlock
+      super(iprot, oprot, name, seqid)
+    end
+
+    def ready
+      Thread.new { read_goaway }
+      super
+    end
+
+    def close
+      wait_ready
+
+      write_shell @message_type + 1
+
+      @mutex.synchronize do
+        loop do
+          @cond.wait(@mutex)
+
+          return nil if @closed
+        end
+      end
+    end
+
+    def send(arg)
+      raise EOFError if @closed
+
+      unless @oprot.trans.open?
+        mark_as_closed
+        raise EOFError
+      end
+
+      wait_ready
+      write @message_type, @klass.new(arg: arg)
+    end
+
+    private
+
+    def read_goaway
+      mtype = read_shell
+
+      write_shell @message_type + 2 if mtype == (@message_type + 1)
+    ensure
+      mark_as_closed
+      @unlock&.call
+    end
+  end
+end

data/lib/thrift/processor.rb

@@ -17,6 +17,8 @@
 # under the License.
 #
 
+require 'logger'
+
 module Thrift
   module Processor
     class BaseProcessorFunction
@@ -95,83 +97,321 @@ module Thrift
       end
     end
 
-    def initialize(handler, middlewares = [])
+    def initialize(handler, middlewares = [], logger=nil)
       @handler = handler
+      if logger.nil?
+        @logger = Logger.new(STDERR)
+        @logger.level = Logger::WARN
+      else
+        @logger = logger
+      end
       @middleware = Middleware.wrap(middlewares)
 
-      @functions = if self.class.const_defined?(:METHODS)
-        self.class::METHODS.reduce({}) do |acc, (key, args)|
-          klass = args[:oneway] ? UnaryProcessorFunction : BinaryProcessorFunction
+      @processors = if self.class.const_defined? :METHODS
+                      self.class::METHODS.reduce({}) do |acc, (name, info)|
+                        acc.merge(name => build_processor(name, info))
+                      end
+                    else
+                      {}
+                    end
+    end
 
-          acc.merge key => klass.new(
-            key,
-            @middleware,
-            args[:args_klass],
-            method("execute_#{key}")
-          )
-        end
-      end || {}
+    def read_args(iprot, args_class)
+      args = args_class.new
+      args.read(iprot)
+      iprot.read_message_end
+      args
+    end
+
+    def write_result(result, oprot, name, seqid)
+      oprot.write_message_begin(name, MessageTypes::REPLY, seqid)
+      result.write(oprot)
+      oprot.write_message_end
+      oprot.trans.flush
     end
 
     def process(iprot, oprot)
       name, _type, seqid = iprot.read_message_begin
 
-      func = @functions[name]
+      mth = "process_#{name}"
+      if respond_to?(mth)
+        send(mth, seqid, iprot, oprot)
+        return true
+      end
 
-      return func.process(seqid, iprot, oprot) if func
+      (
+        @processors[name] || UnkwonFunctionProcessor.new(name)
+      ).process(seqid, iprot, oprot)
+    end
 
-      # 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)
-        rescue => e
-          write_exception(e, oprot, name, seqid)
-        end
-        true
+    def build_processor(name, info)
+      if info[:result_klass].nil?
+        UnaryProcessor
+      elsif info[:stream_klass].nil? && info[:sink_klass].nil?
+        BinaryProcessor
+      elsif info[:sink_klass].nil?
+        OutboundStreamProcessor
+      elsif info[:stream_klass].nil?
+        InboundStreamProcessor
       else
-        iprot.skip(Types::STRUCT)
+        BidiStreamProcessor
+      end.new(name, info, @middleware, @handler)
+    end
+
+    class BaseProcessor
+      def initialize(name, args_class)
+        @name = name
+        @args_class = args_class
+      end
+
+      protected
+
+      def read_args(iprot)
+        args = @args_class.new
+        args.read(iprot)
         iprot.read_message_end
+        args
+      end
+
+      def write_exception(exception, oprot, name, seqid)
+        oprot.write_message_begin(name, MessageTypes::EXCEPTION, seqid)
+
+        unless exception.is_a? ApplicationException
+          exception = ApplicationException.new(
+            ApplicationException::INTERNAL_ERROR,
+            "Internal error processing #{name}: #{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(@name, MessageTypes::REPLY, seqid)
+        result.write(oprot)
+        oprot.write_message_end
+        oprot.trans.flush
+      end
+    end
+
+    class UnaryProcessor < BaseProcessor
+      def initialize(name, info, middleware, handler)
+        @middleware = middleware
+        @handler = handler
+        @arg_keys = info[:args]
+
+        super name, info[:args_klass]
+      end
+
+      def process(_seqid, iprot, _oprot)
+        @middleware.handle_unary(@name, read_args(iprot)) do |args|
+          @handler.send(@name, *@arg_keys.map { |k| args.send k })
+          nil
+        end
+
+        true
+      end
+    end
+
+    class BinaryProcessor < BaseProcessor
+      def initialize(name, info, middleware, handler)
+        @middleware = middleware
+        @handler = handler
+        @arg_keys = info[:args]
+        @void_result = info[:void_result]
+        @result_klass = info[:result_klass]
+        @exceptions = info[:exceptions]
+
+        super name, info[:args_klass]
+      end
+
+      def process(seqid, iprot, oprot)
+        res = @middleware.handle_binary(@name, read_args(iprot)) do |args|
+          execute(args)
+        end
+
+        write_result(res, oprot, seqid)
+
+        true
+      rescue ApplicationException => e
+        write_exception(e, oprot, @name, seqid)
+        true
+      rescue StandardError => e
         write_exception(
           ApplicationException.new(
-            ApplicationException::UNKNOWN_METHOD,
-            'Unknown function ' + name,
+            ApplicationException::INTERNAL_ERROR,
+            "Internal error processing #{@name}: #{e.class}: #{e}"
           ),
           oprot,
-          name,
+          @name,
           seqid
         )
-        false
+        true
+      end
+
+      protected
+
+      def execute(args, *extra_args)
+        res = @result_klass.new
+
+        begin
+          s = @handler.send(
+            @name,
+            *@arg_keys.map { |k| args.send k },
+            *extra_args
+          )
+
+          res.success = s unless @void_result
+        rescue => e
+          k, = @exceptions.find { |(_, klass)| e.is_a? klass }
+
+          raise e unless k
+
+          res.send "#{k}=", e
+        end
+
+        res
       end
     end
 
-    def read_args(iprot, args_class)
-      args = args_class.new
-      args.read(iprot)
-      iprot.read_message_end
-      args
+    class BaseStreamProcessor < BinaryProcessor
+      def initialize(name, info, middleware, handler)
+        @mutex = Mutex.new
+        @cond = ConditionVariable.new
+        @closed = false
+
+        super name, info, middleware, handler
+      end
+
+      protected
+
+      def close
+        @mutex.synchronize do
+          @closed = true
+          @cond.broadcast
+        end
+      end
+
+      def wait
+        @mutex.synchronize do
+          @cond.wait(@mutex) unless @closed
+        end
+      end
     end
 
-    def write_exception(exception, oprot, name, seqid)
-      oprot.write_message_begin(name, MessageTypes::EXCEPTION, seqid)
+    class OutboundStreamProcessor < BaseStreamProcessor
+      def initialize(name, info, middleware, handler)
+        @stream_klass = info[:stream_klass]
+
+        super name, info, middleware, handler
+      end
 
-      unless exception.is_a? ApplicationException
-        exception = ApplicationException.new(
-          ApplicationException::INTERNAL_ERROR,
-          "Internal error processing #{name}: #{exception.class}: #{exception}"
+      def process(seqid, iprot, oprot)
+        @closed = false
+        stream = TOutboundStream.new(
+          iprot, oprot, @name, seqid, @stream_klass,
+          MessageTypes::SERVER_STREAM_MESSAGE,
+          method(:close)
         )
+        res = @middleware.handle_outbound_stream(
+          @name, read_args(iprot), stream
+        ) { |args, stream| execute(args, stream) }
+
+        write_result(res, oprot, seqid)
+        stream.ready
+
+        wait
+
+        true
       end
+    end
 
-      exception.write(oprot)
-      oprot.write_message_end
-      oprot.trans.flush
+    class InboundStreamProcessor < BaseStreamProcessor
+      def initialize(name, info, middleware, handler)
+        @sink_klass = info[:sink_klass]
+
+        super name, info, middleware, handler
+      end
+
+      def process(seqid, iprot, oprot)
+        @closed = false
+        stream = TInboundStream.new(
+          iprot, oprot, @name, seqid, @sink_klass,
+          MessageTypes::CLIENT_STREAM_MESSAGE,
+          method(:close)
+        )
+
+        res = @middleware.handle_inbound_stream(
+          @name, read_args(iprot), stream
+        ) { |args, sink| execute(args, sink) }
+
+
+        write_result(res, oprot, seqid)
+        stream.ready
+
+        wait
+
+        true
+      end
     end
 
-    def write_result(result, oprot, name, seqid)
-      oprot.write_message_begin(name, MessageTypes::REPLY, seqid)
-      result.write(oprot)
-      oprot.write_message_end
-      oprot.trans.flush
+    class BidiStreamProcessor < BaseStreamProcessor
+      def initialize(name, info, middleware, handler)
+        @sink_klass = info[:sink_klass]
+        @stream_klass = info[:stream_klass]
+
+        super name, info, middleware, handler
+      end
+
+      def process(seqid, iprot, oprot)
+        @closed = false
+
+        bidi_stream = TBidiStream.new(
+          iprot, oprot, @name, seqid, @sink_klass, @stream_klass,
+          MessageTypes::CLIENT_STREAM_MESSAGE,
+          MessageTypes::SERVER_STREAM_MESSAGE,
+          method(:close)
+        )
+
+        res = @middleware.handle_bidi_stream(
+          @name, read_args(iprot),
+          TBidiInboundStream.new(bidi_stream),
+          TBidiOutboundStream.new(bidi_stream)
+        ) do |args, sink, stream|
+          execute(args, stream, sink)
+        end
+
+        write_result(res, oprot, seqid)
+        bidi_stream.ready
+
+        wait
+
+        true
+      end
+    end
+
+    class UnkwonFunctionProcessor < BaseProcessor
+      def initialize(name)
+        super name, nil
+      end
+
+      def process(seqid, iprot, oprot)
+        iprot.skip(Types::STRUCT)
+        iprot.read_message_end
+        write_exception(
+          ApplicationException.new(
+            ApplicationException::UNKNOWN_METHOD,
+            'Unknown function ' + @name,
+          ),
+          oprot,
+          @name,
+          seqid
+        )
+
+        false
+      end
     end
   end
 end

data/lib/thrift/protocol/base_protocol.rb

@@ -1,4 +1,4 @@
-# 
+#
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements. See the NOTICE file
 # distributed with this work for additional information
@@ -6,23 +6,21 @@
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License. You may obtain a copy of the License at
-# 
+#
 #   http://www.apache.org/licenses/LICENSE-2.0
-# 
+#
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# 
+#
 
 # this require is to make generated struct definitions happy
-require 'set'
 
 module Thrift
   class ProtocolException < Exception
-
     UNKNOWN = 0
     INVALID_DATA = 1
     NEGATIVE_SIZE = 2
@@ -33,14 +31,13 @@ module Thrift
 
     attr_reader :type
 
-    def initialize(type=UNKNOWN, message=nil)
+    def initialize(type = UNKNOWN, message = nil)
       super(message)
       @type = type
     end
   end
 
   class BaseProtocol
-
     attr_reader :trans
 
     def initialize(trans)
@@ -48,7 +45,7 @@ module Thrift
     end
 
     def native?
-      puts "wrong method is being called!"
+      puts 'wrong method is being called!'
       false
     end
 
@@ -56,19 +53,19 @@ module Thrift
       raise NotImplementedError
     end
 
-    def write_message_end; nil; end
+    def write_message_end = nil
 
     def write_struct_begin(name)
       raise NotImplementedError
     end
 
-    def write_struct_end; nil; end
+    def write_struct_end = nil
 
     def write_field_begin(name, type, id)
       raise NotImplementedError
     end
 
-    def write_field_end; nil; end
+    def write_field_end = nil
 
     def write_field_stop
       raise NotImplementedError
@@ -78,19 +75,19 @@ module Thrift
       raise NotImplementedError
     end
 
-    def write_map_end; nil; end
+    def write_map_end = nil
 
     def write_list_begin(etype, size)
       raise NotImplementedError
     end
 
-    def write_list_end; nil; end
+    def write_list_end = nil
 
     def write_set_begin(etype, size)
       raise NotImplementedError
     end
 
-    def write_set_end; nil; end
+    def write_set_end = nil
 
     def write_bool(bool)
       raise NotImplementedError
@@ -141,37 +138,37 @@ module Thrift
       raise NotImplementedError
     end
 
-    def read_message_end; nil; end
+    def read_message_end = nil
 
     def read_struct_begin
       raise NotImplementedError
     end
 
-    def read_struct_end; nil; end
+    def read_struct_end = nil
 
     def read_field_begin
       raise NotImplementedError
     end
 
-    def read_field_end; nil; end
+    def read_field_end = nil
 
     def read_map_begin
       raise NotImplementedError
     end
 
-    def read_map_end; nil; end
+    def read_map_end = nil
 
     def read_list_begin
       raise NotImplementedError
     end
 
-    def read_list_end; nil; end
+    def read_list_end = nil
 
     def read_set_begin
       raise NotImplementedError
     end
 
-    def read_set_end; nil; end
+    def read_set_end = nil
 
     def read_bool
       raise NotImplementedError
@@ -230,7 +227,7 @@ module Thrift
         value = args[2]
       elsif args.size == 4
         # handles the deprecated method signature - write_field(name, type, fid, value)
-        field_info = {:name => args[0], :type => args[1]}
+        field_info = { name: args[0], type: args[1] }
         fid = args[2]
         value = args[3]
       else
@@ -253,9 +250,7 @@ module Thrift
     def write_type(field_info, value)
       # if field_info is a Fixnum, assume it is a Thrift::Types constant
       # convert it into a field_info Hash for backwards compatibility
-      if field_info.is_a? Fixnum
-        field_info = {:type => field_info}
-      end
+      field_info = { type: field_info } if field_info.is_a? Integer
 
       case field_info[:type]
       when Types::BOOL
@@ -291,11 +286,9 @@ module Thrift
     #
     # Returns the value read; object type varies based on field_info[:type].
     def read_type(field_info)
-      # if field_info is a Fixnum, assume it is a Thrift::Types constant
+      # if field_info is an Integer, assume it is a Thrift::Types constant
       # convert it into a field_info Hash for backwards compatibility
-      if field_info.is_a? Fixnum
-        field_info = {:type => field_info}
-      end
+      field_info = { type: field_info } if field_info.is_a? Integer
 
       case field_info[:type]
       when Types::BOOL
@@ -323,8 +316,6 @@ module Thrift
 
     def skip(type)
       case type
-      when Types::STOP
-        nil
       when Types::BOOL
         read_bool
       when Types::BYTE
@@ -342,8 +333,9 @@ module Thrift
       when Types::STRUCT
         read_struct_begin
         while true
-          name, type, id = read_field_begin
+          _, type, = read_field_begin
           break if type == Types::STOP
+
           skip(type)
           read_field_end
         end
@@ -367,13 +359,23 @@ module Thrift
           skip(etype)
         end
         read_list_end
+      else
+        raise ProtocolException.new(ProtocolException::INVALID_DATA, 'Invalid data')
       end
     end
+
+    def to_s
+      "#{trans}"
+    end
   end
 
   class BaseProtocolFactory
     def get_protocol(trans)
       raise NotImplementedError
     end
+
+    def to_s
+      'base'
+    end
   end
-end
+end