backport partitions creation (#355)

Author: mensfeld

CHANGELOG.md

@@ -1,6 +1,8 @@
 # Rdkafka Changelog
 
 ## 0.15.0 (Unreleased)
+- **[Feature]** Add `Admin#metadata` (mensfeld)
+- **[Feature]** Add `Admin#create_partitions` (mensfeld)
 - **[Feature]** Add `Admin#delete_group` utility (piotaixr)
 - **[Feature]** Add Create and Delete ACL Feature To Admin Functions (vgnanasekaran)
 - [Enhancement] Bump librdkafka to 2.3.0 (mensfeld)

lib/rdkafka.rb

@@ -15,6 +15,8 @@
 require "rdkafka/admin/delete_groups_report"
 require "rdkafka/admin/delete_topic_handle"
 require "rdkafka/admin/delete_topic_report"
+require "rdkafka/admin/create_partitions_handle"
+require "rdkafka/admin/create_partitions_report"
 require "rdkafka/admin/create_acl_handle"
 require "rdkafka/admin/create_acl_report"
 require "rdkafka/admin/delete_acl_handle"

lib/rdkafka/admin.rb

@@ -14,6 +14,19 @@ def finalizer
       ->(_) { close }
     end
 
+    # Performs the metadata request using admin
+    #
+    # @param topic_name [String, nil] metadat about particular topic or all if nil
+    # @param timeout_ms [Integer] metadata request timeout
+    # @return [Metadata] requested metadata
+    def metadata(topic_name = nil, timeout_ms = 2_000)
+      closed_admin_check(__method__)
+
+      @native_kafka.with_inner do |inner|
+        Metadata.new(inner, topic_name, timeout_ms)
+      end
+    end
+
     # Close this admin instance
     def close
       return if closed?
@@ -216,6 +229,71 @@ def delete_topic(topic_name)
       delete_topic_handle
     end
 
+    # Creates extra partitions for a given topic
+    #
+    # @param topic_name [String]
+    # @param partition_count [Integer] how many partitions we want to end up with for given topic
+    #
+    # @raise [ConfigError] When the partition count or replication factor are out of valid range
+    # @raise [RdkafkaError] When the topic name is invalid or the topic already exists
+    # @raise [RdkafkaError] When the topic configuration is invalid
+    #
+    # @return [CreateTopicHandle] Create topic handle that can be used to wait for the result of creating the topic
+    def create_partitions(topic_name, partition_count)
+      closed_admin_check(__method__)
+
+      @native_kafka.with_inner do |inner|
+        error_buffer = FFI::MemoryPointer.from_string(" " * 256)
+        new_partitions_ptr = Rdkafka::Bindings.rd_kafka_NewPartitions_new(
+          FFI::MemoryPointer.from_string(topic_name),
+          partition_count,
+          error_buffer,
+          256
+        )
+        if new_partitions_ptr.null?
+          raise Rdkafka::Config::ConfigError.new(error_buffer.read_string)
+        end
+
+        pointer_array = [new_partitions_ptr]
+        topics_array_ptr = FFI::MemoryPointer.new(:pointer)
+        topics_array_ptr.write_array_of_pointer(pointer_array)
+
+        # Get a pointer to the queue that our request will be enqueued on
+        queue_ptr = Rdkafka::Bindings.rd_kafka_queue_get_background(inner)
+        if queue_ptr.null?
+          Rdkafka::Bindings.rd_kafka_NewPartitions_destroy(new_partitions_ptr)
+          raise Rdkafka::Config::ConfigError.new("rd_kafka_queue_get_background was NULL")
+        end
+
+        # Create and register the handle we will return to the caller
+        create_partitions_handle = CreatePartitionsHandle.new
+        create_partitions_handle[:pending] = true
+        create_partitions_handle[:response] = -1
+        CreatePartitionsHandle.register(create_partitions_handle)
+        admin_options_ptr = Rdkafka::Bindings.rd_kafka_AdminOptions_new(inner, Rdkafka::Bindings::RD_KAFKA_ADMIN_OP_CREATEPARTITIONS)
+        Rdkafka::Bindings.rd_kafka_AdminOptions_set_opaque(admin_options_ptr, create_partitions_handle.to_ptr)
+
+        begin
+          Rdkafka::Bindings.rd_kafka_CreatePartitions(
+            inner,
+            topics_array_ptr,
+            1,
+            admin_options_ptr,
+            queue_ptr
+          )
+        rescue Exception
+          CreatePartitionsHandle.remove(create_partitions_handle.to_ptr.address)
+          raise
+        ensure
+          Rdkafka::Bindings.rd_kafka_AdminOptions_destroy(admin_options_ptr)
+          Rdkafka::Bindings.rd_kafka_queue_destroy(queue_ptr)
+          Rdkafka::Bindings.rd_kafka_NewPartitions_destroy(new_partitions_ptr)
+        end
+
+        create_partitions_handle
+      end
+    end
+
     # Create acl
     # @param resource_type - values of type rd_kafka_ResourceType_t
     #        https://github.com/confluentinc/librdkafka/blob/292d2a66b9921b783f08147807992e603c7af059/src/rdkafka.h#L7307
@@ -528,6 +606,7 @@ def describe_acl(resource_type:, resource_name:, resource_pattern_type:, princip
     end
 
     private
+
     def closed_admin_check(method)
       raise Rdkafka::ClosedAdminError.new(method) if closed?
     end

lib/rdkafka/admin/create_partitions_handle.rb

@@ -0,0 +1,27 @@
+module Rdkafka
+  class Admin
+    class CreatePartitionsHandle < AbstractHandle
+      layout :pending, :bool,
+             :response, :int,
+             :error_string, :pointer,
+             :result_name, :pointer
+
+      # @return [String] the name of the operation
+      def operation_name
+        "create partitions"
+      end
+
+      # @return [Boolean] whether the create topic was successful
+      def create_result
+        CreatePartitionsReport.new(self[:error_string], self[:result_name])
+      end
+
+      def raise_error
+        raise RdkafkaError.new(
+            self[:response],
+            broker_message: CreateTopicReport.new(self[:error_string], self[:result_name]).error_string
+        )
+      end
+    end
+  end
+end

lib/rdkafka/admin/create_partitions_report.rb

@@ -0,0 +1,6 @@
+module Rdkafka
+  class Admin
+    class CreatePartitionsReport < CreateTopicReport
+    end
+  end
+end

lib/rdkafka/bindings.rb

@@ -302,6 +302,16 @@ def self.partitioner(str, partition_count, partitioner_name = "consistent_random
     attach_function :rd_kafka_event_DeleteTopics_result, [:pointer], :pointer, blocking: true
     attach_function :rd_kafka_DeleteTopics_result_topics, [:pointer, :pointer], :pointer, blocking: true
 
+    # Create partitions
+    RD_KAFKA_ADMIN_OP_CREATEPARTITIONS = 3
+    RD_KAFKA_ADMIN_OP_CREATEPARTITIONS_RESULT = 102
+
+    attach_function :rd_kafka_CreatePartitions, [:pointer, :pointer, :size_t, :pointer, :pointer], :void
+    attach_function :rd_kafka_NewPartitions_new, %i[pointer size_t pointer size_t], :pointer
+    attach_function :rd_kafka_NewPartitions_destroy, [:pointer], :void
+    attach_function :rd_kafka_event_CreatePartitions_result, [:pointer], :pointer
+    attach_function :rd_kafka_CreatePartitions_result_topics, [:pointer, :pointer], :pointer
+
     # Delete Group
 
     RD_KAFKA_ADMIN_OP_DELETEGROUPS = 7   # rd_kafka_admin_op_t

lib/rdkafka/callbacks.rb

@@ -128,6 +128,8 @@ def self.call(_, event_ptr, _)
           process_create_topic(event_ptr)
         elsif event_type == Rdkafka::Bindings::RD_KAFKA_EVENT_DELETETOPICS_RESULT
           process_delete_topic(event_ptr)
+        elsif event_type == Rdkafka::Bindings::RD_KAFKA_ADMIN_OP_CREATEPARTITIONS_RESULT
+          process_create_partitions(event_ptr)
         elsif event_type == Rdkafka::Bindings::RD_KAFKA_EVENT_CREATEACLS_RESULT
           process_create_acl(event_ptr)
         elsif event_type == Rdkafka::Bindings::RD_KAFKA_EVENT_DELETEACLS_RESULT
@@ -192,6 +194,23 @@ def self.process_delete_topic(event_ptr)
         end
       end
 
+      def self.process_create_partitions(event_ptr)
+        create_partitionss_result = Rdkafka::Bindings.rd_kafka_event_CreatePartitions_result(event_ptr)
+
+        # Get the number of create topic results
+        pointer_to_size_t = FFI::MemoryPointer.new(:int32)
+        create_partitions_result_array = Rdkafka::Bindings.rd_kafka_CreatePartitions_result_topics(create_partitionss_result, pointer_to_size_t)
+        create_partitions_results = TopicResult.create_topic_results_from_array(pointer_to_size_t.read_int, create_partitions_result_array)
+        create_partitions_handle_ptr = Rdkafka::Bindings.rd_kafka_event_opaque(event_ptr)
+
+        if create_partitions_handle = Rdkafka::Admin::CreatePartitionsHandle.remove(create_partitions_handle_ptr.address)
+          create_partitions_handle[:response] = create_partitions_results[0].result_error
+          create_partitions_handle[:error_string] = create_partitions_results[0].error_string
+          create_partitions_handle[:result_name] = create_partitions_results[0].result_name
+          create_partitions_handle[:pending] = false
+        end
+      end
+
       def self.process_create_acl(event_ptr)
         create_acls_result = Rdkafka::Bindings.rd_kafka_event_CreateAcls_result(event_ptr)
 

spec/rdkafka/admin_spec.rb

@@ -9,6 +9,7 @@
   after do
     # Registry should always end up being empty
     expect(Rdkafka::Admin::CreateTopicHandle::REGISTRY).to be_empty
+    expect(Rdkafka::Admin::CreatePartitionsHandle::REGISTRY).to be_empty
     expect(Rdkafka::Admin::DescribeAclHandle::REGISTRY).to be_empty
     expect(Rdkafka::Admin::CreateAclHandle::REGISTRY).to be_empty
     expect(Rdkafka::Admin::DeleteAclHandle::REGISTRY).to be_empty
@@ -366,4 +367,41 @@
 
     end
   end
+
+  describe '#create_partitions' do
+    let(:metadata) { admin.metadata(topic_name).topics.first }
+
+    context 'when topic does not exist' do
+      it 'expect to fail due to unknown partition' do
+        expect { admin.create_partitions(topic_name, 10).wait }.to raise_error(Rdkafka::RdkafkaError, /unknown_topic_or_part/)
+      end
+    end
+
+    context 'when topic already has the desired number of partitions' do
+      before { admin.create_topic(topic_name, 2, 1).wait }
+
+      it 'expect not to change number of partitions' do
+        expect { admin.create_partitions(topic_name, 2).wait }.to raise_error(Rdkafka::RdkafkaError, /invalid_partitions/)
+        expect(metadata[:partition_count]).to eq(2)
+      end
+    end
+
+    context 'when topic has more than the requested number of partitions' do
+      before { admin.create_topic(topic_name, 5, 1).wait }
+
+      it 'expect not to change number of partitions' do
+        expect { admin.create_partitions(topic_name, 2).wait }.to raise_error(Rdkafka::RdkafkaError, /invalid_partitions/)
+        expect(metadata[:partition_count]).to eq(5)
+      end
+    end
+
+    context 'when topic has less then desired number of partitions' do
+      before { admin.create_topic(topic_name, 1, 1).wait }
+
+      it 'expect to change number of partitions' do
+        admin.create_partitions(topic_name, 10).wait
+        expect(metadata[:partition_count]).to eq(10)
+      end
+    end
+  end
 end

spec/spec_helper.rb

@@ -11,6 +11,7 @@
 require "rspec"
 require "rdkafka"
 require "timeout"
+require "securerandom"
 
 def rdkafka_base_config
   {